From 5cb9791562fa1ad68fba0d3538eca3a785a63841 Mon Sep 17 00:00:00 2001 From: Yu-Yun Chang Date: Wed, 14 Oct 2026 13:31:01 +0000 Subject: [PATCH 1/3] fix: drain last K tile before cross-wave LDS reads in FlyDSL 8-wave a8w8 GEMM - The peeled K_ITERS - 2 step read the K_ITERS - 1 b0/a0/b1 LDS tiles with no vmcnt wait + barrier pairing. - Waves 4..7 lag one barrier, so waves 0..3 could read rows/cols waves 4..7 had not finished DMA-ing into LDS. - Symptom: intermittent 16-64 row (or column) bands per tile off by one K tile, max rel err ~0.15, any M. - Replace the barrier after the final a_next1 load with wait_barrier(N_LDS_STEPS_A); perf is unchanged. - Add op_tests/test_flydsl_gemm_a8w8_8wave.py, which checks every one of --repeat launches per shape/kernel. --- aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py | 7 +- op_tests/test_flydsl_gemm_a8w8_8wave.py | 165 ++++++++++++++++++++ 2 files changed, 171 insertions(+), 1 deletion(-) create mode 100644 op_tests/test_flydsl_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..b47b3bb37bf 100644 --- a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py +++ b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py @@ -490,7 +490,12 @@ def kernel_gemm( # Main loop prefetches a_next1 one step behind; issue the final # K_ITERS - 1 tile here, otherwise c10 / c11 read stale A1 data. a_g2s.load(a_next1, A1_gl_offset + (K_ITERS - 1) * BLOCK_K) - rocdl.s_barrier() + # Every wave loads a slice of each LDS tile, and waves 4..7 run one + # barrier behind waves 0..3. Drain the K_ITERS - 1 b0/a0/b1 tiles here + # (only the a_next1 load above stays in flight), two barriers before + # waves 0..3 read them; otherwise they can read LDS that waves 4..7 + # have not finished filling. + wait_barrier(N_LDS_STEPS_A) c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) diff --git a/op_tests/test_flydsl_gemm_a8w8_8wave.py b/op_tests/test_flydsl_gemm_a8w8_8wave.py new file mode 100644 index 00000000000..5cc7ec91eba --- /dev/null +++ b/op_tests/test_flydsl_gemm_a8w8_8wave.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Correctness + perf for the FlyDSL 8-wave a8w8 ptpc bpreshuffle GEMM (gfx950). + +The model path (``gemm_a8w8_bpreshuffle`` -> tuned ``flydsl_bpreshuffle_8w_*`` row): + + xq : [M, K] fp8, per-token fp32 scale [M, 1] + wq : [N, K] fp8 in shuffle_weight (16, 16) layout, per-channel scale [N, 1] + out : [M, N] bf16, preallocated + +A cross-wave LDS race on the last K tile used to corrupt 16-64 row (or column) +bands of a tile on some calls only, so every kernel is launched ``--repeat`` +times on the same inputs and each result is checked. + +Run: + python op_tests/test_flydsl_gemm_a8w8_8wave.py + python op_tests/test_flydsl_gemm_a8w8_8wave.py -s 12288,6144,7168 -r 50 +""" + +import argparse +import itertools + +import pandas as pd +import torch + +import aiter +from aiter import dtypes +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.gemm_op_a8w8 import gemm_a8w8_bpreshuffle_flydsl +from aiter.ops.shuffle import shuffle_weight +from aiter.test_common import ( + benchmark, + checkAllclose, + run_perftest, +) + +torch.set_default_device("cuda") + +# The 8-wave kernel uses the CDNA4 MFMA_Scale atom: gfx950 only. +SUPPORTED_GFX = ["gfx950"] +SEED = 0 + +_KERNELS = [ + "flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_2x4", + "flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_2x0", + "flydsl_bpreshuffle_8w_128x256x128_F8_F8_B16_1x8", +] + +# Tuned kimik3 shapes; before the fix (32768, 6400, 7168) failed on ~50% of calls. +_MNK = [ + (32768, 6400, 7168), + (12288, 6144, 7168), + (1000, 7168, 3584), +] + + +def run_torch(xq, wq, x_scale, w_scale, dtype=dtypes.bf16): + # Reference only: fp32 dequant matmul, cast back. Not timed, not in the table. + out = (xq.to(dtypes.fp32) @ wq.to(dtypes.fp32).t()) * x_scale * w_scale.t() + return out.to(dtype) + + +@benchmark() +def test_gemm_8wave(m, n, k, kernel, repeat): + torch.manual_seed(SEED) + x = torch.randn((m, k), dtype=dtypes.bf16) + w = torch.randn((n, k), dtype=dtypes.bf16) + xq, x_scale = aiter.pertoken_quant(x, quant_dtype=dtypes.fp8) + wq, w_scale = aiter.pertoken_quant(w, quant_dtype=dtypes.fp8) + wq_shuffled = shuffle_weight(wq, layout=(16, 16)) + out = torch.empty((m, n), dtype=dtypes.bf16) + ref = run_torch(xq, wq, x_scale, w_scale).to(dtypes.fp32) + + def _run(): + return gemm_a8w8_bpreshuffle_flydsl( + xq, wq_shuffled, x_scale, w_scale, out, {"kernelName": kernel} + ) + + # The race is intermittent: check every launch, not just the last one. + bad_calls = 0 + for _ in range(repeat): + out.fill_(float("nan")) + y = _run().to(dtypes.fp32) + if ( + checkAllclose( + ref, y, rtol=2e-2, atol=5e-2, tol_err_ratio=0.0, printLog=False + ) + > 0 + ): + bad_calls += 1 + + y, us = run_perftest(_run) + err = checkAllclose( + ref, y.to(dtypes.fp32), rtol=2e-2, atol=5e-2, msg=f"{kernel}: 8wave gemm" + ) + flops = 2 * m * n * k + nbytes = m * k + n * k + m * n * out.element_size() + return { + "gfx": get_gfx(), + "8wave us": us, + "8wave TFLOPS": flops / us / 1e6, + "8wave TB/s": nbytes / us / 1e6, + "8wave err": err, + "8wave bad calls": bad_calls, + } + + +def main(): + if get_gfx() not in SUPPORTED_GFX: + aiter.logger.warning("flydsl 8wave gemm unsupported on %s; skipping", get_gfx()) + return + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="config input of test", + ) + parser.add_argument( + "-s", + "--mnk", + type=dtypes.str2tuple, + nargs="*", + default=_MNK, + help="""Shape of mnk. + e.g.: -s 32768,6400,7168""", + ) + parser.add_argument( + "-k", + "--kernel", + type=str, + nargs="*", + default=_KERNELS, + help="""8-wave kernelName(s). + e.g.: -k flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_2x4""", + ) + parser.add_argument( + "-r", + "--repeat", + type=int, + nargs="*", + default=[20], + help="""Launches checked per shape and kernel. + e.g.: -r 50""", + ) + args = parser.parse_args() + + rows = [ + test_gemm_8wave(m, n, k, kernel, repeat) + for (m, n, k), kernel, repeat in itertools.product( + args.mnk, args.kernel, args.repeat + ) + ] + df = pd.DataFrame(rows) + aiter.logger.info( + "flydsl 8wave gemm summary (markdown):\n%s", df.to_markdown(index=False) + ) + failed = df[(df["8wave bad calls"] > 0) | (df["8wave err"] > 0)] + if not failed.empty: + raise AssertionError( + f"flydsl 8wave gemm mismatched the reference:\n{failed.to_string()}" + ) + + +if __name__ == "__main__": + main() From b39b16d0d608e362ece0180d77182400dc9d92cd Mon Sep 17 00:00:00 2001 From: Yu-Yun Chang Date: Wed, 14 Oct 2026 14:08:36 +0000 Subject: [PATCH 2/3] fix: drain last K tile before cross-wave LDS reads in FlyDSL 8-wave gather_kv_b_proj - gather_gemm_8wave.py has the 8-wave GEMM's tail: a plain barrier after the final a_next1 load. - Its SAFE_VMCNT loop wait can leave the K_ITERS - 1 b1 loads (and at BLOCK_M >= 256 the last a0 step) in flight. - Waves then read those b1 / a0 rows before any vmcnt wait of the wave that loaded them covers the LDS DMA. - Not reproduced on gfx950: 0 bad of 30-100 launches per config, also with an HBM hog and with fault injection. - Same fix as the GEMM: wait_barrier(N_LDS_STEPS_A) replaces that barrier; one s_waitcnt in the ISA, perf unchanged. - Extend op_tests/test_flydsl_gemm_a8w8_8wave.py with a repeated-launch gather_kv_b_proj check and its own table. --- aiter/ops/flydsl/kernels/gather_gemm_8wave.py | 8 +- op_tests/test_flydsl_gemm_a8w8_8wave.py | 149 ++++++++++++++++-- 2 files changed, 145 insertions(+), 12 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gather_gemm_8wave.py b/aiter/ops/flydsl/kernels/gather_gemm_8wave.py index e3c3b48cfc7..c1cc9d7a72d 100644 --- a/aiter/ops/flydsl/kernels/gather_gemm_8wave.py +++ b/aiter/ops/flydsl/kernels/gather_gemm_8wave.py @@ -871,7 +871,13 @@ def kernel_gather( # The main loop prefetches a_next1 one step behind; issue the final # K_ITERS - 1 tile here or c10 / c11 read stale A1 data. a1_g2s.load(a_next1, (K_ITERS - 1) * BLOCK_K) - rocdl.s_barrier() + # Every wave loads a slice of each LDS tile, and waves 4..7 run one + # barrier behind waves 0..3. The main loop's SAFE_VMCNT wait can leave + # the K_ITERS - 1 a0 / b1 loads in flight, and no later wait of the + # loading wave covers them before they are read. Drain them here (only + # the a_next1 load above stays in flight); otherwise waves can read LDS + # that waves 4..7 have not finished filling. + wait_barrier(N_LDS_STEPS_A) c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) diff --git a/op_tests/test_flydsl_gemm_a8w8_8wave.py b/op_tests/test_flydsl_gemm_a8w8_8wave.py index 5cc7ec91eba..ed191b5232e 100644 --- a/op_tests/test_flydsl_gemm_a8w8_8wave.py +++ b/op_tests/test_flydsl_gemm_a8w8_8wave.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Correctness + perf for the FlyDSL 8-wave a8w8 ptpc bpreshuffle GEMM (gfx950). +"""Correctness + perf for the FlyDSL 8-wave kernels (gfx950): the a8w8 ptpc +bpreshuffle GEMM and the fused gather + kv_b_proj built on its pipeline. The model path (``gemm_a8w8_bpreshuffle`` -> tuned ``flydsl_bpreshuffle_8w_*`` row): @@ -9,13 +10,22 @@ wq : [N, K] fp8 in shuffle_weight (16, 16) layout, per-channel scale [N, 1] out : [M, N] bf16, preallocated +The gather path (``gather_kv_b_proj_flydsl``, MLA prefix expansion): + + k_buffer : [num_blocks, 1, 576] fp8 cache, gathered by kv_indices [num_tokens] + wq : [heads * (nope + v), 512] fp8 in shuffle_weight (16, 16) layout + k_prefix : [num_tokens, heads, nope + 64] bf16, preallocated + v_prefix : [num_tokens, heads, v] bf16, preallocated + A cross-wave LDS race on the last K tile used to corrupt 16-64 row (or column) -bands of a tile on some calls only, so every kernel is launched ``--repeat`` -times on the same inputs and each result is checked. +bands of a GEMM tile on some calls only; the gather kernel shares that pipeline +tail. So every kernel is launched ``--repeat`` times on the same inputs into a +NaN-filled output and each result is checked. Run: python op_tests/test_flydsl_gemm_a8w8_8wave.py python op_tests/test_flydsl_gemm_a8w8_8wave.py -s 12288,6144,7168 -r 50 + python op_tests/test_flydsl_gemm_a8w8_8wave.py -g 16384,16,128,128 -r 100 """ import argparse @@ -27,6 +37,7 @@ import aiter from aiter import dtypes from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.flydsl.gather_kv_b_proj import gather_kv_b_proj_flydsl from aiter.ops.gemm_op_a8w8 import gemm_a8w8_bpreshuffle_flydsl from aiter.ops.shuffle import shuffle_weight from aiter.test_common import ( @@ -54,6 +65,18 @@ (1000, 7168, 3584), ] +KV_C_DIM = 512 +KV_PE_DIM = 64 + +# (num_tokens, heads, qk_nope_head_dim, v_head_dim): DeepSeek MLA 128+128 at 16 +# heads per rank, on both the 128-row (2048) and 256-row (16384) M tile, and the +# GLM-5.2 192+256 head, which splits each head over two workgroups. +_GATHER = [ + (2048, 16, 128, 128), + (16384, 16, 128, 128), + (8192, 16, 192, 256), +] + def run_torch(xq, wq, x_scale, w_scale, dtype=dtypes.bf16): # Reference only: fp32 dequant matmul, cast back. Not timed, not in the table. @@ -61,6 +84,15 @@ def run_torch(xq, wq, x_scale, w_scale, dtype=dtypes.bf16): return out.to(dtype) +def run_torch_gather(k_buffer, kv_indices, wq, w_scale, heads, nope, v_dim): + # Reference only: fp32 gather + dequant matmul, k_scale = 1. Not timed. + latent = k_buffer[kv_indices.long()].reshape(-1, KV_C_DIM + KV_PE_DIM).float() + kv_c, k_pe = latent.split([KV_C_DIM, KV_PE_DIM], dim=-1) + proj = (kv_c @ (wq.to(dtypes.fp32) * w_scale).t()).view(-1, heads, nope + v_dim) + rope = k_pe.unsqueeze(1).expand(-1, heads, -1) + return torch.cat([proj[..., :nope], rope], dim=-1), proj[..., nope:] + + @benchmark() def test_gemm_8wave(m, n, k, kernel, repeat): torch.manual_seed(SEED) @@ -106,6 +138,78 @@ def _run(): } +@benchmark() +def test_gather_kv_b_proj_8wave(num_tokens, heads, nope, v_dim, repeat): + torch.manual_seed(SEED) + num_blocks = 2 * num_tokens + k_buffer = torch.randn((num_blocks, 1, KV_C_DIM + KV_PE_DIM)).to(dtypes.fp8) + kv_indices = torch.randperm(num_blocks)[:num_tokens].to(torch.int32) + kv_indptr = torch.tensor([0, num_tokens], dtype=torch.int32) + k_scale = torch.ones(1, dtype=dtypes.fp32) + weight_n = heads * (nope + v_dim) + wq = torch.randn((weight_n, KV_C_DIM)).to(dtypes.fp8) + w_scale = torch.rand((weight_n, 1), dtype=dtypes.fp32) + 0.5 + wq_shuffled = shuffle_weight(wq, layout=(16, 16)) + k_prefix = torch.empty((num_tokens, heads, nope + KV_PE_DIM), dtype=dtypes.bf16) + v_prefix = torch.empty((num_tokens, heads, v_dim), dtype=dtypes.bf16) + k_ref, v_ref = run_torch_gather( + k_buffer, kv_indices, wq, w_scale, heads, nope, v_dim + ) + + def _run(): + gather_kv_b_proj_flydsl( + k_buffer, + k_scale, + kv_indptr, + kv_indices, + kv_indptr, + wq_shuffled, + w_scale, + k_prefix, + v_prefix, + ) + return k_prefix, v_prefix + + def _check(out, printLog=False): + return max( + checkAllclose( + ref, + y.to(dtypes.fp32), + rtol=2e-2, + atol=5e-2, + tol_err_ratio=0.0, + printLog=printLog, + msg=f"gather_kv_b_proj 8wave {name}", + ) + for name, ref, y in zip(("k_prefix", "v_prefix"), (k_ref, v_ref), out) + ) + + # The race is intermittent: check every launch, not just the last one. + bad_calls = 0 + for _ in range(repeat): + k_prefix.fill_(float("nan")) + v_prefix.fill_(float("nan")) + if _check(_run()) > 0: + bad_calls += 1 + + out, us = run_perftest(_run) + err = _check(out, printLog=True) + flops = 2 * num_tokens * weight_n * KV_C_DIM + nbytes = ( + num_tokens * (KV_C_DIM + KV_PE_DIM) + + weight_n * KV_C_DIM + + (k_prefix.numel() + v_prefix.numel()) * k_prefix.element_size() + ) + return { + "gfx": get_gfx(), + "gather us": us, + "gather TFLOPS": flops / us / 1e6, + "gather TB/s": nbytes / us / 1e6, + "gather err": err, + "gather bad calls": bad_calls, + } + + def main(): if get_gfx() not in SUPPORTED_GFX: aiter.logger.warning("flydsl 8wave gemm unsupported on %s; skipping", get_gfx()) @@ -133,6 +237,15 @@ def main(): help="""8-wave kernelName(s). e.g.: -k flydsl_bpreshuffle_8w_256x256x128_F8_F8_B16_2x4""", ) + parser.add_argument( + "-g", + "--gather", + type=dtypes.str2tuple, + nargs="*", + default=_GATHER, + help="""Gather shape: num_tokens,heads,qk_nope_head_dim,v_head_dim. + e.g.: -g 16384,16,128,128""", + ) parser.add_argument( "-r", "--repeat", @@ -144,20 +257,34 @@ def main(): ) args = parser.parse_args() - rows = [ + gemm_rows = [ test_gemm_8wave(m, n, k, kernel, repeat) for (m, n, k), kernel, repeat in itertools.product( args.mnk, args.kernel, args.repeat ) ] - df = pd.DataFrame(rows) - aiter.logger.info( - "flydsl 8wave gemm summary (markdown):\n%s", df.to_markdown(index=False) - ) - failed = df[(df["8wave bad calls"] > 0) | (df["8wave err"] > 0)] - if not failed.empty: + gather_rows = [ + test_gather_kv_b_proj_8wave(*shape, repeat) + for shape, repeat in itertools.product(args.gather, args.repeat) + ] + + failed = [] + for name, rows, prefix in ( + ("gemm", gemm_rows, "8wave"), + ("gather_kv_b_proj", gather_rows, "gather"), + ): + if not rows: + continue + df = pd.DataFrame(rows) + aiter.logger.info( + "flydsl 8wave %s summary (markdown):\n%s", name, df.to_markdown(index=False) + ) + bad = df[(df[f"{prefix} bad calls"] > 0) | (df[f"{prefix} err"] > 0)] + if not bad.empty: + failed.append(f"{name}:\n{bad.to_string()}") + if failed: raise AssertionError( - f"flydsl 8wave gemm mismatched the reference:\n{failed.to_string()}" + "flydsl 8wave kernels mismatched the reference:\n" + "\n".join(failed) ) From 68f6188655a37489c06aad34ed8e56cad4845a3d Mon Sep 17 00:00:00 2001 From: Yu-Yun Chang Date: Wed, 14 Oct 2026 14:51:58 +0000 Subject: [PATCH 3/3] docs: spell out the barrier pairing behind the FlyDSL 8-wave tail drain - mfma.call ends in s_barrier: waves 0..3 pass c10's only after waves 4..7 ran the tail vmcnt(N_LDS_STEPS_A). - So b_next0 is read after both wave groups drained the K_ITERS - 1 b0/a0/b1 loads; the drain needs no move. - gather_kv_b_proj: only the last a0 load (N_LDS_STEPS_A >= N_LDS_STEPS_B) is at risk; b1 is read after vmcnt(0). - Comment-only: final ISA is byte-identical for the 256x256/128x256 GEMM and two gather_kv_b_proj variants. --- aiter/ops/flydsl/kernels/gather_gemm_8wave.py | 11 ++++++----- aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py | 7 ++++--- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gather_gemm_8wave.py b/aiter/ops/flydsl/kernels/gather_gemm_8wave.py index c1cc9d7a72d..28d7076aafc 100644 --- a/aiter/ops/flydsl/kernels/gather_gemm_8wave.py +++ b/aiter/ops/flydsl/kernels/gather_gemm_8wave.py @@ -872,11 +872,12 @@ def kernel_gather( # K_ITERS - 1 tile here or c10 / c11 read stale A1 data. a1_g2s.load(a_next1, (K_ITERS - 1) * BLOCK_K) # Every wave loads a slice of each LDS tile, and waves 4..7 run one - # barrier behind waves 0..3. The main loop's SAFE_VMCNT wait can leave - # the K_ITERS - 1 a0 / b1 loads in flight, and no later wait of the - # loading wave covers them before they are read. Drain them here (only - # the a_next1 load above stays in flight); otherwise waves can read LDS - # that waves 4..7 have not finished filling. + # barrier behind waves 0..3 (mfma.call ends in a barrier too). When + # N_LDS_STEPS_A >= N_LDS_STEPS_B the main loop's SAFE_VMCNT wait can + # leave the last K_ITERS - 1 a0 load in flight, and waves 0..3 read + # a_cur0 once waves 4..7 reach the plain barrier after b_next0. Drain + # it here (only the a_next1 load above stays in flight); otherwise + # waves can read LDS that waves 4..7 have not finished filling. wait_barrier(N_LDS_STEPS_A) c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) diff --git a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py index b47b3bb37bf..16109853b71 100644 --- a/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py +++ b/aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py @@ -492,9 +492,10 @@ def kernel_gemm( a_g2s.load(a_next1, A1_gl_offset + (K_ITERS - 1) * BLOCK_K) # Every wave loads a slice of each LDS tile, and waves 4..7 run one # barrier behind waves 0..3. Drain the K_ITERS - 1 b0/a0/b1 tiles here - # (only the a_next1 load above stays in flight), two barriers before - # waves 0..3 read them; otherwise they can read LDS that waves 4..7 - # have not finished filling. + # (only the a_next1 load above stays in flight); otherwise waves 0..3 + # can read LDS that waves 4..7 have not finished filling. mfma.call + # ends in a barrier, so waves 0..3 pass c10's -- and then read + # b_next0 -- only once waves 4..7 have executed this wait. wait_barrier(N_LDS_STEPS_A) c10_frag = mfma.call(a1_frag, b0_frag, c10_frag)