chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Opt moe align block size kernel (#14133)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-12-02 19:13:55 +08:00
committed by GitHub
parent 9530b76630
commit c5947ecd85
3 changed files with 96 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/benchmark/bench_moe_align_block_size.py
+29 -9
View File
@@ -162,14 +162,17 @@ def calculate_diff(num_tokens, num_experts=256, block_size=128, topk=8):
]
)
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
# SGL kernel uses dynamic padding optimization
max_num_tokens_padded_sgl = topk_ids.numel() + num_experts * (block_size - 1)
if topk_ids.numel() < num_experts + 1:
max_num_tokens_padded_sgl = topk_ids.numel() * block_size
sorted_ids_cuda = torch.empty(
(max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
(max_num_tokens_padded_sgl,), dtype=torch.int32, device=topk_ids.device
)
sorted_ids_cuda.fill_(topk_ids.numel())
max_num_m_blocks = max_num_tokens_padded // block_size
max_num_m_blocks_sgl = max_num_tokens_padded_sgl // block_size
expert_ids_cuda = torch.zeros(
(max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device
(max_num_m_blocks_sgl,), dtype=torch.int32, device=topk_ids.device
)
num_tokens_post_pad_cuda = torch.empty(
(1), dtype=torch.int32, device=topk_ids.device
@@ -178,14 +181,21 @@ def calculate_diff(num_tokens, num_experts=256, block_size=128, topk=8):
num_experts + 1, dtype=torch.int32, device=topk_ids.device
)
sorted_ids_triton = torch.empty_like(sorted_ids_cuda)
# Triton and vLLM use original padding calculation
max_num_tokens_padded_triton = topk_ids.numel() + num_experts * (block_size - 1)
max_num_m_blocks_triton = max_num_tokens_padded_triton // block_size
sorted_ids_triton = torch.empty(
(max_num_tokens_padded_triton,), dtype=torch.int32, device=topk_ids.device
)
sorted_ids_triton.fill_(topk_ids.numel())
expert_ids_triton = torch.zeros_like(expert_ids_cuda)
expert_ids_triton = torch.zeros(
(max_num_m_blocks_triton,), dtype=torch.int32, device=topk_ids.device
)
num_tokens_post_pad_triton = torch.empty_like(num_tokens_post_pad_cuda)
sorted_ids_vllm = torch.empty_like(sorted_ids_cuda)
sorted_ids_vllm = torch.empty_like(sorted_ids_triton)
sorted_ids_vllm.fill_(topk_ids.numel())
expert_ids_vllm = torch.zeros_like(expert_ids_cuda)
expert_ids_vllm = torch.zeros_like(expert_ids_triton)
num_tokens_post_pad_vllm = torch.empty_like(num_tokens_post_pad_cuda)
# compare the performance of cuda, triton and vllm implementation
@@ -326,7 +336,17 @@ def benchmark(num_tokens, num_experts, topk, provider):
device="cuda",
)
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
# Calculate max_num_tokens_padded based on provider
if provider == "sgl" or provider == "sgl_fusion":
# Apply dynamic padding optimization for SGL kernel
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
if topk_ids.numel() < num_experts:
max_num_tokens_padded = topk_ids.numel() * block_size
else: # triton
# Use original padding calculation for Triton
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
# Create tensors
sorted_ids = torch.empty(
(max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
)