From 2463e41fbe3893b552fdc6a1c26c6b4841a169b9 Mon Sep 17 00:00:00 2001 From: xiangch Date: Sat, 10 Oct 2026 08:17:14 +0000 Subject: [PATCH 1/9] [FlyDSL] [gfx950] Fix 8-wave a8w8 GEMM dropping LDS writes at the K-step 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) --- aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py | 4 +- op_tests/flydsl_tests/test_gemm_a8w8_8wave.py | 98 +++++++++++++++++++ 2 files changed, 101 insertions(+), 1 deletion(-) create mode 100644 op_tests/flydsl_tests/test_gemm_a8w8_8wave.py diff --git a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py index e4b6ddf3d07..4504868e7f9 100644 --- a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py +++ b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py @@ -463,7 +463,9 @@ def kernel_gemm( c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) b_g2s.load(b_cur1, B1_gl_offset + (k + 2) * B_K_STEP) - wait_barrier(2 * N_LDS_STEPS_A + N_LDS_STEPS_B) + # Letting 2 * A + B loads stay outstanding here crosses the barrier + # with LDS writes still in flight that the next iteration reads. + wait_barrier(N_LDS_STEPS_A + N_LDS_STEPS_B - 1) c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) diff --git a/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py b/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py new file mode 100644 index 00000000000..ae5f22b18cb --- /dev/null +++ b/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Correctness of the 8-wave a8w8 bpreshuffle GEMM at large grids. + +The main-loop barrier that ends each K step admits a bounded number of +outstanding global->LDS loads. Admitting too many lets the barrier pass while +writes the next step reads are still in flight, which shows up as a handful of +wrong rows rather than a NaN -- small enough to pass for quantisation noise. +Whether it happens at all depends on how the waves interleave, so a shape can +be clean on one launch and wrong on the next; each shape is therefore run +several times. +""" + +import argparse + +import torch + +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.flydsl.gemm_a8w8_bpreshuffle_8wave import flydsl_8wave_gemm_a8 +from aiter.ops.shuffle import shuffle_weight + +torch.set_default_device("cuda") + +SUPPORTED_GFX = ["gfx950"] + +# (M, N, K, block_m, block_n), drawn from the tuned configs that select this +# kernel. The two large ones are where the barrier count actually bites. +SHAPES = [ + (1472, 3584, 7168, 128, 256), + (3392, 2176, 7168, 128, 256), + (3456, 7168, 4224, 128, 256), + (1920, 7168, 1792, 256, 256), + (2624, 8448, 7168, 256, 256), + (32768, 7168, 1792, 256, 256), +] +REPS = 6 +# A dropped load leaves whole rows wrong by far more than this, so counting +# offending elements separates cleanly: zero when correct, thousands when not. +# An L2 norm would not -- the bad rows are a small fraction of a large output. +REL_TOL = 0.05 + + +def run_shape(m, n, k, block_m, block_n, reps, verbose=False): + gen = torch.Generator(device="cuda").manual_seed(0) + a = (torch.randn(m, k, generator=gen) / 8).to(torch.float8_e4m3fn) + b = (torch.randn(n, k, generator=gen) / 8).to(torch.float8_e4m3fn) + a_scale = torch.rand(m, 1, generator=gen) + 0.5 + b_scale = torch.rand(n, 1, generator=gen) + 0.5 + b_shuffled = shuffle_weight(b, (16, 16)) + + # rounded to the kernel's output precision, so the comparison sees the + # kernel's error rather than bf16's + ref = ((a.float() * a_scale) @ (b.float() * b_scale).T).to(torch.bfloat16).float() + + # Floor the per-element tolerance well below a typical output but above + # zero: elements that cancel to near nothing disagree with torch on + # summation order alone, and that is not what this test is looking for. + floor = 0.01 * ref.square().mean().sqrt() + + worst = 0 + for _ in range(reps): + out = torch.empty(m, n, dtype=torch.bfloat16) + flydsl_8wave_gemm_a8(a, b_shuffled, a_scale, b_scale, out, block_m, block_n) + diff = (out.float() - ref).abs() + bad = int((diff > REL_TOL * ref.abs().clamp(min=floor)).sum()) + worst = max(worst, bad) + if verbose: + print(f" bad_elems {bad}") + return worst + + +def test_gemm_a8w8_8wave(): + if get_gfx() not in SUPPORTED_GFX: + return + for m, n, k, block_m, block_n in SHAPES: + worst = run_shape(m, n, k, block_m, block_n, REPS) + assert worst == 0, ( + f"M={m} N={n} K={k} tile={block_m}x{block_n}: " + f"{worst} elements off by more than {REL_TOL:.0%} in {REPS} runs" + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reps", type=int, default=REPS) + args = parser.parse_args() + + failed = [] + for m, n, k, block_m, block_n in SHAPES: + print(f"M={m} N={n} K={k} tile={block_m}x{block_n}") + worst = run_shape(m, n, k, block_m, block_n, args.reps, verbose=True) + print(f" worst bad_elems {worst} {'ok' if worst == 0 else 'FAIL'}") + if worst: + failed.append((m, n, k, block_m, block_n)) + print(f"\n{len(SHAPES) - len(failed)}/{len(SHAPES)} shapes within tolerance") + for shape in failed: + print(f" failed: {shape}") From 65d103763b7220fa4f73b68958bd3a0f3b6964c9 Mon Sep 17 00:00:00 2001 From: xiangch Date: Sat, 10 Oct 2026 08:23:43 +0000 Subject: [PATCH 2/9] [FlyDSL] [gfx950] Note the divergent prologue barrier in the 8-wave GEMM 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) --- aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py index 4504868e7f9..f53c7b7607c 100644 --- a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py +++ b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py @@ -431,6 +431,19 @@ def kernel_gemm( b_g2s.load(b_cur1, B1_gl_offset + 0 * B_K_STEP) a_g2s.load(a_cur1, A1_gl_offset + 0 * BLOCK_K) + # FIXME: this barrier is divergent and has no counterpart on the other + # half of the workgroup. wave_m is wave_id // 4, so waves 4-7 execute it + # and waves 0-3 do not; waves 0-3 reach the barrier below instead and + # the hardware releases the two groups together. Every later barrier is + # then pairing waves that are one barrier apart, which is not what any + # of them was written to express. + # + # Left alone here because it is not the cause of the corruption the + # wait count below fixes -- making it unconditional while keeping the + # old count still produces wrong results (3 of 6 runs on M=32768 N=7168 + # K=1792), and with the corrected count both forms are clean. Removing + # the divergence is a separate change that needs its own analysis of + # what the prologue barriers are meant to order. if wave_m == 1: rocdl.s_barrier() From 9b0de95479def0b9d2a6721cc5fa74067b23b759 Mon Sep 17 00:00:00 2001 From: xiangch Date: Sat, 10 Oct 2026 08:29:46 +0000 Subject: [PATCH 3/9] [FlyDSL] [gfx950] Correct the note on the prologue's half-wave barrier 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) --- aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py | 33 +++++++++++++-------- 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py index f53c7b7607c..019c24a47ce 100644 --- a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py +++ b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py @@ -431,19 +431,14 @@ def kernel_gemm( b_g2s.load(b_cur1, B1_gl_offset + 0 * B_K_STEP) a_g2s.load(a_cur1, A1_gl_offset + 0 * BLOCK_K) - # FIXME: this barrier is divergent and has no counterpart on the other - # half of the workgroup. wave_m is wave_id // 4, so waves 4-7 execute it - # and waves 0-3 do not; waves 0-3 reach the barrier below instead and - # the hardware releases the two groups together. Every later barrier is - # then pairing waves that are one barrier apart, which is not what any - # of them was written to express. + # Opens a deliberate half-wave stagger: this gives waves 4-7 one extra + # barrier, and since s_barrier is a counting rendezvous, waves 0-3 run + # one phase ahead from here on -- which is what the double-buffered + # main loop below wants. Do not make it unconditional; that destroys + # the stagger without fixing anything (still 3 of 6 runs wrong on + # M=32768 N=7168 K=1792 with the old wait count). # - # Left alone here because it is not the cause of the corruption the - # wait count below fixes -- making it unconditional while keeping the - # old count still produces wrong results (3 of 6 runs on M=32768 N=7168 - # K=1792), and with the corrected count both forms are clean. Removing - # the divergence is a separate change that needs its own analysis of - # what the prologue barriers are meant to order. + # It is never closed, though. See the note before the epilogue. if wave_m == 1: rocdl.s_barrier() @@ -540,6 +535,20 @@ def kernel_gemm( rocdl.s_setprio(0) rocdl.s_barrier() + # TODO: the half-wave stagger the prologue opened is never closed. The + # matching `if wave_m == 0: rocdl.s_barrier()` belongs here, before the + # epilogue, so the two halves are level again -- gcnasm closes it at + # the same point (opus_gemm_a2a_lsa, quad-subtile template) and so does + # mori's fused copy of this pipeline. + # + # Nothing follows store_c here, so the imbalance is not observable + # today and this is left as a note rather than a change. It does not + # stay unobservable: with the counts unbalanced, a `wait_barrier(0)` + # added after store_c does not mean "every wave's C tile has retired", + # it rendezvouses waves 0-3 with waves 4-7 still one barrier back and + # before their stores. Anything appended to this epilogue is silently + # wrong until the pair is closed. + # Scale and store back to gmem wave_n_offset = wave_n * (N_TILES_B * 16) wave_m_offset = wave_m * (N_TILES_A * 16) From d7c2e395a835fb532f6c8ff4aaddb7dfc645fe7c Mon Sep 17 00:00:00 2001 From: xiangch Date: Fri, 16 Oct 2026 02:10:49 +0000 Subject: [PATCH 4/9] Move the fix into the tail, where the dependency is 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) --- aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py | 40 ++++++++++++++++--- op_tests/flydsl_tests/test_gemm_a8w8_8wave.py | 27 ++++++++----- 2 files changed, 51 insertions(+), 16 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py index 019c24a47ce..56adb6c409d 100644 --- a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py +++ b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py @@ -65,6 +65,12 @@ def pack_i32x4_i32x8(lo, hi): def wait_barrier(count): + """``s_waitcnt vmcnt(count)`` then ``s_barrier``. + + The VM wait constrains only the requests *this* wave issued; the barrier is + what the workgroup rendezvouses on. A bare ``s_barrier`` is neither -- it + does not drain anything, so it cannot stand in for a completion wait. + """ _llvm.inline_asm( res=None, operands_=[], @@ -74,6 +80,25 @@ def wait_barrier(count): ) +def read_done_barrier(): + """``s_waitcnt lgkmcnt(0)`` then ``s_barrier``: reads done before refill. + + The leading wave group refills B0 one stage after the barrier the lagging + group's S0 pairs with, so the two groups need an explicit + read-complete -> rendezvous -> overwrite order on that slot. LDS reads + retire under LGKM, which the VM waits elsewhere in this kernel say nothing + about. The memory clobber keeps the compiler from sinking a DS read past + the barrier. + """ + _llvm.inline_asm( + res=None, + operands_=[], + asm_string="s_waitcnt lgkmcnt(0)\ns_barrier", + constraints="~{memory}", + has_side_effects=True, + ) + + class G2SLoader: def __init__(self, gl_src, gl_offsets, n_load_steps, lds_dtype, wave_id): self.g2lds_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128) @@ -454,7 +479,9 @@ def kernel_gemm( b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled) a0_frag = a_s2r.load(a_cur0) a_g2s.load(a_next1, A1_gl_offset + (k + 1) * BLOCK_K) - rocdl.s_barrier() + # The B0 slot just read is refilled by the peer group right after + # this rendezvous; make the reads complete before advertising it. + read_done_barrier() c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) @@ -471,9 +498,7 @@ def kernel_gemm( c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) b_g2s.load(b_cur1, B1_gl_offset + (k + 2) * B_K_STEP) - # Letting 2 * A + B loads stay outstanding here crosses the barrier - # with LDS writes still in flight that the next iteration reads. - wait_barrier(N_LDS_STEPS_A + N_LDS_STEPS_B - 1) + wait_barrier(2 * N_LDS_STEPS_A + N_LDS_STEPS_B) c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) @@ -487,7 +512,12 @@ def kernel_gemm( k = K_ITERS - 2 b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled) a0_frag = a_s2r.load(a_cur0) - rocdl.s_barrier() + # The peeled steps below pre-read step K_ITERS - 1's B0 and A0 -- B0 at + # the end of this step, A0 at the top of the next -- both *before* the + # last step's wait_barrier(0). Those tiles were issued two steps back, + # so drain them here, where the peer group's reads pair. The last + # step's wait stays: it still covers the final A1, issued below. + wait_barrier(0) c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) diff --git a/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py b/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py index ae5f22b18cb..3bd4ab990a0 100644 --- a/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py +++ b/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py @@ -1,15 +1,17 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Correctness of the 8-wave a8w8 bpreshuffle GEMM at large grids. - -The main-loop barrier that ends each K step admits a bounded number of -outstanding global->LDS loads. Admitting too many lets the barrier pass while -writes the next step reads are still in flight, which shows up as a handful of -wrong rows rather than a NaN -- small enough to pass for quantisation noise. -Whether it happens at all depends on how the waves interleave, so a shape can -be clean on one launch and wrong on the next; each shape is therefore run -several times. +"""Correctness of the 8-wave a8w8 bpreshuffle GEMM. + +The peeled tail pre-reads the last K step's B0 and A0 before the wait that +covers their producers, so it can read LDS the producing loads have not +finished writing. It shows up as a handful of wrong rows rather than a NaN -- +small enough to pass for quantisation noise. Whether it happens at all depends +on how the waves interleave, so a shape can be clean on one launch and wrong on +the next; each shape is therefore run several times. + +``K=256`` is the shape that pins the fix to the tail: ``K_ITERS == 2`` makes the +main loop run zero times, so no amount of tightening *there* can cover it. """ import argparse @@ -24,9 +26,12 @@ SUPPORTED_GFX = ["gfx950"] -# (M, N, K, block_m, block_n), drawn from the tuned configs that select this -# kernel. The two large ones are where the barrier count actually bites. +# (M, N, K, block_m, block_n). Most are drawn from the tuned configs that +# select this kernel; K=256 and K=512 are not tuned shapes but are legal, and +# the first is the one the main loop cannot reach. SHAPES = [ + (4096, 7168, 256, 256, 256), + (4096, 7168, 512, 256, 256), (1472, 3584, 7168, 128, 256), (3392, 2176, 7168, 128, 256), (3456, 7168, 4224, 128, 256), From f6fa9ac62407129534bb0a834701c74bfbcf26f7 Mon Sep 17 00:00:00 2001 From: xiangch Date: Mon, 19 Oct 2026 02:57:30 +0000 Subject: [PATCH 5/9] [FlyDSL] [gfx950] Trim comments in the 8-wave GEMM sync fix 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 --- aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py | 48 ++----------------- op_tests/flydsl_tests/test_gemm_a8w8_8wave.py | 26 +++------- 2 files changed, 10 insertions(+), 64 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py index 56adb6c409d..bc7b9809ecb 100644 --- a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py +++ b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py @@ -65,12 +65,6 @@ def pack_i32x4_i32x8(lo, hi): def wait_barrier(count): - """``s_waitcnt vmcnt(count)`` then ``s_barrier``. - - The VM wait constrains only the requests *this* wave issued; the barrier is - what the workgroup rendezvouses on. A bare ``s_barrier`` is neither -- it - does not drain anything, so it cannot stand in for a completion wait. - """ _llvm.inline_asm( res=None, operands_=[], @@ -81,15 +75,7 @@ def wait_barrier(count): def read_done_barrier(): - """``s_waitcnt lgkmcnt(0)`` then ``s_barrier``: reads done before refill. - - The leading wave group refills B0 one stage after the barrier the lagging - group's S0 pairs with, so the two groups need an explicit - read-complete -> rendezvous -> overwrite order on that slot. LDS reads - retire under LGKM, which the VM waits elsewhere in this kernel say nothing - about. The memory clobber keeps the compiler from sinking a DS read past - the barrier. - """ + """Finish this wave's LDS reads, then barrier.""" _llvm.inline_asm( res=None, operands_=[], @@ -456,14 +442,7 @@ def kernel_gemm( b_g2s.load(b_cur1, B1_gl_offset + 0 * B_K_STEP) a_g2s.load(a_cur1, A1_gl_offset + 0 * BLOCK_K) - # Opens a deliberate half-wave stagger: this gives waves 4-7 one extra - # barrier, and since s_barrier is a counting rendezvous, waves 0-3 run - # one phase ahead from here on -- which is what the double-buffered - # main loop below wants. Do not make it unconditional; that destroys - # the stagger without fixing anything (still 3 of 6 runs wrong on - # M=32768 N=7168 K=1792 with the old wait count). - # - # It is never closed, though. See the note before the epilogue. + # Intentional: waves 0-3 run one barrier ahead of waves 4-7. if wave_m == 1: rocdl.s_barrier() @@ -479,8 +458,7 @@ def kernel_gemm( b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled) a0_frag = a_s2r.load(a_cur0) a_g2s.load(a_next1, A1_gl_offset + (k + 1) * BLOCK_K) - # The B0 slot just read is refilled by the peer group right after - # this rendezvous; make the reads complete before advertising it. + # The other wave group refills B0 right after this barrier. read_done_barrier() c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) @@ -512,11 +490,7 @@ def kernel_gemm( k = K_ITERS - 2 b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled) a0_frag = a_s2r.load(a_cur0) - # The peeled steps below pre-read step K_ITERS - 1's B0 and A0 -- B0 at - # the end of this step, A0 at the top of the next -- both *before* the - # last step's wait_barrier(0). Those tiles were issued two steps back, - # so drain them here, where the peer group's reads pair. The last - # step's wait stays: it still covers the final A1, issued below. + # Drain the last tile's loads: its B0/A0 are read before the final wait. wait_barrier(0) c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) @@ -565,20 +539,6 @@ def kernel_gemm( rocdl.s_setprio(0) rocdl.s_barrier() - # TODO: the half-wave stagger the prologue opened is never closed. The - # matching `if wave_m == 0: rocdl.s_barrier()` belongs here, before the - # epilogue, so the two halves are level again -- gcnasm closes it at - # the same point (opus_gemm_a2a_lsa, quad-subtile template) and so does - # mori's fused copy of this pipeline. - # - # Nothing follows store_c here, so the imbalance is not observable - # today and this is left as a note rather than a change. It does not - # stay unobservable: with the counts unbalanced, a `wait_barrier(0)` - # added after store_c does not mean "every wave's C tile has retired", - # it rendezvouses waves 0-3 with waves 4-7 still one barrier back and - # before their stores. Anything appended to this epilogue is silently - # wrong until the pair is closed. - # Scale and store back to gmem wave_n_offset = wave_n * (N_TILES_B * 16) wave_m_offset = wave_m * (N_TILES_A * 16) diff --git a/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py b/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py index 3bd4ab990a0..a950f284991 100644 --- a/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py +++ b/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py @@ -3,15 +3,8 @@ """Correctness of the 8-wave a8w8 bpreshuffle GEMM. -The peeled tail pre-reads the last K step's B0 and A0 before the wait that -covers their producers, so it can read LDS the producing loads have not -finished writing. It shows up as a handful of wrong rows rather than a NaN -- -small enough to pass for quantisation noise. Whether it happens at all depends -on how the waves interleave, so a shape can be clean on one launch and wrong on -the next; each shape is therefore run several times. - -``K=256`` is the shape that pins the fix to the tail: ``K_ITERS == 2`` makes the -main loop run zero times, so no amount of tightening *there* can cover it. +A missing LDS sync shows up as a few wrong rows on some launches only, so each +shape runs several times. K=256 has no main-loop iterations. """ import argparse @@ -26,9 +19,7 @@ SUPPORTED_GFX = ["gfx950"] -# (M, N, K, block_m, block_n). Most are drawn from the tuned configs that -# select this kernel; K=256 and K=512 are not tuned shapes but are legal, and -# the first is the one the main loop cannot reach. +# (M, N, K, block_m, block_n); mostly from the tuned configs. SHAPES = [ (4096, 7168, 256, 256, 256), (4096, 7168, 512, 256, 256), @@ -40,9 +31,7 @@ (32768, 7168, 1792, 256, 256), ] REPS = 6 -# A dropped load leaves whole rows wrong by far more than this, so counting -# offending elements separates cleanly: zero when correct, thousands when not. -# An L2 norm would not -- the bad rows are a small fraction of a large output. +# Count elements off by >5%; an L2 norm hides a few wrong rows. REL_TOL = 0.05 @@ -54,13 +43,10 @@ def run_shape(m, n, k, block_m, block_n, reps, verbose=False): b_scale = torch.rand(n, 1, generator=gen) + 0.5 b_shuffled = shuffle_weight(b, (16, 16)) - # rounded to the kernel's output precision, so the comparison sees the - # kernel's error rather than bf16's + # Rounded to bf16 like the kernel output. ref = ((a.float() * a_scale) @ (b.float() * b_scale).T).to(torch.bfloat16).float() - # Floor the per-element tolerance well below a typical output but above - # zero: elements that cancel to near nothing disagree with torch on - # summation order alone, and that is not what this test is looking for. + # Tolerance floor, so near-zero outputs are not judged on summation order. floor = 0.01 * ref.square().mean().sqrt() worst = 0 From e8a6929d0b5683b5910ba7e5d08e971e1f90c186 Mon Sep 17 00:00:00 2001 From: xiangch Date: Mon, 19 Oct 2026 03:02:39 +0000 Subject: [PATCH 6/9] [FlyDSL] [gfx950] Drop the dedicated 8-wave GEMM test 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 --- op_tests/flydsl_tests/test_gemm_a8w8_8wave.py | 89 ------------------- 1 file changed, 89 deletions(-) delete mode 100644 op_tests/flydsl_tests/test_gemm_a8w8_8wave.py diff --git a/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py b/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py deleted file mode 100644 index a950f284991..00000000000 --- a/op_tests/flydsl_tests/test_gemm_a8w8_8wave.py +++ /dev/null @@ -1,89 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Correctness of the 8-wave a8w8 bpreshuffle GEMM. - -A missing LDS sync shows up as a few wrong rows on some launches only, so each -shape runs several times. K=256 has no main-loop iterations. -""" - -import argparse - -import torch - -from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.flydsl.gemm_a8w8_bpreshuffle_8wave import flydsl_8wave_gemm_a8 -from aiter.ops.shuffle import shuffle_weight - -torch.set_default_device("cuda") - -SUPPORTED_GFX = ["gfx950"] - -# (M, N, K, block_m, block_n); mostly from the tuned configs. -SHAPES = [ - (4096, 7168, 256, 256, 256), - (4096, 7168, 512, 256, 256), - (1472, 3584, 7168, 128, 256), - (3392, 2176, 7168, 128, 256), - (3456, 7168, 4224, 128, 256), - (1920, 7168, 1792, 256, 256), - (2624, 8448, 7168, 256, 256), - (32768, 7168, 1792, 256, 256), -] -REPS = 6 -# Count elements off by >5%; an L2 norm hides a few wrong rows. -REL_TOL = 0.05 - - -def run_shape(m, n, k, block_m, block_n, reps, verbose=False): - gen = torch.Generator(device="cuda").manual_seed(0) - a = (torch.randn(m, k, generator=gen) / 8).to(torch.float8_e4m3fn) - b = (torch.randn(n, k, generator=gen) / 8).to(torch.float8_e4m3fn) - a_scale = torch.rand(m, 1, generator=gen) + 0.5 - b_scale = torch.rand(n, 1, generator=gen) + 0.5 - b_shuffled = shuffle_weight(b, (16, 16)) - - # Rounded to bf16 like the kernel output. - ref = ((a.float() * a_scale) @ (b.float() * b_scale).T).to(torch.bfloat16).float() - - # Tolerance floor, so near-zero outputs are not judged on summation order. - floor = 0.01 * ref.square().mean().sqrt() - - worst = 0 - for _ in range(reps): - out = torch.empty(m, n, dtype=torch.bfloat16) - flydsl_8wave_gemm_a8(a, b_shuffled, a_scale, b_scale, out, block_m, block_n) - diff = (out.float() - ref).abs() - bad = int((diff > REL_TOL * ref.abs().clamp(min=floor)).sum()) - worst = max(worst, bad) - if verbose: - print(f" bad_elems {bad}") - return worst - - -def test_gemm_a8w8_8wave(): - if get_gfx() not in SUPPORTED_GFX: - return - for m, n, k, block_m, block_n in SHAPES: - worst = run_shape(m, n, k, block_m, block_n, REPS) - assert worst == 0, ( - f"M={m} N={n} K={k} tile={block_m}x{block_n}: " - f"{worst} elements off by more than {REL_TOL:.0%} in {REPS} runs" - ) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--reps", type=int, default=REPS) - args = parser.parse_args() - - failed = [] - for m, n, k, block_m, block_n in SHAPES: - print(f"M={m} N={n} K={k} tile={block_m}x{block_n}") - worst = run_shape(m, n, k, block_m, block_n, args.reps, verbose=True) - print(f" worst bad_elems {worst} {'ok' if worst == 0 else 'FAIL'}") - if worst: - failed.append((m, n, k, block_m, block_n)) - print(f"\n{len(SHAPES) - len(failed)}/{len(SHAPES)} shapes within tolerance") - for shape in failed: - print(f" failed: {shape}") From d2c1981a23c1e82fcd42017aa38b9b4b1df2e811 Mon Sep 17 00:00:00 2001 From: xiangch Date: Tue, 20 Oct 2026 12:17:08 +0000 Subject: [PATCH 7/9] [FlyDSL] [gfx950] Use rocdl.s_waitcnt instead of inline asm for the S0 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 --- aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py index bc7b9809ecb..4e76d96e9b2 100644 --- a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py +++ b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py @@ -74,17 +74,6 @@ def wait_barrier(count): ) -def read_done_barrier(): - """Finish this wave's LDS reads, then barrier.""" - _llvm.inline_asm( - res=None, - operands_=[], - asm_string="s_waitcnt lgkmcnt(0)\ns_barrier", - constraints="~{memory}", - has_side_effects=True, - ) - - class G2SLoader: def __init__(self, gl_src, gl_offsets, n_load_steps, lds_dtype, wave_id): self.g2lds_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128) @@ -458,8 +447,9 @@ def kernel_gemm( b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled) a0_frag = a_s2r.load(a_cur0) a_g2s.load(a_next1, A1_gl_offset + (k + 1) * BLOCK_K) - # The other wave group refills B0 right after this barrier. - read_done_barrier() + # The other wave group refills B0 right after this barrier: finish this wave's LDS reads first. + rocdl.s_waitcnt(lgkmcnt=0) + rocdl.s_barrier() c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) From 8842a0aeabe1177a5ae6ae7544a1b16d6903dba0 Mon Sep 17 00:00:00 2001 From: xiangch Date: Tue, 20 Oct 2026 12:31:13 +0000 Subject: [PATCH 8/9] [FlyDSL] [gfx950] Close the 8-wave GEMM's half-wave barrier stagger before 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 --- aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py index 4e76d96e9b2..e5c6cc74ac7 100644 --- a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py +++ b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py @@ -527,7 +527,13 @@ def kernel_gemm( c10_frag = mfma.call(a1_frag, b0_frag, c10_frag, set_prio=False) c11_frag = mfma.call(a1_frag, b1_frag, c11_frag, set_prio=False) rocdl.s_setprio(0) - rocdl.s_barrier() + # Balance the prologue's extra wave_m == 1 barrier so both groups' epilogues overlap. + # Not at BLOCK_N=512: already at 256 VGPRs, the branch raises spills from 4 to 11. + if const_expr(BLOCK_N > 256): + rocdl.s_barrier() + else: + if wave_m == 0: + rocdl.s_barrier() # Scale and store back to gmem wave_n_offset = wave_n * (N_TILES_B * 16) From 27d2909152a10633a4e00bf892701857eaaf5b7d Mon Sep 17 00:00:00 2001 From: xiangch Date: Tue, 20 Oct 2026 12:55:47 +0000 Subject: [PATCH 9/9] [FlyDSL] [gfx950] Balance the epilogue barrier on every tile; drop the 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 --- aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py index e5c6cc74ac7..ff53a9be915 100644 --- a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py +++ b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py @@ -431,7 +431,6 @@ def kernel_gemm( b_g2s.load(b_cur1, B1_gl_offset + 0 * B_K_STEP) a_g2s.load(a_cur1, A1_gl_offset + 0 * BLOCK_K) - # Intentional: waves 0-3 run one barrier ahead of waves 4-7. if wave_m == 1: rocdl.s_barrier() @@ -447,7 +446,6 @@ def kernel_gemm( b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled) a0_frag = a_s2r.load(a_cur0) a_g2s.load(a_next1, A1_gl_offset + (k + 1) * BLOCK_K) - # The other wave group refills B0 right after this barrier: finish this wave's LDS reads first. rocdl.s_waitcnt(lgkmcnt=0) rocdl.s_barrier() @@ -480,7 +478,6 @@ def kernel_gemm( k = K_ITERS - 2 b0_frag = b_s2r.load(b_cur0, preshuffled=b_preshuffled) a0_frag = a_s2r.load(a_cur0) - # Drain the last tile's loads: its B0/A0 are read before the final wait. wait_barrier(0) c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) @@ -527,13 +524,8 @@ def kernel_gemm( c10_frag = mfma.call(a1_frag, b0_frag, c10_frag, set_prio=False) c11_frag = mfma.call(a1_frag, b1_frag, c11_frag, set_prio=False) rocdl.s_setprio(0) - # Balance the prologue's extra wave_m == 1 barrier so both groups' epilogues overlap. - # Not at BLOCK_N=512: already at 256 VGPRs, the branch raises spills from 4 to 11. - if const_expr(BLOCK_N > 256): + if wave_m == 0: rocdl.s_barrier() - else: - if wave_m == 0: - rocdl.s_barrier() # Scale and store back to gmem wave_n_offset = wave_n * (N_TILES_B * 16)