Fuse sorted_token_ids padding to moe_align_block_size kernel (#7437)

2025-06-25 08:44:27 +08:00
parent 112b496a6c
commit 57ab776910
7 changed files with 163 additions and 70 deletions
--- a/sgl-kernel/python/sgl_kernel/moe.py
+++ b/sgl-kernel/python/sgl_kernel/moe.py
@@ -12,6 +12,7 @@ def moe_align_block_size(
    num_tokens_post_pad,
    token_cnts_buffer,
    cumsum_buffer,
+    pad_sorted_token_ids=False,
 ):
    torch.ops.sgl_kernel.moe_align_block_size.default(
        topk_ids,
@@ -22,6 +23,7 @@ def moe_align_block_size(
        num_tokens_post_pad,
        token_cnts_buffer,
        cumsum_buffer,
+        pad_sorted_token_ids,
    )