From 6da8f5f69e0f74fa4555234f8bd016b092af9592 Mon Sep 17 00:00:00 2001 From: Ma Mingfei Date: Mon, 30 Mar 2026 14:43:16 +0800 Subject: [PATCH] fix topk softmax performance issue (#14702) --- sgl-kernel/csrc/cpu/topk.cpp | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/sgl-kernel/csrc/cpu/topk.cpp b/sgl-kernel/csrc/cpu/topk.cpp index 0471661e5..100e87a7c 100644 --- a/sgl-kernel/csrc/cpu/topk.cpp +++ b/sgl-kernel/csrc/cpu/topk.cpp @@ -227,11 +227,9 @@ void topk_softmax_kernel_impl( queue[e] = {scores[e], e}; } - std::partial_sort( - queue.begin(), - queue.begin() + num_experts_per_group, - queue.end(), - [](const elem_t& x, const elem_t& y) -> bool { return x.first > y.first; }); + std::partial_sort(queue.begin(), queue.begin() + topk, queue.end(), [](const elem_t& x, const elem_t& y) -> bool { + return x.first > y.first; + }); for (int64_t j = 0; j < topk; ++j) { topk_weights[i * topk + j] = queue[j].first;