A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 11:15 UTC
[Bugfix] Fix persistent_topk cooperative deadlock at TopK=1024 (#41189)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Yongye Zhu committed
Apr 30, 2026 at 00:03 UTC
a749a33d8d05acdd3ab346bd3f0c6b5c9c80474f
1 file changed
+59
-4
csrc/topk.cu
+59
-4
@@ -82,18 +82,73 @@ void launch_persistent_topk(const torch::Tensor& logits,
82
size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
83
if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;
84
85
+ // Query occupancy for the instantiation that will actually launch;
86
+ // overestimating it deadlocks the cooperative barrier.
87
int occupancy = 1;
86
- cudaOccupancyMaxActiveBlocksPerMultiprocessor(
87
- &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
88
- smem_size);
88
+ cudaError_t occ_err = cudaSuccess;
89
+ if (vec_size == 4) {
90
+ occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
91
+ &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
92
+ smem_size);
93
+ } else if (vec_size == 2) {
94
+ occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
95
+ &occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
96
+ smem_size);
97
+ } else {
98
+ occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
99
+ &occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
100
+ smem_size);
101
+ }
102
+ TORCH_CHECK(occ_err == cudaSuccess,
103
+ "persistent_topk occupancy query failed: ",
104
+ cudaGetErrorString(occ_err));
105
if (occupancy < 1) occupancy = 1;
106
91
- uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
107
+ // The cooperative spin-wait barrier only runs when at least one row hits
108
+ // the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs
109
+ // early-exit, so oversubscription can't deadlock and headroom is wasted.
110
+ const bool needs_cooperative =
111
+ static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;
112
+
113
+ const uint32_t hw_resident_cap =
114
+ static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy);
115
+ uint32_t max_resident_ctas = hw_resident_cap;
116
+ if (needs_cooperative) {
117
+ // Reserve one CTA per SM when occupancy allows; fall back to a single
118
+ // CTA when occupancy == 1 (the most deadlock-prone case — any straggler
119
+ // kernel that takes the only slot on one SM hangs the barrier). Never
120
+ // drop below one full group's worth.
121
+ uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u;
122
+ if (max_resident_ctas >= headroom + ctas_per_group) {
123
+ max_resident_ctas -= headroom;
124
+ }
125
+ }
126
uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
127
static_cast<uint32_t>(num_rows));
128
if (num_groups == 0) num_groups = 1;
129
uint32_t total_ctas = num_groups * ctas_per_group;
130
131
+ // If the cooperative launch wouldn't fit, fall back to FilteredTopK
132
+ // instead of deadlocking. Only relevant when needs_cooperative.
133
+ if (needs_cooperative && total_ctas > hw_resident_cap) {
134
+ TORCH_CHECK(max_smem_per_block >= 128 * 1024,
135
+ "persistent_topk would oversubscribe and the FilteredTopK "
136
+ "fallback requires >=128KB smem per block (have ",
137
+ max_smem_per_block, "). total_ctas=", total_ctas,
138
+ " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
139
+ ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
140
+ ", smem=", smem_size, ").");
141
+ cudaError_t status =
142
+ vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
143
+ logits.data_ptr<float>(), output.data_ptr<int32_t>(),
144
+ lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
145
+ static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
146
+ stream);
147
+ TORCH_CHECK(status == cudaSuccess,
148
+ "FilteredTopK fallback failed: ", cudaGetErrorString(status));
149
+ return;
150
+ }
151
+
152
size_t state_bytes = num_groups * sizeof(P::RadixRowState);
153
TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
154
"workspace too small, need ", state_bytes, " bytes");