[BugFix] Restore full tokens for Qwen MTP When MoE SP (#48429)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Canlin Guo committed
Jul 13, 2026 at 13:29 UTC
36484e464a6cf763c5b4c8af7be8e19df324997a
2 files changed
+20
-2
vllm/model_executor/models/qwen3_5_mtp.py
+10
-1
@@ -26,6 +26,7 @@ from vllm.model_executor.models.qwen3_5 import (
26
)
27
from vllm.model_executor.models.qwen3_next import (
28
QwenNextMixtureOfExperts,
29
+ _all_gather_hidden_and_residual,
30
_is_shared_expert_fse_compatible,
31
)
32
from vllm.sequence import IntermediateTensors
@@ -150,7 +151,8 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
151
residual = intermediate_tensors["residual"]
152
153
current_step_idx = spec_step_idx % self.num_mtp_layers
153
- hidden_states, residual = self.layers[current_step_idx](
154
+ mtp_layer = self.layers[current_step_idx]
155
+ hidden_states, residual = mtp_layer(
156
positions=positions,
157
hidden_states=hidden_states,
158
residual=residual,
@@ -161,6 +163,13 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
163
{"hidden_states": hidden_states, "residual": residual}
164
)
165
166
+ if mtp_layer.use_attn_reduce_scatter_for_moe:
167
+ hidden_states, residual = _all_gather_hidden_and_residual(
168
+ hidden_states,
169
+ residual,
170
+ positions.shape[-1],
171
+ self.config.hidden_size,
172
+ )
173
hidden_states, _ = self.norm(hidden_states, residual)
174
return hidden_states
175
vllm/model_executor/models/qwen3_next_mtp.py
+10
-1
@@ -23,6 +23,7 @@ from vllm.model_executor.models.qwen3_next import (
23
Qwen3NextModel,
24
Qwen3NextRMSNorm,
25
QwenNextMixtureOfExperts,
26
+ _all_gather_hidden_and_residual,
27
)
28
from vllm.sequence import IntermediateTensors
29
from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
@@ -130,7 +131,8 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
131
residual = intermediate_tensors["residual"]
132
133
current_step_idx = spec_step_idx % self.num_mtp_layers
133
- hidden_states, residual = self.layers[current_step_idx](
134
+ mtp_layer = self.layers[current_step_idx]
135
+ hidden_states, residual = mtp_layer(
136
positions=positions,
137
hidden_states=hidden_states,
138
residual=residual,
@@ -141,6 +143,13 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
143
{"hidden_states": hidden_states, "residual": residual}
144
)
145
146
+ if mtp_layer.use_attn_reduce_scatter_for_moe:
147
+ hidden_states, residual = _all_gather_hidden_and_residual(
148
+ hidden_states,
149
+ residual,
150
+ positions.shape[-1],
151
+ self.config.hidden_size,
152
+ )
153
hidden_states, _ = self.norm(hidden_states, residual)
154
return hidden_states
155