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

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

Author: yaochengji

vllm/model_executor/layers/fused_moe/layer.py

@@ -1320,8 +1320,13 @@ def maybe_all_reduce_tensor_model_parallel(
 
     def forward(self, hidden_states: torch.Tensor,
                 router_logits: torch.Tensor):
-        return torch.ops.vllm.moe_forward(hidden_states, router_logits,
-                                          self.layer_name)
+        # 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)
 
     def forward_impl_chunked(self, full_hidden_states: torch.Tensor,
                              full_router_logits: torch.Tensor):