Opt moe align block size kernel (#14133)

2025-12-02 19:13:55 +08:00
parent 9530b76630
commit c5947ecd85
3 changed files with 96 additions and 55 deletions
@@ -165,6 +165,8 @@ def test_moe_align_block_size_compare_implementations(
    ]

    max_num_tokens_padded = topk_ids.numel() + (num_experts + 1) * (block_size - 1)
+    if topk_ids.numel() < num_experts + 1:
+        max_num_tokens_padded = topk_ids.numel() * block_size

    sorted_ids_cuda = torch.empty(
        (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device