A high-throughput and memory-efficient inference and serving engine for LLMs
[XPU][DeepSeek-V4] Fix MTP: sync with upstream fixes #44821 and #43746 (#45240)
Signed-off-by: Ma Jian <jian1.ma@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Ma Jian committed
Jun 12, 2026 at 15:41 UTC
04cec9e4d846947e70cc9beebce0a51230905c68
1 file changed
+29
-18
vllm/models/deepseek_v4/xpu/mtp.py
+29
-18
@@ -18,7 +18,6 @@ import regex as re
18
import torch
19
import torch.nn as nn
20
21
-from vllm.compilation.decorators import support_torch_compile
21
from vllm.config import VllmConfig
22
from vllm.distributed import (
23
get_tensor_model_parallel_rank,
@@ -39,6 +38,10 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
38
from vllm.model_executor.models.deepseek_mtp import SharedHead
39
from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
40
from vllm.model_executor.models.utils import maybe_prefix
41
+from vllm.models.deepseek_v4.common.ops import (
42
+ fused_mtp_input_rmsnorm,
43
+ mtp_shared_head_rmsnorm,
44
+)
45
from vllm.platforms import current_platform
46
from vllm.sequence import IntermediateTensors
47
@@ -87,6 +90,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
90
bias=False,
91
return_bias=False,
92
quant_config=quant_config,
93
+ prefix=f"{prefix}.e_proj",
94
)
95
self.h_proj = ReplicatedLinear(
96
config.hidden_size,
@@ -94,6 +98,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
98
bias=False,
99
return_bias=False,
100
quant_config=quant_config,
101
+ prefix=f"{prefix}.h_proj",
102
)
103
104
self.hc_eps = config.hc_eps
@@ -133,22 +138,31 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
138
spec_step_index: int = 0,
139
) -> torch.Tensor:
140
assert inputs_embeds is not None
136
- # masking inputs at position 0, as not needed by MTP
137
- inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
138
- inputs_embeds = self.enorm(inputs_embeds)
139
-
141
# Target stashes pre-hc_head residual as flat (T, hc_mult * D);
141
- # reshape to (T, hc_mult, D) — the training-time layout.
142
+ # reshape to (T, hc_mult, D) — the training-time layout — before
143
+ # the fused norm pass so both inputs are 3D-friendly.
144
previous_hidden_states = previous_hidden_states.view(
145
-1, self.hc_mult, self.config.hidden_size
146
)
145
- previous_hidden_states = self.hnorm(previous_hidden_states)
147
+ # Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm.
148
+ inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm(
149
+ inputs_embeds,
150
+ positions,
151
+ previous_hidden_states,
152
+ self.enorm.weight.data,
153
+ self.hnorm.weight.data,
154
+ self.enorm.variance_epsilon,
155
+ self.hc_mult,
156
+ )
157
hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(
158
inputs_embeds
159
).unsqueeze(-2)
160
hidden_states, residual, post_mix, res_mix = self.mtp_block(
161
positions=positions, x=hidden_states, input_ids=None
162
)
163
+ hidden_states = self.mtp_block.hc_post(
164
+ hidden_states, residual, post_mix, res_mix
165
+ )
166
# Return the flat pre-hc_head residual so it can be re-fed as the
167
# next spec step's `previous_hidden_states` when
168
# num_speculative_tokens > 1. hc_head is deferred to compute_logits.
@@ -238,13 +252,15 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
252
mtp_layer.rms_norm_eps,
253
mtp_layer.hc_eps,
254
)
241
- logits = self.logits_processor(
242
- mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states)
255
+ hidden_states = mtp_shared_head_rmsnorm(
256
+ hidden_states,
257
+ mtp_layer.shared_head.norm.weight.data,
258
+ mtp_layer.shared_head.norm.variance_epsilon,
259
)
260
+ logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states)
261
return logits
262
263
247
-@support_torch_compile
264
class DeepSeekV4MTP(nn.Module):
265
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
266
super().__init__()
@@ -285,11 +301,6 @@ class DeepSeekV4MTP(nn.Module):
301
".emb.tok_emb.weight": ".embed_tokens.weight",
302
".head.weight": ".shared_head.head.weight",
303
".norm.weight": ".shared_head.norm.weight",
288
- # Pre-MoE norm + gate are now owned by
289
- # ``DeepseekV4MoE.norm_gate`` (see NormGatedLinear).
290
- ".ffn_norm.weight": ".ffn.norm_gate.norm.weight",
291
- ".ffn.gate.weight": ".ffn.norm_gate.gate.weight",
292
- ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid",
304
}
305
306
def _remap_weight_name(name: str) -> str:
@@ -437,11 +448,11 @@ class DeepSeekV4MTP(nn.Module):
448
".shared_experts.w2", ".shared_experts.down_proj"
449
)
450
if name.endswith(".ffn.gate.bias"):
440
- # ``e_score_correction_bias`` lives on
441
- # ``norm_gate`` directly (not on the inner gate).
451
+ # ``e_score_correction_bias`` lives on the gate
452
+ # under a different attribute name.
453
name = name.replace(
454
".ffn.gate.bias",
444
- ".ffn.norm_gate.e_score_correction_bias",
455
+ ".ffn.gate.e_score_correction_bias",
456
)
457
param = params_dict[name]
458
weight_loader = getattr(