A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Model] Supplement to PR 24862: Pass param prefix to LLMHead (#25805)
Signed-off-by: whx-sjtu <2952154980@qq.com>
whx committed
Oct 3, 2025 at 21:34 UTC
cbf9221992bd17b1335632c4d9e1afb04c406200
8 files changed
+35
-12
vllm/model_executor/models/deepseek_mtp.py
+6
-2
@@ -28,13 +28,15 @@ class SharedHead(nn.Module):
28
def __init__(
29
self,
30
config: PretrainedConfig,
31
+ prefix: str,
32
quant_config: Optional[QuantizationConfig] = None,
33
) -> None:
34
super().__init__()
35
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
36
self.head = ParallelLMHead(config.vocab_size,
37
config.hidden_size,
37
- quant_config=quant_config)
38
+ quant_config=quant_config,
39
+ prefix=maybe_prefix(prefix, "head"))
40
41
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
42
return self.norm(hidden_states)
@@ -64,7 +66,9 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
66
device="cuda")
67
else:
68
topk_indices_buffer = None
67
- self.shared_head = SharedHead(config=config, quant_config=quant_config)
69
+ self.shared_head = SharedHead(config=config,
70
+ prefix=prefix,
71
+ quant_config=quant_config)
72
self.mtp_block = DeepseekV2DecoderLayer(vllm_config, prefix,
73
topk_indices_buffer)
74
vllm/model_executor/models/glm4_moe_mtp.py
+6
-2
@@ -50,13 +50,15 @@ class SharedHead(nn.Module):
50
def __init__(
51
self,
52
config: PretrainedConfig,
53
+ prefix: str,
54
quant_config: Optional[QuantizationConfig] = None,
55
) -> None:
56
super().__init__()
57
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
58
self.head = ParallelLMHead(config.vocab_size,
59
config.hidden_size,
59
- quant_config=quant_config)
60
+ quant_config=quant_config,
61
+ prefix=maybe_prefix(prefix, "head"))
62
63
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
64
return self.norm(hidden_states)
@@ -77,7 +79,9 @@ class Glm4MoeMultiTokenPredictorLayer(nn.Module):
79
self.eh_proj = nn.Linear(config.hidden_size * 2,
80
config.hidden_size,
81
bias=False)
80
- self.shared_head = SharedHead(config=config, quant_config=quant_config)
82
+ self.shared_head = SharedHead(config=config,
83
+ prefix=prefix,
84
+ quant_config=quant_config)
85
self.mtp_block = Glm4MoeDecoderLayer(config=config,
86
cache_config=cache_config,
87
quant_config=quant_config,
vllm/model_executor/models/gpt_neox.py
+1
@@ -296,6 +296,7 @@ class GPTNeoXForCausalLM(nn.Module, SupportsPP):
296
config.vocab_size,
297
config.hidden_size,
298
quant_config=quant_config,
299
+ prefix=maybe_prefix(prefix, "embed_out"),
300
)
301
if self.config.tie_word_embeddings:
302
self.embed_out.weight = self.gpt_neox.embed_in.weight
vllm/model_executor/models/longcat_flash_mtp.py
+1
@@ -140,6 +140,7 @@ class LongCatFlashMTP(nn.Module, SupportsPP):
140
self.config.vocab_size,
141
self.config.hidden_size,
142
quant_config=self.quant_config,
143
+ prefix=maybe_prefix(prefix, "lm_head"),
144
)
145
self.logits_processor = LogitsProcessor(self.config.vocab_size)
146
vllm/model_executor/models/medusa.py
+2
-1
@@ -82,7 +82,8 @@ class Medusa(nn.Module):
82
config.hidden_size,
83
org_num_embeddings=self.truncated_vocab_size,
84
padding_size=DEFAULT_VOCAB_PADDING_SIZE,
85
- ) for _ in range(self.config.num_heads)
85
+ prefix=maybe_prefix(prefix, f"lm_heads.{i}"),
86
+ ) for i in range(self.config.num_heads)
87
])
88
89
logit_scale = getattr(config, "logit_scale", 1.0)
vllm/model_executor/models/mlp_speculator.py
+14
-4
@@ -13,6 +13,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
13
ParallelLMHead, VocabParallelEmbedding)
14
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
15
16
+from .utils import maybe_prefix
17
+
18
SQRT2 = 2**0.5
19
20
@@ -97,8 +99,13 @@ class MLPSpeculator(nn.Module):
99
self.proj = nn.ModuleList([proj_first] + [proj_tied] *
100
(self.max_speculative_tokens - 1))
101
100
- head = ParallelLMHead(self.vocab_size, self.inner_dim, bias=False)
101
- self.head = nn.ModuleList([head] * self.max_speculative_tokens)
102
+ self.head = nn.ModuleList([
103
+ ParallelLMHead(self.vocab_size,
104
+ self.inner_dim,
105
+ bias=False,
106
+ prefix=maybe_prefix(prefix, f"head.{i}"))
107
+ for i in range(self.max_speculative_tokens)
108
+ ])
109
110
ln = MLPSpeculatorLayerNorm(self.inner_dim,
111
elementwise_scale_and_shift=True)
@@ -120,8 +127,11 @@ class MLPSpeculator(nn.Module):
127
])
128
129
self.head = nn.ModuleList([
123
- ParallelLMHead(self.vocab_size, self.inner_dim, bias=False)
124
- for _ in range(self.max_speculative_tokens)
130
+ ParallelLMHead(self.vocab_size,
131
+ self.inner_dim,
132
+ bias=False,
133
+ prefix=maybe_prefix(prefix, f"head.{i}"))
134
+ for i in range(self.max_speculative_tokens)
135
])
136
self.ln = nn.ModuleList([
137
MLPSpeculatorLayerNorm(self.inner_dim,
vllm/model_executor/models/qwen3_vl_moe.py
+2
-1
@@ -296,7 +296,8 @@ class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM):
296
prefix=maybe_prefix(prefix, "model"))
297
self.lm_head = ParallelLMHead(self.config.vocab_size,
298
self.config.hidden_size,
299
- quant_config=self.quant_config)
299
+ quant_config=self.quant_config,
300
+ prefix=maybe_prefix(prefix, "lm_head"))
301
if self.config.tie_word_embeddings:
302
self.lm_head.weight = self.model.embed_tokens.weight
303
self.logits_processor = LogitsProcessor(self.config.vocab_size)
vllm/model_executor/models/whisper.py
+3
-2
@@ -45,7 +45,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
45
from .interfaces import (MultiModalEmbeddings, SupportsMultiModal,
46
SupportsTranscription)
47
from .utils import (AutoWeightsLoader, WeightsMapper, cast_overflow_tensors,
48
- make_layers)
48
+ make_layers, maybe_prefix)
49
50
logger = init_logger(__name__)
51
@@ -885,7 +885,8 @@ class WhisperForConditionalGeneration(nn.Module, SupportsTranscription,
885
self.unpadded_vocab_size = config.vocab_size
886
self.proj_out = ParallelLMHead(config.vocab_size,
887
config.d_model,
888
- quant_config=quant_config)
888
+ quant_config=quant_config,
889
+ prefix=maybe_prefix(prefix, "proj_out"))
890
self.proj_out = self.proj_out.tie_weights(
891
self.model.decoder.embed_tokens)
892
logit_scale = getattr(config, "logit_scale", 1.0)