Update fused_moe.py

zwd003 · Jun 28, 2024 · 0746b4f · 0746b4f
1 parent 4cf44a5
commit 0746b4f
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -346,11 +346,15 @@ def fused_topk(
                                topk,
                                dtype=torch.float32,
                                device=hidden_states.device)
+    topk_ids = torch.empty(M,
+                           topk,
+                           dtype=torch.int32,
+                           device=hidden_states.device)
     token_expert_indicies = torch.empty(M,
                                         topk,
                                         dtype=torch.int32,
                                         device=hidden_states.device)
-    moe_kernels.topk_softmax(
+    ops.topk_softmax(
         topk_weights,
         topk_ids,
         token_expert_indicies,