[optimize] fuse renormalize into moe_topk_softmax (#7744)

Co-authored-by: ispobock <ispobaoke@gmail.com>
2025-07-04 03:42:44 +08:00
parent 6840a7bbb2
commit 2998c4bdf4
7 changed files with 254 additions and 101 deletions
--- a/sgl-kernel/python/sgl_kernel/moe.py
+++ b/sgl-kernel/python/sgl_kernel/moe.py
@@ -30,11 +30,11 @@ def moe_align_block_size(
 def topk_softmax(
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
-    token_expert_indices: torch.Tensor,
    gating_output: float,
+    renormalize: bool = False,
 ) -> None:
    torch.ops.sgl_kernel.topk_softmax.default(
-        topk_weights, topk_ids, token_expert_indices, gating_output
+        topk_weights, topk_ids, gating_output, renormalize
    )