iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Kernel] Fall back from persistent top-k on low-shared-memory GPUs (#54110)

Browse Source 
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>
This commit is contained in:
lucamotz
2026-09-05 08:58:37 -04:00
committed by GitHub
co-authored by Luca Motz
parent 7fbd44cbe0
commit bc96d76aa9
1 changed files with 9 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/topk.cu
+9 -10
View File
@@ -4,6 +4,7 @@
#include <cuda_runtime.h>
#include <algorithm>
#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<float, int32_t, TopK>(
logits.const_data_ptr<float>(),