diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index be9aaf5c19..63deab6015 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -1179,6 +1179,18 @@ __device__ void radix_topk(const float* __restrict__ row_input, } } + // The arrival publication below is made by thread 0 + // ALONE, and `red_release`'s `fence.acq_rel.gpu` orders only the issuing + // thread's own prior accesses. Without this barrier thread 0 could signal + // arrival while other warps of this CTA had not yet merged their bins into + // `current_hist`, or (in CTA 0) had not finished zeroing `next_hist`. A + // peer CTA that observes the completed arrival count then leaves its own + // `wait_ge` -- whose `__syncthreads()` is AFTER arrival, so it orders + // nothing here -- and reads a half-written histogram, which changes the + // pivot. The barrier must be between the producing stores and the leader's + // publication; program order in one thread is not cross-thread ordering. + __syncthreads(); + if (tx == 0) { red_release(&state->arrival_counter, 1); }