[Bugfix][Model][Spec Decode] Defer disposable GLM MTP head

tests/v1/spec_decode/test_mtp.py

@@ -5,6 +5,7 @@
 
 import pytest
 import torch
+import torch.nn as nn
 
 from tests.v1.attention.utils import (
     BatchSpec,
@@ -22,6 +23,7 @@
     VllmConfig,
 )
 from vllm.config.load import LoadConfig
+from vllm.model_executor.models.deepseek_mtp import DeferredLMHead, SharedHead
 from vllm.model_executor.models.llama import LlamaForCausalLM
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -31,6 +33,53 @@
 DEVICE_TYPE = current_platform.device_type
 
 
+def test_shared_head_can_defer_lm_head():
+    config = mock.MagicMock(hidden_size=16, rms_norm_eps=1e-5, vocab_size=32)
+
+    with mock.patch(
+        "vllm.model_executor.models.deepseek_mtp.ParallelLMHead"
+    ) as parallel_lm_head:
+        regular = SharedHead(config, "mtp")
+        deferred = SharedHead(config, "mtp", defer_lm_head=True)
+
+    parallel_lm_head.assert_called_once()
+    assert regular.head is parallel_lm_head.return_value
+    assert isinstance(deferred.head, DeferredLMHead)
+    assert not tuple(deferred.head.parameters())
+    with pytest.raises(RuntimeError, match="was not replaced"):
+        deferred.head(torch.zeros(1, 16))
+
+
+def test_glm_mtp_defers_shared_head():
+    from vllm.models.glm5next.nvidia import mtp
+
+    config = mock.MagicMock(
+        hidden_size=16,
+        rms_norm_eps=1e-5,
+        index_topk=8,
+        index_kpool=4,
+    )
+    vllm_config = mock.MagicMock()
+    vllm_config.speculative_config.draft_model_config.hf_config = config
+    vllm_config.scheduler_config.max_num_batched_tokens = 4
+
+    with (
+        mock.patch.object(
+            mtp, "RMSNorm", side_effect=lambda *args, **kwargs: nn.Identity()
+        ),
+        mock.patch.object(mtp, "Glm5NextDecoderLayer", return_value=nn.Identity()),
+        mock.patch.object(mtp, "SharedHead", return_value=nn.Identity()) as shared_head,
+        mock.patch.object(mtp.current_platform, "device_type", "cpu"),
+    ):
+        mtp.Glm5NextMultiTokenPredictorLayer(vllm_config, "model.layers.1")
+
+    shared_head.assert_called_once_with(
+        config=config,
+        prefix="model.layers.1",
+        defer_lm_head=True,
+    )
+
+
 def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer:
     """Create an MTP proposer with unified model configuration."""
     model_config = ModelConfig(
@@ -70,6 +119,10 @@ def test_mtp_load_model_unified(mock_get_model, mock_get_layers, mock_get_pp_gro
     # Setup mocks
     mock_model = mock.MagicMock()
     mock_model.model.embed_tokens.weight.shape = (131072, 4096)
+    draft_head = mock.MagicMock()
+    mock_model.model.layers = [
+        mock.MagicMock(shared_head=mock.MagicMock(head=draft_head))
+    ]
     mock_get_model.return_value = mock_model
     # MTP does not have its own embed_tokens or lm_head
     # so it should share them with the target model
@@ -111,6 +164,7 @@ class _TargetModelStub(LlamaForCausalLM):
     mock_get_model.assert_called_once()
     # MTP shares lm_head with target model
     assert proposer.model.lm_head == target_model.lm_head
+    assert proposer.model.model.layers[0].shared_head.head is target_model.lm_head
     # MTP shares embed_tokens with target model
     assert proposer.model.model.embed_tokens == target_model.model.embed_tokens
 

vllm/model_executor/models/deepseek_mtp.py

@@ -46,21 +46,31 @@
 )
 
 
+class DeferredLMHead(nn.Module):
+    def forward(self, *args, **kwargs):
+        raise RuntimeError("deferred MTP head was not replaced by the target lm_head")
+
+
 class SharedHead(nn.Module):
     def __init__(
         self,
         config: PretrainedConfig,
         prefix: str,
         quant_config: QuantizationConfig | None = None,
+        *,
+        defer_lm_head: bool = False,
     ) -> None:
         super().__init__()
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.head = ParallelLMHead(
-            config.vocab_size,
-            config.hidden_size,
-            quant_config=quant_config,
-            prefix=maybe_prefix(prefix, "head"),
-        )
+        if defer_lm_head:
+            self.head = DeferredLMHead()
+        else:
+            self.head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "head"),
+            )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return self.norm(hidden_states)

vllm/models/glm5next/nvidia/mtp.py

@@ -42,7 +42,6 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
         assert vllm_config.speculative_config is not None
         config = vllm_config.speculative_config.draft_model_config.hf_config
         self.config = config
-        quant_config = vllm_config.quant_config
 
         self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -66,7 +65,9 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
             device=current_platform.device_type,
         )
         self.shared_head = SharedHead(
-            config=config, prefix=prefix, quant_config=quant_config
+            config=config,
+            prefix=prefix,
+            defer_lm_head=True,
         )
         # MTP layers sit past the base model's hidden layers; parse the index
         # from the prefix (e.g. "...layers.32") so the decoder builds an MLA