From 7b3d595eb197d714052ce296cc8b124f0dc8af31 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Wed, 24 Jun 2026 15:51:11 +0800 Subject: [PATCH] [CI/Build] Fix topk histogram build on SM75 (#46550) Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> (cherry picked from commit 191826ec612dc6648b176ed4e94c3e6e551e9c09) --- csrc/libtorch_stable/topk_histogram_4096.cuh | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/csrc/libtorch_stable/topk_histogram_4096.cuh b/csrc/libtorch_stable/topk_histogram_4096.cuh index 71c6c2cdf01..5f9f823a339 100644 --- a/csrc/libtorch_stable/topk_histogram_4096.cuh +++ b/csrc/libtorch_stable/topk_histogram_4096.cuh @@ -97,13 +97,22 @@ __device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane, } // Returns the sum of a value across all 32 threads in the warp, and every -// thread gets the same result SM90+ PTX instruction that does a hardware -// warp-wide reduction in a single instruction w.r.t. warp::reduce_sum(), which -// uses a __shfl_xor_sync butterfly tree (5 shuffles for 32 lanes) +// thread gets the same result. SM80+ uses redux.sync.add.u32, a single PTX +// instruction for hardware warp-wide reduction. Older targets use the +// __shfl_xor_sync butterfly tree, like warp::reduce_sum() (5 shuffles for 32 +// lanes). __device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) uint32_t r; asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v)); return r; +#else + #pragma unroll + for (uint32_t mask = kWarpSize >> 1; mask > 0; mask >>= 1) { + v += __shfl_xor_sync(0xFFFFFFFF, v, mask); + } + return v; +#endif } // ============================================================================ @@ -412,7 +421,7 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores, if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc; __syncthreads(); - // Step 3: Inter-warp prefix via redux.sync + // Step 3: Inter-warp prefix across warp sums. const auto tmp = smem->warp_sum[lane_id]; uint32_t prefix = warp_reduce_sum_full( lane_id < warp_id ? tmp : 0); // sum of all prior warps