[Bugfix][MLA] Gather nvfp4_ds_mla chunked-context KV with the NVFP4 kernel

[stu-cao]

Status (2026-10-15)

Reworked after @zyongye's review (thanks — every point was right) and @LucasWilkinson's stu-cao#2, which proposed the same seq_starts change; rebased onto main 19b0efce30. Lucas's kernel-level test (random cache bytes, non-block-aligned starts, bit-exact against the same slots gathered without starts) is included as his commit in tests/kernels/attention/test_cache.py. Where this differs from #2: the BF16 workspace hunk is dropped (zyongye's point 1, it was unreachable) and DCP is rejected at startup instead of being left to fail in the generic gather (point 4).

• seq_starts on the NVFP4 gather kernel instead of a block-table shim. cp_gather_and_upconvert_nvfp4_kv_cache gains an optional seq_starts like its fp8 twin (token_offset = out_token - workspace_starts[req] + seq_starts[req]), and the dense dispatch passes chunk.starts exactly as the fp8_ds_mla branch does. shift_block_table_to_starts, its alignment precondition, the clamp, the per-chunk gather and its CPU test are gone. Other callers are unchanged (the argument defaults to None).
• Dropped the MLACommonMetadataBuilder workspace hunk and the fp8_query test arm. nvfp4_ds_mla is only accepted by FlashMLASparseBackend, whose builder (SparseMLACommonMetadataBuilder) allocates its own BF16 workspace and hard-codes a BF16 prefill query, so neither was reachable.
• DCP is rejected at startup (VllmConfig.validate_nvfp4_kv_cache_with_mla) instead of in the per-layer forward on the first request with context; covered by tests/test_config.py::test_nvfp4_ds_mla_rejects_decode_context_parallelism.
• The raw kv_c_and_k_pe_cache pass is intentional: the nvfp4_ds_mla cache is allocated as torch.uint8. The three other gather dispatch sites that still send nvfp4_ds_mla to the generic gather (mla_attention.py DCP path, kimi_k3/nvidia/mla.py, sparse_mla_attention.py) are unreachable today and left for a follow-up.

Verification on the rebased head (dc6aa4ae53, B300 SM103, from-source build): 17 passed, 0 failed — test_mla_prefill_context_nvfp4.py (1, the seq_starts path on real NVFP4 rows), test_mla_context_chunks.py (14), test_config.py::test_nvfp4_ds_mla_rejects_decode_context_parallelism (2); registered schema cp_gather_and_upconvert_nvfp4_kv_cache(..., int batch_size, Tensor? seq_starts=None); torch 2.14.0+cu130, TORCH_CUDA_ARCH_LIST=10.3a. tests/kernels/attention/test_cache.py -k nvfp4: 27 passed including test_cp_gather_and_upconvert_nvfp4_kv_cache_with_seq_starts.

Purpose

--kv-cache-dtype nvfp4_ds_mla crashes in the MLA chunked-context prefill:

mla_attention.py  _compute_prefill_context
_custom_ops.py    gather_and_maybe_dequant_cache
RuntimeError: get_fp8_kv_cache_data_type, csrc/attention/dtype_fp8.cuh:34,
              Unsupported fp8 kv cache data type: nvfp4_ds_mla

MLACommonBaseImpl._compute_prefill_context gathers each chunk's cached context KV into a BF16 workspace. It has a branch for fp8_ds_mla; every other dtype goes to gather_and_maybe_dequant_cache, which only knows the fp8 variants. The NVFP4 upconvert kernel added with the nvfp4_ds_mla format, cp_gather_and_upconvert_nvfp4_kv_cache, is never called here.

Sparse MLA reaches this path through its dense-MHA prefill: a prefill whose total length is <= index_topk and that has cached context. In practice:

• speculative-decoding warm-up: the dummy batch takes this path, so an engine with nvfp4_ds_mla KV and a DFlash draft never finishes starting;
• chunked prefill splitting a short prompt: the continuation step has context;
• prefix-cache hits on short prompts.

Longer prefills use the sparse FlashMLA kernels, which already handle nvfp4_ds_mla.

Fix:

• Route nvfp4_ds_mla to cp_gather_and_upconvert_nvfp4_kv_cache.
• Give that kernel an optional seq_starts (csrc + bindings + _custom_ops), mirroring cp_gather_and_upconvert_fp8_kv_cache: a continuation chunk reads from its context start instead of the request's first token. Without it, every chunk after the first would silently re-read the request's first tokens.
• Reject nvfp4_ds_mla with decode context parallelism at config validation; that combination is not validated.

Other dtypes are unaffected. I found no existing issue or PR for this.

Test Plan

New tests:

• tests/v1/attention/test_mla_prefill_context_nvfp4.py (SM100). Writes random KV into an nvfp4_ds_mla cache with concat_and_cache_mla, then runs _compute_prefill_context on it. The reference is the same method on a BF16 cache holding the dequantized values, through the generic gather. Every chunk must see bit-identical Q/K/V; the batch includes a chunk that continues one request and starts another.
• tests/kernels/attention/test_cache.py::test_cp_gather_and_upconvert_nvfp4_kv_cache_with_seq_starts (SM100): the kernel's seq_starts with non-block-aligned starts, bit-exact against a full gather.
• tests/test_config.py::test_nvfp4_ds_mla_rejects_decode_context_parallelism (CPU).

pytest tests/v1/attention/test_mla_prefill_context_nvfp4.py tests/v1/attention/test_mla_context_chunks.py tests/kernels/attention/test_cache.py tests/test_config.py -k "nvfp4 or context_chunks"

End-to-end on 1 × GB200 node (4 GPUs), run with the first revision of this change (block-table shift; same gather semantics):

• Model: GLM-5.3 (NVFP4 weights) + DFlash2 draft (fp8 KV, k=4).
• Engine: aggregated TP4, FLASHMLA_SPARSE, --kv-cache-dtype nvfp4_ds_mla, --block-size 1024, --max-model-len 1048576, prefix caching on.
• Arms: the same build with and without this change, plus an fp8-KV reference at --max-model-len 262144, since fp8 cannot fit one 1M-token request on TP4.
• Checks: startup; 6 distinct verifiable prompts, sequential then 12-way concurrent × 3 rounds; 12 concurrent ~1.7K-token prompts with --max-num-seqs 8 and an 8192-token budget so some are necessarily split and their continuation runs this path, compared with the same prompts served unsplit; draft acceptance over 3 × 400 tokens on a short and a ~10K-token prompt.

Test Result

Unit tests: see Status above for the current head. The first revision, on GB200 against v0.30.1rc0: the SM100 test failed before the change (Unsupported fp8 kv cache data type: nvfp4_ds_mla) and passed with it.

End-to-end (first revision):

	without the change	with the change	fp8 KV reference
startup	crashes in warm-up (trace above)	serves	serves (at 256K)
KV capacity at 1M, TP4	—	1,646,592 tokens (1.57×)	does not fit one 1M request
correctness, 42 requests	—	42 / 42	42 / 42
forced split prefill, 36 requests	—	36 / 36, identical to unsplit	36 / 36, identical to unsplit
draft acceptance, short prompt	—	39.0%	37.7%
draft acceptance, ~10K-token prompt	—	48.1%	47.4%

AI assistance (Claude) was used to investigate, write the change and tests, and run the tests above. I reviewed every changed line.

🤖 Generated with Claude Code

[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.

[zyongye]

The core fix looks right. I traced the chunk planner and the NVFP4 kernel and simulated both tests' chunk plans: starts are block-aligned on every reachable path, the shift lands on the right physical slots, and no column that is actually read ever gets clamped. The reference-vs-dequantized-cache test design is nice.

A few things worth addressing:

1. The workspace hunk looks dead. nvfp4_ds_mla appears in exactly one supported_kv_cache_dtypes list — FlashMLASparseBackend (flashmla_sparse.py:142) — whose builder derives from SparseMLACommonMetadataBuilder, not the MLACommonMetadataBuilder changed at mla_attention.py:2446. The workspace that actually feeds this gather is sparse_mla_attention.py:217 (dtype=self.model_config.dtype). The fp8-prefill-query rationale can't apply either: sparse metadata hardcodes q_data_type=self.model_config.dtype (sparse_mla_attention.py:399). The same reasoning makes the fp8_query arm of the new test a config no backend can produce — and since _build_prefill_metadata hardcodes a bf16 workspace, that arm doesn't exercise the hunk either.

2. Would seq_starts on the NVFP4 kernel be simpler than the block-table shim? cp_gather_and_upconvert_nvfp4_kv_cache_kernel computes token_offset = out_token_id - workspace_starts[req_id]; a + seq_starts[req_id] plus the arg plumbing would delete shift_block_table_to_starts, its alignment precondition, the clamp, the per-chunk gather, and the DCP NotImplementedError — and match the fp8 twin, which already handles arbitrary (even unaligned) starts via map_gather_page_task (cache_kernels.cu:1229).

3. shift_block_table_to_starts fails silently. The "start is a multiple of block_size" precondition isn't enforced, and clamp_ aliases an out-of-range column onto the last block, so a violation yields wrong KV rather than an error. Two unaligned sources already exist in tree: align_chunk_to_block=current_platform.is_cuda() (sparse_mla_attention.py:320), and DCP, where chunk.starts is overwritten with local offsets cdiv(start, virtual) * local_block_size (mla_attention.py:2152) — kept out only by the new raise in a different method. An assert (or dropping the clamp) would make it loud.

4. The DCP rejection fails late. The raise at mla_attention.py:3076 sits inside the per-layer forward and only fires once a request carries context, so the engine starts, serves context-free prompts, then dies on the first prefix-cache hit. A validator next to validate_nvfp4_kv_cache_with_mla (config/vllm.py:3499) would reject at startup. It's also untested as written.

Minor:

• The gather dispatch now has four near-identical copies: mla_attention.py:2959 (fixed here), mla_attention.py:3106, kimi_k3/nvidia/mla.py:940, sparse_mla_attention.py:919. The other three still route nvfp4_ds_mla into gather_and_maybe_dequant_cache; they're unreachable today by accident (_is_masked_mha_available returns False for any quantized cache), not by contract.
• flashmla_sparse.py:1114 passes kv_c_and_k_pe_cache.view(torch.uint8); the new call passes it raw. This kernel hard-checks for uint8 (cache_kernels.cu:1880) while the fp8 one doesn't — either the view there is redundant or this call needs it.
• The shift builds a full (num_reqs, max_blocks) int64 index and table copy per continuation chunk, though plan_mla_context_chunks guarantees only row 0 can ever shift.

Review assisted by Claude Code.

[LucasWilkinson]

Sorry will do a proper review tomorrow but my agentic review seems to think there might be an easier way to do this? stu-cao#2

Similar to @zyongye comment 2

[stu-cao]

@zyongye Thank you for the review. I have addressed your comments in the new revision.

1. Removed the change to MLACommonMetadataBuilder which is a code path never executed for nvfp4_ds_mla. I reverted this change.
2. Use seq_starts in the kernel instead of block table shift in original implementation.
3. DCP fail early. As I didn't explicitly implement and validate DCP, for now, I will leave DCP as unsupported. In the new revision, we added VllmConfig.validate_nvfp4_kv_cache_with_mla now rejects nvfp4_ds_mla with decode_context_parallel_size > 1 at startup.
4. kv_c_and_k_pe_cache pass is intentional as the proper way would be create one helper _gather_context_kv() that owns the dtype dispatch. This would touch the shared MLA code touching Kimi-K3 and the sparse path. I can handle this in a follow-up PR.

[stu-cao]

@LucasWilkinson Hi Lucas, thank you so much for the review and suggestion PR! It make complete sense to me. I incorporate the change from your PR in this PR. The seq_starts kernel argument is added. I validated the setup on a B300.

Here's other changes I made based on Yongye's comment. Could you please take another look?

@zyongye Thank you for the review. I have addressed your comments in the new revision.

1. Removed the change to MLACommonMetadataBuilder which is a code path never executed for nvfp4_ds_mla. I reverted this change.
2. Use seq_starts in the kernel instead of block table shift in original implementation.
3. DCP fail early. As I didn't explicitly implement and validate DCP, for now, I will leave DCP as unsupported. In the new revision, we added VllmConfig.validate_nvfp4_kv_cache_with_mla now rejects nvfp4_ds_mla with decode_context_parallel_size > 1 at startup.
4. kv_c_and_k_pe_cache pass is intentional as the proper way would be create one helper _gather_context_kv() that owns the dtype dispatch. This would touch the shared MLA code touching Kimi-K3 and the sparse path. I can handle this in a follow-up PR.