[Bugfix][TPU] fix moe custom kernel import (#18853)

Signed-off-by: Chengji Yao <chengjiyao@google.com>
2025-05-28 19:34:19 -07:00
parent 26b4fa45be
commit 8e882ffdc0
1 changed files with 1 additions and 0 deletions
--- a/vllm/model_executor/layers/fused_moe/moe_pallas.py
+++ b/vllm/model_executor/layers/fused_moe/moe_pallas.py
@@ -2,6 +2,7 @@

 import torch
 import torch.nn.functional as F
+import torch_xla.experimental.custom_kernel  # noqa: F401


 def _histogram(input: torch.Tensor, min: int, max: int) -> torch.Tensor: