[Kernel] Fall back from persistent top-k on low-shared-memory GPUs (#54110)
Signed-off-by: Luca Motz <321921718+lucamotz@users.noreply.github.com> Signed-off-by: Luca Motz <luca.motz@icloud.com> Co-authored-by: Luca Motz <321921718+lucamotz@users.noreply.github.com>
lucamotz committed
Sep 5, 2026 at 14:58 UTC
bc96d76aa9946f3af81fabe508ae396088974b98
1 file changed
+9
-10
csrc/libtorch_stable/topk.cu
+9
-10
@@ -4,6 +4,7 @@
4
#include <cuda_runtime.h>
5
#include <algorithm>
6
7
+#include "ops.h"
8
#include "torch_utils.h"
9
10
#ifndef USE_ROCM
@@ -132,17 +133,15 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
133
if (num_groups == 0) num_groups = 1;
134
uint32_t total_ctas = num_groups * ctas_per_group;
135
135
- // If the cooperative launch wouldn't fit, fall back to FilteredTopK
136
- // instead of deadlocking. Only relevant when needs_cooperative.
136
+ // If the cooperative launch wouldn't fit, use the generic decode kernel on
137
+ // low-smem devices or FilteredTopK where its 128 KiB requirement is met.
138
if (needs_cooperative && total_ctas > hw_resident_cap) {
138
- STD_TORCH_CHECK(
139
- max_smem_per_block >= 128 * 1024,
140
- "persistent_topk would oversubscribe and the FilteredTopK "
141
- "fallback requires >=128KB smem per block (have ",
142
- max_smem_per_block, "). total_ctas=", total_ctas,
143
- " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
144
- ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
145
- ", smem=", smem_size, ").");
139
+ if (max_smem_per_block < 128 * 1024) {
140
+ const int64_t next_n = lengths.dim() == 2 ? lengths.size(1) : 1;
141
+ top_k_per_row_decode(logits, next_n, lengths, output, num_rows,
142
+ logits.stride(0), logits.stride(1), TopK);
143
+ return;
144
+ }
145
cudaError_t status =
146
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
147
logits.const_data_ptr<float>(),