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, 03:30 UTC
[TPU][Bugfix] fix the MoE OOM issue (#20339)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Jul 5, 2025 at 21:19 UTC
4548c03c50d8ef067b296fb8d610f9d8a8178482
1 file changed
+7
-2
vllm/model_executor/layers/fused_moe/layer.py
+7
-2
@@ -1320,8 +1320,13 @@ class FusedMoE(torch.nn.Module):
1320
1321
def forward(self, hidden_states: torch.Tensor,
1322
router_logits: torch.Tensor):
1323
- return torch.ops.vllm.moe_forward(hidden_states, router_logits,
1324
- self.layer_name)
1323
+ # TODO: Once the OOM issue for the TPU backend is resolved, we will
1324
+ # switch to using the moe_forward custom op.
1325
+ if current_platform.is_tpu():
1326
+ return self.forward_impl(hidden_states, router_logits)
1327
+ else:
1328
+ return torch.ops.vllm.moe_forward(hidden_states, router_logits,
1329
+ self.layer_name)
1330
1331
def forward_impl_chunked(self, full_hidden_states: torch.Tensor,
1332
full_router_logits: torch.Tensor):