[Bugfix] Sync waves before the RMSNorm/LayerNorm quant kernels reuse their reduction slot

[jhinpan]

Summary

The RMSNorm and LayerNorm quant kernels run two block reductions per row through one LDS buffer: a sum, then the row max for the quant scale. block_reduce_max lets wave 0 store its partial max into slot 0 before every wave has loaded the sum from that slot. This PR adds a barrier before that store.

Motivation

The end of the sum reduction and the start of the max reduction are:

s_red[0] = total                     # wave 0, lane 0
barrier
sum = s_red[0]                       # every wave
...                                  # normalize the row and find the thread max
s_red[wave] = wave_max(thread_max)   # lane 0 of every wave; wave 0 writes slot 0
barrier

No barrier orders a wave's sum = s_red[0] before wave 0's store into slot 0. A wave that loads slot 0 late uses wave 0's partial max as the sum. That wave normalizes its part of the row with a wrong rrms (LayerNorm: a wrong mean), so its outputs, the row max, and the quant scale of the whole row are wrong.

block_reduce in softmax_kernel.py reuses one buffer for its max and sum reductions in the same way. It already starts with a barrier for this reason (#1022).

Affected paths, when the workgroup has more than one wave:

• RMSNorm quant and fused-add RMSNorm quant (dynamic and smooth), on the generic path, which runs when the 128-bit tile does not divide N. The whole-tile path reads its sum of squares from s_red2 and is not affected.
• LayerNorm quant and fused-add LayerNorm quant (dynamic and smooth), on every path. They read the sum from s_sum[0], which block_reduce_max reuses.

The race shows only when wave 0 finishes the whole second pass over the row and its wave reduction before another wave runs the load that follows the barrier.

Changes

• block_reduce_max in _build_rmsnorm_quant_module, _build_fused_add_rmsnorm_quant_module, _build_layernorm_quant_module, and _build_fused_add_layernorm_quant_module runs gpu.barrier() before lane 0 of each wave stores its partial max.
• The kernels run one more workgroup barrier per row. Single-wave workgroups (RED_SLOTS == 1) reduce in registers and do not change.

Testing

• Unit tests added/updated: the existing quant tests in tests/kernels/test_rmsnorm.py and tests/kernels/test_layernorm.py run every changed builder. A test cannot force the wave order that triggers the race.
• Performance benchmarks run
• Tested on MI300X

Dependencies

• No new third-party dependencies added
• If new dependencies added, justification provided

Breaking Changes

None.

Related PRs

• [Kernel][Perf] Tune RMSNorm for gfx1250: BF16 hardware conversion, early loads, small-N kernel choice #1227 adds N = 7168, 2880, and 1000 to the RMSNorm quant test configs. These N run the generic path that this PR fixes.

[cschenjunlin]

This PR adds a gpu.barrier() to block_reduce_max in all four quant builders. It fixes the race caused by reusing the same LDS slot for the sum and max reductions.
However, this PR is based on the main branch as of September 29. Commit 6487994 in #848 has already fixed the same issue, and the current main already contains the following code in rmsnorm_kernel.py (lines 1133–1140):

# s_red still holds the sum reduction's result, which every wave
# broadcast-reads on exit. Only this barrier orders that read
# against the store below; without it a late wave reads this max.
gpu.barrier()
w = wave_reduce_max(val)
if lane == 0:
    fx.memref_store(w, s_red, wave)
gpu.barrier()

I used git merge-tree to simulate the merge. This PR merges cleanly with main, and the redundant barrier could be merged into main without noticing. Therefore, I recommend removing one of duplicative gpu.barrier().

[jhinpan]

Thanks for checking this. #848 (squashed as 1bd7ebd5) already added the same barrier at the head of block_reduce_max in all four quant builders, and a merge with current main would only add a second barrier right after it. Closing as superseded by #848.