[Bugfix] Recycle post-final-norm hidden in GLM MTP (single norm) (#47448)
xiaozhoupy committed
Jul 6, 2026 at 16:07 UTC
d039c171144bdc6e19d14d3dd9b0625f84f16d1c
1 file changed
+16
-7
vllm/models/deepseek_v32/nvidia/mtp.py
+16
-7
@@ -93,10 +93,18 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
93
# main model fuses that all-reduce into the next norm, but here the
94
# recycle hidden is consumed directly, so reduce it now.
95
hidden_states = tensor_model_parallel_all_reduce(hidden_states)
96
- # Return the pre-final-norm recycle hidden (re-fed as the next spec
97
- # step's previous_hidden_states); shared_head norm is applied in
98
- # compute_logits. Matches the V2-runner / deepseek_v4 MTP contract.
99
- return residual + hidden_states
96
+ # Recycle the POST-final-norm hidden into the next draft step. The
97
+ # residual-add is fused into the final RMSNorm so it is computed
98
+ # exactly once, and the result is returned for both tuple positions:
99
+ # the draft-logits hidden (compute_logits applies the LM head only) and
100
+ # the recycled previous_hidden_states. Recycling the pre-final-norm
101
+ # hidden mismatches the draft model's hnorm and lowers MTP acceptance;
102
+ # post-norm recycle matches deepseek_mtp.py (PR #45895). The tuple form
103
+ # is understood by both the V2 speculator (isinstance-tuple check) and
104
+ # the legacy proposer (model_returns_tuple is True for the
105
+ # DeepSeekMTPModel architecture).
106
+ hidden_states, _ = self.shared_head.norm(hidden_states, residual)
107
+ return hidden_states, hidden_states
108
109
110
class DeepseekV32MultiTokenPredictor(nn.Module):
@@ -168,9 +176,10 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
176
) -> torch.Tensor:
177
current_step_idx = spec_step_idx % self.num_mtp_layers
178
mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
171
- return self.logits_processor(
172
- mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states)
173
- )
179
+ # hidden_states is already post-final-norm (produced in the layer
180
+ # forward and recycled as-is); apply the LM head only, without a
181
+ # second RMSNorm.
182
+ return self.logits_processor(mtp_layer.shared_head.head, hidden_states)
183
184
185
class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):