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/csrc/moe/moe_topk_softmax_kernels.cu
+4 -1
View File
@@ -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));