From bc96d76aa9946f3af81fabe508ae396088974b98 Mon Sep 17 00:00:00 2001 From: lucamotz Date: Sat, 5 Sep 2026 14:58:37 +0200 Subject: [PATCH] [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 Co-authored-by: Luca Motz <321921718+lucamotz@users.noreply.github.com> --- csrc/libtorch_stable/topk.cu | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu index 07b7521f86..b3e1a8623f 100644 --- a/csrc/libtorch_stable/topk.cu +++ b/csrc/libtorch_stable/topk.cu @@ -4,6 +4,7 @@ #include #include +#include "ops.h" #include "torch_utils.h" #ifndef USE_ROCM @@ -132,17 +133,15 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, if (num_groups == 0) num_groups = 1; uint32_t total_ctas = num_groups * ctas_per_group; - // If the cooperative launch wouldn't fit, fall back to FilteredTopK - // instead of deadlocking. Only relevant when needs_cooperative. + // If the cooperative launch wouldn't fit, use the generic decode kernel on + // low-smem devices or FilteredTopK where its 128 KiB requirement is met. if (needs_cooperative && total_ctas > hw_resident_cap) { - STD_TORCH_CHECK( - max_smem_per_block >= 128 * 1024, - "persistent_topk would oversubscribe and the FilteredTopK " - "fallback requires >=128KB smem per block (have ", - max_smem_per_block, "). total_ctas=", total_ctas, - " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK, - ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group, - ", smem=", smem_size, ")."); + if (max_smem_per_block < 128 * 1024) { + const int64_t next_n = lengths.dim() == 2 ? lengths.size(1) : 1; + top_k_per_row_decode(logits, next_n, lengths, output, num_rows, + logits.stride(0), logits.stride(1), TopK); + return; + } cudaError_t status = vllm::FilteredTopKRaggedTransform( logits.const_data_ptr(),