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, 05:00 UTC
[Bugfix] Fix condition to clear persistent topk so that it can be captured regardless (#41665)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Yongye Zhu committed
May 6, 2026 at 19:17 UTC
80d5e7d103ee181f66389d0e2580841e2a2f1db5
1 file changed
+17
-8
csrc/topk.cu
+17
-8
@@ -153,14 +153,23 @@ void launch_persistent_topk(const torch::Tensor& logits,
153
TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
154
"workspace too small, need ", state_bytes, " bytes");
155
156
- // Zero the per-group RadixRowState region before launch — only when the
157
- // radix path will actually run (max_seq_len > RADIX_THRESHOLD). The
158
- // RadixRowState fields (arrival_counter, histograms) are only touched by
159
- // radix_topk; the decode/medium paths inside the persistent kernel
160
- // operate purely in shared memory and never read these globals, so a
161
- // stale workspace is harmless for them.
156
+ // Zero the per-group RadixRowState region before launch.
157
//
163
- // Why we need the memset (when needs_cooperative is true):
158
+ // Issued UNCONDITIONALLY so the memset is captured as its own node in
159
+ // the cudagraph (a separate cudaMemsetAsync node, sequenced before the
160
+ // persistent_topk_kernel launch on the same stream). The previous
161
+ // host-side guard `if (needs_cooperative)` was evaluated at capture time;
162
+ // when capture-time max_seq_len <= RADIX_THRESHOLD (always true under
163
+ // FULL_DECODE_ONLY with max_model_len < 32 K) the memset would NOT be
164
+ // captured, leaving the workspace state to accumulate across replays.
165
+ // That's a latent correctness bug if the runtime data ever takes the
166
+ // radix path, and removes one variable while debugging hangs in the
167
+ // decode/medium paths.
168
+ //
169
+ // Cost is sub-microsecond: state_bytes = num_groups * sizeof(RadixRowState)
170
+ // is ~3 KB per group, ~100 KB for the largest grids on this hardware.
171
+ //
172
+ // Why the memset is required (regardless of which path the kernel takes):
173
// 1. arrival_counter accumulates within a launch and is never reset,
174
// so a prior call leaves it at a large positive value. Without this
175
// reset, the very first wait_ge in the next call sees counter >>
@@ -169,7 +178,7 @@ void launch_persistent_topk(const torch::Tensor& logits,
178
// __syncthreads(), so it had no happens-before edge to CTA-1+'s
179
// first red_release. cudaMemsetAsync is stream-ordered: the zero
180
// is globally visible before any CTA runs.
172
- if (needs_cooperative) {
181
+ {
182
cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr<uint8_t>(), 0,
183
state_bytes, stream);
184
TORCH_CHECK(mz_err == cudaSuccess,