@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / edd60ac93a

Fork 0

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, 08:45 UTC

[Bugfix] Fix persistent_topk inter-CTA init race on RadixRowState (#41444)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

Yongye Zhu committed May 1, 2026 at 17:42 UTC edd60ac93a3247c7ef1bf1e2a3e9c0e95bc83bf6

2 files changed +29 -19

csrc/persistent_topk.cuh

+6 -19

@@ -887,27 +887,14 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)

887 uint32_t* shared_ordered =

888 reinterpret_cast<uint32_t*>(smem_raw + kFixedSmemLarge);

889

890 - // RadixRowState for multi-CTA cooperative radix

890 + // RadixRowState for multi-CTA cooperative radix.

891 + // Zero-initialization is done host-side via cudaMemsetAsync in topk.cu

892 + // before launch — that gives a stream-ordered happens-before edge for all

893 + // CTAs, which the previous in-kernel init (CTA-0 only + intra-CTA

894 + // __syncthreads) did not provide and which manifested as a race against

895 + // CTA-1+'s first red_release on arrival_counter.

896 RadixRowState* state = &params.row_states[group_id];

897

893 - // -- Initialize RadixRowState (only needed if large rows exist) --

894 - if (params.max_seq_len > RADIX_THRESHOLD) {

895 - if (cta_in_group == 0) {

896 - for (uint32_t buf = 0; buf < 3; buf++) {

897 - for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) {

898 - state->histogram[buf][i] = 0;

899 - }

900 - }

901 - if (tx == 0) {

902 - state->remaining_k = 0;

903 - state->prefix = 0;

904 - state->arrival_counter = 0;

905 - state->output_counter = 0;

906 - }

907 - }

908 - __syncthreads();

909 - }

910 -

898 int barrier_phase = 0;

899 const uint32_t total_iters = (params.num_rows + num_groups - 1) / num_groups;

900

csrc/topk.cu

+23

@@ -153,6 +153,29 @@ 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.

162 + //

163 + // Why we need the memset (when needs_cooperative is true):

164 + // 1. arrival_counter accumulates within a launch and is never reset,

165 + // so a prior call leaves it at a large positive value. Without this

166 + // reset, the very first wait_ge in the next call sees counter >>

167 + // target and returns instantly, breaking the barrier.

168 + // 2. The previous in-kernel init only ran in CTA-0 with intra-CTA

169 + // __syncthreads(), so it had no happens-before edge to CTA-1+'s

170 + // first red_release. cudaMemsetAsync is stream-ordered: the zero

171 + // is globally visible before any CTA runs.

172 + if (needs_cooperative) {

173 + cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr<uint8_t>(), 0,

174 + state_bytes, stream);

175 + TORCH_CHECK(mz_err == cudaSuccess,

176 + "row_states memset failed: ", cudaGetErrorString(mz_err));

177 + }

178 +

179 P::PersistentTopKParams params;

180 params.input = logits.data_ptr<float>();

181 params.output = output.data_ptr<int32_t>();