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
[Compilation Bug] Fix Inductor Graph Output with Shape Issue (#24772)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 12, 2025 at 17:23 UTC
3beadc2f25d1111db6a5efee52d5cfe2ccf28d45
1 file changed
+6
-3
vllm/model_executor/models/qwen3_moe.py
+6
-3
@@ -170,8 +170,9 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
170
return quant_config
171
172
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
173
- # NOTE: hidden_states can have either 1D or 2D shape.
174
- orig_shape = hidden_states.shape
173
+ assert hidden_states.dim(
174
+ ) <= 2, "Qwen3MoeSparseMoeBlock only supports 1D or 2D inputs"
175
+ is_input_1d = hidden_states.dim() == 1
176
hidden_dim = hidden_states.shape[-1]
177
hidden_states = hidden_states.view(-1, hidden_dim)
178
@@ -180,7 +181,9 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
181
final_hidden_states = self.experts(hidden_states=hidden_states,
182
router_logits=router_logits)
183
183
- return final_hidden_states.view(orig_shape)
184
+ # return to 1d if input is 1d
185
+ return final_hidden_states.squeeze(0) if is_input_1d else \
186
+ final_hidden_states
187
188
189
class Qwen3MoeAttention(nn.Module):