Skip to content

GitLab

Unverified Commit 8e882ffd authored by Chengji Yao's avatar Chengji Yao Committed by GitHub
Browse files

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


Signed-off-by: default avatarChengji Yao <chengjiyao@google.com>
parent 26b4fa45
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
vllm/model_executor/layers/fused_moe/moe_pallas.py
View file @ 8e882ffd
......@@ -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:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment