Repository navigation
[FlyDSL] [gfx950] Fix 8-wave a8w8 GEMM reading LDS before the last tile lands - #6368
Conversation
…tep barrier The barrier that ends each main-loop K step admits `2 * N_LDS_STEPS_A + N_LDS_STEPS_B` outstanding global->LDS loads. That is more than the step can have in flight and still be safe: the barrier passes while writes the next step reads have not landed, and the next step consumes stale LDS. The result is a few wrong rows, not a NaN -- rel_l2 lands around 1e-3, close enough to fp8 quantisation noise to go unnoticed. Whether it fires depends on how the waves interleave, so a shape can be clean on one launch and wrong on the next, and larger grids hit it more often. Shapes in the tuned configs are affected. `M=32768 N=7168 K=1792` on the 256x256 tile is wrong on 4 of 4 runs, and `M=3456 N=7168 K=4224` on 128x256 on 1 of 4. Lowering the count to `N_LDS_STEPS_A + N_LDS_STEPS_B - 1` clears it: both shapes above are then correct and bit-identical across runs. A sweep of the 128x256 tile (where the two counts are 4 and 2) shows 4 and 3 corrupting and 2 and 0 clean, so this is the loosest count that is safe there. Waiting on more loads costs throughput, mostly on 128x256 and mostly at small grids where the K loop is short enough for the barrier to matter: tile grid before after 128x256 168 55.6us 65.0us +16.9% 128x256 224 60.6us 67.8us +11.9% 128x256 243 64.9us 68.0us +4.8% 128x256 252 63.7us 67.9us +6.6% 256x256 144 84.7us 87.0us +2.7% 256x256 363 167.8us 177.3us +5.7% 256x256 495 204.9us 204.7us -0.1% 256x256 3584 451.7us 479.0us +6.0% Measured on MI355X, min of per-arm medians over two alternating passes. Recovering this would need a tighter count per tile rather than one formula; correctness comes first here. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
The prologue's `if wave_m == 1: rocdl.s_barrier()` is executed by waves 4-7 only. Waves 0-3 skip it and reach the next barrier instead, so the hardware pairs the two groups one barrier apart and keeps them that way for the rest of the kernel. Documented rather than changed: making it unconditional while keeping the old wait count still corrupts, and with the corrected count both forms are clean, so it is not what this PR is fixing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The previous commit read the divergent `if wave_m == 1: s_barrier()` as a bug in itself. It is not: it deliberately puts waves 0-3 one phase ahead, which is the stagger the double-buffered main loop is built around. Making it unconditional destroys that and fixes nothing. What is missing is its counterpart. The stagger is never closed, so the two halves are still a barrier apart when the epilogue starts. gcnasm closes it before the same epilogue (opus_gemm_a2a_lsa quad-subtile template) and so does mori's fused copy of this pipeline, with `if wave_m == 0: rocdl.s_barrier()`. Still a note rather than a change: nothing follows store_c here, so the imbalance cannot be observed today. It stops being unobservable the moment anything is appended to the epilogue -- a `wait_barrier(0)` placed after store_c would rendezvous waves 0-3 with waves 4-7 still one barrier back and before their stores, rather than meaning what it says. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Tightening the main loop's end-of-step wait from 2A+B to A+B-1 was the wrong place for it, for two reasons. It does not cover every legal K. At K=256, K_ITERS is 2 and the main loop body runs zero times, so the changed line never executes: M=4096 N=7168 K=256 tile=256x256 is still wrong 6 runs out of 6. And it pays per K step for a dependency the tail owns. The peeled steps pre-read the last K step's B0 (end of step K_ITERS-2) and A0 (top of step K_ITERS-1), both before that step's wait_barrier(0). Tightening the main loop only made those tiles land in LDS earlier, which is why it happened to fix the long-K shapes. So: restore the main loop's original count and drain the final-tile producers in the tail, where the pre-reads are. The last step's wait stays -- it still covers the final A1, which is issued after it. Also order the WAR on B0 explicitly. The leading group refills that slot one stage after the barrier the lagging group's S0 pairs with, and LDS reads retire under LGKM, which this kernel's VM waits say nothing about. read_done_barrier() is lgkmcnt(0) + s_barrier. This one is a statically found gap: the tail drain alone already clears every shape here, so it is not separately demonstrated -- it costs nothing measurable and writes the dependency down. On 8x MI355X, 8 runs per shape with NaN-poisoned output, all clean including the two new K=256/512 entries. At 32768x7168x1792 the main loop's restored count is worth 6%: 445.2us against 476.3us for A+B-1. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Updated: the fix now sits in the tail, and the main loop's original count is restored. Thanks to whoever dug into this — the K=256 counterexample is correct and it changes the diagnosis.
It does not cover every legal K. At And it charges every K step for a dependency that belongs to the tail. The peeled steps pre-read the final K step's B0 (end of What it does now. Restore Separately, Measured, 8x MI355X, ptpc kernel, 8 runs per shape with NaN-poisoned output:
So the tail fix is both more correct and ~6% faster than Full shape list now passing 8/8: Scope I have not covered: this is the ptpc kernel only, on FlyDSL 0.3.4.1. I have not re-read the generated ISA to confirm the barrier pairing, and I have not run a full regression across every tile / K / scale layout. One note for anyone comparing with the mxfp8 kernels: #6407's already has the tail |
|
Re-measured everything and rewrote the PR body. One correction to my earlier comment, and a note on method. Correction. I quoted "445.2us against 476.3us for Method. This box is shared, and I caught a neighbour's job inflating timings by 10-13% mid-session — larger than the effect being measured. The run behind the new tables is gated: it starts only after six consecutive samples show every GPU under 5%, and is thrown away if anything is running when it finishes. Two intermediate runs were discarded that way. What the new tables say. The cost of fixing this in the main loop was up to +9.3% on the 128x256 grids and +7.8% at Correctness is unchanged from what I reported: One observation worth recording: which shapes |
Review asked for shorter comments. Each added block is now one line, the wait_barrier docstring is dropped, and the epilogue TODO about the unclosed half-wave stagger is removed -- it is recorded in the PR description. No code change. Co-Authored-By: Claude <noreply@anthropic.com>
Review preferred not adding a test file per fix. The existing op_tests/test_gemm_a8w8.py already runs every tuned flydsl row through gemm_a8w8_bpreshuffle, which is the only way this kernel is reached on gfx950 (untuned shapes fall back to CK). On 32768x7168x1792 its "ck bpreshuffle err" reads 0.05-0.09% before this fix and 0 after. Co-Authored-By: Claude <noreply@anthropic.com>
yangyuhuiling
changed the title
|
LGTM |
… the loop (#785) The fused GEMM's peeled tail pre-reads the last K step's B0 (end of step K_ITERS-2) and A0 (top of step K_ITERS-1), both before that step's wait_barrier(0). Nothing drains those tiles' global->LDS loads first, so the reads can take stale LDS: a few wrong rows, no NaN, intermittent. The main loop's end-of-step wait was lowered from 2A+B to A+B-1 to fix large-grid corruption, but that only made the tail's tiles land sooner. At K=256 the main loop runs zero times, so it cannot help: ptpc 4096x7168x256 is wrong on 12/12 single-GPU launches with it in place. Drain the producers in the tail itself, and restore the loop's original count. With the tail fixed, 2A+B is clean on [16384, 7168] K=512 -- the shape that motivated A+B-1, and which still corrupts 8/12 under 2A+B without the tail change. The last step's wait_barrier(0) stays; it still covers the final A1, issued after it. Also add read_done_barrier() (lgkmcnt(0) + s_barrier) at the 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; the tail drain alone clears every failure reproduced here, and it costs nothing measurable. Same diagnosis as ROCm/aiter#6368, which this pipeline was copied from. Measured on 8x MI355X, gemm-only, idle-gated, arms alternating: ptpc 32768x7168x1792 437.1 -> 412.8 us -5.5% blockscale 16384x7168x2048 377.9 -> 362.3 us -4.1% mxfp8 32768x7168x1792 476.6 -> 462.2 us -3.0% mxfp8 16384x7168x2048 246.8 -> 246.2 us -0.2% The fused gemm_ar path is unchanged within noise (communication-bound). Co-authored-by: xiangch <xiangch@amd.com> Co-authored-by: Claude <noreply@anthropic.com>
…0 read drain Replace read_done_barrier()'s inline asm with rocdl.s_waitcnt(lgkmcnt=0) + rocdl.s_barrier(), the form FlyDSL#1306 uses for the same fix. Same 54 "s_waitcnt lgkmcnt(0); s_barrier" pairs at K=7168 and the same VGPR/SGPR counts, but 212 (256x256) / 106 (128x256) fewer s_waitcnt overall: the inline asm was opaque to SIInsertWaitcnts, which then re-inserted lgkmcnt waits it could not prove were already satisfied. MI355X, FlyDSL 0.4.0: 4096x7168x256 and 32768x7168x1792 are 8/8 wrong on main and 0/8 with this change, as with the inline asm; timing is within noise of the inline asm on six shapes. Co-Authored-By: Claude <noreply@anthropic.com>
…efore the epilogue The prologue's `if wave_m == 1: s_barrier()` leaves waves 4-7 one barrier ahead in count, and the last unconditional s_barrier only paired with their previous one, so waves 4-7 sat in a barrier until waves 0-3 had stored and exited: the two groups' epilogues ran back to back. Issue the final barrier on wave_m == 0 only, as ROCm/FlyDSL#1306 does for the mxfp8 kernel, so the counts match and both epilogues overlap. Kept unconditional at BLOCK_N=512: that tile is already at 256 VGPRs, the branch raises spills from 4 to 11, and all three tuned 128x512 shapes measured 2.4-2.7% slower with it. The 256-wide tiles have headroom (256x256: 252 VGPRs either way; 128x256: 152 -> 156, no spills). MI355X, FlyDSL 0.4.0, 16 tuned shapes at their tuned tile/wpe/xcd, 7 interleaved rounds: geomean -2.1%, up to -6.1% (1472x7168x1536); large-M multi-wave shapes -0.7% to -2.1%. 18 shapes incl. 4096x7168x256 and 32768x7168x1792 stay 0/8 wrong with NaN-poisoned output. Co-Authored-By: Claude <noreply@anthropic.com>
…e added comments Issue the final barrier on wave_m == 0 for every tile, BLOCK_N=512 included, instead of keeping it unconditional there. Remove the comments this PR added. Co-Authored-By: Claude <noreply@anthropic.com>
* fix(gemm-ar): drain the peeled tail's producers instead of tightening the loop The fused GEMM's peeled tail pre-reads the last K step's B0 (end of step K_ITERS-2) and A0 (top of step K_ITERS-1), both before that step's wait_barrier(0). Nothing drains those tiles' global->LDS loads first, so the reads can take stale LDS: a few wrong rows, no NaN, intermittent. The main loop's end-of-step wait was lowered from 2A+B to A+B-1 to fix large-grid corruption, but that only made the tail's tiles land sooner. At K=256 the main loop runs zero times, so it cannot help: ptpc 4096x7168x256 is wrong on 12/12 single-GPU launches with it in place. Drain the producers in the tail itself, and restore the loop's original count. With the tail fixed, 2A+B is clean on [16384, 7168] K=512 -- the shape that motivated A+B-1, and which still corrupts 8/12 under 2A+B without the tail change. The last step's wait_barrier(0) stays; it still covers the final A1, issued after it. Also add read_done_barrier() (lgkmcnt(0) + s_barrier) at the 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; the tail drain alone clears every failure reproduced here, and it costs nothing measurable. Same diagnosis as ROCm/aiter#6368, which this pipeline was copied from. Measured on 8x MI355X, gemm-only, idle-gated, arms alternating: ptpc 32768x7168x1792 437.1 -> 412.8 us -5.5% blockscale 16384x7168x2048 377.9 -> 362.3 us -4.1% mxfp8 32768x7168x1792 476.6 -> 462.2 us -3.0% mxfp8 16384x7168x2048 246.8 -> 246.2 us -0.2% The fused gemm_ar path is unchanged within noise (communication-bound). Co-Authored-By: Claude <noreply@anthropic.com> * perf(gemm-ar): mxfp8 with row-major scales staged through LDS Add quant="mxfp8_raw": the same 32-wide ue8m0 GEMM, taking A's scales as [M, K/32] and B's as [N/32, K/32] uint8 -- what the quantiser and the checkpoint already hold -- and staging them through LDS, after aiter's ScaleRawLds. "mxfp8" (packed, preshuffle_a_scale) is unchanged and stays the default, so existing callers keep working. Reading the packed layout from global once per K step costs twice: the MFMA's vmcnt wait on the scale also drains the tile prefetches issued before it (63 vmcnt waits at K=2048, 20 here), and the scale registers spill in a kernel already at 256 VGPRs (8 at K=512, none here). Each wave now copies its share of four K steps into LDS with buffer_load ... lds, double-buffered, and the pipeline's own wait_barrier covers the copies. GEMM only, MI355X, CUDA-graph timed, idle-gated, modes alternating, cold weights: -9% to -15% at M=64/1024 on the 256x128 tile, -1% to -12% on the 256x256 tile across 10 shapes, +0.9% at K=640 where five K steps do not amortise two chunks. Fused-sdma at TP4 wo_b is unchanged (+-0.3%): that layer is bound by the collective. No padding: any M goes in. Co-Authored-By: Claude <noreply@anthropic.com> * fix(gemm-ar): stage mxfp8_raw scale chunks with the offset in the VGPR The chunk offset went in through soffset, which the raw buffer range check does not cover, so the last row's bytes past K were read from beyond the scale tensor rather than zero-filled. They are never consumed, but the read could reach past the allocation. With the offset in the per-lane address the check covers it, and the comment that claimed zero-fill is now true. Also qualify "faster than mxfp8": it holds cold and at large M; hot at small M or short K it can be a few percent slower. Co-Authored-By: Claude <noreply@anthropic.com> --------- Co-authored-by: xiangch <xiangch@amd.com> Co-authored-by: Claude <noreply@anthropic.com>
Problem
The 8-wave a8w8 bpreshuffle GEMM can return wrong results: a few wrong rows, no NaN, only on some launches.
The peeled tail reads the last K step's B0 and A0 from LDS before the wait that covers their global→LDS loads:
Fix
s_barrier()withwait_barrier(0), so the final tile has landed before it is read. The last step's existingwait_barrier(0)stays (it covers the final A1).rocdl.s_waitcnt(lgkmcnt=0)+rocdl.s_barrier()(the same form as 1,Balance the prologue's extra wave_m == 1 barrier so both groups' ep… FlyDSL#1306, no inline asm), so LDS reads finish before the other wave group refills B0. This one is a statically found gap; the tail change alone fixes every failure reproduced here, and it costs nothing measurable.if wave_m == 1: rocdl.s_barrier()leaves waves 4-7 one barrier ahead, so the last unconditionals_barrieronly paired with their previous one and waves 4-7 waited until waves 0-3 had stored and exited: the two groups' epilogues ran back to back. The final barrier is now issued onwave_m == 0only, as 1,Balance the prologue's extra wave_m == 1 barrier so both groups' ep… FlyDSL#1306 does for the mxfp8 kernel.An earlier revision of this PR instead lowered the main loop's end-of-step wait from
2A+BtoA+B-1. That only made the final tiles land sooner on long K, charged every K step for it, and does nothing atK=256, where the main loop runs zero times. It has been reverted; the main loop keeps its original count. Thanks to the reviewer who found theK=256counterexample.Correctness
MI355X, 8 launches per shape, output NaN-poisoned before each, counting elements off by >5%:
A+B-1Which shapes fail on
mainvaries with machine load, since it depends on how the waves interleave.Cost
The synchronisation fix alone (before the epilogue change), same harness for all three, idle-gated:
A+B-1−0.5% to +2.2% against
mainacross the ten shapes measured, against up to +9.3% forA+B-1.Net, whole PR against
mainMI355X, FlyDSL 0.4.0, the 16 tuned rows below at their tuned tile /
waves_per_eu/xcd_swizzle, 7 rounds interleaved per shape, idle-gated, median µs (last column: rounds where this PR was faster):Geomean -1.0%. The 256×256 rows gain -0.7% to -5.7% from the epilogue overlap; the long-K 128×256 rows stay +1.6% to +3.1%, where the sync fix's cost is larger than the overlap. The 128×512 row is from a separate 7-round run against
main; all three tuned 128×512 rows are +1.2% to +2.2% (that tile is already at 256 VGPRs, and the epilogue branch raises spills from 4 to 11). The same 16 rows plus4096x7168x256and32768x7168x1792are 0/8 wrong with NaN-poisoned output.Test
No new test file.
op_tests/test_gemm_a8w8.pyalready runs every tuned flydsl row throughgemm_a8w8_bpreshuffle, which is the only way this kernel is reached on gfx950 (untuned shapes fall back to CK). On32768x7168x1792itsck bpreshuffle erris 0.05–0.09% before this fix and 0 after. That column is reported, not asserted, so it will not fail CI on its own.K=256is not reachable through the dispatch; it was run against the kernel directly.Scope
gfx950 ptpc kernel only. The mxfp8 8-wave kernel from #6407 already has the tail
wait_barrier(0).🤖 Generated with Claude Code