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

[Perf][Qwen3.5] Add case 512 to topkGatingSoftmaxKernelLauncher, (#25775)

Browse Source
This commit is contained in:
Jincong Chen
2026-05-25 16:08:21 +08:00
committed by GitHub
parent b0cf01eb85
commit e27d4fb70f
3 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/benchmark/bench_moe_topk_softmax.py
+1 -1
View File
@@ -144,7 +144,7 @@ if IS_CI:
else: else:
num_tokens_range = [128, 512, 1024, 2048, 4096, 8192, 16384, 32768] num_tokens_range = [128, 512, 1024, 2048, 4096, 8192, 16384, 32768]
num_experts_range = [32, 64, 128, 256, 12, 512] num_experts_range = [32, 64, 128, 256, 12, 512]
topk_range = [1, 2, 4, 8] topk_range = [1, 2, 4, 8, 10]
configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range)) configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
sgl-kernel/csrc/moe/moe_topk_softmax_kernels.cu
+4 -1
View File
@@ -696,6 +696,9 @@ void topkGatingSoftmaxKernelLauncher(
case 256: case 256:
LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB);
break; break;
case 512:
LAUNCH_SOFTMAX(T, 512, WARPS_PER_TB);
break;
default: { default: {
TORCH_CHECK( TORCH_CHECK(
softmax_workspace != nullptr, softmax_workspace != nullptr,
@@ -751,7 +754,7 @@ void topk_softmax(
const int topk = static_cast<int>(topk_weights.size(-1)); const int topk = static_cast<int>(topk_weights.size(-1));
const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0); const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
const bool needs_workspace = !is_pow_2 || num_experts > 256; const bool needs_workspace = !is_pow_2 || num_experts > 512;
const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0; const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output)); const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
sgl-kernel/tests/test_moe_topk_softmax.py
+1 -1
View File
@@ -18,7 +18,7 @@ def compare_topk_values(gating_output, topk_indices_ref, topk_indices):
itertools.product( itertools.product(
[1, 16, 128, 512, 1024, 2048], # num_tokens [1, 16, 128, 512, 1024, 2048], # num_tokens
[512], # num_experts [512], # num_experts
[1, 2, 3, 4, 5, 8], # topk [1, 2, 3, 4, 5, 8, 10], # topk
) )
), ),
) )