[TPU][Bugfix] fix the MoE OOM issue (#20339)

Signed-off-by: Chengji Yao <chengjiyao@google.com>

vllm/model_executor/layers/fused_moe/layer.py

@@ -1320,6 +1320,11 @@ class FusedMoE(torch.nn.Module):
 
     def forward(self, hidden_states: torch.Tensor,
                 router_logits: torch.Tensor):
+        # TODO: Once the OOM issue for the TPU backend is resolved, we will
+        # switch to using the moe_forward custom op.
+        if current_platform.is_tpu():
+            return self.forward_impl(hidden_states, router_logits)
+        else:
             return torch.ops.vllm.moe_forward(hidden_states, router_logits,
                                               self.layer_name)