[Bugfix] Make DSV4 sparse MLA work end-to-end for plain decode, MTP, and DSpark - #51538
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
✅ @lucifer1004, CI is now available for this PR.
|
|
/ci run |
|
✅ Triggered Buildkite CI #83052 for commit |
lucifer1004
changed the title
|
/ci retry |
|
✅ Triggered Buildkite CI #83127 for commit |
|
This pull request has merge conflicts that must be resolved before it can be |
|
Validated on SM121 (GB10), two nodes, tensor-parallel. Platform
At TP=2 the sparse-MLA decode shape is The server starts and serves, and speculative decoding stays active. No crash across SpeculationRead from
Throughput divided by mean accepted length gives 13.22, 13.34, 13.42 and 13.60 steps Long contextIsolated single request per point, 320 output tokens.
One change outside this PRThe DeepGEMM revision pinned by FlashInfer requirement, measured rather than assumedWe downgraded to flashinfer 0.6.17 stable on this setup to check whether the nightly is Worth spelling out because the required With What we cannot claimWe measured throughput and acceptance, not accuracy: no evaluation was run. And we have Raw draws, the benchmark harness and the serve/cluster scripts are here, in case any of Happy to re-run anything specific on this hardware if it helps land the PR. |
…ry API Since flashinfer-ai#4380, a decode-form call (num_tokens <= 64) that matches no decode instantiation raises in Python instead of tripping the prefill orchestrator's C++ "num_tokens > 64" assert, but the error is a flat parameter dump pointing at private dispatch tables, so the reader still has to open the source to work out which parameter to change (flashinfer-ai#4541). Make the dispatch-miss error name the mismatch: which of topk, num_heads, page_block_size, or d_qk missed the instantiated set, and which values are available. The old message prefix and shape summary are kept so callers matching on them are unaffected. Add flashinfer.mla.supported_sparse_mla_sm120_configs() so serving frameworks can validate (num_heads, topk, page_block_size) at init time instead of on the first decode request; vLLM currently reads the private tables for this (vllm-project/vllm#51538). Dispatch behavior is unchanged: shapes that previously reached a kernel still do; only the no-kernel error message and the new query API changed. Developed in combination with Claude Fable 5. Closes flashinfer-ai#4541
…ry API Since flashinfer-ai#4380, a decode-form call (num_tokens <= 64) that matches no decode instantiation raises in Python instead of tripping the prefill orchestrator's C++ "num_tokens > 64" assert, but the error is a flat parameter dump pointing at private dispatch tables, so the reader still has to open the source to work out which parameter to change (flashinfer-ai#4541). Make the dispatch-miss error name the mismatch: which of topk, num_heads, page_block_size, or d_qk missed the instantiated set, and which values are available. The old message prefix and shape summary are kept so callers matching on them are unaffected. Add flashinfer.mla.supported_sparse_mla_sm120_configs() so serving frameworks can validate (num_heads, topk, page_block_size) at init time instead of on the first decode request; vLLM currently reads the private tables for this (vllm-project/vllm#51538). Dispatch behavior is unchanged: shapes that previously reached a kernel still do; only the no-kernel error message and the new query API changed. Developed in combination with Claude Fable 5. Closes flashinfer-ai#4541
…ry API Since flashinfer-ai#4380, a decode-form call (num_tokens <= 64) that matches no decode instantiation raises in Python instead of tripping the prefill orchestrator's C++ "num_tokens > 64" assert, but the error is a flat parameter dump pointing at private dispatch tables, so the reader still has to open the source to work out which parameter to change (flashinfer-ai#4541). Make the dispatch-miss error name the mismatch: which of topk, num_heads, page_block_size, or d_qk missed the instantiated set, and which values are available. The old message prefix and shape summary are kept so callers matching on them are unaffected. Add flashinfer.mla.supported_sparse_mla_sm120_configs() so serving frameworks can validate (num_heads, topk, page_block_size) at init time instead of on the first decode request; vLLM currently reads the private tables for this (vllm-project/vllm#51538). Dispatch behavior is unchanged: shapes that previously reached a kernel still do; only the no-kernel error message and the new query API changed. Developed in combination with Claude Fable 5. Closes flashinfer-ai#4541
|
I feel the workspace lane change in MRV2 might not be necessary 🤔 |
…and DSpark (vllm-project#51538) Signed-off-by: Wyett <wyettzeng@gmail.com>
…and DSpark (vllm-project#51538) Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>
…and DSpark (vllm-project#51538) Signed-off-by: khushali9 <khushali.desai9@gmail.com>
Apply the upstream vLLM DSpark fix at build time without changing the pinned vLLM revision. Validate patch applicability against the exact lock in CI. Refs Anemll#11 Upstream: vllm-project/vllm#51538 Signed-off-by: code4me2 <velvetmoon222999@gmail.com>
… PR#51538 dfecbb5) Root cause of production Xid 31 (MMU VIRT_WRITE) on the last PP rank: draft KV slots were computed as block_id*block_size+offset without a null-block guard, so sliding-window block-table entries that carried evicted/padding block 0 produced writes into physical block 0. Also -1-fill sample_idx_mapping so CUDA graph capture padding rows do not scatter into request slot 0. Upstream-PR: vllm-project/vllm#51538 Upstream-Commit: dfecbb52ce1
A FULL-CUDA-graph replay of an MTP batch pads the request count; the padded request has seq_len == 0, and with next_n > 1 the DSA indexer's per-token context length for its first row is seq_len - next_n + 1 < 0 unless clamped. The sparse top-k kernels consume that length as uint32, and the pre-vllm-project#51538 persistent_topk wedged on it (vllm-project#51593). vllm-project#51538 fixed both producer paths and the kernel guard, but the producer clamp had no direct regression test. A downstream fork that predates vllm-project#51538 reproduced the wedge while serving with FlashInfer sparse MLA on SM12x, where it presented as a FlashInfer decode hang (flashinfer-ai/flashinfer#5015): the FlashInfer launch was queued on the same stream behind the already-wedged top-k. Add one parametrised test that builds DeepseekV32IndexerMetadataBuilder for three live requests plus one seq_len-0 padding request with next_n == 2 and requires the padded request's per-token context lengths to be [0, 0] on both the native (B, next_n) path and the flattened uniform-decode kernel path. No production change; this is coverage for the invariant vllm-project#51538 already enforces. Signed-off-by: Linzhang Chen <chenlinzhang1-c@toki.waseda.jp>
7 participants
Purpose
DeepSeek-V4-Flash-0731 could not run reliably through the SM120 sparse MLA backend. This fixes the seven defects that blocked it across all three decode modes -- plain decode, MTP, and DSpark -- verified end-to-end on 8xRTX PRO 6000 Blackwell across in-flight batching and prefill/decode disaggregation.
Commits 1-5 unblock DSpark. Commits 6-7 fix a hang that is not DSpark-specific: it strands any MTP (
next_n > 1) server on this backend once the batch drains, and is a pre-existing defect onmainrather than a regression from this PR. It is filed as #51593, with the full root-cause analysis in this comment.Why this is not a duplicate
This consolidates #51042 after coordination with the maintainers (see #51042 (comment)). @ilmarkov is
Co-authored-byon the SWA-width commit, whoseamd/rocm.pyhunk anddecode_swa_widthfield are his work. Two deliberate design differences are described in that comment.Duplicate checks run for
dspark,sparse MLA SWA width,mxfp4 gemm1_alpha, anddeepseek v4 expert placement; no other open PR covers the remaining four commits.For commits 6-7, checks run for
persistent_topk,indexer seq_lens spec decode, and51593 in:body. Two open PRs are adjacent but do not overlap:seq_len - decode_len=0 - 0=0and is already safe. The uniform and native spec-decode paths that produce the negative length are untouched by it. The two changes are complementary and do not conflict.persistent_topk.cuh(radix histogram reuse after short rows). Worth flagging the interaction: its precondition is a CTA group processing rows that crossRADIX_THRESHOLD, and the out-of-range lengths fixed here cause exactly that spuriously. That PR is still needed for genuinely long rows.What is fixed
SWA widths — non-causal draft batches allocate
decode_swa_indiceswider thanwindow_size, but the FlashInfer DSV4 path reshaped withwindow_sizeand crashed the draft. The dense width is now carried on the metadata. The non-causal width pads to a multiple of 64 (192 for the K=5 shape: 128 sliding-window + 5 draft entries) rather than 128 (256), matching the kernel's 64-entry tile; both dispatch after feat(sm120): consolidate DSV4 sparse MLA top-k 192/256 support flashinfer-ai/flashinfer#4380, and 192 measures 13-16% faster at >=8 tokens.Workspace lanes — the DSpark target and draft CUDA graphs retain workspace views concurrently, so one buffer per ubatch let a resize for one orphan the other's live tensor. The second lane is allocated only for V2 DSpark.
Graph replay and draft KV —
sample_idx_mappingwas zero-filled, so capture executed padding rows that scattered into request slot 0; captured backbone outputs could be freed before replay read their storage; draft KV could be written into physical block 0, the null block. Draft sampling also moves to a disjoint Philox counter range, since the rejection sampler keys both its acceptance uniform and its recovery Gumbel noise by token position.MXFP4 SwiGLU parameters —
FlashInferExpertsinjected the GPT-OSS activation constants (gemm1_alpha=1.702,gemm1_beta=1.0,gemm1_clamp_limit=7.0) for every mxfp4 weight dtype. DeepSeek V4 uses this path under--moe-backend flashinfer_cutlass, so its SwiGLU was evaluated with GPT-OSS constants and generation collapsed. GPT-OSS is unaffected:GptOssMxfp4MoEMethodsupplies the same constants through its quant config, which the added test pins.SM120 gate — a FlashInfer build can expose the sparse MLA decode API without carrying the DSV4 specialization a configuration needs. That now fails at model init with the required
(num_q_heads, top_k)shape instead of an opaque kernel launch failure at the first decode.Negative indexer context lengths under MTP — padded decode slots carry
seq_len == 0, and withnext_n > 1both spec-decode paths computed a negative per-token context length for the first token of each padded request (0 - 2 + 0 + 1 = -1). The sparse-MLA top-k kernels consumelengthsasuint32, so-1is read as ~4.29e9. Clamped at 0 in both paths, matching the variable-length path which already yields0. Withnext_n == 1the expression collapses toseq_len, which is why plain decode never hit this.Top-k kernels hardened against out-of-range lengths —
persistent_topk_kernelcastlengthstouint32before testingRADIX_THRESHOLD, so the bogus ~4.29e9 forced the row onto the multi-CTA radix path. Because thecta_in_group != 0early exit is decided from the host-sidemax_seq_lenwhile the per-row branch reads device memory, the two disagreed and stranded the group leader on the inter-CTA barrier forever — the kernel never retired, the async output-copy event never fired, and the engine hung waiting for a response that was never sent. The row length is now clamped tomin(stride, max_seq_len)before any decision, which also removes an out-of-bounds read into the next row.cooperative_topkcompares signed (so it cannot hang) but then casts touint32, emitting indices0..TopK-1instead of-1padding; clamped as well.Model evaluation
DeepSeek-V4-Flash-0731,
--moe-backend flashinfer_cutlass,--attention-backend FLASHINFER_MLA_SPARSE_DSV4, gsm8kn=1319:Standard rejection sampling is distribution-preserving, and DSpark lands 0.0023 from target-only against a difference-stderr of ~0.009 (0.26 sigma), on the same checkpoint and the same tree.
MTP on the base checkpoint (
num_nextn_predict_layers = 1), the configuration that #51593 hangs on, gsm8kn=1319:That sits inside the 0.9454-0.9492 band spanned by the target-only and DSpark rows above, on the same checkpoint and tree, so the clamp does not move accuracy. No unfixed baseline is quoted for this row: without commits 6-7 the same server wedges under the drain reproducer within 30-180 s (see Test plan), and I did not obtain a clean unfixed eval run to compare against.
Isolating fix 4 — identical configuration with only that commit removed:
The TP4 pair was measured on both FlashInfer revisions with no significant change (DSpark 0.9454 -> 0.9477, target-only 0.9462 -> 0.9454), so the two rows measured on the earlier revision are directly comparable.
Test plan
Result: 71 passed, 141 skipped.
For commits 6-7, a drain-shaped reproducer (waves of concurrent completions that retire unevenly, wedging the unfixed build in 30-180 s) against
deepseek-v4-flashwith--speculative-config '{"method":"mtp","num_speculative_tokens":1}', TP4:The commit-7-only run is the one that shows the kernel guard is load-bearing rather than dead code: with commit 6 in place nothing negative ever reaches the kernel, so it would pass either way.
Root cause was confirmed on a live hang with cuda-gdb: exactly one of 16 CTAs resident, spinning in
wait_ge, witharrival_counter == 1againsttarget_val == 2— a group leader waiting on a peer that had already taken themax_seq_lenearly exit.AI assistance
This work was produced with AI assistance. Every changed line has been reviewed by the submitter, and the tests and evaluations above were executed on the hardware described, with results reported as measured.