remove moe_align_block_size torch.zeros in small batch/expert mode (#5298)

60bcbf2a · Xiaoyu Zhang · GitHub · a0a9f6d6 · 60bcbf2a
Unverified Commit 60bcbf2a authored Apr 12, 2025 by Xiaoyu Zhang Committed by GitHub Apr 11, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py +1 -1

No files found.
--- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
+++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
@@ -707,7 +707,7 @@ def moe_align_block_size(
            dtype=torch.int32,
            device=topk_ids.device,
        )
-        cumsum_buffer = torch.zeros(
+        cumsum_buffer = torch.empty(
            num_experts + 1, dtype=torch.int32, device=topk_ids.device
        )