fix topk softmax performance issue (#14702)
This commit is contained in:
@@ -227,11 +227,9 @@ void topk_softmax_kernel_impl(
|
|||||||
queue[e] = {scores[e], e};
|
queue[e] = {scores[e], e};
|
||||||
}
|
}
|
||||||
|
|
||||||
std::partial_sort(
|
std::partial_sort(queue.begin(), queue.begin() + topk, queue.end(), [](const elem_t& x, const elem_t& y) -> bool {
|
||||||
queue.begin(),
|
return x.first > y.first;
|
||||||
queue.begin() + num_experts_per_group,
|
});
|
||||||
queue.end(),
|
|
||||||
[](const elem_t& x, const elem_t& y) -> bool { return x.first > y.first; });
|
|
||||||
|
|
||||||
for (int64_t j = 0; j < topk; ++j) {
|
for (int64_t j = 0; j < topk; ++j) {
|
||||||
topk_weights[i * topk + j] = queue[j].first;
|
topk_weights[i * topk + j] = queue[j].first;
|
||||||
|
|||||||
Reference in New Issue
Block a user