xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-24 09:26:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Chengji Yao <chengjiyao@google.com>
This commit is contained in:
Chengji Yao 2025-05-28 19:34:19 -07:00 committed by GitHub
parent 26b4fa45be
commit 8e882ffdc0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
vllm/model_executor/layers/fused_moe/moe_pallas.py
View File

@ -2,6 +2,7 @@
import torch import torch
import torch.nn.functional as F 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: def _histogram(input: torch.Tensor, min: int, max: int) -> torch.Tensor: