fix(dspark): preserve target quantization for folded drafts

[alexbi29]

Summary

• keep the target quantization config for folded DSparkDraftModel checkpoints
• continue using the draft quantization config for standalone DSpark checkpoints
• extend the existing DSpark load test to cover both cases

Root cause

load_dspark_model now replaces the copied VllmConfig.quant_config with the draft model quantization config. That is correct for standalone drafts, but DeepSeek V4 uses a folded draft whose weights live in the target checkpoint and whose model-specific DeepseekV4FP8Config selects FP4/MXFP4 experts.

The folded draft ModelConfig resolves a generic FP8 config. Applying it to the whole draft registers w13_weight_scale_inv/w2_weight_scale_inv, while the folded target checkpoint supplies FP4 expert scales and the DSpark loader maps them to w13_weight_scale/w2_weight_scale. Cold startup then fails with a missing w13_weight_scale parameter.

DSparkDraftModel is the explicit architecture marker assigned to this folded DeepSeek path, so use it to retain the copied target quantization config. Other DSpark architectures keep the new standalone-draft behavior.

Tests

• python -m pytest -q tests/v1/spec_decode/test_dspark.py -k load_dspark_model_shares_direct_draft_embedding_and_head (2 passed)
• DeepSeek-V4-Flash-0731 TP2/EP folded-DSpark cold boot and structured tool-call smoke passed on SM120

[alexbi29]

@jasl This is ready for review. The pre-run gate requires a ready/verified label because the author currently has 3 merged PRs in this repo, and I do not have permission to apply repository labels. The focused folded-vs-standalone DSpark test passes (2 cases), and the folded DeepSeek-V4 cold boot was validated on TP2/EP SM120.

[jasl]

Reviewed and reproduced on GB10 (SM121) 2-node TP=2. This is correct and I'm merging it. Thank you — you also found a red test of ours that our own gates were not running.

The bug reproduces exactly

Controlled A/B, same node pair, same runner, same spec shape, patch as the only variable:

arm	patch present	serve rc	V2 runner active	result
72f5a30158	NO	6	yes (2 log lines)	KeyError: 'layers.0.ffn.experts.routed_experts.w13_weight_scale' at param = params_dict[name_mapped], 45/48 shards in
8cd16958d0	YES	0	yes (3 log lines)	serves; #19 instruction-following gate PASS, free-form generation coherent

Word for word the failure you described.

One correction to the PR description

The trigger is not "the draft is folded" — it is "SpeculativeConfig.quantization was not inherited", and those are different sets.

With the implicit shape — {"method":"dspark","num_speculative_tokens":5}, no model key —
speculative.py:775-776 does if not self.quantization: self.quantization = self.target_model_config.quantization, which resolves to deepseek_v4_fp8. The draft then gets a DeepseekV4FP8Config with expert_dtype="fp4" as well, so the pre-patch line was already harmless. I confirmed that on hardware: V2 + implicit DSpark + 0731 boots rc=0 both patched and unpatched.

The failure needs the explicit shape, {"method":"dspark","model":"deepseek-ai/DeepSeek-V4-Flash-0731",...}. There self.quantization stays None (it is assigned only at speculative.py:769 and :776, both inside the self.model is None branch), hf_config_override rewrites model_type to deepseek_mtp before _verify_quantization runs, so DeepseekV4FP8Config.override_quantization_method never fires and the draft falls back to a plain Fp8Config → Fp8MoEMethod → w13_weight_scale_inv, against a checkpoint whose expert scales the DSpark loader maps to w13_weight_scale (dspark.py:1003-1007).

This matters for anyone reading the commit later: your discriminator is a superset of the failing condition. That is fine — it is safe, it is derived from the same place that decides foldedness (speculative.py:1025-1036 assigns architectures=["DSparkDraftModel"] to exactly the DSv4 folded case), and it is a no-op wherever quantization was inherited. But the reason it works is not the reason the description gives. I'll note this in the merge commit rather than ask you to rewrite it.

The test is not vacuous

Checked directly rather than assumed — with your test but utils.py reverted:

FAILED ...test_load_dspark_model_shares_direct_draft_embedding_and_head[DSparkDraftModel-True]
E   AssertionError: assert <object ...> is <object ...>
1 failed, 1 passed

Exactly the right shape: only the folded parametrization fails, Qwen3DSparkModel-False still passes, so the standalone path is provably untouched.

You found a test we had been shipping red

tests/v1/spec_decode/test_dspark.py was already failing on our shipped head 72f5a30158 with AttributeError: 'object' object has no attribute 'hf_config' — broken by upstream ecf4aa5ce2, which made load_dspark_model read draft_model_config.hf_config while the test's stub still passed a bare object(). Your rewritten stub fixes it.

It went unnoticed because that file was not in our unit list at all. It is now (scripts/run_dsv4_unit_suite.sh), along with the suite itself, which until today lived only as a hand-copied scratch file on the test nodes.

One residual, not a blocker

The patch leaves dspark.py:702 untouched, so after it the draft's MoE uses the target's DeepseekV4FP8Config while self.quant_config still comes from get_draft_quant_config. Those two disagree outside the MoE: DeepseekV4FP8Config.is_scale_e8m0 is True for expert_dtype="fp4" (quant_config.py:81-84) and drives scale_dtype=torch.float8_e8m0fnu for FP8 linear layers (fp8.py:366), where a generic Fp8Config has no such attribute and falls back to float32.

So the patch trades a hard crash for a possible subtle inconsistency. That is the right trade, and empirically it does not bite: the patched explicit-model serve passes the #19 gate and generates coherently. I am not asking you to extend the fix to :702 here — that file is shared by the V1 and V2 runners, so it would enlarge the blast radius of a patch that currently cannot affect V1 at all. Worth a separate PR if anyone hits it.

Scope for other users of this branch

load_dspark_model is reached only from vllm/v1/worker/gpu/spec_decode/dspark/speculator.py:94, i.e. the V2 model runner. The V1 runner we default to uses vllm/v1/spec_decode/dspark.py and never touches it. So this is invisible on the default path — which is exactly why it sat here unnoticed, and exactly why "does it affect our config?" is the wrong question on a shared branch.

[jasl]

Merged as e30c240fbb, now on codex/ds4-sm120-min-enable and ds4-sm120-preview-dev at 67f0057da4, tagged sm120-pr-41834-stable-preview-20260802d.

Closing manually rather than by GitHub's auto-close: the branch had moved on (a 41-commit upstream merge plus a FlashInfer 0.6.16 bump landed while this was open), so it went in through wip/merge-234 first and GitHub did not match the commits. Your authorship is preserved on the commit.

Full review is in the comment above — the short version is that the fix is right, the test is non-vacuous, and the root cause is one step to the left of where the description puts it (the trigger is an explicitly-named speculative model, not foldedness itself; the implicit shape inherits deepseek_v4_fp8 from the target and the patch is a no-op there).

Thank you especially for the test rewrite. tests/v1/spec_decode/test_dspark.py turned out to be red on our shipped heads and had been for a while — your stub fix repaired one failure, and running the file for the first time exposed two more (an atol set below one bf16 ULP, since fixed with the ULP measurement written into the commit). None of that was visible to us because the file was not in any of our gates. It is now, along with the suite itself, which until this week lived only as a hand-copied script on the test nodes.