[Bug] DSpark speculative decoding cannot start on SM120: decode-dsv4 has no topk=192 instantiation, so verify falls through to the prefill kernel's num_tokens > 64 assert

[ressl]

Summary

On SM120 (RTX PRO 6000 Blackwell), DSpark speculative decoding for
DeepSeek-V4-Flash can never start. Every launch dies during CUDA graph capture
with:

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

The assertion can never be satisfied, because speculative verify submits
batch x (gamma + 1) tokens, which is below 64 by construction.

Root cause

FlashInfer instantiates sparse_mla_sm120_decode_dsv4 only for a fixed
(num_heads, topk) table (_DECODE_DSV4_DISPATCH). For 64 heads that is
{128, 512, 1024}.

I instrumented _flash_mla_flashinfer in
python/sglang/kernels/ops/attention/flash_mla_sm120.py and logged every call
over a full startup. The split is clean:

topk	_decode_dsv4_dispatchable	calls
128	True	4644
192	False	4

Normal decode always arrives with topk=128 and dispatches to the decode
kernel. DSpark's draft attention arrives with topk=192, which has no
instantiation, so _decode_dsv4_dispatchable returns False and the call falls
through to module.sparse_mla_sm120_paged_attention, the prefill orchestrator,
whose first check is num_tokens > 64.

Every other dispatch input is fine: model_type=_MODEL_TYPE_DSV4, d_qk=512,
kv_pbs=64 after the existing page-split, num_heads=64.

What is not the cause

• Not the CUDA graph ladder. --cuda-graph-max-bs-decode of 1, 2 and 10
  all fail identically, only with different token counts (6, 12, 50).
• Not gamma. --speculative-dspark-block-size 5 matches the checkpoint's
  own dspark_block_size: 5 in config.json.
• Not the attention backend. SGLANG_SM120_FLASHMLA_BACKEND=triton clears
  capture but aborts on the first real prefill.

Reproduction

sglang serve --model-path <DeepSeek-V4-Flash-0731> \
  --tp-size 4 --kv-cache-dtype fp8_e4m3 \
  --max-running-requests 10 --cuda-graph-max-bs-decode 10 \
  --speculative-algorithm DSPARK --speculative-dspark-block-size 5

Hardware: 4x RTX PRO 6000 Blackwell (SM120), PCIe-only, driver 610.57.04,
CUDA 13.3. SGLang 0.0.0.dev1+ga358374ae. The same code path is present on
main today (flash_mla_sm120.py, the _sparse_mla_sm120_paged_attention
call after the B <= _FI_DECODE_MAX_TOKENS block).

Fix

Two options:

1. Instantiate decode-dsv4 for topk=192 in FlashInfer.
2. Have the SGLang shim pad the sparse index up to the next instantiated width
   and bound the real length with topk_length, which the kernel already
   accepts.

I implemented (2) and will open a PR. Measured on the hardware above:

Metric	before	after
single-stream decode, natural text	80.8 tok/s	223.7 tok/s
aggregate decode, 10 concurrent	429.2 tok/s	465.2 tok/s
acceptance rate	n/a	0.62 (accept len 4.10)

Correctness check: padding with two different values (first index vs last
index of the row) produces byte-identical output across five deterministic
fixtures, which confirms topk_length bounds the read and the padding is never
consumed.

[ressl]

Update: I withdrew the PR (#33986).

Padding the index up to an instantiated width does make DSpark start and run fast (up to 3.3x single-stream on this hardware), but after twelve hours in production it turned out to intermittently corrupt generated text -- stray characters in longer outputs. Rolled back.

So the workaround in that PR should not be used. The bug report itself is unchanged and reproducible: _decode_dsv4_dispatchable returns False for topk=192, the call falls through to sparse_mla_sm120_paged_attention, and that asserts num_tokens > 64 -- which speculative verify, submitting batch * (gamma + 1) tokens, can never satisfy.

Instantiating decode-dsv4 for topk=192 still looks like the correct fix.

[secinto]

Data point from 2x RTX PRO 6000 Blackwell (SM120, TP=2, no NVLink), DeepSeek-V4-Flash-0731: the missing-instantiation half of this issue is fixed by flashinfer PR #4380, but DSpark output is still incorrect — the bug appears to be in the SM120 decode kernel's verify-path numerics, not just the dispatch gap.

Setup:

• sglang 0.5.17 (stock — the -1-padding workaround discussed here explicitly removed)
• flashinfer built from tag v0.6.18rc3, the first tag containing when doing inference with mutilple gpus,does sglang auto detect and use nccl lib to increase QTS？ #4380 (note: the 0.6.16 and 0.6.17 PyPI wheels do not contain it; their _DECODE_DSV4_DISPATCH is identical to 0.6.15.post1)
• nvidia-cutlass-dsl[cu13]==4.7.0 (the source install pulls cu12 DSL libs by default, which ICEs in the CuTe-DSL rmsnorm on a cu13 stack — worth knowing for anyone reproducing)
• driver 610.57.04 / CUDA UMD 13.3 / torch 2.11.0+cu130, --tp-size 2 --moe-runner-backend flashinfer_mxfp4 --disable-custom-all-reduce, NCCL_P2P_DISABLE=1, --speculative-algorithm DSPARK --cuda-graph-max-bs-decode 8 --max-running-requests 8 --mem-fraction-static 0.94

What's fixed: the launch failure in the OP is gone. Dispatch hits the native per-rank (32, 192) decode shape, there is no num_tokens > 64 fall-through to the prefill orchestrator, CUDA-graph capture succeeds, and the autotuner profiles sparse_mla_sm120_decode_dsv4 normally.

What's still broken: at temperature 0, generations degenerate into token-level repetition after roughly a paragraph (~100+ tokens), deterministically reproducible:

To understand how a B-tree maintains its balance is balanced, its balance is, we can be,
let's, and why it, ... and the, and the, the, the, the, the, the, [repeats to length limit]

Short generations look fine, and long-context retrieval (needle at ~2.6K and ~31.7K prompt tokens) also passes — the failure only surfaces on longer generations, so short smoke tests will not catch it.

Why I think this indicts the verify kernel rather than the acceptance policy: with an all-greedy batch, dspark_components/dspark_verify.py::accept_draft_tokens routes to AcceptGreedy (a draft token is accepted only iff it equals the target's argmax), which is lossless by construction — temp-0 output should be identical to non-speculative decoding. A SPEC=0 launch of the same venv/build is clean (verified over 400-token greedy reads, ~92 tok/s decode). So the target-verify forward through the sparse-MLA SM120 decode path — verify-sized batches (gamma 5, ≤ 6 tokens/req) at topk 192 — appears to produce wrong logits even with the real #4380 instantiations. Consistent with that, the earlier pad-to-512 workaround produced the same corruption signature, so the padding was likely never the root cause.

Happy to provide full launch logs, the validation transcripts, or run patched builds on this hardware.