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 15, 04:45 UTC
[V1][Speculative Decoding] Fix DeepSeek MTP (#20022)
Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com>
cjackal committed
Jun 26, 2025 at 00:41 UTC
8359f4c8d840d409fe698d9bf428ecd2f7e85b75
2 files changed
+21
-9
vllm/model_executor/models/deepseek_mtp.py
+20
-8
@@ -52,11 +52,6 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
52
quant_config: Optional[QuantizationConfig] = None,
53
) -> None:
54
super().__init__()
55
- self.embed_tokens = VocabParallelEmbedding(
56
- config.vocab_size,
57
- config.hidden_size,
58
- )
59
-
55
self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
56
self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
57
self.eh_proj = nn.Linear(config.hidden_size * 2,
@@ -74,8 +69,6 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
69
inputs_embeds: Optional[torch.Tensor] = None,
70
spec_step_index: int = 0,
71
) -> torch.Tensor:
77
- if inputs_embeds is None:
78
- inputs_embeds = self.embed_tokens(input_ids)
72
assert inputs_embeds is not None
73
# masking inputs at position 0, as not needed by MTP
74
inputs_embeds[positions == 0] = 0
@@ -112,7 +105,10 @@ class DeepSeekMultiTokenPredictor(nn.Module):
105
for idx in range(self.mtp_start_layer_idx,
106
self.mtp_start_layer_idx + self.num_mtp_layers)
107
})
115
-
108
+ self.embed_tokens = VocabParallelEmbedding(
109
+ config.vocab_size,
110
+ config.hidden_size,
111
+ )
112
self.logits_processor = LogitsProcessor(config.vocab_size)
113
114
def forward(
@@ -123,6 +119,8 @@ class DeepSeekMultiTokenPredictor(nn.Module):
119
inputs_embeds: Optional[torch.Tensor] = None,
120
spec_step_idx: int = 0,
121
) -> torch.Tensor:
122
+ if inputs_embeds is None:
123
+ inputs_embeds = self.embed_tokens(input_ids)
124
current_step_idx = (spec_step_idx % self.num_mtp_layers)
125
return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
126
input_ids,
@@ -242,6 +240,12 @@ class DeepSeekMTP(nn.Module, SupportsPP):
240
if name.endswith(".bias") and name not in params_dict:
241
continue
242
243
+ # According to DeepSeek-V3 Technical Report, MTP modules
244
+ # shares embedding layer. We only load the first weights.
245
+ if (spec_layer != self.model.mtp_start_layer_idx
246
+ and ".layers" not in name):
247
+ continue
248
+
249
param = params_dict[name]
250
weight_loader = getattr(param, "weight_loader",
251
default_weight_loader)
@@ -253,17 +257,25 @@ class DeepSeekMTP(nn.Module, SupportsPP):
257
"""
258
Rewrite the weight name to match the format of the original model.
259
Add .mtp_block for modules in transformer layer block for spec layer
260
+ and rename shared layer weights to be top level.
261
"""
262
spec_layer_weight_names = [
263
"embed_tokens", "enorm", "hnorm", "eh_proj", "shared_head"
264
]
265
+ shared_weight_names = ["embed_tokens"]
266
spec_layer_weight = False
267
+ shared_weight = False
268
for weight_name in spec_layer_weight_names:
269
if weight_name in name:
270
spec_layer_weight = True
271
+ if weight_name in shared_weight_names:
272
+ shared_weight = True
273
break
274
if not spec_layer_weight:
275
# treat rest weights as weights for transformer layer block
276
name = name.replace(f"model.layers.{spec_layer}.",
277
f"model.layers.{spec_layer}.mtp_block.")
278
+ elif shared_weight:
279
+ # treat shared weights as top level weights
280
+ name = name.replace(f"model.layers.{spec_layer}.", "model.")
281
return name
vllm/v1/spec_decode/eagle.py
+1
-1
@@ -148,7 +148,7 @@ class EagleProposer:
148
assert self.runner is not None
149
150
# FIXME: need to consider multiple kv_cache_groups
151
- attn_metadata = self.runner.attn_metadata_builder.build(
151
+ attn_metadata = self.runner.attn_metadata_builders[0].build(
152
common_prefix_len=0,
153
common_attn_metadata=common_attn_metadata,
154
)