[Bugfix][Spec Decode] Separate DSpark width from MTP stage validation (#54631)
Signed-off-by: Luca Motz <luca.motz@icloud.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
lucamotz committed
Sep 24, 2026 at 08:08 UTC
ee2f3cdf127e999c2f5531b4e9d2399b08a21433
2 files changed
+96
-16
tests/config/test_speculative_draft_hf_overrides.py
+85
@@ -47,6 +47,91 @@ def test_none_overrides_fall_back_to_arch_mapping():
47
assert composed is SpeculativeConfig.hf_config_override
48
49
50
+def _make_speculative_config(
51
+ hf_kwargs: dict, architecture: str, **speculative_kwargs
52
+) -> SpeculativeConfig:
53
+ hf_config = SpeculativeConfig.hf_config_override(_make_hf_config(**hf_kwargs))
54
+ draft = MagicMock(
55
+ hf_config=hf_config, architectures=hf_config.architectures, max_model_len=128
56
+ )
57
+ draft.registry.inspect_model_cls.return_value = (None, architecture)
58
+ target = MagicMock(max_model_len=128, quantization=None, hf_overrides={})
59
+ target.hf_config.model_type = hf_kwargs["model_type"]
60
+
61
+ with patch("vllm.config.speculative.ModelConfig", return_value=draft):
62
+ return SpeculativeConfig(
63
+ model="draft",
64
+ **speculative_kwargs,
65
+ target_model_config=target,
66
+ target_parallel_config=ParallelConfig(),
67
+ )
68
+
69
+
70
+@pytest.mark.cpu_test
71
+@pytest.mark.parametrize(
72
+ ("hf_kwargs", "n_predict", "architecture"),
73
+ [
74
+ (
75
+ dict(
76
+ model_type="deepseek_v4",
77
+ num_nextn_predict_layers=3,
78
+ dspark_block_size=5,
79
+ ),
80
+ 3,
81
+ "DSparkDraftModel",
82
+ ),
83
+ (
84
+ dict(
85
+ model_type="deepseek_v41",
86
+ num_nextn_predict_layers=3,
87
+ dspark_block_size=5,
88
+ ),
89
+ 3,
90
+ "DSparkV41DraftModel",
91
+ ),
92
+ (
93
+ dict(
94
+ model_type="gemma4_text",
95
+ architectures=["Gemma4DSparkModel"],
96
+ block_size=5,
97
+ ),
98
+ None,
99
+ "Gemma4DSparkModel",
100
+ ),
101
+ ],
102
+)
103
+@pytest.mark.parametrize("speculative_kwargs", [{"num_speculative_tokens": 5}, {}])
104
+def test_dspark_width_is_independent_of_mtp_stages(
105
+ hf_kwargs, n_predict, architecture, speculative_kwargs
106
+):
107
+ config = _make_speculative_config(
108
+ hf_kwargs, architecture, method="dspark", **speculative_kwargs
109
+ )
110
+
111
+ assert config.num_speculative_tokens == 5
112
+ assert getattr(config.draft_model_config.hf_config, "n_predict", None) == n_predict
113
+ assert config.draft_model_config.hf_config.architectures == [architecture]
114
+
115
+
116
+@pytest.mark.cpu_test
117
+def test_mtp_stages_are_independent_of_dspark_width():
118
+ hf_kwargs = dict(
119
+ model_type="deepseek_v4", num_nextn_predict_layers=3, dspark_block_size=5
120
+ )
121
+ config = _make_speculative_config(
122
+ hf_kwargs, "DeepSeekV4MTPModel", method="mtp", num_speculative_tokens=6
123
+ )
124
+
125
+ assert config.num_speculative_tokens == 6
126
+ assert config.draft_model_config.hf_config.n_predict == 3
127
+ assert config.draft_model_config.hf_config.architectures == ["DeepSeekV4MTPModel"]
128
+
129
+ with pytest.raises(ValueError, match="must be divisible by n_predict=3"):
130
+ _make_speculative_config(
131
+ hf_kwargs, "DeepSeekV4MTPModel", method="mtp", num_speculative_tokens=5
132
+ )
133
+
134
+
135
@pytest.mark.cpu_test
136
def test_callable_overrides_reach_the_draft_config():
137
"""A callable override (config-to-config transform) composes with the
vllm/config/speculative.py
+11
-16
@@ -694,8 +694,8 @@ class SpeculativeConfig:
694
if hf_config.model_type in ("deepseek_v4", "deepseek_v41"):
695
# V4.1 has no classic-MTP draft: its checkpoints ship DSpark stages
696
# under ``mtp.*``, so only V4 gets an MTP architecture here. The
697
- # DSpark path rewrites ``architectures`` itself and needs only
698
- # ``n_predict``; ``method="mtp"`` on V4.1 is rejected below.
697
+ # DSpark path rewrites ``architectures`` itself;
698
+ # ``method="mtp"`` on V4.1 is rejected below.
699
is_v41 = hf_config.model_type == "deepseek_v41"
700
hf_config.model_type = "deepseek_mtp"
701
n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
@@ -1440,14 +1440,6 @@ class SpeculativeConfig:
1440
draft_hf_config.architectures = [
1441
"DSparkV41DraftModel" if is_v41 else "DSparkDraftModel"
1442
]
1443
- if is_v41:
1444
- # hf_config_override set n_predict to the number of
1445
- # MTP stages (3), but one DSpark round drafts
1446
- # dspark_block_size tokens; num_speculative_tokens
1447
- # divisibility is checked against n_predict below.
1448
- draft_hf_config.n_predict = getattr(
1449
- draft_hf_config, "dspark_block_size", None
1450
- ) or getattr(draft_hf_config, "n_predict", None)
1443
self.draft_model_config.quantization = (
1444
self.target_model_config.quantization
1445
)
@@ -1464,11 +1456,6 @@ class SpeculativeConfig:
1456
and getattr(hf, "target_layer_ids", None) is not None
1457
):
1458
hf.dspark_target_layer_ids = hf.target_layer_ids
1467
- if (
1468
- getattr(hf, "n_predict", None) is None
1469
- and getattr(hf, "block_size", None) is not None
1470
- ):
1471
- hf.n_predict = hf.block_size
1459
1460
if self.method in ("dflash", "dspark"):
1461
self.parallel_drafting = True
@@ -1483,7 +1470,15 @@ class SpeculativeConfig:
1470
n_predict = getattr(
1471
self.draft_model_config.hf_config, "n_predict", None
1472
)
1486
- if n_predict is not None:
1473
+ if self.use_dspark():
1474
+ if self.num_speculative_tokens is None:
1475
+ # DSpark's parallel width is independent of MTP stages.
1476
+ hf_config = self.draft_model_config.hf_config
1477
+ block_size = getattr(hf_config, "block_size", None)
1478
+ if block_size is None:
1479
+ block_size = getattr(hf_config, "dspark_block_size", None)
1480
+ self.num_speculative_tokens = block_size
1481
+ elif n_predict is not None:
1482
if self.num_speculative_tokens is None:
1483
# Default to max value defined in draft model config.
1484
self.num_speculative_tokens = n_predict