From a67ff89bfdfb20f9a4ad3c097073bff014cbcb6c Mon Sep 17 00:00:00 2001 From: Felix Li Date: Sun, 5 Jul 2026 14:21:03 +0000 Subject: [PATCH 1/2] Refactor fp8 8wave MFMA calls through fx.gemm Wrap the priority/barrier scheduling in the 8wave MFMA helper so call sites stay compact while preserving the existing grouped tail schedule. Co-authored-by: Cursor --- kernels/fp8_gemm_8wave.py | 34 ++-------------------------------- kernels/fp8_gemm_utils.py | 35 ++++++++++++++++++++++++++++------- 2 files changed, 30 insertions(+), 39 deletions(-) diff --git a/kernels/fp8_gemm_8wave.py b/kernels/fp8_gemm_8wave.py index c92d9f908..dfc12377c 100644 --- a/kernels/fp8_gemm_8wave.py +++ b/kernels/fp8_gemm_8wave.py @@ -138,36 +138,24 @@ def kernel_gemm( a_g2s.load(a_next1, A1_gl_offset + (k + 1) * BLOCK_K) rocdl.s_barrier() - rocdl.s_setprio(1) c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) - rocdl.s_setprio(0) - rocdl.s_barrier() b1_frag = b_s2r.load(b_cur1, preshuffled=b_preshuffled) b_g2s.load(b_cur0, B0_gl_offset + (k + 2) * B_K_STEP) rocdl.s_barrier() - rocdl.s_setprio(1) c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) - rocdl.s_setprio(0) - rocdl.s_barrier() a1_frag = a_s2r.load(a_cur1) a_g2s.load(a_cur0, A0_gl_offset + (k + 2) * BLOCK_K) rocdl.s_barrier() - rocdl.s_setprio(1) c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) - rocdl.s_setprio(0) - rocdl.s_barrier() b_g2s.load(b_cur1, B1_gl_offset + (k + 2) * B_K_STEP) wait_barrier(2 * N_LDS_STEPS_A + N_LDS_STEPS_B) - rocdl.s_setprio(1) c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) - rocdl.s_setprio(0) - rocdl.s_barrier() # Swap cur and next a_cur0, a_next0 = a_next0, a_cur0 @@ -181,18 +169,12 @@ def kernel_gemm( a0_frag = a_s2r.load(a_cur0) rocdl.s_barrier() - rocdl.s_setprio(1) c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) - rocdl.s_setprio(0) - rocdl.s_barrier() b1_frag = b_s2r.load(b_cur1, preshuffled=b_preshuffled) rocdl.s_barrier() - rocdl.s_setprio(1) c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) - rocdl.s_setprio(0) - rocdl.s_barrier() a1_frag = a_s2r.load(a_cur1) # Main loop prefetches a_next1 one step behind; issue the final @@ -200,18 +182,12 @@ def kernel_gemm( a_g2s.load(a_next1, A1_gl_offset + (K_ITERS - 1) * BLOCK_K) rocdl.s_barrier() - rocdl.s_setprio(1) c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) - rocdl.s_setprio(0) - rocdl.s_barrier() b0_frag = b_s2r.load(b_next0, preshuffled=b_preshuffled) rocdl.s_barrier() - rocdl.s_setprio(1) c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) - rocdl.s_setprio(0) - rocdl.s_barrier() # Swap cur and next a_cur0, a_next0 = a_next0, a_cur0 a_cur1, a_next1 = a_next1, a_cur1 @@ -223,25 +199,19 @@ def kernel_gemm( a0_frag = a_s2r.load(a_cur0) wait_barrier(0) - rocdl.s_setprio(1) c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) - rocdl.s_setprio(0) - rocdl.s_barrier() b1_frag = b_s2r.load(b_cur1, preshuffled=b_preshuffled) rocdl.s_barrier() - rocdl.s_setprio(1) c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) - rocdl.s_setprio(0) - rocdl.s_barrier() a1_frag = a_s2r.load(a_cur1) rocdl.s_barrier() rocdl.s_setprio(1) - c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) - c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag, with_priority_barrier=False) + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag, with_priority_barrier=False) rocdl.s_setprio(0) rocdl.s_barrier() diff --git a/kernels/fp8_gemm_utils.py b/kernels/fp8_gemm_utils.py index 4b3462320..0973f2d9b 100644 --- a/kernels/fp8_gemm_utils.py +++ b/kernels/fp8_gemm_utils.py @@ -2,10 +2,9 @@ # Copyright (c) 2025 FlyDSL Project Contributors import flydsl.expr as fx -from flydsl._mlir.dialects import fly as fly_dialect from flydsl._mlir.dialects import llvm as _llvm from flydsl._mlir.dialects.fly_rocdl import TargetAddressSpace -from flydsl.expr import arith, const_expr, range_constexpr +from flydsl.expr import arith, const_expr, range_constexpr, rocdl from flydsl.expr.typing import Vector as Vec @@ -207,7 +206,6 @@ def wait_barrier(count): class Mfma16x16x128: def __init__(self, n_tiles_a, n_tiles_b): self.atom = fx.make_mma_atom(fx.rocdl.cdna4.MFMA_Scale(16, 16, 128, fx.Float8E4M3FN)) - self.accum_type = Vec.make_type(4, fx.Float32) self.zero_value = Vec.filled(4, 0.0, fx.Float32) self.n_tiles_a = n_tiles_a self.n_tiles_b = n_tiles_b @@ -215,18 +213,41 @@ def __init__(self, n_tiles_a, n_tiles_b): def idx(self, i, j): return i * self.n_tiles_b + j + def _make_operand_frag(self, value): + frag = fx.make_rmem_tensor(8, fx.Int32) + frag.store(Vec(value)) + return frag + + def _make_accum_frag(self, value): + frag = fx.make_rmem_tensor(4, fx.Float32) + frag.store(Vec(value)) + return frag + def _do_mma(self, a, b, c): - return fly_dialect.mma_atom_call_ssa([self.accum_type], self.atom, a, b, c) + a_frag = self._make_operand_frag(a) + b_frag = self._make_operand_frag(b) + c_frag = self._make_accum_frag(c) + fx.gemm(self.atom, c_frag, a_frag, b_frag, c_frag) + return c_frag.load().ir_value() - def call(self, a, b, c): + def call(self, a, b, c, *, with_priority_barrier=True): assert len(a) == self.n_tiles_a assert len(b) == self.n_tiles_b assert len(c) == self.n_tiles_a * self.n_tiles_b + a_frags = [self._make_operand_frag(a[idx]) for idx in range_constexpr(self.n_tiles_a)] + b_frags = [self._make_operand_frag(b[idx]) for idx in range_constexpr(self.n_tiles_b)] + c_frags = [self._make_accum_frag(c[idx]) for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b)] + if const_expr(with_priority_barrier): + rocdl.s_setprio(1) for i in range_constexpr(self.n_tiles_a): for j in range_constexpr(self.n_tiles_b): - c[self.idx(i, j)] = self._do_mma(a[i], b[j], c[self.idx(i, j)]) - return c + cf = c_frags[self.idx(i, j)] + fx.gemm(self.atom, cf, a_frags[i], b_frags[j], cf) + if const_expr(with_priority_barrier): + rocdl.s_setprio(0) + rocdl.s_barrier() + return [c_frags[idx].load().ir_value() for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b)] def call_one(self, a, b, c, i, j): assert i < self.n_tiles_a and j < self.n_tiles_b From 613ea6b6007de7e8871d7044fc43f231b865534b Mon Sep 17 00:00:00 2001 From: Felix Li Date: Mon, 6 Jul 2026 02:50:12 +0000 Subject: [PATCH 2/2] Rename fp8 8wave priority flag Use the shorter set_prio keyword for the MFMA helper's optional priority/barrier wrapper. Co-authored-by: Cursor --- kernels/fp8_gemm_8wave.py | 4 ++-- kernels/fp8_gemm_utils.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/kernels/fp8_gemm_8wave.py b/kernels/fp8_gemm_8wave.py index dfc12377c..de4ba7402 100644 --- a/kernels/fp8_gemm_8wave.py +++ b/kernels/fp8_gemm_8wave.py @@ -210,8 +210,8 @@ def kernel_gemm( rocdl.s_barrier() rocdl.s_setprio(1) - c10_frag = mfma.call(a1_frag, b0_frag, c10_frag, with_priority_barrier=False) - c11_frag = mfma.call(a1_frag, b1_frag, c11_frag, with_priority_barrier=False) + 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() diff --git a/kernels/fp8_gemm_utils.py b/kernels/fp8_gemm_utils.py index 0973f2d9b..c3eadf3da 100644 --- a/kernels/fp8_gemm_utils.py +++ b/kernels/fp8_gemm_utils.py @@ -230,7 +230,7 @@ def _do_mma(self, a, b, c): fx.gemm(self.atom, c_frag, a_frag, b_frag, c_frag) return c_frag.load().ir_value() - def call(self, a, b, c, *, with_priority_barrier=True): + def call(self, a, b, c, *, set_prio=True): assert len(a) == self.n_tiles_a assert len(b) == self.n_tiles_b assert len(c) == self.n_tiles_a * self.n_tiles_b @@ -238,13 +238,13 @@ def call(self, a, b, c, *, with_priority_barrier=True): a_frags = [self._make_operand_frag(a[idx]) for idx in range_constexpr(self.n_tiles_a)] b_frags = [self._make_operand_frag(b[idx]) for idx in range_constexpr(self.n_tiles_b)] c_frags = [self._make_accum_frag(c[idx]) for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b)] - if const_expr(with_priority_barrier): + if const_expr(set_prio): rocdl.s_setprio(1) for i in range_constexpr(self.n_tiles_a): for j in range_constexpr(self.n_tiles_b): cf = c_frags[self.idx(i, j)] fx.gemm(self.atom, cf, a_frags[i], b_frags[j], cf) - if const_expr(with_priority_barrier): + if const_expr(set_prio): rocdl.s_setprio(0) rocdl.s_barrier() return [c_frags[idx].load().ir_value() for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b)]