[Attention][MLA] Support nvfp4_ds_mla on FLASHINFER_MLA_SPARSE, with native NVFP4 decode

[stu-cao]

Status (2026-10-16)

• FlashInfer dependency resolved upstream. FlashInfer v0.7.2 (2026-10-15) contains the native NVFP4 sparse MLA decode kernel (feat(cake_backend): generated backend="cake" for nvfp4_sparse_mla_decode on SM100 / SM103 flashinfer-ai/flashinfer#5970), and vLLM main pins flashinfer-python==0.7.2 since [Dependency] Upgrade FlashInfer to 0.7.2 and restore KDA auto backend #61913. This PR keeps the tested cuda kernel (backend="cuda", FlashInfer's default).
• But the 0.7.2 wheel ships the Python entry point without the kernel's JIT source (flashinfer/experimental/nvfp4_sparse_mla_decode/csrc/* is missing from the package-data allowlist), so nvfp4_sparse_mla_decode imports yet cannot compile. Fix: fix(nvfp4_sparse_mla_decode): ship the JIT sources in wheels flashinfer-ai/flashinfer#6454. Until a release carries it, has_flashinfer_nvfp4_sparse_mla_decode() (fifth commit) also requires the source file to be present; on a wheel without it the backend logs once and uses the FP8-staged path, so the dtype works everywhere and the native kernel switches on by itself once FlashInfer ships the source.
• Rebased onto vLLM main at fb3fec970e (head 02f5bbf1ad, 5 commits). [Attention] Enable FA4 sparse MQA decode on Blackwell #55866 moved the generic sparse-MQA path into SparseMLACommonImpl.forward_mqa; this backend's forward_mqa is now a thin override that defers to the base for non-NVFP4 caches, and the NVFP4 staged-row kernel calls go through the new _forward_mqa_kernel(layer=, block_size=FP8_STAGING_PAGE_SIZE, is_decode=) entry. No NVFP4 method logic changed.
• Unit tests on the rebased head (B300, 2026-10-16), selection test_flashmla_nope_sm90_backend_selection.py + test_flashinfer_mla_sparse_nvfp4_native_decode.py + test_sparse_mla_backends.py -k 'nvfp4 or sparse or FLASHINFER or sm100':
  • with the pinned flashinfer-python==0.7.2 wheel: 639 passed, 0 failed; the 10 native-decode cases skip ("needs SM100/SM103 and FlashInfer's native NVFP4 sparse MLA decode") because the probe sees no JIT source — the staged path serves every case;
  • with an editable FlashInfer main checkout (source present): 649 passed, 0 failed, native-decode 10/10.
• GSM8K rerun on the rebased head (2026-10-17, B300): nvfp4_ds_mla 90.5 % vs fp8 90.3 %, 0 invalid responses in both, same flags as the table below (the server logged "decodes nvfp4_ds_mla with FlashInfer's native NVFP4 kernel"). KV cache at --gpu-memory-utilization 0.9, TP4: 4,512,384 vs 2,861,440 tokens (1.58x). The end-to-end cuda vs cake comparison below is unchanged.

Purpose

FLASHINFER_MLA_SPARSE (the TRTLLM-gen sparse MLA backend on SM10x) now accepts --kv-cache-dtype nvfp4_ds_mla, which only FLASHMLA_SPARSE accepted. The NVFP4 cache holds 1.6x the tokens of an FP8 one (352 vs 576 bytes per token).

Four commits:

1. FP8 staging, based on Chi McIsaac's work (Co-authored-by in the commit). The TRTLLM-gen kernel reads FP8 rows only, so each layer stages the rows it attends to as FP8 (dequant / k_scale): decode and spec-decode tokens gather their top-k rows, prefill tokens share one context gather per request and chunk (Triton kernels in nvfp4_ds_mla_fp8_gather.py). The staging workspace is reserved up front so the memory profile accounts for it. The dtype needs kv_lora_rank 512 and qk_rope_head_dim 64 and is rejected with context parallelism and with HiSparse. The rebase preserves the current TRTLLM builder's UNIFORM_BATCH CUDA-graph support and varlen decode limit; host-planned staged prefill does not widen graph capture. The impl and builder retain class-level defaults for tests that build them with object.__new__.
2. Native decode. When FlashInfer provides flashinfer.mla.nvfp4_sparse_mla_decode (experimental; the original feat(nvfp4_sparse_mla_decode): add experimental NVFP4 sparse MLA decode on SM100/SM103 flashinfer-ai/flashinfer#5717 implementation is now included in merged feat(cake_backend): generated backend="cake" for nvfp4_sparse_mla_decode on SM100 / SM103 flashinfer-ai/flashinfer#5970), the layer has 16 heads per rank, the top-k width is 512, 1024 or 2048 and the GPU is SM100 or SM103, decode reads the cache rows as stored: no staging copy, no decode staging workspace, and no e4m3 re-rounding of the dequantized values. Both fp8 bmm scales are divided by k_scale, because the staged rows they were built for hold dequant / k_scale while the native kernel reads exact values. Everything else keeps the staged path. The runtime capability probe (has_flashinfer_nvfp4_sparse_mla_decode()) keeps older FlashInfer packages on the fallback; the op is on FlashInfer main, but this PR does not assume a particular release. The native call explicitly selects backend="cuda" to preserve the B300-tested behavior; see Status for the cake comparison.
3. Tests. test_sparse_backend_decode_correctness already parametrizes nvfp4_ds_mla and skipped it for this backend. Once the backend accepts the dtype those cases run, and they failed: the mock layer quantized the query to fp8 only for fp8* caches, while MLAAttention does it for every quantized cache. The mock now does the same, and the FlashInfer nvfp4_ds_mla cases use the fp8 tolerance (the backend reads them with an fp8 query, native NVFP4 decode where supported, and staged FP8 rows elsewhere).

Backend selection is unchanged (part of the first commit). On SM10x, auto-selection prefers FLASHINFER_MLA_SPARSE for every quantized KV dtype, so accepting nvfp4_ds_mla would have moved existing users of that dtype from FLASHMLA_SPARSE to the new path without an end-to-end benchmark. nvfp4_ds_mla keeps FLASHMLA_SPARSE first and FLASHINFER_MLA_SPARSE as the fallback; this path is opt-in with --attention-backend FLASHINFER_MLA_SPARSE, and test_sm100_quantized_sparse_priority pins the order. Making it the default needs numbers against FlashMLA's NVFP4 kernels, which this PR does not have.

Not a duplicate. Checks run as AGENTS.md prescribes, 2026-09-29: gh pr list --repo vllm-project/vllm --state open --search with "nvfp4_ds_mla", "FLASHINFER_MLA_SPARSE", "nvfp4 flashinfer sparse mla" and "FlashInfer nvfp4 decode MLA". No open PR adds nvfp4_ds_mla to FLASHINFER_MLA_SPARSE, and flashinfer_mla_sparse.py on main has no NVFP4 support. Related, separate: #58744 (mine; touches mla_attention.py's chunked-context gather, not this backend) and #53770 (a fp8_k_nvfp4_v KV format for the non-MLA FlashInfer backend).

AI assistance: the original PR was prepared with Claude Code (Claude Opus 5.5 and Sonnet 5.5); the rebase and upstream-FlashInfer alignment were assisted by OpenAI Codex. The commits carry assistance/attribution trailers.

Test Plan

Historical hardware run: one B300 SXM6 (SM103, driver 580.173.02), CUDA 13.0, torch 2.13.0+cu130, FlashInfer at PR flashinfer-ai/flashinfer#5717. These GPU tests have not yet been rerun after the rebase (see Status).

pytest tests/v1/attention/test_flashinfer_mla_sparse_nvfp4_native_decode.py
pytest tests/v1/attention/test_sparse_mla_backends.py
pytest tests/v1/attention/test_mla_backends.py
pytest tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py
pytest tests/v1/attention/test_sparse_mla_kv_cache_layout.py tests/v1/attention/test_flashinfer_sparse_mla_workspace.py

The first file is new. test_native_decode_matches_staged_decode runs FlashInfer's native decode and the FP8-staged decode over a per-layer cache whose blocks alternate with another layer's, with k_scale 1 and 0.5, and checks both against an exact FP32 reference (native under 1 %, staged under 6 %, native vs staged under 6 %). test_backend_native_decode_folds_k_scale calls the backend's own _nvfp4_decode (impl built without __init__ and the index conversion stubbed, as the neighbouring tests do) and checks it against the same reference with k_scale 1 and 0.5, so a missing 1 / k_scale on either bmm scale fails it.

Test Result

Historical B300 results: PR head 0e2b2810e8 against upstream main 4e0a414c44 (before the rebase):

Selection	main	this PR
test_flashinfer_mla_sparse_nvfp4_native_decode.py (new)	n/a	10 passed
test_sparse_mla_backends.py	396 passed, 378 skipped	480 passed, 294 skipped
test_sparse_mla_kv_cache_layout.py, test_flashinfer_sparse_mla_workspace.py	46 passed	46 passed
7 other files that mention the backend	206 passed, 27 skipped	212 passed, 27 skipped

The 84 extra passes in test_sparse_mla_backends.py are the FlashInfer x nvfp4_ds_mla cases of test_sparse_backend_decode_correctness, which the backend used to skip. The 6 extra passes in the other files are the new test_sm100_quantized_sparse_priority cases. The 7 other files are test_dspark_noncausal_sparse_mla.py, test_flashinfer_sparse_mla_sm120_api.py, test_flashmla_nope_sm90_backend_selection.py, test_cuda_backend_probe_errors.py, test_flashinfer_mla_sparse_sm90.py, tests/kernels/attention/test_attention_selector.py and tests/kernels/attention/test_flashinfer_mla_decode.py. test_mla_backends.py (335 dense-MLA tests) was not run to completion; none of its tests reach code this PR changes.

GSM8K, tests/evals/gsm8k/gsm8k_eval.py, 20-shot, 1319 questions, temperature 0, GLM-5.3 with NVFP4 weights, TP4 on four B300 GPUs, FLASHINFER_MLA_SPARSE, no speculative decoding:

--kv-cache-dtype	accuracy	95 % interval	invalid responses	rebased head 02f5bbf1ad (2026-10-17)
fp8	89.69 % (1183/1319)	87.9 to 91.2 %	0	90.3 %, 0 invalid
nvfp4_ds_mla	89.76 % (1184/1319)	88.0 to 91.3 %	1	90.5 %, 0 invalid

The NVFP4 server logged that decode used FlashInfer's native kernel. The 20-shot prompts are 3,545 to 3,708 tokens, above the 2,048 top-k, so prefill took the staged sparse path. At the same memory budget the KV cache held 4,526,592 tokens with nvfp4_ds_mla and 2,870,464 with fp8 (1.58x). One run per arm, each server on a different node. The servers ran an earlier revision of this branch, before the class-level defaults, the None guard, the selection-order change and two tests. The eval names --attention-backend FLASHINFER_MLA_SPARSE explicitly, and none of those changes alters an instance built through __init__.

Limits of the evidence:

• The FlashInfer-backend nvfp4_ds_mla cases in test_sparse_backend_decode_correctness use 128 total heads at TP 1, 2 and 4 and top-k 128, so they run the staged path end to end. The native branch is covered by the backend-level test above (top-k index conversion stubbed) and by the GSM8K run (16 heads per rank at TP4).
• GSM8K shows no accuracy loss on this task. It does not cover long-context retrieval or tool calls, and speculative decoding was off.

🤖 Generated with Claude Code

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[mergify]

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

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

[gaby]

@stu-cao Merge conflicts

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

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

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

[mergify]

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

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

[stu-cao]

GSM8K rerun on the rebased head 02f5bbf1ad (B300, GLM-5.3 NVFP4 weights, TP4, --attention-backend FLASHINFER_MLA_SPARSE, 20-shot, 1,319 questions, temperature 0, no speculative decoding; same flags as the earlier run in the description):

--kv-cache-dtype	accuracy	invalid responses	KV cache at --gpu-memory-utilization 0.9
nvfp4_ds_mla	90.5 %	0	4,512,384 tokens
fp8	90.3 %	0	2,861,440 tokens

The NVFP4 server ran the native decode kernel (it logged "FLASHINFER_MLA_SPARSE decodes nvfp4_ds_mla with FlashInfer's native NVFP4 kernel; prefill is staged to FP8"), with FlashInfer main installed from source; the pinned 0.7.2 wheel lacks the kernel's JIT source (flashinfer-ai/flashinfer#6454), in which case the probe keeps the FP8-staged path. Pre-rebase results were 89.76 % vs 89.69 %.