[Bugfix][Model][Spec Decode] Defer disposable GLM MTP head (#55442)
Signed-off-by: Luca Motz <luca.motz@icloud.com>
lucamotz committed
Sep 22, 2026 at 22:08 UTC
276db94fb22f032a933c3c17e6ed377a8fc612fb
3 files changed
+40
-8
tests/v1/spec_decode/test_mtp.py
+26
@@ -5,6 +5,7 @@ from unittest import mock
5
6
import pytest
7
import torch
8
+import torch.nn as nn
9
10
from tests.v1.attention.utils import (
11
BatchSpec,
@@ -31,6 +32,31 @@ mimo_7b_dir = "XiaomiMiMo/MiMo-7B-Base"
32
DEVICE_TYPE = current_platform.device_type
33
34
35
+def test_glm_mtp_defers_lm_head(default_vllm_config):
36
+ from vllm.model_executor.models import deepseek_mtp
37
+ from vllm.models.glm5next.common import mtp
38
+
39
+ config = mock.MagicMock(
40
+ hidden_size=16,
41
+ rms_norm_eps=1e-5,
42
+ index_topk=8,
43
+ index_kpool=4,
44
+ )
45
+ vllm_config = mock.MagicMock()
46
+ vllm_config.speculative_config.draft_model_config.hf_config = config
47
+ vllm_config.scheduler_config.max_num_batched_tokens = 4
48
+
49
+ with (
50
+ mock.patch.object(mtp, "Glm5NextDecoderLayer", return_value=nn.Identity()),
51
+ mock.patch.object(deepseek_mtp, "ParallelLMHead") as parallel_lm_head,
52
+ mock.patch.object(mtp.current_platform, "device_type", "cpu"),
53
+ ):
54
+ layer = mtp.Glm5NextMultiTokenPredictorLayer(vllm_config, "model.layers.1")
55
+
56
+ parallel_lm_head.assert_not_called()
57
+ assert layer.shared_head.head is None
58
+
59
+
60
def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer:
61
"""Create an MTP proposer with unified model configuration."""
62
model_config = ModelConfig(
vllm/model_executor/models/deepseek_mtp.py
+11
-6
@@ -54,15 +54,20 @@ class SharedHead(nn.Module):
54
config: PretrainedConfig,
55
prefix: str,
56
quant_config: QuantizationConfig | None = None,
57
+ *,
58
+ defer_lm_head: bool = False,
59
) -> None:
60
super().__init__()
61
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
60
- self.head = ParallelLMHead(
61
- config.vocab_size,
62
- config.hidden_size,
63
- quant_config=quant_config,
64
- prefix=maybe_prefix(prefix, "head"),
65
- )
62
+ if defer_lm_head:
63
+ self.head = None
64
+ else:
65
+ self.head = ParallelLMHead(
66
+ config.vocab_size,
67
+ config.hidden_size,
68
+ quant_config=quant_config,
69
+ prefix=maybe_prefix(prefix, "head"),
70
+ )
71
72
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
73
return self.norm(hidden_states)
vllm/models/glm5next/common/mtp.py
+3
-2
@@ -42,7 +42,6 @@ class Glm5NextMultiTokenPredictorLayer(nn.Module):
42
assert vllm_config.speculative_config is not None
43
config = vllm_config.speculative_config.draft_model_config.hf_config
44
self.config = config
45
- quant_config = vllm_config.quant_config
45
46
self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
47
self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -66,7 +65,9 @@ class Glm5NextMultiTokenPredictorLayer(nn.Module):
65
device=current_platform.device_type,
66
)
67
self.shared_head = SharedHead(
69
- config=config, prefix=prefix, quant_config=quant_config
68
+ config=config,
69
+ prefix=prefix,
70
+ defer_lm_head=True,
71
)
72
# MTP layers sit past the base model's hidden layers; parse the index
73
# from the prefix (e.g. "...layers.32") so the decoder builds an MLA