From 20ae9e83f95eea6059106717c2f4a3f91d820dad Mon Sep 17 00:00:00 2001 From: Benson Ma Date: Fri, 12 Jun 2026 12:58:14 -0700 Subject: [PATCH] Fix ROCm __syncthreads deadlock in compute_amax_and_quantize_kernel Summary: Fixes a ROCm/AMD GPU deadlock in the FP4 fused amax+quantize kernel, found by auditing fbgemm_gpu for the same barrier-divergence pattern fixed in D107554507. ## The bug In compute_amax_and_quantize_kernel (quantize.cu), threads whose idx >= n hit an early `return` (quantize.cu:1958-1959) before reaching a block-wide reduction. The launch fp4_fused_amax_quantize uses block = dim3(blocksize, blocks_per_cta=4) and blocks = ceil_div(numel, blocksize*4), so when numel is not a multiple of blocksize*4 the tail block has some threads that return early while the rest enter the block reduction. The surviving threads stall forever on a block-wide barrier -> deadlock on ROCm (latent UB on NVIDIA). ## Where the hanging barrier actually is (not visible in the diff) The blocking barriers are NOT in the kernel body -- they are reached transitively through a device helper, which is why the early return was easy to miss. Call chain: 1. Kernel calls compute_max() quantize.cu:1964 2. compute_max() branches on THREAD_X: quantize.cu:1864-1872 - THREAD_X == 32 -> compute_max_warp() (warp-only, no block barrier) - else -> compute_max_block() quantize.cu:1870 The sole instantiation is <__nv_bfloat16, 16, 4> (quantize.cu:1993), so THREAD_X == 16 and it ALWAYS takes the compute_max_block() path. 3. compute_max_block() (quantize.cu:1825-1846) has two block-wide barriers: - implicit __syncthreads() inside cub::BlockReduce::Reduce() quantize.cu:1837 - explicit __syncthreads() quantize.cu:1843 Both gate the whole physical 64-thread block (dim3(blocksize, 4)), even though cub is declared BlockReduce with per-row temp_storage[threadIdx.y]. So a thread that returns early in ANY row stalls the entire block. ## The fix Remove the early return and mask the work with `active = idx < n`: - Inactive lanes load a neutral 0.0f. The reduction is a fabsf-max, and fabsf(0.0f) = 0 <= any real |x|, so inactive lanes cannot perturb block_amax. - compute_max() is called by ALL threads -> full block participation, both __syncthreads() are reached by every thread. - Only active lanes (idx < n) store y[idx]; no OOB read or write. Behavior is bit-identical for active lanes. blocks = ceil_div(...) guarantees the tail block always has >= 1 active lane, so block_amax is never degenerate. ## Caveat This deadlock exists only because the sole launch uses THREAD_X = 16. With THREAD_X = 32, compute_max() would take the compute_max_warp() path (no block __syncthreads()) and the early return would have been safe. The fix is correct for the code as it exists today and remains correct if a non-32 THREAD_X is added. Reviewed By: henrylhtsang Differential Revision: D107946896 --- .../experimental/gen_ai/src/quantize/quantize.cu | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cu b/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cu index 720c0cf125..e54665399a 100644 --- a/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cu +++ b/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cu @@ -1955,12 +1955,15 @@ __global__ void compute_amax_and_quantize_kernel( const auto warp_id = threadIdx.y; const int64_t idx = (block_idx + warp_id) * blocksize + thread_id; - if (idx >= n) - return; + // Do NOT early-return for out-of-range threads: compute_max() below uses a + // block-wide cub::BlockReduce + __syncthreads() that every thread must reach, + // otherwise the block deadlocks on ROCm/AMD. Mask the work with `active`. + const bool active = idx < n; float scale = 1.f, unscale = 1.f; - float thread_val = x[idx]; + // Inactive lanes feed 0 into the fabsf-max reduction (neutral). + float thread_val = active ? static_cast(x[idx]) : 0.0f; float block_amax = compute_max(thread_val, blocksize); if (global_amax) { @@ -1970,7 +1973,9 @@ __global__ void compute_amax_and_quantize_kernel( compute_scale(scale, unscale, block_amax); } - y[idx] = quantize(thread_val, scale, unscale); + if (active) { + y[idx] = quantize(thread_val, scale, unscale); + } } int ceil_div(const int a, const int b) {