[katec846]

Summary

Add an SM90 FP8 MMA forward kernel for the head-dimension-512 attention layer used by Gemma-4. Both the paged KV cache and QK/PV MMA use FP8.

The kernel uses TMA loads for FP8 K, overlaps V transposition with QK computation, and promotes PV partial sums to FP32 after each KV block.

On H200 with CuTe DSL 4.8.0, the kernel achieves 1.078–1.155× speedup over FP8 KV with FP16 MMA. Against native FP16, it achieves 1.020–1.223× speedup across the tested shapes except batch 1 / 4K, where speedup is 0.960×.

Thanks to @simon-veitner-redhat for the SM90 FP8 MMA performance, numerical, and synchronization fixes included in this PR

Performance

H200 decode microbenchmark with FP16 Q, query length 1, 8 Q heads, 4 KV heads, head dimension 512, page size 64, and num_splits=1. Results are median kernel latency after 200 warmup launches, measured over seven windows of 100 launches each. Lower is better. Using cutedsl==4.8.0

Batch	KV length	FP8 KV + FP8 MMA	FP8 KV + FP16 MMA (PR #164)	FP16 KV + FP16 MMA	Speedup vs PR #164	Speedup vs native FP16
1	4K	0.1992 ms	0.2147 ms	0.1911 ms	1.078×	0.960×
1	16K	0.7703 ms	0.8779 ms	0.7859 ms	1.140×	1.020×
1	28K	1.3397 ms	1.5270 ms	1.3739 ms	1.140×	1.026×
16	4K	0.2020 ms	0.2316 ms	0.2076 ms	1.147×	1.028×
16	16K	0.7725 ms	0.8914 ms	0.7966 ms	1.154×	1.031×
16	28K	1.3428 ms	1.5509 ms	1.3866 ms	1.155×	1.033×
64	4K	0.4343 ms	0.4797 ms	0.5166 ms	1.105×	1.189×
64	16K	1.7179 ms	1.8701 ms	2.0655 ms	1.089×	1.202×
64	28K	3.0736 ms	3.3133 ms	3.7605 ms	1.078×	1.223×

would be nice to get a review from simon

[jhaotingc]

Thanks @LucasWilkinson, I'll share with Simon when some of the accuracy scores are validated thoroughly. Thanks a lot!!

[simon-veitner-redhat]

Thanks for this PR. I reviewed it with GPU runs on an H200 and put the fixes I suggest on a branch you can cherry-pick: https://github.com/simon-veitner-redhat/flash-attention/tree/fa199-review-fixes (7 signed-off commits on 8dbd390).

The main issue: FP8 WGMMA accumulates below FP32 (PTX ISA 9.7.16.5.2), and the PV sum of a whole split stays in one WGMMA accumulator. With q = 0 and V = 1 the output is 0.500 at 32K keys and 0.125 at 128K with num_splits=1. On real Gemma 4 activations at 16K, the relative error is 5.5e-2 (layer 5) and 2.3e-1 (layer 59). Commit 3 promotes PV to FP32 every KV block and adds a test for the long-split case.

Commits 1 and 2 give bitwise-identical output and make the kernel faster. The whole branch runs at 0.981x of your head (geomean, causal). Commit 1 also closes a cross-warpgroup WAR on sK that today only an injected delay triggers. Commit 7 fences the LDSM reads of sStage before releasing it to the next V TMA (PTX ISA 8.6, 9.7.9.26.2, 9.7.14.4); the release-then-fence order comes from the PR 164 dequant path and 8dbd390 kept it. It does not change the output. Picked alone, it fixes the V path only, because the K copy on 8dbd390 has the same order until commit 1 removes that copy.

after commit	time vs 8dbd390	test_flash_attn_fp8_kv_cache.py failing	Gemma 4 16K error, layer 5 / 59
(8dbd390)	1	10 of 14	5.5e-2 / 2.3e-1
1 K in two sK slots	0.909x	10	unchanged (bitwise)
2 V transpose between QK halves	0.896x	10	unchanged (bitwise)
3 PV promoted to FP32 per block	0.958x	10	3.5e-3 / 2.1e-3
4 residual Q term first	0.958x	4	3.0e-3 / 1.7e-3
5 split QK sum (causal, local)	0.982x	0	2.0e-3 / 1.7e-3
6 comments and docs	0.982x	0	2.0e-3 / 1.7e-3
7 fence sStage before release	0.981x	0	2.0e-3 / 1.7e-3

Time is the geomean over 14 vLLM-style shapes (decode b1 to b256 at 4K to 64K, prefill 2K and 8K), all causal, with CUDA graphs and a flushed L2. The error is relative L2 against float64 on the same e4m3 bytes. Commits 1, 3, 4, 6 and 7 apply alone; 2 needs 1; 5 needs 1 to 4. Each commit message has its numbers.

Things to know when picking:

• The error and the 14 of 14 in the table are for causal and local builds, which is what vLLM runs (it always passes causal=True). Commit 5 splits the QK sum only there, because in non-causal builds ptxas serializes the WGMMAs (C7511). Non-causal builds keep the commit-4 accuracy: with the 14 tests switched to causal=False, 4 still fail at the branch head (10 on 8dbd390).
• Commit 3 alone costs 1.101x, because the K copy addresses then spill. After commits 1 and 2 it costs nothing over your head.
• On commits 1 to 3, one fp16-Q case of test_flash_attn_kvcache_fp8_dequant_sm90 misses its max-error bound by one fp16 ulp (3.17e-3 against 2.94e-3). Your head already reaches 2.93e-3 on the num_splits=2 variant of the same case. The error comes from QK: commit 4 halves it, and from commit 4 on all 39 FP8 tests in test_flash_attn.py pass.

What each commit changes

1. Load FP8 K with TMA into a two-slot sK

On 8dbd390, K is loaded by TMA into sStage and then copied by all 256 threads into sK behind a named barrier. The two buffers have the same layout, so the copy only moves bytes, and the next block's copy can overwrite sK while the other warpgroup's QK WGMMA still reads it. I load K with TMA straight into sK, which now has two slots with their own TMA pipeline and per-slot mbarriers; each warpgroup waits for its slot and releases it after its own QK wait_group. sStage then holds only V, and a trace-time assert rejects persistent and CLC schedulers because the slot parity restarts per tile. It applies alone, the output is bitwise identical, and it gives most of the speedup, at the cost of shared memory growing to about 225 of 227 KiB.

2. Transpose FP8 V between the two QK halves

A cycle trace showed the V[n-1] transpose, issued after all 32 QK QGMMAs, never overlapped them. I issue it through a new mid() hook between the two Q terms in _gemm_two_fp8_a, inside the same commit group, so it runs while the first 16 QGMMAs execute. _gemm_qk_fp8_compensated now names its operands first and second, which commit 4 uses. It needs commit 1, the output is bitwise identical, and it is a small further speedup.

3. Promote FP8 PV sums to FP32 every KV block

This is the accuracy fix from the top of the comment: the PV sum of a whole split stayed in one WGMMA accumulator, which sums below FP32. _gemm_pv_fp8_promoted replaces _gemm_pv_fp8_compensated and runs each block's PV as four 64-column passes into a zero-initialized 32-register scratch, and _pv_fold adds the scratch into acc_O with FP32 adds. The q = 0, V = 1 output becomes 1.000 at every length, and a new test in tests/cute/test_flash_attn.py checks 32K and 128K. It applies alone but costs about 10% there because the K copy addresses spill; after commits 1 and 2 it costs nothing over your head.

4. Issue the residual Q term first in FP8 QK

Each QGMMA truncates its addends relative to the running sum, so the 16 residual Q k-steps lost most of their bits when they were added after the high terms. I pass the residual Q fragment as the first operand, so it is summed while the accumulator is still small. It is a two-line change with the same instructions and registers and no time change. It applies alone, brings test_flash_attn_fp8_kv_cache.py from 10 to 4 failures, and fixes the fp16-Q case that commit 3 misses by one ulp.

5. Split the FP8 QK sum in causal and local builds

The 4 remaining failures (logits with std about 22) come from summing all 32 QGMMAs of QK into one accumulator. _gemm_qk_fp8_split sums k-steps 0 to 7 into acc_S and 8 to 15 into the PV scratch, which is idle during QK, and _fold_qk_split adds the scratch into acc_S in FP32 before the first PV pass; the first block uses its own 32-register temporary, and both keep the residual-first order. It is built only for causal and local kernels, because in non-causal builds ptxas serializes the WGMMAs (C7511) and they ran 1.33 to 2.03x slower. It needs commits 1 to 4, makes all 14 tests pass, and costs about 2.5% over commit 4.

6. Document FP8 MMA accumulation and fp16 Q limits

The comments in flash_fwd_sm90.py and interface.py said the FP8 path accumulates in FP32. I state that FP8 WGMMA sums below FP32 within an instruction chain and that only the per-block PV fold adds in FP32, that fp16 Q keeps 8 of its 11 significant bits in the E4M3 pair, and the real range of the large-Q test. Comments only; it applies alone.

7. Fence sStage reads before releasing it to TMA

The V transpose reads sStage with LDSM (generic proxy) and the next V tile is written by TMA (async proxy), but the proxy fence ran after the release, and only lane 0 of each warp arrives on the empty barrier. I issue fence.proxy.async and a __syncwarp before the elected sStage release; the fence still comes before the per-warpgroup barrier that makes the transposed sV visible to the PV WGMMA. I saw no hang or corruption without it, so this closes a gap in the memory model rather than a reproduced bug. It applies alone (V path only, as noted above), with bitwise-identical output and no time change.

I used AI assistance (Claude) for this review and checked the results myself.

[katec846]

Thanks for the detailed review! I’ve cherry-picked all seven fixes and updated the description with the rerun microbenchmark results.

[simon-veitner-redhat]

Please adjust according to below comment.

[simon-veitner-redhat]

Requesting changes: please integrate the three fixes below before merging, F1 (accuracy) above all. Everything else looks good. I ran this on an H200 at 6d71c3f.

These pass at 6d71c3f:

• FA FP8 suite: 39 passed, 2 skipped (SM100-only).
• hdim512 tests: 51 passed.
• Generic kvcache suite: 2496 passed, 576 skipped.
• vLLM's call shapes (varlen Q, strided paged K/V, 8 or 16 q heads per KV head, MTP decode, prefix prefill, num_splits 0, 1 and 32): all 18 cases pass, FP8 rel_l2 at most 1.9e-3.
• GSM8K through vllm#60372, TP2, 1319 questions: 0.9697 (0.9719 with the merge-base FA).
• The BF16 path is unchanged: O is bitwise identical to the merge base in all 11 BF16 and FP16 cases.

Three things need fixing. I put a fix for each on one branch, ready to cherry-pick:
simon-veitner-redhat/flash-attention fa199-review-fixes-final, three commits on 6d71c3f.

git fetch https://github.com/simon-veitner-redhat/flash-attention.git fa199-review-fixes-final
git cherry-pick 383a9fdd35fe 0431b106963c 8a783b9c84df

With all three on 6d71c3f, on an H200: the FP8 suite gives 41 passed, 2 skipped; the hdim512 tests 51 passed; the new tail test 2 passed.

F1. Small softmax probabilities never reach PV.
_prepare_fp8_probabilities stores P as e4m3(256p) + e4m3(256p - hi). Both terms round to 0 for p below 2^-18 of the running row max, but row_sum still counts those keys in FP32, so the output is pulled toward 0. KV blocks run last to first, so a key that dominates late in the sequence cuts the earlier context out. vLLM's prefill chunks resolve to num_splits=1, which is the worst case.

Measured at 6d71c3f:

• Synthetic, 131K keys, one late key ahead by Δ: Δ=12 gives +24%, Δ=13 gives -100% (merge base: 0.17% and 0.02%).
• Gemma-4 31B FP8, 121K-token prefill, 10 global layers: per-row error p99 0.86-2.97% (merge base 0.26-0.31%), worst row 7.9%.
• End to end I could not see it: 80-prompt multi-value recall at 66K and 121K gives 0.923 against 0.925 on the merge base; GSM8K 0.970 against 0.972.

Fix: 383a9fd. It scales the residual by 32 (|residual| <= 8, so 256 < 448) and runs it as a second PV pass through the same FP32 scratch. Keys reach PV down to 2^-23 of the row max. It brings the Gemma-4 p99 to 0.29-0.42% and the worst row to 0.98%. It adds test_flash_attn_kvcache_fp8_mma_keeps_softmax_tail_sm90, which fails at 6d71c3f (+24%, -100%) and passes with the fix.

The fix costs 6-13% kernel time in my runs (prefill 8K x 123K +7%, 2K x 33K +13%, decode b64 x 33K +10%, b1 x 123K +6%), so feel free to schedule the second pass more cheaply, but the accuracy fix itself should go into this PR. Two other fixes did not work: raising fp8_p_scale to 448 still fails the new test (+29% and +93%), and a separate residual accumulator was 18-32% slower.

F2. Q above 448 loses accuracy.
Q is split into two E4M3 terms at a fixed scale of 1.0, with no range check. With bf16 Q and 1024 keys, the relative error is about 0.2% up to |q| = 448, 1.9% at 768, 9.1% at 1024 and 58% at 2048. The merge base stays at 0.16%. Gemma-4 is far from this (max |q| is 22 over a whole 121K request), but vLLM routes every SM90 FP8-KV d512 model here. A runtime check would need a device sync, so the fix documents the range on fp8_kv_dequant in interface.py.

Fix: 0431b10, docstring only.

F4. Rows with no visible key now return LSE -inf when a sink is set.
The empty-split override in mma is no longer gated on is_split_kv. A non-split launch with a learnable sink now returns LSE = -inf instead of the sink logit for rows with no visible key. This also applies to BF16 (d512 and d128). O is unchanged, and nothing in vLLM reads this LSE today.

Fix: 8a783b9. It gates the row_sum fill on self.is_split_kv or self.fp8_kv_dequant. With it, LSE and O equal the merge base in all 11 BF16 cases. The kvcache suite still passes (2496 passed, 576 skipped). FP8 non-split rows still get -inf, because FP8 consumes a dummy K/V block there.

This review was prepared with Claude Code and reviewed by a human (@simon-veitner-redhat).

[katec846]

Thanks! Cherry-picked and pushed all three fixes. H200 regression tests passed, and the full TP2 GSM8K rerun scored 96.89% (1,278/1,319), with zero invalid answers.

[simon-veitner-redhat]

Thanks! Cherry-picked and pushed all three fixes. H200 regression tests passed, and the full TP2 GSM8K rerun scored 96.89% (1,278/1,319), with zero invalid answers.

thanks @katec846 could you remeasure E2E perf?