From cf54ce65c7f662ff35423fa2fae8557179731ee3 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 00:33:21 +0000 Subject: [PATCH 01/17] Add optional FlyDSL dependency for ROCm PyTorch builds via NVTE_USE_FLYDSL --- setup.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/setup.py b/setup.py index 2f4ae06e0..ebe581cdc 100644 --- a/setup.py +++ b/setup.py @@ -156,6 +156,14 @@ def setup_requirements() -> Tuple[List[str], List[str]]: ] test_reqs: List[str] = ["pytest>=8.2.1"] + # Optional FlyDSL dependency for ROCm PyTorch builds. + if ( + rocm_build() + and "pytorch" in frameworks + and bool(int(os.getenv("NVTE_USE_FLYDSL", "0"))) + ): + install_reqs.extend(["flydsl"]) + # Framework-specific requirements if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): if "pytorch" in frameworks: From 2ed8285c40a21dddc58b237df5b90be748ba34be Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 03:15:01 +0000 Subject: [PATCH 02/17] Wire FlyDSL into the TransformerEngine MXFP8 GEMM dispatch path --- .../pytorch/cpp_extensions/gemm.py | 21 +- .../pytorch/flydsl_kernels/__init__.py | 3 + .../pytorch/flydsl_kernels/gemm/__init__.py | 13 + .../flydsl_kernels/gemm/fp8_gemm_utils.py | 262 ++++ .../flydsl_kernels/gemm/gemm_wrappers.py | 123 ++ .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 1242 +++++++++++++++++ 6 files changed, 1663 insertions(+), 1 deletion(-) create mode 100644 transformer_engine/pytorch/flydsl_kernels/__init__.py create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 3e787f820..211861425 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -460,7 +460,26 @@ def general_gemm( "beta": beta, } - out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) + # FlyDSL is currently an opt-in MXFP8-only backend. Keep every other + # datatype/recipe on the existing C++ generic_gemm path. + use_gemm_flydsl = ( + IS_HIP_EXTENSION + and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0"))) + and isinstance(A, MXFP8TensorStorage) + and isinstance(B, MXFP8TensorStorage) + ) + + if use_gemm_flydsl: + # Lazy import keeps FlyDSL off the normal Transformer Engine import path. + from ..flydsl_kernels.gemm import te_generic_gemm_flydsl + + out, bias_grad, gelu_input, extra_output = te_generic_gemm_flydsl( + *args, **kwargs + ) + else: + out, bias_grad, gelu_input, extra_output = tex.generic_gemm( + *args, **kwargs + ) if IS_HIP_EXTENSION and use_bf16_tn_output_workaround: out = cast_if_needed(out, torch.float32) diff --git a/transformer_engine/pytorch/flydsl_kernels/__init__.py b/transformer_engine/pytorch/flydsl_kernels/__init__.py new file mode 100644 index 000000000..92fa250e8 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/__init__.py @@ -0,0 +1,3 @@ +from . import gemm + +__all__ = ["gemm"] \ No newline at end of file diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py new file mode 100644 index 000000000..784d17d2f --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL GEMM kernels (dense, non-grouped) for BF16/FP16/FP32/FP8/MXFP8.""" + +from .gemm_wrappers import ( + te_generic_gemm_flydsl, +) + +__all__ = [ + "te_generic_gemm_flydsl", +] \ No newline at end of file diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py new file mode 100644 index 000000000..a8bdfb717 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects.fly_rocdl import TargetAddressSpace +from flydsl.expr import arith, const_expr, range_constexpr, rocdl +from flydsl.expr.typing import Vector as Vec + +# ceildiv is the canonical cdiv from the shared layer +def cdiv(numer: int, denom: int) -> int: + return (numer + denom - 1) // denom + + +ceildiv = cdiv + +def divmod(a, b): + """Integer divmod that works on DSL values (e.g. ``Int32``). + + The builtin ``divmod`` rejects DSL scalar types, so this uses the overloaded + ``//`` / ``%`` operators to emit the corresponding ops. + """ + return (a // b, a % b) + + +def preshuffle_b(b_t): + """Permute row-major ``B_T`` ``(N, K)`` for ``b_preshuffled=True``.""" + n, k = b_t.shape[-2:] + assert n % 16 == 0 and k % 64 == 0, f"need N%16==0 and K%64==0, got N={n} K={k}" + return b_t.reshape(n // 16, 16, k // 64, 4, 16).permute(0, 2, 3, 1, 4).contiguous() + + +def make_fp8_buffer_tensor(arg_i8, fp8_ir_t): + # max_size=False with no num_records_bytes: cosize(layout) becomes a + # runtime expression because TensorAdaptor defaults to layout-dynamic + # memref (post #554), so the descriptor adapts to the actual tensor + # extent and no longer bakes the first-call's shape into IR. + t_i8 = fx.rocdl.make_buffer_tensor(arg_i8, max_size=False) + iter_i8 = fx.get_iter(t_i8) + f8_buf_ptr_ty = fx.PointerType.get( + elem_ty=fp8_ir_t, + address_space=TargetAddressSpace.BufferDesc, + alignment=fx.PointerType(iter_i8.type).alignment, + ) + iter_f8 = fx.recast_iter(f8_buf_ptr_ty, iter_i8) + return fx.Tensor(fx.make_view(iter_f8, fx.get_layout(t_i8))) + + +def swizzle_128(row, col): + offset = row * 128 + col + swizzle = ((offset % (16 * 128)) >> 8) << 4 + swizzled_offset = offset ^ swizzle + return swizzled_offset // 128, swizzled_offset % 128 + + +def compute_global_swizzle(lane_id, wave_id, K, n_rounds, preshuffled): + offsets = [] + n_waves = fx.block_dim.x // 64 + for round in range_constexpr(n_rounds): + if const_expr(preshuffled): + row = lane_id % 8 + wave_id * 8 + round * (n_waves * 8) + col = (lane_id // 8) * 16 + offsets.append( + (row // 16) * (K * 16) + (row % 16) * 16 + (col // 64) * 1024 + ((col % 64) // 16) * 256 + (col % 16) + ) + else: + row = lane_id // 8 + wave_id * 8 + round * (n_waves * 8) + col = (lane_id % 8) * 16 + r, c = swizzle_128(row, col) + offsets.append(r * K + c) + return offsets + + +class G2SLoader: + def __init__(self, gl_src, gl_offsets, n_load_steps, lds_dtype, wave_id): + self.g2lds_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128) + self.LdsPtr_t = fx.PointerType.get(lds_dtype, 2, 512) + self.gl_src = gl_src + self.gl_offsets = gl_offsets + self.n_load_steps = n_load_steps + self.wave_id = wave_id + self.n_waves = fx.block_dim.x // 64 + + def _lds_dst_at(self, lds_dst, step): + step_off = self.wave_id * 1024 + step * (self.n_waves * 1024) + base_i32 = fx.Int32(fx.ptrtoint(lds_dst.ptr)) + sum_i32 = base_i32 + fx.Int32(step_off) + lds_ptr = fx.inttoptr(self.LdsPtr_t, sum_i32) + return fx.make_view(lds_ptr, fx.make_layout(1, 1)) + + def load(self, lds_dst, k_offset): + for step in range_constexpr(self.n_load_steps): + src = fx.slice(self.gl_src, (None, fx.Int32(self.gl_offsets[step]))) + dst = self._lds_dst_at(lds_dst, step) + fx.copy(self.g2lds_atom, src, dst, soffset=fx.Int32(k_offset)) + + def load_one(self, lds_dst, k_offset, step): + src = fx.slice(self.gl_src, (None, fx.Int32(self.gl_offsets[step]))) + dst = self._lds_dst_at(lds_dst, step) + fx.copy(self.g2lds_atom, src, dst, soffset=fx.Int32(k_offset)) + + +def pack_i32x4_i32x8(lo, hi): + # Pack two i32x4 as one i32x8 + return lo.shuffle(hi, list(range(8))) + + +class S2RLoader: + def __init__(self, wave_idx, n_tiles): + self.lane_id = fx.thread_idx.x % 64 + self.wave_idx = wave_idx + self.n_tiles = n_tiles + + def _vec_load_16xf8(self, lds_src, offset): + off_tup = fx.make_int_tuple(offset) + ptr_off = fx.add_offset(lds_src.ptr, off_tup) + i8_iter = fx.recast_iter(fx.Uint8, ptr_off) + view = fx.make_view(i8_iter, fx.make_layout(16, 1)) + return view.load() + + def load(self, lds_src, preshuffled=False): + frag = [] + for i in range_constexpr(self.n_tiles): + halves = [] + row = self.wave_idx * (self.n_tiles * 16) + i * 16 + self.lane_id % 16 + for step in range_constexpr(2): + col = (self.lane_id // 16) * 16 + step * 64 + if const_expr(preshuffled): + offset = (row // 8) * 1024 + (row % 8) * 16 + (col // 16) * 128 + else: + row_swz, col_swz = swizzle_128(row, col) + offset = row_swz * 128 + col_swz + v = self._vec_load_16xf8(lds_src, offset) + halves.append(v.bitcast(fx.Int32)) + frag.append(pack_i32x4_i32x8(halves[0], halves[1])) + return frag + + def load_one(self, lds_src, lds_offset): + v = self._vec_load_16xf8(lds_src, lds_offset) + return v.bitcast(fx.Int32) + + +class StoreC: + def __init__(self, A_scale, B_scale, C, c_rows, c_cols, c_idx_fn, n_tiles_a, n_tiles_b): + self.c_rows = c_rows + self.c_cols = c_cols + self.lane_id = fx.thread_idx.x % 64 + self.c_idx_fn = c_idx_fn + self.n_tiles_a = n_tiles_a + self.n_tiles_b = n_tiles_b + # Exact byte counts from compile-time shape (BF16 C output, FP32 scales). + # ``num_records_bytes`` is required when ``max_size=False`` -- see + # ``make_buffer_tensor`` docstring for the silent-OOB rationale. + c_nbytes = c_rows * c_cols * 2 # BFloat16 = 2 bytes + sa_nbytes = c_rows * 4 # Float32 row-wise scale + sb_nbytes = c_cols * 4 # Float32 col-wise scale + gC = fx.rocdl.make_buffer_tensor(C, max_size=False, num_records_bytes=c_nbytes) + gSA = fx.rocdl.make_buffer_tensor(A_scale, max_size=False, num_records_bytes=sa_nbytes) + gSB = fx.rocdl.make_buffer_tensor(B_scale, max_size=False, num_records_bytes=sb_nbytes) + self.c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + self.sa_div = fx.logical_divide(gSA, fx.make_layout(1, 1)) + self.sb_div = fx.logical_divide(gSB, fx.make_layout(1, 1)) + + self.scale_atom_4 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32) + self.scale_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + self.out_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), fx.BFloat16) + self.reg_f32_4 = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Float32) + self.reg_f32_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + self.reg_bf16_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.BFloat16) + + def _load_scale_vec4(self, row): + fx.copy(self.scale_atom_4, fx.slice(self.sa_div, (None, fx.Int32(row))), self.reg_f32_4) + return Vec(fx.memref_load_vec(self.reg_f32_4)) + + def _load_scale_scalar(self, col): + fx.copy(self.scale_atom_1, fx.slice(self.sb_div, (None, fx.Int32(col))), self.reg_f32_1) + return Vec(fx.memref_load_vec(self.reg_f32_1))[0] + + def _store_bf16(self, value_bf16, c_index): + fx.memref_store_vec(Vec.filled(1, value_bf16, fx.BFloat16), self.reg_bf16_1) + fx.copy(self.out_atom_1, self.reg_bf16_1, fx.slice(self.c_div, (None, fx.Int32(c_index)))) + + def store(self, c_frag, base_row, base_col): + a_scales = [ + self._load_scale_vec4(base_row + i * 16 + (self.lane_id // 16) * 4) for i in range_constexpr(self.n_tiles_a) + ] + b_scales = [ + self._load_scale_scalar(base_col + i * 16 + self.lane_id % 16) for i in range_constexpr(self.n_tiles_b) + ] + for ti in range_constexpr(self.n_tiles_a): + row = base_row + ti * 16 + (self.lane_id // 16) * 4 + for tj in range_constexpr(self.n_tiles_b): + col = base_col + tj * 16 + self.lane_id % 16 + col_valid = col < self.c_cols + oob = fx.Int32(self.c_rows * self.c_cols) + vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)]) + for i in range_constexpr(4): + scaled = (vec_f32[i] * (a_scales[ti][i] * b_scales[tj])).to(fx.BFloat16) + c_index = (row + i) * self.c_cols + col + self._store_bf16(scaled, arith.select(col_valid, c_index, oob)) + + +def wait_barrier(count): + _llvm.inline_asm( + res=None, + operands_=[], + asm_string=f"s_waitcnt vmcnt({count})\ns_barrier", + constraints="", + has_side_effects=True, + ) + + +class Mfma16x16x128: + def __init__(self, n_tiles_a, n_tiles_b): + self.atom = fx.make_mma_atom(fx.rocdl.cdna4.MFMA_Scale(16, 16, 128, fx.Float8E4M3FN)) + self.zero_value = Vec.filled(4, 0.0, fx.Float32) + self.n_tiles_a = n_tiles_a + self.n_tiles_b = n_tiles_b + + def idx(self, i, j): + return i * self.n_tiles_b + j + + def _make_operand_frag(self, value): + frag = fx.make_rmem_tensor(8, fx.Int32) + frag.store(Vec(value)) + return frag + + def _make_accum_frag(self, value): + frag = fx.make_rmem_tensor(4, fx.Float32) + frag.store(Vec(value)) + return frag + + def _do_mma(self, a, b, c): + a_frag = self._make_operand_frag(a) + b_frag = self._make_operand_frag(b) + c_frag = self._make_accum_frag(c) + fx.gemm(self.atom, c_frag, a_frag, b_frag, c_frag) + return c_frag.load().ir_value() + + def call(self, a, b, c, *, set_prio=True): + assert len(a) == self.n_tiles_a + assert len(b) == self.n_tiles_b + assert len(c) == self.n_tiles_a * self.n_tiles_b + + a_frags = [self._make_operand_frag(a[idx]) for idx in range_constexpr(self.n_tiles_a)] + b_frags = [self._make_operand_frag(b[idx]) for idx in range_constexpr(self.n_tiles_b)] + c_frags = [self._make_accum_frag(c[idx]) for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b)] + if const_expr(set_prio): + rocdl.s_setprio(1) + for i in range_constexpr(self.n_tiles_a): + for j in range_constexpr(self.n_tiles_b): + cf = c_frags[self.idx(i, j)] + fx.gemm(self.atom, cf, a_frags[i], b_frags[j], cf) + if const_expr(set_prio): + rocdl.s_setprio(0) + rocdl.s_barrier() + return [c_frags[idx].load().ir_value() for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b)] + + def call_one(self, a, b, c, i, j): + assert i < self.n_tiles_a and j < self.n_tiles_b + + return self._do_mma(a[i], b[j], c[self.idx(i, j)]) \ No newline at end of file diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py new file mode 100644 index 000000000..90a12130c --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -0,0 +1,123 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Minimal TE entry point for the FlyDSL MXFP8 TN backend.""" + +import torch +import transformer_engine_torch as tex + +from .mxfp8_gemm import mxfp8_matmul + + +def te_generic_gemm_flydsl( + A, + transa, + B, + transb, + D, + quantizer, + output_dtype, + bias=None, + bias_type=None, + gelu=False, + gelu_in=None, + grad=False, + workspace=None, + workspaceSize=0, + accumulate=False, + use_split_accumulator=False, + comm_overlap=None, + comm_type=None, + extra_output=None, + bulk_overlap=False, + alpha=1.0, + beta=0.0, +): + """Run the FlyDSL MXFP8 kernel for TE's TN path.""" + if not transa or transb: + raise NotImplementedError( + "FlyDSL MXFP8 currently supports only transa=True, transb=False" + ) + + if output_dtype not in (None, tex.DType.kFloat16): + raise NotImplementedError( + f"FlyDSL MXFP8 currently supports only FP16 output, got {output_dtype}" + ) + + if quantizer is not None: + raise NotImplementedError("FlyDSL MXFP8 output quantization is not implemented") + + if float(alpha) != 1.0 or float(beta) != 0.0: + raise NotImplementedError("FlyDSL MXFP8 supports only alpha=1 and beta=0") + + if accumulate: + raise NotImplementedError("FlyDSL MXFP8 accumulation is not implemented") + + if bias is not None and bias.numel() != 0: + raise NotImplementedError("FlyDSL MXFP8 bias is not implemented") + + if gelu or grad: + raise NotImplementedError("FlyDSL MXFP8 GELU/gradient epilogues are not implemented") + + # TE TN path: + # A rowwise payload: weight [N, K] + # B rowwise payload: activation [..., K] + A_data = A._rowwise_data + A_scale = A._rowwise_scale_inv + B_data = B._rowwise_data + B_scale = B._rowwise_scale_inv + + if A_data is None or A_scale is None: + raise RuntimeError("A does not contain rowwise MXFP8 data and scales") + + if B_data is None or B_scale is None: + raise RuntimeError("B does not contain rowwise MXFP8 data and scales") + + n, k = A_data.shape + B_flat = B_data.reshape(-1, B_data.shape[-1]) + m, kb = B_flat.shape + + if kb != k: + raise ValueError(f"MXFP8 inner dimensions do not match: {k} and {kb}") + + A_scale = A_scale.reshape(n, -1) + B_scale = B_scale.reshape(m, -1) + + output_shape = (*B_data.shape[:-1], n) + + if D is None: + D = torch.empty( + output_shape, + dtype=torch.float16, + device=B_data.device, + ) + else: + if tuple(D.shape) != output_shape: + raise ValueError( + f"D shape {tuple(D.shape)} does not match expected {output_shape}" + ) + + if D.dtype != torch.float16: + raise TypeError( + f"FlyDSL MXFP8 requires FP16 output, got {D.dtype}" + ) + + if not D.is_contiguous(): + raise ValueError("FlyDSL MXFP8 requires contiguous output storage") + + # Public mxfp8_matmul contract: + # a: [M, K] + # a_scale: [M, K/32] + # b: [K, N] + # b_scale: [N, K/32] + # c: [M, N] FP16 + mxfp8_matmul( + B_flat, + B_scale, + A_data.transpose(0, 1), + A_scale, + D.view(m, n), + ) + + return D, None, None, None diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py new file mode 100644 index 000000000..bde328ced --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -0,0 +1,1242 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL MXFP8 4-wave GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K128 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes A and B as FP8 E4M3 tensors shaped [M, K] and [N, K], and writes +float16 C shaped [M, N]. The public ``mxfp8_matmul`` entry point accepts the +Transformer Engine TN contract and performs the required private adaptation. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +# Transformer Engine-local FlyDSL utilities. +from .fp8_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_fp8_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 128 + +# Public metadata consumed by wrappers — keep. +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K +SCALE_GROUP_SIZE = 32 + + +def pack_mx32_scales_iter(scales_u8: torch.Tensor) -> torch.Tensor: + """Pack raw [Rows, K/32] E8M0 uint8 scales as [K/128, Rows] uint32. + + This is the intermediate HK/TE iteration-major form: each word contains + four consecutive K32 scale bytes for one K128 iteration and one matrix row. + It is *not* the final MFMA operand layout. + """ + assert scales_u8.dtype == torch.uint8 + rows, qk = scales_u8.shape + assert qk % 4 == 0 + s32 = scales_u8.contiguous().view(rows, qk // 4, 4).to(torch.int32) + packed = ( + s32[:, :, 0] + | (s32[:, :, 1] << 8) + | (s32[:, :, 2] << 16) + | (s32[:, :, 3] << 24) + ) + return packed.transpose(0, 1).contiguous() + + +def pack_mx32_scales_for_hk(scales_u8: torch.Tensor) -> torch.Tensor: + """True HK MFMA scale packing: raw [Rows, K/32] -> [K/128, Rows] i32. + + HK's GEMM hot loop loads one uint32 scale operand per lane for each 64-row + A/B half. The four bytes in that operand correspond to the four 16-row + MFMA slices inside the 64-row half; the scaled-MFMA op_sel/op_sel_hi bits + select the byte. With this layout the GEMM kernel does no hot-loop byte + extraction or broadcast. + """ + assert scales_u8.dtype == torch.uint8 + rows, qk = scales_u8.shape + assert qk % 4 == 0 + assert rows % 64 == 0, f"rows={rows} must be a multiple of 64 for HK MFMA scale packing" + + scale_iter = pack_mx32_scales_iter(scales_u8) # [K/128, Rows], int32 + device = scales_u8.device + + row = torch.arange(rows, device=device, dtype=torch.int64) + r16 = row % 16 + k_sub = (row // 16) % 4 + tile = row // 64 + + packed = torch.zeros_like(scale_iter) + for g in range(4): + src_row = tile * 64 + g * 16 + r16 + src_val = scale_iter[:, src_row] + byte_val = (src_val >> (k_sub * 8).view(1, rows)) & 0xFF + packed |= byte_val << (g * 8) + + return packed.contiguous() + +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] + ) + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + + +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 _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + 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), + ) + + +def _compile_kernel(K: int): + """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_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + NUM_THREADS = 256 + WARP_SIZE = 64 + + 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 = 16 + + 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_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, As: fx.Tensor, B: fx.Tensor, Bs: fx.Tensor, C: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32 + ): + 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") + + num_blocks_m = c_m // BLOCK_M + num_blocks_n = c_n // BLOCK_N + + pid_m, pid_n = _xcd_swizzle(num_blocks_m, 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) + + # 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) + 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 _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 _to_raw_inline_asm_operand(value): + # TODO: Replace arith._to_raw once FlyDSL exposes a supported public + # API for passing wrapped values to llvm.InlineAsmOp. _to_raw is + # deprecated, but remains heavily used internally by FlyDSL. + return arith._to_raw(value) + + 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") + 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(): + # 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, + 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_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_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_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_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, 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 + + 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, + [ + _to_raw_inline_asm_operand(a_frag), + _to_raw_inline_asm_operand(b_frag), + _to_raw_inline_asm_operand(a_scale), + _to_raw_inline_asm_operand(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 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, + [ + _to_raw_inline_asm_operand(a_frag), + _to_raw_inline_asm_operand(b_frag), + _to_raw_inline_asm_operand(a_scale), + _to_raw_inline_asm_operand(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 + 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_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_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) + + + # 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, sn) + 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_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. + a_scales = scale_tile[0] if sm == 0 else scale_tile[1] + 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, + refill_b, + 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 + # 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 + + 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) + + # 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) + 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) + + 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) + + 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) + + 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) + + 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) + + 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) + + 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) + + 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_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 = 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) + 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) + mfma_2n(_acc_idx(1, 0, 2), a00, as00, b12, b13, bs12, bs13, 2) + + 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) + mfma_2n(_acc_idx(1, 1, 2), a01, as01, b12, b13, bs12, bs13, 2) + + 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) + mfma_2n(_acc_idx(1, 2, 2), a02, as02, b12, b13, bs12, bs13, 2) + + 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) + 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_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_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 + + 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_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_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 + + # 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) + + 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) + + 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) + + # 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], + # 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) + 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) + 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) + 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. + + # 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) + + # 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. + # 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, + lds_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, + lds_b1, + a0_regs, + b0_regs, + scales1, + refill_scales, + ) + + # 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: + 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) + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + As: fx.Tensor, + B: fx.Tensor, + Bs: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + As, + B, + Bs, + C, + 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), stream=stream) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch(K: int): + return _compile_kernel(K) + + + +def mxfp8_matmul( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing TN MXFP8 adapter. + + Public/backend contract: + a: [M, K] FP8 payload + a_scale: [M, K/32] raw E8M0 bytes + b: [K, N] FP8 payload + b_scale: [N, K/32] raw E8M0 bytes + c: [M, N] float16 output + + The optimized HK core currently consumes B as row-major [N, K] and consumes + MFMA-ready packed int32 scales. Keep those implementation details behind + this adapter so the TE-facing contract matches the Triton/TE TN contract. + """ + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL MXFP8 TN expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + m, k = a.shape + kb, n = b.shape + if kb != k: + raise ValueError(f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}") + + expected_a_scale = (m, k // SCALE_GROUP_SIZE) + expected_b_scale = (n, k // SCALE_GROUP_SIZE) + if tuple(a_scale.shape) != expected_a_scale: + raise ValueError( + f"A scale shape {tuple(a_scale.shape)} != expected {expected_a_scale}" + ) + if tuple(b_scale.shape) != expected_b_scale: + raise ValueError( + f"B scale shape {tuple(b_scale.shape)} != expected {expected_b_scale}" + ) + if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + raise TypeError( + "FlyDSL MXFP8 expects raw E8M0 scales stored as torch.uint8" + ) + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype != torch.float16: + raise TypeError( + f"The current FlyDSL MXFP8 kernel stores float16 output, got {c.dtype}" + ) + + # TE/Triton expose B logically as [K, N]. The existing optimized HK core + # streams contiguous K rows, so adapt B to its private [N, K] representation. + # In the normal TE TN path, b is itself a transpose view of contiguous + # rowwise weight storage, so b.T is already contiguous and this is not a + # physical transpose/copy. + b_hk = b.transpose(0, 1).contiguous() + + # Convert TE's raw per-K32 E8M0 scales into the MFMA-ready words consumed by + # the optimized scaled-MFMA hot loop. + a_scale_hk = pack_mx32_scales_for_hk(a_scale) + b_scale_hk = pack_mx32_scales_for_hk(b_scale) + + doGemm(a, a_scale_hk, b_hk, b_scale_hk, c, stream=stream) + +def doGemm( + A: torch.Tensor, + As: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + C: torch.Tensor, + stream=None, +): + """Launch the K-specialized kernel with runtime M/N. + + A and B are shaped [M, K] and [N, K]. As/Bs are preshuffled packed + uint32 scale words shaped [K/128, M] and [K/128, N]. C is shaped [M, N]. + M and N are not hardcoded; K is used only to choose/cache the compile-time + specialized launch function. + """ + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + assert M_runtime % _BLOCK_M == 0, f"M={M_runtime} must be a multiple of {_BLOCK_M}" + assert N_runtime % _BLOCK_N == 0, f"N={N_runtime} must be a multiple of {_BLOCK_N}" + assert K_runtime % _BLOCK_K == 0, f"K={K_runtime} must be a multiple of {_BLOCK_K}" + num_k_tiles = K_runtime // _BLOCK_K + assert num_k_tiles >= 4, f"K={K_runtime} gives {num_k_tiles} K128 tiles; need at least 4" + expected_as = (K_runtime // _BLOCK_K, M_runtime) + expected_bs = (K_runtime // _BLOCK_K, N_runtime) + assert As.dtype == torch.int32, f"As dtype {As.dtype} != torch.int32 packed scales" + assert Bs.dtype == torch.int32, f"Bs dtype {Bs.dtype} != torch.int32 packed scales" + assert As.shape == expected_as, f"As shape {tuple(As.shape)} != {expected_as}" + assert Bs.shape == expected_bs, f"Bs shape {tuple(Bs.shape)} != {expected_bs}" + assert C.shape == (M_runtime, N_runtime), ( + f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" + ) + if stream is None: + stream = torch.cuda.current_stream() + # Match the Transformer Engine integration descriptor contract exactly. The optimized + # G2SLoader path consumes flat byte-addressed A/B tensors; scales and C are + # likewise passed as flat contiguous storage. Passing the original 2-D + # torch tensors changes the tensor descriptor/layout seen by + # make_fp8_buffer_tensor() and causes the loader's linear offsets to address + # the wrong elements. + A_arg = A.view(torch.uint8).contiguous().view(-1) + B_arg = B.view(torch.uint8).contiguous().view(-1) + As_arg = As.contiguous().view(-1) + Bs_arg = Bs.contiguous().view(-1) + C_arg = C.contiguous().view(-1) + + launch = _cached_launch(int(K_runtime)) + launch( + A_arg, + As_arg, + B_arg, + Bs_arg, + C_arg, + M_runtime, + N_runtime, + stream=stream, + ) From 17b9b7442749bb5c5cd14392c95cdea0e3287cc5 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 13:21:02 +0000 Subject: [PATCH 03/17] Add initial support for TN BF16 FlyDSL GEMM backend --- .../pytorch/cpp_extensions/gemm.py | 18 +- .../pytorch/flydsl_kernels/gemm/bf16_gemm.py | 1065 +++++++++++++++++ .../flydsl_kernels/gemm/fp16_gemm_utils.py | 93 ++ .../flydsl_kernels/gemm/gemm_wrappers.py | 258 +++- 4 files changed, 1386 insertions(+), 48 deletions(-) create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 211861425..30496df8d 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -460,13 +460,25 @@ def general_gemm( "beta": beta, } - # FlyDSL is currently an opt-in MXFP8-only backend. Keep every other + # FlyDSL is currently an opt-in BF16/MXFP8-only backend. Keep every other # datatype/recipe on the existing C++ generic_gemm path. + + is_mxfp8_gemm = ( + isinstance(A, MXFP8TensorStorage) + and isinstance(B, MXFP8TensorStorage) + ) + + is_bf16_gemm = ( + isinstance(A, torch.Tensor) + and isinstance(B, torch.Tensor) + and A.dtype == torch.bfloat16 + and B.dtype == torch.bfloat16 + ) + use_gemm_flydsl = ( IS_HIP_EXTENSION and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0"))) - and isinstance(A, MXFP8TensorStorage) - and isinstance(B, MXFP8TensorStorage) + and (is_mxfp8_gemm or is_bf16_gemm) ) if use_gemm_flydsl: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py new file mode 100644 index 000000000..ea489c38a --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py @@ -0,0 +1,1065 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL BF16 4-wave GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K64 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes A and B as BF16 tensors shaped [M, K] and [N, K], and writes BF16 C +shaped [M, N]. The public ``bf16_matmul`` entry point accepts Transformer +Engine's TN contract and performs the required private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +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 + +# Transformer Engine-local FlyDSL utilities. +from .fp16_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_bf16_byte_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 64 + +# Public metadata consumed by wrappers. +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +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 = 2 +VEC_BYTES = 16 + +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 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "bf16_pp_smem_a0" +LDS_SYM_A1 = "bf16_pp_smem_a1" +LDS_SYM_B0 = "bf16_pp_smem_b0" +LDS_SYM_B1 = "bf16_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 64 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +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] + ) + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + + +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 _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + 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), + ) + + +def _compile_kernel(K: int, use_xcd_remap: bool = True): + """Build the specialized 4-wave kernel for compile-time ``K``. + + ``K`` must contain at least four K64 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + NUM_THREADS = 256 + WARP_SIZE = 64 + + 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 = 2 + VEC_BYTES = 16 + + 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 + + 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} K64 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LDS_BYTES_HALF = LDS_ELEMS_HALF * ELEM_BYTES + LOAD_PASSES_HALF = LDS_BYTES_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x64 BF16 page is two independent 128x64 half-pages. + # Store LDS as bytes so BufferCopyLDS128b sees i8 on both source and + # destination. Each half-page remains exactly 16 KiB. + a0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + 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) + + # A/B arrive as contiguous uint8 byte views. Keeping staging byte-addressed + # preserves the original 16-byte G2L instruction cadence and vmcnt values. + gA = make_bf16_byte_buffer_tensor(A) + gB = make_bf16_byte_buffer_tensor(B) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + 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) + + # 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 * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + gl_off_b = compute_global_swizzle(lane, wave_id, K * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, fx.Uint8.ir_type, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, fx.Uint8.ir_type, 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) + 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 BF16. + 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 _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) + + 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") + 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) + + def hot_loop_scheduler_q_refill_2n(): + # Eight refill VMEM operations overlap four independent 8-MFMA + # groups (two K32 slices x two N-halves). + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_mfma(8) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # Eight refill VMEM operations and eight distributed A-bottom LDS + # reads overlap four independent 8-MFMA K32 groups. + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_dsrd(2) + rocdl.sched_mfma(8) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # Eight two-read prefetch groups overlap four complete-quadrant + # 16-MFMA groups (two K32 slices for each of Q2 and Q3). + for _ in range_constexpr(4): + rocdl.sched_dsrd(4) + rocdl.sched_mfma(16) + rocdl.sched_barrier(0) + + 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 + # 128x64 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 * ELEM_BYTES) + k_base * fx.Index(ELEM_BYTES) + 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 * ELEM_BYTES) + k_base * fx.Index(ELEM_BYTES) + 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_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 16-byte half of the wave operand tile. + # 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_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, 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 * ELEM_BYTES)) + + 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 _bf16_k32_frag(full_frag, k32): + # A/B 16x64 BF16 wave fragments are i32x8. Each K32 MFMA + # consumes one contiguous i32x4 slice (eight BF16 values/lane). + lo = k32 * 4 + v = Vec(full_frag) + return Vec.from_elements( + [v[lo], v[lo + 1], v[lo + 2], v[lo + 3]], + fx.Int32, + ) + + def _pinned_bf16_mfma_once(acc_idx, a_k32, b_k32): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [arith._to_raw(a_k32), arith._to_raw(b_k32)], + ( + f"v_mfma_f32_16x16x32_bf16 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Accumulate one logical 16x16x64 BF16 product into pinned AGPRs.""" + for k32 in range_constexpr(2): + _pinned_bf16_mfma_once( + acc_idx, + _bf16_k32_frag(a_frag, k32), + _bf16_k32_frag(b_frag, k32), + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + # The final logical K64 update is two in-place K32 MFMAs. + assert dst_slot == old_acc_idx + pinned_mfma(old_acc_idx, a_frag, b_frag) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + def mfma_2n_4mi_k32(subtile_id, n_base, k32, a0, a1, a2, a3, b0, b1): + """Issue one K32 slice for a 4x2 accumulator slab.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1) + for mi in range_constexpr(4): + a_k32 = _bf16_k32_frag(a_frags[mi], k32) + for nj in range_constexpr(2): + _pinned_bf16_mfma_once( + _acc_idx(subtile_id, mi, n_base + nj), + a_k32, + _bf16_k32_frag(b_frags[nj], k32), + ) + + def mfma_4n_4mi_k32(subtile_id, k32, a0, a1, a2, a3, b0, b1, b2, b3): + """Issue one K32 slice for a complete 4x4 quadrant.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1, b2, b3) + for mi in range_constexpr(4): + a_k32 = _bf16_k32_frag(a_frags[mi], k32) + for ni in range_constexpr(4): + _pinned_bf16_mfma_once( + _acc_idx(subtile_id, mi, ni), + a_k32, + _bf16_k32_frag(b_frags[ni], k32), + ) + + 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_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.BFloat16), 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, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + 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 * ELEM_BYTES) + return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): + + # 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) + + # 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 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + # 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. Compute is K32-major across all 16 + # independent accumulators, eliminating the two-deep same-AGPR + # dependency chains produced by pinned_mfma(). + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n_4mi_k32(0, 0, 0, a00, a01, a02, a03, b00, b01) + + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n_4mi_k32(0, 2, 0, a00, a01, a02, a03, b02, b03) + + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + # K32 slice 0 already covers K[0:32]. + + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + # Keep this refill/LDS-read slot compute-free. + + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n_4mi_k32(0, 0, 1, a00, a01, a02, a03, b00, b01) + + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n_4mi_k32(0, 2, 1, a00, a01, a02, a03, b02, b03) + + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + # K32 slice 1 already covers K[32:64]. + + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + # Keep this refill/LDS-read slot compute-free. + + 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 = 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) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n_4mi_k32(1, 0, 0, a00, a01, a02, a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + # K32 slice 0 already covers K[0:32]. + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + # Keep this refill slot compute-free. + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n_4mi_k32(1, 0, 1, a00, a01, a02, a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + # K32 slice 1 already covers K[32:64]. + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + # Keep this refill slot compute-free. + 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, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n_4mi_k32(2, 0, a10, a11, a12, a13, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + # K32 slice 0 already covers K[0:32]. + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n_4mi_k32(2, 1, a10, a11, a12, a13, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + # K32 slice 1 already covers K[32:64]. + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n_4mi_k32(3, 0, a10, a11, a12, a13, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + # K32 slice 0 already covers K[0:32]. + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n_4mi_k32(3, 1, a10, a11, a12, a13, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + # K32 slice 1 already covers K[32:64]. + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # 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], + ) + 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. + 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) + 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) + rocdl.sched_barrier(0) + + a0_regs = load_a_subtile_regs(lds_a0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K64 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + else: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # 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, + 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), stream=stream) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch(K: int, use_xcd_remap: bool = True): + return _compile_kernel(K, use_xcd_remap=use_xcd_remap) + + +def bf16_matmul( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing TN BF16 GEMM adapter. + + Public/backend contract: + a: [M, K] BF16 + b: [K, N] BF16 + c: [M, N] BF16 output + + The optimized core streams both operands with K contiguous and therefore + privately consumes B as [N, K]. In the normal TE TN path, ``b`` is a + transpose view of contiguous rowwise weight storage, so ``b.T`` is already + contiguous and does not require a physical transpose. + """ + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL BF16 TN expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + m, k = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + if a.dtype != torch.bfloat16 or b.dtype != torch.bfloat16: + raise TypeError( + "FlyDSL BF16 GEMM expects both operands to have torch.bfloat16 dtype, " + f"got {a.dtype} and {b.dtype}" + ) + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype != torch.bfloat16: + raise TypeError( + f"The current FlyDSL BF16 kernel stores torch.bfloat16 output, got {c.dtype}" + ) + if a.device != b.device or a.device != c.device: + raise ValueError( + f"A, B, and C must be on the same device, got {a.device}, {b.device}, and {c.device}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL BF16 GEMM requires contiguous output storage") + + b_hk = b.transpose(0, 1).contiguous() + doGemm(a, b_hk, c, stream=stream) + + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch the private K-specialized BF16 core. + + A and B are shaped [M, K] and [N, K]; C is shaped [M, N]. M and N + remain runtime values, while K selects the cached compile-time specialization. + """ + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + assert A.dtype == torch.bfloat16 and B.dtype == torch.bfloat16 + assert C.dtype == torch.bfloat16 + assert M_runtime % _BLOCK_M == 0, f"M={M_runtime} must be a multiple of {_BLOCK_M}" + assert N_runtime % _BLOCK_N == 0, f"N={N_runtime} must be a multiple of {_BLOCK_N}" + assert K_runtime % _BLOCK_K == 0, f"K={K_runtime} must be a multiple of {_BLOCK_K}" + num_k_tiles = K_runtime // _BLOCK_K + assert num_k_tiles >= 4, f"K={K_runtime} gives {num_k_tiles} K64 tiles; need at least 4" + assert C.shape == (M_runtime, N_runtime) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.contiguous().view(torch.uint8).view(-1) + B_arg = B.contiguous().view(torch.uint8).view(-1) + C_arg = C.view(-1) + launch = _cached_launch(int(K_runtime), bool(use_xcd_remap)) + launch(A_arg, B_arg, C_arg, M_runtime, N_runtime, stream=stream) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py new file mode 100644 index 000000000..5aaeab3b1 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors +"""Minimal byte-staging helpers for the first-pass BF16 four-wave GEMM.""" + +import flydsl.expr as fx +from flydsl.expr import const_expr, range_constexpr + +# ceildiv is the canonical cdiv from the shared layer +def cdiv(numer: int, denom: int) -> int: + return (numer + denom - 1) // denom + + +ceildiv = cdiv + +def divmod(a, b): + return (a // b, a % b) + + +def swizzle_128(row, col_in_bytes): + """HK 128-byte row XOR swizzle; ``col_in_bytes`` is a byte coordinate.""" + offset = row * 128 + col_in_bytes + swizzle = ((offset % (16 * 128)) >> 8) << 4 + swizzled_offset = offset ^ swizzle + return swizzled_offset // 128, swizzled_offset % 128 + + +def make_bf16_byte_buffer_tensor(arg_u8): + """Create a byte-addressed buffer tensor from a contiguous BF16 uint8 view.""" + return fx.rocdl.make_buffer_tensor(arg_u8, max_size=False) + + +def compute_global_swizzle(lane_id, wave_id, row_stride_bytes, n_rounds, preshuffled=False): + offsets = [] + n_waves = fx.block_dim.x // 64 + for round in range_constexpr(n_rounds): + if const_expr(preshuffled): + raise AssertionError("BF16 first-pass port does not support preshuffled operands") + row = lane_id // 8 + wave_id * 8 + round * (n_waves * 8) + col_bytes = (lane_id % 8) * 16 + r, c = swizzle_128(row, col_bytes) + offsets.append(r * row_stride_bytes + c) + return offsets + + +class G2SLoader: + """Issue raw 16-byte buffer-to-LDS copies. + + Both the global source and LDS destination must be byte-addressed. Fly's copy lowering does not legalize an i8 buffer source paired with a bf16 LDS + destination even when the transfer width is the same 128 bits. + """ + def __init__(self, gl_src, gl_offsets, n_load_steps, lds_dtype, wave_id): + self.g2lds_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128) + self.LdsPtr_t = fx.PointerType.get(lds_dtype, 2, 512) + self.gl_src = gl_src + self.gl_offsets = gl_offsets + self.n_load_steps = n_load_steps + self.wave_id = wave_id + self.n_waves = fx.block_dim.x // 64 + + def _lds_dst_at(self, lds_dst, step): + step_off = self.wave_id * 1024 + step * (self.n_waves * 1024) + base_i32 = fx.Int32(fx.ptrtoint(lds_dst.ptr)) + lds_ptr = fx.inttoptr(self.LdsPtr_t, base_i32 + fx.Int32(step_off)) + return fx.make_view(lds_ptr, fx.make_layout(1, 1)) + + def load(self, lds_dst, byte_offset): + for step in range_constexpr(self.n_load_steps): + src = fx.slice(self.gl_src, (None, fx.Int32(self.gl_offsets[step]))) + fx.copy(self.g2lds_atom, src, self._lds_dst_at(lds_dst, step), soffset=fx.Int32(byte_offset)) + + def load_one(self, lds_dst, byte_offset, step): + src = fx.slice(self.gl_src, (None, fx.Int32(self.gl_offsets[step]))) + fx.copy(self.g2lds_atom, src, self._lds_dst_at(lds_dst, step), soffset=fx.Int32(byte_offset)) + + +def pack_i32x4_i32x8(lo, hi): + return lo.shuffle(hi, list(range(8))) + + +class S2RLoader: + """Raw 16-byte LDS reader used to assemble an i32x8 BF16 K64 fragment.""" + def __init__(self, wave_idx, n_tiles): + self.lane_id = fx.thread_idx.x % 64 + self.wave_idx = wave_idx + self.n_tiles = n_tiles + + def _vec_load_16bytes(self, lds_src, offset): + ptr_off = fx.add_offset(lds_src.ptr, fx.make_int_tuple(offset)) + i8_iter = fx.recast_iter(fx.Uint8, ptr_off) + return fx.make_view(i8_iter, fx.make_layout(16, 1)).load() + + def load_one(self, lds_src, lds_offset): + return self._vec_load_16bytes(lds_src, lds_offset).bitcast(fx.Int32) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 90a12130c..cfe5a1006 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -2,67 +2,59 @@ # # See LICENSE for license information. -"""Minimal TE entry point for the FlyDSL MXFP8 TN backend.""" +"""TE entry points for the FlyDSL GEMM backend.""" import torch import transformer_engine_torch as tex +from .bf16_gemm import bf16_matmul from .mxfp8_gemm import mxfp8_matmul -def te_generic_gemm_flydsl( - A, - transa, - B, - transb, - D, +def _validate_common_epilogue( + *, quantizer, - output_dtype, - bias=None, - bias_type=None, - gelu=False, - gelu_in=None, - grad=False, - workspace=None, - workspaceSize=0, - accumulate=False, - use_split_accumulator=False, - comm_overlap=None, - comm_type=None, - extra_output=None, - bulk_overlap=False, - alpha=1.0, - beta=0.0, + bias, + gelu, + grad, + accumulate, + alpha, + beta, ): - """Run the FlyDSL MXFP8 kernel for TE's TN path.""" - if not transa or transb: + """Validate features not yet implemented by the FlyDSL GEMM backend.""" + if quantizer is not None: raise NotImplementedError( - "FlyDSL MXFP8 currently supports only transa=True, transb=False" + "FlyDSL GEMM output quantization is not implemented" ) - if output_dtype not in (None, tex.DType.kFloat16): + if float(alpha) != 1.0 or float(beta) != 0.0: raise NotImplementedError( - f"FlyDSL MXFP8 currently supports only FP16 output, got {output_dtype}" + "FlyDSL GEMM currently supports only alpha=1 and beta=0" ) - if quantizer is not None: - raise NotImplementedError("FlyDSL MXFP8 output quantization is not implemented") - - if float(alpha) != 1.0 or float(beta) != 0.0: - raise NotImplementedError("FlyDSL MXFP8 supports only alpha=1 and beta=0") - if accumulate: - raise NotImplementedError("FlyDSL MXFP8 accumulation is not implemented") + raise NotImplementedError( + "FlyDSL GEMM accumulation is not implemented" + ) if bias is not None and bias.numel() != 0: - raise NotImplementedError("FlyDSL MXFP8 bias is not implemented") + raise NotImplementedError( + "FlyDSL GEMM bias is not implemented" + ) if gelu or grad: - raise NotImplementedError("FlyDSL MXFP8 GELU/gradient epilogues are not implemented") + raise NotImplementedError( + "FlyDSL GEMM GELU/gradient epilogues are not implemented" + ) + + +def _is_mxfp8_operand(t): + """Return whether ``t`` exposes TE MXFP8 rowwise storage.""" + return hasattr(t, "_rowwise_data") and hasattr(t, "_rowwise_scale_inv") - # TE TN path: - # A rowwise payload: weight [N, K] - # B rowwise payload: activation [..., K] + +def _run_mxfp8_tn(A, B, D): + """Run the existing FlyDSL MXFP8 TN path.""" A_data = A._rowwise_data A_scale = A._rowwise_scale_inv B_data = B._rowwise_data @@ -83,7 +75,6 @@ def te_generic_gemm_flydsl( A_scale = A_scale.reshape(n, -1) B_scale = B_scale.reshape(m, -1) - output_shape = (*B_data.shape[:-1], n) if D is None: @@ -97,14 +88,14 @@ def te_generic_gemm_flydsl( raise ValueError( f"D shape {tuple(D.shape)} does not match expected {output_shape}" ) - if D.dtype != torch.float16: raise TypeError( f"FlyDSL MXFP8 requires FP16 output, got {D.dtype}" ) - if not D.is_contiguous(): - raise ValueError("FlyDSL MXFP8 requires contiguous output storage") + raise ValueError( + "FlyDSL MXFP8 requires contiguous output storage" + ) # Public mxfp8_matmul contract: # a: [M, K] @@ -120,4 +111,181 @@ def te_generic_gemm_flydsl( D.view(m, n), ) - return D, None, None, None + return D + + +def _run_bf16_tn(A, B, D): + """Run FlyDSL BF16 for TE's TN operand convention. + + TE supplies: + A: weight [N, K] + B: activation [..., K] + + ``bf16_matmul`` consumes: + a: activation [M, K] + b: weight.T [K, N] + c: output [M, N] + """ + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError( + "FlyDSL BF16 GEMM expects plain torch.Tensor operands" + ) + + if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: + raise TypeError( + "FlyDSL BF16 GEMM requires BF16 inputs, " + f"got A={A.dtype} and B={B.dtype}" + ) + + if A.ndim != 2: + raise ValueError( + f"FlyDSL BF16 TN expects weight A to be rank 2, got {tuple(A.shape)}" + ) + if B.ndim < 2: + raise ValueError( + f"FlyDSL BF16 TN expects activation B to have rank >= 2, got {tuple(B.shape)}" + ) + + n, k = A.shape + B_flat = B.reshape(-1, B.shape[-1]) + m, kb = B_flat.shape + + if kb != k: + raise ValueError( + f"BF16 inner dimensions do not match: A{tuple(A.shape)} and " + f"B{tuple(B.shape)}" + ) + + output_shape = (*B.shape[:-1], n) + + if D is None: + D = torch.empty( + output_shape, + dtype=torch.bfloat16, + device=B.device, + ) + else: + if tuple(D.shape) != output_shape: + raise ValueError( + f"D shape {tuple(D.shape)} does not match expected {output_shape}" + ) + if D.dtype != torch.bfloat16: + raise TypeError( + f"FlyDSL BF16 requires BF16 output, got {D.dtype}" + ) + if D.device != B.device: + raise ValueError( + f"D must be on {B.device}, got {D.device}" + ) + if not D.is_contiguous(): + raise ValueError( + "FlyDSL BF16 requires contiguous output storage" + ) + + if A.device != B.device: + raise ValueError( + f"A and B must be on the same device, got {A.device} and {B.device}" + ) + + bf16_matmul( + B_flat, + A.transpose(0, 1), + D.view(m, n), + ) + + return D + + +def te_generic_gemm_flydsl( + A, + transa, + B, + transb, + D, + quantizer, + output_dtype, + bias=None, + bias_type=None, + gelu=False, + gelu_in=None, + grad=False, + workspace=None, + workspaceSize=0, + accumulate=False, + use_split_accumulator=False, + comm_overlap=None, + comm_type=None, + extra_output=None, + bulk_overlap=False, + alpha=1.0, + beta=0.0, +): + """Run a supported FlyDSL GEMM through TE's generic GEMM interface. + + Currently supported: + - MXFP8 TN input with FP16 output + - BF16 TN input with BF16 output + """ + del bias_type + del gelu_in + del workspace + del workspaceSize + del use_split_accumulator + del comm_overlap + del comm_type + del extra_output + del bulk_overlap + + if not transa or transb: + raise NotImplementedError( + "FlyDSL GEMM currently supports only transa=True, transb=False" + ) + + _validate_common_epilogue( + quantizer=quantizer, + bias=bias, + gelu=gelu, + grad=grad, + accumulate=accumulate, + alpha=alpha, + beta=beta, + ) + + a_is_mxfp8 = _is_mxfp8_operand(A) + b_is_mxfp8 = _is_mxfp8_operand(B) + + if a_is_mxfp8 or b_is_mxfp8: + if not (a_is_mxfp8 and b_is_mxfp8): + raise ValueError( + "Mixed MXFP8 and non-MXFP8 FlyDSL GEMM inputs are not supported" + ) + + if output_dtype not in (None, tex.DType.kFloat16): + raise NotImplementedError( + "FlyDSL MXFP8 currently supports only FP16 output, " + f"got {output_dtype}" + ) + + D = _run_mxfp8_tn(A, B, D) + return D, None, None, None + + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError( + "Unsupported FlyDSL GEMM operand types: " + f"{type(A).__name__} and {type(B).__name__}" + ) + + if A.dtype == torch.bfloat16 and B.dtype == torch.bfloat16: + if output_dtype not in (None, tex.DType.kBFloat16): + raise NotImplementedError( + "FlyDSL BF16 currently supports only BF16 output, " + f"got {output_dtype}" + ) + + D = _run_bf16_tn(A, B, D) + return D, None, None, None + + raise NotImplementedError( + "FlyDSL GEMM currently supports only MXFP8 or BF16 inputs; " + f"got A={A.dtype} and B={B.dtype}" + ) From 155e81e737bdd1e76505aa29f0fd206765ae43ac Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 13:38:05 +0000 Subject: [PATCH 04/17] Add initial support for TN FP16 FlyDSL GEMM backend --- .../pytorch/cpp_extensions/gemm.py | 10 +- .../pytorch/flydsl_kernels/gemm/fp16_gemm.py | 1072 +++++++++++++++++ .../flydsl_kernels/gemm/gemm_wrappers.py | 88 +- 3 files changed, 1163 insertions(+), 7 deletions(-) create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 30496df8d..a385e9e6d 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -460,7 +460,7 @@ def general_gemm( "beta": beta, } - # FlyDSL is currently an opt-in BF16/MXFP8-only backend. Keep every other + # FlyDSL is currently an opt-in FP16/BF16/MXFP8-only backend. Keep every other # datatype/recipe on the existing C++ generic_gemm path. is_mxfp8_gemm = ( @@ -468,17 +468,17 @@ def general_gemm( and isinstance(B, MXFP8TensorStorage) ) - is_bf16_gemm = ( + is_fp16_bf16_gemm = ( isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor) - and A.dtype == torch.bfloat16 - and B.dtype == torch.bfloat16 + and A.dtype == torch.bfloat16 or A.dtype == torch.float16 + and B.dtype == torch.bfloat16 or B.dtype == torch.float16 ) use_gemm_flydsl = ( IS_HIP_EXTENSION and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0"))) - and (is_mxfp8_gemm or is_bf16_gemm) + and (is_mxfp8_gemm or is_fp16_bf16_gemm) ) if use_gemm_flydsl: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py new file mode 100644 index 000000000..66f68816e --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py @@ -0,0 +1,1072 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL FP16 4-wave GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K64 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes A and B as FP16 tensors shaped [M, K] and [N, K], and writes FP16 C +shaped [M, N]. The public ``fp16_matmul`` entry point accepts Transformer +Engine's TN contract and performs the required private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +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 + +# Transformer Engine-local FlyDSL utilities. +from .fp16_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_bf16_byte_buffer_tensor as make_fp16_byte_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 64 + +# Public metadata consumed by wrappers. +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +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 = 2 +VEC_BYTES = 16 + +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 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp16_pp_smem_a0" +LDS_SYM_A1 = "fp16_pp_smem_a1" +LDS_SYM_B0 = "fp16_pp_smem_b0" +LDS_SYM_B1 = "fp16_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 64 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + +def make_fp16_inputs(M, N, K, device="cuda"): + """Generate FP16 A[M,K] and B[N,K] inputs.""" + A = (torch.randn(M, K, device=device) * 0.5).to(torch.float16) + B = (torch.randn(N, K, device=device) * 0.5).to(torch.float16) + return A, B + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +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] + ) + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + + +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 _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + 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), + ) + + +def _compile_kernel(K: int, use_xcd_remap: bool = True): + """Build the specialized 4-wave kernel for compile-time ``K``. + + ``K`` must contain at least four K64 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + NUM_THREADS = 256 + WARP_SIZE = 64 + + 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 = 2 + VEC_BYTES = 16 + + 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 + + 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} K64 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LDS_BYTES_HALF = LDS_ELEMS_HALF * ELEM_BYTES + LOAD_PASSES_HALF = LDS_BYTES_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x64 FP16 page is two independent 128x64 half-pages. + # Store LDS as bytes so BufferCopyLDS128b sees i8 on both source and + # destination. Each half-page remains exactly 16 KiB. + a0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + 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) + + # A/B arrive as contiguous uint8 byte views. Keeping staging byte-addressed + # preserves the original 16-byte G2L instruction cadence and vmcnt values. + gA = make_fp16_byte_buffer_tensor(A) + gB = make_fp16_byte_buffer_tensor(B) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + 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) + + # 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 * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + gl_off_b = compute_global_swizzle(lane, wave_id, K * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, fx.Uint8.ir_type, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, fx.Uint8.ir_type, 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) + 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 FP16. + 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 _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) + + 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") + 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) + + def hot_loop_scheduler_q_refill_2n(): + # Eight refill VMEM operations overlap four independent 8-MFMA + # groups (two K32 slices x two N-halves). + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_mfma(8) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # Eight refill VMEM operations and eight distributed A-bottom LDS + # reads overlap four independent 8-MFMA K32 groups. + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_dsrd(2) + rocdl.sched_mfma(8) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # Eight two-read prefetch groups overlap four complete-quadrant + # 16-MFMA groups (two K32 slices for each of Q2 and Q3). + for _ in range_constexpr(4): + rocdl.sched_dsrd(4) + rocdl.sched_mfma(16) + rocdl.sched_barrier(0) + + 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 + # 128x64 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 * ELEM_BYTES) + k_base * fx.Index(ELEM_BYTES) + 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 * ELEM_BYTES) + k_base * fx.Index(ELEM_BYTES) + 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_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 16-byte half of the wave operand tile. + # 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_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, 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 * ELEM_BYTES)) + + 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 _fp16_k32_frag(full_frag, k32): + # A/B 16x64 FP16 wave fragments are i32x8. Each K32 MFMA + # consumes one contiguous i32x4 slice (eight FP16 values/lane). + lo = k32 * 4 + v = Vec(full_frag) + return Vec.from_elements( + [v[lo], v[lo + 1], v[lo + 2], v[lo + 3]], + fx.Int32, + ) + + def _pinned_fp16_mfma_once(acc_idx, a_k32, b_k32): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [arith._to_raw(a_k32), arith._to_raw(b_k32)], + ( + f"v_mfma_f32_16x16x32_f16 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Accumulate one logical 16x16x64 FP16 product into pinned AGPRs.""" + for k32 in range_constexpr(2): + _pinned_fp16_mfma_once( + acc_idx, + _fp16_k32_frag(a_frag, k32), + _fp16_k32_frag(b_frag, k32), + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + # The final logical K64 update is two in-place K32 MFMAs. + assert dst_slot == old_acc_idx + pinned_mfma(old_acc_idx, a_frag, b_frag) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + def mfma_2n_4mi_k32(subtile_id, n_base, k32, a0, a1, a2, a3, b0, b1): + """Issue one K32 slice for a 4x2 accumulator slab.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1) + for mi in range_constexpr(4): + a_k32 = _fp16_k32_frag(a_frags[mi], k32) + for nj in range_constexpr(2): + _pinned_fp16_mfma_once( + _acc_idx(subtile_id, mi, n_base + nj), + a_k32, + _fp16_k32_frag(b_frags[nj], k32), + ) + + def mfma_4n_4mi_k32(subtile_id, k32, a0, a1, a2, a3, b0, b1, b2, b3): + """Issue one K32 slice for a complete 4x4 quadrant.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1, b2, b3) + for mi in range_constexpr(4): + a_k32 = _fp16_k32_frag(a_frags[mi], k32) + for ni in range_constexpr(4): + _pinned_fp16_mfma_once( + _acc_idx(subtile_id, mi, ni), + a_k32, + _fp16_k32_frag(b_frags[ni], k32), + ) + + 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_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) + + + # 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, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + 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 * ELEM_BYTES) + return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): + + # 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) + + # 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 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + # 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. Compute is K32-major across all 16 + # independent accumulators, eliminating the two-deep same-AGPR + # dependency chains produced by pinned_mfma(). + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n_4mi_k32(0, 0, 0, a00, a01, a02, a03, b00, b01) + + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n_4mi_k32(0, 2, 0, a00, a01, a02, a03, b02, b03) + + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + # K32 slice 0 already covers K[0:32]. + + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + # Keep this refill/LDS-read slot compute-free. + + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n_4mi_k32(0, 0, 1, a00, a01, a02, a03, b00, b01) + + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n_4mi_k32(0, 2, 1, a00, a01, a02, a03, b02, b03) + + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + # K32 slice 1 already covers K[32:64]. + + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + # Keep this refill/LDS-read slot compute-free. + + 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 = 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) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n_4mi_k32(1, 0, 0, a00, a01, a02, a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + # K32 slice 0 already covers K[0:32]. + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + # Keep this refill slot compute-free. + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n_4mi_k32(1, 0, 1, a00, a01, a02, a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + # K32 slice 1 already covers K[32:64]. + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + # Keep this refill slot compute-free. + 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, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n_4mi_k32(2, 0, a10, a11, a12, a13, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + # K32 slice 0 already covers K[0:32]. + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n_4mi_k32(2, 1, a10, a11, a12, a13, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + # K32 slice 1 already covers K[32:64]. + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n_4mi_k32(3, 0, a10, a11, a12, a13, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + # K32 slice 0 already covers K[0:32]. + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n_4mi_k32(3, 1, a10, a11, a12, a13, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + # K32 slice 1 already covers K[32:64]. + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # 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], + ) + 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. + 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) + 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) + rocdl.sched_barrier(0) + + a0_regs = load_a_subtile_regs(lds_a0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K64 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + else: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # 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, + 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), stream=stream) + + return launch_gemm + + +@functools.lru_cache(maxsize=None) +def _cached_launch(K: int, use_xcd_remap: bool = True): + return _compile_kernel(K, use_xcd_remap=use_xcd_remap) + + +def fp16_matmul( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing TN FP16 GEMM adapter. + + Public/backend contract: + a: [M, K] FP16 + b: [K, N] FP16 + c: [M, N] FP16 output + + The optimized core streams both operands with K contiguous and therefore + privately consumes B as [N, K]. In the normal TE TN path, ``b`` is a + transpose view of contiguous rowwise weight storage, so ``b.T`` is already + contiguous and does not require a physical transpose. + """ + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL FP16 TN expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + m, k = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + if a.dtype != torch.float16 or b.dtype != torch.float16: + raise TypeError( + "FlyDSL FP16 GEMM expects both operands to have torch.float16 dtype, " + f"got {a.dtype} and {b.dtype}" + ) + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype != torch.float16: + raise TypeError( + f"The current FlyDSL FP16 kernel stores torch.float16 output, got {c.dtype}" + ) + if a.device != b.device or a.device != c.device: + raise ValueError( + f"A, B, and C must be on the same device, got {a.device}, {b.device}, and {c.device}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL FP16 GEMM requires contiguous output storage") + + b_hk = b.transpose(0, 1).contiguous() + doGemm(a, b_hk, c, stream=stream) + + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch the private K-specialized FP16 core. + + A and B are shaped [M, K] and [N, K]; C is shaped [M, N]. M and N + remain runtime values, while K selects the cached compile-time specialization. + """ + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + assert A.dtype == torch.float16 and B.dtype == torch.float16 + assert C.dtype == torch.float16 + assert M_runtime % _BLOCK_M == 0, f"M={M_runtime} must be a multiple of {_BLOCK_M}" + assert N_runtime % _BLOCK_N == 0, f"N={N_runtime} must be a multiple of {_BLOCK_N}" + assert K_runtime % _BLOCK_K == 0, f"K={K_runtime} must be a multiple of {_BLOCK_K}" + num_k_tiles = K_runtime // _BLOCK_K + assert num_k_tiles >= 4, f"K={K_runtime} gives {num_k_tiles} K64 tiles; need at least 4" + assert C.shape == (M_runtime, N_runtime) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.contiguous().view(torch.uint8).view(-1) + B_arg = B.contiguous().view(torch.uint8).view(-1) + C_arg = C.view(-1) + launch = _cached_launch(int(K_runtime), bool(use_xcd_remap)) + launch(A_arg, B_arg, C_arg, M_runtime, N_runtime, stream=stream) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index cfe5a1006..65ca7a913 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -8,6 +8,7 @@ import transformer_engine_torch as tex from .bf16_gemm import bf16_matmul +from .fp16_gemm import fp16_matmul from .mxfp8_gemm import mxfp8_matmul @@ -196,6 +197,78 @@ def _run_bf16_tn(A, B, D): return D +def _run_fp16_tn(A, B, D): + """Run FlyDSL FP16 for TE's TN operand convention.""" + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError( + "FlyDSL FP16 GEMM expects plain torch.Tensor operands" + ) + + if A.dtype != torch.float16 or B.dtype != torch.float16: + raise TypeError( + "FlyDSL FP16 GEMM requires FP16 inputs, " + f"got A={A.dtype} and B={B.dtype}" + ) + + if A.ndim != 2: + raise ValueError( + f"FlyDSL FP16 TN expects weight A to be rank 2, got {tuple(A.shape)}" + ) + if B.ndim < 2: + raise ValueError( + f"FlyDSL FP16 TN expects activation B to have rank >= 2, got {tuple(B.shape)}" + ) + + n, k = A.shape + B_flat = B.reshape(-1, B.shape[-1]) + m, kb = B_flat.shape + + if kb != k: + raise ValueError( + f"FP16 inner dimensions do not match: A{tuple(A.shape)} and " + f"B{tuple(B.shape)}" + ) + + output_shape = (*B.shape[:-1], n) + + if D is None: + D = torch.empty( + output_shape, + dtype=torch.float16, + device=B.device, + ) + else: + if tuple(D.shape) != output_shape: + raise ValueError( + f"D shape {tuple(D.shape)} does not match expected {output_shape}" + ) + if D.dtype != torch.float16: + raise TypeError( + f"FlyDSL FP16 requires FP16 output, got {D.dtype}" + ) + if D.device != B.device: + raise ValueError( + f"D must be on {B.device}, got {D.device}" + ) + if not D.is_contiguous(): + raise ValueError( + "FlyDSL FP16 requires contiguous output storage" + ) + + if A.device != B.device: + raise ValueError( + f"A and B must be on the same device, got {A.device} and {B.device}" + ) + + fp16_matmul( + B_flat, + A.transpose(0, 1), + D.view(m, n), + ) + + return D + + def te_generic_gemm_flydsl( A, transa, @@ -225,6 +298,7 @@ def te_generic_gemm_flydsl( Currently supported: - MXFP8 TN input with FP16 output - BF16 TN input with BF16 output + - FP16 TN input with FP16 output """ del bias_type del gelu_in @@ -235,7 +309,7 @@ def te_generic_gemm_flydsl( del comm_type del extra_output del bulk_overlap - + if not transa or transb: raise NotImplementedError( "FlyDSL GEMM currently supports only transa=True, transb=False" @@ -285,7 +359,17 @@ def te_generic_gemm_flydsl( D = _run_bf16_tn(A, B, D) return D, None, None, None + if A.dtype == torch.float16 and B.dtype == torch.float16: + if output_dtype not in (None, tex.DType.kFloat16): + raise NotImplementedError( + "FlyDSL FP16 currently supports only FP16 output, " + f"got {output_dtype}" + ) + + D = _run_fp16_tn(A, B, D) + return D, None, None, None + raise NotImplementedError( - "FlyDSL GEMM currently supports only MXFP8 or BF16 inputs; " + "FlyDSL GEMM currently supports only MXFP8, BF16, or FP16 inputs; " f"got A={A.dtype} and B={B.dtype}" ) From 492a29c3fc7e12a5d503cd9332da7b77ffa890f4 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 14:14:24 +0000 Subject: [PATCH 05/17] Add initial support for TN FP8 FlyDSL GEMM backend --- .../pytorch/cpp_extensions/gemm.py | 28 +- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 1091 +++++++++++++++++ .../flydsl_kernels/gemm/gemm_wrappers.py | 225 +++- 3 files changed, 1335 insertions(+), 9 deletions(-) create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index a385e9e6d..df11bf277 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -460,25 +460,37 @@ def general_gemm( "beta": beta, } - # FlyDSL is currently an opt-in FP16/BF16/MXFP8-only backend. Keep every other - # datatype/recipe on the existing C++ generic_gemm path. - + # FlyDSL is currently an opt-in TN backend for: + # - MXFP8 + # - tensor-wise E4M3 x E4M3 FP8 + # - matching BF16 or FP16 inputs + # Keep every other datatype, recipe, and layout on the existing C++ path. + from ..tensor.storage.float8_tensor_storage import Float8TensorStorage + is_mxfp8_gemm = ( isinstance(A, MXFP8TensorStorage) and isinstance(B, MXFP8TensorStorage) ) + is_fp8_gemm = ( + isinstance(A, Float8TensorStorage) + and isinstance(B, Float8TensorStorage) + and A._fp8_dtype == tex.DType.kFloat8E4M3 + and B._fp8_dtype == tex.DType.kFloat8E4M3 + ) + is_fp16_bf16_gemm = ( - isinstance(A, torch.Tensor) - and isinstance(B, torch.Tensor) - and A.dtype == torch.bfloat16 or A.dtype == torch.float16 - and B.dtype == torch.bfloat16 or B.dtype == torch.float16 + type(A) is torch.Tensor + and type(B) is torch.Tensor + and A.dtype == B.dtype + and A.dtype in (torch.bfloat16, torch.float16) ) use_gemm_flydsl = ( IS_HIP_EXTENSION + and layout == "TN" and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0"))) - and (is_mxfp8_gemm or is_fp16_bf16_gemm) + and (is_mxfp8_gemm or is_fp8_gemm or is_fp16_bf16_gemm) ) if use_gemm_flydsl: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py new file mode 100644 index 000000000..a85b8ea91 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -0,0 +1,1091 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL tensor-wise FP8 4-wave GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K128 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes A and B as FP8 E4M3 tensors shaped [M, K] and [N, K], one FP32 inverse +scale per operand, and writes float16 C shaped [M, N]. The public ``fp8_matmul`` +entry point accepts Transformer Engine's TN contract and performs the required +private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +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 + +# Transformer Engine-local FlyDSL utilities. +from .fp8_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_fp8_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 128 + +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +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 = 16 + +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 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp8_pp_smem_a0" +LDS_SYM_A1 = "fp8_pp_smem_a1" +LDS_SYM_B0 = "fp8_pp_smem_b0" +LDS_SYM_B1 = "fp8_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 128 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +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] + ) + + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + +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 _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + 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), + ) + + +def _compile_kernel(K: int, 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_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + NUM_THREADS = 256 + WARP_SIZE = 64 + + 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 = 16 + + 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 + + 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_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, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + 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)) + a_scale_rsrc = buffer_ops.create_buffer_resource(A_scale_inv, max_size=True) + b_scale_rsrc = buffer_ops.create_buffer_resource(B_scale_inv, max_size=True) + output_scale = ( + buffer_ops.buffer_load(a_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + * buffer_ops.buffer_load(b_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + ) + 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) + + # 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) + 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 _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) + + 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") + 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) + + def hot_loop_scheduler_q_refill_2n(): + 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(): + 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(): + for _ in range_constexpr(8): + rocdl.sched_dsrd(2) + rocdl.sched_mfma(4) + rocdl.sched_barrier(0) + + 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_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_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_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, 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 + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Issue ordinary FP8 MFMA into the fixed physical accumulator bank.""" + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + """Final-page ordinary FP8 MFMA with independently named AGPR source/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), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + f"$0, $1, " + f"a[{old_pin}:{old_pin + 3}]" + ), + ( + f"v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}}," + f"~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}" + ), + has_side_effects=True, + ) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + 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_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] * output_scale).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: + # 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, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + 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, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): + + # 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) + + # 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 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + # 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) + 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, b00, b01) + + 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, b02, b03) + + 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, b00, b01) + + 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, b02, b03) + + 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, b00, b01) + + 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, b02, b03) + + 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, b00, b01) + + 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, b02, b03) + + 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 = 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) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n(_acc_idx(1, 0, 0), a00, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n(_acc_idx(1, 0, 2), a00, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + mfma_2n(_acc_idx(1, 1, 0), a01, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + mfma_2n(_acc_idx(1, 1, 2), a01, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n(_acc_idx(1, 2, 0), a02, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n(_acc_idx(1, 2, 2), a02, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + mfma_2n(_acc_idx(1, 3, 0), a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + mfma_2n(_acc_idx(1, 3, 2), a03, b12, b13) + 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, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # 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], + ) + 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. + 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) + 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) + rocdl.sched_barrier(0) + + a0_regs = load_a_subtile_regs(lds_a0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K128 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + else: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # 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_inv, + B_scale_inv, + 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), stream=stream) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch(K: int, use_xcd_remap: bool = True): + return _compile_kernel(K, use_xcd_remap=use_xcd_remap) + + + +def fp8_matmul( + a: torch.Tensor, + a_scale_inv: torch.Tensor, + b: torch.Tensor, + b_scale_inv: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing TN tensor-wise FP8 adapter. + + Public/backend contract: + a: [M, K] FP8 E4M3 activation payload + a_scale_inv: one-element FP32 inverse quantization scale + b: [K, N] FP8 E4M3 weight payload + b_scale_inv: one-element FP32 inverse quantization scale + c: [M, N] float16 output + + The optimized private core streams both operands as row-major [Rows, K], + so B is adapted from TE's logical [K, N] representation to [N, K]. + """ + if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): + raise TypeError("FlyDSL FP8 GEMM expects plain torch.Tensor payloads") + + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL FP8 TN expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + if a.dtype != torch.float8_e4m3fn or b.dtype != torch.float8_e4m3fn: + raise TypeError( + "FlyDSL FP8 GEMM requires torch.float8_e4m3fn payloads, " + f"got A={a.dtype} and B={b.dtype}" + ) + + m, k = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + + for name, scale in ( + ("A_scale_inv", a_scale_inv), + ("B_scale_inv", b_scale_inv), + ): + if not isinstance(scale, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if scale.dtype != torch.float32 or scale.numel() != 1: + raise TypeError( + f"{name} must contain exactly one FP32 value, got " + f"dtype={scale.dtype}, shape={tuple(scale.shape)}" + ) + + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype != torch.float16: + raise TypeError( + f"The current FlyDSL FP8 kernel stores float16 output, got {c.dtype}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL FP8 requires contiguous output storage") + + tensors = (a, b, a_scale_inv, b_scale_inv, c) + if any(t.device != a.device for t in tensors[1:]): + raise ValueError( + "A, B, inverse scales, and C must be on the same device" + ) + + # In the normal TE TN path, b is a transpose view of contiguous rowwise + # weight storage, so b.T is already contiguous and this does not require a + # physical transpose/copy. + b_hk = b.transpose(0, 1).contiguous() + doGemm( + a, + b_hk, + c, + a_scale_inv, + b_scale_inv, + stream=stream, + ) + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + A_scale_inv: torch.Tensor, + B_scale_inv: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + assert A.dtype == torch.float8_e4m3fn, f"A dtype {A.dtype} != torch.float8_e4m3fn" + assert B.dtype == torch.float8_e4m3fn, f"B dtype {B.dtype} != torch.float8_e4m3fn" + assert C.dtype == torch.float16, f"C dtype {C.dtype} != torch.float16" + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + assert M_runtime % _BLOCK_M == 0, f"M={M_runtime} must be a multiple of {_BLOCK_M}" + assert N_runtime % _BLOCK_N == 0, f"N={N_runtime} must be a multiple of {_BLOCK_N}" + assert K_runtime % _BLOCK_K == 0, f"K={K_runtime} must be a multiple of {_BLOCK_K}" + num_k_tiles = K_runtime // _BLOCK_K + assert num_k_tiles >= 4, f"K={K_runtime} gives {num_k_tiles} K128 tiles; need at least 4" + assert A_scale_inv.dtype == torch.float32 and A_scale_inv.numel() == 1 + assert B_scale_inv.dtype == torch.float32 and B_scale_inv.numel() == 1 + assert C.shape == (M_runtime, N_runtime), ( + f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" + ) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.view(torch.uint8).contiguous().view(-1) + B_arg = B.view(torch.uint8).contiguous().view(-1) + C_arg = C.contiguous().view(-1) + A_scale_arg = A_scale_inv.contiguous().view(-1) + B_scale_arg = B_scale_inv.contiguous().view(-1) + + launch = _cached_launch(int(K_runtime), bool(use_xcd_remap)) + launch( + A_arg, + B_arg, + C_arg, + A_scale_arg, + B_scale_arg, + M_runtime, + N_runtime, + stream=stream, + ) + diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 65ca7a913..8245eeef5 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -7,11 +7,41 @@ import torch import transformer_engine_torch as tex +from transformer_engine.pytorch.utils import get_device_compute_capability + from .bf16_gemm import bf16_matmul from .fp16_gemm import fp16_matmul +from .fp8_gemm import fp8_matmul from .mxfp8_gemm import mxfp8_matmul +def reinterpret_as_fp8_tensor( + a: torch.Tensor, + dtype: tex.DType, +) -> torch.Tensor: + """View TE's uint8 payload as the native torch FP8 dtype for this GPU.""" + capability = get_device_compute_capability() + + # gfx950 uses OCP FP8. gfx942 and earlier ROCm architectures use FNUZ. + use_ocp_fp8 = capability == (9, 5) + + if dtype == tex.DType.kFloat8E4M3: + torch_dtype = ( + torch.float8_e4m3fn + if use_ocp_fp8 + else torch.float8_e4m3fnuz + ) + elif dtype == tex.DType.kFloat8E5M2: + torch_dtype = ( + torch.float8_e5m2 + if use_ocp_fp8 + else torch.float8_e5m2fnuz + ) + else: + raise TypeError(f"Unsupported TE FP8 dtype: {dtype}") + + return a.view(torch_dtype) + def _validate_common_epilogue( *, quantizer, @@ -54,6 +84,180 @@ def _is_mxfp8_operand(t): return hasattr(t, "_rowwise_data") and hasattr(t, "_rowwise_scale_inv") +def _is_fp8_operand(t): + """Return whether ``t`` is a regular TE tensor-wise FP8 operand.""" + try: + from transformer_engine.pytorch import Float8Tensor + from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import ( + Float8TensorStorage, + ) + except ImportError: + return False + + return isinstance(t, (Float8Tensor, Float8TensorStorage)) + + +def _reinterpret_fp8_payload(data, fp8_dtype, name): + """Reinterpret TE's uint8 payload using its ``tex.DType`` metadata.""" + if data is None: + raise RuntimeError(f"{name} does not contain the required FP8 payload") + + if fp8_dtype not in ( + tex.DType.kFloat8E4M3, + tex.DType.kFloat8E5M2, + ): + raise TypeError( + f"{name} has unsupported TE FP8 dtype metadata: {fp8_dtype}" + ) + + # TE stores Float8Tensor payloads as uint8. Use TE's shared conversion + # helper so ROCm's correct native torch FP8 type is selected from tex.DType. + if data.dtype == torch.uint8: + return reinterpret_as_fp8_tensor(data, fp8_dtype) + + # A materialized payload may already have been reinterpreted. Accept it + # only when its TE metadata is one of the recognized FP8 enum values. + if data.element_size() == 1 and data.dtype.is_floating_point: + return data + + raise TypeError( + f"{name} FP8 storage must be uint8 or an already reinterpreted " + f"one-byte floating-point tensor, got {data.dtype}" + ) + + +def _valid_fp8_transpose(t): + """Return whether a TE Float8 operand has usable columnwise storage.""" + return ( + hasattr(t, "_transpose") + and t._transpose is not None + and not getattr(t, "_transpose_invalid", False) + ) + + +def _run_fp8_tn(A, B, D): + """Run tensor-wise E4M3 x E4M3 FlyDSL FP8 for TE's TN convention. + + TE supplies: + A: weight [N, K], transa=True + B: activation [..., K], transb=False + + ``fp8_matmul`` consumes: + a: activation [M, K] + b: weight.T [K, N] + c: output [M, N] + """ + if not (_is_fp8_operand(A) and _is_fp8_operand(B)): + raise TypeError( + "FlyDSL FP8 GEMM expects Float8Tensor or Float8TensorStorage operands" + ) + + a_fp8_dtype = getattr(A, "_fp8_dtype", None) + b_fp8_dtype = getattr(B, "_fp8_dtype", None) + if ( + a_fp8_dtype != tex.DType.kFloat8E4M3 + or b_fp8_dtype != tex.DType.kFloat8E4M3 + ): + raise NotImplementedError( + "The current FlyDSL FP8 kernel supports only " + "tex.DType.kFloat8E4M3 x tex.DType.kFloat8E4M3; " + f"got A={a_fp8_dtype} and B={b_fp8_dtype}" + ) + + # A is transposed by the TE TN call. Prefer its already-materialized + # columnwise payload, which has the exact [K, N] layout consumed by + # fp8_matmul. Fall back to a transpose view of rowwise [N, K] storage. + if _valid_fp8_transpose(A): + A_t = _reinterpret_fp8_payload(A._transpose, a_fp8_dtype, "A._transpose") + if A_t.ndim != 2: + raise ValueError( + f"FlyDSL FP8 TN expects transposed weight storage to be rank 2, " + f"got {tuple(A_t.shape)}" + ) + k, n = A_t.shape + else: + A_data = _reinterpret_fp8_payload(getattr(A, "_data", None), a_fp8_dtype, "A._data") + if A_data.ndim != 2: + raise ValueError( + f"FlyDSL FP8 TN expects weight A to be rank 2, " + f"got {tuple(A_data.shape)}" + ) + n, k = A_data.shape + A_t = A_data.transpose(0, 1) + + # B is not transposed by TE, so rowwise storage is required. Flatten any + # leading activation dimensions into M while retaining the K dimension. + B_data = _reinterpret_fp8_payload(getattr(B, "_data", None), b_fp8_dtype, "B._data") + if B_data.ndim < 2: + raise ValueError( + f"FlyDSL FP8 TN expects activation B to have rank >= 2, " + f"got {tuple(B_data.shape)}" + ) + + B_flat = B_data.reshape(-1, B_data.shape[-1]) + m, kb = B_flat.shape + if kb != k: + raise ValueError( + f"FP8 inner dimensions do not match: weight K={k} and " + f"activation K={kb}" + ) + + A_scale_inv = getattr(A, "_scale_inv", None) + B_scale_inv = getattr(B, "_scale_inv", None) + for name, scale in ( + ("A._scale_inv", A_scale_inv), + ("B._scale_inv", B_scale_inv), + ): + if not isinstance(scale, torch.Tensor): + raise RuntimeError(f"{name} is not populated") + if scale.dtype != torch.float32 or scale.numel() != 1: + raise ValueError( + f"{name} must contain exactly one FP32 tensor-wise inverse " + f"scale, got dtype={scale.dtype}, shape={tuple(scale.shape)}" + ) + + output_shape = (*B_data.shape[:-1], n) + if D is None: + D = torch.empty( + output_shape, + dtype=torch.float16, + device=B_data.device, + ) + else: + if tuple(D.shape) != output_shape: + raise ValueError( + f"D shape {tuple(D.shape)} does not match expected {output_shape}" + ) + if D.dtype != torch.float16: + raise TypeError( + f"FlyDSL FP8 requires FP16 output, got {D.dtype}" + ) + if D.device != B_data.device: + raise ValueError( + f"D must be on {B_data.device}, got {D.device}" + ) + if not D.is_contiguous(): + raise ValueError( + "FlyDSL FP8 requires contiguous output storage" + ) + + if A_t.device != B_data.device: + raise ValueError( + f"A and B must be on the same device, got {A_t.device} " + f"and {B_data.device}" + ) + + fp8_matmul( + B_flat, + B_scale_inv, + A_t, + A_scale_inv, + D.view(m, n), + ) + + return D + + def _run_mxfp8_tn(A, B, D): """Run the existing FlyDSL MXFP8 TN path.""" A_data = A._rowwise_data @@ -297,6 +501,7 @@ def te_generic_gemm_flydsl( Currently supported: - MXFP8 TN input with FP16 output + - tensor-wise E4M3 x E4M3 FP8 TN input with FP16 output - BF16 TN input with BF16 output - FP16 TN input with FP16 output """ @@ -343,6 +548,24 @@ def te_generic_gemm_flydsl( D = _run_mxfp8_tn(A, B, D) return D, None, None, None + a_is_fp8 = _is_fp8_operand(A) + b_is_fp8 = _is_fp8_operand(B) + + if a_is_fp8 or b_is_fp8: + if not (a_is_fp8 and b_is_fp8): + raise ValueError( + "Mixed regular FP8 and non-FP8 FlyDSL GEMM inputs are not supported" + ) + + if output_dtype not in (None, tex.DType.kFloat16): + raise NotImplementedError( + "FlyDSL tensor-wise FP8 currently supports only FP16 output, " + f"got {output_dtype}" + ) + + D = _run_fp8_tn(A, B, D) + return D, None, None, None + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): raise TypeError( "Unsupported FlyDSL GEMM operand types: " @@ -370,6 +593,6 @@ def te_generic_gemm_flydsl( return D, None, None, None raise NotImplementedError( - "FlyDSL GEMM currently supports only MXFP8, BF16, or FP16 inputs; " + "FlyDSL GEMM currently supports only MXFP8, tensor-wise E4M3 FP8, BF16, or FP16 inputs; " f"got A={A.dtype} and B={B.dtype}" ) From 6ef76b3ec2f5430808fd7f57d20308705d1635ec Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 14:24:48 +0000 Subject: [PATCH 06/17] Add initial support for TN FP32 FlyDSL GEMM backend --- .../pytorch/cpp_extensions/gemm.py | 27 - .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 1079 +++++++++++++++++ .../flydsl_kernels/gemm/gemm_wrappers.py | 86 +- 3 files changed, 1164 insertions(+), 28 deletions(-) create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index df11bf277..147b4644b 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -460,37 +460,10 @@ def general_gemm( "beta": beta, } - # FlyDSL is currently an opt-in TN backend for: - # - MXFP8 - # - tensor-wise E4M3 x E4M3 FP8 - # - matching BF16 or FP16 inputs - # Keep every other datatype, recipe, and layout on the existing C++ path. - from ..tensor.storage.float8_tensor_storage import Float8TensorStorage - - is_mxfp8_gemm = ( - isinstance(A, MXFP8TensorStorage) - and isinstance(B, MXFP8TensorStorage) - ) - - is_fp8_gemm = ( - isinstance(A, Float8TensorStorage) - and isinstance(B, Float8TensorStorage) - and A._fp8_dtype == tex.DType.kFloat8E4M3 - and B._fp8_dtype == tex.DType.kFloat8E4M3 - ) - - is_fp16_bf16_gemm = ( - type(A) is torch.Tensor - and type(B) is torch.Tensor - and A.dtype == B.dtype - and A.dtype in (torch.bfloat16, torch.float16) - ) - use_gemm_flydsl = ( IS_HIP_EXTENSION and layout == "TN" and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0"))) - and (is_mxfp8_gemm or is_fp8_gemm or is_fp16_bf16_gemm) ) if use_gemm_flydsl: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py new file mode 100644 index 000000000..c18cde73c --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -0,0 +1,1079 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL FP32 4-wave GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K32 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes A and B as FP32 tensors shaped [M, K] and [N, K], and writes FP32 C +shaped [M, N]. The public ``fp32_matmul`` entry point accepts Transformer +Engine's TN contract and performs the required private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +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 + +# Transformer Engine-local FlyDSL utilities. +from .fp16_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_bf16_byte_buffer_tensor as make_fp32_byte_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 32 + +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +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 = 4 +VEC_BYTES = 16 + +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 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp32_pp_smem_a0" +LDS_SYM_A1 = "fp32_pp_smem_a1" +LDS_SYM_B0 = "fp32_pp_smem_b0" +LDS_SYM_B1 = "fp32_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 32 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + +def make_fp32_inputs(M, N, K, device="cuda"): + """Generate FP32 A[M,K] and B[N,K] inputs.""" + A = (torch.randn(M, K, device=device) * 0.5).to(torch.float32) + B = (torch.randn(N, K, device=device) * 0.5).to(torch.float32) + return A, B + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +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] + ) + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + + +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 _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + 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), + ) + + +def _compile_kernel(K: int, use_xcd_remap: bool = True): + """Build the specialized 4-wave kernel for compile-time ``K``. + + ``K`` must contain at least four K32 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + NUM_THREADS = 256 + WARP_SIZE = 64 + + 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 = 4 + VEC_BYTES = 16 + + 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 + + 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} K32 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LDS_BYTES_HALF = LDS_ELEMS_HALF * ELEM_BYTES + LOAD_PASSES_HALF = LDS_BYTES_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x64 FP32 page is two independent 128x64 half-pages. + # Store LDS as bytes so BufferCopyLDS128b sees i8 on both source and + # destination. Each half-page remains exactly 16 KiB. + a0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + 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) + + # A/B arrive as contiguous uint8 byte views. Keeping staging byte-addressed + # preserves the original 16-byte G2L instruction cadence and vmcnt values. + gA = make_fp32_byte_buffer_tensor(A) + gB = make_fp32_byte_buffer_tensor(B) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + 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) + + # 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 * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + gl_off_b = compute_global_swizzle(lane, wave_id, K * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, fx.Uint8.ir_type, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, fx.Uint8.ir_type, 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) + 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(4) # C is FP32. + 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 _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) + + 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") + 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) + + def hot_loop_scheduler_q_refill_2n(): + # Eight refill VMEM operations overlap four independent 8-MFMA + # groups (two K16 halves x two N-halves). + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_mfma(32) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # Eight refill VMEM operations and eight distributed A-bottom LDS + # reads overlap four independent 8-MFMA K32 groups. + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_dsrd(2) + rocdl.sched_mfma(32) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # Eight two-read prefetch groups overlap four complete-quadrant + # 16-MFMA groups (two K16 halves for each of Q2 and Q3). + for _ in range_constexpr(4): + rocdl.sched_dsrd(4) + rocdl.sched_mfma(64) + rocdl.sched_barrier(0) + + 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 + # 128x64 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 * ELEM_BYTES) + k_base * fx.Index(ELEM_BYTES) + 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 * ELEM_BYTES) + k_base * fx.Index(ELEM_BYTES) + 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_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 16-byte half of the wave operand tile. + # 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_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, 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 * ELEM_BYTES)) + + 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 _fp32_k4_operand(full_frag, k32_half, k4): + # A/B 16x32 FP32 wave fragments are i32x8: one FP32 value per + # VGPR and eight K4 MFMA steps per logical K32 tile. Keep the + # existing two-half schedule by grouping four K4 steps per half. + return Vec(full_frag)[k32_half * 4 + k4] + + def _pinned_fp32_mfma_once(acc_idx, a_k4, b_k4): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [arith._to_raw(a_k4), arith._to_raw(b_k4)], + ( + f"v_mfma_f32_16x16x4_f32 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Accumulate one logical 16x16x32 FP32 product into pinned AGPRs.""" + for k32_half in range_constexpr(2): + for k4 in range_constexpr(4): + _pinned_fp32_mfma_once( + acc_idx, + _fp32_k4_operand(a_frag, k32_half, k4), + _fp32_k4_operand(b_frag, k32_half, k4), + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + # The final logical K32 update is eight in-place K4 FP32 MFMAs. + assert dst_slot == old_acc_idx + pinned_mfma(old_acc_idx, a_frag, b_frag) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + def mfma_2n_4mi_k32(subtile_id, n_base, k32_half, a0, a1, a2, a3, b0, b1): + """Issue four K4 steps (one K16 half) for a 4x2 accumulator slab.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1) + for k4 in range_constexpr(4): + for mi in range_constexpr(4): + a_k4 = _fp32_k4_operand(a_frags[mi], k32_half, k4) + for nj in range_constexpr(2): + _pinned_fp32_mfma_once( + _acc_idx(subtile_id, mi, n_base + nj), + a_k4, + _fp32_k4_operand(b_frags[nj], k32_half, k4), + ) + + def mfma_4n_4mi_k32(subtile_id, k32_half, a0, a1, a2, a3, b0, b1, b2, b3): + """Issue four K4 steps (one K16 half) for a complete 4x4 quadrant.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1, b2, b3) + for k4 in range_constexpr(4): + for mi in range_constexpr(4): + a_k4 = _fp32_k4_operand(a_frags[mi], k32_half, k4) + for ni in range_constexpr(4): + _pinned_fp32_mfma_once( + _acc_idx(subtile_id, mi, ni), + a_k4, + _fp32_k4_operand(b_frags[ni], k32_half, k4), + ) + + 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_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], 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, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + 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 * ELEM_BYTES) + return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): + + # 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) + + # 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 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + # 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. Compute is K16-half-major across all 16 + # independent accumulators, eliminating the two-deep same-AGPR + # dependency chains produced by pinned_mfma(). + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n_4mi_k32(0, 0, 0, a00, a01, a02, a03, b00, b01) + + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n_4mi_k32(0, 2, 0, a00, a01, a02, a03, b02, b03) + + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + # K32 slice 0 already covers K[0:16]. + + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + # Keep this refill/LDS-read slot compute-free. + + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n_4mi_k32(0, 0, 1, a00, a01, a02, a03, b00, b01) + + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n_4mi_k32(0, 2, 1, a00, a01, a02, a03, b02, b03) + + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + # K32 slice 1 already covers K[16:32]. + + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + # Keep this refill/LDS-read slot compute-free. + + 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 = 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) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n_4mi_k32(1, 0, 0, a00, a01, a02, a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + # K32 slice 0 already covers K[0:16]. + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + # Keep this refill slot compute-free. + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n_4mi_k32(1, 0, 1, a00, a01, a02, a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + # K32 slice 1 already covers K[16:32]. + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + # Keep this refill slot compute-free. + 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, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n_4mi_k32(2, 0, a10, a11, a12, a13, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + # K32 slice 0 already covers K[0:16]. + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n_4mi_k32(2, 1, a10, a11, a12, a13, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + # K32 slice 1 already covers K[16:32]. + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n_4mi_k32(3, 0, a10, a11, a12, a13, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + # K32 slice 0 already covers K[0:16]. + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n_4mi_k32(3, 1, a10, a11, a12, a13, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + # K32 slice 1 already covers K[16:32]. + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # 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], + ) + 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. + 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) + 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) + rocdl.sched_barrier(0) + + a0_regs = load_a_subtile_regs(lds_a0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K32 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + else: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # 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, + 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), stream=stream) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch(K: int, use_xcd_remap: bool = True): + return _compile_kernel(K, use_xcd_remap=use_xcd_remap) + + + +def fp32_matmul( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing TN FP32 GEMM adapter. + + Public/backend contract: + a: [M, K] FP32 + b: [K, N] FP32 + c: [M, N] FP32 output + + The optimized core streams both operands with K contiguous and therefore + privately consumes B as [N, K]. In the normal TE TN path, ``b`` is a + transpose view of contiguous rowwise weight storage, so ``b.T`` is already + contiguous and does not require a physical transpose. + """ + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL FP32 TN expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + m, k = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + if a.dtype != torch.float32 or b.dtype != torch.float32: + raise TypeError( + "FlyDSL FP32 GEMM expects both operands to have torch.float32 dtype, " + f"got {a.dtype} and {b.dtype}" + ) + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype != torch.float32: + raise TypeError( + f"The current FlyDSL FP32 kernel stores torch.float32 output, got {c.dtype}" + ) + if a.device != b.device or a.device != c.device: + raise ValueError( + f"A, B, and C must be on the same device, got " + f"{a.device}, {b.device}, and {c.device}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL FP32 GEMM requires contiguous output storage") + + b_hk = b.transpose(0, 1).contiguous() + doGemm(a, b_hk, c, stream=stream) + + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch the private K-specialized FP32 core. + + A and B are shaped [M, K] and [N, K]; C is shaped [M, N]. M and N + remain runtime values, while K selects the cached compile-time specialization. + """ + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + assert A.dtype == torch.float32 and B.dtype == torch.float32 + assert C.dtype == torch.float32 + assert M_runtime % _BLOCK_M == 0, ( + f"M={M_runtime} must be a multiple of {_BLOCK_M}" + ) + assert N_runtime % _BLOCK_N == 0, ( + f"N={N_runtime} must be a multiple of {_BLOCK_N}" + ) + assert K_runtime % _BLOCK_K == 0, ( + f"K={K_runtime} must be a multiple of {_BLOCK_K}" + ) + num_k_tiles = K_runtime // _BLOCK_K + assert num_k_tiles >= 4, ( + f"K={K_runtime} gives {num_k_tiles} K32 tiles; need at least 4" + ) + assert C.shape == (M_runtime, N_runtime) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.contiguous().view(torch.uint8).view(-1) + B_arg = B.contiguous().view(torch.uint8).view(-1) + C_arg = C.view(-1) + launch = _cached_launch(int(K_runtime), bool(use_xcd_remap)) + launch(A_arg, B_arg, C_arg, M_runtime, N_runtime, stream=stream) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 8245eeef5..c04e12d90 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -11,6 +11,7 @@ from .bf16_gemm import bf16_matmul from .fp16_gemm import fp16_matmul +from .fp32_gemm import fp32_matmul from .fp8_gemm import fp8_matmul from .mxfp8_gemm import mxfp8_matmul @@ -473,6 +474,78 @@ def _run_fp16_tn(A, B, D): return D + +def _run_fp32_tn(A, B, D): + """Run FlyDSL FP32 for TE's TN operand convention.""" + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError( + "FlyDSL FP32 GEMM expects plain torch.Tensor operands" + ) + + if A.dtype != torch.float32 or B.dtype != torch.float32: + raise TypeError( + "FlyDSL FP32 GEMM requires FP32 inputs, " + f"got A={A.dtype} and B={B.dtype}" + ) + + if A.ndim != 2: + raise ValueError( + f"FlyDSL FP32 TN expects weight A to be rank 2, got {tuple(A.shape)}" + ) + if B.ndim < 2: + raise ValueError( + f"FlyDSL FP32 TN expects activation B to have rank >= 2, got {tuple(B.shape)}" + ) + + n, k = A.shape + B_flat = B.reshape(-1, B.shape[-1]) + m, kb = B_flat.shape + + if kb != k: + raise ValueError( + f"FP32 inner dimensions do not match: A{tuple(A.shape)} and " + f"B{tuple(B.shape)}" + ) + + output_shape = (*B.shape[:-1], n) + + if D is None: + D = torch.empty( + output_shape, + dtype=torch.float32, + device=B.device, + ) + else: + if tuple(D.shape) != output_shape: + raise ValueError( + f"D shape {tuple(D.shape)} does not match expected {output_shape}" + ) + if D.dtype != torch.float32: + raise TypeError( + f"FlyDSL FP32 requires FP32 output, got {D.dtype}" + ) + if D.device != B.device: + raise ValueError( + f"D must be on {B.device}, got {D.device}" + ) + if not D.is_contiguous(): + raise ValueError( + "FlyDSL FP32 requires contiguous output storage" + ) + + if A.device != B.device: + raise ValueError( + f"A and B must be on the same device, got {A.device} and {B.device}" + ) + + fp32_matmul( + B_flat, + A.transpose(0, 1), + D.view(m, n), + ) + + return D + def te_generic_gemm_flydsl( A, transa, @@ -504,6 +577,7 @@ def te_generic_gemm_flydsl( - tensor-wise E4M3 x E4M3 FP8 TN input with FP16 output - BF16 TN input with BF16 output - FP16 TN input with FP16 output + - FP32 TN input with FP32 output """ del bias_type del gelu_in @@ -592,7 +666,17 @@ def te_generic_gemm_flydsl( D = _run_fp16_tn(A, B, D) return D, None, None, None + if A.dtype == torch.float32 and B.dtype == torch.float32: + if output_dtype not in (None, tex.DType.kFloat32): + raise NotImplementedError( + "FlyDSL FP32 currently supports only FP32 output, " + f"got {output_dtype}" + ) + + D = _run_fp32_tn(A, B, D) + return D, None, None, None + raise NotImplementedError( - "FlyDSL GEMM currently supports only MXFP8, tensor-wise E4M3 FP8, BF16, or FP16 inputs; " + "FlyDSL GEMM currently supports only MXFP8, tensor-wise E4M3 FP8, BF16, FP16, or FP32 inputs; " f"got A={A.dtype} and B={B.dtype}" ) From 9a9462e4574f1fcabd34c2f241a4689580f0d80b Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 16:15:02 +0000 Subject: [PATCH 07/17] add layout support for NN/NT FlyDSL GEMM --- .../pytorch/cpp_extensions/gemm.py | 6 +- .../flydsl_kernels/gemm/gemm_wrappers.py | 807 ++++++++++-------- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 272 +++--- 3 files changed, 615 insertions(+), 470 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 147b4644b..46f0b2ee9 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -460,11 +460,7 @@ def general_gemm( "beta": beta, } - use_gemm_flydsl = ( - IS_HIP_EXTENSION - and layout == "TN" - and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0"))) - ) + use_gemm_flydsl = IS_HIP_EXTENSION and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0"))) if use_gemm_flydsl: # Lazy import keeps FlyDSL off the normal Transformer Engine import path. diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index c04e12d90..3dd8a648f 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -4,6 +4,8 @@ """TE entry points for the FlyDSL GEMM backend.""" +import os + import torch import transformer_engine_torch as tex @@ -80,22 +82,43 @@ def _validate_common_epilogue( ) -def _is_mxfp8_operand(t): - """Return whether ``t`` exposes TE MXFP8 rowwise storage.""" - return hasattr(t, "_rowwise_data") and hasattr(t, "_rowwise_scale_inv") - - -def _is_fp8_operand(t): - """Return whether ``t`` is a regular TE tensor-wise FP8 operand.""" +def _classify_input(t): + """Classify a GEMM operand for the FlyDSL backend.""" try: - from transformer_engine.pytorch import Float8Tensor + from transformer_engine.pytorch.float8_tensor import Float8Tensor from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import ( Float8TensorStorage, ) + if isinstance(t, (Float8Tensor, Float8TensorStorage)): + return "fp8", t + except ImportError: + pass + + try: + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor + from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import ( + MXFP8TensorStorage, + ) + if isinstance(t, (MXFP8Tensor, MXFP8TensorStorage)): + return "mxfp8", t + except ImportError: + pass + + try: + from transformer_engine.pytorch.quantized_tensor import ( + QuantizedTensorStorage, + ) + if isinstance(t, QuantizedTensorStorage): + raise ValueError( + f"The FlyDSL GEMM backend does not support " + f"{type(t).__name__}. Only Float8Tensor / " + f"Float8TensorStorage and MXFP8Tensor / " + f"MXFP8TensorStorage are implemented." + ) except ImportError: - return False + pass - return isinstance(t, (Float8Tensor, Float8TensorStorage)) + return "regular", None def _reinterpret_fp8_payload(data, fp8_dtype, name): @@ -136,416 +159,432 @@ def _valid_fp8_transpose(t): ) -def _run_fp8_tn(A, B, D): - """Run tensor-wise E4M3 x E4M3 FlyDSL FP8 for TE's TN convention. - TE supplies: - A: weight [N, K], transa=True - B: activation [..., K], transb=False +def _mxfp8_debug_enabled() -> bool: + value = os.getenv("DEBUG_FLYDSL_MXFP8_GEMM", "") + return value.lower() not in ("", "0", "false", "no", "off") - ``fp8_matmul`` consumes: - a: activation [M, K] - b: weight.T [K, N] - c: output [M, N] - """ - if not (_is_fp8_operand(A) and _is_fp8_operand(B)): - raise TypeError( - "FlyDSL FP8 GEMM expects Float8Tensor or Float8TensorStorage operands" - ) - a_fp8_dtype = getattr(A, "_fp8_dtype", None) - b_fp8_dtype = getattr(B, "_fp8_dtype", None) - if ( - a_fp8_dtype != tex.DType.kFloat8E4M3 - or b_fp8_dtype != tex.DType.kFloat8E4M3 - ): - raise NotImplementedError( - "The current FlyDSL FP8 kernel supports only " - "tex.DType.kFloat8E4M3 x tex.DType.kFloat8E4M3; " - f"got A={a_fp8_dtype} and B={b_fp8_dtype}" - ) - - # A is transposed by the TE TN call. Prefer its already-materialized - # columnwise payload, which has the exact [K, N] layout consumed by - # fp8_matmul. Fall back to a transpose view of rowwise [N, K] storage. - if _valid_fp8_transpose(A): - A_t = _reinterpret_fp8_payload(A._transpose, a_fp8_dtype, "A._transpose") - if A_t.ndim != 2: - raise ValueError( - f"FlyDSL FP8 TN expects transposed weight storage to be rank 2, " - f"got {tuple(A_t.shape)}" - ) - k, n = A_t.shape - else: - A_data = _reinterpret_fp8_payload(getattr(A, "_data", None), a_fp8_dtype, "A._data") - if A_data.ndim != 2: - raise ValueError( - f"FlyDSL FP8 TN expects weight A to be rank 2, " - f"got {tuple(A_data.shape)}" - ) - n, k = A_data.shape - A_t = A_data.transpose(0, 1) - - # B is not transposed by TE, so rowwise storage is required. Flatten any - # leading activation dimensions into M while retaining the K dimension. - B_data = _reinterpret_fp8_payload(getattr(B, "_data", None), b_fp8_dtype, "B._data") - if B_data.ndim < 2: - raise ValueError( - f"FlyDSL FP8 TN expects activation B to have rank >= 2, " - f"got {tuple(B_data.shape)}" - ) +def _mxfp8_debug(message: str) -> None: + if _mxfp8_debug_enabled(): + print(f"[DEBUG_FLYDSL_MXFP8_GEMM] {message}") - B_flat = B_data.reshape(-1, B_data.shape[-1]) - m, kb = B_flat.shape - if kb != k: - raise ValueError( - f"FP8 inner dimensions do not match: weight K={k} and " - f"activation K={kb}" - ) - A_scale_inv = getattr(A, "_scale_inv", None) - B_scale_inv = getattr(B, "_scale_inv", None) - for name, scale in ( - ("A._scale_inv", A_scale_inv), - ("B._scale_inv", B_scale_inv), - ): - if not isinstance(scale, torch.Tensor): - raise RuntimeError(f"{name} is not populated") - if scale.dtype != torch.float32 or scale.numel() != 1: - raise ValueError( - f"{name} must contain exactly one FP32 tensor-wise inverse " - f"scale, got dtype={scale.dtype}, shape={tuple(scale.shape)}" - ) +def _canonicalize_blas_pair( + A_data: torch.Tensor, + transa: bool, + B_data: torch.Tensor, + transb: bool, +): + """Swap TE BLAS operands and apply their original transpose flags.""" + a_flydsl = B_data.transpose(0, 1) if transb else B_data + b_flydsl = A_data.transpose(0, 1) if transa else A_data + return a_flydsl, b_flydsl - output_shape = (*B_data.shape[:-1], n) - if D is None: - D = torch.empty( - output_shape, - dtype=torch.float16, - device=B_data.device, - ) - else: - if tuple(D.shape) != output_shape: - raise ValueError( - f"D shape {tuple(D.shape)} does not match expected {output_shape}" - ) - if D.dtype != torch.float16: - raise TypeError( - f"FlyDSL FP8 requires FP16 output, got {D.dtype}" - ) - if D.device != B_data.device: - raise ValueError( - f"D must be on {B_data.device}, got {D.device}" - ) - if not D.is_contiguous(): - raise ValueError( - "FlyDSL FP8 requires contiguous output storage" - ) - if A_t.device != B_data.device: +def _flatten_rowwise(t: torch.Tensor, name: str) -> torch.Tensor: + """Flatten all leading dimensions while preserving the final dimension.""" + if t.ndim < 2: raise ValueError( - f"A and B must be on the same device, got {A_t.device} " - f"and {B_data.device}" + f"FlyDSL GEMM expects {name} to have rank >= 2, got {tuple(t.shape)}" ) + return t.reshape(-1, t.shape[-1]) - fp8_matmul( - B_flat, - B_scale_inv, - A_t, - A_scale_inv, - D.view(m, n), - ) - return D - - -def _run_mxfp8_tn(A, B, D): - """Run the existing FlyDSL MXFP8 TN path.""" - A_data = A._rowwise_data - A_scale = A._rowwise_scale_inv - B_data = B._rowwise_data - B_scale = B._rowwise_scale_inv - - if A_data is None or A_scale is None: - raise RuntimeError("A does not contain rowwise MXFP8 data and scales") - - if B_data is None or B_scale is None: - raise RuntimeError("B does not contain rowwise MXFP8 data and scales") +def _canonicalize_blas_operands( + A_data: torch.Tensor, + transa: bool, + B_data: torch.Tensor, + transb: bool, +): + """Convert TE's BLAS-shaped operands to FlyDSL row-major operands. - n, k = A_data.shape - B_flat = B_data.reshape(-1, B_data.shape[-1]) - m, kb = B_flat.shape + TE's generic GEMM interface follows BLAS column-major interpretation. + FlyDSL kernels consume ordinary row-major matrices: - if kb != k: - raise ValueError(f"MXFP8 inner dimensions do not match: {k} and {kb}") + a_flydsl: [M, K] + b_flydsl: [K, N] - A_scale = A_scale.reshape(n, -1) - B_scale = B_scale.reshape(m, -1) - output_shape = (*B_data.shape[:-1], n) + The standard conversion is to swap A/B and apply the original transpose + flags to the swapped operands: - if D is None: - D = torch.empty( - output_shape, - dtype=torch.float16, - device=B_data.device, + a_flydsl = op(B) + b_flydsl = op(A) + """ + if transa and transb: + raise NotImplementedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" ) - else: - if tuple(D.shape) != output_shape: - raise ValueError( - f"D shape {tuple(D.shape)} does not match expected {output_shape}" - ) - if D.dtype != torch.float16: - raise TypeError( - f"FlyDSL MXFP8 requires FP16 output, got {D.dtype}" - ) - if not D.is_contiguous(): - raise ValueError( - "FlyDSL MXFP8 requires contiguous output storage" - ) - # Public mxfp8_matmul contract: - # a: [M, K] - # a_scale: [M, K/32] - # b: [K, N] - # b_scale: [N, K/32] - # c: [M, N] FP16 - mxfp8_matmul( + A_flat = _flatten_rowwise(A_data, "A") + B_flat = _flatten_rowwise(B_data, "B") + + a_flydsl, b_flydsl = _canonicalize_blas_pair( + A_flat, + transa, B_flat, - B_scale, - A_data.transpose(0, 1), - A_scale, - D.view(m, n), + transb, ) - return D + m, k = a_flydsl.shape + kb, n = b_flydsl.shape + if kb != k: + layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" + raise ValueError( + f"FlyDSL {layout} canonicalization produced incompatible operands: " + f"{tuple(a_flydsl.shape)} @ {tuple(b_flydsl.shape)}" + ) + return a_flydsl, b_flydsl, m, n, k -def _run_bf16_tn(A, B, D): - """Run FlyDSL BF16 for TE's TN operand convention. - TE supplies: - A: weight [N, K] - B: activation [..., K] +def _validate_or_allocate_output( + D, + *, + shape, + dtype, + device, + backend_name, +): + if D is None: + return torch.empty(shape, dtype=dtype, device=device) - ``bf16_matmul`` consumes: - a: activation [M, K] - b: weight.T [K, N] - c: output [M, N] - """ - if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): - raise TypeError( - "FlyDSL BF16 GEMM expects plain torch.Tensor operands" + if tuple(D.shape) != tuple(shape): + raise ValueError( + f"D shape {tuple(D.shape)} does not match expected {tuple(shape)}" ) - - if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: + if D.dtype != dtype: raise TypeError( - "FlyDSL BF16 GEMM requires BF16 inputs, " - f"got A={A.dtype} and B={B.dtype}" + f"FlyDSL {backend_name} requires {dtype} output, got {D.dtype}" ) - - if A.ndim != 2: + if D.device != device: raise ValueError( - f"FlyDSL BF16 TN expects weight A to be rank 2, got {tuple(A.shape)}" + f"D must be on {device}, got {D.device}" ) - if B.ndim < 2: + if not D.is_contiguous(): raise ValueError( - f"FlyDSL BF16 TN expects activation B to have rank >= 2, got {tuple(B.shape)}" + f"FlyDSL {backend_name} requires contiguous output storage" ) + return D - n, k = A.shape - B_flat = B.reshape(-1, B.shape[-1]) - m, kb = B_flat.shape - if kb != k: - raise ValueError( - f"BF16 inner dimensions do not match: A{tuple(A.shape)} and " - f"B{tuple(B.shape)}" +def _run_regular_gemm( + A, + transa, + B, + transb, + D, + *, + dtype, + matmul, + backend_name, +): + """Run FP16/BF16/FP32 through shared TN/NN/NT shape handling.""" + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError( + f"FlyDSL {backend_name} GEMM expects plain torch.Tensor operands" ) - - output_shape = (*B.shape[:-1], n) - - if D is None: - D = torch.empty( - output_shape, - dtype=torch.bfloat16, - device=B.device, + if A.dtype != dtype or B.dtype != dtype: + raise TypeError( + f"FlyDSL {backend_name} GEMM requires {dtype} inputs, " + f"got A={A.dtype} and B={B.dtype}" ) - else: - if tuple(D.shape) != output_shape: - raise ValueError( - f"D shape {tuple(D.shape)} does not match expected {output_shape}" - ) - if D.dtype != torch.bfloat16: - raise TypeError( - f"FlyDSL BF16 requires BF16 output, got {D.dtype}" - ) - if D.device != B.device: - raise ValueError( - f"D must be on {B.device}, got {D.device}" - ) - if not D.is_contiguous(): - raise ValueError( - "FlyDSL BF16 requires contiguous output storage" - ) - if A.device != B.device: raise ValueError( f"A and B must be on the same device, got {A.device} and {B.device}" ) - bf16_matmul( - B_flat, - A.transpose(0, 1), - D.view(m, n), + a_flydsl, b_flydsl, m, n, _ = _canonicalize_blas_operands( + A, transa, B, transb ) + D = _validate_or_allocate_output( + D, + shape=(m, n), + dtype=dtype, + device=A.device, + backend_name=backend_name, + ) + + matmul( + a_flydsl, + b_flydsl, + D.view(m, n), + ) return D -def _run_fp16_tn(A, B, D): - """Run FlyDSL FP16 for TE's TN operand convention.""" - if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): - raise TypeError( - "FlyDSL FP16 GEMM expects plain torch.Tensor operands" - ) +def _get_fp8_logical_rowwise_payload(t, name): + """Return logical rowwise FP8 data, matching the Triton wrapper. - if A.dtype != torch.float16 or B.dtype != torch.float16: - raise TypeError( - "FlyDSL FP16 GEMM requires FP16 inputs, " - f"got A={A.dtype} and B={B.dtype}" - ) + Prefer TE's rowwise ``_data``. If only valid columnwise ``_transpose`` + storage exists, materialize a rowwise copy once for canonicalization. + """ + fp8_dtype = getattr(t, "_fp8_dtype", None) + data = getattr(t, "_data", None) - if A.ndim != 2: - raise ValueError( - f"FlyDSL FP16 TN expects weight A to be rank 2, got {tuple(A.shape)}" + if data is not None: + return _reinterpret_fp8_payload( + data, + fp8_dtype, + f"{name}._data", ) - if B.ndim < 2: - raise ValueError( - f"FlyDSL FP16 TN expects activation B to have rank >= 2, got {tuple(B.shape)}" + + if not _valid_fp8_transpose(t): + raise RuntimeError( + f"{name} has neither valid rowwise (_data) nor " + f"columnwise (_transpose) FP8 storage" ) - n, k = A.shape - B_flat = B.reshape(-1, B.shape[-1]) - m, kb = B_flat.shape + transpose_data = _reinterpret_fp8_payload( + t._transpose, + fp8_dtype, + f"{name}._transpose", + ) - if kb != k: + if transpose_data.ndim < 2: raise ValueError( - f"FP16 inner dimensions do not match: A{tuple(A.shape)} and " - f"B{tuple(B.shape)}" + f"{name}._transpose must have rank >= 2, " + f"got {tuple(transpose_data.shape)}" ) - output_shape = (*B.shape[:-1], n) + # TE's columnwise payload represents the transpose of the logical rowwise + # tensor. Materialize rowwise storage before applying BLAS transpose flags, + # exactly as the Triton wrapper's materialize_rowwise_from_columnwise path. + return transpose_data.transpose(-2, -1).contiguous() - if D is None: - D = torch.empty( - output_shape, - dtype=torch.float16, - device=B.device, - ) + +def _select_mxfp8_data_and_scale( + t, + *, + will_transpose: bool, + name: str, +): + """Select the TE MXFP8 representation required by BLAS semantics.""" + if will_transpose: + data = getattr(t, "_columnwise_data", None) + scale = getattr(t, "_columnwise_scale_inv", None) + orientation = "columnwise" else: - if tuple(D.shape) != output_shape: - raise ValueError( - f"D shape {tuple(D.shape)} does not match expected {output_shape}" - ) - if D.dtype != torch.float16: - raise TypeError( - f"FlyDSL FP16 requires FP16 output, got {D.dtype}" - ) - if D.device != B.device: - raise ValueError( - f"D must be on {B.device}, got {D.device}" - ) - if not D.is_contiguous(): - raise ValueError( - "FlyDSL FP16 requires contiguous output storage" - ) + data = getattr(t, "_rowwise_data", None) + scale = getattr(t, "_rowwise_scale_inv", None) + orientation = "rowwise" + + _mxfp8_debug( + f"{name}: will_transpose={will_transpose}, " + f"selected={orientation}, data_present={data is not None}, " + f"scale_present={scale is not None}" + ) - if A.device != B.device: - raise ValueError( - f"A and B must be on the same device, got {A.device} and {B.device}" + if data is None or scale is None: + raise RuntimeError( + f"{name} does not contain required {orientation} MXFP8 data and scales" ) - - fp16_matmul( - B_flat, - A.transpose(0, 1), - D.view(m, n), + + _mxfp8_debug( + f"{name} selected data shape={tuple(data.shape)}, " + f"dtype={data.dtype}, stride={tuple(data.stride())}; " + f"scale shape={tuple(scale.shape)}, dtype={scale.dtype}, " + f"stride={tuple(scale.stride())}" ) + return data, scale - return D +def _flatten_mxfp8_scale(t: torch.Tensor, name: str) -> torch.Tensor: + if t.ndim < 2: + raise ValueError( + f"FlyDSL MXFP8 expects {name} scale rank >= 2, " + f"got {tuple(t.shape)}" + ) + original_shape = tuple(t.shape) + if t.ndim > 2: + t = t.reshape(-1, t.shape[-1]) + _mxfp8_debug( + f"{name} scale flatten: {original_shape} -> {tuple(t.shape)}, " + f"contiguous={t.is_contiguous()}" + ) + return t -def _run_fp32_tn(A, B, D): - """Run FlyDSL FP32 for TE's TN operand convention.""" - if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): - raise TypeError( - "FlyDSL FP32 GEMM expects plain torch.Tensor operands" - ) +def _run_mxfp8( + A, + transa, + B, + transb, + D, +): + """Canonicalize TE MXFP8 operands, then launch the fused backend.""" + layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" + _mxfp8_debug( + f"entry: layout={layout}, A_type={type(A).__name__}, " + f"B_type={type(B).__name__}, D_provided={D is not None}" + ) - if A.dtype != torch.float32 or B.dtype != torch.float32: - raise TypeError( - "FlyDSL FP32 GEMM requires FP32 inputs, " - f"got A={A.dtype} and B={B.dtype}" - ) + # Match TE CanonicalizeGemmInput / Triton data_and_scale_for_transpose: + # A: transa=True -> rowwise, transa=False -> columnwise + # B: transb=True -> columnwise, transb=False -> rowwise + A_data, A_scale = _select_mxfp8_data_and_scale( + A, + will_transpose=not transa, + name="A", + ) + B_data, B_scale = _select_mxfp8_data_and_scale( + B, + will_transpose=transb, + name="B", + ) + + a_flydsl, b_flydsl, m, n, k = _canonicalize_blas_operands( + A_data, + transa, + B_data, + transb, + ) - if A.ndim != 2: + A_scale = _flatten_mxfp8_scale(A_scale, "A") + B_scale = _flatten_mxfp8_scale(B_scale, "B") + a_scale, b_scale = _canonicalize_blas_pair( + A_scale, + transa, + B_scale, + transb, + ) + + _mxfp8_debug( + f"canonicalized layout={layout}: " + f"a={tuple(a_flydsl.shape)}, stride={tuple(a_flydsl.stride())}; " + f"b={tuple(b_flydsl.shape)}, stride={tuple(b_flydsl.stride())}" + ) + _mxfp8_debug( + f"canonicalized scales: " + f"a_scale={tuple(a_scale.shape)}, stride={tuple(a_scale.stride())}; " + f"b_scale={tuple(b_scale.shape)}, stride={tuple(b_scale.stride())}" + ) + _mxfp8_debug(f"derived GEMM dimensions: M={m}, N={n}, K={k}") + + if a_flydsl.device != b_flydsl.device: raise ValueError( - f"FlyDSL FP32 TN expects weight A to be rank 2, got {tuple(A.shape)}" + f"A and B must be on the same device, got " + f"{a_flydsl.device} and {b_flydsl.device}" ) - if B.ndim < 2: + + scale_group_size = 32 + if k % scale_group_size != 0: raise ValueError( - f"FlyDSL FP32 TN expects activation B to have rank >= 2, got {tuple(B.shape)}" + f"K={k} must be divisible by MXFP8 scale group size " + f"{scale_group_size}" ) - n, k = A.shape - B_flat = B.reshape(-1, B.shape[-1]) - m, kb = B_flat.shape - - if kb != k: + # Shared BLAS canonicalization yields: + # a_scale [M, K/32] + # b_scale [K/32, N] + expected_a_scale = (m, k // scale_group_size) + expected_b_scale = (k // scale_group_size, n) + if tuple(a_scale.shape) != expected_a_scale: raise ValueError( - f"FP32 inner dimensions do not match: A{tuple(A.shape)} and " - f"B{tuple(B.shape)}" + f"A scale shape {tuple(a_scale.shape)} != expected " + f"{expected_a_scale}" ) + if tuple(b_scale.shape) != expected_b_scale: + raise ValueError( + f"B scale shape {tuple(b_scale.shape)} != expected " + f"{expected_b_scale}" + ) + if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + raise TypeError("FlyDSL MXFP8 expects raw E8M0 scales as torch.uint8") + + D = _validate_or_allocate_output( + D, + shape=(m, n), + dtype=torch.float16, + device=a_flydsl.device, + backend_name="MXFP8", + ) + + return mxfp8_matmul( + a_flydsl, + a_scale, + b_flydsl, + b_scale, + D.view(m, n), + ) - output_shape = (*B.shape[:-1], n) - if D is None: - D = torch.empty( - output_shape, - dtype=torch.float32, - device=B.device, +def _run_fp8( + A, + transa, + B, + transb, + D, +): + """Run tensor-wise E4M3 x E4M3 FP8 for TN/NN/NT.""" + a_fp8_dtype = getattr(A, "_fp8_dtype", None) + b_fp8_dtype = getattr(B, "_fp8_dtype", None) + if ( + a_fp8_dtype != tex.DType.kFloat8E4M3 + or b_fp8_dtype != tex.DType.kFloat8E4M3 + ): + raise NotImplementedError( + "The current FlyDSL FP8 kernel supports only " + "tex.DType.kFloat8E4M3 x tex.DType.kFloat8E4M3; " + f"got A={a_fp8_dtype} and B={b_fp8_dtype}" ) - else: - if tuple(D.shape) != output_shape: - raise ValueError( - f"D shape {tuple(D.shape)} does not match expected {output_shape}" - ) - if D.dtype != torch.float32: - raise TypeError( - f"FlyDSL FP32 requires FP32 output, got {D.dtype}" - ) - if D.device != B.device: - raise ValueError( - f"D must be on {B.device}, got {D.device}" - ) - if not D.is_contiguous(): + + if transa and transb: + raise NotImplementedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) + + # Match Triton's regular-FP8 handling: establish logical rowwise + # payloads first, then apply the same shared BLAS-to-row-major + # canonicalization used for FP16/BF16/FP32. + A_data = _get_fp8_logical_rowwise_payload(A, "A") + B_data = _get_fp8_logical_rowwise_payload(B, "B") + + A_scale_inv = getattr(A, "_scale_inv", None) + B_scale_inv = getattr(B, "_scale_inv", None) + for name, scale in ( + ("A._scale_inv", A_scale_inv), + ("B._scale_inv", B_scale_inv), + ): + if not isinstance(scale, torch.Tensor): + raise RuntimeError(f"{name} is not populated") + if scale.dtype != torch.float32 or scale.numel() != 1: raise ValueError( - "FlyDSL FP32 requires contiguous output storage" + f"{name} must contain exactly one FP32 tensor-wise inverse " + f"scale, got dtype={scale.dtype}, shape={tuple(scale.shape)}" ) - if A.device != B.device: + a_flydsl, b_flydsl, m, n, _ = _canonicalize_blas_operands( + A_data, transa, B_data, transb + ) + + if a_flydsl.device != b_flydsl.device: raise ValueError( - f"A and B must be on the same device, got {A.device} and {B.device}" + f"A and B must be on the same device, got " + f"{a_flydsl.device} and {b_flydsl.device}" ) - fp32_matmul( - B_flat, - A.transpose(0, 1), - D.view(m, n), + D = _validate_or_allocate_output( + D, + shape=(m, n), + dtype=torch.float16, + device=a_flydsl.device, + backend_name="FP8", ) + # Operand swap means B's tensor-wise scale belongs to a_flydsl and A's + # tensor-wise scale belongs to b_flydsl. + fp8_matmul( + a_flydsl, + B_scale_inv, + b_flydsl, + A_scale_inv, + D.view(m, n), + ) return D + def te_generic_gemm_flydsl( A, transa, @@ -572,12 +611,19 @@ def te_generic_gemm_flydsl( ): """Run a supported FlyDSL GEMM through TE's generic GEMM interface. - Currently supported: - - MXFP8 TN input with FP16 output - - tensor-wise E4M3 x E4M3 FP8 TN input with FP16 output - - BF16 TN input with BF16 output - - FP16 TN input with FP16 output - - FP32 TN input with FP32 output + Supported layouts: + - TN: transa=True, transb=False + - NN: transa=False, transb=False + - NT: transa=False, transb=True + + TT is intentionally rejected. + + Supported dtypes: + - MXFP8 input with FP16 output + - tensor-wise E4M3 x E4M3 FP8 input with FP16 output + - BF16 input with BF16 output + - FP16 input with FP16 output + - FP32 input with FP32 output """ del bias_type del gelu_in @@ -588,10 +634,10 @@ def te_generic_gemm_flydsl( del comm_type del extra_output del bulk_overlap - - if not transa or transb: + + if transa and transb: raise NotImplementedError( - "FlyDSL GEMM currently supports only transa=True, transb=False" + "FlyDSL GEMM does not support transa=True, transb=True (TT)" ) _validate_common_epilogue( @@ -604,47 +650,55 @@ def te_generic_gemm_flydsl( beta=beta, ) - a_is_mxfp8 = _is_mxfp8_operand(A) - b_is_mxfp8 = _is_mxfp8_operand(B) + a_kind, _ = _classify_input(A) + b_kind, _ = _classify_input(B) - if a_is_mxfp8 or b_is_mxfp8: - if not (a_is_mxfp8 and b_is_mxfp8): + if a_kind == "mxfp8" or b_kind == "mxfp8": + # Validate both are MXFP8 + if a_kind != b_kind: raise ValueError( "Mixed MXFP8 and non-MXFP8 FlyDSL GEMM inputs are not supported" ) + # Sanity: both operands must have at least one pre-quantized copy. + if getattr(A, '_rowwise_data', None) is None and getattr(A, '_columnwise_data', None) is None: + raise RuntimeError("MXFP8Tensor has neither rowwise nor columnwise data") + if getattr(B, '_rowwise_data', None) is None and getattr(B, '_columnwise_data', None) is None: + raise RuntimeError("MXFP8Tensor has neither rowwise nor columnwise data") + + # Only supports FP16 output for now. if output_dtype not in (None, tex.DType.kFloat16): raise NotImplementedError( "FlyDSL MXFP8 currently supports only FP16 output, " f"got {output_dtype}" ) - D = _run_mxfp8_tn(A, B, D) + D = _run_mxfp8(A, transa, B, transb, D) return D, None, None, None - a_is_fp8 = _is_fp8_operand(A) - b_is_fp8 = _is_fp8_operand(B) - - if a_is_fp8 or b_is_fp8: - if not (a_is_fp8 and b_is_fp8): + if a_kind == "fp8" or b_kind == "fp8": + if a_kind != b_kind: raise ValueError( "Mixed regular FP8 and non-FP8 FlyDSL GEMM inputs are not supported" ) - if output_dtype not in (None, tex.DType.kFloat16): raise NotImplementedError( "FlyDSL tensor-wise FP8 currently supports only FP16 output, " f"got {output_dtype}" ) - D = _run_fp8_tn(A, B, D) + D = _run_fp8(A, transa, B, transb, D) return D, None, None, None - if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + if a_kind != "regular" or b_kind != "regular": raise TypeError( "Unsupported FlyDSL GEMM operand types: " f"{type(A).__name__} and {type(B).__name__}" ) + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError( + "FlyDSL regular GEMM expects plain torch.Tensor operands" + ) if A.dtype == torch.bfloat16 and B.dtype == torch.bfloat16: if output_dtype not in (None, tex.DType.kBFloat16): @@ -652,8 +706,16 @@ def te_generic_gemm_flydsl( "FlyDSL BF16 currently supports only BF16 output, " f"got {output_dtype}" ) - - D = _run_bf16_tn(A, B, D) + D = _run_regular_gemm( + A, + transa, + B, + transb, + D, + dtype=torch.bfloat16, + matmul=bf16_matmul, + backend_name="BF16", + ) return D, None, None, None if A.dtype == torch.float16 and B.dtype == torch.float16: @@ -662,8 +724,16 @@ def te_generic_gemm_flydsl( "FlyDSL FP16 currently supports only FP16 output, " f"got {output_dtype}" ) - - D = _run_fp16_tn(A, B, D) + D = _run_regular_gemm( + A, + transa, + B, + transb, + D, + dtype=torch.float16, + matmul=fp16_matmul, + backend_name="FP16", + ) return D, None, None, None if A.dtype == torch.float32 and B.dtype == torch.float32: @@ -672,11 +742,20 @@ def te_generic_gemm_flydsl( "FlyDSL FP32 currently supports only FP32 output, " f"got {output_dtype}" ) - - D = _run_fp32_tn(A, B, D) + D = _run_regular_gemm( + A, + transa, + B, + transb, + D, + dtype=torch.float32, + matmul=fp32_matmul, + backend_name="FP32", + ) return D, None, None, None raise NotImplementedError( - "FlyDSL GEMM currently supports only MXFP8, tensor-wise E4M3 FP8, BF16, FP16, or FP32 inputs; " + "FlyDSL GEMM currently supports only MXFP8, tensor-wise E4M3 FP8, " + "BF16, FP16, or FP32 inputs; " f"got A={A.dtype} and B={B.dtype}" ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index bde328ced..09405a2f2 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -2,16 +2,23 @@ # # See LICENSE for license information. -"""FlyDSL MXFP8 4-wave GEMM kernel for Transformer Engine. +"""FlyDSL MXFP8 GEMM implementation. -The kernel specializes on K at compile time because the K128 loop is fully -hand-unrolled. M/N are runtime launch dimensions. The private optimized core -consumes A and B as FP8 E4M3 tensors shaped [M, K] and [N, K], and writes -float16 C shaped [M, N]. The public ``mxfp8_matmul`` entry point accepts the -Transformer Engine TN contract and performs the required private adaptation. +This module contains both the HK-derived optimized 4-wave kernel and its +MXFP8-specific launch preparation. Transformer Engine BLAS canonicalization +is performed by ``gemm_wrappers.py`` before entering ``mxfp8_matmul``. + +Canonical launch inputs: + + a: [M, K] FP8 payload + a_scale: [M, K/32] raw E8M0 bytes + b: [K, N] FP8 payload + b_scale: [K/32, N] raw E8M0 bytes + D: [M, N] float16 output """ import functools +import os import torch @@ -44,16 +51,33 @@ SCALE_GROUP_SIZE = 32 +def _debug_enabled() -> bool: + value = os.getenv("DEBUG_FLYDSL_MXFP8_GEMM", "") + return value.lower() not in ("", "0", "false", "no", "off") + + +def _debug(message: str) -> None: + if _debug_enabled(): + print(f"[DEBUG_FLYDSL_MXFP8_GEMM] {message}") + + def pack_mx32_scales_iter(scales_u8: torch.Tensor) -> torch.Tensor: - """Pack raw [Rows, K/32] E8M0 uint8 scales as [K/128, Rows] uint32. + """Pack raw [Rows, K/32] E8M0 scales as [K/128, Rows] uint32.""" + if scales_u8.dtype != torch.uint8: + raise TypeError( + f"MXFP8 scales must be torch.uint8 E8M0 bytes, got {scales_u8.dtype}" + ) + if scales_u8.ndim != 2: + raise ValueError( + f"MXFP8 scales must be rank 2, got shape {tuple(scales_u8.shape)}" + ) - This is the intermediate HK/TE iteration-major form: each word contains - four consecutive K32 scale bytes for one K128 iteration and one matrix row. - It is *not* the final MFMA operand layout. - """ - assert scales_u8.dtype == torch.uint8 rows, qk = scales_u8.shape - assert qk % 4 == 0 + if qk % 4 != 0: + raise ValueError( + f"Scale K dimension must be divisible by 4 K32 groups, got {qk}" + ) + s32 = scales_u8.contiguous().view(rows, qk // 4, 4).to(torch.int32) packed = ( s32[:, :, 0] @@ -65,36 +89,33 @@ def pack_mx32_scales_iter(scales_u8: torch.Tensor) -> torch.Tensor: def pack_mx32_scales_for_hk(scales_u8: torch.Tensor) -> torch.Tensor: - """True HK MFMA scale packing: raw [Rows, K/32] -> [K/128, Rows] i32. + """Convert raw rowwise E8M0 scales to [K/128, Rows] MFMA-ready words.""" + scale_iter = pack_mx32_scales_iter(scales_u8) + rows = scales_u8.shape[0] - HK's GEMM hot loop loads one uint32 scale operand per lane for each 64-row - A/B half. The four bytes in that operand correspond to the four 16-row - MFMA slices inside the 64-row half; the scaled-MFMA op_sel/op_sel_hi bits - select the byte. With this layout the GEMM kernel does no hot-loop byte - extraction or broadcast. - """ - assert scales_u8.dtype == torch.uint8 - rows, qk = scales_u8.shape - assert qk % 4 == 0 - assert rows % 64 == 0, f"rows={rows} must be a multiple of 64 for HK MFMA scale packing" + if rows % 64 != 0: + raise ValueError( + f"Rows={rows} must be a multiple of 64 for HK MFMA scale packing" + ) - scale_iter = pack_mx32_scales_iter(scales_u8) # [K/128, Rows], int32 device = scales_u8.device - row = torch.arange(rows, device=device, dtype=torch.int64) - r16 = row % 16 - k_sub = (row // 16) % 4 + row_within_16 = row % 16 + k_subgroup = (row // 16) % 4 tile = row // 64 packed = torch.zeros_like(scale_iter) - for g in range(4): - src_row = tile * 64 + g * 16 + r16 - src_val = scale_iter[:, src_row] - byte_val = (src_val >> (k_sub * 8).view(1, rows)) & 0xFF - packed |= byte_val << (g * 8) + for group in range(4): + source_row = tile * 64 + group * 16 + row_within_16 + source_value = scale_iter[:, source_row] + byte_value = ( + source_value >> (k_subgroup * 8).view(1, rows) + ) & 0xFF + packed |= byte_value << (group * 8) return packed.contiguous() + def _encode_waitcnt(vmcnt=63, lgkmcnt=15): """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. @@ -1116,74 +1137,7 @@ def _cached_launch(K: int): -def mxfp8_matmul( - a: torch.Tensor, - a_scale: torch.Tensor, - b: torch.Tensor, - b_scale: torch.Tensor, - c: torch.Tensor, - stream=None, -): - """TE-facing TN MXFP8 adapter. - - Public/backend contract: - a: [M, K] FP8 payload - a_scale: [M, K/32] raw E8M0 bytes - b: [K, N] FP8 payload - b_scale: [N, K/32] raw E8M0 bytes - c: [M, N] float16 output - - The optimized HK core currently consumes B as row-major [N, K] and consumes - MFMA-ready packed int32 scales. Keep those implementation details behind - this adapter so the TE-facing contract matches the Triton/TE TN contract. - """ - if a.ndim != 2 or b.ndim != 2: - raise ValueError( - f"FlyDSL MXFP8 TN expects rank-2 operands, got A{tuple(a.shape)} " - f"and B{tuple(b.shape)}" - ) - - m, k = a.shape - kb, n = b.shape - if kb != k: - raise ValueError(f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}") - - expected_a_scale = (m, k // SCALE_GROUP_SIZE) - expected_b_scale = (n, k // SCALE_GROUP_SIZE) - if tuple(a_scale.shape) != expected_a_scale: - raise ValueError( - f"A scale shape {tuple(a_scale.shape)} != expected {expected_a_scale}" - ) - if tuple(b_scale.shape) != expected_b_scale: - raise ValueError( - f"B scale shape {tuple(b_scale.shape)} != expected {expected_b_scale}" - ) - if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: - raise TypeError( - "FlyDSL MXFP8 expects raw E8M0 scales stored as torch.uint8" - ) - if tuple(c.shape) != (m, n): - raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") - if c.dtype != torch.float16: - raise TypeError( - f"The current FlyDSL MXFP8 kernel stores float16 output, got {c.dtype}" - ) - - # TE/Triton expose B logically as [K, N]. The existing optimized HK core - # streams contiguous K rows, so adapt B to its private [N, K] representation. - # In the normal TE TN path, b is itself a transpose view of contiguous - # rowwise weight storage, so b.T is already contiguous and this is not a - # physical transpose/copy. - b_hk = b.transpose(0, 1).contiguous() - - # Convert TE's raw per-K32 E8M0 scales into the MFMA-ready words consumed by - # the optimized scaled-MFMA hot loop. - a_scale_hk = pack_mx32_scales_for_hk(a_scale) - b_scale_hk = pack_mx32_scales_for_hk(b_scale) - - doGemm(a, a_scale_hk, b_hk, b_scale_hk, c, stream=stream) - -def doGemm( +def do_gemm( A: torch.Tensor, As: torch.Tensor, B: torch.Tensor, @@ -1240,3 +1194,119 @@ def doGemm( N_runtime, stream=stream, ) + + +__all__ = [ + "BLOCK_M", + "BLOCK_N", + "BLOCK_K", + "do_gemm", +] + + + +def mxfp8_matmul( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + D: torch.Tensor, + stream=None, +): + """Launch the fused MXFP8 kernel from canonical row-major operands. + + BLAS operand canonicalization, shape derivation, and output allocation are + intentionally owned by ``gemm_wrappers.py``. This function only validates + the MXFP8-specific scale contract, converts B to the HK [N, K] convention, + packs E8M0 scales, and launches the optimized 4-wave implementation. + """ + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL MXFP8 expects rank-2 canonical operands, got " + f"a={tuple(a.shape)} and b={tuple(b.shape)}" + ) + + m, k = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Incompatible canonical MXFP8 operands: " + f"{tuple(a.shape)} @ {tuple(b.shape)}" + ) + + if a.device != b.device: + raise ValueError( + f"a and b must be on the same device, got {a.device} and {b.device}" + ) + if D.device != a.device: + raise ValueError(f"D must be on {a.device}, got {D.device}") + if tuple(D.shape) != (m, n): + raise ValueError( + f"D shape {tuple(D.shape)} does not match expected {(m, n)}" + ) + if D.dtype != torch.float16: + raise TypeError( + f"FlyDSL MXFP8 requires torch.float16 output, got {D.dtype}" + ) + if not D.is_contiguous(): + raise ValueError("FlyDSL MXFP8 requires contiguous output storage") + + if k % SCALE_GROUP_SIZE != 0: + raise ValueError( + f"K={k} must be divisible by MXFP8 scale group size " + f"{SCALE_GROUP_SIZE}" + ) + + # Canonical scale contract: + # a_scale [M, K/32] + # b_scale [K/32, N] + expected_a_scale = (m, k // SCALE_GROUP_SIZE) + expected_b_scale = (k // SCALE_GROUP_SIZE, n) + if tuple(a_scale.shape) != expected_a_scale: + raise ValueError( + f"a_scale shape {tuple(a_scale.shape)} != expected " + f"{expected_a_scale}" + ) + if tuple(b_scale.shape) != expected_b_scale: + raise ValueError( + f"b_scale shape {tuple(b_scale.shape)} != expected " + f"{expected_b_scale}" + ) + if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + raise TypeError("FlyDSL MXFP8 expects raw E8M0 scales as torch.uint8") + + # The HK core consumes B and its scales in row-oriented [N, K] form. + b_hk = b.transpose(0, 1).contiguous() + b_scale_rows = b_scale.transpose(0, 1).contiguous() + a_scale_hk = pack_mx32_scales_for_hk(a_scale) + b_scale_hk = pack_mx32_scales_for_hk(b_scale_rows) + + _debug( + f"private kernel inputs: a={tuple(a.shape)}, " + f"contiguous={a.is_contiguous()}; " + f"b_hk={tuple(b_hk.shape)}, contiguous={b_hk.is_contiguous()}; " + f"a_scale_hk={tuple(a_scale_hk.shape)}, " + f"b_scale_hk={tuple(b_scale_hk.shape)}, D={tuple(D.shape)}" + ) + _debug("launching fused MXFP8 4-wave kernel") + + do_gemm( + a, + a_scale_hk, + b_hk, + b_scale_hk, + D.view(m, n), + stream=stream, + ) + + _debug("launch complete") + return D + + +__all__ = [ + "BLOCK_M", + "BLOCK_N", + "BLOCK_K", + "SCALE_GROUP_SIZE", + "mxfp8_matmul", +] From 4c9543f68a98f88a9bd306f8f9584923503032d6 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 16:27:03 +0000 Subject: [PATCH 08/17] add support for bf16/fp32 output types for flydsl mxfp8 gemm --- .../flydsl_kernels/gemm/gemm_wrappers.py | 26 +++++++--- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 49 ++++++++++++++----- 2 files changed, 57 insertions(+), 18 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 3dd8a648f..ba7c84afb 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -412,6 +412,8 @@ def _run_mxfp8( B, transb, D, + *, + output_dtype: torch.dtype, ): """Canonicalize TE MXFP8 operands, then launch the fused backend.""" layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" @@ -496,7 +498,7 @@ def _run_mxfp8( D = _validate_or_allocate_output( D, shape=(m, n), - dtype=torch.float16, + dtype=output_dtype, device=a_flydsl.device, backend_name="MXFP8", ) @@ -619,7 +621,7 @@ def te_generic_gemm_flydsl( TT is intentionally rejected. Supported dtypes: - - MXFP8 input with FP16 output + - MXFP8 input with FP16, BF16, or FP32 output - tensor-wise E4M3 x E4M3 FP8 input with FP16 output - BF16 input with BF16 output - FP16 input with FP16 output @@ -666,14 +668,26 @@ def te_generic_gemm_flydsl( if getattr(B, '_rowwise_data', None) is None and getattr(B, '_columnwise_data', None) is None: raise RuntimeError("MXFP8Tensor has neither rowwise nor columnwise data") - # Only supports FP16 output for now. - if output_dtype not in (None, tex.DType.kFloat16): + mxfp8_output_dtypes = { + None: torch.float16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in mxfp8_output_dtypes: raise NotImplementedError( - "FlyDSL MXFP8 currently supports only FP16 output, " + "FlyDSL MXFP8 supports FP16, BF16, or FP32 output, " f"got {output_dtype}" ) - D = _run_mxfp8(A, transa, B, transb, D) + D = _run_mxfp8( + A, + transa, + B, + transb, + D, + output_dtype=mxfp8_output_dtypes[output_dtype], + ) return D, None, None, None if a_kind == "fp8" or b_kind == "fp8": diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 09405a2f2..532712ede 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -14,7 +14,7 @@ a_scale: [M, K/32] raw E8M0 bytes b: [K, N] FP8 payload b_scale: [K/32, N] raw E8M0 bytes - D: [M, N] float16 output + D: [M, N] float16, bfloat16, or float32 output """ import functools @@ -199,13 +199,29 @@ def _xcd_swizzle(num_pid_m, num_pid_n): ) -def _compile_kernel(K: int): - """Build the specialized 4-wave kernel for compile-time ``K``. +def _compile_kernel(K: int, output_dtype: torch.dtype): + """Build the specialized 4-wave kernel for compile-time ``K`` and output dtype. ``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_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL MXFP8 supports only float16, bfloat16, and float32 " + f"outputs, got {output_dtype}" + ) + NUM_THREADS = 256 WARP_SIZE = 64 @@ -315,7 +331,7 @@ def kernel_gemm( # 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_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) c_rsrc = buffer_ops.create_buffer_resource( C, max_size=True, @@ -589,7 +605,10 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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) + value = Vec(acc)[ii] + if output_dtype != torch.float32: + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, c_idx) # Explicit register coordinates for HK-style four-quadrant mapping. @@ -1132,8 +1151,8 @@ def launch_gemm( return launch_gemm @functools.lru_cache(maxsize=None) -def _cached_launch(K: int): - return _compile_kernel(K) +def _cached_launch(K: int, output_dtype: torch.dtype): + return _compile_kernel(K, output_dtype) @@ -1169,6 +1188,10 @@ def do_gemm( assert C.shape == (M_runtime, N_runtime), ( f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" ) + assert C.dtype in (torch.float16, torch.bfloat16, torch.float32), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) if stream is None: stream = torch.cuda.current_stream() # Match the Transformer Engine integration descriptor contract exactly. The optimized @@ -1183,7 +1206,7 @@ def do_gemm( Bs_arg = Bs.contiguous().view(-1) C_arg = C.contiguous().view(-1) - launch = _cached_launch(int(K_runtime)) + launch = _cached_launch(int(K_runtime), C.dtype) launch( A_arg, As_arg, @@ -1218,7 +1241,7 @@ def mxfp8_matmul( BLAS operand canonicalization, shape derivation, and output allocation are intentionally owned by ``gemm_wrappers.py``. This function only validates the MXFP8-specific scale contract, converts B to the HK [N, K] convention, - packs E8M0 scales, and launches the optimized 4-wave implementation. + packs E8M0 scales, and launches the output-dtype-specialized 4-wave implementation. """ if a.ndim != 2 or b.ndim != 2: raise ValueError( @@ -1244,9 +1267,10 @@ def mxfp8_matmul( raise ValueError( f"D shape {tuple(D.shape)} does not match expected {(m, n)}" ) - if D.dtype != torch.float16: + if D.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise TypeError( - f"FlyDSL MXFP8 requires torch.float16 output, got {D.dtype}" + "FlyDSL MXFP8 supports torch.float16, torch.bfloat16, or " + f"torch.float32 output, got {D.dtype}" ) if not D.is_contiguous(): raise ValueError("FlyDSL MXFP8 requires contiguous output storage") @@ -1286,7 +1310,8 @@ def mxfp8_matmul( f"contiguous={a.is_contiguous()}; " f"b_hk={tuple(b_hk.shape)}, contiguous={b_hk.is_contiguous()}; " f"a_scale_hk={tuple(a_scale_hk.shape)}, " - f"b_scale_hk={tuple(b_scale_hk.shape)}, D={tuple(D.shape)}" + f"b_scale_hk={tuple(b_scale_hk.shape)}, D={tuple(D.shape)}, " + f"D_dtype={D.dtype}" ) _debug("launching fused MXFP8 4-wave kernel") From b9eaa8a42cf622e87b309a90e2b4ec49ebe21fef Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 16:39:54 +0000 Subject: [PATCH 09/17] add support for bf16/fp32 output types for flydsl fp8 gemm --- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 70 +++++++++++++++---- .../flydsl_kernels/gemm/gemm_wrappers.py | 26 +++++-- 2 files changed, 77 insertions(+), 19 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index a85b8ea91..099a015f0 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -7,9 +7,9 @@ The kernel specializes on K at compile time because the K128 loop is fully hand-unrolled. M/N are runtime launch dimensions. The private optimized core consumes A and B as FP8 E4M3 tensors shaped [M, K] and [N, K], one FP32 inverse -scale per operand, and writes float16 C shaped [M, N]. The public ``fp8_matmul`` -entry point accepts Transformer Engine's TN contract and performs the required -private adaptation. +scale per operand, and writes float16, bfloat16, or float32 C shaped [M, N]. The public +``fp8_matmul`` entry point accepts Transformer Engine's TN contract and +performs the required private adaptation. This module imports ``flydsl`` at import time and must therefore be imported lazily only after FlyDSL availability has been confirmed. @@ -183,13 +183,32 @@ def _xcd_swizzle(num_pid_m, num_pid_n): ) -def _compile_kernel(K: int, use_xcd_remap: bool = True): - """Build the specialized 4-wave kernel for compile-time ``K``. +def _compile_kernel( + K: int, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + """Build the specialized 4-wave kernel for compile-time K/output dtype. ``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_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {output_dtype}" + ) NUM_THREADS = 256 WARP_SIZE = 64 @@ -311,7 +330,7 @@ def kernel_gemm( # 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_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) c_rsrc = buffer_ops.create_buffer_resource( C, max_size=True, @@ -512,7 +531,10 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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] * output_scale).to(fx.Float16), c_rsrc, c_idx) + value = Vec(acc)[ii] * output_scale + if output_dtype != torch.float32: + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, c_idx) # Explicit register coordinates for HK-style four-quadrant mapping. @@ -955,8 +977,16 @@ def launch_gemm( return launch_gemm @functools.lru_cache(maxsize=None) -def _cached_launch(K: int, use_xcd_remap: bool = True): - return _compile_kernel(K, use_xcd_remap=use_xcd_remap) +def _cached_launch( + K: int, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + output_dtype, + use_xcd_remap=use_xcd_remap, + ) @@ -975,7 +1005,7 @@ def fp8_matmul( a_scale_inv: one-element FP32 inverse quantization scale b: [K, N] FP8 E4M3 weight payload b_scale_inv: one-element FP32 inverse quantization scale - c: [M, N] float16 output + c: [M, N] float16, bfloat16, or float32 output The optimized private core streams both operands as row-major [Rows, K], so B is adapted from TE's logical [K, N] representation to [N, K]. @@ -1016,9 +1046,10 @@ def fp8_matmul( if tuple(c.shape) != (m, n): raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") - if c.dtype != torch.float16: + if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise TypeError( - f"The current FlyDSL FP8 kernel stores float16 output, got {c.dtype}" + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {c.dtype}" ) if not c.is_contiguous(): raise ValueError("FlyDSL FP8 requires contiguous output storage") @@ -1056,7 +1087,14 @@ def doGemm( N_runtime, Kb_runtime = B.shape assert A.dtype == torch.float8_e4m3fn, f"A dtype {A.dtype} != torch.float8_e4m3fn" assert B.dtype == torch.float8_e4m3fn, f"B dtype {B.dtype} != torch.float8_e4m3fn" - assert C.dtype == torch.float16, f"C dtype {C.dtype} != torch.float16" + assert C.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" assert M_runtime % _BLOCK_M == 0, f"M={M_runtime} must be a multiple of {_BLOCK_M}" assert N_runtime % _BLOCK_N == 0, f"N={N_runtime} must be a multiple of {_BLOCK_N}" @@ -1077,7 +1115,11 @@ def doGemm( A_scale_arg = A_scale_inv.contiguous().view(-1) B_scale_arg = B_scale_inv.contiguous().view(-1) - launch = _cached_launch(int(K_runtime), bool(use_xcd_remap)) + launch = _cached_launch( + int(K_runtime), + C.dtype, + bool(use_xcd_remap), + ) launch( A_arg, B_arg, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index ba7c84afb..9a4397b7d 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -518,6 +518,8 @@ def _run_fp8( B, transb, D, + *, + output_dtype: torch.dtype, ): """Run tensor-wise E4M3 x E4M3 FP8 for TN/NN/NT.""" a_fp8_dtype = getattr(A, "_fp8_dtype", None) @@ -570,7 +572,7 @@ def _run_fp8( D = _validate_or_allocate_output( D, shape=(m, n), - dtype=torch.float16, + dtype=output_dtype, device=a_flydsl.device, backend_name="FP8", ) @@ -622,7 +624,7 @@ def te_generic_gemm_flydsl( Supported dtypes: - MXFP8 input with FP16, BF16, or FP32 output - - tensor-wise E4M3 x E4M3 FP8 input with FP16 output + - tensor-wise E4M3 x E4M3 FP8 input with FP16, BF16, or FP32 output - BF16 input with BF16 output - FP16 input with FP16 output - FP32 input with FP32 output @@ -695,13 +697,27 @@ def te_generic_gemm_flydsl( raise ValueError( "Mixed regular FP8 and non-FP8 FlyDSL GEMM inputs are not supported" ) - if output_dtype not in (None, tex.DType.kFloat16): + + fp8_output_dtypes = { + None: torch.float16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in fp8_output_dtypes: raise NotImplementedError( - "FlyDSL tensor-wise FP8 currently supports only FP16 output, " + "FlyDSL tensor-wise FP8 supports FP16, BF16, or FP32 output, " f"got {output_dtype}" ) - D = _run_fp8(A, transa, B, transb, D) + D = _run_fp8( + A, + transa, + B, + transb, + D, + output_dtype=fp8_output_dtypes[output_dtype], + ) return D, None, None, None if a_kind != "regular" or b_kind != "regular": From 71c4ef453a8dd77dd440ea2b1070570505e6fde7 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 17:10:03 +0000 Subject: [PATCH 10/17] add broad output dtype support for flydsl fp8/fp16/bf16 gemms --- .../pytorch/flydsl_kernels/gemm/bf16_gemm.py | 73 +++++++++++++++---- .../pytorch/flydsl_kernels/gemm/fp16_gemm.py | 71 +++++++++++++++--- .../flydsl_kernels/gemm/gemm_wrappers.py | 32 ++++++-- 3 files changed, 144 insertions(+), 32 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py index ea489c38a..b5c6045c2 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py @@ -6,8 +6,9 @@ The kernel specializes on K at compile time because the K64 loop is fully hand-unrolled. M/N are runtime launch dimensions. The private optimized core -consumes A and B as BF16 tensors shaped [M, K] and [N, K], and writes BF16 C -shaped [M, N]. The public ``bf16_matmul`` entry point accepts Transformer +consumes A and B as BF16 tensors shaped [M, K] and [N, K], and writes FP16, +BF16, or FP32 C shaped [M, N]. The public ``bf16_matmul`` entry point accepts +Transformer Engine's TN contract and performs the required private adaptation. This module imports ``flydsl`` at import time and must therefore be imported @@ -183,8 +184,12 @@ def _xcd_swizzle(num_pid_m, num_pid_n): ) -def _compile_kernel(K: int, use_xcd_remap: bool = True): - """Build the specialized 4-wave kernel for compile-time ``K``. +def _compile_kernel( + K: int, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + """Build the specialized 4-wave kernel for compile-time ``K`` and output dtype. ``K`` must contain at least four K64 tiles. Runtime M/N are expected to be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. @@ -207,6 +212,21 @@ def _compile_kernel(K: int, use_xcd_remap: bool = True): ELEM_BYTES = 2 VEC_BYTES = 16 + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL BF16 GEMM output dtype must be torch.float16, " + f"torch.bfloat16, or torch.float32, got {output_dtype}" + ) + LDS_ELEMS_A = BLOCK_M * BLOCK_K LDS_ELEMS_B = BLOCK_N * BLOCK_K LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES @@ -306,7 +326,7 @@ def kernel_gemm( # 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 BF16. + c_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) c_rsrc = buffer_ops.create_buffer_resource( C, max_size=True, @@ -536,7 +556,10 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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.BFloat16), c_rsrc, c_idx) + value = Vec(acc)[ii] + if const_expr(output_dtype != torch.float32): + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, c_idx) # Explicit register coordinates for HK-style four-quadrant mapping. @@ -976,8 +999,16 @@ def launch_gemm( return launch_gemm @functools.lru_cache(maxsize=None) -def _cached_launch(K: int, use_xcd_remap: bool = True): - return _compile_kernel(K, use_xcd_remap=use_xcd_remap) +def _cached_launch( + K: int, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + output_dtype, + use_xcd_remap=use_xcd_remap, + ) def bf16_matmul( @@ -991,7 +1022,7 @@ def bf16_matmul( Public/backend contract: a: [M, K] BF16 b: [K, N] BF16 - c: [M, N] BF16 output + c: [M, N] FP16, BF16, or FP32 output The optimized core streams both operands with K contiguous and therefore privately consumes B as [N, K]. In the normal TE TN path, ``b`` is a @@ -1017,9 +1048,14 @@ def bf16_matmul( ) if tuple(c.shape) != (m, n): raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") - if c.dtype != torch.bfloat16: + if c.dtype not in ( + torch.float16, + torch.bfloat16, + torch.float32, + ): raise TypeError( - f"The current FlyDSL BF16 kernel stores torch.bfloat16 output, got {c.dtype}" + "FlyDSL BF16 GEMM output dtype must be torch.float16, " + f"torch.bfloat16, or torch.float32, got {c.dtype}" ) if a.device != b.device or a.device != c.device: raise ValueError( @@ -1048,7 +1084,14 @@ def doGemm( N_runtime, Kb_runtime = B.shape assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" assert A.dtype == torch.bfloat16 and B.dtype == torch.bfloat16 - assert C.dtype == torch.bfloat16 + assert C.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) assert M_runtime % _BLOCK_M == 0, f"M={M_runtime} must be a multiple of {_BLOCK_M}" assert N_runtime % _BLOCK_N == 0, f"N={N_runtime} must be a multiple of {_BLOCK_N}" assert K_runtime % _BLOCK_K == 0, f"K={K_runtime} must be a multiple of {_BLOCK_K}" @@ -1061,5 +1104,9 @@ def doGemm( A_arg = A.contiguous().view(torch.uint8).view(-1) B_arg = B.contiguous().view(torch.uint8).view(-1) C_arg = C.view(-1) - launch = _cached_launch(int(K_runtime), bool(use_xcd_remap)) + launch = _cached_launch( + int(K_runtime), + C.dtype, + bool(use_xcd_remap), + ) launch(A_arg, B_arg, C_arg, M_runtime, N_runtime, stream=stream) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py index 66f68816e..7ad7ed31f 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py @@ -6,8 +6,9 @@ The kernel specializes on K at compile time because the K64 loop is fully hand-unrolled. M/N are runtime launch dimensions. The private optimized core -consumes A and B as FP16 tensors shaped [M, K] and [N, K], and writes FP16 C -shaped [M, N]. The public ``fp16_matmul`` entry point accepts Transformer +consumes A and B as FP16 tensors shaped [M, K] and [N, K], and writes FP16, +BF16, or FP32 C shaped [M, N]. The public ``fp16_matmul`` entry point accepts +Transformer Engine's TN contract and performs the required private adaptation. This module imports ``flydsl`` at import time and must therefore be imported @@ -189,7 +190,11 @@ def _xcd_swizzle(num_pid_m, num_pid_n): ) -def _compile_kernel(K: int, use_xcd_remap: bool = True): +def _compile_kernel( + K: int, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): """Build the specialized 4-wave kernel for compile-time ``K``. ``K`` must contain at least four K64 tiles. Runtime M/N are expected to @@ -213,6 +218,21 @@ def _compile_kernel(K: int, use_xcd_remap: bool = True): ELEM_BYTES = 2 VEC_BYTES = 16 + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL FP16 GEMM output dtype must be torch.float16, " + f"torch.bfloat16, or torch.float32, got {output_dtype}" + ) + LDS_ELEMS_A = BLOCK_M * BLOCK_K LDS_ELEMS_B = BLOCK_N * BLOCK_K LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES @@ -312,7 +332,7 @@ def kernel_gemm( # 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 FP16. + c_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) c_rsrc = buffer_ops.create_buffer_resource( C, max_size=True, @@ -542,7 +562,10 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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) + value = Vec(acc)[ii] + if const_expr(output_dtype != torch.float32): + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, c_idx) # Explicit register coordinates for HK-style four-quadrant mapping. @@ -983,8 +1006,16 @@ def launch_gemm( @functools.lru_cache(maxsize=None) -def _cached_launch(K: int, use_xcd_remap: bool = True): - return _compile_kernel(K, use_xcd_remap=use_xcd_remap) +def _cached_launch( + K: int, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + output_dtype, + use_xcd_remap=use_xcd_remap, + ) def fp16_matmul( @@ -998,7 +1029,7 @@ def fp16_matmul( Public/backend contract: a: [M, K] FP16 b: [K, N] FP16 - c: [M, N] FP16 output + c: [M, N] FP16, BF16, or FP32 output The optimized core streams both operands with K contiguous and therefore privately consumes B as [N, K]. In the normal TE TN path, ``b`` is a @@ -1024,9 +1055,14 @@ def fp16_matmul( ) if tuple(c.shape) != (m, n): raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") - if c.dtype != torch.float16: + if c.dtype not in ( + torch.float16, + torch.bfloat16, + torch.float32, + ): raise TypeError( - f"The current FlyDSL FP16 kernel stores torch.float16 output, got {c.dtype}" + "FlyDSL FP16 GEMM output dtype must be torch.float16, " + f"torch.bfloat16, or torch.float32, got {c.dtype}" ) if a.device != b.device or a.device != c.device: raise ValueError( @@ -1055,7 +1091,14 @@ def doGemm( N_runtime, Kb_runtime = B.shape assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" assert A.dtype == torch.float16 and B.dtype == torch.float16 - assert C.dtype == torch.float16 + assert C.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) assert M_runtime % _BLOCK_M == 0, f"M={M_runtime} must be a multiple of {_BLOCK_M}" assert N_runtime % _BLOCK_N == 0, f"N={N_runtime} must be a multiple of {_BLOCK_N}" assert K_runtime % _BLOCK_K == 0, f"K={K_runtime} must be a multiple of {_BLOCK_K}" @@ -1068,5 +1111,9 @@ def doGemm( A_arg = A.contiguous().view(torch.uint8).view(-1) B_arg = B.contiguous().view(torch.uint8).view(-1) C_arg = C.view(-1) - launch = _cached_launch(int(K_runtime), bool(use_xcd_remap)) + launch = _cached_launch( + int(K_runtime), + C.dtype, + bool(use_xcd_remap), + ) launch(A_arg, B_arg, C_arg, M_runtime, N_runtime, stream=stream) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 9a4397b7d..49bfce512 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -278,6 +278,7 @@ def _run_regular_gemm( dtype, matmul, backend_name, + output_dtype=None, ): """Run FP16/BF16/FP32 through shared TN/NN/NT shape handling.""" if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): @@ -298,10 +299,13 @@ def _run_regular_gemm( A, transa, B, transb ) + if output_dtype is None: + output_dtype = dtype + D = _validate_or_allocate_output( D, shape=(m, n), - dtype=dtype, + dtype=output_dtype, device=A.device, backend_name=backend_name, ) @@ -625,8 +629,8 @@ def te_generic_gemm_flydsl( Supported dtypes: - MXFP8 input with FP16, BF16, or FP32 output - tensor-wise E4M3 x E4M3 FP8 input with FP16, BF16, or FP32 output - - BF16 input with BF16 output - - FP16 input with FP16 output + - BF16 input with FP16, BF16, or FP32 output + - FP16 input with FP16, BF16, or FP32 output - FP32 input with FP32 output """ del bias_type @@ -731,9 +735,15 @@ def te_generic_gemm_flydsl( ) if A.dtype == torch.bfloat16 and B.dtype == torch.bfloat16: - if output_dtype not in (None, tex.DType.kBFloat16): + bf16_output_dtypes = { + None: torch.bfloat16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in bf16_output_dtypes: raise NotImplementedError( - "FlyDSL BF16 currently supports only BF16 output, " + "FlyDSL BF16 supports FP16, BF16, or FP32 output, " f"got {output_dtype}" ) D = _run_regular_gemm( @@ -745,13 +755,20 @@ def te_generic_gemm_flydsl( dtype=torch.bfloat16, matmul=bf16_matmul, backend_name="BF16", + output_dtype=bf16_output_dtypes[output_dtype], ) return D, None, None, None if A.dtype == torch.float16 and B.dtype == torch.float16: - if output_dtype not in (None, tex.DType.kFloat16): + fp16_output_dtypes = { + None: torch.float16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in fp16_output_dtypes: raise NotImplementedError( - "FlyDSL FP16 currently supports only FP16 output, " + "FlyDSL FP16 supports FP16, BF16, or FP32 output, " f"got {output_dtype}" ) D = _run_regular_gemm( @@ -763,6 +780,7 @@ def te_generic_gemm_flydsl( dtype=torch.float16, matmul=fp16_matmul, backend_name="FP16", + output_dtype=fp16_output_dtypes[output_dtype], ) return D, None, None, None From aa19610cc717df562cfb9d82b10cc8e947179a15 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 17:26:38 +0000 Subject: [PATCH 11/17] add mixed e4m3/e5m2 fp8 dtype support for flydsl fp8 GEMM --- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 79 +++++++++++++------ .../flydsl_kernels/gemm/gemm_wrappers.py | 15 ++-- 2 files changed, 65 insertions(+), 29 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 099a015f0..c11c6765e 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -6,7 +6,8 @@ The kernel specializes on K at compile time because the K128 loop is fully hand-unrolled. M/N are runtime launch dimensions. The private optimized core -consumes A and B as FP8 E4M3 tensors shaped [M, K] and [N, K], one FP32 inverse +consumes independently typed FP8 E4M3 or E5M2 A/B tensors shaped [M, K] and +[N, K], one FP32 inverse scale per operand, and writes float16, bfloat16, or float32 C shaped [M, N]. The public ``fp8_matmul`` entry point accepts Transformer Engine's TN contract and performs the required private adaptation. @@ -185,16 +186,31 @@ def _xcd_swizzle(num_pid_m, num_pid_n): def _compile_kernel( K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, output_dtype: torch.dtype, use_xcd_remap: bool = True, ): - """Build the specialized 4-wave kernel for compile-time K/output dtype. + """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. ``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_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + fp8_input_types = { + torch.float8_e4m3fn: (fx.Float8E4M3FN, 0), + torch.float8_e5m2: (fx.Float8E5M2, 1), + } + try: + a_fx_dtype, a_matrix_format = fp8_input_types[a_fp8_dtype] + b_fx_dtype, b_matrix_format = fp8_input_types[b_fp8_dtype] + except KeyError as exc: + raise TypeError( + "FlyDSL FP8 input dtype must be torch.float8_e4m3fn or " + f"torch.float8_e5m2, got A={a_fp8_dtype}, B={b_fp8_dtype}" + ) from exc + if output_dtype == torch.float16: output_element_bytes = 2 output_fx_dtype = fx.Float16 @@ -248,14 +264,14 @@ def _compile_kernel( 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] + a0_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a0_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + b0_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b0_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) def kernel_gemm( @@ -273,9 +289,10 @@ def kernel_gemm( 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_f8_ir_t = a_fx_dtype.ir_type + b_f8_ir_t = b_fx_dtype.ir_type + gA = make_fp8_buffer_tensor(A, a_f8_ir_t) + gB = make_fp8_buffer_tensor(B, 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)) a_scale_rsrc = buffer_ops.create_buffer_resource(A_scale_inv, max_size=True) @@ -315,8 +332,8 @@ def kernel_gemm( # 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) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) s2r = S2RLoader(fx.Int32(0), 1) layout_lane16 = fx.make_layout((4, 16), (16, 1)) @@ -474,7 +491,8 @@ def pinned_mfma(acc_idx, a_frag, b_frag): f"v_mfma_f32_16x16x128_f8f6f4 " f"a[{acc_pin}:{acc_pin + 3}], " f"$0, $1, " - f"a[{acc_pin}:{acc_pin + 3}]" + f"a[{acc_pin}:{acc_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" ), ( f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," @@ -497,7 +515,8 @@ def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): f"v_mfma_f32_16x16x128_f8f6f4 " f"a[{dst_pin}:{dst_pin + 3}], " f"$0, $1, " - f"a[{old_pin}:{old_pin + 3}]" + f"a[{old_pin}:{old_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" ), ( f"v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}}," @@ -979,11 +998,15 @@ def launch_gemm( @functools.lru_cache(maxsize=None) def _cached_launch( K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, output_dtype: torch.dtype, use_xcd_remap: bool = True, ): return _compile_kernel( K, + a_fp8_dtype, + b_fp8_dtype, output_dtype, use_xcd_remap=use_xcd_remap, ) @@ -1001,9 +1024,9 @@ def fp8_matmul( """TE-facing TN tensor-wise FP8 adapter. Public/backend contract: - a: [M, K] FP8 E4M3 activation payload + a: [M, K] FP8 E4M3 or E5M2 activation payload a_scale_inv: one-element FP32 inverse quantization scale - b: [K, N] FP8 E4M3 weight payload + b: [K, N] FP8 E4M3 or E5M2 weight payload b_scale_inv: one-element FP32 inverse quantization scale c: [M, N] float16, bfloat16, or float32 output @@ -1019,9 +1042,13 @@ def fp8_matmul( f"and B{tuple(b.shape)}" ) - if a.dtype != torch.float8_e4m3fn or b.dtype != torch.float8_e4m3fn: + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + if a.dtype not in supported_fp8_dtypes or b.dtype not in supported_fp8_dtypes: raise TypeError( - "FlyDSL FP8 GEMM requires torch.float8_e4m3fn payloads, " + "FlyDSL FP8 GEMM expects E4M3 or E5M2 payloads, " f"got A={a.dtype} and B={b.dtype}" ) @@ -1085,8 +1112,12 @@ def doGemm( """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" M_runtime, K_runtime = A.shape N_runtime, Kb_runtime = B.shape - assert A.dtype == torch.float8_e4m3fn, f"A dtype {A.dtype} != torch.float8_e4m3fn" - assert B.dtype == torch.float8_e4m3fn, f"B dtype {B.dtype} != torch.float8_e4m3fn" + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + assert A.dtype in supported_fp8_dtypes, f"unsupported A FP8 dtype: {A.dtype}" + assert B.dtype in supported_fp8_dtypes, f"unsupported B FP8 dtype: {B.dtype}" assert C.dtype in ( torch.float16, torch.bfloat16, @@ -1117,6 +1148,8 @@ def doGemm( launch = _cached_launch( int(K_runtime), + A.dtype, + B.dtype, C.dtype, bool(use_xcd_remap), ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 49bfce512..c39f58063 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -525,16 +525,19 @@ def _run_fp8( *, output_dtype: torch.dtype, ): - """Run tensor-wise E4M3 x E4M3 FP8 for TN/NN/NT.""" + """Run tensor-wise E4M3/E5M2 FP8 combinations for TN/NN/NT.""" a_fp8_dtype = getattr(A, "_fp8_dtype", None) b_fp8_dtype = getattr(B, "_fp8_dtype", None) + supported_fp8_dtypes = ( + tex.DType.kFloat8E4M3, + tex.DType.kFloat8E5M2, + ) if ( - a_fp8_dtype != tex.DType.kFloat8E4M3 - or b_fp8_dtype != tex.DType.kFloat8E4M3 + a_fp8_dtype not in supported_fp8_dtypes + or b_fp8_dtype not in supported_fp8_dtypes ): raise NotImplementedError( - "The current FlyDSL FP8 kernel supports only " - "tex.DType.kFloat8E4M3 x tex.DType.kFloat8E4M3; " + "FlyDSL FP8 supports E4M3 and E5M2 independently for A/B; " f"got A={a_fp8_dtype} and B={b_fp8_dtype}" ) @@ -628,7 +631,7 @@ def te_generic_gemm_flydsl( Supported dtypes: - MXFP8 input with FP16, BF16, or FP32 output - - tensor-wise E4M3 x E4M3 FP8 input with FP16, BF16, or FP32 output + - tensor-wise E4M3/E5M2 FP8 A/B combinations with FP16, BF16, or FP32 output - BF16 input with FP16, BF16, or FP32 output - FP16 input with FP16, BF16, or FP32 output - FP32 input with FP32 output From 2ee10d3ce6a7b2272fbb5fcd1e0725ae78cec3fd Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 17:39:36 +0000 Subject: [PATCH 12/17] add mixed e4m3/e5m2 fp8 dtype support for flydsl mxfp8 GEMM --- .../flydsl_kernels/gemm/gemm_wrappers.py | 33 ++++++- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 96 ++++++++++++++----- 2 files changed, 103 insertions(+), 26 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index c39f58063..1b60a6a05 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -419,7 +419,22 @@ def _run_mxfp8( *, output_dtype: torch.dtype, ): - """Canonicalize TE MXFP8 operands, then launch the fused backend.""" + """Canonicalize independently typed E4M3/E5M2 MXFP8 operands.""" + a_fp8_dtype = getattr(A, "_fp8_dtype", None) + b_fp8_dtype = getattr(B, "_fp8_dtype", None) + supported_fp8_dtypes = ( + tex.DType.kFloat8E4M3, + tex.DType.kFloat8E5M2, + ) + if ( + a_fp8_dtype not in supported_fp8_dtypes + or b_fp8_dtype not in supported_fp8_dtypes + ): + raise NotImplementedError( + "FlyDSL MXFP8 supports E4M3 and E5M2 independently for A/B; " + f"got A={a_fp8_dtype} and B={b_fp8_dtype}" + ) + layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" _mxfp8_debug( f"entry: layout={layout}, A_type={type(A).__name__}, " @@ -440,6 +455,14 @@ def _run_mxfp8( name="B", ) + # MXFP8Tensor stores rowwise/columnwise payloads as raw uint8. Reinterpret + # those exact bytes using each operand's own FP8 metadata before applying + # BLAS canonicalization. No copy or numerical conversion is performed here. + if A_data.dtype == torch.uint8: + A_data = reinterpret_as_fp8_tensor(A_data, a_fp8_dtype) + if B_data.dtype == torch.uint8: + B_data = reinterpret_as_fp8_tensor(B_data, b_fp8_dtype) + a_flydsl, b_flydsl, m, n, k = _canonicalize_blas_operands( A_data, transa, @@ -458,8 +481,10 @@ def _run_mxfp8( _mxfp8_debug( f"canonicalized layout={layout}: " - f"a={tuple(a_flydsl.shape)}, stride={tuple(a_flydsl.stride())}; " - f"b={tuple(b_flydsl.shape)}, stride={tuple(b_flydsl.stride())}" + f"a={tuple(a_flydsl.shape)}, dtype={a_flydsl.dtype}, " + f"stride={tuple(a_flydsl.stride())}; " + f"b={tuple(b_flydsl.shape)}, dtype={b_flydsl.dtype}, " + f"stride={tuple(b_flydsl.stride())}" ) _mxfp8_debug( f"canonicalized scales: " @@ -630,7 +655,7 @@ def te_generic_gemm_flydsl( TT is intentionally rejected. Supported dtypes: - - MXFP8 input with FP16, BF16, or FP32 output + - MXFP8 E4M3/E5M2 A/B combinations with FP16, BF16, or FP32 output - tensor-wise E4M3/E5M2 FP8 A/B combinations with FP16, BF16, or FP32 output - BF16 input with FP16, BF16, or FP32 output - FP16 input with FP16, BF16, or FP32 output diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 532712ede..4e83f271b 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -10,9 +10,9 @@ Canonical launch inputs: - a: [M, K] FP8 payload + a: [M, K] FP8 E4M3 or E5M2 payload a_scale: [M, K/32] raw E8M0 bytes - b: [K, N] FP8 payload + b: [K, N] FP8 E4M3 or E5M2 payload b_scale: [K/32, N] raw E8M0 bytes D: [M, N] float16, bfloat16, or float32 output """ @@ -199,14 +199,32 @@ def _xcd_swizzle(num_pid_m, num_pid_n): ) -def _compile_kernel(K: int, output_dtype: torch.dtype): - """Build the specialized 4-wave kernel for compile-time ``K`` and output dtype. +def _compile_kernel( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, +): + """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. ``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_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + fp8_input_types = { + torch.float8_e4m3fn: (fx.Float8E4M3FN, 0), + torch.float8_e5m2: (fx.Float8E5M2, 1), + } + try: + a_fx_dtype, a_matrix_format = fp8_input_types[a_fp8_dtype] + b_fx_dtype, b_matrix_format = fp8_input_types[b_fp8_dtype] + except KeyError as exc: + raise TypeError( + "FlyDSL MXFP8 input dtype must be torch.float8_e4m3fn or " + f"torch.float8_e5m2, got A={a_fp8_dtype}, B={b_fp8_dtype}" + ) from exc + if output_dtype == torch.float16: output_element_bytes = 2 output_fx_dtype = fx.Float16 @@ -262,14 +280,14 @@ def _compile_kernel(K: int, output_dtype: torch.dtype): 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] + a0_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a0_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + b0_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b0_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) def kernel_gemm( @@ -281,9 +299,10 @@ def kernel_gemm( 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_f8_ir_t = a_fx_dtype.ir_type + b_f8_ir_t = b_fx_dtype.ir_type + gA = make_fp8_buffer_tensor(A, a_f8_ir_t) + gB = make_fp8_buffer_tensor(B, 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) @@ -316,8 +335,8 @@ def kernel_gemm( # 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) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) s2r = S2RLoader(fx.Int32(0), 1) layout_lane16 = fx.make_layout((4, 16), (16, 1)) @@ -544,7 +563,8 @@ def pinned_mfma(acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): 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"op_sel_hi:[{mi >> 1},{ni >> 1},0] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" ), (f"v,v,v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}"), has_side_effects=True, @@ -571,7 +591,8 @@ def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag, a_scale, b_scale, m 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"op_sel_hi:[{mi >> 1},{ni >> 1},0] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" ), (f"v,v,v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}},~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}"), has_side_effects=True, @@ -1151,8 +1172,18 @@ def launch_gemm( return launch_gemm @functools.lru_cache(maxsize=None) -def _cached_launch(K: int, output_dtype: torch.dtype): - return _compile_kernel(K, output_dtype) +def _cached_launch( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, +): + return _compile_kernel( + K, + a_fp8_dtype, + b_fp8_dtype, + output_dtype, + ) @@ -1173,6 +1204,12 @@ def do_gemm( """ M_runtime, K_runtime = A.shape N_runtime, Kb_runtime = B.shape + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + assert A.dtype in supported_fp8_dtypes, f"unsupported A MXFP8 dtype: {A.dtype}" + assert B.dtype in supported_fp8_dtypes, f"unsupported B MXFP8 dtype: {B.dtype}" assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" assert M_runtime % _BLOCK_M == 0, f"M={M_runtime} must be a multiple of {_BLOCK_M}" assert N_runtime % _BLOCK_N == 0, f"N={N_runtime} must be a multiple of {_BLOCK_N}" @@ -1206,7 +1243,12 @@ def do_gemm( Bs_arg = Bs.contiguous().view(-1) C_arg = C.contiguous().view(-1) - launch = _cached_launch(int(K_runtime), C.dtype) + launch = _cached_launch( + int(K_runtime), + A.dtype, + B.dtype, + C.dtype, + ) launch( A_arg, As_arg, @@ -1257,6 +1299,16 @@ def mxfp8_matmul( f"{tuple(a.shape)} @ {tuple(b.shape)}" ) + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + if a.dtype not in supported_fp8_dtypes or b.dtype not in supported_fp8_dtypes: + raise TypeError( + "FlyDSL MXFP8 expects E4M3 or E5M2 payloads independently, " + f"got a={a.dtype} and b={b.dtype}" + ) + if a.device != b.device: raise ValueError( f"a and b must be on the same device, got {a.device} and {b.device}" From 979f38ce7a04435cf20a4d278220ecd7b687ee29 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 18:08:15 +0000 Subject: [PATCH 13/17] add pytorch flydsl gemm tests --- transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 1b60a6a05..343624d6c 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -159,7 +159,6 @@ def _valid_fp8_transpose(t): ) - def _mxfp8_debug_enabled() -> bool: value = os.getenv("DEBUG_FLYDSL_MXFP8_GEMM", "") return value.lower() not in ("", "0", "false", "no", "off") From e1896cdd7e7ec88a30b5e49a8062bbee0214e2fd Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 18:09:58 +0000 Subject: [PATCH 14/17] add the actual pytorch flydsl gemm tests --- tests/pytorch/flydsl_kernels/test_gemm.py | 551 ++++++++++++++++++++++ 1 file changed, 551 insertions(+) create mode 100644 tests/pytorch/flydsl_kernels/test_gemm.py diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py new file mode 100644 index 000000000..48510880f --- /dev/null +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -0,0 +1,551 @@ +# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# +# License for AMD contributions = MIT. See LICENSE for more information + +"""User-facing FlyDSL GEMM tests -- ``general_gemm()`` under ``NVTE_USE_FLYDSL=1``. + +Exercises the same public entry point used by TE ``Linear`` / +``LayerNormLinear``. Coverage mirrors the Triton user-facing GEMM tests for the +currently supported FlyDSL surface: + +- fp32 / fp16 / bf16 regular tensors +- same-format and mixed-format tensor-wise FP8 +- same-format and mixed-format MXFP8 +- TN / NN / NT layouts +- batched multidimensional FP8 flattening + +Fused BIAS and BGRADB epilogues are intentionally not included yet because the +FlyDSL GEMM path does not currently support them. + +Each test compares the FlyDSL path against two independent references: + +1. ``torch.matmul`` on dequantized inputs, independent of hipBLASLt behavior. +2. The native C++ ``tex.generic_gemm`` backend through the same + ``general_gemm`` public surface. + +FlyDSL kernels currently require tile-aligned launch dimensions, so the test +shapes are aligned to the 256x256x128 kernel contract rather than reusing the +odd-sized Triton edge-mask cases. +""" + +import os + +import pytest +import torch + +from transformer_engine.pytorch import Float8Tensor +from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm +from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer +from transformer_engine.pytorch.tensor.mxfp8_tensor import ( + MXFP8Quantizer, + MXFP8Tensor, +) +import transformer_engine_torch as tex + + +# --- Feature detection -------------------------------------------------------- + +major, minor = torch.cuda.get_device_capability() + +# The current FlyDSL MXFP8 implementation uses the gfx950 fp8-scaled MFMA. +has_mxfp8_support = major == 9 and minor >= 5 + +requires_mxfp8_support = pytest.mark.skipif( + not has_mxfp8_support, + reason="FlyDSL MXFP8 requires gfx950+ fp8-scaled MFMA support", +) + + +# --- Test parameters ---------------------------------------------------------- + +# The current FlyDSL kernels have no M/N edge masks and specialize K in K128 +# tiles. Keep all dimensions aligned to exercise the supported production path. +FLYDSL_SHAPES = [ + (512, 512, 512), + (512, 1024, 512), + (1024, 512, 1024), +] + +MXFP8_SHAPES = [ + (512, 512, 512), + (512, 1024, 512), +] + +LAYOUTS = ["TN", "NN", "NT"] + +FP8_FORMAT_COMBOS = [ + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2), + (tex.DType.kFloat8E5M2, tex.DType.kFloat8E4M3), + (tex.DType.kFloat8E5M2, tex.DType.kFloat8E5M2), +] + +FP8_FORMAT_IDS = [ + "e4m3_e4m3", + "e4m3_e5m2", + "e5m2_e4m3", + "e5m2_e5m2", +] + +REGULAR_DTYPES = [torch.float32, torch.float16, torch.bfloat16] + + +# --- Fixtures ----------------------------------------------------------------- + +@pytest.fixture(autouse=True) +def cleanup_env(): + """Save and restore FlyDSL-related environment variables between tests.""" + old_flydsl = os.environ.get("NVTE_USE_FLYDSL") + old_mxfp8 = os.environ.get("NVTE_ROCM_ENABLE_MXFP8") + + yield + + if old_flydsl is None: + os.environ.pop("NVTE_USE_FLYDSL", None) + else: + os.environ["NVTE_USE_FLYDSL"] = old_flydsl + + if old_mxfp8 is None: + os.environ.pop("NVTE_ROCM_ENABLE_MXFP8", None) + else: + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = old_mxfp8 + + +# --- Helpers ------------------------------------------------------------------ + +def get_shapes(layout, M, K, N): + """Return the A/B storage shapes used by TE's public GEMM tests.""" + if layout == "TN": + return (M, K), (N, K) + if layout == "NN": + return (M, K), (K, M) + if layout == "NT": + return (M, K), (M, K) + raise ValueError(f"Unsupported layout: {layout}") + + +def compute_pytorch_reference(A_ref, B_ref, layout): + """Compute the equivalent public-layout GEMM with ``torch.matmul``.""" + if layout == "TN": + return torch.matmul(B_ref, A_ref.T) + if layout == "NN": + return torch.matmul(B_ref, A_ref) + if layout == "NT": + return torch.matmul(B_ref.T, A_ref) + raise ValueError(f"Unsupported layout: {layout}") + + +def create_fp8_tensors(M, K, N, layout, fp8_dtype_a, fp8_dtype_b): + """Create independently typed Float8Tensor inputs and references.""" + A_shape, B_shape = get_shapes(layout, M, K, N) + A_f32 = torch.randn(A_shape, dtype=torch.float32, device="cuda") * 0.5 + B_f32 = torch.randn(B_shape, dtype=torch.float32, device="cuda") * 0.5 + + A_fp8 = Float8Quantizer( + scale=torch.full((1,), 1.0, dtype=torch.float32, device="cuda"), + amax=torch.empty((1,), dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype_a, + )(A_f32) + B_fp8 = Float8Quantizer( + scale=torch.full((1,), 1.0, dtype=torch.float32, device="cuda"), + amax=torch.empty((1,), dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype_b, + )(B_f32) + + return A_fp8, B_fp8, A_fp8.dequantize(), B_fp8.dequantize() + + +def _make_mxfp8_quantizer(fp8_dtype): + """Create one independently typed MXFP8 quantizer with both orientations.""" + quantizer = MXFP8Quantizer(fp8_dtype=fp8_dtype) + quantizer.set_usage(rowwise=True, columnwise=True) + return quantizer + + +def create_mxfp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, +): + """Create independently typed MXFP8Tensor inputs and references.""" + A_shape, B_shape = get_shapes(layout, M, K, N) + A_f32 = torch.randn(A_shape, dtype=torch.float32, device="cuda") * 0.5 + B_f32 = torch.randn(B_shape, dtype=torch.float32, device="cuda") * 0.5 + + A_mxfp8 = _make_mxfp8_quantizer(fp8_dtype_a)(A_f32) + B_mxfp8 = _make_mxfp8_quantizer(fp8_dtype_b)(B_f32) + + return ( + A_mxfp8, + B_mxfp8, + A_mxfp8.dequantize(), + B_mxfp8.dequantize(), + ) + + +def call_gemm(A, B, layout, out_dtype, use_flydsl=True): + """Call ``general_gemm`` through either FlyDSL or the native C++ path.""" + os.environ["NVTE_USE_FLYDSL"] = "1" if use_flydsl else "0" + + output, bias_grad, gelu_input, extra_output = general_gemm( + A=A, + B=B, + out_dtype=out_dtype, + layout=layout, + bias=None, + quantization_params=None, + gelu=False, + grad=False, + accumulate=False, + ) + + assert bias_grad is None + assert gelu_input is None + assert extra_output is None + return output + + +def assert_gemm_close(actual, expected, *, atol, rtol): + """Compare through FP32 so output narrowing does not hide diagnostics.""" + torch.testing.assert_close( + actual.float(), + expected.float(), + atol=atol, + rtol=rtol, + equal_nan=False, + ) + + +# ============================================================================== +# Approach 1: FlyDSL vs PyTorch torch.matmul reference +# ============================================================================== + +@pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "dtype", + REGULAR_DTYPES, + ids=["fp32", "fp16", "bf16"], +) +def test_flydsl_vs_pytorch_regular(M, K, N, layout, dtype): + """Test regular FlyDSL GEMM against an FP32 PyTorch reference.""" + torch.manual_seed(42) + + A_shape, B_shape = get_shapes(layout, M, K, N) + A = torch.randn(A_shape, dtype=dtype, device="cuda") * 0.5 + B = torch.randn(B_shape, dtype=dtype, device="cuda") * 0.5 + + output = call_gemm( + A, + B, + layout, + out_dtype=dtype, + use_flydsl=True, + ) + expected = compute_pytorch_reference(A.float(), B.float(), layout) + + assert_gemm_close(output, expected, atol=1e-3, rtol=1e-2) + + +@pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_pytorch_fp8(M, K, N, layout, fp8_format): + """Test same-format and mixed-format tensor-wise FP8 FlyDSL GEMMs.""" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + A_fp8, B_fp8, A_deq, B_deq = create_fp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, + ) + + output = call_gemm( + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + expected = compute_pytorch_reference( + A_deq.float(), + B_deq.float(), + layout, + ) + + assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) + + +@requires_mxfp8_support +@pytest.mark.parametrize("M, K, N", MXFP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_pytorch_mxfp8(M, K, N, layout, fp8_format): + """Test same-format and mixed-format MXFP8 FlyDSL GEMMs.""" + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = "1" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + A_mxfp8, B_mxfp8, A_deq, B_deq = create_mxfp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, + ) + + output = call_gemm( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + expected = compute_pytorch_reference( + A_deq.float(), + B_deq.float(), + layout, + ) + + assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) + + +# ============================================================================== +# Approach 2: FlyDSL vs native C++ ``generic_gemm`` reference +# ============================================================================== + +@pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "dtype", + REGULAR_DTYPES, + ids=["fp32", "fp16", "bf16"], +) +def test_flydsl_vs_cpp_regular(M, K, N, layout, dtype): + """Test regular FlyDSL GEMM against the native C++ backend.""" + torch.manual_seed(42) + + A_shape, B_shape = get_shapes(layout, M, K, N) + A = torch.randn(A_shape, dtype=dtype, device="cuda") * 0.5 + B = torch.randn(B_shape, dtype=dtype, device="cuda") * 0.5 + + flydsl_out = call_gemm( + A, + B, + layout, + out_dtype=dtype, + use_flydsl=True, + ) + cpp_out = call_gemm( + A, + B, + layout, + out_dtype=dtype, + use_flydsl=False, + ) + + assert_gemm_close(flydsl_out, cpp_out, atol=1e-3, rtol=1e-2) + + +@pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_cpp_fp8(M, K, N, layout, fp8_format): + """Test same-format and mixed-format FP8 against native C++.""" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + A_fp8, B_fp8, _, _ = create_fp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, + ) + + flydsl_out = call_gemm( + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + cpp_out = call_gemm( + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=False, + ) + + assert_gemm_close(flydsl_out, cpp_out, atol=5e-3, rtol=1e-2) + + +@requires_mxfp8_support +@pytest.mark.parametrize("M, K, N", MXFP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_cpp_mxfp8(M, K, N, layout, fp8_format): + """Test same-format and mixed-format MXFP8 against native C++.""" + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = "1" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + A_mxfp8, B_mxfp8, _, _ = create_mxfp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, + ) + + flydsl_out = call_gemm( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + cpp_out = call_gemm( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=False, + ) + + assert_gemm_close(flydsl_out, cpp_out, atol=5e-3, rtol=1e-2) + + +# ============================================================================== +# Batched multidimensional FP8 coverage +# ============================================================================== + +@pytest.mark.parametrize( + "batch_size, M, K, N", + [ + (2, 256, 512, 256), + (4, 256, 512, 256), + ], +) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_pytorch_fp8_multidim( + batch_size, + M, + K, + N, + fp8_format, +): + """Exercise flatten-leading-dim semantics for multidimensional FP8.""" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + + # TN layout: the wrapper flattens all leading dimensions into rows. + A_f32 = ( + torch.randn( + batch_size, + M, + K, + dtype=torch.float32, + device="cuda", + ) + * 0.5 + ) + B_f32 = ( + torch.randn( + batch_size, + N, + K, + dtype=torch.float32, + device="cuda", + ) + * 0.5 + ) + + A_fp8 = Float8Quantizer( + scale=torch.full((1,), 1.0, dtype=torch.float32, device="cuda"), + amax=torch.empty((1,), dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype_a, + )(A_f32) + B_fp8 = Float8Quantizer( + scale=torch.full((1,), 1.0, dtype=torch.float32, device="cuda"), + amax=torch.empty((1,), dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype_b, + )(B_f32) + + output = call_gemm( + A_fp8, + B_fp8, + layout="TN", + out_dtype=torch.float32, + use_flydsl=True, + ) + + A_flat = A_fp8.dequantize().reshape(-1, K) + B_flat = B_fp8.dequantize().reshape(-1, K) + expected = torch.matmul(B_flat, A_flat.T) + + assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) + + +if __name__ == "__main__": + # Quick smoke tests using one case from each supported input family. + os.environ["NVTE_USE_FLYDSL"] = "1" + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = "1" + + test_flydsl_vs_pytorch_regular( + 256, + 512, + 256, + "TN", + torch.float16, + ) + test_flydsl_vs_pytorch_fp8( + 256, + 512, + 256, + "TN", + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2), + ) + + if has_mxfp8_support: + test_flydsl_vs_pytorch_mxfp8( + 256, + 512, + 256, + "TN", + (tex.DType.kFloat8E5M2, tex.DType.kFloat8E4M3), + ) + + print("All FlyDSL GEMM smoke tests passed!") From 838e64f6c804ca1e91a10e1c0ae5b838718c524c Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 19:31:48 +0000 Subject: [PATCH 15/17] add e2e flydsl test in test_numerics --- tests/pytorch/test_numerics.py | 139 ++++++++++++++++++ .../flydsl_kernels/gemm/gemm_wrappers.py | 109 +++++++++++--- 2 files changed, 229 insertions(+), 19 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 5c9686f15..d6b8a59c0 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -1344,6 +1344,145 @@ def test_linear_accuracy(dtype, bs, model, return_bias, bias): assert_allclose(te_output, torch_output, tolerance, rtol[dtype]) +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("bs", batch_sizes) +@pytest.mark.parametrize("model", ["126m"]) +@pytest.mark.parametrize( + "fp8_recipe", + [ + None, + recipe.Float8CurrentScaling(), + recipe.DelayedScaling(), + recipe.MXFP8BlockScaling(), + ], +) +def test_linear_accuracy_flydsl( + dtype, + bs, + model, + fp8_recipe, +): + """Compare FlyDSL and native TE Linear forward, dgrad, and wgrad.""" + + if not IS_HIP_EXTENSION: + pytest.skip("FlyDSL GEMM is only supported on HIP.") + + fp8 = fp8_recipe is not None + config = model_configs[model] + + if isinstance(fp8_recipe, recipe.MXFP8BlockScaling): + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + elif fp8 and not fp8_available: + pytest.skip(reason_for_no_fp8) + + if config.max_seqlen_q % 16 != 0 and fp8: + pytest.skip("FP8 requires sequence length to be divisible by 16.") + + # Validate the GEMM backend, not quantized parameter storage. + # FlyDSL GEMM does not currently support bias. + with quantized_model_init(enabled=False, recipe=fp8_recipe): + linear_ref = Linear( + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + device="cuda", + ).eval() + + linear_flydsl = Linear( + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + device="cuda", + ).eval() + + with torch.no_grad(): + linear_flydsl.weight.copy_(linear_ref.weight) + + input_shape = ( + config.max_seqlen_q, + bs, + config.hidden_size, + ) + + inp_ref = torch.randn( + input_shape, + dtype=dtype, + device="cuda", + requires_grad=True, + ) + inp_flydsl = inp_ref.detach().clone().requires_grad_(True) + + try: + # Native TE backend. + os.environ.pop("NVTE_USE_FLYDSL", None) + + reset_rng_states() + FP8GlobalStateManager.reset() + + with autocast(enabled=fp8, recipe=fp8_recipe): + out_ref = linear_ref(inp_ref) + + out_ref.sum().backward() + torch.cuda.synchronize() + + # FlyDSL backend. + os.environ["NVTE_USE_FLYDSL"] = "1" + + reset_rng_states() + FP8GlobalStateManager.reset() + + with autocast(enabled=fp8, recipe=fp8_recipe): + out_flydsl = linear_flydsl(inp_flydsl) + + out_flydsl.sum().backward() + torch.cuda.synchronize() + + finally: + os.environ.pop("NVTE_USE_FLYDSL", None) + FP8GlobalStateManager.reset() + + atol, rtol = get_tolerances(dtype) + + if fp8: + atol = max(atol, 1e-2) + rtol = max(rtol, 1e-2) + + torch.testing.assert_close( + out_flydsl, + out_ref, + atol=atol, + rtol=rtol, + ) + + torch.testing.assert_close( + inp_flydsl.grad, + inp_ref.grad, + atol=atol, + rtol=rtol, + ) + + # Wgrad is an NT GEMM with a reduction over the flattened + # sequence/batch dimension. FlyDSL and the native TE backend may use + # different FP32 accumulation orders, so allow the small expected + # non-associative rounding difference. + wgrad_atol = atol + wgrad_rtol = rtol + + if dtype == torch.float32 and not fp8: + wgrad_atol = max(wgrad_atol, 1e-4) + wgrad_rtol = max(wgrad_rtol, 1e-4) + + torch.testing.assert_close( + linear_flydsl.weight.grad, + linear_ref.weight.grad, + atol=wgrad_atol, + rtol=wgrad_rtol, + ) + + @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("bs", batch_sizes) @pytest.mark.parametrize("model", ["small"]) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 343624d6c..ebd53d926 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -18,6 +18,39 @@ from .mxfp8_gemm import mxfp8_matmul +def _product(shape): + """Return the product of dimensions in ``shape``.""" + result = 1 + for dim in shape: + result *= dim + return result + + +def _get_gemm_output_shape(A, transa, B, transb) -> torch.Size: + """Compute TE's logical GEMM output shape. + + This matches ``getGemmOutputShape`` in the C++/Triton backends: the + physical GEMM is flattened to ``[M, N]``, while the returned tensor keeps + B's leading dimensions when ``transb`` is false. + """ + A_shape = A if isinstance(A, torch.Size) else A.shape + B_shape = B if isinstance(B, torch.Size) else B.shape + + if len(A_shape) < 2 or len(B_shape) < 2: + raise ValueError( + "FlyDSL GEMM expects both logical operands to have rank >= 2, " + f"got A={tuple(A_shape)} and B={tuple(B_shape)}" + ) + + A0 = _product(A_shape[:-1]) + A1 = A_shape[-1] + B1 = B_shape[-1] + + output_shape = [B1] if transb else list(B_shape[:-1]) + output_shape.append(A0 if transa else A1) + return torch.Size(output_shape) + + def reinterpret_as_fp8_tensor( a: torch.Tensor, dtype: tex.DType, @@ -76,11 +109,6 @@ def _validate_common_epilogue( "FlyDSL GEMM bias is not implemented" ) - if gelu or grad: - raise NotImplementedError( - "FlyDSL GEMM GELU/gradient epilogues are not implemented" - ) - def _classify_input(t): """Classify a GEMM operand for the FlyDSL backend.""" @@ -294,16 +322,23 @@ def _run_regular_gemm( f"A and B must be on the same device, got {A.device} and {B.device}" ) + output_shape = _get_gemm_output_shape(A, transa, B, transb) + a_flydsl, b_flydsl, m, n, _ = _canonicalize_blas_operands( A, transa, B, transb ) + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL {backend_name} logical output shape {tuple(output_shape)} " + f"does not match flattened GEMM shape {(m, n)}" + ) if output_dtype is None: output_dtype = dtype D = _validate_or_allocate_output( D, - shape=(m, n), + shape=output_shape, dtype=output_dtype, device=A.device, backend_name=backend_name, @@ -317,6 +352,29 @@ def _run_regular_gemm( return D +def _materialize_rowwise_from_columnwise( + transpose_data: torch.Tensor, + name: str, +) -> torch.Tensor: + """Reconstruct logical rowwise FP8 data from TE columnwise storage. + + This matches Triton's ``materialize_rowwise_from_columnwise`` exactly. + TE stores an n-D rowwise tensor ``[D0, ..., Dn-2, K]`` columnwise as + ``[K, D0, ..., Dn-2]``. Recover rowwise storage by rotating the leading + K dimension back to the tail. + """ + if transpose_data.ndim < 2: + raise ValueError( + f"{name} must have rank >= 2, got {tuple(transpose_data.shape)}" + ) + + if transpose_data.ndim == 2: + return transpose_data.transpose(0, 1).contiguous() + + perm = list(range(1, transpose_data.ndim)) + [0] + return transpose_data.permute(*perm).contiguous() + + def _get_fp8_logical_rowwise_payload(t, name): """Return logical rowwise FP8 data, matching the Triton wrapper. @@ -345,16 +403,10 @@ def _get_fp8_logical_rowwise_payload(t, name): f"{name}._transpose", ) - if transpose_data.ndim < 2: - raise ValueError( - f"{name}._transpose must have rank >= 2, " - f"got {tuple(transpose_data.shape)}" - ) - - # TE's columnwise payload represents the transpose of the logical rowwise - # tensor. Materialize rowwise storage before applying BLAS transpose flags, - # exactly as the Triton wrapper's materialize_rowwise_from_columnwise path. - return transpose_data.transpose(-2, -1).contiguous() + return _materialize_rowwise_from_columnwise( + transpose_data, + f"{name}._transpose", + ) def _select_mxfp8_data_and_scale( @@ -462,12 +514,21 @@ def _run_mxfp8( if B_data.dtype == torch.uint8: B_data = reinterpret_as_fp8_tensor(B_data, b_fp8_dtype) + output_shape = _get_gemm_output_shape( + A_data.shape, transa, B_data.shape, transb + ) + a_flydsl, b_flydsl, m, n, k = _canonicalize_blas_operands( A_data, transa, B_data, transb, ) + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL MXFP8 logical output shape {tuple(output_shape)} " + f"does not match flattened GEMM shape {(m, n)}" + ) A_scale = _flatten_mxfp8_scale(A_scale, "A") B_scale = _flatten_mxfp8_scale(B_scale, "B") @@ -525,19 +586,20 @@ def _run_mxfp8( D = _validate_or_allocate_output( D, - shape=(m, n), + shape=output_shape, dtype=output_dtype, device=a_flydsl.device, backend_name="MXFP8", ) - return mxfp8_matmul( + mxfp8_matmul( a_flydsl, a_scale, b_flydsl, b_scale, D.view(m, n), ) + return D def _run_fp8( @@ -590,9 +652,18 @@ def _run_fp8( f"scale, got dtype={scale.dtype}, shape={tuple(scale.shape)}" ) + output_shape = _get_gemm_output_shape( + A_data.shape, transa, B_data.shape, transb + ) + a_flydsl, b_flydsl, m, n, _ = _canonicalize_blas_operands( A_data, transa, B_data, transb ) + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL FP8 logical output shape {tuple(output_shape)} " + f"does not match flattened GEMM shape {(m, n)}" + ) if a_flydsl.device != b_flydsl.device: raise ValueError( @@ -602,7 +673,7 @@ def _run_fp8( D = _validate_or_allocate_output( D, - shape=(m, n), + shape=output_shape, dtype=output_dtype, device=a_flydsl.device, backend_name="FP8", From 9670650980077732083119cce7d274733b11424d Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 20:05:17 +0000 Subject: [PATCH 16/17] add flydsl gemm fallback support --- tests/pytorch/test_numerics.py | 5 ++- .../pytorch/cpp_extensions/gemm.py | 40 ++++++++++++++++--- .../pytorch/flydsl_kernels/gemm/__init__.py | 6 +-- .../pytorch/flydsl_kernels/gemm/bf16_gemm.py | 28 +++++++++++-- .../pytorch/flydsl_kernels/gemm/exceptions.py | 6 +++ .../pytorch/flydsl_kernels/gemm/fp16_gemm.py | 28 +++++++++++-- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 36 +++++++++++------ .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 26 ++++++++++-- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 25 ++++++++++-- 9 files changed, 162 insertions(+), 38 deletions(-) create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index d6b8a59c0..c9d9a5563 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -1346,7 +1346,7 @@ def test_linear_accuracy(dtype, bs, model, return_bias, bias): @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("bs", batch_sizes) -@pytest.mark.parametrize("model", ["126m"]) +@pytest.mark.parametrize("model", ["small", "126m"]) @pytest.mark.parametrize( "fp8_recipe", [ @@ -1418,6 +1418,7 @@ def test_linear_accuracy_flydsl( try: # Native TE backend. os.environ.pop("NVTE_USE_FLYDSL", None) + os.environ.pop("NVTE_FLYDSL_GEMM_WARN_FALLBACK", None) reset_rng_states() FP8GlobalStateManager.reset() @@ -1430,6 +1431,7 @@ def test_linear_accuracy_flydsl( # FlyDSL backend. os.environ["NVTE_USE_FLYDSL"] = "1" + os.environ["NVTE_FLYDSL_GEMM_WARN_FALLBACK"] = "1" reset_rng_states() FP8GlobalStateManager.reset() @@ -1442,6 +1444,7 @@ def test_linear_accuracy_flydsl( finally: os.environ.pop("NVTE_USE_FLYDSL", None) + os.environ.pop("NVTE_FLYDSL_GEMM_WARN_FALLBACK", None) FP8GlobalStateManager.reset() atol, rtol = get_tolerances(dtype) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 46f0b2ee9..52ca77379 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -10,6 +10,7 @@ import ctypes import os import functools +import warnings import torch from torch.utils.cpp_extension import IS_HIP_EXTENSION import transformer_engine_torch as tex @@ -460,18 +461,45 @@ def general_gemm( "beta": beta, } - use_gemm_flydsl = IS_HIP_EXTENSION and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0"))) + use_gemm_flydsl = ( + IS_HIP_EXTENSION + and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0"))) + ) if use_gemm_flydsl: # Lazy import keeps FlyDSL off the normal Transformer Engine import path. - from ..flydsl_kernels.gemm import te_generic_gemm_flydsl - - out, bias_grad, gelu_input, extra_output = te_generic_gemm_flydsl( - *args, **kwargs + from ..flydsl_kernels.gemm import ( + FlyDSLUnsupportedError, + te_generic_gemm_flydsl, ) + + try: + out, bias_grad, gelu_input, extra_output = te_generic_gemm_flydsl( + *args, + **kwargs, + ) + except FlyDSLUnsupportedError as exc: + warn_fallback = os.environ.get( + "NVTE_FLYDSL_GEMM_WARN_FALLBACK", + "0", + ).lower() not in ("", "0", "false", "no", "off") + + if warn_fallback: + warnings.warn( + "[FLYDSL WARNING]: FlyDSL GEMM does not support this configuration; " + f"falling back to the default backend. Reason: {exc}", + UserWarning, + stacklevel=2, + ) + + out, bias_grad, gelu_input, extra_output = tex.generic_gemm( + *args, + **kwargs, + ) else: out, bias_grad, gelu_input, extra_output = tex.generic_gemm( - *args, **kwargs + *args, + **kwargs, ) if IS_HIP_EXTENSION and use_bf16_tn_output_workaround: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py index 784d17d2f..5acdce6a2 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py @@ -4,10 +4,10 @@ """FlyDSL GEMM kernels (dense, non-grouped) for BF16/FP16/FP32/FP8/MXFP8.""" -from .gemm_wrappers import ( - te_generic_gemm_flydsl, -) +from .exceptions import FlyDSLUnsupportedError +from .gemm_wrappers import te_generic_gemm_flydsl __all__ = [ + "FlyDSLUnsupportedError", "te_generic_gemm_flydsl", ] \ No newline at end of file diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py index b5c6045c2..4201d571d 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py @@ -27,6 +27,7 @@ from flydsl.expr.typing import Vector as Vec # Transformer Engine-local FlyDSL utilities. +from .exceptions import FlyDSLUnsupportedError from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -1092,11 +1093,30 @@ def doGemm( "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " f"got {C.dtype}" ) - assert M_runtime % _BLOCK_M == 0, f"M={M_runtime} must be a multiple of {_BLOCK_M}" - assert N_runtime % _BLOCK_N == 0, f"N={N_runtime} must be a multiple of {_BLOCK_N}" - assert K_runtime % _BLOCK_K == 0, f"K={K_runtime} must be a multiple of {_BLOCK_K}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + num_k_tiles = K_runtime // _BLOCK_K - assert num_k_tiles >= 4, f"K={K_runtime} gives {num_k_tiles} K64 tiles; need at least 4" + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) assert C.shape == (M_runtime, N_runtime) if stream is None: stream = torch.cuda.current_stream() diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py b/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py new file mode 100644 index 000000000..1ae38569a --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py @@ -0,0 +1,6 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +class FlyDSLUnsupportedError(RuntimeError): + """The GEMM request is valid but unsupported by the available FlyDSL kernels.""" \ No newline at end of file diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py index 7ad7ed31f..709f76484 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py @@ -27,6 +27,7 @@ from flydsl.expr.typing import Vector as Vec # Transformer Engine-local FlyDSL utilities. +from .exceptions import FlyDSLUnsupportedError from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -1099,11 +1100,30 @@ def doGemm( "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " f"got {C.dtype}" ) - assert M_runtime % _BLOCK_M == 0, f"M={M_runtime} must be a multiple of {_BLOCK_M}" - assert N_runtime % _BLOCK_N == 0, f"N={N_runtime} must be a multiple of {_BLOCK_N}" - assert K_runtime % _BLOCK_K == 0, f"K={K_runtime} must be a multiple of {_BLOCK_K}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + num_k_tiles = K_runtime // _BLOCK_K - assert num_k_tiles >= 4, f"K={K_runtime} gives {num_k_tiles} K64 tiles; need at least 4" + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) assert C.shape == (M_runtime, N_runtime) if stream is None: stream = torch.cuda.current_stream() diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index c18cde73c..6cbd61102 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -26,6 +26,7 @@ from flydsl.expr.typing import Vector as Vec # Transformer Engine-local FlyDSL utilities. +from .exceptions import FlyDSLUnsupportedError from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -1055,19 +1056,30 @@ def doGemm( assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" assert A.dtype == torch.float32 and B.dtype == torch.float32 assert C.dtype == torch.float32 - assert M_runtime % _BLOCK_M == 0, ( - f"M={M_runtime} must be a multiple of {_BLOCK_M}" - ) - assert N_runtime % _BLOCK_N == 0, ( - f"N={N_runtime} must be a multiple of {_BLOCK_N}" - ) - assert K_runtime % _BLOCK_K == 0, ( - f"K={K_runtime} must be a multiple of {_BLOCK_K}" - ) + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + num_k_tiles = K_runtime // _BLOCK_K - assert num_k_tiles >= 4, ( - f"K={K_runtime} gives {num_k_tiles} K32 tiles; need at least 4" - ) + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) assert C.shape == (M_runtime, N_runtime) if stream is None: stream = torch.cuda.current_stream() diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index c11c6765e..4a0578f60 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -27,6 +27,8 @@ from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec +from .exceptions import FlyDSLUnsupportedError + # Transformer Engine-local FlyDSL utilities. from .fp8_gemm_utils import ( G2SLoader, @@ -1127,11 +1129,27 @@ def doGemm( f"got {C.dtype}" ) assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" - assert M_runtime % _BLOCK_M == 0, f"M={M_runtime} must be a multiple of {_BLOCK_M}" - assert N_runtime % _BLOCK_N == 0, f"N={N_runtime} must be a multiple of {_BLOCK_N}" - assert K_runtime % _BLOCK_K == 0, f"K={K_runtime} must be a multiple of {_BLOCK_K}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) num_k_tiles = K_runtime // _BLOCK_K - assert num_k_tiles >= 4, f"K={K_runtime} gives {num_k_tiles} K128 tiles; need at least 4" + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) assert A_scale_inv.dtype == torch.float32 and A_scale_inv.numel() == 1 assert B_scale_inv.dtype == torch.float32 and B_scale_inv.numel() == 1 assert C.shape == (M_runtime, N_runtime), ( diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 4e83f271b..a54dc43d7 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -30,6 +30,7 @@ from flydsl.expr.typing import Vector as Vec # Transformer Engine-local FlyDSL utilities. +from .exceptions import FlyDSLUnsupportedError from .fp8_gemm_utils import ( G2SLoader, S2RLoader, @@ -1211,11 +1212,27 @@ def do_gemm( assert A.dtype in supported_fp8_dtypes, f"unsupported A MXFP8 dtype: {A.dtype}" assert B.dtype in supported_fp8_dtypes, f"unsupported B MXFP8 dtype: {B.dtype}" assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" - assert M_runtime % _BLOCK_M == 0, f"M={M_runtime} must be a multiple of {_BLOCK_M}" - assert N_runtime % _BLOCK_N == 0, f"N={N_runtime} must be a multiple of {_BLOCK_N}" - assert K_runtime % _BLOCK_K == 0, f"K={K_runtime} must be a multiple of {_BLOCK_K}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) num_k_tiles = K_runtime // _BLOCK_K - assert num_k_tiles >= 4, f"K={K_runtime} gives {num_k_tiles} K128 tiles; need at least 4" + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) expected_as = (K_runtime // _BLOCK_K, M_runtime) expected_bs = (K_runtime // _BLOCK_K, N_runtime) assert As.dtype == torch.int32, f"As dtype {As.dtype} != torch.int32 packed scales" From ede7ace8bc264c12032d36eb04fdf39a218a2fad Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Thu, 23 Jul 2026 04:31:01 +0000 Subject: [PATCH 17/17] Add direct FP8 NN/NT FlyDSL GEMM specializations --- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 1 - .../flydsl_kernels/gemm/fp8_gemm_nn.py | 1208 ++++++++++++++++ .../flydsl_kernels/gemm/fp8_gemm_nt.py | 1225 +++++++++++++++++ .../flydsl_kernels/gemm/fp8_gemm_utils.py | 53 +- .../flydsl_kernels/gemm/gemm_wrappers.py | 228 ++- 5 files changed, 2704 insertions(+), 11 deletions(-) create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 4a0578f60..b29fa67d8 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -1181,4 +1181,3 @@ def doGemm( N_runtime, stream=stream, ) - diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py new file mode 100644 index 000000000..122d1ad95 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py @@ -0,0 +1,1208 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL tensor-wise FP8 4-wave NN GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K128 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes independently typed FP8 E4M3 or E5M2 A/B tensors shaped [K, M] and +[N, K], one FP32 inverse +scale per operand, and writes float16, bfloat16, or float32 C shaped [M, N]. The public +``fp8_matmul`` entry point accepts an NN contract and +performs the required private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +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 .exceptions import FlyDSLUnsupportedError + +# Transformer Engine-local FlyDSL utilities. +from .fp8_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_linear_128x128, + compute_global_swizzle, + make_fp8_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 128 + +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +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 = 16 + +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 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp8_pp_smem_a0" +LDS_SYM_A1 = "fp8_pp_smem_a1" +LDS_SYM_B0 = "fp8_pp_smem_b0" +LDS_SYM_B1 = "fp8_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 128 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +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] + ) + + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + +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 _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + 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), + ) + + +def _compile_kernel( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. + + ``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_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + + fp8_input_types = { + torch.float8_e4m3fn: (fx.Float8E4M3FN, 0), + torch.float8_e5m2: (fx.Float8E5M2, 1), + } + try: + a_fx_dtype, a_matrix_format = fp8_input_types[a_fp8_dtype] + b_fx_dtype, b_matrix_format = fp8_input_types[b_fp8_dtype] + except KeyError as exc: + raise TypeError( + "FlyDSL FP8 input dtype must be torch.float8_e4m3fn or " + f"torch.float8_e5m2, got A={a_fp8_dtype}, B={b_fp8_dtype}" + ) from exc + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {output_dtype}" + ) + NUM_THREADS = 256 + WARP_SIZE = 64 + + 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 = 16 + + 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 + + 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_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[a_fx_dtype, LDS_ELEMS_HALF, 16] + a0_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + b0_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b0_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_1: fx.Array[b_fx_dtype, 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, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + 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) + + a_f8_ir_t = a_fx_dtype.ir_type + b_f8_ir_t = b_fx_dtype.ir_type + gA = make_fp8_buffer_tensor(A, a_f8_ir_t) + gB = make_fp8_buffer_tensor(B, 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)) + a_scale_rsrc = buffer_ops.create_buffer_resource(A_scale_inv, max_size=True) + b_scale_rsrc = buffer_ops.create_buffer_resource(B_scale_inv, max_size=True) + output_scale = ( + buffer_ops.buffer_load(a_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + * buffer_ops.buffer_load(b_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + ) + 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) + + # 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_linear_128x128(lane, wave_id, c_m, LOAD_PASSES_HALF) + 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, a_f8_ir_t, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_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) + 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(output_element_bytes) + 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 _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) + + 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") + 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) + + def hot_loop_scheduler_q_refill_2n(): + 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(): + # One logical A K64 half is two ds_read_b64_tr_b8 instructions. + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(2) + rocdl.sched_mfma(2) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + for _ in range_constexpr(8): + rocdl.sched_dsrd(2) + rocdl.sched_mfma(4) + rocdl.sched_barrier(0) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # NN A is contiguous [K, M]. Stage a physical row-major [K128, M128] + # half-page without the TN XOR swizzle; S2R performs the transpose. + m_base = bx_m_idx + fx.Index(subtile * (BLOCK_M // 2)) + global_base = k_base * fx.Index(c_m) + m_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 + 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_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_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, 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 + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Issue ordinary FP8 MFMA into the fixed physical accumulator bank.""" + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + """Final-page ordinary FP8 MFMA with independently named AGPR source/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), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + f"$0, $1, " + f"a[{old_pin}:{old_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" + ), + ( + f"v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}}," + f"~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}" + ), + has_side_effects=True, + ) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + 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_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 + value = Vec(acc)[ii] * output_scale + if output_dtype != torch.float32: + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, 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, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + # Physical LDS A is [K128, M128]. The CDNA4 transpose-read lane + # mapping addresses 8 K rows x 16 M columns per K64 operand half. + # + # The wave-level base follows the documented transpose-load layout: + # k_row = lane_div_32 * 4 + (lane_mod_16 // 4) -> 0..7 + # m_col = m_tile + n-group/lane-in-quad offset -> 0..15 + # Two addresses separated by 32 K rows produce the complementary + # halves required for the complete K64 i32x4 operand. + local_m_tile = ( + (reg_subtile_m_idx0 + fx.Index(sm * 2)) * fx.Index(SUBTILE_M) + + fx.Index(mi * MFMA_M) + - fx.Index(sm * (BLOCK_M // 2)) + ) + k_row = (fx.Index(lane) // fx.Index(32)) * fx.Index(4) + (lane_mod_16 // fx.Index(4)) + m_col = local_m_tile + ((fx.Index(lane) // fx.Index(16)) % fx.Index(2)) * fx.Index(8) + (lane_mod_16 % fx.Index(4)) * fx.Index(2) + k_half_base = fx.Index(half * 64) + first = (k_half_base + k_row) * fx.Index(BLOCK_M // 2) + m_col + second = (k_half_base + k_row + fx.Index(32)) * fx.Index(BLOCK_M // 2) + m_col + return s2r.load_one_transpose(lds_a[sm], fx.Int32(first), fx.Int32(second)) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): + + # 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) + + # 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 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + # 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) + 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, b00, b01) + + 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, b02, b03) + + 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, b00, b01) + + 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, b02, b03) + + 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, b00, b01) + + 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, b02, b03) + + 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, b00, b01) + + 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, b02, b03) + + 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 = 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) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n(_acc_idx(1, 0, 0), a00, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n(_acc_idx(1, 0, 2), a00, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + mfma_2n(_acc_idx(1, 1, 0), a01, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + mfma_2n(_acc_idx(1, 1, 2), a01, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n(_acc_idx(1, 2, 0), a02, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n(_acc_idx(1, 2, 2), a02, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + mfma_2n(_acc_idx(1, 3, 0), a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + mfma_2n(_acc_idx(1, 3, 2), a03, b12, b13) + 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, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # 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], + ) + 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. + 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) + 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) + rocdl.sched_barrier(0) + + a0_regs = load_a_subtile_regs(lds_a0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K128 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + else: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # 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_inv, + B_scale_inv, + 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), stream=stream) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch_nn( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + a_fp8_dtype, + b_fp8_dtype, + output_dtype, + use_xcd_remap=use_xcd_remap, + ) + + + +def fp8_matmul( + a: torch.Tensor, + a_scale_inv: torch.Tensor, + b: torch.Tensor, + b_scale_inv: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing NN tensor-wise FP8 adapter. + + Public/backend contract: + a: [K, M] FP8 E4M3 or E5M2 activation payload + a_scale_inv: one-element FP32 inverse quantization scale + b: [N, K] FP8 E4M3 or E5M2 weight payload + b_scale_inv: one-element FP32 inverse quantization scale + c: [M, N] float16, bfloat16, or float32 output + + The NN core consumes TE's existing physical payloads directly: + A is contiguous columnwise storage [K, M] and B is contiguous rowwise + storage [N, K]. No transpose or materialization is performed. + """ + if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): + raise TypeError("FlyDSL FP8 GEMM expects plain torch.Tensor payloads") + + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL FP8 NN expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + if a.dtype not in supported_fp8_dtypes or b.dtype not in supported_fp8_dtypes: + raise TypeError( + "FlyDSL FP8 GEMM expects E4M3 or E5M2 payloads, " + f"got A={a.dtype} and B={b.dtype}" + ) + + k, m = a.shape + n, kb = b.shape + if kb != k: + raise ValueError( + f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + + for name, scale in ( + ("A_scale_inv", a_scale_inv), + ("B_scale_inv", b_scale_inv), + ): + if not isinstance(scale, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if scale.dtype != torch.float32 or scale.numel() != 1: + raise TypeError( + f"{name} must contain exactly one FP32 value, got " + f"dtype={scale.dtype}, shape={tuple(scale.shape)}" + ) + + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {c.dtype}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL FP8 requires contiguous output storage") + + tensors = (a, b, a_scale_inv, b_scale_inv, c) + if any(t.device != a.device for t in tensors[1:]): + raise ValueError( + "A, B, inverse scales, and C must be on the same device" + ) + + if not a.is_contiguous(): + raise ValueError( + "FlyDSL FP8 NN requires contiguous A [K, M] storage; " + "refusing to materialize a replacement" + ) + if not b.is_contiguous(): + raise ValueError( + "FlyDSL FP8 NN requires contiguous B [N, K] storage; " + "refusing to materialize a replacement" + ) + + doGemm( + a, + b, + c, + a_scale_inv, + b_scale_inv, + stream=stream, + ) + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + A_scale_inv: torch.Tensor, + B_scale_inv: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" + K_runtime, M_runtime = A.shape + N_runtime, Kb_runtime = B.shape + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + assert A.dtype in supported_fp8_dtypes, f"unsupported A FP8 dtype: {A.dtype}" + assert B.dtype in supported_fp8_dtypes, f"unsupported B FP8 dtype: {B.dtype}" + assert C.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + assert A_scale_inv.dtype == torch.float32 and A_scale_inv.numel() == 1 + assert B_scale_inv.dtype == torch.float32 and B_scale_inv.numel() == 1 + assert C.shape == (M_runtime, N_runtime), ( + f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" + ) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.view(torch.uint8).contiguous().view(-1) + B_arg = B.view(torch.uint8).contiguous().view(-1) + C_arg = C.contiguous().view(-1) + A_scale_arg = A_scale_inv.contiguous().view(-1) + B_scale_arg = B_scale_inv.contiguous().view(-1) + + launch = _cached_launch_nn( + int(K_runtime), + A.dtype, + B.dtype, + C.dtype, + bool(use_xcd_remap), + ) + launch( + A_arg, + B_arg, + C_arg, + A_scale_arg, + B_scale_arg, + M_runtime, + N_runtime, + stream=stream, + ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py new file mode 100644 index 000000000..975fd898d --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py @@ -0,0 +1,1225 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL tensor-wise FP8 4-wave NT GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K128 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes independently typed FP8 E4M3 or E5M2 A/B tensors shaped [K, M] and +[K, N], one FP32 inverse +scale per operand, and writes float16, bfloat16, or float32 C shaped [M, N]. The public +``fp8_matmul`` entry point accepts an NT contract and +performs the required private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +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 .exceptions import FlyDSLUnsupportedError + +# Transformer Engine-local FlyDSL utilities. +from .fp8_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_linear_128x128, + compute_global_swizzle, + make_fp8_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 128 + +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +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 = 16 + +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 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp8_pp_smem_a0" +LDS_SYM_A1 = "fp8_pp_smem_a1" +LDS_SYM_B0 = "fp8_pp_smem_b0" +LDS_SYM_B1 = "fp8_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 128 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +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] + ) + + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + +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 _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + 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), + ) + + +def _compile_kernel( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. + + ``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_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + + fp8_input_types = { + torch.float8_e4m3fn: (fx.Float8E4M3FN, 0), + torch.float8_e5m2: (fx.Float8E5M2, 1), + } + try: + a_fx_dtype, a_matrix_format = fp8_input_types[a_fp8_dtype] + b_fx_dtype, b_matrix_format = fp8_input_types[b_fp8_dtype] + except KeyError as exc: + raise TypeError( + "FlyDSL FP8 input dtype must be torch.float8_e4m3fn or " + f"torch.float8_e5m2, got A={a_fp8_dtype}, B={b_fp8_dtype}" + ) from exc + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {output_dtype}" + ) + NUM_THREADS = 256 + WARP_SIZE = 64 + + 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 = 16 + + 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 + + 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_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[a_fx_dtype, LDS_ELEMS_HALF, 16] + a0_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + b0_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b0_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_1: fx.Array[b_fx_dtype, 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, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + 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) + + a_f8_ir_t = a_fx_dtype.ir_type + b_f8_ir_t = b_fx_dtype.ir_type + gA = make_fp8_buffer_tensor(A, a_f8_ir_t) + gB = make_fp8_buffer_tensor(B, 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)) + a_scale_rsrc = buffer_ops.create_buffer_resource(A_scale_inv, max_size=True) + b_scale_rsrc = buffer_ops.create_buffer_resource(B_scale_inv, max_size=True) + output_scale = ( + buffer_ops.buffer_load(a_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + * buffer_ops.buffer_load(b_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + ) + 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) + + # 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_linear_128x128(lane, wave_id, c_m, LOAD_PASSES_HALF) + gl_off_b = compute_global_linear_128x128(lane, wave_id, c_n, LOAD_PASSES_HALF) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_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) + 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(output_element_bytes) + 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 _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) + + 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") + 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) + + def hot_loop_scheduler_q_refill_2n(): + 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(): + # One logical transposed K64 half is two ds_read_b64_tr_b8 instructions. + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(2) + rocdl.sched_mfma(2) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + for _ in range_constexpr(8): + rocdl.sched_dsrd(2) + rocdl.sched_mfma(4) + rocdl.sched_barrier(0) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # NT A is contiguous [K, M]. Stage a physical row-major [K128, M128] + # half-page without the TN XOR swizzle; S2R performs the transpose. + m_base = bx_m_idx + fx.Index(subtile * (BLOCK_M // 2)) + global_base = k_base * fx.Index(c_m) + m_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): + # NT B is contiguous [K, N]. Stage a physical row-major [K128, N128] + # half-page without XOR swizzling; S2R performs the transpose. + n_base = by_n_idx + fx.Index(subtile * (BLOCK_N // 2)) + global_base = k_base * fx.Index(c_n) + n_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_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_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_half(lds_b, local_row, half, k_half): + # B is physically [K, N]. Each LDS half-page is [K128, N128]. + # One K64 MFMA half requires two ds_read_b64_tr_b8 instructions, + # exactly like the transposed A path. + half_col = local_row - fx.Index(half * (BLOCK_N // 2)) + k_base = fx.Index(k_half * 64) + first = k_base * fx.Index(BLOCK_N // 2) + half_col + second = first + fx.Index(32 * (BLOCK_N // 2)) + return s2r.load_one_transpose( + lds_b[half], + fx.Int32(first), + fx.Int32(second), + ) + + def load_b_frag(lds_b, local_row, half): + x0 = load_b_frag_half(lds_b, local_row, half, 0) + x1 = load_b_frag_half(lds_b, local_row, half, 1) + return pack_frag_halves(x0, x1) + + 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): + """Issue ordinary FP8 MFMA into the fixed physical accumulator bank.""" + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + """Final-page ordinary FP8 MFMA with independently named AGPR source/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), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + f"$0, $1, " + f"a[{old_pin}:{old_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" + ), + ( + f"v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}}," + f"~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}" + ), + has_side_effects=True, + ) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + 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_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 + value = Vec(acc)[ii] * output_scale + if output_dtype != torch.float32: + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, 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, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + # Physical LDS A is [K128, M128]. The CDNA4 transpose-read lane + # mapping addresses 8 K rows x 16 M columns per K64 operand half. + # + # The wave-level base follows the documented transpose-load layout: + # k_row = lane_div_32 * 4 + (lane_mod_16 // 4) -> 0..7 + # m_col = m_tile + n-group/lane-in-quad offset -> 0..15 + # Two addresses separated by 32 K rows produce the complementary + # halves required for the complete K64 i32x4 operand. + local_m_tile = ( + (reg_subtile_m_idx0 + fx.Index(sm * 2)) * fx.Index(SUBTILE_M) + + fx.Index(mi * MFMA_M) + - fx.Index(sm * (BLOCK_M // 2)) + ) + k_row = (fx.Index(lane) // fx.Index(32)) * fx.Index(4) + (lane_mod_16 // fx.Index(4)) + m_col = local_m_tile + ((fx.Index(lane) // fx.Index(16)) % fx.Index(2)) * fx.Index(8) + (lane_mod_16 % fx.Index(4)) * fx.Index(2) + k_half_base = fx.Index(half * 64) + first = (k_half_base + k_row) * fx.Index(BLOCK_M // 2) + m_col + second = (k_half_base + k_row + fx.Index(32)) * fx.Index(BLOCK_M // 2) + m_col + return s2r.load_one_transpose(lds_a[sm], fx.Int32(first), fx.Int32(second)) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): + + # 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) + + # 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 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + # 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) + 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, b00, b01) + + 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, b02, b03) + + 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, b00, b01) + + 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, b02, b03) + + 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, b00, b01) + + 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, b02, b03) + + 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, b00, b01) + + 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, b02, b03) + + 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 = 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) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n(_acc_idx(1, 0, 0), a00, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n(_acc_idx(1, 0, 2), a00, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + mfma_2n(_acc_idx(1, 1, 0), a01, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + mfma_2n(_acc_idx(1, 1, 2), a01, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n(_acc_idx(1, 2, 0), a02, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n(_acc_idx(1, 2, 2), a02, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + mfma_2n(_acc_idx(1, 3, 0), a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + mfma_2n(_acc_idx(1, 3, 2), a03, b12, b13) + 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, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # 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], + ) + 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. + 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) + 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) + rocdl.sched_barrier(0) + + a0_regs = load_a_subtile_regs(lds_a0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K128 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + else: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # 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_inv, + B_scale_inv, + 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), stream=stream) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch_nt( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + a_fp8_dtype, + b_fp8_dtype, + output_dtype, + use_xcd_remap=use_xcd_remap, + ) + + + +def fp8_matmul( + a: torch.Tensor, + a_scale_inv: torch.Tensor, + b: torch.Tensor, + b_scale_inv: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing NT tensor-wise FP8 adapter. + + Public/backend contract: + a: [K, M] FP8 E4M3 or E5M2 activation payload + a_scale_inv: one-element FP32 inverse quantization scale + b: [K, N] FP8 E4M3 or E5M2 weight payload + b_scale_inv: one-element FP32 inverse quantization scale + c: [M, N] float16, bfloat16, or float32 output + + The NT core consumes TE's existing physical payloads directly: + A and B are contiguous columnwise payloads [K, M] and [K, N]. + No transpose or materialization is performed. + """ + if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): + raise TypeError("FlyDSL FP8 GEMM expects plain torch.Tensor payloads") + + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL FP8 NT expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + if a.dtype not in supported_fp8_dtypes or b.dtype not in supported_fp8_dtypes: + raise TypeError( + "FlyDSL FP8 GEMM expects E4M3 or E5M2 payloads, " + f"got A={a.dtype} and B={b.dtype}" + ) + + k, m = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + + for name, scale in ( + ("A_scale_inv", a_scale_inv), + ("B_scale_inv", b_scale_inv), + ): + if not isinstance(scale, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if scale.dtype != torch.float32 or scale.numel() != 1: + raise TypeError( + f"{name} must contain exactly one FP32 value, got " + f"dtype={scale.dtype}, shape={tuple(scale.shape)}" + ) + + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {c.dtype}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL FP8 requires contiguous output storage") + + tensors = (a, b, a_scale_inv, b_scale_inv, c) + if any(t.device != a.device for t in tensors[1:]): + raise ValueError( + "A, B, inverse scales, and C must be on the same device" + ) + + if not a.is_contiguous(): + raise ValueError( + "FlyDSL FP8 NT requires contiguous A [K, M] storage; " + "refusing to materialize a replacement" + ) + if not b.is_contiguous(): + raise ValueError( + "FlyDSL FP8 NT requires contiguous B [K, N] storage; " + "refusing to materialize a replacement" + ) + + doGemm( + a, + b, + c, + a_scale_inv, + b_scale_inv, + stream=stream, + ) + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + A_scale_inv: torch.Tensor, + B_scale_inv: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" + K_runtime, M_runtime = A.shape + Kb_runtime, N_runtime = B.shape + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + assert A.dtype in supported_fp8_dtypes, f"unsupported A FP8 dtype: {A.dtype}" + assert B.dtype in supported_fp8_dtypes, f"unsupported B FP8 dtype: {B.dtype}" + assert C.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + assert A_scale_inv.dtype == torch.float32 and A_scale_inv.numel() == 1 + assert B_scale_inv.dtype == torch.float32 and B_scale_inv.numel() == 1 + assert C.shape == (M_runtime, N_runtime), ( + f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" + ) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.view(torch.uint8).contiguous().view(-1) + B_arg = B.view(torch.uint8).contiguous().view(-1) + C_arg = C.contiguous().view(-1) + A_scale_arg = A_scale_inv.contiguous().view(-1) + B_scale_arg = B_scale_inv.contiguous().view(-1) + + launch = _cached_launch_nt( + int(K_runtime), + A.dtype, + B.dtype, + C.dtype, + bool(use_xcd_remap), + ) + launch( + A_arg, + B_arg, + C_arg, + A_scale_arg, + B_scale_arg, + M_runtime, + N_runtime, + stream=stream, + ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py index a8bdfb717..b8ed21535 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py @@ -2,10 +2,12 @@ # Copyright (c) 2025 FlyDSL Project Contributors import flydsl.expr as fx -from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm as _llvm, vector from flydsl._mlir.dialects.fly_rocdl import TargetAddressSpace from flydsl.expr import arith, const_expr, range_constexpr, rocdl from flydsl.expr.typing import Vector as Vec +from flydsl.expr.utils.arith import _to_raw as as_mlir_value # ceildiv is the canonical cdiv from the shared layer def cdiv(numer: int, denom: int) -> int: @@ -71,6 +73,23 @@ def compute_global_swizzle(lane_id, wave_id, K, n_rounds, preshuffled): return offsets +def compute_global_linear_128x128(lane_id, wave_id, leading_dim, n_rounds): + """Offsets for an unswizzled row-major 128x128 tile. + + This uses the same 16-byte/thread DMA decomposition as + ``compute_global_swizzle`` but does not XOR-permute the logical source + coordinates. It is used by the NN A path, whose LDS page is physically + [K128, M128] for the CDNA4 transpose-read instruction. + """ + offsets = [] + n_waves = fx.block_dim.x // 64 + for round in range_constexpr(n_rounds): + row = lane_id // 8 + wave_id * 8 + round * (n_waves * 8) + col = (lane_id % 8) * 16 + offsets.append(row * leading_dim + col) + return offsets + + class G2SLoader: def __init__(self, gl_src, gl_offsets, n_load_steps, lds_dtype, wave_id): self.g2lds_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128) @@ -139,6 +158,36 @@ def load_one(self, lds_src, lds_offset): v = self._vec_load_16xf8(lds_src, lds_offset) return v.bitcast(fx.Int32) + def _ds_read_b64_tr_b8(self, lds_src, byte_offset): + """Issue one gfx950 ``ds_read_b64_tr_b8`` and return i32x2. + + The inline-asm output uses one even-aligned 64-bit VGPR tuple. The + compiler owns allocation of the ``=v`` tuple; the memory clobber keeps + the operation ordered with respect to LDS traffic. + """ + base_i32 = fx.Int32(fx.ptrtoint(lds_src.ptr)) + addr_i32 = base_i32 + fx.Int32(byte_offset) + raw_type = ir.VectorType.get([2], ir.IntegerType.get_signless(32)) + raw = _llvm.inline_asm( + raw_type, + [as_mlir_value(addr_i32)], + "ds_read_b64_tr_b8 $0, $1\n", + "=v,v,~{memory}", + has_side_effects=True, + ) + return Vec(vector.BitCastOp(raw_type, raw).result, (2,), fx.Int32) + + def load_one_transpose(self, lds_src, first_byte_offset, second_byte_offset): + """Load one K64 FP8 MFMA operand half from physical LDS [K, M]. + + CDNA4 requires two ``ds_read_b64_tr_b8`` instructions for the complete + K64 operand. Each instruction returns i32x2; concatenation preserves the + existing i32x4 half-fragment interface used by the GEMM hot loop. + """ + lo = self._ds_read_b64_tr_b8(lds_src, first_byte_offset) + hi = self._ds_read_b64_tr_b8(lds_src, second_byte_offset) + return lo.shuffle(hi, [0, 1, 2, 3]) + class StoreC: def __init__(self, A_scale, B_scale, C, c_rows, c_cols, c_idx_fn, n_tiles_a, n_tiles_b): @@ -259,4 +308,4 @@ def call(self, a, b, c, *, set_prio=True): def call_one(self, a, b, c, i, j): assert i < self.n_tiles_a and j < self.n_tiles_b - return self._do_mma(a[i], b[j], c[self.idx(i, j)]) \ No newline at end of file + return self._do_mma(a[i], b[j], c[self.idx(i, j)]) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index ebd53d926..5beff5a75 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -15,6 +15,8 @@ from .fp16_gemm import fp16_matmul from .fp32_gemm import fp32_matmul from .fp8_gemm import fp8_matmul +from .fp8_gemm_nn import fp8_matmul as fp8_matmul_nn +from .fp8_gemm_nt import fp8_matmul as fp8_matmul_nt from .mxfp8_gemm import mxfp8_matmul @@ -611,7 +613,17 @@ def _run_fp8( *, output_dtype: torch.dtype, ): - """Run tensor-wise E4M3/E5M2 FP8 combinations for TN/NN/NT.""" + """Run tensor-wise E4M3/E5M2 FP8 combinations for TN/NN/NT. + + NN and NT are dispatched directly from TE's existing physical + representations without transpose kernels or materialized payloads: + + - NN core: [K, M] x [N, K] + - NT core: [K, M] x [K, N] + + TN retains the shared canonicalized path through + ``fp8_gemm.fp8_matmul``. + """ a_fp8_dtype = getattr(A, "_fp8_dtype", None) b_fp8_dtype = getattr(B, "_fp8_dtype", None) supported_fp8_dtypes = ( @@ -632,12 +644,6 @@ def _run_fp8( "FlyDSL GEMM does not support transa=True, transb=True (TT)" ) - # Match Triton's regular-FP8 handling: establish logical rowwise - # payloads first, then apply the same shared BLAS-to-row-major - # canonicalization used for FP16/BF16/FP32. - A_data = _get_fp8_logical_rowwise_payload(A, "A") - B_data = _get_fp8_logical_rowwise_payload(B, "B") - A_scale_inv = getattr(A, "_scale_inv", None) B_scale_inv = getattr(B, "_scale_inv", None) for name, scale in ( @@ -652,6 +658,212 @@ def _run_fp8( f"scale, got dtype={scale.dtype}, shape={tuple(scale.shape)}" ) + # TE exposes GEMM operands in BLAS/column-major convention. The + # row-major FlyDSL result is formed from the swapped operands: + # + # flydsl_a = op(B) + # flydsl_b = op(A) + # + # For NN, the dedicated kernel consumes: + # + # flydsl_a physical [K, M] = B columnwise storage + # flydsl_b physical [N, K] = A columnwise storage + # + # Here FlyDSL M is TE's n and FlyDSL N is TE's m, so the kernel writes + # the existing TE output allocation in its ordinary [M, N] view. Both + # payloads already exist; this path performs no transpose or materialization. + if not transa and not transb: + if not _valid_fp8_transpose(B): + raise RuntimeError( + "FlyDSL FP8 NN requires valid B columnwise (_transpose) storage" + ) + if not _valid_fp8_transpose(A): + raise RuntimeError( + "FlyDSL FP8 NN requires valid A columnwise (_transpose) storage" + ) + + a_flydsl = _reinterpret_fp8_payload( + B._transpose, + b_fp8_dtype, + "B._transpose", + ) + b_flydsl = _reinterpret_fp8_payload( + A._transpose, + a_fp8_dtype, + "A._transpose", + ) + + if a_flydsl.ndim != 2 or b_flydsl.ndim != 2: + raise ValueError( + "FlyDSL FP8 NN direct path expects rank-2 columnwise storage, " + f"got B._transpose={tuple(a_flydsl.shape)} and " + f"A._transpose={tuple(b_flydsl.shape)}" + ) + if not a_flydsl.is_contiguous() or not b_flydsl.is_contiguous(): + raise ValueError( + "FlyDSL FP8 NN requires contiguous TE columnwise storage; " + "refusing to materialize replacement operands" + ) + + k, m = a_flydsl.shape + n, kb = b_flydsl.shape + if kb != k: + raise ValueError( + "FlyDSL FP8 NN storage mismatch after BLAS operand swap: " + f"B._transpose{tuple(a_flydsl.shape)} and " + f"A._transpose{tuple(b_flydsl.shape)}" + ) + + # Float8TensorStorage does not expose a public ``shape`` attribute. + # The direct NN operands already determine the flattened kernel output + # shape exactly. Preserve TE's preallocated logical output shape when + # one is provided; otherwise use the flattened [M, N] shape. + output_shape = ( + D.shape + if D is not None + else torch.Size((m, n)) + ) + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL FP8 NN logical output shape {tuple(output_shape)} " + f"does not match kernel shape {(m, n)}" + ) + + if a_flydsl.device != b_flydsl.device: + raise ValueError( + f"A and B must be on the same device, got " + f"{a_flydsl.device} and {b_flydsl.device}" + ) + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=a_flydsl.device, + backend_name="FP8 NN", + ) + + # Scales follow the swapped FlyDSL operands. + fp8_matmul_nn( + a_flydsl, + B_scale_inv, + b_flydsl, + A_scale_inv, + D.view(m, n), + ) + return D + + # For TE NT (transa=False, transb=True), the dedicated kernel consumes + # both swapped operands directly from TE columnwise storage: + # + # kernel A physical [K, M] = B._transpose + # kernel B physical [K, N] = A._transpose + # + # Both operands are therefore staged as physical K-major tiles and read + # from LDS with ``ds_read_b64_tr_b8``. No torch transpose, + # ``.contiguous()``, or temporary FP8 payload is introduced. + if not transa and transb: + if not _valid_fp8_transpose(B): + raise RuntimeError( + "FlyDSL FP8 NT requires valid B columnwise (_transpose) storage" + ) + if not _valid_fp8_transpose(A): + raise RuntimeError( + "FlyDSL FP8 NT requires valid A columnwise (_transpose) storage" + ) + + # TE columnwise payloads are contiguous transposes of the logical + # rowwise tensors. After the BLAS operand swap, their exposed shapes are: + # + # B._transpose: [M, K] + # A._transpose: [N, K] + # + # The NT kernel consumes the same physical bytes as: + # + # kernel A: [K, M] + # kernel B: [K, N] + # + # Reinterpret only the 2-D shape. ``view`` is zero-copy and preserves + # the exact columnwise allocation; no torch transpose or materialization + # is performed. + b_columnwise = _reinterpret_fp8_payload( + B._transpose, + b_fp8_dtype, + "B._transpose", + ) + a_columnwise = _reinterpret_fp8_payload( + A._transpose, + a_fp8_dtype, + "A._transpose", + ) + + if b_columnwise.ndim != 2 or a_columnwise.ndim != 2: + raise ValueError( + "FlyDSL FP8 NT direct path expects rank-2 columnwise storage, " + f"got B._transpose={tuple(b_columnwise.shape)} and " + f"A._transpose={tuple(a_columnwise.shape)}" + ) + if not b_columnwise.is_contiguous() or not a_columnwise.is_contiguous(): + raise ValueError( + "FlyDSL FP8 NT requires contiguous TE columnwise storage; " + "refusing to materialize replacement operands" + ) + + m, k = b_columnwise.shape + n, ka = a_columnwise.shape + if ka != k: + raise ValueError( + "FlyDSL FP8 NT columnwise K mismatch after BLAS operand swap: " + f"B._transpose{tuple(b_columnwise.shape)} and " + f"A._transpose{tuple(a_columnwise.shape)}" + ) + + a_flydsl = b_columnwise.view(k, m) + b_flydsl = a_columnwise.view(k, n) + + # Float8TensorStorage does not expose a public ``shape`` attribute. + # Preserve TE's preallocated logical output shape when available. + output_shape = ( + D.shape + if D is not None + else torch.Size((m, n)) + ) + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL FP8 NT logical output shape {tuple(output_shape)} " + f"does not match kernel shape {(m, n)}" + ) + + if a_flydsl.device != b_flydsl.device: + raise ValueError( + f"A and B must be on the same device, got " + f"{a_flydsl.device} and {b_flydsl.device}" + ) + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=a_flydsl.device, + backend_name="FP8 NT", + ) + + # Scales follow the BLAS-swapped kernel operands. + fp8_matmul_nt( + a_flydsl, + B_scale_inv, + b_flydsl, + A_scale_inv, + D.view(m, n), + ) + return D + + # Match Triton's regular-FP8 handling: establish logical rowwise + # payloads first, then apply the same shared BLAS-to-row-major + # canonicalization used for FP16/BF16/FP32. + A_data = _get_fp8_logical_rowwise_payload(A, "A") + B_data = _get_fp8_logical_rowwise_payload(B, "B") + output_shape = _get_gemm_output_shape( A_data.shape, transa, B_data.shape, transb ) @@ -904,4 +1116,4 @@ def te_generic_gemm_flydsl( "FlyDSL GEMM currently supports only MXFP8, tensor-wise E4M3 FP8, " "BF16, FP16, or FP32 inputs; " f"got A={A.dtype} and B={B.dtype}" - ) + ) \ No newline at end of file