From 8e882ffdc0b4e9e3e6fc3fabe1c199fb6b8736a2 Mon Sep 17 00:00:00 2001 From: Chengji Yao Date: Wed, 28 May 2025 19:34:19 -0700 Subject: [PATCH] [Bugfix][TPU] fix moe custom kernel import (#18853) Signed-off-by: Chengji Yao --- vllm/model_executor/layers/fused_moe/moe_pallas.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/model_executor/layers/fused_moe/moe_pallas.py b/vllm/model_executor/layers/fused_moe/moe_pallas.py index babeb97308a9f..53945999288db 100644 --- 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: