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

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))
 

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<int>(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));

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
         )
     ),
 )