v0.11.0 landed butterfly_stage_kernel_kahan (Dekker 2Prod compensated complex multiply in the NKI butterfly kernel). Compile + output-match is verified on trn1 (NKI 2.24.5133.0) via test_kahan_butterfly_compiles_and_matches_fast, but the core question — does the kahan compensator actually improve FP32 FFT accuracy on Trainium? — is unmeasured.
On CPU, "kahan" mode is equivalent to "fast" because the compensation is only applied at the chirp multiplies (which are O(N), not the error hotspot); the dominant error source is the 3-FFT butterfly chain (O(N log N)). On NKI the butterfly is now compensated, but whether that translates into measurable FP32 gains on silicon hasn't been characterized.
What's needed
- Precision sweep.
trnfft.fft at power-of-2 N ∈ {1024, 4096, 16384, 65536} with fast vs kahan, comparing against a numpy FP64 reference. Measure max-abs-err and rel-err distributions over ≥ 100 random inputs per N.
- Adversarial sweep. Inputs chosen to maximize catastrophic cancellation in the butterfly complex multiply (e.g., twiddle angles + input magnitudes that make
t_re*o_re ≈ t_im*o_im). These are where kahan should show the largest benefit.
- Bluestein end-to-end. Same sweep for arbitrary-N sizes that trigger Bluestein. Benefit should compound through the 3-FFT chain if kahan works.
Decision criteria
- If kahan ≥ 2× better rel-err on typical inputs: promote to a documented recommendation for precision-sensitive NKI workloads. Consider default on specific size thresholds.
- If kahan ≥ 5× better only on adversarial inputs: document as "available for known-difficult inputs"; no default change.
- If kahan < 2× across both sweeps: document honestly as "kernel is present but the win is modest on current silicon" and deprioritize. Still useful on future hardware with higher-precision intermediates.
Cost
Gated on a trn1 or trn2 provision + benchmark run. ~1-2 hours of silicon time plus analysis.
Follow-up from #52.
v0.11.0 landed
butterfly_stage_kernel_kahan(Dekker 2Prod compensated complex multiply in the NKI butterfly kernel). Compile + output-match is verified on trn1 (NKI 2.24.5133.0) viatest_kahan_butterfly_compiles_and_matches_fast, but the core question — does the kahan compensator actually improve FP32 FFT accuracy on Trainium? — is unmeasured.On CPU,
"kahan"mode is equivalent to"fast"because the compensation is only applied at the chirp multiplies (which are O(N), not the error hotspot); the dominant error source is the 3-FFT butterfly chain (O(N log N)). On NKI the butterfly is now compensated, but whether that translates into measurable FP32 gains on silicon hasn't been characterized.What's needed
trnfft.fftat power-of-2 N ∈ {1024, 4096, 16384, 65536} with fast vs kahan, comparing against a numpy FP64 reference. Measure max-abs-err and rel-err distributions over ≥ 100 random inputs per N.t_re*o_re ≈ t_im*o_im). These are where kahan should show the largest benefit.Decision criteria
Cost
Gated on a trn1 or trn2 provision + benchmark run. ~1-2 hours of silicon time plus analysis.
Follow-up from #52.