[0z5a]

Overview

Every causal row runs in one call on the unsplit sparse decode kernel, reading the fp8 SWA and nvfp4 compressed caches directly; sparse_paged_indices derives each token's slots and lengths from (request, position) only.

Split from #62365 by @aoshen02 and rebased onto current main (190df0b).
Stacked on #65053: only the last commit is new here.

Changes

• tests/kernels/test_compressor_kv_cache.py (+216 / -0)
• tests/kernels/test_top_k_per_row.py (+88 / -0)
• tests/v1/attention/test_deepseek_v4_swa_visible.py (+1087 / -0)
• vllm/model_executor/layers/sparse_attn_indexer.py (+69 / -16)
• vllm/models/deepseek_v41/attention.py (+39 / -0)
• vllm/models/deepseek_v41/common/ops/cache_utils.py (+16 / -2)
• vllm/models/deepseek_v41/compressor.py (+15 / -0)
• vllm/models/deepseek_v41/nvidia/flashmla.py (+138 / -0)
• vllm/models/deepseek_v41/sparse_mla.py (+7 / -0)
• vllm/v1/attention/backends/mla/indexer.py (+7 / -0)

Rebase onto main 190df0b: Sparse_attn_indexer: main moved the prefill top-k call into the per-chunk branch (#60618, #62240), so the BI top-k is selected there; prefill DeepSelect (new on main) is off under BI; decode keeps FP32 logits when main's takes_bf16_logits rule (#64138) says so or under BI, and the decode logits budget counts the FP32 copy under BI. main's chunked SM100 decode (#64806) is off under BI. main's fused prefill top-k already refuses BI. flashmla: main's num_prefill_tokens. Tests: type fixes for main's Optional topk_indices_buffer and ClassVar PREFILL_CHUNK_SIZE; the decode-vs-prefill test's model config stub gets main's hf_text_config.

Validation

One B300 GPU. The tests this PR adds or edits, on the parent commit with this PR's test files (before) and on this PR (after), measured with the stack on main b8c04dd. Rebasing it since then changed one import line (DSpark PR), the MXFP8 kernel tuple next to main's (tensor-core tiers PR) and the placement of new tests; in the indexer PR, main's new chunked SM100 decode (#64806) is off under BI, which keeps the BI decode path that was measured. These rebases were not rerun on GPU.

pytest tests/kernels/test_compressor_kv_cache.py -k "<selection below>"

before:

E       AssertionError: assert False == True
3 failed, 10 passed, 197 deselected in 16.11s

after:

13 passed, 197 deselected in 14.80s

pytest tests/kernels/test_top_k_per_row.py -k "<selection below>"

before:

19 passed, 860 deselected in 15.19s

after:

19 passed, 860 deselected in 12.95s

pytest tests/v1/attention/test_deepseek_v4_swa_visible.py -k "<selection below>"

before:

E       AssertionError: assert ('fp8_ds_mla'... 'fp8_ds_mla') == ('nvfp4_ds_ml...nvfp4_ds_mla')
8 failed, 22 passed, 109 deselected, 4 warnings in 91.81s (0:01:31)

after:

30 passed, 109 deselected in 99.70s (0:01:39)

-k selections

tests/kernels/test_compressor_kv_cache.py:

test_v41_bi_compressor_rejects_unproven_ring_geometry or test_v41_compressor_bi_batch_chunk_pd_graph_cache_bytes or test_v41_compressor_bi_declaration_stays_scoped or test_v41_pcp_compressor_preserves_pairs_and_decode_state

tests/kernels/test_top_k_per_row.py:

test_bi_bf16_indexer_cast_preserves_valid_scores_and_selector or test_candidate_block_selection_is_row_independent or test_sparse_indexer_resolved_persistent_bf16

tests/v1/attention/test_deepseek_v4_swa_visible.py:

test_ds41_bi_indexer_decode_rows_match_prefill_rows or test_ds41_bi_nvfp4_gather_matches_paged_record_and_graph or test_ds41_cr2_indexer_bi_declaration_stays_scoped or test_ds41_flashmla_bi_auto_kv_cache_dtype or test_ds41_flashmla_cr2_bi_builder_constructs_on_cpu or test_ds41_flashmla_cr2_indexer_bi_batch_chunk_pd_graph or test_ds41_flashmla_swa_bi_batch_chunk_pd_graph or test_ds41_flashmla_swa_bi_builder_policy or test_ds41_swa_bi_declaration_stays_scoped or test_prefix_cache_replays_the_swa_window or test_v41_bi_sparse_paged_indices

pre-commit ruff and mypy pass on the changed files.

Not tested: a full DeepSeek-V4.1 model.

---

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.

[mergify]

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

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, @0z5a.

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