From a37e9fbe302f947d6fb1781573e66e0e4c6e55dd Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Fri, 17 Jul 2026 18:25:16 +0000 Subject: [PATCH 1/6] Add optimized 4-wave MXFP8 GEMM kernel Signed-off-by: Aristotle Martin --- kernels/gemm/mxfp8_gemm_4wave.py | 1008 ++++++++++++++++++++++++ tests/kernels/test_mxfp8_gemm_4wave.py | 309 ++++++++ 2 files changed, 1317 insertions(+) create mode 100644 kernels/gemm/mxfp8_gemm_4wave.py create mode 100644 tests/kernels/test_mxfp8_gemm_4wave.py diff --git a/kernels/gemm/mxfp8_gemm_4wave.py b/kernels/gemm/mxfp8_gemm_4wave.py new file mode 100644 index 000000000..bb88b21c4 --- /dev/null +++ b/kernels/gemm/mxfp8_gemm_4wave.py @@ -0,0 +1,1008 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""4-wave MXFP8 GEMM for AMD CDNA4/gfx950. + +Algorithm derived from the HipKittens MXFP8_4wave kernel +(https://github.com/HazyResearch/HipKittens/blob/a288366e4245528f74540b3fe446637cf8345745/kernels/cdna4/gemm/mxfp8/MXFP8_4wave/4_wave.cu#L1). + +The kernel targets ``v_mfma_scale_f32_16x16x128_f8f6f4`` using MXFP8 E4M3 +inputs and E8M0 per-K32 scaling factors. Scale operands are prepacked on the +host into the format consumed by the scaled MFMA instruction and are loaded +through a dedicated scale pipeline separate from the A/B data path. + +A and B operands are staged through XOR-swizzled LDS buffers using a +HipKittens-style 4-wave schedule with overlapped global loads, LDS reads, +and MFMA execution. The implementation uses direct ROCDL scheduling controls +to preserve the intended ordering of memory and compute operations. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec +from kernels.gemm.fp8_gemm_utils import divmod, pack_i32x4_i32x8, swizzle_128 + + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _xcd_swizzle(num_pid_m, num_pid_n): + """Map the linear workgroup ID to a GEMM tile coordinate. + + Uses row-major ordering for small or uneven grids. Otherwise, partitions + workgroups evenly across XCDs and applies grouped-M ordering to improve + operand reuse and cache locality. + """ + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +LDS_VECTOR_BYTES = 16 + + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + + +def compile_mxfp8_gemm_4w(*, K: int, BLOCK_M: int = 256, BLOCK_N: int = 256, use_xcd_remap: bool = True): + """Build the specialized 4-wave kernel for compile-time ``K``. + + ``K`` must contain at least four K128 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_K = 128 + NUM_THREADS = 256 + WARP_SIZE = 64 + NUM_WAVES = NUM_THREADS // WARP_SIZE + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 1 + VEC_BYTES = LDS_VECTOR_BYTES + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + LOAD_PASSES_SCALES = 16 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K128 tiles; the two-page pipeline needs at least 4" + + LDS_SYM_A0 = "mxfp8_pp_smem_a0" + LDS_SYM_A1 = "mxfp8_pp_smem_a1" + LDS_SYM_B0 = "mxfp8_pp_smem_b0" + LDS_SYM_B1 = "mxfp8_pp_smem_b1" + # Model each ping-pong LDS page as a distinct LLVM alias scope. Loads from + # the page being consumed can then be scheduled independently of direct + # global-to-LDS refills targeting a different page. + LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' + SCOPE_IDS = ("a0", "a1", "b0", "b1") + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, B: fx.Tensor, C: fx.Tensor, As: fx.Tensor, Bs: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32 + ): + _LDS_PTR_TY = ir.Type.parse("!llvm.ptr<3>") + _I8_TY = T.i8 + _GEP_DYN = -(2**31) + + def _gep_lds(base_ptr, byte_offset_i32): + return llvm.getelementptr(_LDS_PTR_TY, base_ptr, [byte_offset_i32], [_GEP_DYN], _I8_TY, None) + + def _scope_attr(ids): + """Build an LLVM alias-scope array for the named LDS pages.""" + inner = ", ".join(f'#llvm.alias_scope' for sid in ids) + return ir.Attribute.parse(f"[{inner}]") + + # Tag each access with alias.scope(page) and noalias(other pages). + # These are compile-time promises, justified because A0/A1/B0/B1 are + # separate LDS globals; incorrect metadata here could permit invalid + # memory reordering. + _SCOPE = {sid: _scope_attr((sid,)) for sid in SCOPE_IDS} + _NOALIAS = {sid: _scope_attr(tuple(other for other in SCOPE_IDS if other != sid)) for sid in SCOPE_IDS} + + lds_a0_base_ptr = llvm.mlir_addressof(_LDS_PTR_TY, LDS_SYM_A0) + lds_a1_base_ptr = llvm.mlir_addressof(_LDS_PTR_TY, LDS_SYM_A1) + lds_b0_base_ptr = llvm.mlir_addressof(_LDS_PTR_TY, LDS_SYM_B0) + lds_b1_base_ptr = llvm.mlir_addressof(_LDS_PTR_TY, LDS_SYM_B1) + + def _make_lds_view(base_ptr, my_scope, other_scopes): + """Create a byte-addressed LDS view carrying LLVM alias metadata.""" + vec_t = Vec.make_type(4, fx.Int32) + + def vec_load_i32x4_byte_offset(byte_off_idx): + # Attach the page's alias.scope and the other pages' noalias + # scopes to every LDS read. LLVM can then keep reads from the + # current page independent of refills targeting another page. + byte_off_i32 = arith.index_cast(T.i32, byte_off_idx) + gep = _gep_lds(base_ptr, byte_off_i32) + return llvm.LoadOp( + vec_t, + gep, + alignment=4, + alias_scopes=my_scope, + noalias_scopes=other_scopes, + ).result + + view = type("AliasScopedF8LDSView", (), {})() + view.vec_load_i32x4_byte_offset = vec_load_i32x4_byte_offset + return view + + lds_a0 = _make_lds_view(lds_a0_base_ptr, _SCOPE["a0"], _NOALIAS["a0"]) + lds_a1 = _make_lds_view(lds_a1_base_ptr, _SCOPE["a1"], _NOALIAS["a1"]) + lds_b0 = _make_lds_view(lds_b0_base_ptr, _SCOPE["b0"], _NOALIAS["b0"]) + lds_b1 = _make_lds_view(lds_b1_base_ptr, _SCOPE["b1"], _NOALIAS["b1"]) + a_rsrc = buffer_ops.create_buffer_resource(A, max_size=True) + b_rsrc = buffer_ops.create_buffer_resource(B, max_size=True) + as_rsrc = buffer_ops.create_buffer_resource(As, max_size=True) + bs_rsrc = buffer_ops.create_buffer_resource(Bs, max_size=True) + tx = gpu.thread_id("x") + + num_blocks_m = c_m // BLOCK_M + num_blocks_n = c_n // BLOCK_N + + if const_expr(use_xcd_remap): + pid_m, pid_n = _xcd_swizzle(num_blocks_m, num_blocks_n) + else: + pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) + + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is expressed in MLIR index type. Convert + # once here and use these index-typed tile bases for every address. + bx_m_idx = fx.Index(bx_m) + by_n_idx = fx.Index(by_n) + + layout_wave_lane = fx.make_layout((NUM_WAVES, WARP_SIZE), (WARP_SIZE, 1)) + coord_wave_lane = fx.idx2crd(fx.Int32(tx), layout_wave_lane) + wave_id = fx.get(coord_wave_lane, 0) + lane = fx.get(coord_wave_lane, 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(2) # C is f16. + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) + llvm.InlineAsmOp( + None, + [], + "", + clobbers, + has_side_effects=True, + ) + + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) + + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _unwrap_mlir_value(x): + """Return the bare MLIR value for the known MFMA operand types. + + A/B fragments are ``flydsl.expr.typing.Vector`` wrappers and expose + their underlying MLIR value through ``ir_value()``. Scale operands + are ``flydsl.expr.utils.arith.ArithValue`` instances, which are + already accepted directly by ``llvm.InlineAsmOp``. + """ + if isinstance(x, Vec): + return x.ir_value() + return x + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def read_pinned_accumulator(acc_idx): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + # As/Bs are MFMA-ready packed scale words: [K128, row] uint32. + # Each loaded dword already contains the four 16-row/16-col MFMA scale + # bytes for this lane's 64-row A/B half. The MFMA instruction selects + # the byte via op_sel/op_sel_hi, so there is intentionally no hot-loop + # byte extraction and no 0x01010101 broadcast here. + c_m_idx = fx.Index(c_m) + c_n_idx = fx.Index(c_n) + + def hot_loop_scheduler_q_refill_2n(): + # 8 chunks: + # 1 VMEM/LDS-refill pass + # 2 MFMA + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_mfma(2) + + rocdl.sched_barrier(0) + + def load_a_scale_row(k128, row): + packed = buffer_ops.buffer_load( + as_rsrc, + k128 * c_m_idx + bx_m_idx + row, + vec_width=1, + dtype=T.i32, + ) + return packed + + def load_b_scale_row(k128, row): + packed = buffer_ops.buffer_load( + bs_rsrc, + k128 * c_n_idx + by_n_idx + row, + vec_width=1, + dtype=T.i32, + ) + return packed + + def load_a_scale_subtile(k128, sm): + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + a_row = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(lane) + a_scale = load_a_scale_row(k128, a_row) + return (a_scale, a_scale, a_scale, a_scale) + + def load_b_scale_subtile(k128, sn): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(lane) + b_scale = load_b_scale_row(k128, b_row) + return (b_scale, b_scale, b_scale, b_scale) + + def load_scale_tile(k128): + # Load all scale VGPRs needed by this wave for this K128 tile once. + # Return order: A-top, A-bottom, B-left, B-right. + return ( + load_a_scale_subtile(k128, 0), + load_a_scale_subtile(k128, 1), + load_b_scale_subtile(k128, 0), + load_b_scale_subtile(k128, 1), + ) + + def stage_a_pass(k_base, load_i, lds_a_base_ptr, lds_a_scope, lds_a_noalias): + linear = fx.Index(tx) + fx.Index(load_i * NUM_THREADS) + byte_base = linear * fx.Index(VEC_BYTES) + + row = byte_base // fx.Index(BLOCK_K) + col = byte_base % fx.Index(BLOCK_K) + + # Keep the LDS destination physically contiguous for + # raw_ptr_buffer_load_lds, but load the logical K column that maps + # to this physical slot under the XOR layout. + a_phys_col_bytes = col + _, a_log_col_bytes = swizzle_128(row, a_phys_col_bytes) + a_global_byte = (bx_m_idx + row) * fx.Index(K) + (k_base + a_log_col_bytes) + a_lds_byte_i32 = arith.index_cast(T.i32, byte_base) + a_lds_byte_uniform = rocdl.readfirstlane(T.i32, a_lds_byte_i32) + a_lds_ptr = _gep_lds(lds_a_base_ptr, a_lds_byte_uniform) + + rocdl.raw_ptr_buffer_load_lds( + a_rsrc, + a_lds_ptr, + fx.Int32(VEC_BYTES), + fx.Int32(a_global_byte), + fx.Int32(0), + fx.Int32(0), + fx.Int32(1), + alias_scopes=lds_a_scope, + noalias_scopes=lds_a_noalias, + ) + + def stage_b_pass(k_base, load_i, lds_b_base_ptr, lds_b_scope, lds_b_noalias): + linear = fx.Index(tx) + fx.Index(load_i * NUM_THREADS) + byte_base = linear * fx.Index(VEC_BYTES) + + row = byte_base // fx.Index(BLOCK_K) + col = byte_base % fx.Index(BLOCK_K) + + # Same physical-contiguous LDS write / logical-swizzled global + # column scheme as A. + b_phys_col_bytes = col + _, b_log_col_bytes = swizzle_128(row, b_phys_col_bytes) + b_global_byte = (by_n_idx + row) * fx.Index(K) + (k_base + b_log_col_bytes) + b_lds_byte_i32 = arith.index_cast(T.i32, byte_base) + b_lds_byte_uniform = rocdl.readfirstlane(T.i32, b_lds_byte_i32) + b_lds_ptr = _gep_lds(lds_b_base_ptr, b_lds_byte_uniform) + + rocdl.raw_ptr_buffer_load_lds( + b_rsrc, + b_lds_ptr, + fx.Int32(VEC_BYTES), + fx.Int32(b_global_byte), + fx.Int32(0), + fx.Int32(0), + fx.Int32(1), + alias_scopes=lds_b_scope, + noalias_scopes=lds_b_noalias, + ) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a_base_ptr, lds_a_scope, lds_a_noalias): + load_i = subtile * LOAD_PASSES_A_SUBTILE + pass_in_subtile + stage_a_pass(k_base, load_i, lds_a_base_ptr, lds_a_scope, lds_a_noalias) + + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b_base_ptr, lds_b_scope, lds_b_noalias): + load_i = subtile * LOAD_PASSES_B_SUBTILE + pass_in_subtile + stage_b_pass(k_base, load_i, lds_b_base_ptr, lds_b_scope, lds_b_noalias) + + def stage_a_subtile(k_base, subtile, lds_a_base_ptr, lds_a_scope, lds_a_noalias): + start_pass = subtile * LOAD_PASSES_A_SUBTILE + stop_pass = start_pass + LOAD_PASSES_A_SUBTILE + for load_i in range_constexpr(start_pass, stop_pass): + stage_a_pass(k_base, load_i, lds_a_base_ptr, lds_a_scope, lds_a_noalias) + + def stage_b_subtile(k_base, subtile, lds_b_base_ptr, lds_b_scope, lds_b_noalias): + start_pass = subtile * LOAD_PASSES_B_SUBTILE + stop_pass = start_pass + LOAD_PASSES_B_SUBTILE + for load_i in range_constexpr(start_pass, stop_pass): + stage_b_pass(k_base, load_i, lds_b_base_ptr, lds_b_scope, lds_b_noalias) + + def load_frag_at_byte_base(lds_view, row_byte_base): + # The two K fragments use lane-invariant physical LDS columns. + x0 = lds_view.vec_load_i32x4_byte_offset(row_byte_base + reg_lds_k_col0) + x1 = lds_view.vec_load_i32x4_byte_offset(row_byte_base + reg_lds_k_col1) + return pack_i32x4_i32x8(x0, x1) + + def load_a_frag(lds_a, local_row): + return load_frag_at_byte_base(lds_a, local_row * fx.Index(BLOCK_K)) + + def load_b_frag(lds_b, local_row): + # B is stored as [N, K], but after staging it uses the same + # XOR-swizzled LDS fragment addressing as A. + return load_frag_at_byte_base(lds_b, local_row * fx.Index(BLOCK_K)) + + def _acc_idx(subtile_id, mi, ni): + return subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + + def pinned_mfma(acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): + # Fixed physical accumulator bank, visible SSA A/B/scale operands. + # acc_idx maps directly to a[PIN_ACC_BASE + 4*acc_idx : +3]. + # The scale operands are MFMA-ready packed dwords. mi/ni choose + # which of the four bytes inside the A/B scale dword the MFMA uses. + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + _unwrap_mlir_value(a_frag), + _unwrap_mlir_value(b_frag), + _unwrap_mlir_value(a_scale), + _unwrap_mlir_value(b_scale), + ], + ( + f"v_mfma_scale_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$2, $3 " + f"op_sel:[{mi & 1},{ni & 1},0] " + f"op_sel_hi:[{mi >> 1},{ni >> 1},0]" + ), + (f"v,v,v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}"), + has_side_effects=True, + ) + + def mfma_4n(acc_base, a_frag, a_scale, b0, b1, b2, b3, bs0, bs1, bs2, bs3): + """Emit four N-direction scaled MFMAs into fixed physical AGPR accumulators.""" + mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE + pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, 0) + pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, 1) + pinned_mfma(acc_base + 2, a_frag, b2, a_scale, bs2, mi, 2) + pinned_mfma(acc_base + 3, a_frag, b3, a_scale, bs3, mi, 3) + + def mfma_2n(acc_base, a_frag, a_scale, b0, b1, bs0, bs1, ni_base): + mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE + pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, ni_base + 0) + pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, ni_base + 1) + + def store_output(): + c_n_idx = fx.Index(c_n) + for sm in range_constexpr(2): + # HK store mapping: each wave owns the same 64x64 warp-position + # inside all four 128x128 quadrants, not one contiguous 128x128 tile. + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + + for sn in range_constexpr(2): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + subtile_id = sm * 2 + sn + + subtile_m_base = subtile_m_idx * SUBTILE_M + subtile_n_base = subtile_n_idx * SUBTILE_N + + for mi in range_constexpr(MFMA_M_PER_SUBTILE): + row_base = subtile_m_base + fx.Index(mi * MFMA_M) + lane_div_16 * 4 + + for ni in range_constexpr(MFMA_N_PER_SUBTILE): + col = subtile_n_base + fx.Index(ni * MFMA_N) + lane_mod_16 + acc_idx = _acc_idx(subtile_id, mi, ni) + acc = read_pinned_accumulator(acc_idx) + + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * c_n_idx + col + val = Vec(acc)[ii].to(fx.Float16) + buffer_ops.buffer_store(val, c_rsrc, c_idx) + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + + # Every fragment row differs only by multiples of 16, so row % 16 is + # always lane_mod_16. Hoist the logical->physical XOR mapping once. + _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) + + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 + + reserve_pinned_accumulators() + zero_pinned_accumulators() + + def load_b_subtile_ni_regs(lds_b, scale_tile, sn, ni): + # Fine-grained B register load for one 16-row N-direction MFMA slice. + # Return one packed B fragment and its matching scale operand. + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_scales = scale_tile[2] if sn == 0 else scale_tile[3] + + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + b_ni = load_b_frag(lds_b, b_row_addr) + b_scale_ni = b_scales[ni] + return b_ni, b_scale_ni + + def load_b_subtile_regs(lds_b, scale_tile, sn): + b0, bs0 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 0) + b1, bs1 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 1) + b2, bs2 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 2) + b3, bs3 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 3) + return b0, b1, b2, b3, bs0, bs1, bs2, bs3 + + def load_a_subtile_mi_regs(lds_a, scale_tile, sm, mi): + # Fine-grained A register load for one 16-row M-direction MFMA slice. + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + a_scales = scale_tile[0] if sm == 0 else scale_tile[1] + a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 + a_mi = load_a_frag(lds_a, a_row_addr) + a_scale_mi = a_scales[mi] + return a_mi, a_scale_mi + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a_base, + refill_b_base, + refill_a_scope, + refill_a_noalias, + refill_b_scope, + refill_b_noalias, + a0_regs, + b0_regs, + cur_scales, + prev_refill_scales, + ): + # Scale invariant: + # cur_scales is HK MFMA-ready for K. + # prev_refill_scales is HK MFMA-ready for K+1. + # This iteration issues K+2 scale loads and returns them for the + # next steady iteration or final tail. + + # Wait only far enough for the current page; the next-page refill may remain in flight. + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + # Immediately issue MFMA-ready K+2 scale loads. + # They are returned for the next iteration without any in-kernel + # byte extraction or broadcast. + refill_scales = load_scale_tile(fx.Index(k128 + 2)) + next_scales_ready = prev_refill_scales + # B-left is carried as a complete register tile, so its LDS page can be refilled + # immediately. Only A-top mi=0 is carried to limit the A-fragment live range. + # Split the carried B tile into individual MFMA slices for refill interleaving. + b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + a00, as00 = a0_regs + # Read the remaining A-top slices before refilling this LDS page; afterwards the + # old contents may be overwritten. Keeping the reads together also preserves contiguous MFMA groups. + a01, as01 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 1) + a02, as02 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 2) + a03, as03 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) + b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) + b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) + b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Index((k128 + 2) * BLOCK_K) + + rocdl.sched_barrier(0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a_base, refill_a_scope, refill_a_noalias) + mfma_2n(_acc_idx(0, 0, 0), a00, as00, b00, b01, bs00, bs01, 0) + + stage_b_subtile_pass(k_refill, 0, 0, refill_b_base, refill_b_scope, refill_b_noalias) + mfma_2n(_acc_idx(0, 0, 2), a00, as00, b02, b03, bs02, bs03, 2) + + stage_a_subtile_pass(k_refill, 0, 1, refill_a_base, refill_a_scope, refill_a_noalias) + mfma_2n(_acc_idx(0, 1, 0), a01, as01, b00, b01, bs00, bs01, 0) + + stage_b_subtile_pass(k_refill, 0, 1, refill_b_base, refill_b_scope, refill_b_noalias) + mfma_2n(_acc_idx(0, 1, 2), a01, as01, b02, b03, bs02, bs03, 2) + + stage_a_subtile_pass(k_refill, 0, 2, refill_a_base, refill_a_scope, refill_a_noalias) + mfma_2n(_acc_idx(0, 2, 0), a02, as02, b00, b01, bs00, bs01, 0) + + stage_b_subtile_pass(k_refill, 0, 2, refill_b_base, refill_b_scope, refill_b_noalias) + mfma_2n(_acc_idx(0, 2, 2), a02, as02, b02, b03, bs02, bs03, 2) + + stage_a_subtile_pass(k_refill, 0, 3, refill_a_base, refill_a_scope, refill_a_noalias) + mfma_2n(_acc_idx(0, 3, 0), a03, as03, b00, b01, bs00, bs01, 0) + + stage_b_subtile_pass(k_refill, 0, 3, refill_b_base, refill_b_scope, refill_b_noalias) + mfma_2n(_acc_idx(0, 3, 2), a03, as03, b02, b03, bs02, bs03, 2) + + hot_loop_scheduler_q_refill_2n() + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) + a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) + a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) + a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + stage_b_subtile_pass(k_refill, 1, 0, refill_b_base, refill_b_scope, refill_b_noalias) + mfma_2n(_acc_idx(1, 0, 0), a00, as00, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a_base, refill_a_scope, refill_a_noalias) + mfma_2n(_acc_idx(1, 0, 2), a00, as00, b12, b13, bs12, bs13, 2) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b_base, refill_b_scope, refill_b_noalias) + mfma_2n(_acc_idx(1, 1, 0), a01, as01, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 1, refill_a_base, refill_a_scope, refill_a_noalias) + mfma_2n(_acc_idx(1, 1, 2), a01, as01, b12, b13, bs12, bs13, 2) + + stage_b_subtile_pass(k_refill, 1, 2, refill_b_base, refill_b_scope, refill_b_noalias) + mfma_2n(_acc_idx(1, 2, 0), a02, as02, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a_base, refill_a_scope, refill_a_noalias) + mfma_2n(_acc_idx(1, 2, 2), a02, as02, b12, b13, bs12, bs13, 2) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b_base, refill_b_scope, refill_b_noalias) + mfma_2n(_acc_idx(1, 3, 0), a03, as03, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 3, refill_a_base, refill_a_scope, refill_a_noalias) + mfma_2n(_acc_idx(1, 3, 2), a03, as03, b12, b13, bs12, bs13, 2) + hot_loop_scheduler_q_refill_2n() + + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE + LOAD_PASSES_SCALES, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a0_regs = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 0) + + mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_b0_regs = load_b_subtile_regs(next_b, next_scales_ready, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + return next_a0_regs, next_b0_regs, next_scales_ready, refill_scales + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs, cur_scales, next_scales): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + a00, as00 = a0_regs + a01, as01 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 1) + a02, as02 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 2) + a03, as03 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) + b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) + b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) + b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, as00, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 1, 0), a01, as01, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 2, 0), a02, as02, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 3, 0), a03, as03, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) + a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) + a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) + a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, as00, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 1, 0), a01, as01, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 2, 0), a02, as02, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 3, 0), a03, as03, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a0_regs = load_a_subtile_mi_regs(next_a, next_scales, 0, 0) + + mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_b0_regs = load_b_subtile_regs(next_b, next_scales, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): + _barrier(vmcnt=0, lgkmcnt=0) + + b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + a00, as00 = a0_regs + a01, as01 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 1) + a02, as02 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 2) + a03, as03 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 3) + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) + b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) + b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) + b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, as00, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 1, 0), a01, as01, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 2, 0), a02, as02, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 3, 0), a03, as03, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) + a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) + a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) + a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, as00, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 1, 0), a01, as01, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 2, 0), a02, as02, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 3, 0), a03, as03, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + # Prologue: stage K0/K1 data into ping-pong LDS pages. Scales are not staged in + # LDS: As/Bs are already MFMA-ready preshuffled packed uint32 [K128, row], + # and load_scale_tile returns the current wave's scale operands in VGPRs. + + # Load scales first, so that they become the oldest VMEM ops. + scales0 = load_scale_tile(fx.Index(0)) + scales1 = load_scale_tile(fx.Index(1)) + + stage_a_subtile(fx.Index(0), 0, lds_a0_base_ptr, _SCOPE["a0"], _NOALIAS["a0"]) + stage_b_subtile(fx.Index(0), 0, lds_b0_base_ptr, _SCOPE["b0"], _NOALIAS["b0"]) + stage_b_subtile(fx.Index(0), 1, lds_b0_base_ptr, _SCOPE["b0"], _NOALIAS["b0"]) + stage_a_subtile(fx.Index(0), 1, lds_a0_base_ptr, _SCOPE["a0"], _NOALIAS["a0"]) + + stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1_base_ptr, _SCOPE["a1"], _NOALIAS["a1"]) + stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1_base_ptr, _SCOPE["b1"], _NOALIAS["b1"]) + stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1_base_ptr, _SCOPE["b1"], _NOALIAS["b1"]) + stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1_base_ptr, _SCOPE["a1"], _NOALIAS["a1"]) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + # scales0 is already MFMA-ready; no byte extraction or broadcast is needed. + # Keep the hot loop consistent for k=0 and k>0: + # K0 is consumed directly. K1 MFMA-ready scales are carried as + # prev_refill_scales and become next_scales_ready at loop entry. + + # Carry only A-top mi=0 across iterations to limit VGPR pressure; + # b0_regs carries the complete B-left 64-row register tile. + a0_regs = load_a_subtile_mi_regs(lds_a0, scales0, 0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + # Preload B-left for K0 to registers. + b0_regs = load_b_subtile_regs(lds_b0, scales0, 0) + + # Main HK loop: exactly one logical K128 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + # Scale tiles follow the same K128 progression but remain in VGPRs. + refill_scales = scales1 # K1 scales become the next ready scale tile at loop entry + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs, scales1, refill_scales = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0_base_ptr, + lds_b0_base_ptr, + _SCOPE["a0"], + _NOALIAS["a0"], + _SCOPE["b0"], + _NOALIAS["b0"], + a0_regs, + b0_regs, + scales0, + refill_scales, + ) + else: + a0_regs, b0_regs, scales0, refill_scales = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1_base_ptr, + lds_b1_base_ptr, + _SCOPE["a1"], + _NOALIAS["a1"], + _SCOPE["b1"], + _NOALIAS["b1"], + a0_regs, + b0_regs, + scales1, + refill_scales, + ) + + # Common two-page tail. After the steady loop, a0_regs/b0_regs are + # for the next page after the last consumed K128 tile. The final scale + # tile returned by the last refill belongs to the LDS page that was just + # refilled, so the tail page order depends on parity: + # even NUM_K_TILES: consume LDS0 then final LDS1 + # odd NUM_K_TILES: consume LDS1 then final LDS0 + if (NUM_K_TILES % 2) == 0: + scales1 = refill_scales + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + scales0, + scales1, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs, scales1) + else: + scales0 = refill_scales + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + scales1, + scales0, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs, scales0) + + rocdl.sched_barrier(0) + _barrier() + rocdl.sched_barrier(0) + + store_output() + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + A_scale: fx.Tensor, + B_scale: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + linkage = ir.Attribute.parse("#llvm.linkage") + for sym, size in ( + (LDS_SYM_A0, LDS_BYTES_A), + (LDS_SYM_A1, LDS_BYTES_A), + (LDS_SYM_B0, LDS_BYTES_B), + (LDS_SYM_B1, LDS_BYTES_B), + ): + llvm.GlobalOp( + global_type=ir.Type.parse(f"!llvm.array<{size} x i8>"), + sym_name=sym, + linkage=linkage, + addr_space=3, + alignment=1024, + ) + + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + B, + C, + A_scale, + B_scale, + c_m, + c_n, + value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, + ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), smem=0, stream=stream) + + return launch_gemm diff --git a/tests/kernels/test_mxfp8_gemm_4wave.py b/tests/kernels/test_mxfp8_gemm_4wave.py new file mode 100644 index 000000000..654f053e7 --- /dev/null +++ b/tests/kernels/test_mxfp8_gemm_4wave.py @@ -0,0 +1,309 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""MXFP8 4-wave GEMM correctness + perf harness. + +Kernel implementation: ``kernels/gemm/mxfp8_gemm_4wave.py``. +""" + +import os +import sys + +import pytest +import torch + +import flydsl.compiler as flyc + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +from flydsl.runtime.device import get_rocm_arch # noqa: E402 +from kernels.gemm.mxfp8_gemm_4wave import compile_mxfp8_gemm_4w # noqa: E402 +from tests.test_common import run_perftest # noqa: E402 + +FP8_DTYPE = torch.float8_e4m3fn +OUT_DTYPE = torch.float16 +MX_BLOCK_K = 32 +ARCH = str(get_rocm_arch()) + +if not torch.cuda.is_available(): + pytest.skip("CUDA/ROCm not available. Skipping GPU tests.", allow_module_level=True) + + +def _make_mxfp8_inputs(M: int, N: int, K: int, device: torch.device): + """Create MXFP8 inputs and raw per-K32 E8M0 scale bytes.""" + assert K % MX_BLOCK_K == 0 + + a = (torch.randn(M, K, device=device) * 2).clamp(-6, 6).to(FP8_DTYPE) + b = (torch.randn(N, K, device=device) * 2).clamp(-6, 6).to(FP8_DTYPE) + a_scale_raw = torch.randint( + 125, + 130, + (M, K // MX_BLOCK_K), + dtype=torch.uint8, + device=device, + ) + b_scale_raw = torch.randint( + 125, + 130, + (N, K // MX_BLOCK_K), + dtype=torch.uint8, + device=device, + ) + return a, b, a_scale_raw, b_scale_raw + + +def _dequantize_mxfp8(q: torch.Tensor, scales_u8: torch.Tensor) -> torch.Tensor: + scale = torch.exp2(scales_u8.float() - 127.0).repeat_interleave( + MX_BLOCK_K, + dim=1, + ) + return q.float() * scale + + +def _pack_scale_words(scales_u8: torch.Tensor) -> torch.Tensor: + """Pack raw ``[rows, K/32]`` E8M0 bytes as ``[K/128, rows]`` int32.""" + assert scales_u8.dtype == torch.uint8 + rows, k32_tiles = scales_u8.shape + assert k32_tiles % 4 == 0 + assert rows % 64 == 0 + + s32 = scales_u8.contiguous().view(rows, k32_tiles // 4, 4).to(torch.int32) + iteration_major = ( + (s32[:, :, 0] | (s32[:, :, 1] << 8) | (s32[:, :, 2] << 16) | (s32[:, :, 3] << 24)).transpose(0, 1).contiguous() + ) + + row = torch.arange(rows, device=scales_u8.device, dtype=torch.int64) + row_in_16 = row % 16 + k32_in_word = (row // 16) % 4 + tile_64 = row // 64 + + packed = torch.zeros_like(iteration_major) + for byte_index in range(4): + source_row = tile_64 * 64 + byte_index * 16 + row_in_16 + source_word = iteration_major[:, source_row] + source_byte = (source_word >> (k32_in_word * 8).view(1, rows)) & 0xFF + packed |= source_byte << (byte_index * 8) + + return packed.contiguous() + + +def _as_u8(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(torch.uint8) if "float8" in str(tensor.dtype) else tensor + + +def _verify_mxfp8_output(actual: torch.Tensor, expected: torch.Tensor) -> None: + """Apply the HipKittens MXFP8 absolute-error criterion.""" + actual_f32 = actual.float() + expected_f32 = expected.float() + difference = (actual_f32 - expected_f32).abs() + + tolerance = expected_f32.abs().max() * 1.0e-3 + assert torch.isfinite(actual_f32).all(), "MXFP8 kernel produced non-finite output" + assert torch.all(difference <= tolerance), ( + f"MXFP8 mismatch: max_abs={difference.max().item():.6g}, " + f"mean_abs={difference.mean().item():.6g}, " + f"atol={tolerance.item():.6g}" + ) + + +def _bench_mxfp8_gemm_4wave( + M: int, + N: int, + K: int, + *, + tile_m: int = 256, + tile_n: int = 256, + disable_xcd_remap: bool = False, + num_warmups: int = 2, + num_iters: int = 10, + static_weight_scale: bool = True, +): + """Run + verify one MXFP8 GEMM configuration. Returns TFLOPS.""" + if "gfx95" not in ARCH: + pytest.skip("MXFP8 GEMM requires CDNA4 (gfx95*)") + + if M % tile_m != 0 or N % tile_n != 0 or K % 128 != 0 or K < 512: + pytest.skip("MXFP8 4-wave kernel requires M/N divisible by the selected tile and K >= 512 divisible by 128") + + device = torch.device("cuda") + a, b, a_scale_raw, b_scale_raw = _make_mxfp8_inputs(M, N, K, device) + a_scale = _pack_scale_words(a_scale_raw) + b_scale = _pack_scale_words(b_scale_raw) + c_out = torch.zeros((M, N), dtype=OUT_DTYPE, device=device) + + a_dequantized = _dequantize_mxfp8(a, a_scale_raw) + b_dequantized = _dequantize_mxfp8(b, b_scale_raw) + c_ref = (a_dequantized @ b_dequantized.T).to(OUT_DTYPE) + + launch_fn = compile_mxfp8_gemm_4w( + K=K, + BLOCK_M=tile_m, + BLOCK_N=tile_n, + use_xcd_remap=not disable_xcd_remap, + ) + + print( + f"\n[mxfp8_gemm_4wave] M={M} N={N} K={K} " + f"BLOCK_M={tile_m} BLOCK_N={tile_n} " + f"xcd_remap={not disable_xcd_remap} " + f"static_weight_scale={static_weight_scale}" + ) + + def _args(c, a_input, b_input, a_scale_input, b_scale_input): + b_flat = _as_u8(b_input).contiguous().view(-1) + a_scale_flat = a_scale_input.contiguous().view(-1) + b_scale_flat = b_scale_input.contiguous().view(-1) + + if static_weight_scale: + b_flat = flyc.from_torch_tensor(b_flat) + a_scale_flat = flyc.from_torch_tensor(a_scale_flat) + b_scale_flat = flyc.from_torch_tensor(b_scale_flat) + + return ( + _as_u8(a_input).contiguous().view(-1), + b_flat, + c.contiguous().view(-1), + a_scale_flat, + b_scale_flat, + M, + N, + torch.cuda.current_stream(), + ) + + compiled = flyc.compile( + launch_fn, + *_args(c_out, a, b, a_scale, b_scale), + ) + + def _launch(c, a_input, b_input, a_scale_input, b_scale_input): + compiled(*_args(c, a_input, b_input, a_scale_input, b_scale_input)) + + num_iters = max(2, int(num_iters)) + _, microseconds = run_perftest( + _launch, + c_out, + a, + b, + a_scale, + b_scale, + num_iters=num_iters, + num_warmup=num_warmups, + ) + torch.cuda.synchronize() + + _verify_mxfp8_output(c_out, c_ref) + + flops = 2 * M * N * K + bytes_moved = ( + M * K + + N * K + + M * N * torch.tensor([], dtype=OUT_DTYPE).element_size() + + M * (K // MX_BLOCK_K) + + N * (K // MX_BLOCK_K) + ) + tflops = flops / (microseconds / 1.0e6) / 1.0e12 + tbps = bytes_moved / 1.0e12 / (microseconds / 1.0e6) + print(f"[flyc] Throughput: {microseconds:.1f} us, {tflops:.2f} TFLOPS, BW: {tbps:.3f} TB/s") + return tflops + + +@pytest.mark.parametrize( + "M, N, K, tile_m, tile_n", + [ + pytest.param(512, 512, 512, 256, 256, id="512x512x512"), + pytest.param(5120, 5120, 8320, 256, 256, id="5120x5120x8320"), + pytest.param( + 8192, + 8192, + 8192, + 256, + 256, + marks=pytest.mark.large_shape, + id="8192x8192x8192", + ), + pytest.param( + 9728, + 8192, + 8320, + 256, + 256, + marks=pytest.mark.large_shape, + id="9728x8192x8320", + ), + pytest.param( + 16384, + 16384, + 16384, + 256, + 256, + marks=pytest.mark.large_shape, + id="16384x16384x16384", + ), + ], +) +def test_mxfp8_gemm_4wave(M, N, K, tile_m, tile_n): + _bench_mxfp8_gemm_4wave( + M=M, + N=N, + K=K, + tile_m=tile_m, + tile_n=tile_n, + ) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="MXFP8 4-wave GEMM benchmark") + parser.add_argument("-M", type=int, default=8192) + parser.add_argument("-N", type=int, default=8192) + parser.add_argument("-K", type=int, default=8192) + parser.add_argument("--tile_m", type=int, default=256) + parser.add_argument("--tile_n", type=int, default=256) + parser.add_argument( + "--disable_xcd_remap", + action="store_true", + default=False, + ) + parser.add_argument( + "--num_iters", + type=int, + default=100, + help="Benchmark iterations.", + ) + parser.add_argument( + "--num_warmups", + type=int, + default=10, + help="Warmup iterations.", + ) + parser.add_argument( + "--dynamic_weight_scale", + action="store_true", + default=False, + help=("Use dynamic tensor arguments for weight and scales instead of static DLPack adaptors."), + ) + args = parser.parse_args() + + torch.set_default_device("cuda") + + try: + _bench_mxfp8_gemm_4wave( + M=args.M, + N=args.N, + K=args.K, + tile_m=args.tile_m, + tile_n=args.tile_n, + disable_xcd_remap=args.disable_xcd_remap, + num_warmups=args.num_warmups, + num_iters=args.num_iters, + static_weight_scale=not args.dynamic_weight_scale, + ) + except pytest.skip.Exception as error: + print(f"Skipped: {error}") From 2fc00523c508070501039fa5436e5b2f26b7a4f8 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Sat, 18 Jul 2026 00:33:15 +0000 Subject: [PATCH 2/6] Updates to MXFP8 4-wave gemm including WAR hazard fix and add test suite --- kernels/gemm/mxfp8_gemm_4wave.py | 567 ++++++++++++------------- tests/kernels/test_mxfp8_gemm_4wave.py | 163 ++++++- 2 files changed, 424 insertions(+), 306 deletions(-) diff --git a/kernels/gemm/mxfp8_gemm_4wave.py b/kernels/gemm/mxfp8_gemm_4wave.py index bb88b21c4..9de97cc8b 100644 --- a/kernels/gemm/mxfp8_gemm_4wave.py +++ b/kernels/gemm/mxfp8_gemm_4wave.py @@ -19,13 +19,19 @@ import flydsl.compiler as flyc import flydsl.expr as fx -from flydsl._mlir import ir from flydsl._mlir.dialects import llvm -from flydsl.compiler.kernel_function import CompilationContext from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec -from kernels.gemm.fp8_gemm_utils import divmod, pack_i32x4_i32x8, swizzle_128 +from kernels.gemm.fp8_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + divmod, + make_fp8_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) def _min(a, b): @@ -122,7 +128,6 @@ def compile_mxfp8_gemm_4w(*, K: int, BLOCK_M: int = 256, BLOCK_N: int = 256, use BLOCK_K = 128 NUM_THREADS = 256 WARP_SIZE = 64 - NUM_WAVES = NUM_THREADS // WARP_SIZE SUBTILE_M = 64 SUBTILE_N = 64 @@ -153,72 +158,38 @@ def compile_mxfp8_gemm_4w(*, K: int, BLOCK_M: int = 256, BLOCK_N: int = 256, use NUM_K_TILES = K // BLOCK_K assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K128 tiles; the two-page pipeline needs at least 4" - LDS_SYM_A0 = "mxfp8_pp_smem_a0" - LDS_SYM_A1 = "mxfp8_pp_smem_a1" - LDS_SYM_B0 = "mxfp8_pp_smem_b0" - LDS_SYM_B1 = "mxfp8_pp_smem_b1" - # Model each ping-pong LDS page as a distinct LLVM alias scope. Loads from - # the page being consumed can then be scheduled independently of direct - # global-to-LDS refills targeting a different page. - LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' - SCOPE_IDS = ("a0", "a1", "b0", "b1") + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LOAD_PASSES_HALF = LDS_ELEMS_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x128 page is two independent 128x128 half-pages. + # The hot loop refills one 16-byte pass of one half-page at a time. + a0_0: fx.Array[fx.Float8E4M3FN, LDS_ELEMS_HALF, 16] + a0_1: fx.Array[fx.Float8E4M3FN, LDS_ELEMS_HALF, 16] + a1_0: fx.Array[fx.Float8E4M3FN, LDS_ELEMS_HALF, 16] + a1_1: fx.Array[fx.Float8E4M3FN, LDS_ELEMS_HALF, 16] + b0_0: fx.Array[fx.Float8E4M3FN, LDS_ELEMS_HALF, 16] + b0_1: fx.Array[fx.Float8E4M3FN, LDS_ELEMS_HALF, 16] + b1_0: fx.Array[fx.Float8E4M3FN, LDS_ELEMS_HALF, 16] + b1_1: fx.Array[fx.Float8E4M3FN, LDS_ELEMS_HALF, 16] @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) def kernel_gemm( A: fx.Tensor, B: fx.Tensor, C: fx.Tensor, As: fx.Tensor, Bs: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32 ): - _LDS_PTR_TY = ir.Type.parse("!llvm.ptr<3>") - _I8_TY = T.i8 - _GEP_DYN = -(2**31) - - def _gep_lds(base_ptr, byte_offset_i32): - return llvm.getelementptr(_LDS_PTR_TY, base_ptr, [byte_offset_i32], [_GEP_DYN], _I8_TY, None) - - def _scope_attr(ids): - """Build an LLVM alias-scope array for the named LDS pages.""" - inner = ", ".join(f'#llvm.alias_scope' for sid in ids) - return ir.Attribute.parse(f"[{inner}]") - - # Tag each access with alias.scope(page) and noalias(other pages). - # These are compile-time promises, justified because A0/A1/B0/B1 are - # separate LDS globals; incorrect metadata here could permit invalid - # memory reordering. - _SCOPE = {sid: _scope_attr((sid,)) for sid in SCOPE_IDS} - _NOALIAS = {sid: _scope_attr(tuple(other for other in SCOPE_IDS if other != sid)) for sid in SCOPE_IDS} - - lds_a0_base_ptr = llvm.mlir_addressof(_LDS_PTR_TY, LDS_SYM_A0) - lds_a1_base_ptr = llvm.mlir_addressof(_LDS_PTR_TY, LDS_SYM_A1) - lds_b0_base_ptr = llvm.mlir_addressof(_LDS_PTR_TY, LDS_SYM_B0) - lds_b1_base_ptr = llvm.mlir_addressof(_LDS_PTR_TY, LDS_SYM_B1) - - def _make_lds_view(base_ptr, my_scope, other_scopes): - """Create a byte-addressed LDS view carrying LLVM alias metadata.""" - vec_t = Vec.make_type(4, fx.Int32) - - def vec_load_i32x4_byte_offset(byte_off_idx): - # Attach the page's alias.scope and the other pages' noalias - # scopes to every LDS read. LLVM can then keep reads from the - # current page independent of refills targeting another page. - byte_off_i32 = arith.index_cast(T.i32, byte_off_idx) - gep = _gep_lds(base_ptr, byte_off_i32) - return llvm.LoadOp( - vec_t, - gep, - alignment=4, - alias_scopes=my_scope, - noalias_scopes=other_scopes, - ).result - - view = type("AliasScopedF8LDSView", (), {})() - view.vec_load_i32x4_byte_offset = vec_load_i32x4_byte_offset - return view - - lds_a0 = _make_lds_view(lds_a0_base_ptr, _SCOPE["a0"], _NOALIAS["a0"]) - lds_a1 = _make_lds_view(lds_a1_base_ptr, _SCOPE["a1"], _NOALIAS["a1"]) - lds_b0 = _make_lds_view(lds_b0_base_ptr, _SCOPE["b0"], _NOALIAS["b0"]) - lds_b1 = _make_lds_view(lds_b1_base_ptr, _SCOPE["b1"], _NOALIAS["b1"]) - a_rsrc = buffer_ops.create_buffer_resource(A, max_size=True) - b_rsrc = buffer_ops.create_buffer_resource(B, max_size=True) + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + f8_ir_t = fx.Float8E4M3FN.ir_type + gA = make_fp8_buffer_tensor(A, f8_ir_t) + gB = make_fp8_buffer_tensor(B, f8_ir_t) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) as_rsrc = buffer_ops.create_buffer_resource(As, max_size=True) bs_rsrc = buffer_ops.create_buffer_resource(Bs, max_size=True) tx = gpu.thread_id("x") @@ -240,10 +211,21 @@ def vec_load_i32x4_byte_offset(byte_off_idx): bx_m_idx = fx.Index(bx_m) by_n_idx = fx.Index(by_n) - layout_wave_lane = fx.make_layout((NUM_WAVES, WARP_SIZE), (WARP_SIZE, 1)) - coord_wave_lane = fx.idx2crd(fx.Int32(tx), layout_wave_lane) - wave_id = fx.get(coord_wave_lane, 0) - lane = fx.get(coord_wave_lane, 1) + # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # these values with i32 constants, so Index-typed coordinates would make + # arith.addi receive mixed operand types. + tx_i32 = fx.Int32(tx) + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # The utility mapping is identical to the previous manual staging: + # each step contributes one contiguous 16-byte vector per thread, while + # the global K coordinate is XOR-unswizzled for the physical LDS slot. + gl_off_a = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) + gl_off_b = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, f8_ir_t, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, f8_ir_t, wave_id) + s2r = S2RLoader(fx.Int32(0), 1) layout_lane16 = fx.make_layout((4, 16), (16, 1)) coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) @@ -306,18 +288,6 @@ def _inline_asm_i32(asm_string, constraints, operands=None): ) return _one_i32_result(op) - def _unwrap_mlir_value(x): - """Return the bare MLIR value for the known MFMA operand types. - - A/B fragments are ``flydsl.expr.typing.Vector`` wrappers and expose - their underlying MLIR value through ``ir_value()``. Scale operands - are ``flydsl.expr.utils.arith.ArithValue`` instances, which are - already accepted directly by ``llvm.InlineAsmOp``. - """ - if isinstance(x, Vec): - return x.ir_value() - return x - def _one_i32_result(op): # Accept the result attribute names exposed by the supported MLIR Python bindings. return getattr(op, "result", getattr(op, "res", op.results[0])) @@ -339,15 +309,39 @@ def read_pinned_accumulator(acc_idx): c_n_idx = fx.Index(c_n) def hot_loop_scheduler_q_refill_2n(): - # 8 chunks: - # 1 VMEM/LDS-refill pass - # 2 MFMA + # Steady-state Q1 schedule: eight chunks of one K+2 VMEM/LDS + # refill pass followed by two MFMAs. for _ in range_constexpr(8): rocdl.sched_vmem(1) rocdl.sched_mfma(2) rocdl.sched_barrier(0) + def hot_loop_scheduler_q0_refill_a1_2n(): + # Steady-state Q0 schedule. Each chunk contains exactly: + # 1 K+2 VMEM/LDS refill pass + # 1 current-tile A-bottom K64 ds_read_b128 + # 2 current-tile Q0 MFMAs + # Repeated eight times, this distributes all eight A-bottom LDS reads + # across Q0 and maximizes their distance from reuse of that half-page. + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(2) + + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # Q2/Q3 carry-prefetch schedule used by both the steady loop and the + # penultimate tail tile. Each of eight chunks contains: + # 2 LDS reads for one complete next-tile A-top or B-left fragment + # 4 MFMAs using the current tile + for _ in range_constexpr(8): + rocdl.sched_dsrd(2) + rocdl.sched_mfma(4) + + rocdl.sched_barrier(0) + def load_a_scale_row(k128, row): packed = buffer_ops.buffer_load( as_rsrc, @@ -388,96 +382,50 @@ def load_scale_tile(k128): load_b_scale_subtile(k128, 1), ) - def stage_a_pass(k_base, load_i, lds_a_base_ptr, lds_a_scope, lds_a_noalias): - linear = fx.Index(tx) + fx.Index(load_i * NUM_THREADS) - byte_base = linear * fx.Index(VEC_BYTES) - - row = byte_base // fx.Index(BLOCK_K) - col = byte_base % fx.Index(BLOCK_K) - - # Keep the LDS destination physically contiguous for - # raw_ptr_buffer_load_lds, but load the logical K column that maps - # to this physical slot under the XOR layout. - a_phys_col_bytes = col - _, a_log_col_bytes = swizzle_128(row, a_phys_col_bytes) - a_global_byte = (bx_m_idx + row) * fx.Index(K) + (k_base + a_log_col_bytes) - a_lds_byte_i32 = arith.index_cast(T.i32, byte_base) - a_lds_byte_uniform = rocdl.readfirstlane(T.i32, a_lds_byte_i32) - a_lds_ptr = _gep_lds(lds_a_base_ptr, a_lds_byte_uniform) - - rocdl.raw_ptr_buffer_load_lds( - a_rsrc, - a_lds_ptr, - fx.Int32(VEC_BYTES), - fx.Int32(a_global_byte), - fx.Int32(0), - fx.Int32(0), - fx.Int32(1), - alias_scopes=lds_a_scope, - noalias_scopes=lds_a_noalias, + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one + # 128x128 half-page (16 KiB). Each half has its own LDS base. + global_base = ( + (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) + + k_base ) + a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) - def stage_b_pass(k_base, load_i, lds_b_base_ptr, lds_b_scope, lds_b_noalias): - linear = fx.Index(tx) + fx.Index(load_i * NUM_THREADS) - byte_base = linear * fx.Index(VEC_BYTES) - - row = byte_base // fx.Index(BLOCK_K) - col = byte_base % fx.Index(BLOCK_K) - - # Same physical-contiguous LDS write / logical-swizzled global - # column scheme as A. - b_phys_col_bytes = col - _, b_log_col_bytes = swizzle_128(row, b_phys_col_bytes) - b_global_byte = (by_n_idx + row) * fx.Index(K) + (k_base + b_log_col_bytes) - b_lds_byte_i32 = arith.index_cast(T.i32, byte_base) - b_lds_byte_uniform = rocdl.readfirstlane(T.i32, b_lds_byte_i32) - b_lds_ptr = _gep_lds(lds_b_base_ptr, b_lds_byte_uniform) - - rocdl.raw_ptr_buffer_load_lds( - b_rsrc, - b_lds_ptr, - fx.Int32(VEC_BYTES), - fx.Int32(b_global_byte), - fx.Int32(0), - fx.Int32(0), - fx.Int32(1), - alias_scopes=lds_b_scope, - noalias_scopes=lds_b_noalias, + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + global_base = ( + (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + + k_base ) + b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) - def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a_base_ptr, lds_a_scope, lds_a_noalias): - load_i = subtile * LOAD_PASSES_A_SUBTILE + pass_in_subtile - stage_a_pass(k_base, load_i, lds_a_base_ptr, lds_a_scope, lds_a_noalias) - - def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b_base_ptr, lds_b_scope, lds_b_noalias): - load_i = subtile * LOAD_PASSES_B_SUBTILE + pass_in_subtile - stage_b_pass(k_base, load_i, lds_b_base_ptr, lds_b_scope, lds_b_noalias) - - def stage_a_subtile(k_base, subtile, lds_a_base_ptr, lds_a_scope, lds_a_noalias): - start_pass = subtile * LOAD_PASSES_A_SUBTILE - stop_pass = start_pass + LOAD_PASSES_A_SUBTILE - for load_i in range_constexpr(start_pass, stop_pass): - stage_a_pass(k_base, load_i, lds_a_base_ptr, lds_a_scope, lds_a_noalias) - - def stage_b_subtile(k_base, subtile, lds_b_base_ptr, lds_b_scope, lds_b_noalias): - start_pass = subtile * LOAD_PASSES_B_SUBTILE - stop_pass = start_pass + LOAD_PASSES_B_SUBTILE - for load_i in range_constexpr(start_pass, stop_pass): - stage_b_pass(k_base, load_i, lds_b_base_ptr, lds_b_scope, lds_b_noalias) - - def load_frag_at_byte_base(lds_view, row_byte_base): - # The two K fragments use lane-invariant physical LDS columns. - x0 = lds_view.vec_load_i32x4_byte_offset(row_byte_base + reg_lds_k_col0) - x1 = lds_view.vec_load_i32x4_byte_offset(row_byte_base + reg_lds_k_col1) + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): return pack_i32x4_i32x8(x0, x1) - def load_a_frag(lds_a, local_row): - return load_frag_at_byte_base(lds_a, local_row * fx.Index(BLOCK_K)) + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) - def load_b_frag(lds_b, local_row): - # B is stored as [N, K], but after staging it uses the same - # XOR-swizzled LDS fragment addressing as A. - return load_frag_at_byte_base(lds_b, local_row * fx.Index(BLOCK_K)) + def load_b_frag(lds_b, local_row, half): + # B is [N, K]. Each 128-row half-page has a local row origin of 0. + half_row = local_row - fx.Index(half * (BLOCK_N // 2)) + return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K)) def _acc_idx(subtile_id, mi, ni): return subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni @@ -491,10 +439,10 @@ def pinned_mfma(acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): llvm.InlineAsmOp( None, [ - _unwrap_mlir_value(a_frag), - _unwrap_mlir_value(b_frag), - _unwrap_mlir_value(a_scale), - _unwrap_mlir_value(b_scale), + arith._to_raw(a_frag), + arith._to_raw(b_frag), + arith._to_raw(a_scale), + arith._to_raw(b_scale), ], ( f"v_mfma_scale_f32_16x16x128_f8f6f4 " @@ -578,7 +526,7 @@ def load_b_subtile_ni_regs(lds_b, scale_tile, sn, ni): b_scales = scale_tile[2] if sn == 0 else scale_tile[3] b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 - b_ni = load_b_frag(lds_b, b_row_addr) + b_ni = load_b_frag(lds_b, b_row_addr, sn) b_scale_ni = b_scales[ni] return b_ni, b_scale_ni @@ -589,27 +537,38 @@ def load_b_subtile_regs(lds_b, scale_tile, sn): b3, bs3 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 3) return b0, b1, b2, b3, bs0, bs1, bs2, bs3 + def load_a_subtile_mi_half(lds_a, sm, mi, half): + # One ds_read_b128 for one K64 half of one A MFMA slice. + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Index(BLOCK_K) + return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) + def load_a_subtile_mi_regs(lds_a, scale_tile, sm, mi): # Fine-grained A register load for one 16-row M-direction MFMA slice. - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) a_scales = scale_tile[0] if sm == 0 else scale_tile[1] - a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 - a_mi = load_a_frag(lds_a, a_row_addr) + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + a_mi = pack_frag_halves(x0, x1) a_scale_mi = a_scales[mi] return a_mi, a_scale_mi + def load_a_subtile_regs(lds_a, scale_tile, sm): + a0, as0 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 0) + a1, as1 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 1) + a2, as2 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 2) + a3, as3 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 3) + return a0, a1, a2, a3, as0, as1, as2, as3 + def hk_one_k_with_refill( k128, cur_a, cur_b, next_a, next_b, - refill_a_base, - refill_b_base, - refill_a_scope, - refill_a_noalias, - refill_b_scope, - refill_b_noalias, + refill_a, + refill_b, a0_regs, b0_regs, cur_scales, @@ -630,20 +589,10 @@ def hk_one_k_with_refill( # byte extraction or broadcast. refill_scales = load_scale_tile(fx.Index(k128 + 2)) next_scales_ready = prev_refill_scales - # B-left is carried as a complete register tile, so its LDS page can be refilled - # immediately. Only A-top mi=0 is carried to limit the A-fragment live range. - # Split the carried B tile into individual MFMA slices for refill interleaving. + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs - a00, as00 = a0_regs - # Read the remaining A-top slices before refilling this LDS page; afterwards the - # old contents may be overwritten. Keeping the reads together also preserves contiguous MFMA groups. - a01, as01 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 1) - a02, as02 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 2) - a03, as03 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 3) - - rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) - rocdl.sched_barrier(0) b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) @@ -653,68 +602,83 @@ def hk_one_k_with_refill( # Refill the current ping-pong page with K+2, alternating A and B passes. k_refill = fx.Index((k128 + 2) * BLOCK_K) + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Each complete A-bottom fragment is assembled + # from two independently scheduled K64 halves. rocdl.sched_barrier(0) - stage_a_subtile_pass(k_refill, 0, 0, refill_a_base, refill_a_scope, refill_a_noalias) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) mfma_2n(_acc_idx(0, 0, 0), a00, as00, b00, b01, bs00, bs01, 0) - stage_b_subtile_pass(k_refill, 0, 0, refill_b_base, refill_b_scope, refill_b_noalias) + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) mfma_2n(_acc_idx(0, 0, 2), a00, as00, b02, b03, bs02, bs03, 2) - stage_a_subtile_pass(k_refill, 0, 1, refill_a_base, refill_a_scope, refill_a_noalias) + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) mfma_2n(_acc_idx(0, 1, 0), a01, as01, b00, b01, bs00, bs01, 0) - stage_b_subtile_pass(k_refill, 0, 1, refill_b_base, refill_b_scope, refill_b_noalias) + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) mfma_2n(_acc_idx(0, 1, 2), a01, as01, b02, b03, bs02, bs03, 2) - stage_a_subtile_pass(k_refill, 0, 2, refill_a_base, refill_a_scope, refill_a_noalias) + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) mfma_2n(_acc_idx(0, 2, 0), a02, as02, b00, b01, bs00, bs01, 0) - stage_b_subtile_pass(k_refill, 0, 2, refill_b_base, refill_b_scope, refill_b_noalias) + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) mfma_2n(_acc_idx(0, 2, 2), a02, as02, b02, b03, bs02, bs03, 2) - stage_a_subtile_pass(k_refill, 0, 3, refill_a_base, refill_a_scope, refill_a_noalias) + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) mfma_2n(_acc_idx(0, 3, 0), a03, as03, b00, b01, bs00, bs01, 0) - stage_b_subtile_pass(k_refill, 0, 3, refill_b_base, refill_b_scope, refill_b_noalias) + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) mfma_2n(_acc_idx(0, 3, 2), a03, as03, b02, b03, bs02, bs03, 2) - hot_loop_scheduler_q_refill_2n() + hot_loop_scheduler_q0_refill_a1_2n() + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. rocdl.sched_barrier(0) _barrier(lgkmcnt=0) rocdl.sched_barrier(0) - a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) - a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) - a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) - a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) + as10 = cur_scales[1][0] + as11 = cur_scales[1][1] + as12 = cur_scales[1][2] + as13 = cur_scales[1][3] rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) - rocdl.sched_barrier(0) - - stage_b_subtile_pass(k_refill, 1, 0, refill_b_base, refill_b_scope, refill_b_noalias) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) mfma_2n(_acc_idx(1, 0, 0), a00, as00, b10, b11, bs10, bs11, 0) - stage_a_subtile_pass(k_refill, 1, 0, refill_a_base, refill_a_scope, refill_a_noalias) + stage_a_subtile_pass(k_refill, 1, 0, refill_a) mfma_2n(_acc_idx(1, 0, 2), a00, as00, b12, b13, bs12, bs13, 2) - stage_b_subtile_pass(k_refill, 1, 1, refill_b_base, refill_b_scope, refill_b_noalias) + stage_b_subtile_pass(k_refill, 1, 1, refill_b) mfma_2n(_acc_idx(1, 1, 0), a01, as01, b10, b11, bs10, bs11, 0) - stage_a_subtile_pass(k_refill, 1, 1, refill_a_base, refill_a_scope, refill_a_noalias) + stage_a_subtile_pass(k_refill, 1, 1, refill_a) mfma_2n(_acc_idx(1, 1, 2), a01, as01, b12, b13, bs12, bs13, 2) - stage_b_subtile_pass(k_refill, 1, 2, refill_b_base, refill_b_scope, refill_b_noalias) + stage_b_subtile_pass(k_refill, 1, 2, refill_b) mfma_2n(_acc_idx(1, 2, 0), a02, as02, b10, b11, bs10, bs11, 0) - stage_a_subtile_pass(k_refill, 1, 2, refill_a_base, refill_a_scope, refill_a_noalias) + stage_a_subtile_pass(k_refill, 1, 2, refill_a) mfma_2n(_acc_idx(1, 2, 2), a02, as02, b12, b13, bs12, bs13, 2) - stage_b_subtile_pass(k_refill, 1, 3, refill_b_base, refill_b_scope, refill_b_noalias) + stage_b_subtile_pass(k_refill, 1, 3, refill_b) mfma_2n(_acc_idx(1, 3, 0), a03, as03, b10, b11, bs10, bs11, 0) - stage_a_subtile_pass(k_refill, 1, 3, refill_a_base, refill_a_scope, refill_a_noalias) + stage_a_subtile_pass(k_refill, 1, 3, refill_a) mfma_2n(_acc_idx(1, 3, 2), a03, as03, b12, b13, bs12, bs13, 2) hot_loop_scheduler_q_refill_2n() @@ -724,33 +688,60 @@ def hk_one_k_with_refill( _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE + LOAD_PASSES_SCALES, lgkmcnt=0) rocdl.sched_barrier(0) - next_a0_regs = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 0) - + next_a00, next_as00 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 0) mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_a01, next_as01 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 1) mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_a02, next_as02 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 2) mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_a03, next_as03 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 3) mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - next_b0_regs = load_b_subtile_regs(next_b, next_scales_ready, 0) + next_b00, next_bs00 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 0) mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + next_b01, next_bs01 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 1) mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + next_b02, next_bs02 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 2) mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + next_b03, next_bs03 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 3) mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = ( + next_a00, + next_a01, + next_a02, + next_a03, + next_as00, + next_as01, + next_as02, + next_as03, + ) + next_b0_regs = ( + next_b00, + next_b01, + next_b02, + next_b03, + next_bs00, + next_bs01, + next_bs02, + next_bs03, + ) + return next_a0_regs, next_b0_regs, next_scales_ready, refill_scales def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs, cur_scales, next_scales): _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs - a00, as00 = a0_regs - a01, as01 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 1) - a02, as02 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 2) - a03, as03 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 3) - - rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) - rocdl.sched_barrier(0) b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) @@ -780,32 +771,60 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs, cur_ _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) rocdl.sched_barrier(0) - next_a0_regs = load_a_subtile_mi_regs(next_a, next_scales, 0, 0) - + next_a00, next_as00 = load_a_subtile_mi_regs(next_a, next_scales, 0, 0) mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_a01, next_as01 = load_a_subtile_mi_regs(next_a, next_scales, 0, 1) mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_a02, next_as02 = load_a_subtile_mi_regs(next_a, next_scales, 0, 2) mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_a03, next_as03 = load_a_subtile_mi_regs(next_a, next_scales, 0, 3) mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - next_b0_regs = load_b_subtile_regs(next_b, next_scales, 0) + next_b00, next_bs00 = load_b_subtile_ni_regs(next_b, next_scales, 0, 0) mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + next_b01, next_bs01 = load_b_subtile_ni_regs(next_b, next_scales, 0, 1) mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + next_b02, next_bs02 = load_b_subtile_ni_regs(next_b, next_scales, 0, 2) mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + next_b03, next_bs03 = load_b_subtile_ni_regs(next_b, next_scales, 0, 3) mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = ( + next_a00, + next_a01, + next_a02, + next_a03, + next_as00, + next_as01, + next_as02, + next_as03, + ) + next_b0_regs = ( + next_b00, + next_b01, + next_b02, + next_b03, + next_bs00, + next_bs01, + next_bs02, + next_bs03, + ) + return next_a0_regs, next_b0_regs def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): _barrier(vmcnt=0, lgkmcnt=0) + a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs - a00, as00 = a0_regs - a01, as01 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 1) - a02, as02 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 2) - a03, as03 = load_a_subtile_mi_regs(cur_a, cur_scales, 0, 3) - rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) - rocdl.sched_barrier(0) b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) @@ -853,15 +872,15 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): scales0 = load_scale_tile(fx.Index(0)) scales1 = load_scale_tile(fx.Index(1)) - stage_a_subtile(fx.Index(0), 0, lds_a0_base_ptr, _SCOPE["a0"], _NOALIAS["a0"]) - stage_b_subtile(fx.Index(0), 0, lds_b0_base_ptr, _SCOPE["b0"], _NOALIAS["b0"]) - stage_b_subtile(fx.Index(0), 1, lds_b0_base_ptr, _SCOPE["b0"], _NOALIAS["b0"]) - stage_a_subtile(fx.Index(0), 1, lds_a0_base_ptr, _SCOPE["a0"], _NOALIAS["a0"]) + stage_a_subtile(fx.Index(0), 0, lds_a0) + stage_b_subtile(fx.Index(0), 0, lds_b0) + stage_b_subtile(fx.Index(0), 1, lds_b0) + stage_a_subtile(fx.Index(0), 1, lds_a0) - stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1_base_ptr, _SCOPE["a1"], _NOALIAS["a1"]) - stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1_base_ptr, _SCOPE["b1"], _NOALIAS["b1"]) - stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1_base_ptr, _SCOPE["b1"], _NOALIAS["b1"]) - stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1_base_ptr, _SCOPE["a1"], _NOALIAS["a1"]) + stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) rocdl.sched_barrier(0) _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) @@ -872,15 +891,16 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): # K0 is consumed directly. K1 MFMA-ready scales are carried as # prev_refill_scales and become next_scales_ready at loop entry. - # Carry only A-top mi=0 across iterations to limit VGPR pressure; - # b0_regs carries the complete B-left 64-row register tile. - a0_regs = load_a_subtile_mi_regs(lds_a0, scales0, 0, 0) + # Seed the carried-register pipeline with K0 A-top. In later steady-state + # iterations, Q2/Q3 of the preceding iteration prefetch the next tile's + # A-top and B-left register tiles before their LDS half-pages are reused. + a0_regs = load_a_subtile_regs(lds_a0, scales0, 0) rocdl.sched_barrier(0) _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) rocdl.sched_barrier(0) - # Preload B-left for K0 to registers. + # Complete the K0 carried-register seed with B-left. b0_regs = load_b_subtile_regs(lds_b0, scales0, 0) # Main HK loop: exactly one logical K128 per iteration. @@ -895,12 +915,8 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): lds_b0, lds_a1, lds_b1, - lds_a0_base_ptr, - lds_b0_base_ptr, - _SCOPE["a0"], - _NOALIAS["a0"], - _SCOPE["b0"], - _NOALIAS["b0"], + lds_a0, + lds_b0, a0_regs, b0_regs, scales0, @@ -913,22 +929,20 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): lds_b1, lds_a0, lds_b0, - lds_a1_base_ptr, - lds_b1_base_ptr, - _SCOPE["a1"], - _NOALIAS["a1"], - _SCOPE["b1"], - _NOALIAS["b1"], + lds_a1, + lds_b1, a0_regs, b0_regs, scales1, refill_scales, ) - # Common two-page tail. After the steady loop, a0_regs/b0_regs are - # for the next page after the last consumed K128 tile. The final scale - # tile returned by the last refill belongs to the LDS page that was just - # refilled, so the tail page order depends on parity: + # Common two-page tail. The penultimate tile still uses the Q2/Q3 + # carry-prefetch scheduler to prepare A-top/B-left for the final tile, + # but it performs no K+2 data or scale refill. The final tile performs + # compute only. After the steady loop, a0_regs/b0_regs belong to the + # next tile to consume, while refill_scales belongs to the page most + # recently refilled; therefore tail page order depends on parity: # even NUM_K_TILES: consume LDS0 then final LDS1 # odd NUM_K_TILES: consume LDS1 then final LDS0 if (NUM_K_TILES % 2) == 0: @@ -975,23 +989,6 @@ def launch_gemm( c_n: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - ctx = CompilationContext.get_current() - with ir.InsertionPoint(ctx.gpu_module_body): - linkage = ir.Attribute.parse("#llvm.linkage") - for sym, size in ( - (LDS_SYM_A0, LDS_BYTES_A), - (LDS_SYM_A1, LDS_BYTES_A), - (LDS_SYM_B0, LDS_BYTES_B), - (LDS_SYM_B1, LDS_BYTES_B), - ): - llvm.GlobalOp( - global_type=ir.Type.parse(f"!llvm.array<{size} x i8>"), - sym_name=sym, - linkage=linkage, - addr_space=3, - alignment=1024, - ) - # The integration only dispatches aligned shapes; no partial-tile masking exists. grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) kernel_gemm( @@ -1003,6 +1000,6 @@ def launch_gemm( c_m, c_n, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, - ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), smem=0, stream=stream) + ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) return launch_gemm diff --git a/tests/kernels/test_mxfp8_gemm_4wave.py b/tests/kernels/test_mxfp8_gemm_4wave.py index 654f053e7..cc4dc5c3b 100644 --- a/tests/kernels/test_mxfp8_gemm_4wave.py +++ b/tests/kernels/test_mxfp8_gemm_4wave.py @@ -31,6 +31,77 @@ MX_BLOCK_K = 32 ARCH = str(get_rocm_arch()) +SWEEP_SHAPES = [ + (4096, 12288, 4096), + (4096, 4096, 4096), + (4096, 22016, 4096), + (4096, 4096, 11008), + (8192, 12288, 4096), + (8192, 4096, 4096), + (8192, 22016, 4096), + (8192, 4096, 11008), + (16384, 12288, 4096), + (16384, 4096, 4096), + (16384, 22016, 4096), + (16384, 4096, 11008), + (4096, 10240, 8192), + (4096, 8192, 8192), + (4096, 57344, 8192), + (4096, 8192, 28672), + (8192, 10240, 8192), + (8192, 8192, 8192), + (8192, 57344, 8192), + (8192, 8192, 28672), + (16384, 10240, 8192), + (16384, 8192, 8192), + (16384, 57344, 8192), + (16384, 8192, 28672), + (8192, 6144, 4096), + (8192, 28672, 4096), + (8192, 4096, 14336), + (16384, 6144, 4096), + (16384, 28672, 4096), + (16384, 4096, 14336), + (32768, 6144, 4096), + (32768, 4096, 4096), + (32768, 28672, 4096), + (32768, 4096, 14336), + (8192, 18432, 16384), + (8192, 16384, 16384), + (8192, 106496, 16384), + (8192, 16384, 53248), + (16384, 18432, 16384), + (16384, 16384, 16384), + (16384, 106496, 16384), + (16384, 16384, 53248), + (32768, 18432, 16384), + (32768, 16384, 16384), + (32768, 16384, 53248), + (8192, 4608, 3584), + (8192, 3584, 3584), + (8192, 37888, 3584), + (8192, 3584, 18944), + (16384, 4608, 3584), + (16384, 3584, 3584), + (16384, 37888, 3584), + (16384, 3584, 18944), + (32768, 4608, 3584), + (32768, 3584, 3584), + (32768, 37888, 3584), + (32768, 3584, 18944), + (8192, 59136, 8192), + (8192, 8192, 29568), + (16384, 59136, 8192), + (16384, 8192, 29568), + (32768, 10240, 8192), + (32768, 8192, 8192), + (32768, 59136, 8192), + (32768, 8192, 29568), + (4096, 6144, 4096), + (4096, 28672, 4096), + (4096, 4096, 14336), +] + if not torch.cuda.is_available(): pytest.skip("CUDA/ROCm not available. Skipping GPU tests.", allow_module_level=True) @@ -75,7 +146,14 @@ def _pack_scale_words(scales_u8: torch.Tensor) -> torch.Tensor: s32 = scales_u8.contiguous().view(rows, k32_tiles // 4, 4).to(torch.int32) iteration_major = ( - (s32[:, :, 0] | (s32[:, :, 1] << 8) | (s32[:, :, 2] << 16) | (s32[:, :, 3] << 24)).transpose(0, 1).contiguous() + ( + s32[:, :, 0] + | (s32[:, :, 1] << 8) + | (s32[:, :, 2] << 16) + | (s32[:, :, 3] << 24) + ) + .transpose(0, 1) + .contiguous() ) row = torch.arange(rows, device=scales_u8.device, dtype=torch.int64) @@ -120,8 +198,8 @@ def _bench_mxfp8_gemm_4wave( tile_m: int = 256, tile_n: int = 256, disable_xcd_remap: bool = False, - num_warmups: int = 2, - num_iters: int = 10, + num_warmups: int = 100, + num_iters: int = 100, static_weight_scale: bool = True, ): """Run + verify one MXFP8 GEMM configuration. Returns TFLOPS.""" @@ -129,7 +207,10 @@ def _bench_mxfp8_gemm_4wave( pytest.skip("MXFP8 GEMM requires CDNA4 (gfx95*)") if M % tile_m != 0 or N % tile_n != 0 or K % 128 != 0 or K < 512: - pytest.skip("MXFP8 4-wave kernel requires M/N divisible by the selected tile and K >= 512 divisible by 128") + pytest.skip( + "MXFP8 4-wave kernel requires M/N divisible by the selected tile " + "and K >= 512 divisible by 128" + ) device = torch.device("cuda") a, b, a_scale_raw, b_scale_raw = _make_mxfp8_inputs(M, N, K, device) @@ -184,6 +265,15 @@ def _args(c, a_input, b_input, a_scale_input, b_scale_input): def _launch(c, a_input, b_input, a_scale_input, b_scale_input): compiled(*_args(c, a_input, b_input, a_scale_input, b_scale_input)) + # Validate with a standalone, untimed launch. Correctness is intentionally + # independent of the warmup/timed performance run below, so every default + # or sweep shape must pass validation before its TFLOPS result is measured. + c_out.zero_() + _launch(c_out, a, b, a_scale, b_scale) + torch.cuda.synchronize() + _verify_mxfp8_output(c_out, c_ref) + print("[validation] PASS") + num_iters = max(2, int(num_iters)) _, microseconds = run_perftest( _launch, @@ -197,8 +287,6 @@ def _launch(c, a_input, b_input, a_scale_input, b_scale_input): ) torch.cuda.synchronize() - _verify_mxfp8_output(c_out, c_ref) - flops = 2 * M * N * K bytes_moved = ( M * K @@ -209,7 +297,10 @@ def _launch(c, a_input, b_input, a_scale_input, b_scale_input): ) tflops = flops / (microseconds / 1.0e6) / 1.0e12 tbps = bytes_moved / 1.0e12 / (microseconds / 1.0e6) - print(f"[flyc] Throughput: {microseconds:.1f} us, {tflops:.2f} TFLOPS, BW: {tbps:.3f} TB/s") + print( + f"[flyc] Throughput: {microseconds:.1f} us, " + f"{tflops:.2f} TFLOPS, BW: {tbps:.3f} TB/s" + ) return tflops @@ -266,6 +357,12 @@ def test_mxfp8_gemm_4wave(M, N, K, tile_m, tile_n): parser.add_argument("-K", type=int, default=8192) parser.add_argument("--tile_m", type=int, default=256) parser.add_argument("--tile_n", type=int, default=256) + parser.add_argument( + "--sweep", + action="store_true", + default=False, + help="Benchmark the full unique shape sweep.", + ) parser.add_argument( "--disable_xcd_remap", action="store_true", @@ -287,23 +384,47 @@ def test_mxfp8_gemm_4wave(M, N, K, tile_m, tile_n): "--dynamic_weight_scale", action="store_true", default=False, - help=("Use dynamic tensor arguments for weight and scales instead of static DLPack adaptors."), + help=( + "Use dynamic tensor arguments for weight and scales instead of " + "static DLPack adaptors." + ), ) args = parser.parse_args() torch.set_default_device("cuda") - try: - _bench_mxfp8_gemm_4wave( - M=args.M, - N=args.N, - K=args.K, - tile_m=args.tile_m, - tile_n=args.tile_n, - disable_xcd_remap=args.disable_xcd_remap, - num_warmups=args.num_warmups, - num_iters=args.num_iters, - static_weight_scale=not args.dynamic_weight_scale, + shapes = SWEEP_SHAPES if args.sweep else [(args.M, args.N, args.K)] + results = [] + + for shape_index, (M, N, K) in enumerate(shapes, start=1): + if args.sweep: + print( + f"\n[sweep {shape_index:02d}/{len(shapes):02d}] " + f"M={M} N={N} K={K}" + ) + + try: + tflops = _bench_mxfp8_gemm_4wave( + M=M, + N=N, + K=K, + tile_m=args.tile_m, + tile_n=args.tile_n, + disable_xcd_remap=args.disable_xcd_remap, + num_warmups=args.num_warmups, + num_iters=args.num_iters, + static_weight_scale=not args.dynamic_weight_scale, + ) + results.append((M, N, K, tflops)) + except pytest.skip.Exception as error: + print(f"Skipped: {error}") + + if args.sweep: + assert len(results) == len(shapes), ( + f"Sweep completed only {len(results)}/{len(shapes)} shapes; " + "one or more shapes were skipped or did not produce a result" ) - except pytest.skip.Exception as error: - print(f"Skipped: {error}") + print("\nMXFP8 4-wave sweep results:") + for M, N, K, tflops in results: + print(f"M={M:5d} N={N:6d} K={K:5d} TFLOPS={tflops:.2f}") + print(f"\nSweep validation: PASS ({len(results)}/{len(shapes)} shapes)") From ca9711496009988676843debb14bd04fc860bf26 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Sat, 18 Jul 2026 00:42:40 +0000 Subject: [PATCH 3/6] formatting changes --- kernels/gemm/mxfp8_gemm_4wave.py | 10 ++------- tests/kernels/test_mxfp8_gemm_4wave.py | 29 +++++--------------------- 2 files changed, 7 insertions(+), 32 deletions(-) diff --git a/kernels/gemm/mxfp8_gemm_4wave.py b/kernels/gemm/mxfp8_gemm_4wave.py index 9de97cc8b..3b5b90a04 100644 --- a/kernels/gemm/mxfp8_gemm_4wave.py +++ b/kernels/gemm/mxfp8_gemm_4wave.py @@ -385,17 +385,11 @@ def load_scale_tile(k128): def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one # 128x128 half-page (16 KiB). Each half has its own LDS base. - global_base = ( - (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) - + k_base - ) + global_base = (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) + k_base a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - global_base = ( - (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) - + k_base - ) + global_base = (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + k_base b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) def stage_a_subtile(k_base, subtile, lds_a): diff --git a/tests/kernels/test_mxfp8_gemm_4wave.py b/tests/kernels/test_mxfp8_gemm_4wave.py index cc4dc5c3b..ec0839d41 100644 --- a/tests/kernels/test_mxfp8_gemm_4wave.py +++ b/tests/kernels/test_mxfp8_gemm_4wave.py @@ -146,14 +146,7 @@ def _pack_scale_words(scales_u8: torch.Tensor) -> torch.Tensor: s32 = scales_u8.contiguous().view(rows, k32_tiles // 4, 4).to(torch.int32) iteration_major = ( - ( - s32[:, :, 0] - | (s32[:, :, 1] << 8) - | (s32[:, :, 2] << 16) - | (s32[:, :, 3] << 24) - ) - .transpose(0, 1) - .contiguous() + (s32[:, :, 0] | (s32[:, :, 1] << 8) | (s32[:, :, 2] << 16) | (s32[:, :, 3] << 24)).transpose(0, 1).contiguous() ) row = torch.arange(rows, device=scales_u8.device, dtype=torch.int64) @@ -207,10 +200,7 @@ def _bench_mxfp8_gemm_4wave( pytest.skip("MXFP8 GEMM requires CDNA4 (gfx95*)") if M % tile_m != 0 or N % tile_n != 0 or K % 128 != 0 or K < 512: - pytest.skip( - "MXFP8 4-wave kernel requires M/N divisible by the selected tile " - "and K >= 512 divisible by 128" - ) + pytest.skip("MXFP8 4-wave kernel requires M/N divisible by the selected tile and K >= 512 divisible by 128") device = torch.device("cuda") a, b, a_scale_raw, b_scale_raw = _make_mxfp8_inputs(M, N, K, device) @@ -297,10 +287,7 @@ def _launch(c, a_input, b_input, a_scale_input, b_scale_input): ) tflops = flops / (microseconds / 1.0e6) / 1.0e12 tbps = bytes_moved / 1.0e12 / (microseconds / 1.0e6) - print( - f"[flyc] Throughput: {microseconds:.1f} us, " - f"{tflops:.2f} TFLOPS, BW: {tbps:.3f} TB/s" - ) + print(f"[flyc] Throughput: {microseconds:.1f} us, {tflops:.2f} TFLOPS, BW: {tbps:.3f} TB/s") return tflops @@ -384,10 +371,7 @@ def test_mxfp8_gemm_4wave(M, N, K, tile_m, tile_n): "--dynamic_weight_scale", action="store_true", default=False, - help=( - "Use dynamic tensor arguments for weight and scales instead of " - "static DLPack adaptors." - ), + help=("Use dynamic tensor arguments for weight and scales instead of static DLPack adaptors."), ) args = parser.parse_args() @@ -398,10 +382,7 @@ def test_mxfp8_gemm_4wave(M, N, K, tile_m, tile_n): for shape_index, (M, N, K) in enumerate(shapes, start=1): if args.sweep: - print( - f"\n[sweep {shape_index:02d}/{len(shapes):02d}] " - f"M={M} N={N} K={K}" - ) + print(f"\n[sweep {shape_index:02d}/{len(shapes):02d}] M={M} N={N} K={K}") try: tflops = _bench_mxfp8_gemm_4wave( From 725099936b66bada51e505f2d1b52b6bbc3d99df Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 21 Jul 2026 14:55:58 +0000 Subject: [PATCH 4/6] add pipelined epilogue --- kernels/gemm/mxfp8_gemm_4wave.py | 151 +++++++++++++++++++++---------- 1 file changed, 101 insertions(+), 50 deletions(-) diff --git a/kernels/gemm/mxfp8_gemm_4wave.py b/kernels/gemm/mxfp8_gemm_4wave.py index 3b5b90a04..ed59c2fe1 100644 --- a/kernels/gemm/mxfp8_gemm_4wave.py +++ b/kernels/gemm/mxfp8_gemm_4wave.py @@ -292,8 +292,8 @@ def _one_i32_result(op): # Accept the result attribute names exposed by the supported MLIR Python bindings. return getattr(op, "result", getattr(op, "res", op.results[0])) - def read_pinned_accumulator(acc_idx): - acc_pin = PIN_ACC_BASE + acc_idx * 4 + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") @@ -451,6 +451,33 @@ def pinned_mfma(acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): has_side_effects=True, ) + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): + # Final-page form used by HK: destination and previous partial sum + # may be different AGPR ranges. Once old_acc_idx is consumed, its + # physical slot is dead and can be reused as a later destination. + dst_pin = PIN_ACC_BASE + dst_slot * 4 + old_pin = PIN_ACC_BASE + old_acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + arith._to_raw(a_scale), + arith._to_raw(b_scale), + ], + ( + f"v_mfma_scale_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + f"$0, $1, " + f"a[{old_pin}:{old_pin + 3}], " + f"$2, $3 " + f"op_sel:[{mi & 1},{ni & 1},0] " + f"op_sel_hi:[{mi >> 1},{ni >> 1},0]" + ), + (f"v,v,v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}},~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}"), + has_side_effects=True, + ) + def mfma_4n(acc_base, a_frag, a_scale, b0, b1, b2, b3, bs0, bs1, bs2, bs3): """Emit four N-direction scaled MFMAs into fixed physical AGPR accumulators.""" mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE @@ -464,33 +491,23 @@ def mfma_2n(acc_base, a_frag, a_scale, b0, b1, bs0, bs1, ni_base): pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, ni_base + 0) pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, ni_base + 1) - def store_output(): - c_n_idx = fx.Index(c_n) - for sm in range_constexpr(2): - # HK store mapping: each wave owns the same 64x64 warp-position - # inside all four 128x128 quadrants, not one contiguous 128x128 tile. - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - - for sn in range_constexpr(2): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - subtile_id = sm * 2 + sn + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE - subtile_m_base = subtile_m_idx * SUBTILE_M - subtile_n_base = subtile_n_idx * SUBTILE_N - - for mi in range_constexpr(MFMA_M_PER_SUBTILE): - row_base = subtile_m_base + fx.Index(mi * MFMA_M) + lane_div_16 * 4 - - for ni in range_constexpr(MFMA_N_PER_SUBTILE): - col = subtile_n_base + fx.Index(ni * MFMA_N) + lane_mod_16 - acc_idx = _acc_idx(subtile_id, mi, ni) - acc = read_pinned_accumulator(acc_idx) + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + buffer_ops.buffer_store(Vec(acc)[ii].to(fx.Float16), c_rsrc, c_idx) - for ii in range_constexpr(4): - row = row_base + fx.Index(ii) - c_idx = row * c_n_idx + col - val = Vec(acc)[ii].to(fx.Float16) - buffer_ops.buffer_store(val, c_rsrc, c_idx) # Explicit register coordinates for HK-style four-quadrant mapping. # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions @@ -820,16 +837,13 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) - mfma_4n(_acc_idx(0, 0, 0), a00, as00, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - mfma_4n(_acc_idx(0, 1, 0), a01, as01, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - mfma_4n(_acc_idx(0, 2, 0), a02, as02, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - mfma_4n(_acc_idx(0, 3, 0), a03, as03, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - rocdl.sched_barrier(0) _barrier(lgkmcnt=0) rocdl.sched_barrier(0) @@ -839,24 +853,66 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) - mfma_4n(_acc_idx(1, 0, 0), a00, as00, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(1, 1, 0), a01, as01, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(1, 2, 0), a02, as02, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(1, 3, 0), a03, as03, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - rocdl.sched_barrier(0) _barrier(lgkmcnt=0) rocdl.sched_barrier(0) - mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) - mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + a_scales = (as00, as01, as02, as03, as10, as11, as12, as13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + b_scales = (bs00, bs01, bs02, bs03, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) - mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + a_scales[a_frag_idx], + b_scales[b_frag_idx], + mi, + ni, + ) + pending.append(old_acc_idx) + + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) # Prologue: stage K0/K1 data into ping-pong LDS pages. Scales are not staged in # LDS: As/Bs are already MFMA-ready preshuffled packed uint32 [K128, row], @@ -966,11 +1022,6 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): ) hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs, scales0) - rocdl.sched_barrier(0) - _barrier() - rocdl.sched_barrier(0) - - store_output() @flyc.jit def launch_gemm( From 8b5ab266fb0a39d4eb86acccfb5792656167aa61 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 21 Jul 2026 14:59:11 +0000 Subject: [PATCH 5/6] fix syntax issues --- kernels/gemm/mxfp8_gemm_4wave.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kernels/gemm/mxfp8_gemm_4wave.py b/kernels/gemm/mxfp8_gemm_4wave.py index ed59c2fe1..86aa4f897 100644 --- a/kernels/gemm/mxfp8_gemm_4wave.py +++ b/kernels/gemm/mxfp8_gemm_4wave.py @@ -866,7 +866,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): # # Finalize accumulators in their own physical AGPR slots, but delay # each AGPR read/store until several independent final MFMAs have - # been issued. + # been issued. # # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, # MFMA 4, drain 1, MFMA 5, drain 2, ... From 5cd8b3437da799d9008d4ecf5c10ff923d3ed272 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 21 Jul 2026 15:13:12 +0000 Subject: [PATCH 6/6] remove extra blank lines --- kernels/gemm/mxfp8_gemm_4wave.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/kernels/gemm/mxfp8_gemm_4wave.py b/kernels/gemm/mxfp8_gemm_4wave.py index 86aa4f897..9a8d74fb4 100644 --- a/kernels/gemm/mxfp8_gemm_4wave.py +++ b/kernels/gemm/mxfp8_gemm_4wave.py @@ -508,7 +508,6 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): c_idx = row * fx.Index(c_n) + col buffer_ops.buffer_store(Vec(acc)[ii].to(fx.Float16), c_rsrc, c_idx) - # Explicit register coordinates for HK-style four-quadrant mapping. # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions # inside each 128x128 quadrant: @@ -1022,7 +1021,6 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): ) hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs, scales0) - @flyc.jit def launch_gemm( A: fx.Tensor,