From 2e135cf7703961a0c6524b9a4d2e879aa9df6d74 Mon Sep 17 00:00:00 2001
From: Elias Stehle <3958403+elstehle@users.noreply.github.com>
Date: Wed, 5 Aug 2026 12:38:28 +0200
Subject: [PATCH] Fixes a race condition in `DeviceTopK` (#10627)

* fixes race in DeviceTopK

* changes order

(cherry picked from commit 879ab26a112fbbbeb86ad96e1d2cc5c657c890a9)
---
 cub/cub/agent/agent_topk.cuh | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/cub/cub/agent/agent_topk.cuh b/cub/cub/agent/agent_topk.cuh
index c6cf74ad59d..ebdebb11e02 100644
--- a/cub/cub/agent/agent_topk.cuh
+++ b/cub/cub/agent/agent_topk.cuh
@@ -645,10 +645,12 @@ struct AgentTopK
     bool is_last_pass,
     CounterUpdateFn counter_update_fn)
   {
-    // Ensure all writes to the global memory-histogram are visible to all threads before
-    // proceeding to compute the prefix sum over the histogram.
+    // Make this block's histogram contributions visible device-wide before publishing retirement
     __threadfence();
 
+    // Wait for all threads in this block to finish merge_histograms() before thread 0 signals completion
+    __syncthreads();
+
     // Identify the last block in the grid to perform the prefix sum over the histogram
     bool is_last_block = false;
     if (threadIdx.x == 0)