[yangyuhuiling]

Problem

The fused GEMM in kernels_fused.py can return silently wrong results: a few wrong rows, no NaN, and only on some launches.

The last two K steps are peeled out of the main loop, and they pre-read the final step's operands before the wait that covers their producers:

# step K_ITERS - 2
b0_frag = b_s2r.load(b_cur0, ...)
a0_frag = a_s2r.load(a_cur0)
rocdl.s_barrier()                    # bare barrier: drains nothing
...
b0_frag = b_s2r.load(b_next0, ...)   # pre-reads the FINAL B0

# step K_ITERS - 1
a0_frag = a_s2r.load(a_cur0)         # pre-reads the FINAL A0
wait_barrier(0)                      # too late for either

wait_barrier(n) is s_waitcnt vmcnt(n) + s_barrier. A bare s_barrier is a rendezvous, not a completion wait, so those reads can take LDS that the global→LDS loads have not finished writing.

Why the existing A+B-1 did not fix it

The main loop's end-of-step wait was lowered from 2A+B to A+B-1, with a long comment calling that the fix for large-grid corruption. It made the tail's tiles reach LDS sooner, which hid the race on long K. It does not express the tail's dependency:

• At K=256 the main loop runs zero times (range_constexpr(K_ITERS - 2)), so no count inside it can help. ptpc 4096x7168x256 is wrong on 12/12 single-GPU launches with the current code.
• With the tail drained, the original 2A+B is clean on [16384, 7168] K=512, the shape that motivated A+B-1. Without the tail change, 2A+B still corrupts there 8/12.

Fix

• Drain the final-tile producers in the tail, where the pre-reads are (wait_barrier(0) in place of the bare barrier). The last step's existing wait_barrier(0) stays; it still covers the final A1, which is issued after it.
• Restore the main loop's original count, 2A+B, and replace the old comment.
• Add read_done_barrier() (lgkmcnt(0) + s_barrier) at the main loop's S0. The leading wave group refills B0 one stage after the barrier the lagging group's S0 pairs with, and LDS reads retire under LGKM, which the VM waits do not cover. This one is a statically found gap, not a demonstrated one: the tail drain alone clears every failure reproduced here. It costs nothing measurable.

Same diagnosis as ROCm/aiter#6368, which this pipeline was copied from.

Test

New test_ptpc_tail_reads_complete_lds in tests/python/cco/test_gemm_ar.py. It runs one GPU with no cco and calls compile_fused_gemm_scatter(..., fuse=False) directly, the same way test_ptpc_matches_an_fp32_reference does. It launches 12 times per shape, NaN-poisons the output before each launch, and counts elements off by more than 5% against a bf16-rounded fp32 reference. It checks element-wise rather than by relL2 because the corruption can sit close to the fp8 floor, where the bench's relL2 < 3e-3 gate can miss it.

kernel	4096x7168x256	16384x7168x512
main today (A+B-1)	fails 12/12	passes
2A+B, tail unchanged	fails	fails 8/12
this PR	passes	passes

Reproduced on two MI355X hosts with different stacks (torch 2.9.1 / ROCm 7.2 and torch 2.12 / ROCm 7.14). Worst single launch had 436,944 and 971,116 wrong elements respectively.

Full suites: test_gemm_ar.py + test_gemm_ar_op.py, 89 passed.

Cost

8x MI355X. Each run starts only after the box has been idle across several samples and is discarded if anything else is running when it ends. The arms alternate within each pass; the table is the min over 3 passes.

mode	quant	shape	before	after
gemm-only	ptpc	32768x7168x1792	437.1	412.8	−5.5%
gemm-only	blockscale	16384x7168x2048	377.9	362.3	−4.1%
gemm-only	mxfp8	32768x7168x1792	476.6	462.2	−3.0%
gemm-only	mxfp8	16384x7168x2048	246.8	246.2	−0.2%
fused-sdma	mxfp8	16384x7168x2048	1108.8	1108.1	−0.1%
fused-sdma	mxfp8	8192x7168x2048	587.1	587.0	0.0%

The GEMM gets back what A+B-1 charged on every K step. At 16384x7168x2048 there is no change, matching the old comment's note that A+B-1 was free at that shape. The fused path is unchanged within noise because it is communication-bound, so for the fused op this is a correctness fix rather than a speedup.

Scope

kernels_fused.py only, which is the pinned copy of the main loop that every gemm_ar mode compiles. _gemm_a8w8_8wave.py gains the helper, but its own compile_fp8_gemm_8w is not called anywhere in mori, so its tail is left as is. I have not re-read the generated ISA, and I have not swept every tile size.

Under mxfp8 the race did not reproduce even before this change. The scale loads in that path add a data dependency that makes the compiler emit a wait ahead of the final B0/A0 reads. That masks the race by accident, so the fix applies to all quant modes rather than relying on it.

🤖 Generated with Claude Code

[jhchouuu]

LGTM. Tested on MI355X (ROCm 7.14, flydsl 0.3.0), single GPU, compile_fused_gemm_scatter(..., fuse=False) with ptpc, using the same element-wise check as test_ptpc_tail_reads_complete_lds (NaN-poisoned output, >5% against a bf16-rounded fp32 reference). Three trees: main, this PR, and this PR with only the main-loop count put back to A+B-1:

	main	this PR	PR + A+B-1
4096x7168x256, x12	12/12 bad (worst 1,123,943 elements)	0/12	0/12
16384x7168x512, x12	0/12	0/12	0/12
5 large shapes (K 1792/2048, M up to 65536, 256x256 and 128x256 tiles), x40, idle and with a copy loop on the same GPU	0	0	0
K=8192, 65536x7168x2048 and 32768x7168x4096 (128x256), x100 on 3 GPUs each, with an elementwise-kernel loop on the same GPU	—	0/900	0/900

So the tail drain is what fixes the reported bug, and it is needed. I couldn't make the main loop fail at 2A+B under any load I could generate. Not covered: fuse=True, and mxfp8 with the narrow 256x128 tile.

Two non-blocking comments inline, both about what the code comments claim rather than about the result.

[jhchouuu]

Non-blocking. I tested this count hard (table in the review) and it didn't fail, but no wait in the loop guarantees it, and the comment's explanation of A+B-1 is incomplete.

Waves 0-3 run one barrier phase ahead of 4-7 (the prologue's if wave_m == 1: s_barrier()), so the barrier just before the leading group's S0 pairs with the lagging group's S2 barrier, which waits on nothing. Every LDS tile is written by all eight waves (G2SLoader: row r comes from wave (r % 64) // 8), so the leading group's S0 reads of B0(k+1) (wave_n 1 and 3), and at BLOCK_M=256 of A0(k+1) step 0, include rows the lagging group loaded. The only wait covering those loads is the lagging group's previous wait_barrier. In issue order, it guarantees B0(k+1) when the count is ≤ A+B, and A0(k+1)'s first step when it is ≤ A+B-1. So A+B-1 covers both and 2A+B covers neither; at 2A+B those reads are covered by elapsed time, about one iteration since the lagging group issued the loads. The tail's first two reads, before the new wait_barrier(0), are in the same position. Separately, 2A+B only keeps a wave's own A1(k) complete before S2 when A ≤ B, and mxfp8's narrow 256x128 tile has A > B, where again only the scale-load waits cover it.

Either way is a reasonable call: keep A+B-1 and give up the 3-5% on gemm-only (the fused path doesn't move), or keep 2A+B and say in this comment that the leading group's S0 reads rely on that slack. In both cases I'd drop "it only made the peeled tail's tiles land sooner", since in this loop the lower count is also the only thing that orders those reads.

[jhchouuu]

Non-blocking, about the comment. The lgkmcnt(0) makes this group's own S0 reads finish before any wave of the same group refills B0 at S1, which is a real ordering. It doesn't order the refill against the other group's reads, which is the case described here: the leading group issues the B0 refill in S1, before its S1 barrier, and that barrier is the one paired with the lagging group's S0 barrier. So the refill can already be in flight while the lagging group is still reading B0(k); what keeps it safe is that a global-to-LDS load can't land within the few cycles a DS read takes. Worth saying that in the comment, or moving the refill after the S1 barrier if you want it ordered.