[CuTe, SM100] Fix FP8 e4m3 accuracy: make max_offset dtype-aware to avoid P saturation

[yunweili3]

Proposed fix #2716.

Problem

On the SM100 FP8 forward, inputs for e4m3 consistently produced higher error than e5m2 (up to 1.6× worse rel_l2, growing with seqlen), while FP8 quantization theory — and the reporter's CPU emulation in #2716 — predicts e4m3 should actually be around 2× better.

Root cause

Two FP8 constants interact:

• max_offset = 8: P is computed as exp2(scale·s − scale·row_max + 8), which scales scaling probabilities into (0, 256] before the f32→fp8 convert.
• rescale_threshold = 4.0: when a new KV block raises the row max by less than 4 (log2-scaled units), the kernel will keep the stale row max to skip an O-rescale.

With a stale max, P can reach 2^(8+4) = 4096. The f32→fp8 conversion is satfinite: e5m2 (max 57344) absorbs the overshoot, but e4m3fn (max 448) clamps the largest attention weights — precisely the entries that dominate each row — by up to 9×.

More KV blocks → more skip events → error grows with seqlen.

A secondary effect: row_sum accumulates the unclamped f32 P while the P·V MMA consumes the clamped fp8 P, so normalization is also inconsistent.

Fix

Keep the rescale-skip (it's a perf optimization and is innocent once the budget is respected) and enforce max_offset + rescale_threshold ≤ log2(fp8_max) per dtype: max_offset = 4 for e4m3fn (worst case 2^8 = 256 ≤ 448), unchanged 8 for e5m2.

Both definitions (softmax warp and correction warp, with max_offset_scale = 2^offset) would be updated consistently, and the LSE formula is self-consistent w.r.t. the offset. The only cost is P's flush-to-zero floor moving from max·2^-17 to max·2^-13 — negligible softmax tail mass.

A second commit un-rots the fp8 dtype path in test_flash_attn_output (it set requires_grad on fp8 tensors, now rejected by the forward-only interface, and generated descales that only the reference applied — flash_attn_func has no descale kwargs). fp8 stays out of the default parametrize; this just makes the sweep runnable.

Evidence (1× B200, CUDA 13.0, cutlass-dsl 4.6.0.dev0, base 77aacb6)

The following is the reporter's geometry (b=2, h=32, hkv=2, d=128, GQA 16:1, per-(b,hkv) amax scaling, rel_l2 vs unquantized-fp32 reference, via _flash_attn_fwd with descales):

config	e4m3 before	e4m3 after	e5m2 (unchanged)	quantization-only emulation (e4m3)
s256 uniform	0.1148	0.0515	0.1020	0.0518
s1024 uniform	0.1637	0.0529	0.1046	0.0528
s4096 uniform	0.1704	0.0532	0.1064	0.0537
s1024 peaked	0.3184	0.1208	0.2354	0.1193
s4096 peaked	0.3361	0.1281	0.2506	0.1260

After the fix, e4m3 lands exactly on the quantization-error floor and is approximately 2 times better than e5m2, as expected. LSE max error ≤ 1.1e-4 both formats.

Test suite (420-case fp8-e4m3 test_flash_attn_output sweep: seqlens 64/128 … 4096/4096, d 64–256, causal × sink × mha/gqa/mqa, pack_gqa × num_splits inner sweep):

branch	result
main (77aacb6)	190 failed / 188 passed / 42 skipped (all tolerance asserts)
this PR	378 passed / 0 failed / 42 skipped

bf16 unaffected (dtype-gated change): the CI FA4_TEST_FILTER 8-case suite passes.

Perf parity (hd128, s4096, b2 h32, median of 200): e4m3 0.387 ms (fix) vs 0.390 ms (main); e5m2 0.386 ms both — the fix is a different constant in an existing FMA.

Notes

• The FP8 MLA path (flash_fwd_mla_sm100.py) is unaffected, since it uses rescale_threshold=0 and no max offset.
• An alternative fix — rescale_threshold=0 for FP8 — recovers identical accuracy but costs the skip optimization, meaning the measured numbers were equivalent.