From 4b351f6b9582155263678b66086f9353ee15af7b Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Sun, 21 Dec 2025 10:13:32 +0800 Subject: [PATCH] Apply new moe align block size kernel (#14134) --- .../srt/layers/moe/fused_moe_triton/moe_align_block_size.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py b/python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py index ce1cae66e..2840eb8fc 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py @@ -54,7 +54,10 @@ def moe_align_block_size( - The padding ensures that the total number of tokens is now divisible by block_size for proper block matrix operations. """ - 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 + else: + max_num_tokens_padded = topk_ids.numel() + (num_experts + 1) * (block_size - 1) sorted_ids = torch.empty( (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device )