[Bugfix][Spec Decode] Avoid temporary Qwen3.5 MTP vocab allocations

tests/model_executor/test_qwen3_5_mtp.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import Mock, patch
+
+import pytest
+from torch import nn
+
+
+def _make_vllm_config():
+    from vllm.config import CompilationMode
+
+    hf_config = Mock()
+    hf_config.tie_word_embeddings = False
+    hf_config.vocab_size = 128
+    hf_config.hidden_size = 64
+    hf_config.num_hidden_layers = 2
+    hf_config.mtp_num_hidden_layers = 1
+    hf_config.rms_norm_eps = 1e-6
+    hf_config.num_experts = 0
+
+    vllm_config = Mock()
+    vllm_config.model_config.hf_text_config = hf_config
+    vllm_config.cache_config.mamba_cache_mode = "align"
+    vllm_config.compilation_config.mode = CompilationMode.NONE
+    vllm_config.quant_config = None
+    return vllm_config
+
+
+@pytest.mark.parametrize(("pp_size", "should_allocate"), [(1, False), (2, True)])
+def test_qwen3_5_mtp_embedding_allocation_depends_on_pp(
+    pp_size: int, should_allocate: bool
+):
+    from vllm.model_executor.models import qwen3_5_mtp
+
+    mock_embedding = Mock(return_value=nn.Identity())
+    with patch.multiple(
+        qwen3_5_mtp,
+        VocabParallelEmbedding=mock_embedding,
+        ColumnParallelLinear=Mock(return_value=nn.Identity()),
+        Qwen3_5DecoderLayer=Mock(return_value=nn.Identity()),
+        Qwen3_5RMSNorm=Mock(return_value=nn.Identity()),
+        get_pp_group=Mock(return_value=Mock(world_size=pp_size)),
+        is_model_fused_shared_expert_compatible=Mock(return_value=False),
+        make_empty_intermediate_tensors_factory=Mock(),
+    ):
+        predictor = qwen3_5_mtp.Qwen3_5MultiTokenPredictor(
+            vllm_config=_make_vllm_config()
+        )
+
+    assert (predictor.embed_tokens is not None) is should_allocate
+    assert mock_embedding.call_count == int(should_allocate)
+
+
+def test_qwen3_5_mtp_skips_shared_vocab_modules_and_weights_for_pp1():
+    from vllm.model_executor.models import qwen3_5_mtp
+
+    mock_lm_head = Mock()
+    with patch.multiple(
+        qwen3_5_mtp,
+        Qwen3_5MultiTokenPredictor=Mock(return_value=nn.Module()),
+        ParallelLMHead=mock_lm_head,
+        LogitsProcessor=Mock(),
+        get_pp_group=Mock(return_value=Mock(world_size=1, is_last_rank=True)),
+    ):
+        model = qwen3_5_mtp.Qwen3_5MTP(vllm_config=_make_vllm_config())
+
+    assert model.lm_head is None
+    mock_lm_head.assert_not_called()
+
+    loader = Mock()
+    loader.load_weights.side_effect = lambda weights: {name for name, _ in weights}
+    weights = [
+        ("language_model.model.embed_tokens.weight", Mock()),
+        ("lm_head.weight", Mock()),
+        ("mtp.fc.weight", Mock()),
+    ]
+    with patch(
+        "vllm.model_executor.models.qwen3_5_mtp.AutoWeightsLoader",
+        return_value=loader,
+    ):
+        assert model.load_weights(weights) == {"model.fc.weight"}

tests/model_executor/test_qwen3_5_quantization.py

@@ -61,6 +61,7 @@ def test_qwen3_5_mtp_lm_head_receives_quant_config():
 
     mock_pp_group = Mock()
     mock_pp_group.is_last_rank = True
+    mock_pp_group.world_size = 2
 
     with (
         patch("vllm.model_executor.models.qwen3_5_mtp.Qwen3_5MultiTokenPredictor"),

vllm/model_executor/models/qwen3_5_mtp.py

@@ -81,10 +81,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.mtp_start_layer_idx = config.num_hidden_layers
         self.num_mtp_layers = getattr(config, "mtp_num_hidden_layers", 1)
 
-        self.embed_tokens = VocabParallelEmbedding(
-            self.vocab_size,
-            config.hidden_size,
-        )
+        # With no pipeline parallelism, this is attached from the target model
+        # after loading. Avoid materializing a temporary full-vocabulary copy.
+        self.embed_tokens: VocabParallelEmbedding | None
+        if get_pp_group().world_size == 1:
+            self.embed_tokens = None
+        else:
+            self.embed_tokens = VocabParallelEmbedding(
+                self.vocab_size,
+                config.hidden_size,
+            )
 
         # Workaround: mtp.fc is stored as BF16 in NVFP4 checkpoints but is
         # missing from hf_quant_config.json exclude_modules. Force unquantized.
@@ -141,6 +147,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         )
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        assert self.embed_tokens is not None, (
+            "embed_tokens must be shared from the target model before inference"
+        )
         return self.embed_tokens(input_ids)
 
     def forward(
@@ -240,14 +249,19 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp")
         )
 
-        if get_pp_group().is_last_rank:
+        self._share_target_vocab_weights = get_pp_group().world_size == 1
+        self.lm_head: ParallelLMHead | PPMissingLayer | None
+        if self._share_target_vocab_weights:
+            self.lm_head = None
+        elif get_pp_group().is_last_rank:
             self.lm_head = ParallelLMHead(
                 config.vocab_size,
                 config.hidden_size,
                 quant_config=self.quant_config,
                 prefix=maybe_prefix(prefix, "lm_head"),
             )
             if config.tie_word_embeddings:
+                assert self.model.embed_tokens is not None
                 self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
         else:
             self.lm_head = PPMissingLayer()
@@ -299,14 +313,23 @@ def compute_logits(
         hidden_states: torch.Tensor,
         spec_step_idx: int = 0,
     ) -> torch.Tensor | None:
+        assert self.lm_head is not None, (
+            "lm_head must be shared from the target model before inference"
+        )
         return self.logits_processor(self.lm_head, hidden_states)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        shared_weight_names = ("embed_tokens", "lm_head")
+
         def remap_weight_names(weights):
             for name, weight in weights:
+                if self._share_target_vocab_weights and any(
+                    key in name for key in shared_weight_names
+                ):
+                    continue
                 if name.startswith("mtp."):
                     name = name.replace("mtp.", "model.")
-                elif any(key in name for key in ["embed_tokens", "lm_head"]):
+                elif any(key in name for key in shared_weight_names):
                     if "embed_tokens" in name:
                         name = name.replace("language_model.", "")
                 else: