diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py index 92d527e625..42accbef85 100644 --- a/tests/model_executor/test_qwen3_omni.py +++ b/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 diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index c2c252280d..7318d225f8 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/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 "