[Cute,Fwd,Sm90] Native FP8 (E4M3 Q/K/V) forward for Hopper by jhaotingc · Pull Request #206 · vllm-project/flash-attention

jhaotingc · 2026-10-22T15:20:25Z

Summary

Adds an SM90 FA4 forward path that keeps Q, K and V in E4M3 and runs both GEMMs on FP8 WGMMA, writing fp16 or bf16 output. It is the FA4 counterpart of FA3's Hopper FP8 forward, so FP8 callers no longer need FA3 on Hopper.

On H200, over ten vLLM decode shapes (Gemma-4, Llama-3.1-8B/70B, Qwen3-32B) under CUDA graphs, it is 1.65x faster than FA3 FP8 and 1.97x faster than FA4 BF16 (geomean). Its output error matches FA3 FP8.

Every change is gated on native_fp8 or generates the same code for the other paths, so BF16/FP16 and other architectures compile exactly as before.

Design

V transpose in the producer warpgroup. As in FA3's Transpose_V. FP8 WGMMA needs K-major operands and TMA cannot transpose, so V is loaded row-major into a second smem buffer and producer warps transpose it into the K-major sV. They use LdMatrix8x8x16bOp(transpose=True) and StMatrix8x8x16bOp tiled copies with a byte permute in between (FA3's ldmatrix.trans + __byte_perm + stmatrix), overlapped with the MMAs of the previous block.
Two-level PV accumulation. As in FA3, vllm-project/flash-attention Add two-level accumulation for SM90 FP8 FWD to mitigate long-context accuracy drops #104.
- FP8 WGMMA accumulates at reduced precision, so each PV GEMM covers one KV tile from zero. It is then folded into fp32 O as O = O * scale + O_tile.
- Without this, output with nonzero-mean V shrinks with KV length: 10% at 8k keys, 71% at 128k.
- Cost: 1-3% on d128 decode, 5-7% on d256 decode, 3-11% on prefill.
P in FP8. Scores are cast to E4M3 with a 2**8 bias, which keeps mantissa bits (as in FA3 and FA4 SM100). As in FA3, the FP32 scores are permuted in registers (permute_Cregs_fp8) before the cast, so the FP8 P is already in the PV WGMMA's A-register layout (convert_layout_acc_Aregs_fp8).
Descales. The q/k descales fold into the softmax scale and the v descale into the final normalization. Each is 0-d or (batch, kv_heads).
Paged K/V. When the page size is not tile_n, each tile loads as page-aligned TMA boxes instead of cp.async.
Decode launches.
- Short decode (query tokens x q heads per KV head <= 64) uses 64-row tiles.
- Single-warpgroup tiles run 2 CTAs/SM when the grid exceeds one wave.
- Decode uses a direct (head, split, batch) grid.
- Launch choices depend only on host-known values, so they hold under vLLM's CUDA-graph capture (max_seqlen_k = max_model_len).
- Tiles, RS PV, PackGQA for paged non-TMA K/V, and the split caps are chosen in config.py's select_fwd_config. The kernel keeps its own register split, since its producer warpgroup also transposes V.

Supported:

(d, dv) in {64, 96, 128, 192, 256} x same, and (192, 128);
causal, sliding window, varlen, paged K/V, PackGQA, SplitKV, learnable sink, fp16/bf16 output.

Rejected with an error: E5M2, mixed dtypes, Qv/MLA, block sparsity.

Not covered yet: the tensormap-bounded V loads that main uses for non-paged varlen K/V.

Commits

Add native FP8 (E4M3 Q/K/V) forward
Load paged K/V as page-sized TMA boxes
Decode tiles, fused exp2 bias
Two-level PV accumulation (+ long-KV test that fails without it)
Two CTAs per SM for single-warpgroup tiles
Direct (head, split, batch) grid for decode
Test the decode launch paths
Address review: FA3's FP8 register helpers and V transpose in CuTeDSL, the base O-store copy and one softmax exp2 for all dtypes, simpler interface.py helpers

Performance

H200, clocks locked, all kernels in one session. CUDA graphs replay vLLM's decode launches: paged KV, vLLM's split arguments, max_seqlen_k = 32768, and a 64 MB L2 flush between layers. Each time is the attention span per layer, including the split-KV combine.

Decode step (requests x query tokens, context)	FA3 FP8 (us)	FA4 BF16 (us)	This PR (us)	vs FA3 FP8	vs FA4 BF16
Gemma-4-31B TP4 (d256, 1024 window), 4 x 4, 28k	23.6	16.4	11.4	2.07x	1.44x
Gemma-4-31B TP4, 16 x 4, 28k	44.1	29.6	17.0	2.59x	1.74x
Llama-3.1-8B TP1, 1 x 1, 4k	17.0	14.7	11.4	1.49x	1.29x
Llama-3.1-8B TP1, 32 x 1, 2k	62.6	81.6	36.2	1.73x	2.25x
Llama-3.1-8B TP1, 128 x 1, 1k	123.1	205.2	74.8	1.65x	2.74x
Llama-3.1-8B TP1, 16 x 1, 16k	189.4	303.6	143.4	1.32x	2.12x
Llama-3.1-8B TP1, 32 x 4, 4k	106.1	162.5	66.1	1.61x	2.46x
Llama-3.1-70B TP4, 64 x 1, 4k	58.5	79.9	40.4	1.45x	1.98x
Llama-3.1-70B TP4, 8 x 1, 32k	57.9	78.6	39.4	1.47x	1.99x
Qwen3-32B TP2, 32 x 1, 8k	102.8	154.5	70.6	1.46x	2.19x
Geomean				1.65x	1.97x

Long-KV accuracy. Error is relative to the fp64 reference RMS, with mean-0.5 V, one split.

KV length	1k	8k	32k	128k
FA3 FP8	0.17%	0.15%	0.13%	0.12%
Without two-level accumulation	0.77%	10.4%	34.7%	71.0%
This PR	0.17%	0.15%	0.13%	0.12%

Tests

tests/cute/test_flash_attn.py and the FP8 test files on H200, after rebasing on main:

Suite	Result
`-k fp8` (native FP8 and the SM90 FP8 MMA tests)	117 passed, 12 skipped
`test_flash_attn_kvcache` FP8, full	2504 passed, 576 skipped
`test_flash_attn_fp8_kv_cache.py`, `test_flash_attn_fp8_output.py`, `test_flash_attn_race_condition.py`	651 passed
SM90 regression set	56 passed
`-k "fp8_mma or fp8_dequant"`	41 passed, as on `main`

Native FP8 outputs are bit-identical to this branch before the rebase (30 shapes), and BF16/FP16 and head-dim-512 FP8 MMA outputs are bit-identical to main (21 shapes).

Also checked before the rebase:

compute-sanitizer memcheck/racecheck/synccheck are clean on decode, SplitKV, prefill and mixed launches.
96 extra paged prefill shapes against an fp64 reference produced no wrong outputs.
There are no register spills in any native FP8 kernel.

Based on Simon Veitner's fa3-to-fa4-dense-mla-fp8 branch (co-author on the first commit).

FA4's SM90 native FP8 forward (vllm-project/flash-attention#206) runs attention on E4M3 Q/K/V with FP8 WGMMA. Use it for FP8 KV cache on SM90. - get_flash_attn_version upgrades FA3 to FA4 on SM90 when the KV cache is FP8 E4M3 and FA4 supports the head size (64, 96, 128, 192 or 256, same for V), without softcap or batch invariance. An explicitly requested flash_attn_version is kept. The FlashAttention metadata builder and impl pass fp8_kv_cache, and the builder takes aot_schedule from the layer's FlashAttention version instead of the global one. - Gemma4Config uses FA4 for all layers by default, including SM90 with FP8 KV cache, so its d256 sliding-window layers no longer run on FA3. - SM90 FA4 supports FP8 KV for equal head sizes 64, 96, 128, 192, 256 and 512. - Only d512 keeps BF16/FP16 Q for the FP8-KV dequant path. Smaller heads quantize Q, as FA3 does, and run the native FP8 forward. The FlashAttention pin is updated once vllm-project/flash-attention#206 is merged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

FA4's SM90 native FP8 forward (vllm-project/flash-attention#206) runs attention on E4M3 Q/K/V with FP8 WGMMA. Use it for FP8 KV cache on SM90. - get_flash_attn_version upgrades FA3 to FA4 on SM90 when the KV cache is FP8 E4M3 and FA4 supports the head size (64, 96, 128, 192 or 256, same for V), without softcap or batch invariance. An explicitly requested flash_attn_version is kept. The FlashAttention metadata builder and impl pass fp8_kv_cache, and the builder takes aot_schedule from the layer's FlashAttention version instead of the global one. - Gemma4Config uses FA4 for all layers by default, including SM90 with FP8 KV cache, so its d256 sliding-window layers no longer run on FA3. - SM90 FA4 supports FP8 KV for equal head sizes 64, 96, 128, 192, 256 and 512. - Only d512 keeps BF16/FP16 Q for the FP8-KV dequant path. Smaller heads quantize Q, as FA3 does, and run the native FP8 forward. - The FlashAttention FP8 KV checks use KVQuantMode.FP8_PER_TENSOR instead of matching the cache dtype string, as the metadata builder only has the KV cache spec. The FlashAttention pin is updated once vllm-project/flash-attention#206 is merged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

Add an SM90 FA4 forward path that keeps Q, K and V in E4M3 and runs QK and PV on FP8 WGMMA, writing an fp16 or bf16 output. It is the FA4 counterpart of FA3's Hopper FP8 forward, so FP8 callers no longer need FA3 on Hopper. Kernel (flash_fwd_sm90.py). Every change is gated on native_fp8, so the other dtypes compile as before: - FP8 PV needs a K-major V operand. As in FA3 (Transpose_V), the producer warpgroup loads V row-major into sVLoad (TMA or cp.async), transposes it into the K-major sV with LDSM.T/STSM, and hands it to the MMA warpgroups through a second pipeline, so the transpose overlaps the MMAs. With TMA K/V one producer warp only issues loads while the other three transpose. Because the producer does this extra work, native FP8 sets its own register split instead of the configured one. - Scores are converted to E4M3 in the FP8 A-register layout with a 2**8 bias so the cast keeps mantissa bits. The transpose writes V^T in the score accumulator's K order, so this takes in-register byte permutes only. Softmax.finalize removes the bias from the LSE. - The q and k descales fold into the softmax scale and the v descale into the final normalization. Each descale may be 0-d or (batch, kv_heads), contiguous or broadcast, and is read in place without a copy. - O is written in fp16/bf16, so sQ storage grows to hold it. Rescale-O-before- GEMM is disabled; intra-warpgroup overlap with paged non-TMA K/V is kept only for hdim_v <= 128. Host: - interface.py: E4M3 Q/K/V on SM90 selects native FP8 for (d, dv) in {64, 96, 128, 192, 256} and (192, 128). E5M2, mixed dtypes and block sparsity are rejected. New out_dtype argument. The descale layout classes join the compile key. - config.py: FA3's Hopper FP8 tile choices, RS PV, PackGQA for paged non-TMA K/V, and static splits capped to the N tiles and SM occupancy. Not included: per-request (dynamic) split counts and the persistent varlen schedulers, compile-only (AOT) FP8 specs, MLA and context parallelism, descale arguments on the public flash_attn_func API, and the tensormap-bounded V loads of non-paged varlen K/V. Tests (tests/cute/test_flash_attn.py): - test_flash_attn_output and test_flash_attn_kvcache also run in FP8 on SM90, using their existing FP8 error bar. The output test now builds its learnable sink in the reference dtype (torch.randn cannot produce FP8) and, as for q/k/v, without requires_grad when the inputs are FP8 (forward-only). - test_flash_attn_kvcache_fp8_qkv_sm90: paged decode with 4 query tokens per request, d128/d256, MHA/GQA, 16- and 64-token pages (cp.async and TMA K/V), causal with and without a left window, and SplitKV. - test_flash_attn_fp8_qkv_descales_sm90: scalar, contiguous and broadcast descales, each of q/k/v omitted, and fp16/bf16 output. Based on Simon Veitner's fa3-to-fa4-dense-mla-fp8 branch (e322bf5, b3d5343). Co-authored-by: Simon Veitner <sveitner@redhat.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

Paged K/V whose page size differs from tile_n (e.g. vLLM's 16-token pages) used per-thread cp.async, and the address generation kept the producer warpgroup busy next to the V transpose. Load each KV tile with TMA instead, as tile_n / box_n boxes of box_n = gcd(page_size, tile_n) rows: lane b of the load warp looks up the page of box b, and the boxes are issued from lane-shuffled indices. The box size joins the compile key. Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

- Short decode (query tokens x q heads per KV head <= 64) uses 64-row tiles (one MMA warpgroup) for every head dim; 128-row decode tiles were mostly padding. - The max_offset bias is folded into the exp2 argument as one FFMA. Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

FP8 WGMMA accumulates at reduced precision, so accumulating P @ V over many KV tiles in the tensor-core accumulator truncates each tile's contribution and shrinks the output as the KV range grows (by about 10% at 8k keys and 70% at 128k keys in one split, with nonzero-mean V). As in FA3's FP8 forward, each PV GEMM now accumulates one KV tile from zero and the result is added to the fp32 O together with the rescale (O = O * scale + O_tile, one FFMA). The second O accumulator does not fit next to the next tile's scores, so native FP8 runs the non-overlapped mainloop (intra-warpgroup overlap was worth about 1% here), and hdim_v 256 runs the PV GEMM as two halves through one half-size accumulator. No register spills. Adds a long-KV test with nonzero-mean V, which fails without this change. Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

When a grid exceeds one wave and a 2-stage pipeline of a 64-row tile fits in half of the SM's shared memory, run two CTAs per SM (min_blocks_per_mp=2, 200/56 MMA/producer registers) so one CTA's loads overlap the other's math. With the split heuristic (num_splits=0), splits also fill both CTA slots. Adds the native FP8 shared-memory model used by the later launch rules. Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

When every batch has exactly one M block (decode), launch a (head, split, batch) grid and read the tile coordinates from blockIdx instead of searching the cumulative M-block counts in global memory. The first tile's seqlens are read at kernel entry so their latency overlaps the barrier setup. Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

Decode with vLLM-style arguments (varlen Q, 16-token pages, per-request KV lengths, default split heuristic) over shapes that reach one and two CTAs per SM and the direct (head, split, batch) grid, checked against an fp32 reference with nonzero-mean V. Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

FP8 register helpers and V transpose: - Port the FP8 branch of FA3's convert_layout_acc_Aregs and FA3's permute_Cregs_fp8 to CuTeDSL in utils.py, with links to the FA3 sources. The FP32 scores are permuted before the FP8 conversion, as FA3 does, instead of permuting the converted bytes with PRMT afterwards. - Express the producer V transpose with LdMatrix8x8x16bOp / StMatrix8x8x16bOp tiled copies, the CuTeDSL form of FA3's transpose_V, instead of inline PTX with hand-computed swizzled addresses. - Name the as-loaded V "VLoad" (pipeline_vload, mbar_ptr_VLoad and the transpose's source copy and tensors) and the transposed V "V", since FA4 uses sVt for the MMA's transposed view of sV, and say why native FP8 keeps this second, row-major V buffer. - Drop _scores_to_pv_fragment and the one-line intra-warpgroup-overlap helper; mma_one_n_block builds the FP8 P view inline, as FA3 does. Shared code: - Move SingleMBlockVarlenScheduler into tile_scheduler.py with the other schedulers. - Build the non-TMA O store copy from o_dtype in FlashAttentionForwardBase._setup_attributes, which resolves its TODO about an O dtype that differs from V's, and drop the SM90 native FP8 override. Paths whose output dtype has the input dtype's width get the same copy as before. - Softmax: apply max_offset by lowering the scaled row max, so FP8 and the other dtypes share one exp2 expression, and point the field comment at Note [Low Precision Scaling] in flash_fwd_sm100.py. interface.py: - Check the supported head-dim pairs inline instead of through a module global. - Key the descale compile on which dimensions are broadcast instead of six layout-name strings. Broadcast descales are still used without a copy. The non-FP8 kernels are unchanged. FP8 outputs are bit-identical except in d256 prefill, where the compiler now fuses the two-level rescale O * row_scale + O_pv into one FFMA: a few BF16 outputs move by one ulp, with the same LSE and the same error against an FP32 reference. Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

FA4's SM90 native FP8 forward (vllm-project/flash-attention#206) runs attention on E4M3 Q/K/V with FP8 WGMMA. Use it for FP8 KV cache on SM90. - get_flash_attn_version upgrades FA3 to FA4 on SM90 when the KV cache is FP8 E4M3 and FA4 supports the head size (64, 96, 128, 192 or 256, same for V), without softcap or batch invariance. An explicitly requested flash_attn_version is kept. The FlashAttention metadata builder and impl pass fp8_kv_cache. - Gemma4Config uses FA4 for all layers by default, including SM90 with FP8 KV cache, so its d256 sliding-window layers no longer run on FA3. - SM90 FA4 supports FP8 KV for equal head sizes 64, 96, 128, 192, 256 and 512. - Only d512 keeps BF16/FP16 Q for the FP8-KV dequant path. Smaller heads quantize Q, as FA3 does, and run the native FP8 forward. The FA4 wrapper applies its d512 FP8-KV checks only to BF16/FP16 Q. - The FlashAttention FP8 KV checks use KVQuantMode.FP8_PER_TENSOR instead of matching the cache dtype string, as the metadata builder only has the KV cache spec. - test_flash_attn_sm90_fp8_mma_selection now expects FA4 for d256 with an FP8 KV cache, and no longer expects the FA4 wrapper to reject FP8 Q. The FlashAttention pin is updated once vllm-project/flash-attention#206 is merged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

jhaotingc force-pushed the fa4-sm90-native-fp8-b branch from 3674cc7 to 5f41328 Compare October 23, 2026 14:02

jhaotingc marked this pull request as ready for review October 23, 2026 23:50

jhaotingc force-pushed the fa4-sm90-native-fp8-b branch from 8fb3cc5 to fb0d8dc Compare October 23, 2026 23:56

jhaotingc mentioned this pull request Oct 24, 2026

[Attention] Run SM90 FP8-KV attention on FA4 native FP8 vllm-project/vllm#63139

Open

6 tasks

jhaotingc force-pushed the fa4-sm90-native-fp8-b branch from fb0d8dc to 303f735 Compare October 25, 2026 01:17

jobartlett mentioned this pull request Oct 29, 2026

[Cute,Fwd,Sm90] BF16/FP16 decode launch tuning for Hopper #212

Open

jhaotingc force-pushed the fa4-sm90-native-fp8-b branch from 303f735 to 37aa13a Compare October 31, 2026 12:03

jhaotingc force-pushed the fa4-sm90-native-fp8-b branch from 37aa13a to 214ed25 Compare November 2, 2026 15:36

jhaotingc and others added 8 commits November 5, 2026 17:26

jhaotingc force-pushed the fa4-sm90-native-fp8-b branch from 214ed25 to 46a9ab2 Compare November 6, 2026 05:19

Navigation Menu

Repository navigation

[Cute,Fwd,Sm90] Native FP8 (E4M3 Q/K/V) forward for Hopper - #206

[Cute,Fwd,Sm90] Native FP8 (E4M3 Q/K/V) forward for Hopper#206
jhaotingc wants to merge 8 commits into
vllm-project:mainfrom
jhaotingc:fa4-sm90-native-fp8-b

jhaotingc commented Oct 22, 2026 •

edited

Loading

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

jhaotingc commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Design

Commits

Performance

Tests

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

jhaotingc commented Oct 22, 2026 •

edited

Loading