Fix DSPARK SM120 decode dispatch for non-instantiated topk widths

[hassellof]

Independent reproductions (added 2026-08-06)

Three environments beyond the original report now hit this crash, and one open PR carries this patch as a prerequisite:

• 8× RTX 5090 (sm_120), TP=8 — @ryou315 reproduced the verbatim num_tokens=5 vs 64 assert on sglang 0.5.16 + DeepSeek-V4-Flash-0731 (comment). SGLANG_SM120_FLASHMLA_BACKEND=triton avoids the boot crash but routes all sparse-MLA calls — including normal decode at topk=2048 — through Triton, so it costs throughput on the hot path.
• DGX Spark (sm_121), TP=2 — [Bug] DeepSeek-V4-Flash-0731 DSPARK on 2x DGX Spark (sm_121, TP=2): sparse-MLA prefill rejects topk=192 (config index_topk=512; kernel buckets 128/512/1024/2048) #33134 reports the same topk=192 dispatch gap on the prefill side.
• 2×/4× RTX PRO 6000 (sm_120) — the original report here.
• [Spec] Fix Dspark and Dflash state divergence across TP rank #33614 depends on this patch to run at all: "Add Fix DSPARK SM120 decode dispatch for non-instantiated topk widths #33407 code. It fixes an unrelated bug for SM12x sparse-MLA topk dispatch crash at boot. Dspark + DeepSeek V4 Flash cannot start on DGX Sparks without it."

So on consumer/workstation Blackwell, DeepSeek's own recommended DSPARK recipe currently cannot boot without either this change or a throughput-costing env override.

Relationship to flashinfer#4309

flashinfer#4309 adds native topk=192 instantiations and is the preferred fast path for that width once it is released and pinned. This change is complementary rather than superseded:

• it fail-softs any non-instantiated width on already released flashinfer (≤ 0.6.15.post1, sglang's current pin), so users are not blocked on a flashinfer release + pin bump;
• draft indexer widths are model-controlled, so 192 will not be the last odd width — flashinfer#4336 already reports a missing (32, 256) decode instance;
• the test suite here includes a bucket-table canary (assertNotIn(192, ...) over the installed decode dispatch table) that goes red exactly when a pinned flashinfer gains 192, flagging the moment the pad can be retired for that width.

Motivation

DeepSeek-V4-Flash-0731 — the official release of DeepSeek-V4-Flash — ships with the DSpark speculative-decoding module bundled, and its official SGLang launch recipe enables it via --speculative-algorithm DSPARK, with the draft weights taken from the same checkpoint.

Running that officially recommended configuration on SM120 (Blackwell workstation/desktop) GPUs crashes the server at boot. DSPARK's draft indexer emits topk=192, which is in no instantiated CUTLASS sparse-MLA bucket (decode: (num_heads, topk) table with topk ∈ {128, 512, 1024}; prefill orchestrator: {128, 512, 1024, 2048}), and the prefill kernel additionally asserts num_tokens > 64. Two boot-time crashes result:

• the warmup draft pass (num_tokens=5) falls through to the prefill kernel:
  Check failed: num_tokens > 64 (5 vs. 64): Decode (num_tokens <= 64) must go through sparse_mla_sm120_decode_dsv3_2 or sparse_mla_sm120_decode_dsv4
• the draft CUDA-graph capture (the draft head runs unsharded, num_heads=64, and capture batches exceed 64 tokens):
  Unsupported sparse-MLA prefill configuration: model=DSV4 num_heads=64 topk=192 page_block_size=64

Same failure family as #33134 (DGX Spark, sm_121, also topk=192); this change may resolve that report as well, though we could not verify on sm_121 hardware.

Modifications

In python/sglang/kernels/ops/attention/flash_mla_sm120.py (_flash_mla_flashinfer):

• Right-pad the sparse-MLA index tensor with -1 (the kernels' documented "skip" sentinel, see flashinfer csrc/sparse_mla_sm120.cu) up to the next instantiated topk bucket, and cap the scan via topk_length so the padding is never read — padded rows add no scan work.
• If a decode-sized batch is still not dispatchable to the CUTLASS decode kernel after padding (e.g. draft head_dim != 512), fall back to the existing Triton sparse-decode kernel for that call instead of routing into the prefill path, which would assert.

Dispatch behavior for already-instantiated shapes is unchanged.

Accuracy Tests

The change only widens the index tensor with skip sentinels (never read; scan capped via topk_length) and reroutes shapes that previously crashed — it does not alter attention math for configurations that already worked.

Reference hardware: 4× NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition (GB202, sm_120, 96 GB GDDR7 each), TP=4, deepseek-ai/DeepSeek-V4-Flash-0731, official DSPARK launch recipe (--speculative-algorithm DSPARK, --moe-runner-backend flashinfer_mxfp4, --disable-flashinfer-autotune, --kv-cache-dtype fp8_e4m3), sampling at DeepSeek's official calibration (temperature 1.0).

The full environment this was verified in — image build, patch list, serving flags and measured numbers — is published at https://github.com/ombori/deepseek-v4-flash-0731-sglang-4x-rtx-pro-6000 (prebuilt image on ghcr.io) for anyone wanting to reproduce on SM120.

• Both boot-time crashes above no longer reproduce; the server boots (including full decode CUDA-graph capture for target-verify and draft) and serves continuously under DSPARK.
• Chat completions, multi-tool-call requests (valid JSON arguments), and streaming tool calls verified correct through an OpenAI-compatible proxy chain.

Speed Tests and Profiling

python -m sglang.bench_serving --dataset-name random --random-input-len 1024 --random-output-len 512 --random-range-ratio 1.0 on the reference hardware above:

max concurrency	output tok/s	mean ITL	accept length
16	697	19.1 ms	3.45
32	995	26.8 ms	5.27
64	1282	43.2 ms	5.19

For reference, the same hardware and workload previously measured 649 output tok/s @ concurrency 32 with a conservative EAGLE configuration (steps 1 / draft 2), which this recipe replaces.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests. Note: the affected dispatch path requires SM120-class hardware plus the DSV4 checkpoint; validated end-to-end on the reference hardware above. Happy to add a hardware-gated test if maintainers can point at the preferred harness for SM12x-only paths.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #31072115451
Latest PR Test (Extra): ❌ Run #31072115243

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[LingZ315]

Independent reproduction on 8x RTX 5090 (SM120), TP=8.

Environment:

• Model: deepseek-ai/DeepSeek-V4-Flash-0731 (local HF snapshot)
• SGLang: 0.5.16
• FlashInfer: 0.6.14
• PyTorch: 2.11.0
• Triton: 3.6.0
• Python: 3.12.3
• NVIDIA driver: 595.58.03
• GPUs: 8x NVIDIA GeForce RTX 5090

Relevant launch configuration:

SGLANG_JIT_DEEPGEMM_PRECOMPILE=0 \
SGLANG_ENABLE_SPEC_V2=1 \
SGLANG_ENABLE_THINKING=1 \
HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 \
python3 -m sglang.launch_server \
  --trust-remote-code \
  --model-path /path/to/DeepSeek-V4-Flash-0731 \
  --tp 8 \
  --moe-runner-backend flashinfer_mxfp4 \
  --mem-fraction-static 0.92 \
  --speculative-algorithm DSPARK \
  --chunked-prefill-size 4096 \
  --cuda-graph-max-bs-decode 8 \
  --disable-flashinfer-autotune \
  --tool-call-parser deepseekv4 \
  --reasoning-parser deepseek-v4 \
  --enable-metrics \
  --enable-cache-report \
  --host 0.0.0.0 --port 8000

The target and bundled draft weights load successfully. SGLang reports:

Initialized DSpark draft runner. attention_backend=dsv4, model=DeepseekV4ForCausalLMDSpark, gamma=5, verify_num_draft_tokens=6

The warmup then crashes deterministically in _flash_mla_flashinfer -> flashinfer/mla/_sparse_mla_sm120.py:

TVM_FFI_ICHECK_GT(num_tokens, 64)
tvm.error.InternalError: Check failed: num_tokens > 64 (5 vs. 64):
Decode (num_tokens <= 64) must go through sparse_mla_sm120_decode_dsv3_2
or sparse_mla_sm120_decode_dsv4; got num_tokens=5

Forcing SGLANG_SM120_FLASHMLA_BACKEND=triton avoids this dispatch failure on the same host and allows the model to start and serve inference (with memory-related concurrency limits configured separately).

This confirms the PR's first boot-time failure on consumer SM120 at TP=8 as well as the reported TP=4 setup.

[hassellof]

Thanks for the independent reproduction — that's a third SM120 topology now
(4x RTX PRO 6000 TP=4, 2x RTX PRO 6000 TP=2 via #33134, and your 8x RTX 5090
TP=8), and your trace is the same dispatch failure: the DSpark draft indexer
emits topk=192, which is in no instantiated CUTLASS decode bucket, so the
call falls through to the prefill orchestrator and trips its
num_tokens > 64 check with the draft's tiny token count.

Two notes that may save you time:

SGLANG_SM120_FLASHMLA_BACKEND=triton works but costs throughput. It
routes every sparse-MLA call to the Triton fallback, including normal decode
(topk=2048), which is the hot path. This PR keeps CUTLASS for every shape
that has an instantiation and only pads the non-instantiated widths
(192 → next bucket, -1 skip sentinels, scan capped via topk_length),
falling back to Triton for shapes that remain non-dispatchable. So you should
be able to drop the env var once this lands.

Unrelated to this dispatch bug, but it will bite you next on that config:
we found that DSpark draft depth gamma=5 — the value your log shows, which
is the checkpoint's own dspark_block_size default and what you get when the
flag is omitted — corrupts generations on SM120 under real load, while
gamma 3, 4, 6 and 7 are all clean. It's not a depth-monotonic effect;
gamma=5 specifically is the broken value, explicitly set or inferred. If you
see repetition loops or token salad after the dispatch fix, set
--speculative-dspark-block-size 4 explicitly. I'm writing that up as a
separate issue with the full sweep and repro; I'll link it here when it's
filed.

Our serving recipe for this hardware class, including the patch set and the
measured numbers, is at
https://github.com/ombori/deepseek-v4-flash-0731-sglang-4x-rtx-pro-6000 if
it's useful for your 5090 setup.

[hassellof]

Hi @ch-wan — could you add the run-ci label here? The only failing checks are the label gate itself (Missing required label 'run-ci'); no tests have actually run yet.

State of the PR:

• Rebased onto current main — conflict-free; neither touched file (kernels/ops/attention/flash_mla_sm120.py, flash_mla_sm120_triton.py) changed upstream since the merge base.
• Lint/isort clean; the bucket-dispatch test passes 15/15 locally. It is hermetic (CPU only, CUDA_VISIBLE_DEVICES=99, kernels replaced by recorders), so it needs no SM120 card to run in CI.

Since this was filed, three independent environments have hit the same boot crash — 8× RTX 5090 TP=8 (repro above), DGX Spark sm_121 (#33134), and the original 4× RTX PRO 6000 TP=4 — and #33614 carries this patch as a prerequisite ("Dspark + DeepSeek V4 Flash cannot start on DGX Sparks without it").

So on consumer/workstation Blackwell, DeepSeek's own recommended DSPARK recipe currently cannot boot without either this change or SGLANG_SM120_FLASHMLA_BACKEND=triton, which routes every sparse-MLA call — including normal decode at topk=2048 — through Triton and costs throughput on the hot path.

[guqiong96]

@hassellof Thanks for the PR that fixes this issue:
Exception: Capture cuda graph failed: Check failed: num_tokens > 64 (20 vs. 64) : Decode (num_tokens <= 64) must go through sparse_mla_sm120_decode_dsv3_2 or sparse_mla_sm120_decode_dsv4; got num_tokens=20
The fix has been applied in https://github.com/guqiong96/Lsglang/releases#release-lsglang-v.1.4.7 and later versions. The latest sglang 0.5.18 still hasn't fixed this issue, but with this PR it runs well.😅