use default for torch.ops (#4835)

2025-03-27 19:09:58 -07:00
parent 10a9ab7b07
commit 31dfff7da7
7 changed files with 51 additions and 47 deletions
--- a/sgl-kernel/python/sgl_kernel/moe.py
+++ b/sgl-kernel/python/sgl_kernel/moe.py
@@ -11,7 +11,7 @@ def moe_align_block_size(
    token_cnts_buffer,
    cumsum_buffer,
 ):
-    torch.ops.sgl_kernel.moe_align_block_size(
+    torch.ops.sgl_kernel.moe_align_block_size.default(
        topk_ids,
        num_experts,
        block_size,
@@ -29,6 +29,6 @@ def topk_softmax(
    token_expert_indices: torch.Tensor,
    gating_output: float,
 ) -> None:
-    torch.ops.sgl_kernel.topk_softmax(
+    torch.ops.sgl_kernel.topk_softmax.default(
        topk_weights, topk_ids, token_expert_indices, gating_output
    )