From a84423baacf1a9c66eb06b2f02db88f7705fcadd Mon Sep 17 00:00:00 2001 From: Jin Pan Date: Wed, 7 Oct 2026 13:15:19 -0700 Subject: [PATCH] fix(norm): sync waves before block_reduce_max reuses the reduction slot In the RMSNorm and LayerNorm quant kernels, block_reduce_max has lane 0 of each wave store its partial max into s_red[wave] (s_sum[wave] in LayerNorm). Slot 0 still holds the result of the preceding block reduction, which every wave loads after that reduction's last barrier, and nothing orders those loads before wave 0's store. A wave that loads slot 0 late reads wave 0's partial max instead of the sum, so it normalizes its part of the row with a wrong value, and the row max and the quant scale of the whole row are wrong too. Add a barrier before the first store, as block_reduce in softmax_kernel already does for its max and sum reductions. The RMSNorm whole-tile path reads its sum of squares from s_red2 and is not affected; the RMSNorm generic path and both LayerNorm paths read slot 0. Signed-off-by: Jin Pan --- kernels/norm/layernorm_kernel.py | 4 ++++ kernels/norm/rmsnorm_kernel.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/kernels/norm/layernorm_kernel.py b/kernels/norm/layernorm_kernel.py index 150341eb06..a488801541 100644 --- a/kernels/norm/layernorm_kernel.py +++ b/kernels/norm/layernorm_kernel.py @@ -877,6 +877,8 @@ def block_reduce_max(val): lane = tid % WARP_SIZE wave = tid // WARP_SIZE w = wave_reduce_max(val) + # Other waves may still be reading the previous block reduction's result from s_sum[0]. + gpu.barrier() if lane == 0: fx.memref_store(w, s_sum, wave) gpu.barrier() @@ -1226,6 +1228,8 @@ def block_reduce_max(val): lane = tid % WARP_SIZE wave = tid // WARP_SIZE w = wave_reduce_max(val) + # Other waves may still be reading the previous block reduction's result from s_sum[0]. + gpu.barrier() if lane == 0: fx.memref_store(w, s_sum, wave) gpu.barrier() diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index 7c18d43312..b814dbb3f1 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -842,6 +842,8 @@ def block_reduce_max(val): wave = tid // WARP_SIZE w = wave_reduce_max(val) + # Other waves may still be reading the previous block reduction's result from s_red[0]. + gpu.barrier() if lane == 0: fx.memref_store(w, s_red, wave) gpu.barrier() @@ -1211,6 +1213,8 @@ def block_reduce_max(val): wave = tid // WARP_SIZE w = wave_reduce_max(val) + # Other waves may still be reading the previous block reduction's result from s_red[0]. + gpu.barrier() if lane == 0: fx.memref_store(w, s_red, wave) gpu.barrier()