[zyongye]

Overview

For an FP8 KV cache whose per-tensor scales are 1.0, FlashAttention now passes FA4 no descale at all. Before, it passed broadcast descale views that FA4 copied into contiguous tensors, three copy kernels per attention layer and step.

This PR began as the squash of the whole MiniMax-M3.1-Flash tree. That squash has since been split into the PRs below; this PR now holds only the FA4 change.

How #63232 was split

Part of the original squash	Where it went	State
MiniMax-M3.1 model and DSpark draft model, with their kernels (MSA NVFP4 KV cache and indexer, MegaMoE staging, ROCm MXFP8 linears, ...)	#64111	merged
ROCm: NVFP4 MoE re-quantization to MXFP4/MXFP8, and fused shared experts through online quantization	#64378	merged
FA4: no FP8 descale for a unit per-tensor scale	this PR	open
NIXL P/D: re-sharding NVFP4 KV pages across heterogeneous TP (AttentionSpec.kv_head_major_slots)	#64475	open
Scheduler: keep KV cache zeroing for mixed-precision UniformTypeKVCacheSpecs groups	#64472	open
DSpark/DFlash draft step: Markov bias in the Gumbel kernel, wider DFlash padding, is_padding from the positions kernel, top-p fix	#64476	open
Batch-sharded DSpark draft sampling	dropped from #64476; #63886 covers it	open (#63886)
SimpleCPUOffload: one host buffer per rank with HugeTLB pages	already open as #60130; the squash's extra ROCm change (SDMA copy-engine flag for batched copies) was pushed onto #60130	open (#60130)
Scheduler draft-slot change	dropped: main's max_num_new_slots_for_drafting already returns 0 for DFlash/DSpark	–
Sparse paged-attention KV block count fix (ROCm)	dropped: already on main via #63270	–

The full squash as it was before the split is a0ea941 (on 5d35ce3), in this PR's force-push history.

PRs that were opened against this branch (minimax-m31-flash) while it held the squash: #64359, #64273, #64264, #64206, #63528, #63491 and #62509. Their diffs against this branch now include the whole old squash (e.g. #64359 shows 193 files). The model (#64111) and the ROCm re-quantization (#64378) are on main, so they can retarget main.

Claims

• Unit scales: when FA4 runs with a per-tensor Q/K/V scale of exactly 1.0, FlashAttentionImpl passes None for that descale.
  • FA4 reads a missing descale as 1.0, so the output is bit-identical.
  • FA4's maybe_contiguous no longer copies the three broadcast descale views per layer and step.
• Other scales: calibrated per-tensor scales and per-head scales keep their descale tensors. Other FlashAttention versions get the same broadcast descale as before; only the strides of its size-1 dims are made canonical, which changes metadata, not values.
• (1, 1) descale views: a view with one sequence and one KV head per rank now gets canonical strides. FA4 treats such a view as contiguous and skips its copy, and its layout check then rejects the zero stride.

Validation

On GB300:

pytest tests/v1/attention/test_attention_backends.py

148 passed, 9 skipped, 4 failed. The 4 failures are test_sliding_window_encoder_backend_correctness (embeddinggemma-300m). They fail the same way on main (2abde62) in this environment.

The new test_flash_attn_fa4_unit_fp8_descale_is_exact (Blackwell only) checks that:

• a unit scale gives no descale;
• FA4's output with no descale equals its output with the broadcast unit descale, bit for bit;
• a descale of 2 changes the output, so the comparison is not vacuous;
• calibrated and per-head scales keep their descale tensors.

Not covered:

• No end-to-end benchmark on this branch; the three copy kernels saved per layer and step were not timed.
• No accuracy eval: the output is bit-identical by the test above.
• The (1, 1) stride case is not in the new test, which uses 3 sequences and 4 KV heads.

Details

• Why the check is on the host: the scales are fixed once the weights are loaded, so the _scale_float mirrors decide the case on the host with no GPU sync. A per-head scale's float mirror holds only its maximum, so the check also requires a single-element scale.
• Compiled kernels: FA4 compiles a separate kernel without descales, since descale presence is part of its compile key. A model with unit scales always takes that kernel.
• Not a duplicate: the open FA4 PRs change other things. [Perf][FA4] Warm up dense (paged) FA4 attention kernels #48156 warms up the paged FA4 kernels, [Kernel] Integrate FA4 SM90 FP8 MMA for Gemma-4 full attention #60372 adds SM90 FP8 MMA for Gemma-4, [Perf][Attention] FA4 sparse MLA: grouped prefill with an FP8 KV cache #63328 is sparse-MLA FP8 prefill, and [Feature][Model][Kernel] Inkling: fp8 e4m3 KV cache with per-tensor scales on SM100 #54705 covers Inkling per-tensor scales outside flash_attn.py. None of them changes how FlashAttentionImpl passes descales.

Credits: qizixi wrote the FA4 change and is a co-author in the commit. AI assistance (Claude) was used to split the original squash into the PRs above, to write the FA4 test and to draft this description. The submitter reviewed the change and ran the tests above.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[fxmarty-amd]

@zyongye when can we expect it to land?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zyongye.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[amd-dlimpus]

FYI @zyongye: #64264 is a small ROCm follow-up stacked on this branch. With NVFP4 KV on MI355X, the FlyDSL kernels compile lazily on first launch, and minimax_m3_msa_warmup returns early on ROCm. So the first C16 batch paid about 1.5 s extra of mean TTFT inside requests. #64264 prebuilds every variant at startup (about 12 s per rank). It doesn't touch the CUDA path.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98868 for commit 3e48944f896b.