From e27d4fb70f33d6a86533a5dc3e14114c480b4d2f Mon Sep 17 00:00:00 2001 From: Jincong Chen Date: Mon, 25 May 2026 16:08:21 +0800 Subject: [PATCH] [Perf][Qwen3.5] Add case 512 to topkGatingSoftmaxKernelLauncher, (#25775) --- sgl-kernel/benchmark/bench_moe_topk_softmax.py | 2 +- sgl-kernel/csrc/moe/moe_topk_softmax_kernels.cu | 5 ++++- sgl-kernel/tests/test_moe_topk_softmax.py | 2 +- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/sgl-kernel/benchmark/bench_moe_topk_softmax.py b/sgl-kernel/benchmark/bench_moe_topk_softmax.py index 451ae8d80..9ae15d6c9 100644 --- a/sgl-kernel/benchmark/bench_moe_topk_softmax.py +++ b/sgl-kernel/benchmark/bench_moe_topk_softmax.py @@ -144,7 +144,7 @@ if IS_CI: else: num_tokens_range = [128, 512, 1024, 2048, 4096, 8192, 16384, 32768] 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)) diff --git a/sgl-kernel/csrc/moe/moe_topk_softmax_kernels.cu b/sgl-kernel/csrc/moe/moe_topk_softmax_kernels.cu index 82f8b89fc..44cf9b314 100644 --- a/sgl-kernel/csrc/moe/moe_topk_softmax_kernels.cu +++ b/sgl-kernel/csrc/moe/moe_topk_softmax_kernels.cu @@ -696,6 +696,9 @@ void topkGatingSoftmaxKernelLauncher( case 256: LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB); break; + case 512: + LAUNCH_SOFTMAX(T, 512, WARPS_PER_TB); + break; default: { TORCH_CHECK( softmax_workspace != nullptr, @@ -751,7 +754,7 @@ void topk_softmax( const int topk = static_cast(topk_weights.size(-1)); 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 at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output)); diff --git a/sgl-kernel/tests/test_moe_topk_softmax.py b/sgl-kernel/tests/test_moe_topk_softmax.py index 1a8bfb93c..77ffe8a46 100644 --- a/sgl-kernel/tests/test_moe_topk_softmax.py +++ b/sgl-kernel/tests/test_moe_topk_softmax.py @@ -18,7 +18,7 @@ def compare_topk_values(gating_output, topk_indices_ref, topk_indices): itertools.product( [1, 16, 128, 512, 1024, 2048], # num_tokens [512], # num_experts - [1, 2, 3, 4, 5, 8], # topk + [1, 2, 3, 4, 5, 8, 10], # topk ) ), )