Fuse sorted_token_ids padding to moe_align_block_size kernel (#7437)

2025-06-24 17:44:27 -07:00
parent 112b496a6c
commit 57ab776910
7 changed files with 163 additions and 70 deletions
@@ -212,7 +212,8 @@ void moe_align_block_size(
    torch::Tensor experts_ids,
    torch::Tensor num_tokens_post_pad,
    torch::Tensor token_cnts_buffer,
-    torch::Tensor cumsum_buffer);
+    torch::Tensor cumsum_buffer,
+    bool pad_sorted_token_ids);

 void topk_softmax(
    torch::Tensor& topk_weights,