Repository navigation
[Attention] Enable FA4 sparse MQA decode on Blackwell - #55866
Open
simon-veitner-redhat wants to merge 12 commits into
Open
simon-veitner-redhat wants to merge 12 commits into
simon-veitner-redhat wants to merge 12 commits into
Conversation
simon-veitner-redhat
requested review from
Harry-Chen,
LucasWilkinson,
MatthewBonanni,
pavanimajety and
tlrmchlsmth
as code owners
September 8, 2026 08:55
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
simon-veitner-redhat
force-pushed
the
sparse-mqa
branch
2 times, most recently
from
September 9, 2026 05:58
65d3925 to
c811040
Compare
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
simon-veitner-redhat
force-pushed
the
sparse-mqa
branch
from
September 14, 2026 19:16
c811040 to
8925a7b
Compare
simon-veitner-redhat
force-pushed
the
sparse-mqa
branch
4 times, most recently
from
September 16, 2026 14:48
4204d78 to
a0ba252
Compare
simon-veitner-redhat
requested review from
AndreasKaratzas,
WoosukKwon,
alexm-redhat,
njhill,
tjtanaa,
youkaichao and
zhuohan123
as code owners
September 16, 2026 14:48
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
simon-veitner-redhat
force-pushed
the
sparse-mqa
branch
from
October 9, 2026 05:12
f4caddb to
0850ebc
Compare
Drop mock-only tests that pin internals and duplicate coverage, fold the HiSparse decode-only tests into the lanes test, split the shared converter mock in the FA4 routing test, and replace the SM100 priority restatement with a fallback test. Signed-off-by: Simon Veitner <sveitner@redhat.com> Co-Authored-By: Claude <noreply@anthropic.com>
Contributor
|
/ci run |
LucasWilkinson
enabled auto-merge (squash)
October 9, 2026 14:28
|
✅ Triggered Buildkite CI #93859 for commit |
auto-merge was automatically disabled
October 9, 2026 17:08
Head branch was pushed to by a user without write access
The FA4 sparse MLA backend imports vllm.vllm_flash_attn, which is not built on ROCm, so importing it at module level broke test collection there. Parametrize over AttentionBackendEnum and resolve the backend class inside the test, and mark the FA4 case CUDA only, so the CPU cases for the other backends keep running on ROCm. Signed-off-by: Simon Veitner <sveitner@redhat.com>
simon-veitner-redhat
force-pushed
the
sparse-mqa
branch
from
October 9, 2026 17:12
761b25c to
99e55b6
Compare
Contributor
|
/ci retry |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> # Conflicts: # tests/v1/attention/test_sparse_mla_backends.py # tests/v1/attention/test_sparse_mla_kv_cache_layout.py # vllm/v1/attention/backends/mla/flashattn_mla_sparse.py # vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
Contributor
|
/ci run |
|
✅ Triggered Buildkite CI #94036 for commit |
The test imports vllm.v1.attention.backends.mla.flashattn_mla_sparse, which imports vllm.vllm_flash_attn; that extension is not built on ROCm, so the test failed with ImportError on AMD CI. Mark it CUDA only, matching the neighbouring selection test. Signed-off-by: Simon Veitner <sveitner@redhat.com>
Member
|
/ci run |
|
❌ This PR is 57 commits behind upstream |
Member
|
/ci run |
|
✅ Triggered Buildkite CI #94216 for commit |
Contributor
CI selector (shadow): 194 test steps (270 jobs) instead of 195 (271 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Today's rules run everything for this diff. Selector would run (194)
Would skip (today's rules run them) (1)
Would add (today's rules do not run them) (0)none AMD mirrors: would skip (3)
AMD mirrors: would add (0)none 14 changed files · base |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
Contributor
Author
|
/ci run |
|
✅ CI is already running for this commit: https://buildkite.com/vllm/ci/builds/94216 |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
[Attention] Enable FA4 sparse MQA decode on Blackwell
This adds
FLASH_ATTN_MLA_SPARSE_FA4, a sparse MLA backend for Blackwell with a BF16 KV cache. The FA4 CuTe DSL kernel from vllm-project/flash-attention#195 serves decode batches and the decode rows of mixed HiSparse batches. Every other batch runs on the trtllm-gen lane, as inFLASHINFER_MLA_SPARSE. At 16 or fewer query heads per rank the backend is the first choice, with and without HiSparse; above that head count it stays opt-in.The first commit, "Share sparse MLA MQA cache handling", will go up as its own PR and is only here until it lands. It moves index conversion and the HiSparse lane split from the FA3 and FlashInfer sparse MLA backends into
SparseMLACommonImpl.forward_mqa, so a backend implements only_forward_mqa_kernel(..., is_decode). The design is by @LucasWilkinson. One FA3 behaviour changes with it: FA3's plain path converts top-k through the index group's shared buffer, as FlashInfer already did, instead of once per layer. The values are identical while the layers of a group share block stride, block table and request ids, which index groups guarantee. A test pins leader and follower layers.The second commit is the FA4 backend on top of that hook:
flash_attn_varlen_funcforwardsgather_kv_indices,gather_kv_valid_lengthandq_vto FA4.FlashInferMLASparseImpl. Its hook runs the FA4 decode kernel whenis_decode, otherwise FlashInfer's trtllm-gen call with the LSE converted to base e.UNIFORM_BATCHCUDA graph support. Inheriting the builder brings the DCP workspace pre-size: an earlier version of this PR used the default 394 MiB trtllm workspace on its prefill lane and crashed under DCP 2 on a prefill step of about 8192 tokens, where the softmax-stats slab needs 513 MiB.index_topka multiple of 128, and a supported gathered head count under DCP.autotune_hisparse_decodecompiles the FA4 kernel before graph capture._uses_flat_kv_cache, since it reads the same flat KV cache as FlashInfer.The last commit bumps the flash-attention pin to vllm-project/flash-attention 2e3a6694, which has #195 and the sparse decode split-rule tuning for DCP shapes (vllm-project/flash-attention#200).
Performance
GLM-5.2 NVFP4-FP8 with BF16 KV on B300, TP4, FA4 against
FLASHINFER_MLA_SPARSEfrom the same tree (this PR onmaind71f662), change in median ITL (negative favours FA4). Plain and DCP 2 are averaged over two runs. The HiSparse column assumes #59309 is merged: it applies that PR's diff, pins the MoE autotune tactics so both backends run the same MoE kernels, and averages two server starts with two runs each.* One FA4 run measured 35.2 ms against 33.4 to 33.7 ms for the other three; without it the figure is +0.1%.
Under HiSparse FA4 is on par with FlashInfer, which did not justify a separate priority order. Per decode step, FA4's attention kernels take about 190 µs less than FlashInfer's at 4096 / 8, but start-to-start variance in all-reduce time is larger than that difference. On current
main, #59235 makes HiSparse about 0.45 ms per decode step slower for both backends, because it writesrequest_idsin every layer. #59309 moves those writes into the group leader's resolve, which brings the 57 follower layers' 171 kernels per step back out of the decode graph.Validation
pytest tests/v1/attention/test_sparse_mla_backends.py: 475 passed, 672 skipped.pytest tests/v1/e2e/general/test_hisparse.py: 5 passed (FlashInfer and FA4). The FlashInfer, KV cache layout, backend selection and masked-MHA routing tests pass as well.main, plain generations byte-identical.AI-assisted. The September 7 upstream search found no matching open FA4 sparse-MLA integration PR; related platform-specific work does not provide this SM100 backend.