Skip to content

Characterize NKI kahan butterfly precision on silicon #58

Description

@scttfrdmn

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

  1. 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.
  2. 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.
  3. 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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    nkiNKI kernel development (Trainium hardware)phase-2-precisionRoadmap: phase-2-precision

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions