[Spec Decode] Fix Qwen3 DSpark d2t requirement for padded-vocab drafts (#55133)
Signed-off-by: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Orestis Zambounis committed
Sep 14, 2026 at 16:07 UTC
c676e4930bb3c87b97790520e7979d43887ca5dd
2 files changed
+38
-15
tests/model_executor/test_qwen3_omni.py
+33
-13
@@ -293,6 +293,25 @@ def test_qwen3_omni_text_model_collects_post_deepstack_aux_hidden_states():
293
torch.testing.assert_close(aux_hidden_states[0], torch.tensor([[15.0]]))
294
295
296
+def _dspark_vocab_stub(input_vocab_size: int, draft_vocab_size: int):
297
+ """A Qwen3 DSpark model with only what `load_weights` reads."""
298
+ from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM
299
+
300
+ model = Qwen3DSparkForCausalLM.__new__(Qwen3DSparkForCausalLM)
301
+ nn.Module.__init__(model)
302
+ object.__setattr__(
303
+ model,
304
+ "config",
305
+ SimpleNamespace(
306
+ vocab_size=input_vocab_size,
307
+ draft_vocab_size=draft_vocab_size,
308
+ ),
309
+ )
310
+ object.__setattr__(model, "target_vocab_size", 100)
311
+ object.__setattr__(model, "model", Mock(confidence_head=None))
312
+ return model
313
+
314
+
315
@pytest.mark.skip_global_cleanup
316
@pytest.mark.parametrize(
317
("input_vocab_size", "draft_vocab_size", "weights", "error"),
@@ -311,24 +330,25 @@ def test_qwen3_omni_text_model_collects_post_deepstack_aux_hidden_states():
330
def test_qwen3_dspark_rejects_incomplete_vocab_weights(
331
input_vocab_size, draft_vocab_size, weights, error
332
):
314
- from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM
315
-
316
- model = Qwen3DSparkForCausalLM.__new__(Qwen3DSparkForCausalLM)
317
- nn.Module.__init__(model)
318
- object.__setattr__(
319
- model,
320
- "config",
321
- SimpleNamespace(
322
- vocab_size=input_vocab_size,
323
- draft_vocab_size=draft_vocab_size,
324
- ),
325
- )
326
- object.__setattr__(model, "target_vocab_size", 100)
333
+ model = _dspark_vocab_stub(input_vocab_size, draft_vocab_size)
334
335
with pytest.raises(ValueError, match=error):
336
model.load_weights(weights)
337
338
339
+@pytest.mark.skip_global_cleanup
340
+@pytest.mark.parametrize("draft_vocab_size", [100, 101])
341
+def test_qwen3_dspark_accepts_draft_vocab_at_least_target(draft_vocab_size):
342
+ """A draft vocab >= the target's maps identically, so needs no lm_head/d2t.
343
+
344
+ Guards the padded case (draft > target): the draft's unembedding is padded
345
+ to its physical size, which is not a reduced vocabulary.
346
+ """
347
+ model = _dspark_vocab_stub(100, draft_vocab_size)
348
+
349
+ model.load_weights([])
350
+
351
+
352
@pytest.mark.skip_global_cleanup
353
def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary():
354
from vllm.v1.worker.gpu.spec_decode.dspark import utils as dspark_utils
vllm/model_executor/models/qwen3_dspark.py
+5
-2
@@ -254,7 +254,10 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
254
self.config.draft_vocab_size, scale=logit_scale
255
)
256
self.target_vocab_size = vllm_config.model_config.get_vocab_size()
257
- if self.config.draft_vocab_size != self.target_vocab_size:
257
+ # Only a pruned draft (subset of the target vocab) needs a d2t mapping to
258
+ # scatter ids up. A draft >= the target vocab (e.g. padded to the target's
259
+ # physical embedding size) maps identically, so no d2t is needed.
260
+ if self.config.draft_vocab_size < self.target_vocab_size:
261
self.draft_id_to_target_id = nn.Parameter(
262
torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
263
requires_grad=False,
@@ -334,7 +337,7 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
337
# embed_tokens / lm_head are optional; when omitted they are shared from
338
# the target by load_dspark_model, so skip the unloaded params here.
339
uses_expanded_input_vocab = self.config.vocab_size > self.target_vocab_size
337
- uses_reduced_vocab = self.config.draft_vocab_size != self.target_vocab_size
340
+ uses_reduced_vocab = self.config.draft_vocab_size < self.target_vocab_size
341
if uses_expanded_input_vocab and not includes_embed_tokens:
342
raise ValueError(
343
"Qwen3 DSpark checkpoints whose input vocab_size is larger than "