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

tests/model_executor/test_qwen3_omni.py

@@ -293,6 +293,25 @@ def test_qwen3_omni_text_model_collects_post_deepstack_aux_hidden_states():
     torch.testing.assert_close(aux_hidden_states[0], torch.tensor([[15.0]]))
 
 
+def _dspark_vocab_stub(input_vocab_size: int, draft_vocab_size: int):
+    """A Qwen3 DSpark model with only what `load_weights` reads."""
+    from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM
+
+    model = Qwen3DSparkForCausalLM.__new__(Qwen3DSparkForCausalLM)
+    nn.Module.__init__(model)
+    object.__setattr__(
+        model,
+        "config",
+        SimpleNamespace(
+            vocab_size=input_vocab_size,
+            draft_vocab_size=draft_vocab_size,
+        ),
+    )
+    object.__setattr__(model, "target_vocab_size", 100)
+    object.__setattr__(model, "model", Mock(confidence_head=None))
+    return model
+
+
 @pytest.mark.skip_global_cleanup
 @pytest.mark.parametrize(
     ("input_vocab_size", "draft_vocab_size", "weights", "error"),
@@ -311,24 +330,25 @@ def test_qwen3_omni_text_model_collects_post_deepstack_aux_hidden_states():
 def test_qwen3_dspark_rejects_incomplete_vocab_weights(
     input_vocab_size, draft_vocab_size, weights, error
 ):
-    from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM
-
-    model = Qwen3DSparkForCausalLM.__new__(Qwen3DSparkForCausalLM)
-    nn.Module.__init__(model)
-    object.__setattr__(
-        model,
-        "config",
-        SimpleNamespace(
-            vocab_size=input_vocab_size,
-            draft_vocab_size=draft_vocab_size,
-        ),
-    )
-    object.__setattr__(model, "target_vocab_size", 100)
+    model = _dspark_vocab_stub(input_vocab_size, draft_vocab_size)
 
     with pytest.raises(ValueError, match=error):
         model.load_weights(weights)
 
 
+@pytest.mark.skip_global_cleanup
+@pytest.mark.parametrize("draft_vocab_size", [100, 101])
+def test_qwen3_dspark_accepts_draft_vocab_at_least_target(draft_vocab_size):
+    """A draft vocab >= the target's maps identically, so needs no lm_head/d2t.
+
+    Guards the padded case (draft > target): the draft's unembedding is padded
+    to its physical size, which is not a reduced vocabulary.
+    """
+    model = _dspark_vocab_stub(100, draft_vocab_size)
+
+    model.load_weights([])
+
+
 @pytest.mark.skip_global_cleanup
 def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary():
     from vllm.v1.worker.gpu.spec_decode.dspark import utils as dspark_utils

vllm/model_executor/models/qwen3_dspark.py

@@ -254,7 +254,10 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
             self.config.draft_vocab_size, scale=logit_scale
         )
         self.target_vocab_size = vllm_config.model_config.get_vocab_size()
-        if self.config.draft_vocab_size != self.target_vocab_size:
+        # Only a pruned draft (subset of the target vocab) needs a d2t mapping to
+        # scatter ids up. A draft >= the target vocab (e.g. padded to the target's
+        # physical embedding size) maps identically, so no d2t is needed.
+        if self.config.draft_vocab_size < self.target_vocab_size:
             self.draft_id_to_target_id = nn.Parameter(
                 torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
                 requires_grad=False,
@@ -334,7 +337,7 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
         # embed_tokens / lm_head are optional; when omitted they are shared from
         # the target by load_dspark_model, so skip the unloaded params here.
         uses_expanded_input_vocab = self.config.vocab_size > self.target_vocab_size
-        uses_reduced_vocab = self.config.draft_vocab_size != self.target_vocab_size
+        uses_reduced_vocab = self.config.draft_vocab_size < self.target_vocab_size
         if uses_expanded_input_vocab and not includes_embed_tokens:
             raise ValueError(
                 "Qwen3 DSpark checkpoints whose input vocab_size is larger than "