[Misc][Spec Decode] Detect DeepSeek-V4 DSpark checkpoints from config

[xuebwang-amd]

Purpose

To close #52111.

DeepSeek-V4-Flash-0731 and DeepSeek-V4-Pro-0813 do not have MTP heads — their mtp.* tensors are DSpark drafters. vLLM routes them to DeepSeekV4MTPModel anyway and dies deep in the weight loader:

KeyError: 'model.layers.43.mtp_block.main_norm.weight'   # Flash-0731
KeyError: 'model.layers.61.mtp_block.main_norm.weight'   # Pro-0813

This PR resolves the drafter from config.json instead of the repo name, and turns method="mtp" on a DSpark checkpoint into an actionable startup error.
It is platform-independent: reproduced on MI325X, and on stock vllm/vllm-openai v0.23.0 and v0.27.1 on H100.

Test Plan

pytest tests/config/test_deepseek_v4_dspark_config.py
pytest tests/config/ tests/transformers_utils/test_dspark_mla_config.py
ruff check vllm/config/speculative.py tests/config/test_deepseek_v4_dspark_config.py

Routing is decided from config.json alone, so the decision is asserted end to end through a real ModelConfig + SpeculativeConfig over a synthetic checkpoint directory holding nothing but a config.json.

Covered: explicit mtp rejected on the self-speculation path for both DSpark checkpoint shapes and, for Flash-0731, on the named-draft-model path; plain DeepSeek-V4-Flash still resolving to DeepSeekV4MTPModel; an omitted method
auto-detecting dspark; and the pre-existing block-size minimum. The predicate and the variant mapping are pinned separately as unit tests.

Test Result

tests/config/test_deepseek_v4_dspark_config.py            21 passed
tests/config/ + tests/transformers_utils/test_dspark_mla_config.py
                                                         156 passed, 5 failed
ruff check                                               All checks passed!

TODO: 5 failures

Hardware validation

Hardware validation from the contributors in the thread below:

• 8×MI325X (gfx942, ROCm 7.14) — @zzw09773 reproduced the KeyError on both Flash-0731 (layer 43) and Pro-0813 (layer 61); with this change method="mtp" fails at config validation instead, and the suggested method="dspark", num_speculative_tokens=5 serves Pro-0813 at 50 tok/s prose / 91 tok/s code single-stream, 2.1×/3.8× over no spec decode. Pro-0813's differing dspark_target_layer_ids is why detection keys on presence, not contents.
• 8×H100 SXM (SM90) — @kfastino reproduced it on stock vllm/vllm-openai:v0.27.1 (and v0.23.0) at DP4 + EP, fp8 KV cache, where config validation accepts method="mtp" and every DP engine core then dies on the same KeyError. The patch applies to v0.27.1 in 3 hunks with zero rejects and aborts in ~15 s, before any engine process spawns.
• 4×H200 (SM90) — @kfastino, TP4 + EP on v0.26.0, method="dspark" at k=5 and k=7: 0.89–0.93 acceptance, 2–3× speedup over the base model. This is the path the new error message points users at, on the drafter this PR routes to.

No model-eval results from me: this is config validation only. It converts a crash into a startup error and does not alter any decoding path, nor the weights loaded on paths that already worked.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

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

[calvarado2004]

I reproduced one config-flow gap in this PR and prepared a focused follow-up at calvarado2004@c1953d6.

When a caller explicitly supplies both method="mtp" and model=/path/to/DeepSeek-V4-Flash-0731, the early guard is skipped because model is not None. Later method detection then sees the DSpark schema and silently rewrites the explicit MTP request to DSpark. The follow-up moves the DSpark-schema rejection to immediately after draft ModelConfig construction, before heuristic method detection, so both implicit and explicit model paths fail fast with the intended guidance.

It also adds end-to-end config-flow coverage through real ModelConfig + SpeculativeConfig construction for:

• explicit-path MTP rejection
• implicit-path MTP rejection
• omitted-method DSpark auto-routing
• explicit DSpark preservation at K=5
• K=4 rejection for a block-size-5 checkpoint

The routing fields are shared by the official 0731 and Tom Sarihan NVFP4 variants, so the tests cover the checkpoint contract independently of expert quantization. Validation completed: Ruff check, Ruff format check, Python 3.12 py_compile, and git diff --check. I did not open a competing PR because this is directly in #52165's scope. AI assistance was used; I reviewed the config flow and every changed line. Happy for you to cherry-pick the commit or for me to adjust it to your preferred test shape.

[zzw09773]

Validated this PR on real hardware (8×MI325X, gfx942, ROCm 7.14) against both affected checkpoints — filling in the "coming soon" test-result slot:

Before (both checkpoints, method='mtp'): worker dies deep in the weight loader —

• DeepSeek-V4-Flash-0731: KeyError: 'model.layers.43.mtp_block.main_norm.weight'
• DeepSeek-V4-Pro-0813 (re-uploaded snapshot 72e1d323): KeyError: 'model.layers.61.mtp_block.main_norm.weight'

With this PR applied (the vllm/config/speculative.py change cherry-picked onto 0.26.1rc1.dev668+g3ee2df303), method='mtp' on Flash-0731 now fails fast at config validation with the intended message:

deepseek-ai/DeepSeek-V4-Flash-0731 ships a DSpark drafter rather than an MTP head, so method='mtp' cannot load its weights. Use method='dspark' with num_speculative_tokens >= dspark_block_size (5).

No KeyError, no worker start.

The recommended path in that message is also verified end-to-end on this hardware: {"method": "dspark", "num_speculative_tokens": 5} on DeepSeek-V4-Pro-0813 loads the mtp.{0,1,2}.* drafter cleanly and serves in production at 50 tok/s (prose) / 91 tok/s (code) single-stream, 2.1×/3.8× over no spec decode, with tool calling and 4-way concurrency healthy.

One detection detail worth noting: Pro-0813 carries dspark_target_layer_ids=[58,59,60] / dspark_markov_rank=512 (vs Flash's [40,41,42]/256), and num_nextn_predict_layers=1 on both — so keying on dspark_target_layer_ids rather than the layer count is exactly right and covers both shipped configs. +1 for landing; this closes the trap we hit in #52111.

[kfastino]

NVIDIA H100 validation — complementing the MI325X results above. Hardware: 8×H100 SXM (SM90, one node), deepseek-ai/DeepSeek-V4-Flash-0731 @ 7872f01b, DP=4 + expert parallel, fp8 KV cache, max-model-len 1M.

1. The trap reproduces on a stock NVIDIA release build. vllm/vllm-openai:v0.27.1 with {"method":"mtp","num_speculative_tokens":1}: config validation accepts it, and every DP engine core then dies during weight loading on the main_norm KeyError — the same failure as #52111 (we had also hit it on v0.23.0). The issue is labelled rocm, but it is platform-independent: it's the checkpoint/loader mismatch this PR guards against.

2. This PR's guard works as intended. The vllm/config/speculative.py change cherry-picks onto v0.27.1 with zero rejects (3 hunks). The same method=mtp launch then fails fast at config validation in 15 s, before any engine process starts:

Value error, deepseek-ai/DeepSeek-V4-Flash-0731 ships a DSpark drafter rather than an MTP head, so method='mtp' cannot load its weights. Use method='dspark' with num_speculative_tokens >= dspark_block_size (5).

3. One caveat on the guidance message for Hopper + current release. Following the message's advice — {"method":"dspark","num_speculative_tokens":5} — on v0.27.1 + H100 crashes during CUDA graph capture with CUDA_ERROR_ILLEGAL_ADDRESS in DeepGEMM (fp8_gemm_nt / tf32_hc_prenorm_gemm): SM90 support for this spec-decode path landed after the v0.27.1 cut (#52164; plus #47808 for the confidence head — the 0731 drafter loads 96 params on v0.27.1 vs 97 on main). On nightly-ac7509e2b1db the identical config comes up clean and serves: greedy smoke test is coherent, and the spec-decode counters confirm parallel drafting at K=5 (draft tokens = 5 × drafts). So the guard's advice is correct on main, but Hopper users on the latest release will trade one crash for another until the next release — may be worth a line in the error message or release notes.

One deployment note for 80 GB parts: gpu-memory-utilization 0.90 (fine for the AR-only recipe) OOMs during graph capture once the DSpark drafter is resident; 0.85 works.

+1 to land — this closes the same trap we hit independently.

[kfastino]

Correction to my validation comment above, prompted by additional internal benchmarking data on this checkpoint:

I attributed the method='dspark' crash on v0.27.1 + H100 (CUDA_ERROR_ILLEGAL_ADDRESS in DeepGEMM during CUDA graph capture) to "SM90 support landing post-release" via #52164/#47808. That causal claim was too strong and is partly wrong:

• DSpark serves fine on v0.26.0 on 4×H200 — which is the same SM90 compute capability — at TP4+EP, fp8 KV, k=5/k=7 (0.89–0.93 acceptance, 2–3× over base). So SM90 as such has been supported since at least v0.26.0.
• [Attention][DSA] Take the native decode path for MTP=3 on SM90 #52164 is narrowly the native decode path for MTP=3 on SM90, and k=3 is invalid on this checkpoint anyway (dspark_block_size=5).

What I actually measured stands, but more narrowly: v0.27.1 + method='dspark' + --data-parallel-size 4 on 8×H100 crashes during capture, and the same config on nightly ac7509e2b1db serves cleanly. I did not bisect; the cause could be DP-specific, a v0.27.x regression, or config-dependent (the working v0.26.0 run above was TP4, not DP4). Happy to bisect if useful to anyone — but the practical guidance for this PR's error message is unchanged: latest-release users may need a nightly or specific topology for method='dspark' on Hopper, so a pointer in the message or release notes still seems worthwhile.

[mergify]

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

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