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, 08:30 UTC
[Bugfix][TPU] fix moe custom kernel import (#18853)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
May 28, 2025 at 19:34 UTC
8e882ffdc0b4e9e3e6fc3fabe1c199fb6b8736a2
1 file changed
+1
vllm/model_executor/layers/fused_moe/moe_pallas.py
+1
@@ -2,6 +2,7 @@
2
3
import torch
4
import torch.nn.functional as F
5
+import torch_xla.experimental.custom_kernel # noqa: F401
6
7
8
def _histogram(input: torch.Tensor, min: int, max: int) -> torch.Tensor: