[drkm-eng]

Summary

• Adds fp8 e4m3 q/kv support to the sheared (attention-bias) forward path on SM100/SM103 (validated on SM100). Adapted from Add Blackwell plain FP8 attention support #1 (closed unmerged), plus a correctness fix for the interaction between fp8 P pre-scaling and lazy softmax rescaling.
• blackwell_helpers.py: the tcgen05 MMA kind:: string is derived from the MMA op instead of the hardcoded f16; gemm_ptx_precomputed/gemm_ptx_precomputed_varname take a kind parameter (default f16).
• flash_fwd_sm100.py: fp8 q/kv on the sheared path. P is pre-scaled by 2**max_offset (max_offset=8) before the e4m3 tmem store; the offset is carried consistently through row_sum and the attention-sink term, cancels in the O normalization, and is subtracted from the LSE. St32x32b Repetition 8 for the narrower fp8 P store. An fp8 ex2 tuning entry pins the validated config's values (currently equal to the base tuning). The bias dtype check now requires bf16/f16 explicitly (bias is never quantized). The kv_stage reduction condition changes from v_dtype.width == 8 to v_dequant: identical for existing (blockscaled) configs, and it deliberately leaves plain-fp8 KV on the smem-budget-derived stage count instead of the 2-stage mxfp8 clamp. The tBrS fragment construction moves inside the bias-participation guard, so threads that do not copy bias no longer build the fragment.
• All fp8 P pre-scale machinery (max_offset, rescale_threshold=0.75, the St32x32b Repetition-8 P store) keys on a plain-fp8 flag, not dtype width alone: qk-blockscaled Q is also e4m3 but keeps a bf16 P store, and those configs retain stock behavior (rescale_threshold 0.0, max_offset 0, Repetition 16).
• softmax.py: max_offset parameter, folded into the packed FMA that applies scale_log2 and subtracts the scaled row max.
• interface.py: accepts torch.float8_e4m3fn q/k/v on the non-blockscaled path, asserts SM100/SM103 for it, and returns bf16 output for plain fp8.

Why

Downstream (vLLM Inkling, hybrid attention with learned relative position bias) needs fp8 KV cache for long-context serving. The cute score-mod path supports fp8 but is roughly 7x slower than the sheared path at prefill, so the sheared path needs native fp8. A vLLM PR wiring this in will follow and reference this PR.

Correctness fix over #1: rescale_threshold

With fp8, P values are pre-scaled by 2**max_offset (max_offset=8) before the e4m3 store, so the slack tolerated by lazy rescaling is bounded by the e4m3 finite range: threshold <= log2(448) - 8 = 0.807. #1 set rescale_threshold=4.0. With attention bias, row-max increases below the threshold are routine, and accumulated slack past 0.807 clips P against the e4m3 ceiling: six deterministic test failures with 0.3 to 0.7 abs output error.

The root cause was isolated by a zero-bias diagnostic (failures vanish without bias) and bisection: threshold=0 fixes, max_offset=0 fixes, and disabling the ex2 emulation tuning does not, isolating the interaction of the pre-scale with lazy rescaling. An empirical threshold bracket agrees with the bound: 0.75, 0.80, 0.81, 1.0, and 1.5 all pass the fp8-sheared suite, the failure onset lies between 1.5 and 2.0, and the dose-response is monotonic beyond it (2.0: 4 failures, 3.0: 6, 4.0: 6). The bound is a design invariant, not a sharp cliff at test tolerance: a violation of 0.003 (log2 units) clips only values within 0.2% of the e4m3 ceiling, below the 0.15 test tolerance, so passing slightly above 0.807 is expected.

This change sets rescale_threshold=0.75 for fp8: the 0.807 bound guarantees zero P clipping, and 0.75 leaves margin for compiler reassociation. bf16/f16 keep 8.0.

Provenance

Kept from #1: MMA kind parametrization, fp8 q/kv handling in the forward kernel and softmax, interface fp8 acceptance and output dtype. Skipped: hunks that already exist on current main (pre-allocated out params), including the duplicate out parameter that caused #1's SyntaxError. Searched open and closed PRs and issues; #1 is the only prior attempt at plain fp8 on this path.

Validation

The in-tree test (tests/test_forward_only_tree.py) is structural only, so functional evidence comes from the downstream vLLM Inkling integration suite, as in #1. All runs on B200 (SM100).

Correctness: 160/160 tests green across the FA4 attention and KV write-path suites under tests/models/inkling/ in the vLLM tree (test_fa4_rel_attention.py, test_fa4_warmup.py, test_qkvr_prep.py). The fp8 test matrix ships with the follow-up vLLM PR (a cross-link will be added once it posts); upstream vLLM's current tests do not yet contain it. Matrix: both FA4 backends (this sheared-bias path and the cute score-mod path) x fp8 x scale sets {(1,1,1), (0.5,0.5,1.5)} x rel extents {128, 1024}, at atol/rtol 0.15 against a dequantized-input bf16 reference, plus fp8 write-path roundtrip and gate tests.

Latency: 20 cases spanning self and chunked prefill, single-stream and batched decode, TP1/TP2/TP4 head splits, KV lengths to 128K. CUDA events, 3 warmup iterations (includes JIT compile), 30 timed, median reported. fp8-sheared is faster than bf16-sheared in 18 of 20 cases. The two regressions: decode_kv65536_b1_tp4 (105.5 -> 108.2 us, +2.6%) and local_prefill_q8192_tp1 (340.7 -> 341.6 us, +0.3%). Best cases: 1.45x batched decode at kv=64K b=32 (1658 -> 1140 us), 1.41x at kv=32K b=32, 1.23x chunked prefill at 64K (1855 -> 1509 us); single-stream decode is a more modest ~1.05x.

Benchmark configuration: torch 2.13.0, single B200 (183GB), head_dim 128, GQA (32 query / 8 KV heads at TP1; per-rank 16/4 at TP2, 8/2 at TP4), num_splits from the production heuristic, medians of 30 iterations. Wins grow with KV length; short-KV cases are near parity.

• pre-commit run --files flash_attn/cute/blackwell_helpers.py flash_attn/cute/flash_fwd_sm100.py flash_attn/cute/interface.py flash_attn/cute/softmax.py
• pytest tests/test_forward_only_tree.py
• Follow-up: bump the GIT_TAG in vLLM cmake/external_projects/tml_fa4.cmake once this merges. The build fetches from that tag at build time and installs into vllm/third_party/tml_fa4; there is no checked-in vendored copy to update.

Compatibility

• No behavior change for bf16/f16: the MMA kind resolves to f16 for MmaF16BF16Op, so the emitted PTX is unchanged; rescale_threshold stays 8.0 and max_offset=0 keeps the softmax math identical.
• Plain (non-blockscaled) fp8 is gated to SM100/SM103 by an interface assert. All validation was on SM100; SM103 passes the assert but is untested here.
• Attention bias stays bf16/f16 with fp8 q/kv. The previous check required bias dtype == v_mma_dtype, which would force an e4m3 bias under fp8.
• Output dtype for plain fp8 is bf16.
• The fp8 ex2 tuning entry pins the validated config (1-CTA, causal, head_dim 128, SM100) against future base-table retuning; its values currently equal the base tuning, and other configs run the base tuning.

AI assistance

AI assistance (Claude) was used to implement, benchmark, debug, and document this change; the commit carries a Co-authored-by trailer. The human submitter has reviewed every changed line and can explain and defend the implementation.

[drkm-eng]

vLLM integration PR: vllm-project/vllm#54705. RFC: vllm-project/vllm#54704.

[drkm-eng]

@WoosukKwon @mgoin the vLLM-side PR (vllm-project/vllm#54705) is rebased onto current main and conflict-free; this kernel PR is unchanged and still applies cleanly to the b206834 pin. A review pass would be appreciated. The rescale_threshold fix here is the piece that resolves the bias-path failures the PR #1 approach hit; derivation and B200 validation are in the body.