minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Apply new moe align block size kernel (#14134)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-12-21 10:13:32 +08:00
committed by GitHub
parent 7fa4906f4f
commit 4b351f6b95
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py
+4 -1
View File
@@ -54,7 +54,10 @@ def moe_align_block_size(
- The padding ensures that the total number of tokens is now divisible - The padding ensures that the total number of tokens is now divisible
by block_size for proper block matrix operations. 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( sorted_ids = torch.empty(
(max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
) )