From cf54ce65c7f662ff35423fa2fae8557179731ee3 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 00:33:21 +0000 Subject: [PATCH 01/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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 From 842bfa242fb5665057bd76cdef9303aeae4e014d Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 28 Jul 2026 13:31:02 +0000 Subject: [PATCH 18/31] feat(flydsl): add FP8 NN and NT GEMM specializations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add dedicated FlyDSL FP8 NN and NT kernels alongside the existing TN path. * dispatch TN, NN, and NT to specialized kernels * select matching TE rowwise or columnwise storage without copies * preserve operand scales and independent FP8 dtypes * derive M/N/K from each kernel’s physical layout * preserve TE output shapes while flattening only for launch * validate unsupported layouts and shapes for controlled fallback --- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 29 +- .../flydsl_kernels/gemm/fp8_gemm_nn.py | 236 ++++---- .../flydsl_kernels/gemm/fp8_gemm_nt.py | 274 +++++---- .../flydsl_kernels/gemm/fp8_gemm_utils.py | 134 ++++- .../flydsl_kernels/gemm/gemm_wrappers.py | 541 +++++++++--------- 5 files changed, 648 insertions(+), 566 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index b29fa67d8..45c7ad66e 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -5,12 +5,10 @@ """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 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. +hand-unrolled. M/N are runtime launch dimensions. The public entry point and private optimized core consume independently typed +FP8 E4M3 or E5M2 A/B tensors shaped [M, K] and [N, K], one FP32 inverse scale +per operand, and write float16, bfloat16, or float32 C shaped [M, N]. Operand +normalization is performed by the Transformer Engine wrapper. This module imports ``flydsl`` at import time and must therefore be imported lazily only after FlyDSL availability has been confirmed. @@ -1023,17 +1021,18 @@ def fp8_matmul( c: torch.Tensor, stream=None, ): - """TE-facing TN tensor-wise FP8 adapter. + """Launch TN tensor-wise FP8 GEMM using final kernel operand order. - Public/backend contract: + Contract: a: [M, K] 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: [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 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]. + The wrapper is responsible for adapting TE's logical [K, N] operand into + the core's row-major [N, K] representation. No operand swap or transpose + is performed in this module. """ if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): raise TypeError("FlyDSL FP8 GEMM expects plain torch.Tensor payloads") @@ -1055,7 +1054,7 @@ def fp8_matmul( ) m, k = a.shape - kb, n = b.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)}" @@ -1089,13 +1088,9 @@ def fp8_matmul( "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, + b, c, a_scale_inv, b_scale_inv, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py index 122d1ad95..b24d061b6 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py @@ -2,15 +2,20 @@ # # See LICENSE for license information. -"""FlyDSL tensor-wise FP8 4-wave NN GEMM kernel for Transformer Engine. +"""FlyDSL tensor-wise FP8 NN 4-wave GEMM kernel. + +This NN variant preserves the working 4-wave pipeline and the kernel contract +C = A @ B.T. A is physically [M, K] and B is physically [N, K]. During +staging, each B 128x128 half-page is transposed into XOR-swizzled physical LDS +[K128, N128]. The validated four-read ``ds_read_b64_tr_b8`` sequence then +reconstructs exactly the ordinary B[N, K] fragment consumed by the production +FP8 MFMA. 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. +hand-unrolled. M/N are runtime launch dimensions. The public entry point and private optimized core consume independently typed +FP8 E4M3 or E5M2 A/B tensors shaped [M, K] and [N, K], one FP32 inverse scale +per operand, and write float16, bfloat16, or float32 C shaped [M, N]. Operand +normalization is performed by the Transformer Engine wrapper. This module imports ``flydsl`` at import time and must therefore be imported lazily only after FlyDSL availability has been confirmed. @@ -32,8 +37,8 @@ # Transformer Engine-local FlyDSL utilities. from .fp8_gemm_utils import ( G2SLoader, + G2STransposeLoader, S2RLoader, - compute_global_linear_128x128, compute_global_swizzle, make_fp8_buffer_tensor, pack_i32x4_i32x8, @@ -293,11 +298,8 @@ def kernel_gemm( 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 = ( @@ -330,13 +332,26 @@ def kernel_gemm( 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) + # A keeps the ordinary row-major [M, K] direct-to-LDS path. + gl_off_a = 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 arrives row-major [N, K]. Stage each source 16-byte K vector into + # the transposed XOR-swizzled physical LDS image [K128, N128] required + # by the validated ds_read_b64_tr_b8 inverse mapping. + b_g2s = G2STransposeLoader(B, K, wave_id) s2r = S2RLoader(fx.Int32(0), 1) layout_lane16 = fx.make_layout((4, 16), (16, 1)) @@ -427,10 +442,9 @@ def hot_loop_scheduler_q_refill_2n(): 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_dsrd(1) rocdl.sched_mfma(2) rocdl.sched_barrier(0) @@ -441,15 +455,21 @@ def hot_loop_scheduler_q_prefetch_4n(): 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 + # 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) + # Load row-major global B[N, K], but write the half-page as + # XOR-swizzled physical LDS [K128, N128]. + global_n_base = by_n_idx + fx.Index(subtile * (BLOCK_N // 2)) + b_g2s.load_one( + lds_b[subtile], + global_n_base, + k_base, + pass_in_subtile, + ) def stage_a_subtile(k_base, subtile, lds_a): for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): @@ -475,10 +495,43 @@ def load_frag_at_byte_base(lds_page, row_byte_base): 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 load_b_frag_transpose(lds_page, local_n_tile): + # Exact inverse mapping validated against the ordinary B[N, K] + # production MFMA fragment: + # + # source_k = lane_div_16*16 + lane_in_16//2 + # source_n = local_n_tile + (lane_in_16&1)*8 + # + # base^0x440 advances logical K by 8 under the 128-byte XOR + # swizzle. The DS immediate 0x2000 advances logical K by 64. + lane_div16_i32 = fx.Int32(lane_div_16) + lane_in16_i32 = fx.Int32(lane_mod_16) + source_k = ( + lane_div16_i32 * fx.Int32(16) + + lane_in16_i32 // fx.Int32(2) + ) + source_n = ( + fx.Int32(local_n_tile) + + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) + ) + + physical_k, physical_n = swizzle_128(source_k, source_n) + base = physical_k * fx.Int32(BLOCK_N // 2) + physical_n + other = base ^ fx.Int32(0x440) + + x0 = s2r.load_one_transpose( + lds_page, + base, + other, + immediate_offset=0, + ) + x1 = s2r.load_one_transpose( + lds_page, + base, + other, + immediate_offset=0x2000, + ) + 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 @@ -584,8 +637,12 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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) + local_n_tile = ( + subtile_n_idx * fx.Index(SUBTILE_N) + + fx.Index(ni * MFMA_N) + - fx.Index(sn * (BLOCK_N // 2)) + ) + return load_b_frag_transpose(lds_b[sn], local_n_tile) def load_b_subtile_regs(lds_b, sn): return ( @@ -596,25 +653,11 @@ def load_b_subtile_regs(lds_b, sn): ) 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)) + 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) @@ -1015,7 +1058,7 @@ def launch_gemm( return launch_gemm @functools.lru_cache(maxsize=None) -def _cached_launch_nn( +def _cached_launch( K: int, a_fp8_dtype: torch.dtype, b_fp8_dtype: torch.dtype, @@ -1040,49 +1083,43 @@ def fp8_matmul( c: torch.Tensor, stream=None, ): - """TE-facing NN tensor-wise FP8 adapter. + """Launch correctness-first NN tensor-wise FP8 GEMM. - Public/backend contract: - a: [K, M] FP8 E4M3 or E5M2 activation payload + Contract: + a: [M, K] FP8 payload a_scale_inv: one-element FP32 inverse quantization scale - b: [N, K] FP8 E4M3 or E5M2 weight payload + b: [N, K] FP8 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. + B remains [N, K] through GMEM->LDS. The kernel performs a naive scalar + LDS gather along K for a fixed N row, constructing the same MFMA B + fragments as the optimized transpose-read path. This variant intentionally + does not use ds_read_b64_tr_b8. """ if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): - raise TypeError("FlyDSL FP8 GEMM expects plain torch.Tensor payloads") - + raise TypeError("FlyDSL FP8 NN 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, - ) + 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, " + "FlyDSL FP8 NN GEMM expects E4M3 or E5M2 payloads, " f"got A={a.dtype} and B={b.dtype}" ) - k, m = a.shape + m, k = 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), - ): + 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: @@ -1103,29 +1140,10 @@ def fp8_matmul( 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" - ) + 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) - doGemm( - a, - b, - c, - a_scale_inv, - b_scale_inv, - stream=stream, - ) def doGemm( A: torch.Tensor, @@ -1136,43 +1154,33 @@ def doGemm( stream=None, use_xcd_remap: bool = True, ): - """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" - K_runtime, M_runtime = A.shape + """Launch NN FP8 GEMM with C = A @ B.T, A [M,K], B [N,K].""" + M_runtime, K_runtime = A.shape N_runtime, Kb_runtime = B.shape - supported_fp8_dtypes = ( - torch.float8_e4m3fn, - torch.float8_e5m2, - ) + 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, - ), ( + 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}" + f"FlyDSL FP8 NN GEMM requires M to be a multiple of {_BLOCK_M}, 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}" + f"FlyDSL FP8 NN GEMM requires N to be a multiple of {_BLOCK_N}, 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}" + f"FlyDSL FP8 NN GEMM requires K to be a multiple of {_BLOCK_K}, 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"FlyDSL FP8 NN 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 @@ -1189,12 +1197,8 @@ def doGemm( 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 = _cached_launch( + int(K_runtime), A.dtype, B.dtype, C.dtype, bool(use_xcd_remap) ) launch( A_arg, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py index 975fd898d..e3f8b2cf5 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py @@ -2,15 +2,20 @@ # # See LICENSE for license information. -"""FlyDSL tensor-wise FP8 4-wave NT GEMM kernel for Transformer Engine. +"""FlyDSL tensor-wise FP8 NT 4-wave GEMM kernel. + +This NT variant preserves the working 4-wave pipeline while applying the +validated ``ds_read_b64_tr_b8`` contract to both operands. A is physically +[K, M] and B is physically [K, N]. Each 128x128 source tile is staged into an +XOR-swizzled physical LDS image [K128, X128], and four transpose reads rebuild +the exact ordinary MFMA fragment for one fixed M or N coordinate. 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 +hand-unrolled. M/N are runtime launch dimensions. The public entry point 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. +[K, N], one FP32 inverse scale per operand, and writes float16, bfloat16, or +float32 C shaped [M, N]. Operand normalization is performed by the +Transformer Engine wrapper. This module imports ``flydsl`` at import time and must therefore be imported lazily only after FlyDSL availability has been confirmed. @@ -33,7 +38,6 @@ from .fp8_gemm_utils import ( G2SLoader, S2RLoader, - compute_global_linear_128x128, compute_global_swizzle, make_fp8_buffer_tensor, pack_i32x4_i32x8, @@ -323,20 +327,49 @@ def kernel_gemm( bx_m_idx = fx.Index(bx_m) by_n_idx = fx.Index(by_n) - # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # Keep wave/lane arithmetic in i32. The global-offset helpers combine # 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) + # NT storage is K-major for both operands: + # A [K, M] + # B [K, N] + # + # Read each global 128x128 K-by-X tile in XOR-swizzled coordinate order + # and write it linearly to LDS. Because swizzle_128 is self-inverse, + # this produces the physical XOR-swizzled LDS image [K128, X128] + # consumed by ds_read_b64_tr_b8. + gl_off_a = compute_global_swizzle( + lane, + wave_id, + c_m, + LOAD_PASSES_HALF, + preshuffled=False, + ) + gl_off_b = compute_global_swizzle( + lane, + wave_id, + c_n, + 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)) @@ -427,7 +460,6 @@ def hot_loop_scheduler_q_refill_2n(): 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) @@ -436,22 +468,30 @@ def hot_loop_scheduler_q0_refill_a1_2n(): def hot_loop_scheduler_q_prefetch_4n(): for _ in range_constexpr(8): - rocdl.sched_dsrd(2) + rocdl.sched_dsrd(4) 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 is physically [K, M]. Copy + # A[k_base:k_base+128, bx_m+subtile*128:...] + # into one XOR-swizzled physical LDS half-page [K128, M128]. + global_base = ( + k_base * fx.Index(c_m) + + bx_m_idx + + fx.Index(subtile * (BLOCK_M // 2)) + ) 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 is physically [K, N]. Copy + # B[k_base:k_base+128, by_n+subtile*128:...] + # into one XOR-swizzled physical LDS half-page [K128, N128]. + global_base = ( + k_base * fx.Index(c_n) + + by_n_idx + + fx.Index(subtile * (BLOCK_N // 2)) + ) b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) def stage_a_subtile(k_base, subtile, lds_a): @@ -462,39 +502,47 @@ 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_transposed_frag_half(lds_page, local_x_tile, half): + """Load one K64 portion of a fixed-X MFMA fragment. + + This is the inverse mapping validated against the working ordinary + LDS fragment: + + source_k = lane_div_16*16 + lane_in_16//2 + source_x = local_x_tile + (lane_in_16&1)*8 + + ``base ^ 0x440`` advances logical K by 8 under swizzle_128. + The 0x2000 DS immediate advances logical K by 64. + """ + lane_div16_i32 = fx.Int32(lane_div_16) + lane_in16_i32 = fx.Int32(lane_mod_16) + source_k = ( + lane_div16_i32 * fx.Int32(16) + + lane_in16_i32 // fx.Int32(2) + ) + source_x = ( + fx.Int32(local_x_tile) + + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) + ) + + physical_k, physical_x = swizzle_128(source_k, source_x) + base = physical_k * fx.Int32(128) + physical_x + other = base ^ fx.Int32(0x440) + immediate_offset = 0 if half == 0 else 0x2000 - 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), + lds_page, + base, + other, + immediate_offset=immediate_offset, ) - 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) + def load_transposed_frag(lds_page, local_x_tile): + x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) + x1 = load_transposed_frag_half(lds_page, local_x_tile, 1) return pack_frag_halves(x0, x1) def _acc_idx(subtile_id, mi, ni): @@ -585,14 +633,6 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): # 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 @@ -601,8 +641,12 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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) + local_n_tile = ( + subtile_n_idx * fx.Index(SUBTILE_N) + + fx.Index(ni * MFMA_N) + - fx.Index(sn * (BLOCK_N // 2)) + ) + return load_transposed_frag(lds_b[sn], local_n_tile) def load_b_subtile_regs(lds_b, sn): return ( @@ -613,25 +657,17 @@ def load_b_subtile_regs(lds_b, sn): ) 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. + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) local_m_tile = ( - (reg_subtile_m_idx0 + fx.Index(sm * 2)) * fx.Index(SUBTILE_M) + subtile_m_idx * 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)) + return load_transposed_frag_half( + lds_a[sm], + local_m_tile, + half, + ) def load_a_subtile_mi_regs(lds_a, sm, mi): x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) @@ -1032,7 +1068,7 @@ def launch_gemm( return launch_gemm @functools.lru_cache(maxsize=None) -def _cached_launch_nt( +def _cached_launch( K: int, a_fp8_dtype: torch.dtype, b_fp8_dtype: torch.dtype, @@ -1057,35 +1093,31 @@ def fp8_matmul( c: torch.Tensor, stream=None, ): - """TE-facing NT tensor-wise FP8 adapter. + """Launch NT tensor-wise FP8 GEMM with transpose-read A/B fragments. - Public/backend contract: - a: [K, M] FP8 E4M3 or E5M2 activation payload + Contract: + a: [K, M] FP8 payload a_scale_inv: one-element FP32 inverse quantization scale - b: [K, N] FP8 E4M3 or E5M2 weight payload + b: [K, N] FP8 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. + Both operands remain K-major in global memory. Each tile is staged as a + swizzled physical [K128, X128] LDS image and read with the validated + four-instruction ds_read_b64_tr_b8 fragment contract. """ if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): - raise TypeError("FlyDSL FP8 GEMM expects plain torch.Tensor payloads") - + raise TypeError("FlyDSL FP8 NT 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, - ) + 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, " + "FlyDSL FP8 NT GEMM expects E4M3 or E5M2 payloads, " f"got A={a.dtype} and B={b.dtype}" ) @@ -1096,10 +1128,7 @@ def fp8_matmul( 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), - ): + 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: @@ -1120,29 +1149,10 @@ def fp8_matmul( 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" - ) + 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) - doGemm( - a, - b, - c, - a_scale_inv, - b_scale_inv, - stream=stream, - ) def doGemm( A: torch.Tensor, @@ -1153,43 +1163,33 @@ def doGemm( stream=None, use_xcd_remap: bool = True, ): - """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" + """Launch optimized NT FP8 GEMM from K-major A [K,M] and B [K,N].""" K_runtime, M_runtime = A.shape Kb_runtime, N_runtime = B.shape - supported_fp8_dtypes = ( - torch.float8_e4m3fn, - torch.float8_e5m2, - ) + 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, - ), ( + 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}" + f"FlyDSL FP8 NT GEMM requires M to be a multiple of {_BLOCK_M}, 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}" + f"FlyDSL FP8 NT GEMM requires N to be a multiple of {_BLOCK_N}, 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}" + f"FlyDSL FP8 NT GEMM requires K to be a multiple of {_BLOCK_K}, 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"FlyDSL FP8 NT 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 @@ -1206,12 +1206,8 @@ def doGemm( 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 = _cached_launch( + int(K_runtime), A.dtype, B.dtype, C.dtype, bool(use_xcd_remap) ) launch( A_arg, 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 b8ed21535..2c0e0534b 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py @@ -5,7 +5,8 @@ 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 import arith, buffer_ops, const_expr, range_constexpr, rocdl +from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec from flydsl.expr.utils.arith import _to_raw as as_mlir_value @@ -119,6 +120,70 @@ def load_one(self, lds_dst, k_offset, step): fx.copy(self.g2lds_atom, src, dst, soffset=fx.Int32(k_offset)) +class G2STransposeLoader: + """Stage a row-major 128x128 byte tile as swizzled physical [K, N]. + + The source is a row-major byte matrix ``[N, K]``. Each thread loads one + contiguous 16-byte K vector from global memory, then scatters those bytes + into the 128-byte XOR-swizzled LDS image consumed by + ``ds_read_b64_tr_b8``. + + One ``load_one`` call covers one of the four 4-KiB staging passes for a + 128x128 half-page. + """ + + def __init__(self, gl_src, leading_dim, wave_id): + self.gl_rsrc = buffer_ops.create_buffer_resource(gl_src, max_size=True) + self.leading_dim = fx.Int32(leading_dim) + self.wave_id = fx.Int32(wave_id) + self.lane_id = fx.thread_idx.x % 64 + self.n_waves = fx.block_dim.x // 64 + self.i8_lds_ptr_t = fx.PointerType.get( + elem_ty=ir.IntegerType.get_signless(8), + address_space=2, + alignment=1, + ) + + def _store_u8(self, lds_dst, byte_offset, value): + base_i32 = fx.Int32(fx.ptrtoint(lds_dst.ptr)) + addr_i32 = base_i32 + fx.Int32(byte_offset) + i8_ptr = fx.inttoptr(self.i8_lds_ptr_t, addr_i32) + view = fx.make_view(i8_ptr, fx.make_layout(1, 1)) + fx.memref_store_vec(Vec.filled(1, value, fx.Uint8), view) + + def load_one(self, lds_dst, global_n_base, k_base, step): + """Load one 16-byte/thread pass and transpose it into LDS. + + ``global_n_base`` is the first source N row of this 128-row half-page. + ``k_base`` is the first global K byte of the current K128 tile. + """ + row = ( + self.lane_id // fx.Int32(8) + + self.wave_id * fx.Int32(8) + + fx.Int32(step) * fx.Int32(self.n_waves * 8) + ) + col = (self.lane_id % fx.Int32(8)) * fx.Int32(16) + + global_byte = ( + (fx.Int32(global_n_base) + row) * self.leading_dim + + fx.Int32(k_base) + + col + ) + packed_i32x4 = buffer_ops.buffer_load( + self.gl_rsrc, + global_byte // fx.Int32(4), + vec_width=4, + dtype=T.i32, + ) + packed_u8x16 = Vec(packed_i32x4).bitcast(fx.Uint8) + + for byte_i in range_constexpr(16): + logical_k = col + fx.Int32(byte_i) + physical_k, physical_n = swizzle_128(logical_k, row) + lds_byte = physical_k * fx.Int32(128) + physical_n + self._store_u8(lds_dst, lds_byte, packed_u8x16[byte_i]) + + def pack_i32x4_i32x8(lo, hi): # Pack two i32x4 as one i32x8 return lo.shuffle(hi, list(range(8))) @@ -137,6 +202,23 @@ def _vec_load_16xf8(self, lds_src, offset): view = fx.make_view(i8_iter, fx.make_layout(16, 1)) return view.load() + def _vec_load_1xf8(self, lds_src, offset): + """Naive one-byte LDS load with direct dynamic byte addressing. + + Avoid ``make_int_tuple`` entirely because this FlyDSL build cannot + reliably infer tuple types from dynamic Index expressions. + """ + base_i32 = fx.Int32(fx.ptrtoint(lds_src.ptr)) + addr_i32 = base_i32 + fx.Int32(offset) + i8_lds_ptr_t = fx.PointerType.get( + elem_ty=ir.IntegerType.get_signless(8), + address_space=2, + alignment=1, + ) + i8_ptr = fx.inttoptr(i8_lds_ptr_t, addr_i32) + view = fx.make_view(i8_ptr, fx.make_layout(1, 1)) + return view.load() + def load(self, lds_src, preshuffled=False): frag = [] for i in range_constexpr(self.n_tiles): @@ -158,34 +240,58 @@ 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): + def _ds_read_b64_tr_b8(self, lds_src, byte_offset, immediate_offset=0): """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. + ``immediate_offset`` is encoded in the DS instruction itself. The NN + K128 path uses 0 and 0x2000, where 0x2000 advances the logical K row by + 64 in a 128-byte-wide physical LDS image. """ + if immediate_offset == 0: + asm = "ds_read_b64_tr_b8 $0, $1 offset:0\n" + elif immediate_offset == 0x2000: + asm = "ds_read_b64_tr_b8 $0, $1 offset:8192\n" + else: + raise ValueError( + "ds_read_b64_tr_b8 supports immediate offsets 0 and 0x2000, " + f"got {immediate_offset:#x}" + ) + 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", + asm, "=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. + def load_one_transpose( + self, + lds_src, + first_byte_offset, + second_byte_offset, + immediate_offset=0, + ): + """Load one 16-byte portion of a K128 FP8 MFMA operand. + + Two transpose reads return four packed i32 values. Calling this once + with immediate 0 and once with immediate 0x2000 yields the two i32x4 + portions that concatenate into the production i32x8 MFMA fragment. """ - lo = self._ds_read_b64_tr_b8(lds_src, first_byte_offset) - hi = self._ds_read_b64_tr_b8(lds_src, second_byte_offset) + lo = self._ds_read_b64_tr_b8( + lds_src, + first_byte_offset, + immediate_offset, + ) + hi = self._ds_read_b64_tr_b8( + lds_src, + second_byte_offset, + immediate_offset, + ) return lo.shuffle(hi, [0, 1, 2, 3]) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 5beff5a75..fb670b518 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -11,6 +11,8 @@ from transformer_engine.pytorch.utils import get_device_compute_capability +from .exceptions import FlyDSLUnsupportedError + from .bf16_gemm import bf16_matmul from .fp16_gemm import fp16_matmul from .fp32_gemm import fp32_matmul @@ -31,7 +33,7 @@ def _product(shape): 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 + This matches TE's generic GEMM output-shape convention: the physical GEMM is flattened to ``[M, N]``, while the returned tensor keeps B's leading dimensions when ``transb`` is false. """ @@ -154,7 +156,9 @@ def _classify_input(t): 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") + raise FlyDSLUnsupportedError( + f"{name} does not contain the required FP8 payload" + ) if fp8_dtype not in ( tex.DType.kFloat8E4M3, @@ -199,6 +203,36 @@ def _mxfp8_debug(message: str) -> None: print(f"[DEBUG_FLYDSL_MXFP8_GEMM] {message}") +def _fp8_debug_enabled() -> bool: + value = os.getenv("DEBUG_FLYDSL_FP8_GEMM", "") + return value.lower() not in ("", "0", "false", "no", "off") + + +def _fp8_debug(message: str) -> None: + if _fp8_debug_enabled(): + print(f"[DEBUG_FLYDSL_FP8_GEMM] {message}") + + +def _fp8_tensor_debug(name: str, tensor: torch.Tensor) -> None: + if not _fp8_debug_enabled(): + return + _fp8_debug( + f"{name}: shape={tuple(tensor.shape)}, stride={tuple(tensor.stride())}, " + f"dtype={tensor.dtype}, device={tensor.device}, " + f"contiguous={tensor.is_contiguous()}, data_ptr=0x{tensor.data_ptr():x}" + ) + + +def _fp8_scale_debug(name: str, scale: torch.Tensor) -> None: + if not _fp8_debug_enabled(): + return + value = scale.detach().float().reshape(-1).cpu().tolist() + _fp8_debug( + f"{name}: shape={tuple(scale.shape)}, dtype={scale.dtype}, " + f"device={scale.device}, data_ptr=0x{scale.data_ptr():x}, value={value}" + ) + + def _canonicalize_blas_pair( A_data: torch.Tensor, transa: bool, @@ -220,6 +254,15 @@ def _flatten_rowwise(t: torch.Tensor, name: str) -> torch.Tensor: return t.reshape(-1, t.shape[-1]) +def _flatten_columnwise(t: torch.Tensor, name: str) -> torch.Tensor: + """Flatten TE columnwise storage while preserving its leading dimension.""" + if t.ndim < 2: + raise ValueError( + f"FlyDSL GEMM expects {name} to have rank >= 2, got {tuple(t.shape)}" + ) + return t.reshape(t.shape[0], -1) + + def _canonicalize_blas_operands( A_data: torch.Tensor, transa: bool, @@ -354,61 +397,70 @@ 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. - - 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) +def _get_fp8_rowwise_payload(t, name): + """Return TE's existing rowwise ``_data`` payload without copying.""" data = getattr(t, "_data", None) - - if data is not None: - return _reinterpret_fp8_payload( - data, - fp8_dtype, - f"{name}._data", + if data is None: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 requires existing {name} rowwise (_data) storage" ) + return _reinterpret_fp8_payload( + data, + getattr(t, "_fp8_dtype", None), + f"{name}._data", + ) + +def _get_fp8_columnwise_payload(t, name): + """Return TE's existing columnwise ``_transpose`` payload without copying.""" if not _valid_fp8_transpose(t): - raise RuntimeError( - f"{name} has neither valid rowwise (_data) nor " - f"columnwise (_transpose) FP8 storage" + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 requires valid {name} columnwise (_transpose) storage" ) - - transpose_data = _reinterpret_fp8_payload( + return _reinterpret_fp8_payload( t._transpose, - fp8_dtype, + getattr(t, "_fp8_dtype", None), f"{name}._transpose", ) - return _materialize_rowwise_from_columnwise( - transpose_data, - f"{name}._transpose", - ) + +def _validate_fp8_kernel_operands( + kernel_a, + kernel_b, + *, + layout, + a_storage, + b_storage, +): + """Validate zero-copy physical operands before launching an FP8 kernel.""" + if kernel_a.ndim != 2 or kernel_b.ndim != 2: + raise ValueError( + f"FlyDSL FP8 {layout} expects rank-2 kernel operands, got " + f"{a_storage}={tuple(kernel_a.shape)} and " + f"{b_storage}={tuple(kernel_b.shape)}" + ) + if not kernel_a.is_contiguous() or not kernel_b.is_contiguous(): + raise ValueError( + f"FlyDSL FP8 {layout} requires contiguous {a_storage} and " + f"{b_storage}; refusing to materialize replacement operands" + ) + if kernel_a.device != kernel_b.device: + raise ValueError( + f"FlyDSL FP8 {layout} operands must be on the same device, got " + f"{kernel_a.device} and {kernel_b.device}" + ) + + +def _fp8_output_shape(D, m, n): + """Preserve TE's logical output shape when D is preallocated.""" + 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 logical output shape {tuple(output_shape)} does not " + f"match flattened kernel shape {(m, n)}" + ) + return output_shape def _select_mxfp8_data_and_scale( @@ -494,7 +546,7 @@ def _run_mxfp8( f"B_type={type(B).__name__}, D_provided={D is not None}" ) - # Match TE CanonicalizeGemmInput / Triton data_and_scale_for_transpose: + # Select the MXFP8 representation required by the transpose flags: # A: transa=True -> rowwise, transa=False -> columnwise # B: transb=True -> columnwise, transb=False -> rowwise A_data, A_scale = _select_mxfp8_data_and_scale( @@ -604,6 +656,68 @@ def _run_mxfp8( return D +def _select_fp8_storage_for_layout(A, transa, B, transb): + """Select the exact existing TE FP8 backing required by each layout. + + Fixed zero-copy routes selected for the final kernel contracts: + + TN: wrapper swaps B._data/A._data -> [M,K], [N,K] + NN: wrapper swaps B._data/A._transpose -> [M,K], [N,K] + NT: wrapper swaps B._data/A._data -> [K,M], [K,N] + + In particular, NT must use the contiguous rowwise K-major payloads. + Passing ``_transpose.transpose(0, 1)`` would create strided views and + force the NT adapter to materialize them before launch. + """ + layout = (bool(transa), bool(transb)) + + if layout == (True, False): # TN + A_payload = _get_fp8_rowwise_payload(A, "A") + A_storage = "A._data" + A_data = _flatten_rowwise(A_payload, A_storage) + + B_payload = _get_fp8_rowwise_payload(B, "B") + B_storage = "B._data" + B_data = _flatten_rowwise(B_payload, B_storage) + + elif layout == (False, False): # NN + A_payload = _get_fp8_columnwise_payload(A, "A") + A_storage = "A._transpose" + A_data = _flatten_columnwise(A_payload, A_storage) + + B_payload = _get_fp8_rowwise_payload(B, "B") + B_storage = "B._data" + B_data = _flatten_rowwise(B_payload, B_storage) + + elif layout == (False, True): # NT / dW + # fp8_gemm_nt consumes contiguous K-major operands directly: + # kernel a = B._data [K, M] + # kernel b = A._data [K, N] + # Select rowwise storage here so the ownership swap in _run_fp8 is + # zero-copy and no noncontiguous transpose view reaches the kernel. + A_payload = _get_fp8_rowwise_payload(A, "A") + A_storage = "A._data" + A_data = _flatten_rowwise(A_payload, A_storage) + + B_payload = _get_fp8_rowwise_payload(B, "B") + B_storage = "B._data" + B_data = _flatten_rowwise(B_payload, B_storage) + + else: + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) + + return ( + A_data, + A_storage, + torch.Size(A_payload.shape), + B_data, + B_storage, + torch.Size(B_payload.shape), + ) + + def _run_fp8( A, transa, @@ -613,37 +727,22 @@ def _run_fp8( *, output_dtype: torch.dtype, ): - """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) + """Dispatch tensor-wise FP8 using exact per-kernel storage contracts.""" supported_fp8_dtypes = ( tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2, ) + a_fp8_dtype = getattr(A, "_fp8_dtype", None) + b_fp8_dtype = getattr(B, "_fp8_dtype", None) if ( a_fp8_dtype not in supported_fp8_dtypes or b_fp8_dtype not in supported_fp8_dtypes ): - raise NotImplementedError( + raise FlyDSLUnsupportedError( "FlyDSL FP8 supports E4M3 and E5M2 independently for A/B; " f"got A={a_fp8_dtype} and B={b_fp8_dtype}" ) - if transa and transb: - raise NotImplementedError( - "FlyDSL GEMM does not support transa=True, transb=True (TT)" - ) - A_scale_inv = getattr(A, "_scale_inv", None) B_scale_inv = getattr(B, "_scale_inv", None) for name, scale in ( @@ -651,253 +750,135 @@ def _run_fp8( ("B._scale_inv", B_scale_inv), ): if not isinstance(scale, torch.Tensor): - raise RuntimeError(f"{name} is not populated") + raise FlyDSLUnsupportedError(f"{name} is not populated") if scale.dtype != torch.float32 or scale.numel() != 1: - raise ValueError( + raise FlyDSLUnsupportedError( f"{name} must contain exactly one FP32 tensor-wise inverse " 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", - ) + layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" - 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" - ) + ( + A_data, + A_storage, + A_payload_shape, + B_data, + B_storage, + B_payload_shape, + ) = _select_fp8_storage_for_layout( + A, + bool(transa), + B, + bool(transb), + ) - 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)}" - ) + _validate_fp8_kernel_operands( + A_data, + B_data, + layout=layout, + a_storage=A_storage, + b_storage=B_storage, + ) - # 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)}" - ) + a_scale = B_scale_inv + b_scale = A_scale_inv - 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}" - ) + if layout == "TN": + matmul = fp8_matmul + kernel_layout = "TN" - D = _validate_or_allocate_output( - D, - shape=output_shape, - dtype=output_dtype, - device=a_flydsl.device, - backend_name="FP8 NN", - ) + a_flydsl = B_data + b_flydsl = A_data - # 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" - ) + m, k = a_flydsl.shape + n, kb = b_flydsl.shape - # 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", - ) + elif layout == "NN": + matmul = fp8_matmul_nn + kernel_layout = "NN" - 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" - ) + a_flydsl = B_data + b_flydsl = A_data - 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)}" - ) + m, k = a_flydsl.shape + n, kb = b_flydsl.shape - a_flydsl = b_columnwise.view(k, m) - b_flydsl = a_columnwise.view(k, n) + elif layout == "NT": + matmul = fp8_matmul_nt + kernel_layout = "NT" - # 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)}" - ) + # Exact fp8_gemm_nt contract, with no view or materialization: + # a_flydsl = B._data [K, M] + # b_flydsl = A._data [K, N] + a_flydsl = B_data + b_flydsl = A_data - 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}" - ) + k, m = a_flydsl.shape + kb, n = b_flydsl.shape - D = _validate_or_allocate_output( - D, - shape=output_shape, - dtype=output_dtype, - device=a_flydsl.device, - backend_name="FP8 NT", + else: + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" ) - # Scales follow the BLAS-swapped kernel operands. - fp8_matmul_nt( - a_flydsl, - B_scale_inv, - b_flydsl, - A_scale_inv, - D.view(m, n), + if not a_flydsl.is_contiguous() or not b_flydsl.is_contiguous(): + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 {layout} kernel contract requires contiguous final " + f"operands, got a={tuple(a_flydsl.shape)} " + f"stride={tuple(a_flydsl.stride())} and " + f"b={tuple(b_flydsl.shape)} stride={tuple(b_flydsl.stride())}" ) - 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 - ) - 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 kb != k: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 {layout} selected incompatible physical backings: " + f"{B_storage}={tuple(B_data.shape)} and " + f"{A_storage}={tuple(A_data.shape)}; " + f"kernel operands are {tuple(a_flydsl.shape)} and " + f"{tuple(b_flydsl.shape)}" ) - 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}" + if D is not None: + logical_output_shape = torch.Size(D.shape) + elif layout in ("TN", "NN"): + logical_output_shape = torch.Size((*B_payload_shape[:-1], n)) + else: + logical_output_shape = torch.Size((m, n)) + if _product(logical_output_shape) != m * n: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 {layout} logical output shape " + f"{tuple(logical_output_shape)} does not match kernel output " + f"shape {(m, n)}" ) D = _validate_or_allocate_output( D, - shape=output_shape, + shape=logical_output_shape, dtype=output_dtype, device=a_flydsl.device, - backend_name="FP8", + backend_name=f"FP8 {kernel_layout}", ) - # Operand swap means B's tensor-wise scale belongs to a_flydsl and A's - # tensor-wise scale belongs to b_flydsl. - fp8_matmul( + _fp8_debug( + f"dispatch entry: transa={bool(transa)}, transb={bool(transb)}, " + f"layout={layout}, selected_kernel={matmul.__module__}.{matmul.__name__}" + ) + _fp8_debug(f"selected TE storage: A={A_storage}, B={B_storage}") + _fp8_tensor_debug(f"selected/{A_storage}", A_data) + _fp8_tensor_debug(f"selected/{B_storage}", B_data) + _fp8_tensor_debug("a_flydsl", a_flydsl) + _fp8_tensor_debug("b_flydsl", b_flydsl) + _fp8_scale_debug("a_scale", a_scale) + _fp8_scale_debug("b_scale", b_scale) + _fp8_debug(f"derived M={m}, N={n}, K={k}") + _fp8_tensor_debug("output/D", D) + + matmul( a_flydsl, - B_scale_inv, + a_scale, b_flydsl, - A_scale_inv, + b_scale, D.view(m, n), ) return D From 1ba15ed515baa2de4fa0773317579e8deea136df Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 28 Jul 2026 15:10:47 +0000 Subject: [PATCH 19/31] Add direct MXFP8 NN/NT FlyDSL GEMM specializations --- .../flydsl_kernels/gemm/gemm_wrappers.py | 476 ++++-- .../flydsl_kernels/gemm/mxfp8_gemm_nn.py | 1503 +++++++++++++++++ .../flydsl_kernels/gemm/mxfp8_gemm_nt.py | 1474 ++++++++++++++++ 3 files changed, 3297 insertions(+), 156 deletions(-) create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nn.py create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nt.py diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index fb670b518..e9e861631 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -20,6 +20,8 @@ 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 +from .mxfp8_gemm_nn import mxfp8_matmul as mxfp8_matmul_nn +from .mxfp8_gemm_nt import mxfp8_matmul as mxfp8_matmul_nt def _product(shape): @@ -255,7 +257,7 @@ def _flatten_rowwise(t: torch.Tensor, name: str) -> torch.Tensor: def _flatten_columnwise(t: torch.Tensor, name: str) -> torch.Tensor: - """Flatten TE columnwise storage while preserving its leading dimension.""" + """Flatten TE columnwise storage as [last_dim, product(leading_dims)].""" if t.ndim < 2: raise ValueError( f"FlyDSL GEMM expects {name} to have rank >= 2, got {tuple(t.shape)}" @@ -310,6 +312,42 @@ def _canonicalize_blas_operands( return a_flydsl, b_flydsl, m, n, k +def _resolve_output_shape( + A, + transa, + B, + transb, + D, + *, + m, + n, + backend_name, +): + """Resolve TE's public output shape independently of kernel storage. + + FlyDSL kernels always write a flattened row-major ``[M, N]`` matrix. + TE's public tensor may retain leading dimensions (for example + ``[sequence, batch, hidden]``). Quantized rowwise/columnwise payloads are + physical storage views and must never be used to infer that public shape. + + A caller-provided ``D`` is authoritative. Otherwise derive the logical + shape from the original TE operands, before selecting or flattening any + backing storage. + """ + if D is not None: + output_shape = torch.Size(D.shape) + else: + output_shape = _get_gemm_output_shape(A, transa, B, transb) + + if _product(output_shape) != m * n: + raise FlyDSLUnsupportedError( + f"FlyDSL {backend_name} logical output shape " + f"{tuple(output_shape)} does not match flattened kernel shape " + f"{(m, n)}" + ) + return output_shape + + def _validate_or_allocate_output( D, *, @@ -367,16 +405,19 @@ 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)}" - ) + output_shape = _resolve_output_shape( + A, + transa, + B, + transb, + D, + m=m, + n=n, + backend_name=backend_name, + ) if output_dtype is None: output_dtype = dtype @@ -499,22 +540,58 @@ def _select_mxfp8_data_and_scale( return data, scale -def _flatten_mxfp8_scale(t: torch.Tensor, name: str) -> torch.Tensor: +def _mxfp8_logical_shape(t, name: str) -> torch.Size: + """Return MXFP8 logical shape from a populated backing tensor. + + MXFP8TensorStorage is not a torch.Tensor and does not expose ``.shape``. + Rowwise and columnwise MXFP8 payloads retain the same logical row-major + shape, so either populated backing is sufficient for shape derivation. + """ + data = getattr(t, "_rowwise_data", None) + if data is None: + data = getattr(t, "_columnwise_data", None) + if data is None: + raise FlyDSLUnsupportedError( + f"{name} has neither rowwise nor columnwise MXFP8 data" + ) + return torch.Size(data.shape) + + +def _flatten_mxfp8_scale( + t: torch.Tensor, + name: str, + *, + source_colwise: bool, +) -> torch.Tensor: + """Flatten a raw TE MXFP8 scale tensor without changing orientation. + + Rowwise source: + [..., K/32] -> [outer, K/32] + + Columnwise source: + [K/32, ...] -> [K/32, outer] + """ 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: + if source_colwise: + t = t.reshape(t.shape[0], -1) + orientation = "columnwise" + else: t = t.reshape(-1, t.shape[-1]) + orientation = "rowwise" + _mxfp8_debug( - f"{name} scale flatten: {original_shape} -> {tuple(t.shape)}, " + f"{name} {orientation} scale flatten: " + f"{original_shape} -> {tuple(t.shape)}, " f"contiguous={t.is_contiguous()}" ) return t - def _run_mxfp8( A, transa, @@ -524,7 +601,24 @@ def _run_mxfp8( *, output_dtype: torch.dtype, ): - """Canonicalize independently typed E4M3/E5M2 MXFP8 operands.""" + """Dispatch MXFP8 through exact TN/NN/NT physical contracts. + + TE owns BLAS-shaped operands. After the usual ownership swap, FlyDSL + kernels consume: + + TN: a = B.rowwise [M, K] + b = A.rowwise.T [K, N] (validated TN adapter contract) + + NN: a = B.rowwise [M, K] + b = A.columnwise [K, N] + + NT: a = B.columnwise [K, M] + b = A.columnwise [K, N] + + MXFP8 rowwise and columnwise payloads retain the same logical row-major + shape. Columnwise selection changes the quantization axis; the specialized + NN/NT kernels provide the required transpose-read semantics. + """ a_fp8_dtype = getattr(A, "_fp8_dtype", None) b_fp8_dtype = getattr(B, "_fp8_dtype", None) supported_fp8_dtypes = ( @@ -535,118 +629,181 @@ def _run_mxfp8( a_fp8_dtype not in supported_fp8_dtypes or b_fp8_dtype not in supported_fp8_dtypes ): - raise NotImplementedError( + raise FlyDSLUnsupportedError( "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'}" + dispatch = { + (True, False): ("TN", mxfp8_matmul), + (False, False): ("NN", mxfp8_matmul_nn), + (False, True): ("NT", mxfp8_matmul_nt), + } + try: + kernel_layout, matmul = dispatch[(bool(transa), bool(transb))] + except KeyError as exc: + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) from exc + _mxfp8_debug( - f"entry: layout={layout}, A_type={type(A).__name__}, " - f"B_type={type(B).__name__}, D_provided={D is not None}" + f"entry: layout={layout}, selected_kernel=" + f"{matmul.__module__}.{matmul.__name__}, " + f"A_type={type(A).__name__}, B_type={type(B).__name__}, " + f"D_provided={D is not None}" ) - # Select the MXFP8 representation required by the transpose flags: - # A: transa=True -> rowwise, transa=False -> columnwise - # B: transb=True -> columnwise, transb=False -> rowwise + # Resolve public shapes from actual payload tensors. Never access + # MXFP8TensorStorage.shape: the storage wrapper has no such attribute. + A_logical_shape = _mxfp8_logical_shape(A, "A") + B_logical_shape = _mxfp8_logical_shape(B, "B") + + # Match TE/C++ MXFP8 representation selection exactly: + # A: transa=True -> rowwise; transa=False -> columnwise + # B: transb=False -> rowwise; transb=True -> columnwise + A_source_colwise = not bool(transa) + B_source_colwise = bool(transb) + A_data, A_scale = _select_mxfp8_data_and_scale( A, - will_transpose=not transa, + will_transpose=A_source_colwise, name="A", ) B_data, B_scale = _select_mxfp8_data_and_scale( B, - will_transpose=transb, + will_transpose=B_source_colwise, 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. + # Both MXFP8 payload orientations are stored row-major with the original + # logical shape. Flatten leading dimensions only; do not transpose or + # materialize selected columnwise payloads. + A_data = _flatten_rowwise(A_data, "A MXFP8 payload") + B_data = _flatten_rowwise(B_data, "B MXFP8 payload") + 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) - 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") - a_scale, b_scale = _canonicalize_blas_pair( + A_scale = _flatten_mxfp8_scale( A_scale, - transa, + "A", + source_colwise=A_source_colwise, + ) + B_scale = _flatten_mxfp8_scale( B_scale, - transb, + "B", + source_colwise=B_source_colwise, ) - _mxfp8_debug( - f"canonicalized layout={layout}: " - 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: " - 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}") + # Kernel operand ownership is always swapped relative to TE: + # kernel a <- TE B + # kernel b <- TE A + if kernel_layout == "TN": + # Preserve the validated TN adapter contract: + # a [M,K], b [K,N] + a_flydsl = B_data + b_flydsl = A_data.transpose(0, 1) + a_scale = B_scale + b_scale = A_scale.transpose(0, 1) + + m, k = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a_scale = (m, k // 32) + expected_b_scale = (k // 32, n) + + elif kernel_layout == "NN": + # A's columnwise MXFP8 payload is still row-major in its original + # shape, which is exactly the NN kernel's K-major [K,N] source. + a_flydsl = B_data + b_flydsl = A_data + a_scale = B_scale + b_scale = A_scale + + m, k = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a_scale = (m, k // 32) + expected_b_scale = (k // 32, n) + + else: + # Both selected columnwise payloads directly satisfy the NT kernel's + # K-major contracts without tensor transposes or copies. + a_flydsl = B_data + b_flydsl = A_data + a_scale = B_scale + b_scale = A_scale + + k, m = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a_scale = (k // 32, m) + expected_b_scale = (k // 32, n) + + if kb != k: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 {layout} selected incompatible payloads: " + f"a={tuple(a_flydsl.shape)} and b={tuple(b_flydsl.shape)}" + ) if a_flydsl.device != b_flydsl.device: raise ValueError( - f"A and B must be on the same device, got " + f"FlyDSL MXFP8 {layout} operands must be on the same device, got " f"{a_flydsl.device} and {b_flydsl.device}" ) - scale_group_size = 32 - if k % scale_group_size != 0: + if k % 32 != 0: raise ValueError( - f"K={k} must be divisible by MXFP8 scale group size " - f"{scale_group_size}" + f"K={k} must be divisible by MXFP8 scale group size 32" ) - # 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"A scale shape {tuple(a_scale.shape)} != expected " - f"{expected_a_scale}" + f"FlyDSL MXFP8 {layout} a_scale shape " + f"{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 " - f"{expected_b_scale}" + f"FlyDSL MXFP8 {layout} b_scale shape " + f"{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 as torch.uint8") + # Derive the public result shape from payload shapes, not storage wrappers. + if D is not None: + output_shape = torch.Size(D.shape) + else: + output_shape = _get_gemm_output_shape( + A_logical_shape, + transa, + B_logical_shape, + transb, + ) + if _product(output_shape) != m * n: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 {layout} logical output shape " + f"{tuple(output_shape)} does not match kernel output {(m, n)}" + ) + D = _validate_or_allocate_output( D, shape=output_shape, dtype=output_dtype, device=a_flydsl.device, - backend_name="MXFP8", + backend_name=f"MXFP8 {kernel_layout}", + ) + + _mxfp8_debug( + f"dispatch 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_scale={tuple(a_scale.shape)}; " + f"b_scale={tuple(b_scale.shape)}; " + f"M={m}, N={n}, K={k}" ) - mxfp8_matmul( + matmul( a_flydsl, a_scale, b_flydsl, @@ -655,19 +812,14 @@ def _run_mxfp8( ) return D - def _select_fp8_storage_for_layout(A, transa, B, transb): """Select the exact existing TE FP8 backing required by each layout. - Fixed zero-copy routes selected for the final kernel contracts: + Fixed zero-copy routes: - TN: wrapper swaps B._data/A._data -> [M,K], [N,K] - NN: wrapper swaps B._data/A._transpose -> [M,K], [N,K] - NT: wrapper swaps B._data/A._data -> [K,M], [K,N] - - In particular, NT must use the contiguous rowwise K-major payloads. - Passing ``_transpose.transpose(0, 1)`` would create strided views and - force the NT adapter to materialize them before launch. + TN: A._data, B._data + NN: A._transpose, B._data + NT: A._transpose, B._transpose """ layout = (bool(transa), bool(transb)) @@ -690,18 +842,13 @@ def _select_fp8_storage_for_layout(A, transa, B, transb): B_data = _flatten_rowwise(B_payload, B_storage) elif layout == (False, True): # NT / dW - # fp8_gemm_nt consumes contiguous K-major operands directly: - # kernel a = B._data [K, M] - # kernel b = A._data [K, N] - # Select rowwise storage here so the ownership swap in _run_fp8 is - # zero-copy and no noncontiguous transpose view reaches the kernel. - A_payload = _get_fp8_rowwise_payload(A, "A") - A_storage = "A._data" - A_data = _flatten_rowwise(A_payload, A_storage) + A_payload = _get_fp8_columnwise_payload(A, "A") + A_storage = "A._transpose" + A_data = _flatten_columnwise(A_payload, A_storage) - B_payload = _get_fp8_rowwise_payload(B, "B") - B_storage = "B._data" - B_data = _flatten_rowwise(B_payload, B_storage) + B_payload = _get_fp8_columnwise_payload(B, "B") + B_storage = "B._transpose" + B_data = _flatten_columnwise(B_payload, B_storage) else: raise FlyDSLUnsupportedError( @@ -727,7 +874,21 @@ def _run_fp8( *, output_dtype: torch.dtype, ): - """Dispatch tensor-wise FP8 using exact per-kernel storage contracts.""" + """Dispatch tensor-wise FP8 through one canonical operand contract. + + First select the exact existing TE storage required by the BLAS flags. + Then canonicalize both payloads and scales identically: + + a_flydsl, b_flydsl = op(B), op(A) + a_scale, b_scale = B scale, A scale + + Every kernel is called with: + + matmul(a_flydsl, a_scale, b_flydsl, b_scale, D) + + The layout-specific kernels differ only in the physical layouts they + expect for canonicalized ``a_flydsl`` and ``b_flydsl``. + """ supported_fp8_dtypes = ( tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2, @@ -752,20 +913,31 @@ def _run_fp8( if not isinstance(scale, torch.Tensor): raise FlyDSLUnsupportedError(f"{name} is not populated") if scale.dtype != torch.float32 or scale.numel() != 1: - raise FlyDSLUnsupportedError( + raise ValueError( f"{name} must contain exactly one FP32 tensor-wise inverse " f"scale, got dtype={scale.dtype}, shape={tuple(scale.shape)}" ) layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" + dispatch = { + (True, False): ("TN", fp8_matmul), + (False, False): ("NN", fp8_matmul_nn), + (False, True): ("NT", fp8_matmul_nt), + } + try: + kernel_layout, matmul = dispatch[(bool(transa), bool(transb))] + except KeyError as exc: + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) from exc ( A_data, A_storage, - A_payload_shape, + A_physical_shape, B_data, B_storage, - B_payload_shape, + B_physical_shape, ) = _select_fp8_storage_for_layout( A, bool(transa), @@ -781,97 +953,89 @@ def _run_fp8( b_storage=B_storage, ) + # Scales follow the original TE tensors after BLAS operand ownership swap. a_scale = B_scale_inv b_scale = A_scale_inv if layout == "TN": - matmul = fp8_matmul - kernel_layout = "TN" - + # a_flydsl = B._data [M,K] + # b_flydsl = A._data [N,K] a_flydsl = B_data b_flydsl = A_data - m, k = a_flydsl.shape n, kb = b_flydsl.shape elif layout == "NN": - matmul = fp8_matmul_nn - kernel_layout = "NN" - + # a_flydsl = B._data [M,K] + # b_flydsl = A._transpose flattened as [N,K] a_flydsl = B_data b_flydsl = A_data - m, k = a_flydsl.shape n, kb = b_flydsl.shape - elif layout == "NT": - matmul = fp8_matmul_nt - kernel_layout = "NT" - - # Exact fp8_gemm_nt contract, with no view or materialization: - # a_flydsl = B._data [K, M] - # b_flydsl = A._data [K, N] - a_flydsl = B_data - b_flydsl = A_data - - k, m = a_flydsl.shape - kb, n = b_flydsl.shape - else: - raise FlyDSLUnsupportedError( - "FlyDSL GEMM does not support transa=True, transb=True (TT)" - ) - - if not a_flydsl.is_contiguous() or not b_flydsl.is_contiguous(): - raise FlyDSLUnsupportedError( - f"FlyDSL FP8 {layout} kernel contract requires contiguous final " - f"operands, got a={tuple(a_flydsl.shape)} " - f"stride={tuple(a_flydsl.stride())} and " - f"b={tuple(b_flydsl.shape)} stride={tuple(b_flydsl.stride())}" - ) + # TE columnwise backings are contiguous allocations exposed as: + # B._transpose flattened [M,K] + # A._transpose flattened [N,K] + # + # fp8_gemm_nt consumes those same bytes with K-major tensor metadata: + # kernel_a [K,M] aliases B._transpose + # kernel_b [K,N] aliases A._transpose + m, k = B_data.shape + n, kb = A_data.shape + a_flydsl = B_data.view(k, m) + b_flydsl = A_data.view(kb, n) if kb != k: raise FlyDSLUnsupportedError( f"FlyDSL FP8 {layout} selected incompatible physical backings: " f"{B_storage}={tuple(B_data.shape)} and " - f"{A_storage}={tuple(A_data.shape)}; " - f"kernel operands are {tuple(a_flydsl.shape)} and " - f"{tuple(b_flydsl.shape)}" - ) - - if D is not None: - logical_output_shape = torch.Size(D.shape) - elif layout in ("TN", "NN"): - logical_output_shape = torch.Size((*B_payload_shape[:-1], n)) - else: - logical_output_shape = torch.Size((m, n)) - if _product(logical_output_shape) != m * n: - raise FlyDSLUnsupportedError( - f"FlyDSL FP8 {layout} logical output shape " - f"{tuple(logical_output_shape)} does not match kernel output " - f"shape {(m, n)}" + f"{A_storage}={tuple(A_data.shape)}" ) - D = _validate_or_allocate_output( - D, - shape=logical_output_shape, - dtype=output_dtype, - device=a_flydsl.device, - backend_name=f"FP8 {kernel_layout}", - ) - _fp8_debug( f"dispatch entry: transa={bool(transa)}, transb={bool(transb)}, " f"layout={layout}, selected_kernel={matmul.__module__}.{matmul.__name__}" ) - _fp8_debug(f"selected TE storage: A={A_storage}, B={B_storage}") + _fp8_debug( + f"selected TE storage: A={A_storage}, B={B_storage}" + ) _fp8_tensor_debug(f"selected/{A_storage}", A_data) _fp8_tensor_debug(f"selected/{B_storage}", B_data) + _fp8_debug( + "canonical contract: " + "matmul(a_flydsl, a_scale, b_flydsl, b_scale, D)" + ) _fp8_tensor_debug("a_flydsl", a_flydsl) _fp8_tensor_debug("b_flydsl", b_flydsl) _fp8_scale_debug("a_scale", a_scale) _fp8_scale_debug("b_scale", b_scale) - _fp8_debug(f"derived M={m}, N={n}, K={k}") + _fp8_debug( + f"canonical ownership: a_flydsl<-TE B, b_flydsl<-TE A; " + f"derived M={m}, N={n}, K={k}" + ) + + # Kernel storage is always flattened, but the public TE result must retain + # the logical leading dimensions of the original operands when D is not + # preallocated. Never infer the public shape from _data/_transpose. + logical_output_shape = _resolve_output_shape( + A, + transa, + B, + transb, + D, + m=m, + n=n, + backend_name=f"FP8 {layout}", + ) + + D = _validate_or_allocate_output( + D, + shape=logical_output_shape, + dtype=output_dtype, + device=a_flydsl.device, + backend_name=f"FP8 {kernel_layout}", + ) _fp8_tensor_debug("output/D", D) matmul( diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nn.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nn.py new file mode 100644 index 000000000..b7f20ba6b --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nn.py @@ -0,0 +1,1503 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL MXFP8 NN 4-wave GEMM implementation. + +This specialization preserves the validated MXFP8 TN compute, scale, MFMA, +accumulator, and epilogue pipelines. A is physically row-major [M, K]. +B is physically row-major [N, K], staged as XOR-swizzled [K128, N128] LDS, +and reconstructed with the validated four-read ds_read_b64_tr_b8 path. + +Raw scales enter as A rowwise [M, K/32] and B columnwise [K/32, N]. +Orientation-aware prepacking converts both to the common iteration-major +[K/128, dim] uint32 representation consumed by the kernel.""" + +import functools +import os + +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 .exceptions import FlyDSLUnsupportedError +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 _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, + *, + source_colwise: bool = False, +) -> torch.Tensor: + """Pack raw E8M0 scales as iteration-major ``[K/128, dim]`` uint32. + + ``source_colwise=False`` consumes TE rowwise scales ``[dim, K/32]``. + ``source_colwise=True`` consumes TE columnwise scales ``[K/32, dim]``. + + Both paths produce the same packed representation consumed by every + TN/NN/NT MXFP8 kernel specialization. + """ + 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)}" + ) + + if source_colwise: + qk, dim = scales_u8.shape + if qk % 4 != 0: + raise ValueError( + f"Columnwise scale K dimension must be divisible by 4 K32 groups, got {qk}" + ) + s32 = scales_u8.contiguous().view(qk // 4, 4, dim).to(torch.int32) + return ( + s32[:, 0, :] + | (s32[:, 1, :] << 8) + | (s32[:, 2, :] << 16) + | (s32[:, 3, :] << 24) + ).contiguous() + + dim, qk = scales_u8.shape + if qk % 4 != 0: + raise ValueError( + f"Rowwise scale K dimension must be divisible by 4 K32 groups, got {qk}" + ) + + s32 = scales_u8.contiguous().view(dim, 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, + *, + source_colwise: bool = False, +) -> torch.Tensor: + """Convert raw TE E8M0 scales to ``[K/128, dim]`` MFMA-ready words.""" + scale_iter = pack_mx32_scales_iter( + scales_u8, + source_colwise=source_colwise, + ) + dim = scales_u8.shape[1] if source_colwise else scales_u8.shape[0] + + if dim % 64 != 0: + raise ValueError( + f"Scale outer dimension={dim} must be a multiple of 64 for HK MFMA packing" + ) + + device = scales_u8.device + row = torch.arange(dim, device=device, dtype=torch.int64) + row_within_16 = row % 16 + k_subgroup = (row // 16) % 4 + tile = row // 64 + + packed = torch.zeros_like(scale_iter) + 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, dim) + ) & 0xFF + packed |= byte_value << (group * 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, + 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 + 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 + + 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[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, 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) + + 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) + 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) + + tx_i32 = fx.Int32(tx) + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # A remains ordinary row-major [M, K]. + gl_off_a = 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 is the selected MXFP8 columnwise payload, physically [K, N]. + # Load the K-major source directly into the XOR-swizzled physical LDS + # image [K128, N128] consumed by ds_read_b64_tr_b8. + gl_off_b = compute_global_swizzle( + lane, + wave_id, + c_n, + LOAD_PASSES_HALF, + preshuffled=False, + ) + 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 _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): + # B is physically [K, N]. Copy + # B[k_base:k_base+128, by_n+subtile*128:...] + # into one XOR-swizzled physical LDS half-page [K128, N128]. + global_base = ( + k_base * fx.Index(c_n) + + by_n_idx + + fx.Index(subtile * (BLOCK_N // 2)) + ) + 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_transpose(lds_page, local_n_tile): + # Exact inverse mapping validated against the ordinary B[N, K] + # production MFMA fragment: + # + # source_k = lane_div_16*16 + lane_in_16//2 + # source_n = local_n_tile + (lane_in_16&1)*8 + # + # base^0x440 advances logical K by 8 under the 128-byte XOR + # swizzle. The DS immediate 0x2000 advances logical K by 64. + lane_div16_i32 = fx.Int32(lane_div_16) + lane_in16_i32 = fx.Int32(lane_mod_16) + source_k = ( + lane_div16_i32 * fx.Int32(16) + + lane_in16_i32 // fx.Int32(2) + ) + source_n = ( + fx.Int32(local_n_tile) + + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) + ) + + physical_k, physical_n = swizzle_128(source_k, source_n) + base = physical_k * fx.Int32(BLOCK_N // 2) + physical_n + other = base ^ fx.Int32(0x440) + + x0 = s2r.load_one_transpose( + lds_page, + base, + other, + immediate_offset=0, + ) + x1 = s2r.load_one_transpose( + lds_page, + base, + other, + immediate_offset=0x2000, + ) + 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, 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"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, + ) + + 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"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, + ) + + 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 + 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. + # 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): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_scales = scale_tile[2] if sn == 0 else scale_tile[3] + local_n_tile = ( + subtile_n_idx * fx.Index(SUBTILE_N) + + fx.Index(ni * MFMA_N) + - fx.Index(sn * (BLOCK_N // 2)) + ) + b_ni = load_b_frag_transpose(lds_b[sn], local_n_tile) + return b_ni, b_scales[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, + 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, + ) + + + +def do_gemm( + 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 [K, N]. 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 + 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 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}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 NN 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 NN 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 NN 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 MXFP8 NN 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" + 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})" + ) + 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 + # 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), + A.dtype, + B.dtype, + C.dtype, + ) + launch( + A_arg, + As_arg, + B_arg, + Bs_arg, + C_arg, + M_runtime, + 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 MXFP8 NN GEMM with one transpose-read operand. + + Contract: + a: [M, K] row-major FP8 payload + a_scale: [M, K/32] raw rowwise E8M0 scales + b: [K, N] row-major columnwise-quantized FP8 payload + b_scale: [K/32, N] raw columnwise E8M0 scales + D: [M, N] float16, bfloat16, or float32 output + + The B payload remains physically [K, N]. The kernel stages that K-major + source into the XOR-swizzled LDS image and uses ds_read_b64_tr_b8 to + reconstruct the MFMA B fragment. Scale + prepacking resolves the source orientation before launch, so both packed + scale tensors use the common [K/128, dim] kernel representation. + """ + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL MXFP8 NN expects rank-2 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 MXFP8 NN operands: " + f"A{tuple(a.shape)} 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 MXFP8 NN 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}" + ) + 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 not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + "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") + + if k % SCALE_GROUP_SIZE != 0: + raise ValueError( + f"K={k} must be divisible by MXFP8 scale group size " + f"{SCALE_GROUP_SIZE}" + ) + + 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 {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 as torch.uint8") + if a_scale.device != a.device or b_scale.device != a.device: + raise ValueError("A, B, scales, and D must be on the same device") + + a_scale_hk = pack_mx32_scales_for_hk( + a_scale, + source_colwise=False, + ) + b_scale_hk = pack_mx32_scales_for_hk( + b_scale, + source_colwise=True, + ) + + _debug( + f"NN kernel inputs: a={tuple(a.shape)}, b={tuple(b.shape)}, " + f"a_scale_hk={tuple(a_scale_hk.shape)}, " + f"b_scale_hk={tuple(b_scale_hk.shape)}, D={tuple(D.shape)}" + ) + + do_gemm( + a, + a_scale_hk, + b, + b_scale_hk, + D.view(m, n), + stream=stream, + ) + return D + + +__all__ = [ + "BLOCK_M", + "BLOCK_N", + "BLOCK_K", + "SCALE_GROUP_SIZE", + "mxfp8_matmul", +] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nt.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nt.py new file mode 100644 index 000000000..7139edf5b --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nt.py @@ -0,0 +1,1474 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL MXFP8 NT 4-wave GEMM implementation. + +This specialization preserves the validated MXFP8 TN compute, scale, MFMA, +accumulator, and epilogue pipelines while applying the validated +ds_read_b64_tr_b8 path to both operands. A is physically [K, M] and B is +physically [K, N]. Each source tile is staged as XOR-swizzled [K128, X128] LDS. + +Both raw scale tensors are columnwise, [K/32, M] and [K/32, N]. +Orientation-aware prepacking converts them to the common iteration-major +[K/128, dim] uint32 representation consumed by the kernel.""" + +import functools +import os + +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 .exceptions import FlyDSLUnsupportedError +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 _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, + *, + source_colwise: bool = False, +) -> torch.Tensor: + """Pack raw E8M0 scales as iteration-major ``[K/128, dim]`` uint32. + + ``source_colwise=False`` consumes TE rowwise scales ``[dim, K/32]``. + ``source_colwise=True`` consumes TE columnwise scales ``[K/32, dim]``. + + Both paths produce the same packed representation consumed by every + TN/NN/NT MXFP8 kernel specialization. + """ + 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)}" + ) + + if source_colwise: + qk, dim = scales_u8.shape + if qk % 4 != 0: + raise ValueError( + f"Columnwise scale K dimension must be divisible by 4 K32 groups, got {qk}" + ) + s32 = scales_u8.contiguous().view(qk // 4, 4, dim).to(torch.int32) + return ( + s32[:, 0, :] + | (s32[:, 1, :] << 8) + | (s32[:, 2, :] << 16) + | (s32[:, 3, :] << 24) + ).contiguous() + + dim, qk = scales_u8.shape + if qk % 4 != 0: + raise ValueError( + f"Rowwise scale K dimension must be divisible by 4 K32 groups, got {qk}" + ) + + s32 = scales_u8.contiguous().view(dim, 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, + *, + source_colwise: bool = False, +) -> torch.Tensor: + """Convert raw TE E8M0 scales to ``[K/128, dim]`` MFMA-ready words.""" + scale_iter = pack_mx32_scales_iter( + scales_u8, + source_colwise=source_colwise, + ) + dim = scales_u8.shape[1] if source_colwise else scales_u8.shape[0] + + if dim % 64 != 0: + raise ValueError( + f"Scale outer dimension={dim} must be a multiple of 64 for HK MFMA packing" + ) + + device = scales_u8.device + row = torch.arange(dim, device=device, dtype=torch.int64) + row_within_16 = row % 16 + k_subgroup = (row // 16) % 4 + tile = row // 64 + + packed = torch.zeros_like(scale_iter) + 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, dim) + ) & 0xFF + packed |= byte_value << (group * 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, + 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 + 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 + + 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[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, 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) + + 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) + 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. + bx_m_idx = fx.Index(bx_m) + by_n_idx = fx.Index(by_n) + + tx_i32 = fx.Int32(tx) + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # NT storage is K-major for both operands: + # A [K, M] + # B [K, N] + # + # Read each K-by-X source tile in XOR-swizzled coordinate order and + # write it linearly to LDS. swizzle_128 is self-inverse, producing the + # physical [K128, X128] image consumed by ds_read_b64_tr_b8. + gl_off_a = compute_global_swizzle( + lane, + wave_id, + c_m, + LOAD_PASSES_HALF, + preshuffled=False, + ) + gl_off_b = compute_global_swizzle( + lane, + wave_id, + c_n, + 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 _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(): + # A-bottom and the B slices are transpose reads in NT. + 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(): + # Each prefetched A/B fragment uses two DS_READ_TR instructions. + for _ in range_constexpr(8): + rocdl.sched_dsrd(4) + 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): + # A is physically [K, M]. Copy + # A[k_base:k_base+128, bx_m+subtile*128:...] + # into one XOR-swizzled physical LDS half-page [K128, M128]. + global_base = ( + k_base * fx.Index(c_m) + + bx_m_idx + + fx.Index(subtile * (BLOCK_M // 2)) + ) + 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): + # B is physically [K, N]. Copy + # B[k_base:k_base+128, by_n+subtile*128:...] + # into one XOR-swizzled physical LDS half-page [K128, N128]. + global_base = ( + k_base * fx.Index(c_n) + + by_n_idx + + fx.Index(subtile * (BLOCK_N // 2)) + ) + 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 pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + + def load_transposed_frag_half(lds_page, local_x_tile, half): + """Load one K64 portion of a fixed-X MFMA fragment. + + This is the inverse mapping validated against the working ordinary + LDS fragment: + + source_k = lane_div_16*16 + lane_in_16//2 + source_x = local_x_tile + (lane_in_16&1)*8 + + ``base ^ 0x440`` advances logical K by 8 under swizzle_128. + The 0x2000 DS immediate advances logical K by 64. + """ + lane_div16_i32 = fx.Int32(lane_div_16) + lane_in16_i32 = fx.Int32(lane_mod_16) + source_k = ( + lane_div16_i32 * fx.Int32(16) + + lane_in16_i32 // fx.Int32(2) + ) + source_x = ( + fx.Int32(local_x_tile) + + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) + ) + + physical_k, physical_x = swizzle_128(source_k, source_x) + base = physical_k * fx.Int32(128) + physical_x + other = base ^ fx.Int32(0x440) + immediate_offset = 0 if half == 0 else 0x2000 + + return s2r.load_one_transpose( + lds_page, + base, + other, + immediate_offset=immediate_offset, + ) + + + def load_transposed_frag(lds_page, local_x_tile): + x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) + x1 = load_transposed_frag_half(lds_page, local_x_tile, 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, 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"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, + ) + + 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"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, + ) + + 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 + 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. + # 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_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): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_scales = scale_tile[2] if sn == 0 else scale_tile[3] + local_n_tile = ( + subtile_n_idx * fx.Index(SUBTILE_N) + + fx.Index(ni * MFMA_N) + - fx.Index(sn * (BLOCK_N // 2)) + ) + b_ni = load_transposed_frag(lds_b[sn], local_n_tile) + return b_ni, b_scales[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): + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + local_m_tile = ( + subtile_m_idx * fx.Index(SUBTILE_M) + + fx.Index(mi * MFMA_M) + - fx.Index(sm * (BLOCK_M // 2)) + ) + return load_transposed_frag_half( + lds_a[sm], + local_m_tile, + 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, + 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, + ) + + + +def do_gemm( + A: torch.Tensor, + As: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + C: torch.Tensor, + stream=None, +): + """Launch MXFP8 NT core from K-major A [K,M] and B [K,N].""" + K_runtime, M_runtime = A.shape + Kb_runtime, N_runtime = B.shape + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + 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}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 NT 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 NT 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 NT 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 MXFP8 NT 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" + 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})" + ) + 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}" + ) + + tensors = (A, As, B, Bs, C) + if any(t.device != A.device for t in tensors[1:]): + raise ValueError("A, B, packed scales, and C must be on the same device") + + 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) + As_arg = As.contiguous().view(-1) + Bs_arg = Bs.contiguous().view(-1) + C_arg = C.contiguous().view(-1) + + launch = _cached_launch( + int(K_runtime), + A.dtype, + B.dtype, + C.dtype, + ) + launch( + A_arg, + As_arg, + B_arg, + Bs_arg, + C_arg, + M_runtime, + 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 MXFP8 NT GEMM with transpose-read A and B operands. + + Contract: + a: [K, M] row-major FP8 payload + a_scale: [K/32, M] raw columnwise E8M0 scales + b: [K, N] row-major FP8 payload + b_scale: [K/32, N] raw columnwise E8M0 scales + D: [M, N] float16, bfloat16, or float32 output + + Both operands remain K-major. Each is staged as an XOR-swizzled + [K128, X128] LDS image and reconstructed with ds_read_b64_tr_b8. + """ + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL MXFP8 NT expects rank-2 operands, got " + f"a={tuple(a.shape)} and b={tuple(b.shape)}" + ) + + k, m = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Incompatible MXFP8 NT operands: " + f"A{tuple(a.shape)} 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 MXFP8 NT 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}" + ) + 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 not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + "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") + + if k % SCALE_GROUP_SIZE != 0: + raise ValueError( + f"K={k} must be divisible by MXFP8 scale group size " + f"{SCALE_GROUP_SIZE}" + ) + + expected_a_scale = (k // SCALE_GROUP_SIZE, m) + 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 {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 as torch.uint8") + if a_scale.device != a.device or b_scale.device != a.device: + raise ValueError("A, B, scales, and D must be on the same device") + + a_scale_hk = pack_mx32_scales_for_hk( + a_scale, + source_colwise=True, + ) + b_scale_hk = pack_mx32_scales_for_hk( + b_scale, + source_colwise=True, + ) + + _debug( + f"NT kernel inputs: a={tuple(a.shape)}, b={tuple(b.shape)}, " + f"a_scale_hk={tuple(a_scale_hk.shape)}, " + f"b_scale_hk={tuple(b_scale_hk.shape)}, D={tuple(D.shape)}" + ) + + do_gemm( + a, + a_scale_hk, + b, + b_scale_hk, + D.view(m, n), + stream=stream, + ) + return D + + +__all__ = [ + "BLOCK_M", + "BLOCK_N", + "BLOCK_K", + "SCALE_GROUP_SIZE", + "mxfp8_matmul", +] From 2524a087643597bfa7abeef8251ec218dbca57e7 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 28 Jul 2026 15:28:41 +0000 Subject: [PATCH 20/31] gemm wrappers patch --- .../flydsl_kernels/gemm/gemm_wrappers.py | 228 +++++++----------- 1 file changed, 91 insertions(+), 137 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index e9e861631..3b4ddd737 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -257,7 +257,7 @@ def _flatten_rowwise(t: torch.Tensor, name: str) -> torch.Tensor: def _flatten_columnwise(t: torch.Tensor, name: str) -> torch.Tensor: - """Flatten TE columnwise storage as [last_dim, product(leading_dims)].""" + """Flatten TE columnwise storage while preserving its leading dimension.""" if t.ndim < 2: raise ValueError( f"FlyDSL GEMM expects {name} to have rank >= 2, got {tuple(t.shape)}" @@ -312,42 +312,6 @@ def _canonicalize_blas_operands( return a_flydsl, b_flydsl, m, n, k -def _resolve_output_shape( - A, - transa, - B, - transb, - D, - *, - m, - n, - backend_name, -): - """Resolve TE's public output shape independently of kernel storage. - - FlyDSL kernels always write a flattened row-major ``[M, N]`` matrix. - TE's public tensor may retain leading dimensions (for example - ``[sequence, batch, hidden]``). Quantized rowwise/columnwise payloads are - physical storage views and must never be used to infer that public shape. - - A caller-provided ``D`` is authoritative. Otherwise derive the logical - shape from the original TE operands, before selecting or flattening any - backing storage. - """ - if D is not None: - output_shape = torch.Size(D.shape) - else: - output_shape = _get_gemm_output_shape(A, transa, B, transb) - - if _product(output_shape) != m * n: - raise FlyDSLUnsupportedError( - f"FlyDSL {backend_name} logical output shape " - f"{tuple(output_shape)} does not match flattened kernel shape " - f"{(m, n)}" - ) - return output_shape - - def _validate_or_allocate_output( D, *, @@ -405,19 +369,16 @@ 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 ) - output_shape = _resolve_output_shape( - A, - transa, - B, - transb, - D, - m=m, - n=n, - backend_name=backend_name, - ) + 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 @@ -556,7 +517,6 @@ def _mxfp8_logical_shape(t, name: str) -> torch.Size: ) return torch.Size(data.shape) - def _flatten_mxfp8_scale( t: torch.Tensor, name: str, @@ -592,6 +552,7 @@ def _flatten_mxfp8_scale( ) return t + def _run_mxfp8( A, transa, @@ -812,14 +773,19 @@ def _run_mxfp8( ) return D + def _select_fp8_storage_for_layout(A, transa, B, transb): """Select the exact existing TE FP8 backing required by each layout. - Fixed zero-copy routes: + Fixed zero-copy routes selected for the final kernel contracts: - TN: A._data, B._data - NN: A._transpose, B._data - NT: A._transpose, B._transpose + TN: wrapper swaps B._data/A._data -> [M,K], [N,K] + NN: wrapper swaps B._data/A._transpose -> [M,K], [N,K] + NT: wrapper swaps B._data/A._data -> [K,M], [K,N] + + In particular, NT must use the contiguous rowwise K-major payloads. + Passing ``_transpose.transpose(0, 1)`` would create strided views and + force the NT adapter to materialize them before launch. """ layout = (bool(transa), bool(transb)) @@ -842,13 +808,18 @@ def _select_fp8_storage_for_layout(A, transa, B, transb): B_data = _flatten_rowwise(B_payload, B_storage) elif layout == (False, True): # NT / dW - A_payload = _get_fp8_columnwise_payload(A, "A") - A_storage = "A._transpose" - A_data = _flatten_columnwise(A_payload, A_storage) + # fp8_gemm_nt consumes contiguous K-major operands directly: + # kernel a = B._data [K, M] + # kernel b = A._data [K, N] + # Select rowwise storage here so the ownership swap in _run_fp8 is + # zero-copy and no noncontiguous transpose view reaches the kernel. + A_payload = _get_fp8_rowwise_payload(A, "A") + A_storage = "A._data" + A_data = _flatten_rowwise(A_payload, A_storage) - B_payload = _get_fp8_columnwise_payload(B, "B") - B_storage = "B._transpose" - B_data = _flatten_columnwise(B_payload, B_storage) + B_payload = _get_fp8_rowwise_payload(B, "B") + B_storage = "B._data" + B_data = _flatten_rowwise(B_payload, B_storage) else: raise FlyDSLUnsupportedError( @@ -874,21 +845,7 @@ def _run_fp8( *, output_dtype: torch.dtype, ): - """Dispatch tensor-wise FP8 through one canonical operand contract. - - First select the exact existing TE storage required by the BLAS flags. - Then canonicalize both payloads and scales identically: - - a_flydsl, b_flydsl = op(B), op(A) - a_scale, b_scale = B scale, A scale - - Every kernel is called with: - - matmul(a_flydsl, a_scale, b_flydsl, b_scale, D) - - The layout-specific kernels differ only in the physical layouts they - expect for canonicalized ``a_flydsl`` and ``b_flydsl``. - """ + """Dispatch tensor-wise FP8 using exact per-kernel storage contracts.""" supported_fp8_dtypes = ( tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2, @@ -913,31 +870,20 @@ def _run_fp8( if not isinstance(scale, torch.Tensor): raise FlyDSLUnsupportedError(f"{name} is not populated") if scale.dtype != torch.float32 or scale.numel() != 1: - raise ValueError( + raise FlyDSLUnsupportedError( f"{name} must contain exactly one FP32 tensor-wise inverse " f"scale, got dtype={scale.dtype}, shape={tuple(scale.shape)}" ) layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" - dispatch = { - (True, False): ("TN", fp8_matmul), - (False, False): ("NN", fp8_matmul_nn), - (False, True): ("NT", fp8_matmul_nt), - } - try: - kernel_layout, matmul = dispatch[(bool(transa), bool(transb))] - except KeyError as exc: - raise FlyDSLUnsupportedError( - "FlyDSL GEMM does not support transa=True, transb=True (TT)" - ) from exc ( A_data, A_storage, - A_physical_shape, + A_payload_shape, B_data, B_storage, - B_physical_shape, + B_payload_shape, ) = _select_fp8_storage_for_layout( A, bool(transa), @@ -953,89 +899,97 @@ def _run_fp8( b_storage=B_storage, ) - # Scales follow the original TE tensors after BLAS operand ownership swap. a_scale = B_scale_inv b_scale = A_scale_inv if layout == "TN": - # a_flydsl = B._data [M,K] - # b_flydsl = A._data [N,K] + matmul = fp8_matmul + kernel_layout = "TN" + a_flydsl = B_data b_flydsl = A_data + m, k = a_flydsl.shape n, kb = b_flydsl.shape elif layout == "NN": - # a_flydsl = B._data [M,K] - # b_flydsl = A._transpose flattened as [N,K] + matmul = fp8_matmul_nn + kernel_layout = "NN" + a_flydsl = B_data b_flydsl = A_data + m, k = a_flydsl.shape n, kb = b_flydsl.shape + elif layout == "NT": + matmul = fp8_matmul_nt + kernel_layout = "NT" + + # Exact fp8_gemm_nt contract, with no view or materialization: + # a_flydsl = B._data [K, M] + # b_flydsl = A._data [K, N] + a_flydsl = B_data + b_flydsl = A_data + + k, m = a_flydsl.shape + kb, n = b_flydsl.shape + else: - # TE columnwise backings are contiguous allocations exposed as: - # B._transpose flattened [M,K] - # A._transpose flattened [N,K] - # - # fp8_gemm_nt consumes those same bytes with K-major tensor metadata: - # kernel_a [K,M] aliases B._transpose - # kernel_b [K,N] aliases A._transpose - m, k = B_data.shape - n, kb = A_data.shape - a_flydsl = B_data.view(k, m) - b_flydsl = A_data.view(kb, n) + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) + + if not a_flydsl.is_contiguous() or not b_flydsl.is_contiguous(): + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 {layout} kernel contract requires contiguous final " + f"operands, got a={tuple(a_flydsl.shape)} " + f"stride={tuple(a_flydsl.stride())} and " + f"b={tuple(b_flydsl.shape)} stride={tuple(b_flydsl.stride())}" + ) if kb != k: raise FlyDSLUnsupportedError( f"FlyDSL FP8 {layout} selected incompatible physical backings: " f"{B_storage}={tuple(B_data.shape)} and " - f"{A_storage}={tuple(A_data.shape)}" + f"{A_storage}={tuple(A_data.shape)}; " + f"kernel operands are {tuple(a_flydsl.shape)} and " + f"{tuple(b_flydsl.shape)}" ) + if D is not None: + logical_output_shape = torch.Size(D.shape) + elif layout in ("TN", "NN"): + logical_output_shape = torch.Size((*B_payload_shape[:-1], n)) + else: + logical_output_shape = torch.Size((m, n)) + if _product(logical_output_shape) != m * n: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 {layout} logical output shape " + f"{tuple(logical_output_shape)} does not match kernel output " + f"shape {(m, n)}" + ) + + D = _validate_or_allocate_output( + D, + shape=logical_output_shape, + dtype=output_dtype, + device=a_flydsl.device, + backend_name=f"FP8 {kernel_layout}", + ) + _fp8_debug( f"dispatch entry: transa={bool(transa)}, transb={bool(transb)}, " f"layout={layout}, selected_kernel={matmul.__module__}.{matmul.__name__}" ) - _fp8_debug( - f"selected TE storage: A={A_storage}, B={B_storage}" - ) + _fp8_debug(f"selected TE storage: A={A_storage}, B={B_storage}") _fp8_tensor_debug(f"selected/{A_storage}", A_data) _fp8_tensor_debug(f"selected/{B_storage}", B_data) - _fp8_debug( - "canonical contract: " - "matmul(a_flydsl, a_scale, b_flydsl, b_scale, D)" - ) _fp8_tensor_debug("a_flydsl", a_flydsl) _fp8_tensor_debug("b_flydsl", b_flydsl) _fp8_scale_debug("a_scale", a_scale) _fp8_scale_debug("b_scale", b_scale) - _fp8_debug( - f"canonical ownership: a_flydsl<-TE B, b_flydsl<-TE A; " - f"derived M={m}, N={n}, K={k}" - ) - - # Kernel storage is always flattened, but the public TE result must retain - # the logical leading dimensions of the original operands when D is not - # preallocated. Never infer the public shape from _data/_transpose. - logical_output_shape = _resolve_output_shape( - A, - transa, - B, - transb, - D, - m=m, - n=n, - backend_name=f"FP8 {layout}", - ) - - D = _validate_or_allocate_output( - D, - shape=logical_output_shape, - dtype=output_dtype, - device=a_flydsl.device, - backend_name=f"FP8 {kernel_layout}", - ) + _fp8_debug(f"derived M={m}, N={n}, K={k}") _fp8_tensor_debug("output/D", D) matmul( From de6d22ad53a960e076895bf19da68b63c7b71a30 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 28 Jul 2026 16:43:43 +0000 Subject: [PATCH 21/31] correct FP8 NT storage contract --- .../flydsl_kernels/gemm/fp8_gemm_nt.py | 123 +++++++----------- .../flydsl_kernels/gemm/gemm_wrappers.py | 37 +++--- 2 files changed, 62 insertions(+), 98 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py index e3f8b2cf5..7f585fe0d 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py @@ -4,16 +4,17 @@ """FlyDSL tensor-wise FP8 NT 4-wave GEMM kernel. -This NT variant preserves the working 4-wave pipeline while applying the -validated ``ds_read_b64_tr_b8`` contract to both operands. A is physically -[K, M] and B is physically [K, N]. Each 128x128 source tile is staged into an -XOR-swizzled physical LDS image [K128, X128], and four transpose reads rebuild -the exact ordinary MFMA fragment for one fixed M or N coordinate. +This NT variant is the NN transpose-storage path applied to both operands. +The public contract remains C = A @ B.T with A physically [M, K] and B +physically [N, K]. During staging, each operand's 128x128 half-page is +transposed into an XOR-swizzled physical LDS image [K128, X128]. The validated +``ds_read_b64_tr_b8`` sequence then reconstructs the ordinary MFMA fragment +for one fixed M or N coordinate. The kernel specializes on K at compile time because the K128 loop is fully hand-unrolled. M/N are runtime launch dimensions. The public entry point -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 +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]. Operand normalization is performed by the Transformer Engine wrapper. @@ -36,10 +37,8 @@ # Transformer Engine-local FlyDSL utilities. from .fp8_gemm_utils import ( - G2SLoader, + G2STransposeLoader, S2RLoader, - compute_global_swizzle, - make_fp8_buffer_tensor, pack_i32x4_i32x8, swizzle_128, ) @@ -296,12 +295,6 @@ def kernel_gemm( 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 = ( @@ -334,42 +327,16 @@ def kernel_gemm( wave_id = tx_i32 // fx.Int32(WARP_SIZE) lane = tx_i32 % fx.Int32(WARP_SIZE) - # NT storage is K-major for both operands: - # A [K, M] - # B [K, N] + # Both operands arrive in transpose storage: + # A [M, K] + # B [N, K] # - # Read each global 128x128 K-by-X tile in XOR-swizzled coordinate order - # and write it linearly to LDS. Because swizzle_128 is self-inverse, - # this produces the physical XOR-swizzled LDS image [K128, X128] - # consumed by ds_read_b64_tr_b8. - gl_off_a = compute_global_swizzle( - lane, - wave_id, - c_m, - LOAD_PASSES_HALF, - preshuffled=False, - ) - gl_off_b = compute_global_swizzle( - lane, - wave_id, - c_n, - 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, - ) + # Apply the NN global-to-LDS transpose staging path independently to + # each operand. Each row-major [X128, K128] source half-page becomes + # the XOR-swizzled physical LDS image [K128, X128] consumed by + # ds_read_b64_tr_b8. + a_g2s = G2STransposeLoader(A, K, wave_id) + b_g2s = G2STransposeLoader(B, K, wave_id) s2r = S2RLoader(fx.Int32(0), 1) layout_lane16 = fx.make_layout((4, 16), (16, 1)) @@ -473,26 +440,26 @@ def hot_loop_scheduler_q_prefetch_4n(): rocdl.sched_barrier(0) def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): - # A is physically [K, M]. Copy - # A[k_base:k_base+128, bx_m+subtile*128:...] - # into one XOR-swizzled physical LDS half-page [K128, M128]. - global_base = ( - k_base * fx.Index(c_m) - + bx_m_idx - + fx.Index(subtile * (BLOCK_M // 2)) + # Load row-major global A[M, K], but write the half-page as + # XOR-swizzled physical LDS [K128, M128]. + global_m_base = bx_m_idx + fx.Index(subtile * (BLOCK_M // 2)) + a_g2s.load_one( + lds_a[subtile], + global_m_base, + k_base, + pass_in_subtile, ) - 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): - # B is physically [K, N]. Copy - # B[k_base:k_base+128, by_n+subtile*128:...] - # into one XOR-swizzled physical LDS half-page [K128, N128]. - global_base = ( - k_base * fx.Index(c_n) - + by_n_idx - + fx.Index(subtile * (BLOCK_N // 2)) + # Load row-major global B[N, K], but write the half-page as + # XOR-swizzled physical LDS [K128, N128]. + global_n_base = by_n_idx + fx.Index(subtile * (BLOCK_N // 2)) + b_g2s.load_one( + lds_b[subtile], + global_n_base, + k_base, + pass_in_subtile, ) - 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): @@ -1093,18 +1060,18 @@ def fp8_matmul( c: torch.Tensor, stream=None, ): - """Launch NT tensor-wise FP8 GEMM with transpose-read A/B fragments. + """Launch NT tensor-wise FP8 GEMM from both transpose backings. Contract: - a: [K, M] FP8 payload + a: [M, K] FP8 payload a_scale_inv: one-element FP32 inverse quantization scale - b: [K, N] FP8 payload + b: [N, K] FP8 payload b_scale_inv: one-element FP32 inverse quantization scale c: [M, N] float16, bfloat16, or float32 output - Both operands remain K-major in global memory. Each tile is staged as a - swizzled physical [K128, X128] LDS image and read with the validated - four-instruction ds_read_b64_tr_b8 fragment contract. + Both operands remain row-major [outer, K] in global memory. Each tile is + transposed during GMEM-to-LDS staging, then read with the validated + ds_read_b64_tr_b8 fragment contract. """ if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): raise TypeError("FlyDSL FP8 NT GEMM expects plain torch.Tensor payloads") @@ -1121,8 +1088,8 @@ def fp8_matmul( f"got A={a.dtype} and B={b.dtype}" ) - k, m = a.shape - kb, n = b.shape + m, k = 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)}" @@ -1163,9 +1130,9 @@ def doGemm( stream=None, use_xcd_remap: bool = True, ): - """Launch optimized NT FP8 GEMM from K-major A [K,M] and B [K,N].""" - K_runtime, M_runtime = A.shape - Kb_runtime, N_runtime = B.shape + """Launch optimized NT FP8 GEMM with A [M,K] and B [N,K].""" + 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 FP8 dtype: {A.dtype}" assert B.dtype in supported_fp8_dtypes, f"unsupported B FP8 dtype: {B.dtype}" diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 3b4ddd737..4e68d01c2 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -781,11 +781,11 @@ def _select_fp8_storage_for_layout(A, transa, B, transb): TN: wrapper swaps B._data/A._data -> [M,K], [N,K] NN: wrapper swaps B._data/A._transpose -> [M,K], [N,K] - NT: wrapper swaps B._data/A._data -> [K,M], [K,N] + NT: wrapper swaps B._transpose/A._transpose -> [M,K], [N,K] - In particular, NT must use the contiguous rowwise K-major payloads. - Passing ``_transpose.transpose(0, 1)`` would create strided views and - force the NT adapter to materialize them before launch. + NT is the NN transpose-storage path applied to both operands. Both + transpose allocations stay contiguous in their native [outer,K] shapes; + no tensor transpose, reshape reinterpretation, or materialization occurs. """ layout = (bool(transa), bool(transb)) @@ -808,18 +808,16 @@ def _select_fp8_storage_for_layout(A, transa, B, transb): B_data = _flatten_rowwise(B_payload, B_storage) elif layout == (False, True): # NT / dW - # fp8_gemm_nt consumes contiguous K-major operands directly: - # kernel a = B._data [K, M] - # kernel b = A._data [K, N] - # Select rowwise storage here so the ownership swap in _run_fp8 is - # zero-copy and no noncontiguous transpose view reaches the kernel. - A_payload = _get_fp8_rowwise_payload(A, "A") - A_storage = "A._data" - A_data = _flatten_rowwise(A_payload, A_storage) + # NT extends NN's transpose-storage handling to both operands. + # After ownership swap, B._transpose is kernel A [M,K] and + # A._transpose is kernel B [N,K]. + A_payload = _get_fp8_columnwise_payload(A, "A") + A_storage = "A._transpose" + A_data = _flatten_columnwise(A_payload, A_storage) - B_payload = _get_fp8_rowwise_payload(B, "B") - B_storage = "B._data" - B_data = _flatten_rowwise(B_payload, B_storage) + B_payload = _get_fp8_columnwise_payload(B, "B") + B_storage = "B._transpose" + B_data = _flatten_columnwise(B_payload, B_storage) else: raise FlyDSLUnsupportedError( @@ -926,14 +924,13 @@ def _run_fp8( matmul = fp8_matmul_nt kernel_layout = "NT" - # Exact fp8_gemm_nt contract, with no view or materialization: - # a_flydsl = B._data [K, M] - # b_flydsl = A._data [K, N] + # Correct fp8_gemm_nt contract: NN's [outer,K] transpose-storage + # path applied to both operands. a_flydsl = B_data b_flydsl = A_data - k, m = a_flydsl.shape - kb, n = b_flydsl.shape + m, k = a_flydsl.shape + n, kb = b_flydsl.shape else: raise FlyDSLUnsupportedError( From 38ccfb230fa22f7bd46f5160a4304795f04ba863 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 28 Jul 2026 19:31:00 +0000 Subject: [PATCH 22/31] Unify FP8 TN, NN, and NT GEMM paths Route all tensorwise FP8 layouts through the common FP8 GEMM core after wrapper-side storage backing selection. Remove the redundant FP8 NN and NT kernel variants since columnwise FP8 storage already provides the required materialized transpose. --- .../flydsl_kernels/gemm/fp8_gemm_nn.py | 1212 ----------------- .../flydsl_kernels/gemm/fp8_gemm_nt.py | 1188 ---------------- .../flydsl_kernels/gemm/gemm_wrappers.py | 74 +- 3 files changed, 24 insertions(+), 2450 deletions(-) delete mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py delete mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py deleted file mode 100644 index b24d061b6..000000000 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py +++ /dev/null @@ -1,1212 +0,0 @@ -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. - -"""FlyDSL tensor-wise FP8 NN 4-wave GEMM kernel. - -This NN variant preserves the working 4-wave pipeline and the kernel contract -C = A @ B.T. A is physically [M, K] and B is physically [N, K]. During -staging, each B 128x128 half-page is transposed into XOR-swizzled physical LDS -[K128, N128]. The validated four-read ``ds_read_b64_tr_b8`` sequence then -reconstructs exactly the ordinary B[N, K] fragment consumed by the production -FP8 MFMA. - -The kernel specializes on K at compile time because the K128 loop is fully -hand-unrolled. M/N are runtime launch dimensions. The public entry point and private optimized core consume independently typed -FP8 E4M3 or E5M2 A/B tensors shaped [M, K] and [N, K], one FP32 inverse scale -per operand, and write float16, bfloat16, or float32 C shaped [M, N]. Operand -normalization is performed by the Transformer Engine wrapper. - -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, - G2STransposeLoader, - 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, - 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 - gA = make_fp8_buffer_tensor(A, a_f8_ir_t) - a_div = fx.logical_divide(gA, 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) - - # A keeps the ordinary row-major [M, K] direct-to-LDS path. - gl_off_a = 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 arrives row-major [N, K]. Stage each source 16-byte K vector into - # the transposed XOR-swizzled physical LDS image [K128, N128] required - # by the validated ds_read_b64_tr_b8 inverse mapping. - b_g2s = G2STransposeLoader(B, K, 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(): - 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): - # Load row-major global B[N, K], but write the half-page as - # XOR-swizzled physical LDS [K128, N128]. - global_n_base = by_n_idx + fx.Index(subtile * (BLOCK_N // 2)) - b_g2s.load_one( - lds_b[subtile], - global_n_base, - k_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_transpose(lds_page, local_n_tile): - # Exact inverse mapping validated against the ordinary B[N, K] - # production MFMA fragment: - # - # source_k = lane_div_16*16 + lane_in_16//2 - # source_n = local_n_tile + (lane_in_16&1)*8 - # - # base^0x440 advances logical K by 8 under the 128-byte XOR - # swizzle. The DS immediate 0x2000 advances logical K by 64. - lane_div16_i32 = fx.Int32(lane_div_16) - lane_in16_i32 = fx.Int32(lane_mod_16) - source_k = ( - lane_div16_i32 * fx.Int32(16) - + lane_in16_i32 // fx.Int32(2) - ) - source_n = ( - fx.Int32(local_n_tile) - + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) - ) - - physical_k, physical_n = swizzle_128(source_k, source_n) - base = physical_k * fx.Int32(BLOCK_N // 2) + physical_n - other = base ^ fx.Int32(0x440) - - x0 = s2r.load_one_transpose( - lds_page, - base, - other, - immediate_offset=0, - ) - x1 = s2r.load_one_transpose( - lds_page, - base, - other, - immediate_offset=0x2000, - ) - 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) - local_n_tile = ( - subtile_n_idx * fx.Index(SUBTILE_N) - + fx.Index(ni * MFMA_N) - - fx.Index(sn * (BLOCK_N // 2)) - ) - return load_b_frag_transpose(lds_b[sn], local_n_tile) - - 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, - 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, -): - """Launch correctness-first NN tensor-wise FP8 GEMM. - - Contract: - a: [M, K] FP8 payload - a_scale_inv: one-element FP32 inverse quantization scale - b: [N, K] FP8 payload - b_scale_inv: one-element FP32 inverse quantization scale - c: [M, N] float16, bfloat16, or float32 output - - B remains [N, K] through GMEM->LDS. The kernel performs a naive scalar - LDS gather along K for a fixed N row, constructing the same MFMA B - fragments as the optimized transpose-read path. This variant intentionally - does not use ds_read_b64_tr_b8. - """ - if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): - raise TypeError("FlyDSL FP8 NN 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 NN GEMM expects E4M3 or E5M2 payloads, " - f"got A={a.dtype} and B={b.dtype}" - ) - - m, k = 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") - - 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 NN FP8 GEMM with C = A @ B.T, A [M,K], B [N,K].""" - 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 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 NN GEMM requires M to be a multiple of {_BLOCK_M}, got M={M_runtime}" - ) - if N_runtime % _BLOCK_N != 0: - raise FlyDSLUnsupportedError( - f"FlyDSL FP8 NN GEMM requires N to be a multiple of {_BLOCK_N}, got N={N_runtime}" - ) - if K_runtime % _BLOCK_K != 0: - raise FlyDSLUnsupportedError( - f"FlyDSL FP8 NN GEMM requires K to be a multiple of {_BLOCK_K}, got K={K_runtime}" - ) - num_k_tiles = K_runtime // _BLOCK_K - if num_k_tiles < 4: - raise FlyDSLUnsupportedError( - f"FlyDSL FP8 NN 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( - 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 deleted file mode 100644 index 7f585fe0d..000000000 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py +++ /dev/null @@ -1,1188 +0,0 @@ -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. - -"""FlyDSL tensor-wise FP8 NT 4-wave GEMM kernel. - -This NT variant is the NN transpose-storage path applied to both operands. -The public contract remains C = A @ B.T with A physically [M, K] and B -physically [N, K]. During staging, each operand's 128x128 half-page is -transposed into an XOR-swizzled physical LDS image [K128, X128]. The validated -``ds_read_b64_tr_b8`` sequence then reconstructs the ordinary MFMA fragment -for one fixed M or N coordinate. - -The kernel specializes on K at compile time because the K128 loop is fully -hand-unrolled. M/N are runtime launch dimensions. The public entry point -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]. Operand normalization is performed by the -Transformer Engine wrapper. - -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 ( - G2STransposeLoader, - S2RLoader, - 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_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. The global-offset helpers combine - # 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) - - # Both operands arrive in transpose storage: - # A [M, K] - # B [N, K] - # - # Apply the NN global-to-LDS transpose staging path independently to - # each operand. Each row-major [X128, K128] source half-page becomes - # the XOR-swizzled physical LDS image [K128, X128] consumed by - # ds_read_b64_tr_b8. - a_g2s = G2STransposeLoader(A, K, wave_id) - b_g2s = G2STransposeLoader(B, K, 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(): - 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(4) - rocdl.sched_mfma(4) - rocdl.sched_barrier(0) - - def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): - # Load row-major global A[M, K], but write the half-page as - # XOR-swizzled physical LDS [K128, M128]. - global_m_base = bx_m_idx + fx.Index(subtile * (BLOCK_M // 2)) - a_g2s.load_one( - lds_a[subtile], - global_m_base, - k_base, - pass_in_subtile, - ) - - def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): - # Load row-major global B[N, K], but write the half-page as - # XOR-swizzled physical LDS [K128, N128]. - global_n_base = by_n_idx + fx.Index(subtile * (BLOCK_N // 2)) - b_g2s.load_one( - lds_b[subtile], - global_n_base, - k_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 pack_frag_halves(x0, x1): - return pack_i32x4_i32x8(x0, x1) - - def load_transposed_frag_half(lds_page, local_x_tile, half): - """Load one K64 portion of a fixed-X MFMA fragment. - - This is the inverse mapping validated against the working ordinary - LDS fragment: - - source_k = lane_div_16*16 + lane_in_16//2 - source_x = local_x_tile + (lane_in_16&1)*8 - - ``base ^ 0x440`` advances logical K by 8 under swizzle_128. - The 0x2000 DS immediate advances logical K by 64. - """ - lane_div16_i32 = fx.Int32(lane_div_16) - lane_in16_i32 = fx.Int32(lane_mod_16) - source_k = ( - lane_div16_i32 * fx.Int32(16) - + lane_in16_i32 // fx.Int32(2) - ) - source_x = ( - fx.Int32(local_x_tile) - + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) - ) - - physical_k, physical_x = swizzle_128(source_k, source_x) - base = physical_k * fx.Int32(128) + physical_x - other = base ^ fx.Int32(0x440) - immediate_offset = 0 if half == 0 else 0x2000 - - return s2r.load_one_transpose( - lds_page, - base, - other, - immediate_offset=immediate_offset, - ) - - def load_transposed_frag(lds_page, local_x_tile): - x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) - x1 = load_transposed_frag_half(lds_page, local_x_tile, 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_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) - local_n_tile = ( - subtile_n_idx * fx.Index(SUBTILE_N) - + fx.Index(ni * MFMA_N) - - fx.Index(sn * (BLOCK_N // 2)) - ) - return load_transposed_frag(lds_b[sn], local_n_tile) - - 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) - local_m_tile = ( - subtile_m_idx * fx.Index(SUBTILE_M) - + fx.Index(mi * MFMA_M) - - fx.Index(sm * (BLOCK_M // 2)) - ) - return load_transposed_frag_half( - lds_a[sm], - local_m_tile, - 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, - 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, -): - """Launch NT tensor-wise FP8 GEMM from both transpose backings. - - Contract: - a: [M, K] FP8 payload - a_scale_inv: one-element FP32 inverse quantization scale - b: [N, K] FP8 payload - b_scale_inv: one-element FP32 inverse quantization scale - c: [M, N] float16, bfloat16, or float32 output - - Both operands remain row-major [outer, K] in global memory. Each tile is - transposed during GMEM-to-LDS staging, then read with the validated - ds_read_b64_tr_b8 fragment contract. - """ - if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): - raise TypeError("FlyDSL FP8 NT 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 NT GEMM expects E4M3 or E5M2 payloads, " - f"got A={a.dtype} and B={b.dtype}" - ) - - m, k = 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") - - 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 optimized NT FP8 GEMM with A [M,K] and B [N,K].""" - 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 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 NT GEMM requires M to be a multiple of {_BLOCK_M}, got M={M_runtime}" - ) - if N_runtime % _BLOCK_N != 0: - raise FlyDSLUnsupportedError( - f"FlyDSL FP8 NT GEMM requires N to be a multiple of {_BLOCK_N}, got N={N_runtime}" - ) - if K_runtime % _BLOCK_K != 0: - raise FlyDSLUnsupportedError( - f"FlyDSL FP8 NT GEMM requires K to be a multiple of {_BLOCK_K}, got K={K_runtime}" - ) - num_k_tiles = K_runtime // _BLOCK_K - if num_k_tiles < 4: - raise FlyDSLUnsupportedError( - f"FlyDSL FP8 NT 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( - 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/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 4e68d01c2..544e7545f 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -17,8 +17,6 @@ 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 from .mxfp8_gemm_nn import mxfp8_matmul as mxfp8_matmul_nn from .mxfp8_gemm_nt import mxfp8_matmul as mxfp8_matmul_nt @@ -775,17 +773,19 @@ def _run_mxfp8( def _select_fp8_storage_for_layout(A, transa, B, transb): - """Select the exact existing TE FP8 backing required by each layout. + """Select existing TE FP8 backings and normalize to one core contract. - Fixed zero-copy routes selected for the final kernel contracts: + Tensor-wise FP8 columnwise storage is a materialized transpose, unlike + MXFP8 columnwise storage, which denotes a different quantization direction. - TN: wrapper swaps B._data/A._data -> [M,K], [N,K] - NN: wrapper swaps B._data/A._transpose -> [M,K], [N,K] - NT: wrapper swaps B._transpose/A._transpose -> [M,K], [N,K] + After selecting the required TE backing and swapping BLAS operand ownership, + every supported layout produces the same kernel-visible operands: - NT is the NN transpose-storage path applied to both operands. Both - transpose allocations stay contiguous in their native [outer,K] shapes; - no tensor transpose, reshape reinterpretation, or materialization occurs. + TN: B._data [M,K], A._data [N,K] + NN: B._data [M,K], A._transpose [N,K] + NT: B._transpose [M,K], A._transpose [N,K] + + No kernel-side transpose, transpose staging, or transpose-read is needed. """ layout = (bool(transa), bool(transb)) @@ -808,9 +808,8 @@ def _select_fp8_storage_for_layout(A, transa, B, transb): B_data = _flatten_rowwise(B_payload, B_storage) elif layout == (False, True): # NT / dW - # NT extends NN's transpose-storage handling to both operands. - # After ownership swap, B._transpose is kernel A [M,K] and - # A._transpose is kernel B [N,K]. + # Both selected transpose allocations are already materialized + # row-major [outer,K] backings for the normalized common core. A_payload = _get_fp8_columnwise_payload(A, "A") A_storage = "A._transpose" A_data = _flatten_columnwise(A_payload, A_storage) @@ -843,7 +842,7 @@ def _run_fp8( *, output_dtype: torch.dtype, ): - """Dispatch tensor-wise FP8 using exact per-kernel storage contracts.""" + """Normalize tensor-wise FP8 storage and invoke the common FP8 core.""" supported_fp8_dtypes = ( tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2, @@ -900,42 +899,17 @@ def _run_fp8( a_scale = B_scale_inv b_scale = A_scale_inv - if layout == "TN": - matmul = fp8_matmul - kernel_layout = "TN" - - a_flydsl = B_data - b_flydsl = A_data - - m, k = a_flydsl.shape - n, kb = b_flydsl.shape - - elif layout == "NN": - matmul = fp8_matmul_nn - kernel_layout = "NN" - - a_flydsl = B_data - b_flydsl = A_data - - m, k = a_flydsl.shape - n, kb = b_flydsl.shape - - elif layout == "NT": - matmul = fp8_matmul_nt - kernel_layout = "NT" - - # Correct fp8_gemm_nt contract: NN's [outer,K] transpose-storage - # path applied to both operands. - a_flydsl = B_data - b_flydsl = A_data + # Storage selection is layout-specific; execution is not. Tensor-wise FP8 + # transpose backing is already a materialized row-major transpose, so all + # supported layouts normalize to the common [M,K] x [N,K] core contract. + matmul = fp8_matmul + kernel_layout = "common" - m, k = a_flydsl.shape - n, kb = b_flydsl.shape + a_flydsl = B_data + b_flydsl = A_data - else: - raise FlyDSLUnsupportedError( - "FlyDSL GEMM does not support transa=True, transb=True (TT)" - ) + m, k = a_flydsl.shape + n, kb = b_flydsl.shape if not a_flydsl.is_contiguous() or not b_flydsl.is_contiguous(): raise FlyDSLUnsupportedError( @@ -972,12 +946,12 @@ def _run_fp8( shape=logical_output_shape, dtype=output_dtype, device=a_flydsl.device, - backend_name=f"FP8 {kernel_layout}", + backend_name=f"FP8 {layout} via {kernel_layout} core", ) _fp8_debug( f"dispatch entry: transa={bool(transa)}, transb={bool(transb)}, " - f"layout={layout}, selected_kernel={matmul.__module__}.{matmul.__name__}" + f"layout={layout}, normalized_core={matmul.__module__}.{matmul.__name__}" ) _fp8_debug(f"selected TE storage: A={A_storage}, B={B_storage}") _fp8_tensor_debug(f"selected/{A_storage}", A_data) From f1a5213e7660f0dad2cab3f799c9ba459b4cf06c Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 28 Jul 2026 20:12:35 +0000 Subject: [PATCH 23/31] unify mxfp8 gemm shape variants --- .../flydsl_kernels/gemm/gemm_wrappers.py | 17 +- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 710 +++++--- .../flydsl_kernels/gemm/mxfp8_gemm_nn.py | 1503 ----------------- .../flydsl_kernels/gemm/mxfp8_gemm_nt.py | 1474 ---------------- 4 files changed, 517 insertions(+), 3187 deletions(-) delete mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nn.py delete mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nt.py diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 544e7545f..2e8db0b4d 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -18,8 +18,6 @@ from .fp32_gemm import fp32_matmul from .fp8_gemm import fp8_matmul from .mxfp8_gemm import mxfp8_matmul -from .mxfp8_gemm_nn import mxfp8_matmul as mxfp8_matmul_nn -from .mxfp8_gemm_nt import mxfp8_matmul as mxfp8_matmul_nt def _product(shape): @@ -595,20 +593,20 @@ def _run_mxfp8( layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" dispatch = { - (True, False): ("TN", mxfp8_matmul), - (False, False): ("NN", mxfp8_matmul_nn), - (False, True): ("NT", mxfp8_matmul_nt), + (True, False): "TN", + (False, False): "NN", + (False, True): "NT", } try: - kernel_layout, matmul = dispatch[(bool(transa), bool(transb))] + kernel_layout = dispatch[(bool(transa), bool(transb))] except KeyError as exc: raise FlyDSLUnsupportedError( "FlyDSL GEMM does not support transa=True, transb=True (TT)" ) from exc _mxfp8_debug( - f"entry: layout={layout}, selected_kernel=" - f"{matmul.__module__}.{matmul.__name__}, " + f"entry: layout={layout}, common_kernel=" + f"{mxfp8_matmul.__module__}.{mxfp8_matmul.__name__}, " f"A_type={type(A).__name__}, B_type={type(B).__name__}, " f"D_provided={D is not None}" ) @@ -762,12 +760,13 @@ def _run_mxfp8( f"M={m}, N={n}, K={k}" ) - matmul( + mxfp8_matmul( a_flydsl, a_scale, b_flydsl, b_scale, D.view(m, n), + layout=kernel_layout, ) return D diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index a54dc43d7..4450b95b9 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -2,19 +2,23 @@ # # See LICENSE for license information. -"""FlyDSL MXFP8 GEMM implementation. +"""FlyDSL MXFP8 TN/NN/NT 4-wave GEMM implementation. -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``. +All supported MXFP8 layouts share one source-level kernel generator while +remaining separate compile-time specializations: -Canonical launch inputs: + TN: A [M,K] normal read, B [N,K] normal read + NN: A [M,K] normal read, B [K,N] transpose read + NT: A [K,M] transpose read, B [K,N] transpose read - a: [M, K] FP8 E4M3 or E5M2 payload - a_scale: [M, K/32] raw E8M0 bytes - b: [K, N] FP8 E4M3 or E5M2 payload - b_scale: [K/32, N] raw E8M0 bytes - D: [M, N] float16, bfloat16, or float32 output +The layout is a Python-only cache key. It is never passed as a runtime kernel +argument. Global addressing, LDS fragment reads, scheduler directives, and +scale-source orientation are selected while building each specialized kernel, +so generated TN/NN/NT kernels contain no runtime layout branches. + +All operand payloads use direct ``BufferCopyLDS128b`` global-to-LDS staging. +Transpose variants differ only in K-major global addressing and +``ds_read_b64_tr_b8`` LDS-to-register fragment reconstruction. """ import functools @@ -62,8 +66,19 @@ def _debug(message: str) -> None: print(f"[DEBUG_FLYDSL_MXFP8_GEMM] {message}") -def pack_mx32_scales_iter(scales_u8: torch.Tensor) -> torch.Tensor: - """Pack raw [Rows, K/32] E8M0 scales as [K/128, Rows] uint32.""" +def pack_mx32_scales_iter( + scales_u8: torch.Tensor, + *, + source_colwise: bool = False, +) -> torch.Tensor: + """Pack raw E8M0 scales as iteration-major ``[K/128, dim]`` uint32. + + ``source_colwise=False`` consumes TE rowwise scales ``[dim, K/32]``. + ``source_colwise=True`` consumes TE columnwise scales ``[K/32, dim]``. + + Both paths produce the same packed representation consumed by every + TN/NN/NT MXFP8 kernel specialization. + """ if scales_u8.dtype != torch.uint8: raise TypeError( f"MXFP8 scales must be torch.uint8 E8M0 bytes, got {scales_u8.dtype}" @@ -73,13 +88,27 @@ def pack_mx32_scales_iter(scales_u8: torch.Tensor) -> torch.Tensor: f"MXFP8 scales must be rank 2, got shape {tuple(scales_u8.shape)}" ) - rows, qk = scales_u8.shape + if source_colwise: + qk, dim = scales_u8.shape + if qk % 4 != 0: + raise ValueError( + f"Columnwise scale K dimension must be divisible by 4 K32 groups, got {qk}" + ) + s32 = scales_u8.contiguous().view(qk // 4, 4, dim).to(torch.int32) + return ( + s32[:, 0, :] + | (s32[:, 1, :] << 8) + | (s32[:, 2, :] << 16) + | (s32[:, 3, :] << 24) + ).contiguous() + + dim, qk = scales_u8.shape if qk % 4 != 0: raise ValueError( - f"Scale K dimension must be divisible by 4 K32 groups, got {qk}" + f"Rowwise scale K dimension must be divisible by 4 K32 groups, got {qk}" ) - s32 = scales_u8.contiguous().view(rows, qk // 4, 4).to(torch.int32) + s32 = scales_u8.contiguous().view(dim, qk // 4, 4).to(torch.int32) packed = ( s32[:, :, 0] | (s32[:, :, 1] << 8) @@ -89,18 +118,25 @@ def pack_mx32_scales_iter(scales_u8: torch.Tensor) -> torch.Tensor: return packed.transpose(0, 1).contiguous() -def pack_mx32_scales_for_hk(scales_u8: torch.Tensor) -> torch.Tensor: - """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] +def pack_mx32_scales_for_hk( + scales_u8: torch.Tensor, + *, + source_colwise: bool = False, +) -> torch.Tensor: + """Convert raw TE E8M0 scales to ``[K/128, dim]`` MFMA-ready words.""" + scale_iter = pack_mx32_scales_iter( + scales_u8, + source_colwise=source_colwise, + ) + dim = scales_u8.shape[1] if source_colwise else scales_u8.shape[0] - if rows % 64 != 0: + if dim % 64 != 0: raise ValueError( - f"Rows={rows} must be a multiple of 64 for HK MFMA scale packing" + f"Scale outer dimension={dim} must be a multiple of 64 for HK MFMA packing" ) device = scales_u8.device - row = torch.arange(rows, device=device, dtype=torch.int64) + row = torch.arange(dim, device=device, dtype=torch.int64) row_within_16 = row % 16 k_subgroup = (row // 16) % 4 tile = row // 64 @@ -110,7 +146,7 @@ def pack_mx32_scales_for_hk(scales_u8: torch.Tensor) -> torch.Tensor: 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) + source_value >> (k_subgroup * 8).view(1, dim) ) & 0xFF packed |= byte_value << (group * 8) @@ -205,12 +241,20 @@ def _compile_kernel( a_fp8_dtype: torch.dtype, b_fp8_dtype: torch.dtype, output_dtype: torch.dtype, + layout: str, ): - """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. + """Build one compile-time-specialized TN, NN, or NT kernel. - ``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. + ``layout`` is a Python string consumed while constructing the FlyDSL IR. + It is not a runtime kernel argument. Each cache entry therefore contains + only the addressing, LDS reads, and scheduler directives for that layout. """ + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") + + a_transpose_read = layout == "NT" + b_transpose_read = layout in ("NN", "NT") + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K fp8_input_types = { @@ -277,6 +321,153 @@ def _compile_kernel( LOAD_PASSES_HALF = LDS_ELEMS_HALF // (NUM_THREADS * VEC_BYTES) assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + # Resolve every layout-dependent choice before FlyDSL captures kernel_gemm. + # These are ordinary Python callables/constants, so each cached layout emits + # only its selected addressing, fragment-read, and scheduler path. + Q0_SCHED_DSRD = 2 if a_transpose_read else 1 + PREFETCH_SCHED_DSRD = 4 if a_transpose_read else 2 + + if a_transpose_read: + def _a_leading_dim(c_m): + return c_m + + def _a_global_base(k_base, subtile, c_m, bx_m_idx): + return ( + k_base * fx.Index(c_m) + + bx_m_idx + + fx.Index(subtile * (BLOCK_M // 2)) + ) + + def _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ): + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + local_m_tile = ( + subtile_m_idx * fx.Index(SUBTILE_M) + + fx.Index(mi * MFMA_M) + - fx.Index(sm * (BLOCK_M // 2)) + ) + return load_transposed_frag_half( + lds_a[sm], + local_m_tile, + half, + ) + else: + def _a_leading_dim(c_m): + del c_m + return K + + def _a_global_base(k_base, subtile, c_m, bx_m_idx): + del c_m + return ( + (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) + * fx.Index(K) + + k_base + ) + + def _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ): + del load_transposed_frag_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, + ) + + if b_transpose_read: + def _b_leading_dim(c_n): + return c_n + + def _b_global_base(k_base, subtile, c_n, by_n_idx): + return ( + k_base * fx.Index(c_n) + + by_n_idx + + fx.Index(subtile * (BLOCK_N // 2)) + ) + + def _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ): + del load_normal_b_frag, lane_mod_16 + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + local_n_tile = ( + subtile_n_idx * fx.Index(SUBTILE_N) + + fx.Index(ni * MFMA_N) + - fx.Index(sn * (BLOCK_N // 2)) + ) + return load_transposed_frag(lds_b[sn], local_n_tile) + else: + def _b_leading_dim(c_n): + del c_n + return K + + def _b_global_base(k_base, subtile, c_n, by_n_idx): + del c_n + return ( + (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) + * fx.Index(K) + + k_base + ) + + def _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ): + del load_transposed_frag + 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_normal_b_frag(lds_b, b_row_addr, sn) + + if (not a_transpose_read) or (not b_transpose_read): + def _normal_read_columns(lane_div_16, lane_mod_16): + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + _, col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, col1 = swizzle_128(lane_mod_16, reg_k_col1) + return col0, col1 + else: + def _normal_read_columns(lane_div_16, lane_mod_16): + del lane_div_16, lane_mod_16 + return fx.Int32(0), fx.Int32(0) + @fx.struct class SharedStorage: # Each logical 256x128 page is two independent 128x128 half-pages. @@ -319,25 +510,49 @@ def kernel_gemm( 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 + # 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, a_f8_ir_t, wave_id) - b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) + # Compile-time global leading dimensions: + # normal source [X,K] -> leading dimension K + # transpose source [K,X] -> leading dimension X + a_leading_dim = _a_leading_dim(c_m) + b_leading_dim = _b_leading_dim(c_n) + + gl_off_a = compute_global_swizzle( + lane, + wave_id, + a_leading_dim, + LOAD_PASSES_HALF, + preshuffled=False, + ) + gl_off_b = compute_global_swizzle( + lane, + wave_id, + b_leading_dim, + 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)) @@ -437,26 +652,21 @@ def hot_loop_scheduler_q_refill_2n(): 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. + # TN/NN: one normal A-bottom LDS read per chunk. + # NT: one transpose-read A half plus the matching transpose-read + # scheduling pressure retained from the passing NT specialization. for _ in range_constexpr(8): rocdl.sched_vmem(1) - rocdl.sched_dsrd(1) + rocdl.sched_dsrd(Q0_SCHED_DSRD) 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 + # TN/NN retain two scheduled DS reads per chunk. NT retains four + # because both carried operands use two DS_READ_TR instructions. for _ in range_constexpr(8): - rocdl.sched_dsrd(2) + rocdl.sched_dsrd(PREFETCH_SCHED_DSRD) rocdl.sched_mfma(4) rocdl.sched_barrier(0) @@ -502,14 +712,18 @@ def load_scale_tile(k128): ) def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): - # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one - # 128x128 half-page (16 KiB). Each half has its own LDS base. - global_base = (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) + k_base - a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) + a_g2s.load_one( + lds_a[subtile], + fx.Int32(_a_global_base(k_base, subtile, c_m, bx_m_idx)), + 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) + b_g2s.load_one( + lds_b[subtile], + fx.Int32(_b_global_base(k_base, subtile, c_n, by_n_idx)), + pass_in_subtile, + ) def stage_a_subtile(k_base, subtile, lds_a): for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): @@ -535,10 +749,43 @@ def load_frag_at_byte_base(lds_page, row_byte_base): 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. + def load_normal_b_frag(lds_b, local_row, half): + # Physical [N,K] page, ordinary TN-style fixed-row read. 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)) + return load_frag_at_byte_base( + lds_b[half], + half_row * fx.Index(BLOCK_K), + ) + + def load_transposed_frag_half(lds_page, local_x_tile, half): + # Exact inverse mapping validated by the MXFP8 NN fragment probe. + lane_div16_i32 = fx.Int32(lane_div_16) + lane_in16_i32 = fx.Int32(lane_mod_16) + source_k = ( + lane_div16_i32 * fx.Int32(16) + + lane_in16_i32 // fx.Int32(2) + ) + source_x = ( + fx.Int32(local_x_tile) + + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) + ) + + physical_k, physical_x = swizzle_128(source_k, source_x) + base = physical_k * fx.Int32(128) + physical_x + other = base ^ fx.Int32(0x440) + immediate_offset = 0 if half == 0 else 0x2000 + + return s2r.load_one_transpose( + lds_page, + base, + other, + immediate_offset=immediate_offset, + ) + + def load_transposed_frag(lds_page, local_x_tile): + x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) + x1 = load_transposed_frag_half(lds_page, local_x_tile, 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 @@ -640,13 +887,10 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): # 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_lds_k_col0, reg_lds_k_col1 = _normal_read_columns( + lane_div_16, + lane_mod_16, + ) reg_subtile_m_idx0 = wave_id // 2 reg_subtile_n_idx0 = wave_id % 2 @@ -655,15 +899,19 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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 + b_ni = _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ) + return b_ni, b_scales[ni] def load_b_subtile_regs(lds_b, scale_tile, sn): b0, bs0 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 0) @@ -673,12 +921,18 @@ def load_b_subtile_regs(lds_b, scale_tile, sn): 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) + + return _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ) 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. @@ -1172,22 +1426,6 @@ def launch_gemm( return 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, -): - return _compile_kernel( - K, - a_fp8_dtype, - b_fp8_dtype, - output_dtype, - ) - - - def do_gemm( A: torch.Tensor, As: torch.Tensor, @@ -1195,78 +1433,85 @@ def do_gemm( Bs: torch.Tensor, C: torch.Tensor, stream=None, + *, + layout: str = "TN", ): - """Launch the K-specialized kernel with runtime M/N. + """Launch one cached compile-time MXFP8 layout specialization.""" + if layout == "TN": + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + elif layout == "NN": + M_runtime, K_runtime = A.shape + Kb_runtime, N_runtime = B.shape + elif layout == "NT": + K_runtime, M_runtime = A.shape + Kb_runtime, N_runtime = B.shape + else: + raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") - 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 - 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}" + 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}" + 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}" + f"FlyDSL MXFP8 {layout} GEMM requires M to be a multiple of " + f"{_BLOCK_M}, 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}" + f"FlyDSL MXFP8 {layout} GEMM requires N to be a multiple of " + f"{_BLOCK_N}, 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}" + f"FlyDSL MXFP8 {layout} GEMM requires K to be a multiple of " + f"{_BLOCK_K}, got K={K_runtime}" ) num_k_tiles = K_runtime // _BLOCK_K 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)" + f"FlyDSL MXFP8 {layout} GEMM requires at least 4 K{_BLOCK_K} " + f"tiles, 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" 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})" - ) - 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 tuple(As.shape) == expected_as, f"As shape {tuple(As.shape)} != {expected_as}" + assert tuple(Bs.shape) == expected_bs, f"Bs shape {tuple(Bs.shape)} != {expected_bs}" + assert tuple(C.shape) == (M_runtime, N_runtime), ( + f"C shape {tuple(C.shape)} != {(M_runtime, N_runtime)}" ) + if C.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) + + tensors = (A, As, B, Bs, C) + if any(t.device != A.device for t in tensors[1:]): + raise ValueError("A, B, packed scales, and C must be on the same device") + 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. + + # Preserve the exact flat descriptor contract used by the passing kernels. 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), + _cached_launch( + K_runtime, A.dtype, B.dtype, C.dtype, - ) - launch( + layout, + )( A_arg, As_arg, B_arg, @@ -1278,14 +1523,50 @@ def do_gemm( ) -__all__ = [ - "BLOCK_M", - "BLOCK_N", - "BLOCK_K", - "do_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, + layout: str, +): + """Cache independent TN/NN/NT binaries with no runtime layout argument.""" + return _compile_kernel( + K, + a_fp8_dtype, + b_fp8_dtype, + output_dtype, + layout, + ) +def _validate_common_payloads( + a: torch.Tensor, + b: torch.Tensor, + D: torch.Tensor, + *, + layout: str, +): + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL MXFP8 {layout} expects rank-2 operands, got " + f"a={tuple(a.shape)} 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( + f"FlyDSL MXFP8 {layout} expects E4M3 or E5M2 payloads " + f"independently, got a={a.dtype} and b={b.dtype}" + ) + if a.device != b.device or D.device != a.device: + raise ValueError("A, B, and D must be on the same device") + if D.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + "FlyDSL MXFP8 output must be float16, bfloat16, or float32, " + f"got {D.dtype}" + ) + def mxfp8_matmul( a: torch.Tensor, @@ -1294,113 +1575,140 @@ def mxfp8_matmul( b_scale: torch.Tensor, D: torch.Tensor, stream=None, + *, + layout: str = "TN", ): - """Launch the fused MXFP8 kernel from canonical row-major operands. + """Normalize scale orientation and launch a compile-time layout binary. + + Wrapper-visible contracts: + + TN: a [M,K], b [K,N], scales [M,K/32] and [K/32,N] + NN: a [M,K], b [K,N], scales [M,K/32] and [K/32,N] + NT: a [K,M], b [K,N], scales [K/32,M] and [K/32,N] - 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 output-dtype-specialized 4-wave implementation. + TN preserves the existing adapter conversion to the kernel's normal-read + B [N,K] representation. NN and NT preserve K-major payloads and use + ``ds_read_b64_tr_b8`` inside their compile-time-specialized kernels. """ - 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)}" - ) + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") - 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)}" - ) + _validate_common_payloads(a, b, D, layout=layout) - 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 layout in ("TN", "NN"): + m, k = a.shape + kb, n = b.shape + else: + k, m = a.shape + kb, n = b.shape - if a.device != b.device: + if kb != k: raise ValueError( - f"a and b must be on the same device, got {a.device} and {b.device}" + f"Incompatible MXFP8 {layout} operands: " + f"A{tuple(a.shape)} and B{tuple(b.shape)}" ) - 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 not in (torch.float16, torch.bfloat16, torch.float32): - raise TypeError( - "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") - + raise ValueError(f"D shape {tuple(D.shape)} != expected {(m, n)}") 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) + if layout == "NT": + expected_a_scale = (k // SCALE_GROUP_SIZE, m) + else: + 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}" + f"{expected_a_scale} for {layout}" ) if tuple(b_scale.shape) != expected_b_scale: raise ValueError( f"b_scale shape {tuple(b_scale.shape)} != expected " - f"{expected_b_scale}" + f"{expected_b_scale} for {layout}" ) 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) + if a_scale.device != a.device or b_scale.device != a.device: + raise ValueError("A, B, scales, and D must be on the same device") + + if layout == "TN": + # Preserve the passing TN kernel contract exactly: normal-read B [N,K]. + a_kernel = a + b_kernel = b.transpose(0, 1).contiguous() + a_scale_hk = pack_mx32_scales_for_hk( + a_scale, + source_colwise=False, + ) + b_scale_hk = pack_mx32_scales_for_hk( + b_scale.transpose(0, 1).contiguous(), + source_colwise=False, + ) + elif layout == "NN": + a_kernel = a + b_kernel = b + a_scale_hk = pack_mx32_scales_for_hk( + a_scale, + source_colwise=False, + ) + b_scale_hk = pack_mx32_scales_for_hk( + b_scale, + source_colwise=True, + ) + else: + a_kernel = a + b_kernel = b + a_scale_hk = pack_mx32_scales_for_hk( + a_scale, + source_colwise=True, + ) + b_scale_hk = pack_mx32_scales_for_hk( + b_scale, + source_colwise=True, + ) _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"{layout} kernel inputs: a={tuple(a_kernel.shape)}, " + f"b={tuple(b_kernel.shape)}, " f"a_scale_hk={tuple(a_scale_hk.shape)}, " - f"b_scale_hk={tuple(b_scale_hk.shape)}, D={tuple(D.shape)}, " - f"D_dtype={D.dtype}" + f"b_scale_hk={tuple(b_scale_hk.shape)}, D={tuple(D.shape)}" ) - _debug("launching fused MXFP8 4-wave kernel") do_gemm( - a, + a_kernel, a_scale_hk, - b_hk, + b_kernel, b_scale_hk, D.view(m, n), + layout=layout, stream=stream, ) - - _debug("launch complete") return D +def mxfp8_matmul_nn(*args, **kwargs): + """Compatibility entry point for the common NN specialization.""" + kwargs["layout"] = "NN" + return mxfp8_matmul(*args, **kwargs) + + +def mxfp8_matmul_nt(*args, **kwargs): + """Compatibility entry point for the common NT specialization.""" + kwargs["layout"] = "NT" + return mxfp8_matmul(*args, **kwargs) + + __all__ = [ "BLOCK_M", "BLOCK_N", "BLOCK_K", "SCALE_GROUP_SIZE", + "do_gemm", "mxfp8_matmul", + "mxfp8_matmul_nn", + "mxfp8_matmul_nt", ] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nn.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nn.py deleted file mode 100644 index b7f20ba6b..000000000 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nn.py +++ /dev/null @@ -1,1503 +0,0 @@ -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. - -"""FlyDSL MXFP8 NN 4-wave GEMM implementation. - -This specialization preserves the validated MXFP8 TN compute, scale, MFMA, -accumulator, and epilogue pipelines. A is physically row-major [M, K]. -B is physically row-major [N, K], staged as XOR-swizzled [K128, N128] LDS, -and reconstructed with the validated four-read ds_read_b64_tr_b8 path. - -Raw scales enter as A rowwise [M, K/32] and B columnwise [K/32, N]. -Orientation-aware prepacking converts both to the common iteration-major -[K/128, dim] uint32 representation consumed by the kernel.""" - -import functools -import os - -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 .exceptions import FlyDSLUnsupportedError -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 _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, - *, - source_colwise: bool = False, -) -> torch.Tensor: - """Pack raw E8M0 scales as iteration-major ``[K/128, dim]`` uint32. - - ``source_colwise=False`` consumes TE rowwise scales ``[dim, K/32]``. - ``source_colwise=True`` consumes TE columnwise scales ``[K/32, dim]``. - - Both paths produce the same packed representation consumed by every - TN/NN/NT MXFP8 kernel specialization. - """ - 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)}" - ) - - if source_colwise: - qk, dim = scales_u8.shape - if qk % 4 != 0: - raise ValueError( - f"Columnwise scale K dimension must be divisible by 4 K32 groups, got {qk}" - ) - s32 = scales_u8.contiguous().view(qk // 4, 4, dim).to(torch.int32) - return ( - s32[:, 0, :] - | (s32[:, 1, :] << 8) - | (s32[:, 2, :] << 16) - | (s32[:, 3, :] << 24) - ).contiguous() - - dim, qk = scales_u8.shape - if qk % 4 != 0: - raise ValueError( - f"Rowwise scale K dimension must be divisible by 4 K32 groups, got {qk}" - ) - - s32 = scales_u8.contiguous().view(dim, 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, - *, - source_colwise: bool = False, -) -> torch.Tensor: - """Convert raw TE E8M0 scales to ``[K/128, dim]`` MFMA-ready words.""" - scale_iter = pack_mx32_scales_iter( - scales_u8, - source_colwise=source_colwise, - ) - dim = scales_u8.shape[1] if source_colwise else scales_u8.shape[0] - - if dim % 64 != 0: - raise ValueError( - f"Scale outer dimension={dim} must be a multiple of 64 for HK MFMA packing" - ) - - device = scales_u8.device - row = torch.arange(dim, device=device, dtype=torch.int64) - row_within_16 = row % 16 - k_subgroup = (row // 16) % 4 - tile = row // 64 - - packed = torch.zeros_like(scale_iter) - 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, dim) - ) & 0xFF - packed |= byte_value << (group * 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, - 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 - 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 - - 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[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, 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) - - 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) - 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) - - tx_i32 = fx.Int32(tx) - wave_id = tx_i32 // fx.Int32(WARP_SIZE) - lane = tx_i32 % fx.Int32(WARP_SIZE) - - # A remains ordinary row-major [M, K]. - gl_off_a = 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 is the selected MXFP8 columnwise payload, physically [K, N]. - # Load the K-major source directly into the XOR-swizzled physical LDS - # image [K128, N128] consumed by ds_read_b64_tr_b8. - gl_off_b = compute_global_swizzle( - lane, - wave_id, - c_n, - LOAD_PASSES_HALF, - preshuffled=False, - ) - 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 _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): - # B is physically [K, N]. Copy - # B[k_base:k_base+128, by_n+subtile*128:...] - # into one XOR-swizzled physical LDS half-page [K128, N128]. - global_base = ( - k_base * fx.Index(c_n) - + by_n_idx - + fx.Index(subtile * (BLOCK_N // 2)) - ) - 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_transpose(lds_page, local_n_tile): - # Exact inverse mapping validated against the ordinary B[N, K] - # production MFMA fragment: - # - # source_k = lane_div_16*16 + lane_in_16//2 - # source_n = local_n_tile + (lane_in_16&1)*8 - # - # base^0x440 advances logical K by 8 under the 128-byte XOR - # swizzle. The DS immediate 0x2000 advances logical K by 64. - lane_div16_i32 = fx.Int32(lane_div_16) - lane_in16_i32 = fx.Int32(lane_mod_16) - source_k = ( - lane_div16_i32 * fx.Int32(16) - + lane_in16_i32 // fx.Int32(2) - ) - source_n = ( - fx.Int32(local_n_tile) - + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) - ) - - physical_k, physical_n = swizzle_128(source_k, source_n) - base = physical_k * fx.Int32(BLOCK_N // 2) + physical_n - other = base ^ fx.Int32(0x440) - - x0 = s2r.load_one_transpose( - lds_page, - base, - other, - immediate_offset=0, - ) - x1 = s2r.load_one_transpose( - lds_page, - base, - other, - immediate_offset=0x2000, - ) - 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, 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"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, - ) - - 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"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, - ) - - 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 - 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. - # 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): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_scales = scale_tile[2] if sn == 0 else scale_tile[3] - local_n_tile = ( - subtile_n_idx * fx.Index(SUBTILE_N) - + fx.Index(ni * MFMA_N) - - fx.Index(sn * (BLOCK_N // 2)) - ) - b_ni = load_b_frag_transpose(lds_b[sn], local_n_tile) - return b_ni, b_scales[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, - 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, - ) - - - -def do_gemm( - 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 [K, N]. 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 - 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 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}" - if M_runtime % _BLOCK_M != 0: - raise FlyDSLUnsupportedError( - f"FlyDSL MXFP8 NN 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 NN 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 NN 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 MXFP8 NN 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" - 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})" - ) - 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 - # 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), - A.dtype, - B.dtype, - C.dtype, - ) - launch( - A_arg, - As_arg, - B_arg, - Bs_arg, - C_arg, - M_runtime, - 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 MXFP8 NN GEMM with one transpose-read operand. - - Contract: - a: [M, K] row-major FP8 payload - a_scale: [M, K/32] raw rowwise E8M0 scales - b: [K, N] row-major columnwise-quantized FP8 payload - b_scale: [K/32, N] raw columnwise E8M0 scales - D: [M, N] float16, bfloat16, or float32 output - - The B payload remains physically [K, N]. The kernel stages that K-major - source into the XOR-swizzled LDS image and uses ds_read_b64_tr_b8 to - reconstruct the MFMA B fragment. Scale - prepacking resolves the source orientation before launch, so both packed - scale tensors use the common [K/128, dim] kernel representation. - """ - if a.ndim != 2 or b.ndim != 2: - raise ValueError( - f"FlyDSL MXFP8 NN expects rank-2 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 MXFP8 NN operands: " - f"A{tuple(a.shape)} 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 MXFP8 NN 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}" - ) - 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 not in (torch.float16, torch.bfloat16, torch.float32): - raise TypeError( - "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") - - if k % SCALE_GROUP_SIZE != 0: - raise ValueError( - f"K={k} must be divisible by MXFP8 scale group size " - f"{SCALE_GROUP_SIZE}" - ) - - 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 {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 as torch.uint8") - if a_scale.device != a.device or b_scale.device != a.device: - raise ValueError("A, B, scales, and D must be on the same device") - - a_scale_hk = pack_mx32_scales_for_hk( - a_scale, - source_colwise=False, - ) - b_scale_hk = pack_mx32_scales_for_hk( - b_scale, - source_colwise=True, - ) - - _debug( - f"NN kernel inputs: a={tuple(a.shape)}, b={tuple(b.shape)}, " - f"a_scale_hk={tuple(a_scale_hk.shape)}, " - f"b_scale_hk={tuple(b_scale_hk.shape)}, D={tuple(D.shape)}" - ) - - do_gemm( - a, - a_scale_hk, - b, - b_scale_hk, - D.view(m, n), - stream=stream, - ) - return D - - -__all__ = [ - "BLOCK_M", - "BLOCK_N", - "BLOCK_K", - "SCALE_GROUP_SIZE", - "mxfp8_matmul", -] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nt.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nt.py deleted file mode 100644 index 7139edf5b..000000000 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm_nt.py +++ /dev/null @@ -1,1474 +0,0 @@ -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. - -"""FlyDSL MXFP8 NT 4-wave GEMM implementation. - -This specialization preserves the validated MXFP8 TN compute, scale, MFMA, -accumulator, and epilogue pipelines while applying the validated -ds_read_b64_tr_b8 path to both operands. A is physically [K, M] and B is -physically [K, N]. Each source tile is staged as XOR-swizzled [K128, X128] LDS. - -Both raw scale tensors are columnwise, [K/32, M] and [K/32, N]. -Orientation-aware prepacking converts them to the common iteration-major -[K/128, dim] uint32 representation consumed by the kernel.""" - -import functools -import os - -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 .exceptions import FlyDSLUnsupportedError -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 _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, - *, - source_colwise: bool = False, -) -> torch.Tensor: - """Pack raw E8M0 scales as iteration-major ``[K/128, dim]`` uint32. - - ``source_colwise=False`` consumes TE rowwise scales ``[dim, K/32]``. - ``source_colwise=True`` consumes TE columnwise scales ``[K/32, dim]``. - - Both paths produce the same packed representation consumed by every - TN/NN/NT MXFP8 kernel specialization. - """ - 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)}" - ) - - if source_colwise: - qk, dim = scales_u8.shape - if qk % 4 != 0: - raise ValueError( - f"Columnwise scale K dimension must be divisible by 4 K32 groups, got {qk}" - ) - s32 = scales_u8.contiguous().view(qk // 4, 4, dim).to(torch.int32) - return ( - s32[:, 0, :] - | (s32[:, 1, :] << 8) - | (s32[:, 2, :] << 16) - | (s32[:, 3, :] << 24) - ).contiguous() - - dim, qk = scales_u8.shape - if qk % 4 != 0: - raise ValueError( - f"Rowwise scale K dimension must be divisible by 4 K32 groups, got {qk}" - ) - - s32 = scales_u8.contiguous().view(dim, 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, - *, - source_colwise: bool = False, -) -> torch.Tensor: - """Convert raw TE E8M0 scales to ``[K/128, dim]`` MFMA-ready words.""" - scale_iter = pack_mx32_scales_iter( - scales_u8, - source_colwise=source_colwise, - ) - dim = scales_u8.shape[1] if source_colwise else scales_u8.shape[0] - - if dim % 64 != 0: - raise ValueError( - f"Scale outer dimension={dim} must be a multiple of 64 for HK MFMA packing" - ) - - device = scales_u8.device - row = torch.arange(dim, device=device, dtype=torch.int64) - row_within_16 = row % 16 - k_subgroup = (row // 16) % 4 - tile = row // 64 - - packed = torch.zeros_like(scale_iter) - 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, dim) - ) & 0xFF - packed |= byte_value << (group * 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, - 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 - 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 - - 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[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, 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) - - 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) - 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. - bx_m_idx = fx.Index(bx_m) - by_n_idx = fx.Index(by_n) - - tx_i32 = fx.Int32(tx) - wave_id = tx_i32 // fx.Int32(WARP_SIZE) - lane = tx_i32 % fx.Int32(WARP_SIZE) - - # NT storage is K-major for both operands: - # A [K, M] - # B [K, N] - # - # Read each K-by-X source tile in XOR-swizzled coordinate order and - # write it linearly to LDS. swizzle_128 is self-inverse, producing the - # physical [K128, X128] image consumed by ds_read_b64_tr_b8. - gl_off_a = compute_global_swizzle( - lane, - wave_id, - c_m, - LOAD_PASSES_HALF, - preshuffled=False, - ) - gl_off_b = compute_global_swizzle( - lane, - wave_id, - c_n, - 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 _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(): - # A-bottom and the B slices are transpose reads in NT. - 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(): - # Each prefetched A/B fragment uses two DS_READ_TR instructions. - for _ in range_constexpr(8): - rocdl.sched_dsrd(4) - 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): - # A is physically [K, M]. Copy - # A[k_base:k_base+128, bx_m+subtile*128:...] - # into one XOR-swizzled physical LDS half-page [K128, M128]. - global_base = ( - k_base * fx.Index(c_m) - + bx_m_idx - + fx.Index(subtile * (BLOCK_M // 2)) - ) - 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): - # B is physically [K, N]. Copy - # B[k_base:k_base+128, by_n+subtile*128:...] - # into one XOR-swizzled physical LDS half-page [K128, N128]. - global_base = ( - k_base * fx.Index(c_n) - + by_n_idx - + fx.Index(subtile * (BLOCK_N // 2)) - ) - 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 pack_frag_halves(x0, x1): - return pack_i32x4_i32x8(x0, x1) - - - def load_transposed_frag_half(lds_page, local_x_tile, half): - """Load one K64 portion of a fixed-X MFMA fragment. - - This is the inverse mapping validated against the working ordinary - LDS fragment: - - source_k = lane_div_16*16 + lane_in_16//2 - source_x = local_x_tile + (lane_in_16&1)*8 - - ``base ^ 0x440`` advances logical K by 8 under swizzle_128. - The 0x2000 DS immediate advances logical K by 64. - """ - lane_div16_i32 = fx.Int32(lane_div_16) - lane_in16_i32 = fx.Int32(lane_mod_16) - source_k = ( - lane_div16_i32 * fx.Int32(16) - + lane_in16_i32 // fx.Int32(2) - ) - source_x = ( - fx.Int32(local_x_tile) - + (lane_in16_i32 % fx.Int32(2)) * fx.Int32(8) - ) - - physical_k, physical_x = swizzle_128(source_k, source_x) - base = physical_k * fx.Int32(128) + physical_x - other = base ^ fx.Int32(0x440) - immediate_offset = 0 if half == 0 else 0x2000 - - return s2r.load_one_transpose( - lds_page, - base, - other, - immediate_offset=immediate_offset, - ) - - - def load_transposed_frag(lds_page, local_x_tile): - x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) - x1 = load_transposed_frag_half(lds_page, local_x_tile, 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, 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"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, - ) - - 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"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, - ) - - 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 - 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. - # 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_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): - subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) - b_scales = scale_tile[2] if sn == 0 else scale_tile[3] - local_n_tile = ( - subtile_n_idx * fx.Index(SUBTILE_N) - + fx.Index(ni * MFMA_N) - - fx.Index(sn * (BLOCK_N // 2)) - ) - b_ni = load_transposed_frag(lds_b[sn], local_n_tile) - return b_ni, b_scales[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): - subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) - local_m_tile = ( - subtile_m_idx * fx.Index(SUBTILE_M) - + fx.Index(mi * MFMA_M) - - fx.Index(sm * (BLOCK_M // 2)) - ) - return load_transposed_frag_half( - lds_a[sm], - local_m_tile, - 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, - 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, - ) - - - -def do_gemm( - A: torch.Tensor, - As: torch.Tensor, - B: torch.Tensor, - Bs: torch.Tensor, - C: torch.Tensor, - stream=None, -): - """Launch MXFP8 NT core from K-major A [K,M] and B [K,N].""" - K_runtime, M_runtime = A.shape - Kb_runtime, N_runtime = B.shape - assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" - 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}" - if M_runtime % _BLOCK_M != 0: - raise FlyDSLUnsupportedError( - f"FlyDSL MXFP8 NT 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 NT 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 NT 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 MXFP8 NT 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" - 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})" - ) - 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}" - ) - - tensors = (A, As, B, Bs, C) - if any(t.device != A.device for t in tensors[1:]): - raise ValueError("A, B, packed scales, and C must be on the same device") - - 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) - As_arg = As.contiguous().view(-1) - Bs_arg = Bs.contiguous().view(-1) - C_arg = C.contiguous().view(-1) - - launch = _cached_launch( - int(K_runtime), - A.dtype, - B.dtype, - C.dtype, - ) - launch( - A_arg, - As_arg, - B_arg, - Bs_arg, - C_arg, - M_runtime, - 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 MXFP8 NT GEMM with transpose-read A and B operands. - - Contract: - a: [K, M] row-major FP8 payload - a_scale: [K/32, M] raw columnwise E8M0 scales - b: [K, N] row-major FP8 payload - b_scale: [K/32, N] raw columnwise E8M0 scales - D: [M, N] float16, bfloat16, or float32 output - - Both operands remain K-major. Each is staged as an XOR-swizzled - [K128, X128] LDS image and reconstructed with ds_read_b64_tr_b8. - """ - if a.ndim != 2 or b.ndim != 2: - raise ValueError( - f"FlyDSL MXFP8 NT expects rank-2 operands, got " - f"a={tuple(a.shape)} and b={tuple(b.shape)}" - ) - - k, m = a.shape - kb, n = b.shape - if kb != k: - raise ValueError( - f"Incompatible MXFP8 NT operands: " - f"A{tuple(a.shape)} 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 MXFP8 NT 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}" - ) - 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 not in (torch.float16, torch.bfloat16, torch.float32): - raise TypeError( - "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") - - if k % SCALE_GROUP_SIZE != 0: - raise ValueError( - f"K={k} must be divisible by MXFP8 scale group size " - f"{SCALE_GROUP_SIZE}" - ) - - expected_a_scale = (k // SCALE_GROUP_SIZE, m) - 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 {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 as torch.uint8") - if a_scale.device != a.device or b_scale.device != a.device: - raise ValueError("A, B, scales, and D must be on the same device") - - a_scale_hk = pack_mx32_scales_for_hk( - a_scale, - source_colwise=True, - ) - b_scale_hk = pack_mx32_scales_for_hk( - b_scale, - source_colwise=True, - ) - - _debug( - f"NT kernel inputs: a={tuple(a.shape)}, b={tuple(b.shape)}, " - f"a_scale_hk={tuple(a_scale_hk.shape)}, " - f"b_scale_hk={tuple(b_scale_hk.shape)}, D={tuple(D.shape)}" - ) - - do_gemm( - a, - a_scale_hk, - b, - b_scale_hk, - D.view(m, n), - stream=stream, - ) - return D - - -__all__ = [ - "BLOCK_M", - "BLOCK_N", - "BLOCK_K", - "SCALE_GROUP_SIZE", - "mxfp8_matmul", -] From d36c6c2dd54e8c5ad0d216ecbd1edabce169401d Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 28 Jul 2026 22:34:23 +0000 Subject: [PATCH 24/31] FlyDSL MXFP8: fuse scale prepacking into dedicated GPU kernels Replace the eager PyTorch MXFP8 scale-packing path with stride-aware FlyDSL kernels that directly convert TE E8M0 scales into the HK/MFMA-ready [K/128, dim] packed layout. The previous implementation composed packing from arange, indexing, casts, shifts, masks, ORs, transposes, and contiguous copies. PyTorch lowered these into dozens of small GPU kernels around every GEMM, which dominated end-to-end runtime despite the FlyDSL GEMMs themselves being faster. The new path: - launches one fused scale-pack kernel per GEMM operand - supports both rowwise and columnwise TE scale layouts - consumes non-contiguous scale views using their actual strides - eliminates the intermediate iteration-major scale tensor - removes eager transpose/contiguous preparation from the scale path - preserves the existing HK/MFMA-ready packed representation --- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 236 +++++++++++++----- 1 file changed, 173 insertions(+), 63 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 4450b95b9..7d34495c9 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -66,92 +66,193 @@ def _debug(message: str) -> None: print(f"[DEBUG_FLYDSL_MXFP8_GEMM] {message}") -def pack_mx32_scales_iter( - scales_u8: torch.Tensor, - *, - source_colwise: bool = False, -) -> torch.Tensor: - """Pack raw E8M0 scales as iteration-major ``[K/128, dim]`` uint32. +_SCALE_PACK_THREADS = 256 - ``source_colwise=False`` consumes TE rowwise scales ``[dim, K/32]``. - ``source_colwise=True`` consumes TE columnwise scales ``[K/32, dim]``. - Both paths produce the same packed representation consumed by every - TN/NN/NT MXFP8 kernel specialization. +def _compile_mx32_scale_pack_kernel( + dim: int, + qk: int, + source_colwise: bool, + stride0: int, + stride1: int, +): + """Build one fused raw-E8M0 -> HK-scale packing kernel. + + One GPU thread produces one final ``uint32`` word in the GEMM-consumed + ``[K/128, dim]`` layout. There is no intermediate ``scale_iter`` tensor + and no eager PyTorch shift/index/OR kernels. """ - if scales_u8.dtype != torch.uint8: - raise TypeError( - f"MXFP8 scales must be torch.uint8 E8M0 bytes, got {scales_u8.dtype}" + if dim % 64 != 0: + raise ValueError( + f"Scale outer dimension={dim} must be a multiple of 64" ) - if scales_u8.ndim != 2: + if qk % 4 != 0: + raise ValueError( + f"Scale K/32 dimension={qk} must be divisible by 4" + ) + + k128_tiles = qk // 4 + total_words = k128_tiles * dim + if total_words % _SCALE_PACK_THREADS != 0: raise ValueError( - f"MXFP8 scales must be rank 2, got shape {tuple(scales_u8.shape)}" + f"Packed scale words={total_words} must be divisible by " + f"{_SCALE_PACK_THREADS}" ) + # Select source addressing before FlyDSL captures the kernel. The emitted + # rowwise and columnwise binaries contain no runtime orientation branch. if source_colwise: - qk, dim = scales_u8.shape - if qk % 4 != 0: - raise ValueError( - f"Columnwise scale K dimension must be divisible by 4 K32 groups, got {qk}" + def _source_offset(source_k32, source_row): + # Logical source is [K/32, dim], but the underlying TE tensor may + # be a non-contiguous view. Strides are in uint8 elements. + return ( + source_k32 * fx.Index(stride0) + + source_row * fx.Index(stride1) ) - s32 = scales_u8.contiguous().view(qk // 4, 4, dim).to(torch.int32) - return ( - s32[:, 0, :] - | (s32[:, 1, :] << 8) - | (s32[:, 2, :] << 16) - | (s32[:, 3, :] << 24) - ).contiguous() - - dim, qk = scales_u8.shape - if qk % 4 != 0: - raise ValueError( - f"Rowwise scale K dimension must be divisible by 4 K32 groups, got {qk}" + else: + def _source_offset(source_k32, source_row): + # Logical source is [dim, K/32], with arbitrary positive strides. + return ( + source_row * fx.Index(stride0) + + source_k32 * fx.Index(stride1) + ) + + @flyc.kernel(known_block_size=[_SCALE_PACK_THREADS, 1, 1]) + def kernel_pack_mx32_scales(src: fx.Tensor, dst: fx.Tensor): + src_rsrc = buffer_ops.create_buffer_resource(src, max_size=True) + dst_rsrc = buffer_ops.create_buffer_resource(dst, max_size=True) + + linear = ( + fx.Index(fx.block_idx.x) * fx.Index(_SCALE_PACK_THREADS) + + fx.Index(gpu.thread_id("x")) ) + k128 = linear // fx.Index(dim) + dst_row = linear % fx.Index(dim) + + row_within_16 = dst_row % fx.Index(16) + k_subgroup = (dst_row // fx.Index(16)) % fx.Index(4) + tile = dst_row // fx.Index(64) + source_k32 = k128 * fx.Index(4) + k_subgroup + + def load_scale_byte(group): + source_row = ( + tile * fx.Index(64) + + fx.Index(group * 16) + + row_within_16 + ) + value_i8 = buffer_ops.buffer_load( + src_rsrc, + _source_offset(source_k32, source_row), + vec_width=1, + dtype=T.i8, + ) + # Preserve the raw E8M0 byte when widening. Going through Uint8 + # avoids sign extension for scale bytes >= 0x80. + return fx.Int32(fx.Uint8(value_i8)) + + b0 = load_scale_byte(0) + b1 = load_scale_byte(1) + b2 = load_scale_byte(2) + b3 = load_scale_byte(3) + packed = b0 | (b1 << 8) | (b2 << 16) | (b3 << 24) + buffer_ops.buffer_store(packed, dst_rsrc, linear) + + @flyc.jit + def launch_pack_mx32_scales( + src: fx.Tensor, + dst: fx.Tensor, + stream: fx.Stream = fx.Stream(None), + ): + kernel_pack_mx32_scales(src, dst).launch( + grid=(total_words // _SCALE_PACK_THREADS, 1, 1), + block=(_SCALE_PACK_THREADS, 1, 1), + stream=stream, + ) + + return launch_pack_mx32_scales - s32 = scales_u8.contiguous().view(dim, qk // 4, 4).to(torch.int32) - packed = ( - s32[:, :, 0] - | (s32[:, :, 1] << 8) - | (s32[:, :, 2] << 16) - | (s32[:, :, 3] << 24) + +@functools.lru_cache(maxsize=None) +def _cached_mx32_scale_pack_launch( + dim: int, + qk: int, + source_colwise: bool, + stride0: int, + stride1: int, +): + """Cache orientation-and-stride-specialized fused pack binaries.""" + return _compile_mx32_scale_pack_kernel( + dim, qk, source_colwise, stride0, stride1 ) - return packed.transpose(0, 1).contiguous() def pack_mx32_scales_for_hk( scales_u8: torch.Tensor, *, source_colwise: bool = False, + stream=None, ) -> torch.Tensor: - """Convert raw TE E8M0 scales to ``[K/128, dim]`` MFMA-ready words.""" - scale_iter = pack_mx32_scales_iter( - scales_u8, - source_colwise=source_colwise, - ) - dim = scales_u8.shape[1] if source_colwise else scales_u8.shape[0] + """Launch one fused GPU kernel producing HK MFMA-ready scale words. - if dim % 64 != 0: + Input contracts: + * rowwise: ``[dim, K/32]`` + * columnwise: ``[K/32, dim]`` + + Output contract: + * ``[K/128, dim]`` ``torch.int32`` + """ + if scales_u8.dtype != torch.uint8: + raise TypeError( + f"MXFP8 scales must be torch.uint8 E8M0 bytes, got " + f"{scales_u8.dtype}" + ) + if scales_u8.ndim != 2: raise ValueError( - f"Scale outer dimension={dim} must be a multiple of 64 for HK MFMA packing" + f"MXFP8 scales must be rank 2, got {tuple(scales_u8.shape)}" + ) + if not scales_u8.is_cuda: + raise ValueError("MXFP8 scale packing requires a CUDA/ROCm tensor") + if any(stride <= 0 for stride in scales_u8.stride()): + raise ValueError( + f"MXFP8 scale packing requires positive strides, got " + f"{scales_u8.stride()}" ) - device = scales_u8.device - row = torch.arange(dim, device=device, dtype=torch.int64) - row_within_16 = row % 16 - k_subgroup = (row // 16) % 4 - tile = row // 64 - - packed = torch.zeros_like(scale_iter) - 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, dim) - ) & 0xFF - packed |= byte_value << (group * 8) + if source_colwise: + qk, dim = scales_u8.shape + else: + dim, qk = scales_u8.shape - return packed.contiguous() + if qk % 4 != 0: + raise ValueError( + f"Scale K/32 dimension={qk} must be divisible by 4" + ) + if dim % 64 != 0: + raise ValueError( + f"Scale outer dimension={dim} must be a multiple of 64" + ) + packed = torch.empty( + (qk // 4, dim), + dtype=torch.int32, + device=scales_u8.device, + ) + if stream is None: + stream = torch.cuda.current_stream(scales_u8.device) + + stride0, stride1 = (int(x) for x in scales_u8.stride()) + _cached_mx32_scale_pack_launch( + dim, + qk, + bool(source_colwise), + stride0, + stride1, + )( + scales_u8, + packed, + stream=stream, + ) + return packed def _encode_waitcnt(vmcnt=63, lgkmcnt=15): """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. @@ -1643,10 +1744,14 @@ def mxfp8_matmul( a_scale_hk = pack_mx32_scales_for_hk( a_scale, source_colwise=False, + stream=stream, ) + # b_scale is already TE columnwise [K/32,N]. Consume it directly; + # do not launch an eager transpose/contiguous kernel. b_scale_hk = pack_mx32_scales_for_hk( - b_scale.transpose(0, 1).contiguous(), - source_colwise=False, + b_scale, + source_colwise=True, + stream=stream, ) elif layout == "NN": a_kernel = a @@ -1654,10 +1759,12 @@ def mxfp8_matmul( a_scale_hk = pack_mx32_scales_for_hk( a_scale, source_colwise=False, + stream=stream, ) b_scale_hk = pack_mx32_scales_for_hk( b_scale, source_colwise=True, + stream=stream, ) else: a_kernel = a @@ -1665,10 +1772,12 @@ def mxfp8_matmul( a_scale_hk = pack_mx32_scales_for_hk( a_scale, source_colwise=True, + stream=stream, ) b_scale_hk = pack_mx32_scales_for_hk( b_scale, source_colwise=True, + stream=stream, ) _debug( @@ -1707,6 +1816,7 @@ def mxfp8_matmul_nt(*args, **kwargs): "BLOCK_N", "BLOCK_K", "SCALE_GROUP_SIZE", + "pack_mx32_scales_for_hk", "do_gemm", "mxfp8_matmul", "mxfp8_matmul_nn", From fbc686202e1d40182ba05e659b46ccb1b2d247d8 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 29 Jul 2026 14:13:21 +0000 Subject: [PATCH 25/31] Normalize MXFP8 TN to direct rowwise storage --- .../flydsl_kernels/gemm/gemm_wrappers.py | 16 +++++------ .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 28 +++++++++++-------- 2 files changed, 25 insertions(+), 19 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 2e8db0b4d..ebcdbbca9 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -561,10 +561,10 @@ def _run_mxfp8( """Dispatch MXFP8 through exact TN/NN/NT physical contracts. TE owns BLAS-shaped operands. After the usual ownership swap, FlyDSL - kernels consume: + kernels consume the selected backing directly: TN: a = B.rowwise [M, K] - b = A.rowwise.T [K, N] (validated TN adapter contract) + b = A.rowwise [N, K] NN: a = B.rowwise [M, K] b = A.columnwise [K, N] @@ -659,17 +659,17 @@ def _run_mxfp8( # kernel a <- TE B # kernel b <- TE A if kernel_layout == "TN": - # Preserve the validated TN adapter contract: - # a [M,K], b [K,N] + # Selected rowwise backings already match the TN normal-read contract: + # a [M,K], b [N,K] a_flydsl = B_data - b_flydsl = A_data.transpose(0, 1) + b_flydsl = A_data a_scale = B_scale - b_scale = A_scale.transpose(0, 1) + b_scale = A_scale m, k = a_flydsl.shape - kb, n = b_flydsl.shape + n, kb = b_flydsl.shape expected_a_scale = (m, k // 32) - expected_b_scale = (k // 32, n) + expected_b_scale = (n, k // 32) elif kernel_layout == "NN": # A's columnwise MXFP8 payload is still row-major in its original diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 7d34495c9..76e82b5c5 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -1683,20 +1683,23 @@ def mxfp8_matmul( Wrapper-visible contracts: - TN: a [M,K], b [K,N], scales [M,K/32] and [K/32,N] + TN: a [M,K], b [N,K], scales [M,K/32] and [N,K/32] NN: a [M,K], b [K,N], scales [M,K/32] and [K/32,N] NT: a [K,M], b [K,N], scales [K/32,M] and [K/32,N] - TN preserves the existing adapter conversion to the kernel's normal-read - B [N,K] representation. NN and NT preserve K-major payloads and use - ``ds_read_b64_tr_b8`` inside their compile-time-specialized kernels. + TN consumes the selected rowwise payloads directly. NN and NT preserve + K-major payloads and use ``ds_read_b64_tr_b8`` inside their + compile-time-specialized kernels. """ if layout not in ("TN", "NN", "NT"): raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") _validate_common_payloads(a, b, D, layout=layout) - if layout in ("TN", "NN"): + if layout == "TN": + m, k = a.shape + n, kb = b.shape + elif layout == "NN": m, k = a.shape kb, n = b.shape else: @@ -1720,7 +1723,11 @@ def mxfp8_matmul( expected_a_scale = (k // SCALE_GROUP_SIZE, m) else: expected_a_scale = (m, k // SCALE_GROUP_SIZE) - expected_b_scale = (k // SCALE_GROUP_SIZE, n) + + if layout == "TN": + expected_b_scale = (n, k // SCALE_GROUP_SIZE) + else: + expected_b_scale = (k // SCALE_GROUP_SIZE, n) if tuple(a_scale.shape) != expected_a_scale: raise ValueError( @@ -1738,19 +1745,18 @@ def mxfp8_matmul( raise ValueError("A, B, scales, and D must be on the same device") if layout == "TN": - # Preserve the passing TN kernel contract exactly: normal-read B [N,K]. + # TN selected backings already match the normal-read kernel contract: + # a [M,K], b [N,K] a_kernel = a - b_kernel = b.transpose(0, 1).contiguous() + b_kernel = b a_scale_hk = pack_mx32_scales_for_hk( a_scale, source_colwise=False, stream=stream, ) - # b_scale is already TE columnwise [K/32,N]. Consume it directly; - # do not launch an eager transpose/contiguous kernel. b_scale_hk = pack_mx32_scales_for_hk( b_scale, - source_colwise=True, + source_colwise=False, stream=stream, ) elif layout == "NN": From 4d7d273127e428e7b9a259c88b8cfb0ec5bcab37 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 29 Jul 2026 17:07:10 +0000 Subject: [PATCH 26/31] improve backward bf16 gemms with shape specialization --- .../pytorch/flydsl_kernels/gemm/bf16_gemm.py | 476 ++++++++++++++---- .../flydsl_kernels/gemm/fp16_gemm_utils.py | 166 +++++- .../flydsl_kernels/gemm/gemm_wrappers.py | 125 ++++- 3 files changed, 648 insertions(+), 119 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py index 4201d571d..cf267d695 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py @@ -2,14 +2,18 @@ # # See LICENSE for license information. -"""FlyDSL BF16 4-wave GEMM kernel for Transformer Engine. +"""FlyDSL BF16 TN/NN/NT 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 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. +All supported layouts share one source-level kernel generator while compiling +to separate cached binaries: + + TN: A [M,K] normal read, B [N,K] normal read + NN: A [M,K] normal read, B [K,N] transpose read + NT: A [K,M] transpose read, B [K,N] transpose read + +The layout is a Python-only cache key. Global addressing and LDS fragment +reconstruction are selected while building each specialization, so no runtime +layout branch is emitted in the GEMM kernel. This module imports ``flydsl`` at import time and must therefore be imported lazily only after FlyDSL availability has been confirmed. @@ -31,6 +35,7 @@ from .fp16_gemm_utils import ( G2SLoader, S2RLoader, + compute_global_bf16_transpose_swizzle, compute_global_swizzle, make_bf16_byte_buffer_tensor, pack_i32x4_i32x8, @@ -188,13 +193,20 @@ def _xcd_swizzle(num_pid_m, num_pid_n): def _compile_kernel( K: int, output_dtype: torch.dtype, + layout: str, use_xcd_remap: bool = True, ): - """Build the specialized 4-wave kernel for compile-time ``K`` and output dtype. + """Build one compile-time-specialized TN, NN, or NT BF16 kernel. ``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. """ + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported BF16 kernel layout: {layout}") + + a_transpose_read = layout == "NT" + b_transpose_read = layout in ("NN", "NT") + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K NUM_THREADS = 256 WARP_SIZE = 64 @@ -247,11 +259,180 @@ def _compile_kernel( LOAD_PASSES_HALF = LDS_BYTES_HALF // (NUM_THREADS * VEC_BYTES) assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + # Resolve layout-specific addressing and fragment reads before capture. + Q0_SCHED_DSRD = 4 if a_transpose_read else 2 + PREFETCH_SCHED_DSRD = 8 if a_transpose_read else 4 + + if a_transpose_read: + def _a_leading_dim_bytes(c_m): + return c_m * ELEM_BYTES + + def _a_global_base_bytes(k_base, subtile, c_m, bx_m_idx): + return ( + k_base * fx.Index(c_m * ELEM_BYTES) + + (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) + * fx.Index(ELEM_BYTES) + ) + + def _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ): + del load_frag_half_at_byte_base, lane_mod_16 + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + local_m_tile = ( + subtile_m_idx * fx.Index(SUBTILE_M) + + fx.Index(mi * MFMA_M) + - fx.Index(sm * (BLOCK_M // 2)) + ) + return load_transposed_frag_half(lds_a[sm], local_m_tile, half) + else: + def _a_leading_dim_bytes(c_m): + del c_m + return K * ELEM_BYTES + + def _a_global_base_bytes(k_base, subtile, c_m, bx_m_idx): + del c_m + return ( + (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) + * fx.Index(K * ELEM_BYTES) + + k_base * fx.Index(ELEM_BYTES) + ) + + def _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ): + del load_transposed_frag_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)) + return load_frag_half_at_byte_base( + lds_a[sm], + half_row * fx.Index(BLOCK_K * ELEM_BYTES), + half, + ) + + if b_transpose_read: + def _b_leading_dim_bytes(c_n): + return c_n * ELEM_BYTES + + def _b_global_base_bytes(k_base, subtile, c_n, by_n_idx): + return ( + k_base * fx.Index(c_n * ELEM_BYTES) + + (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) + * fx.Index(ELEM_BYTES) + ) + + def _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ): + del load_normal_b_frag, lane_mod_16 + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + local_n_tile = ( + subtile_n_idx * fx.Index(SUBTILE_N) + + fx.Index(ni * MFMA_N) + - fx.Index(sn * (BLOCK_N // 2)) + ) + return load_transposed_frag(lds_b[sn], local_n_tile) + else: + def _b_leading_dim_bytes(c_n): + del c_n + return K * ELEM_BYTES + + def _b_global_base_bytes(k_base, subtile, c_n, by_n_idx): + del c_n + return ( + (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) + * fx.Index(K * ELEM_BYTES) + + k_base * fx.Index(ELEM_BYTES) + ) + + def _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ): + del load_transposed_frag + 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_normal_b_frag(lds_b, b_row_addr, sn) + + # Resolve global staging maps before FlyDSL captures ``kernel_gemm``. + # BF16 uses K64, so each transpose-read half-page is two independent + # [K64, X64] slices with 128-byte physical rows. + if a_transpose_read: + def _a_global_offsets(lane, wave_id, c_m): + return compute_global_bf16_transpose_swizzle( + lane, + wave_id, + _a_leading_dim_bytes(c_m), + LOAD_PASSES_HALF, + ) + else: + def _a_global_offsets(lane, wave_id, c_m): + del c_m + return compute_global_swizzle( + lane, + wave_id, + K * ELEM_BYTES, + LOAD_PASSES_HALF, + preshuffled=False, + ) + + if b_transpose_read: + def _b_global_offsets(lane, wave_id, c_n): + return compute_global_bf16_transpose_swizzle( + lane, + wave_id, + _b_leading_dim_bytes(c_n), + LOAD_PASSES_HALF, + ) + else: + def _b_global_offsets(lane, wave_id, c_n): + del c_n + return compute_global_swizzle( + lane, + wave_id, + K * ELEM_BYTES, + LOAD_PASSES_HALF, + preshuffled=False, + ) + @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. + # Preserve the passing TN byte-staging contract exactly. A BF16 K64 + # half-page is 128 rows x 128 bytes = 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] @@ -275,8 +456,9 @@ def kernel_gemm( 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. + # A/B arrive as contiguous uint8 byte views of the original + # row-major BF16 tensors. This preserves the validated 16-byte + # BufferCopyLDS128b path and byte-based address arithmetic. 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)) @@ -307,13 +489,26 @@ def kernel_gemm( 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) + # Offsets are always bytes. TN uses the original 128-byte XOR + # swizzle. NN/NT stage K-major BF16 data as two [K64, X64] slices for + # ds_read_b64_tr_b16; the layout choice was resolved before capture. + gl_off_a = _a_global_offsets(lane, wave_id, c_m) + gl_off_b = _b_global_offsets(lane, wave_id, c_n) + + 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)) @@ -410,7 +605,7 @@ def hot_loop_scheduler_q0_refill_a1_2n(): # reads overlap four independent 8-MFMA K32 groups. for _ in range_constexpr(4): rocdl.sched_vmem(2) - rocdl.sched_dsrd(2) + rocdl.sched_dsrd(Q0_SCHED_DSRD) rocdl.sched_mfma(8) rocdl.sched_barrier(0) @@ -418,18 +613,22 @@ 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_dsrd(PREFETCH_SCHED_DSRD) 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) + global_base = _a_global_base_bytes( + k_base, subtile, c_m, bx_m_idx + ) 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) + global_base = _b_global_base_bytes( + k_base, subtile, c_n, by_n_idx + ) b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) def stage_a_subtile(k_base, subtile, lds_a): @@ -459,7 +658,47 @@ def load_frag_at_byte_base(lds_page, row_byte_base): 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)) + return load_frag_at_byte_base( + lds_b[half], + half_row * fx.Index(BLOCK_K * ELEM_BYTES), + ) + + def load_transposed_frag_half(lds_page, local_x_tile, half): + # BF16 uses v_mfma_f32_16x16x32_bf16, not the MXFP8 K128 + # instruction. A 128-X half-page is therefore two independent + # swizzled [K64, X64] BF16 slices. One ds_read_b64_tr_b16 returns + # four BF16 values/lane; two reads form one K32 MFMA fragment. + local_x_i32 = fx.Int32(local_x_tile) + slice_idx = local_x_i32 // fx.Int32(64) + x_in_slice = local_x_i32 % fx.Int32(64) + lane_div16_i32 = fx.Int32(lane_div_16) + lane_in16_i32 = fx.Int32(lane_mod_16) + + source_k = ( + lane_div16_i32 * fx.Int32(8) + + lane_in16_i32 // fx.Int32(4) + ) + source_x_byte = ( + x_in_slice * fx.Int32(ELEM_BYTES) + + (lane_in16_i32 % fx.Int32(4)) * fx.Int32(8) + ) + + physical_k, physical_x = swizzle_128(source_k, source_x_byte) + slice_base = slice_idx * fx.Int32(64 * 128) + base = slice_base + physical_k * fx.Int32(128) + physical_x + other = base ^ fx.Int32(0x220) + immediate_offset = 0 if half == 0 else 0x1000 + return s2r.load_one_transpose_bf16( + lds_page, + base, + other, + immediate_offset=immediate_offset, + ) + + def load_transposed_frag(lds_page, local_x_tile): + x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) + x1 = load_transposed_frag_half(lds_page, local_x_tile, 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 @@ -585,9 +824,15 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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) + return _load_b_ni( + load_transposed_frag, + load_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ) def load_b_subtile_regs(lds_b, sn): return ( @@ -598,11 +843,16 @@ def load_b_subtile_regs(lds_b, sn): ) 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) + return _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ) def load_a_subtile_mi_regs(lds_a, sm, mi): x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) @@ -1003,11 +1253,13 @@ def launch_gemm( def _cached_launch( K: int, output_dtype: torch.dtype, + layout: str, use_xcd_remap: bool = True, ): return _compile_kernel( K, output_dtype, + layout, use_xcd_remap=use_xcd_remap, ) @@ -1016,95 +1268,124 @@ def bf16_matmul( a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, + *, + layout: str, + m: int, + n: int, + k: int, stream=None, ): - """TE-facing TN BF16 GEMM adapter. - - Public/backend contract: - a: [M, K] BF16 - b: [K, N] BF16 - 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 - transpose view of contiguous rowwise weight storage, so ``b.T`` is already - contiguous and does not require a physical transpose. - """ + """Launch the wrapper-selected BF16 TN/NN/NT specialization.""" + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported BF16 layout: {layout}") if a.ndim != 2 or b.ndim != 2: raise ValueError( - f"FlyDSL BF16 TN expects rank-2 operands, got A{tuple(a.shape)} " + f"FlyDSL BF16 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}" + "FlyDSL BF16 GEMM expects torch.bfloat16 operands, " + f"got A={a.dtype}, B={b.dtype}" + ) + if not a.is_contiguous() or not b.is_contiguous(): + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 {layout} requires original contiguous row-major " + f"operands, got A stride={tuple(a.stride())}, " + f"B stride={tuple(b.stride())}" + ) + + m = int(m) + n = int(n) + k = int(k) + + expected_shapes = { + "TN": ((m, k), (n, k)), + "NN": ((m, k), (k, n)), + "NT": ((k, m), (k, n)), + } + expected_a, expected_b = expected_shapes[layout] + if tuple(a.shape) != expected_a or tuple(b.shape) != expected_b: + raise ValueError( + f"FlyDSL BF16 {layout} physical operands do not match contract: " + f"A{tuple(a.shape)} expected {expected_a}; " + f"B{tuple(b.shape)} expected {expected_b}" ) + 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, - ): + if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise TypeError( - "FlyDSL BF16 GEMM output dtype must be torch.float16, " - f"torch.bfloat16, or torch.float32, got {c.dtype}" + "FlyDSL BF16 output must be float16, bfloat16, or float32, " + f"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}" + 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 BF16 GEMM requires contiguous output storage") - b_hk = b.transpose(0, 1).contiguous() - doGemm(a, b_hk, c, stream=stream) - + doGemm( + a, + b, + c, + layout=layout, + m=m, + n=n, + k=k, + stream=stream, + ) def doGemm( A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, + *, + layout: str, + m: int, + n: int, + k: int, stream=None, use_xcd_remap: bool = True, ): - """Launch the private K-specialized BF16 core. + """Launch one cached K/output/layout-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. + A and B are passed unchanged from ``gemm_wrappers.py``. Their pointers + reference the original rowwise allocations: + + TN: A backing [M,K], B backing [N,K] + NN: A backing [M,K], B backing [K,N] + NT: A backing [K,M], B backing [K,N] + + NN/NT orientation is implemented by compile-time global addressing and + ``ds_read_b64_tr_b16`` only. """ - 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 in ( - torch.float16, - torch.bfloat16, - torch.float32, - ), ( - "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " - f"got {C.dtype}" - ) + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported BF16 layout: {layout}") + + M_runtime = int(m) + N_runtime = int(n) + K_runtime = int(k) + + if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: + raise TypeError( + f"BF16 {layout} requires BF16 inputs, got {A.dtype} and {B.dtype}" + ) + if C.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError(f"Unsupported BF16 output dtype: {C.dtype}") + 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}, " @@ -1117,16 +1398,33 @@ def doGemm( 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 tuple(C.shape) != (M_runtime, N_runtime): + raise ValueError( + f"C shape {tuple(C.shape)} != expected {(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), + K_runtime, C.dtype, + layout, bool(use_xcd_remap), ) - launch(A_arg, B_arg, C_arg, M_runtime, N_runtime, stream=stream) + # Preserve the original validated byte-addressed G2L path. These are + # metadata-only dtype/flatten views of the already-contiguous row-major + # tensors selected by gemm_wrappers.py; no transpose or copy is performed. + A_arg = A.view(torch.uint8).view(-1) + B_arg = B.view(torch.uint8).view(-1) + C_arg = C.view(-1) + + 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 index 5aaeab3b1..b0629f21a 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py @@ -1,17 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors -"""Minimal byte-staging helpers for the first-pass BF16 four-wave GEMM.""" +"""Byte-staging helpers for the BF16 four-wave GEMM.""" import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm as _llvm, vector from flydsl.expr import const_expr, range_constexpr +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: return (numer + denom - 1) // denom ceildiv = cdiv + def divmod(a, b): return (a // b, a % b) @@ -24,17 +29,30 @@ def swizzle_128(row, col_in_bytes): 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 make_bf16_buffer_tensor(arg_bf16): + """Create a BF16 BufferDesc directly from the wrapper-provided tensor.""" + return fx.rocdl.make_buffer_tensor(arg_bf16, max_size=False) + +# Backward-compatible name used by fp16_gemm.py. +# Keep the exact existing behavior; this is only a symbol alias. +make_bf16_byte_buffer_tensor = make_bf16_buffer_tensor -def compute_global_swizzle(lane_id, wave_id, row_stride_bytes, n_rounds, preshuffled=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") + 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) @@ -42,12 +60,45 @@ def compute_global_swizzle(lane_id, wave_id, row_stride_bytes, n_rounds, preshuf return offsets -class G2SLoader: - """Issue raw 16-byte buffer-to-LDS copies. +def compute_global_bf16_transpose_swizzle( + lane_id, + wave_id, + leading_dim_bytes, + n_rounds, +): + """Offsets for a K-major BF16 source staged for ``ds_read_b64_tr_b16``. - 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. + One 128-row output half-page is represented in LDS as two independent + swizzled ``[K64, X64]`` slices. Each slice is 64 rows by 128 bytes, so the + complete half-page remains 16 KiB and preserves the existing four-pass + 16-byte/thread DMA cadence. + + The returned offsets are relative to the source tile base: + ``source[k, x_base]`` for a contiguous K-major BF16 matrix. """ + offsets = [] + n_waves = fx.block_dim.x // 64 + for round in range_constexpr(n_rounds): + linear_row = lane_id // 8 + wave_id * 8 + round * (n_waves * 8) + col_bytes = (lane_id % 8) * 16 + + slice_idx = linear_row // 64 + physical_k = linear_row % 64 + + # XOR swizzle is self-inverse for this layout. Map the physical LDS + # chunk back to its logical K/X-byte source coordinate. + logical_k, logical_x_bytes = swizzle_128(physical_k, col_bytes) + offsets.append( + logical_k * leading_dim_bytes + + slice_idx * 64 * 2 + + logical_x_bytes + ) + return offsets + + +class G2SLoader: + """Issue native 16-byte BF16 BufferDesc-to-BF16 LDS copies.""" + 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) @@ -60,17 +111,36 @@ def __init__(self, gl_src, gl_offsets, n_load_steps, lds_dtype, wave_id): 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)) + 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)) + 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)) + 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): @@ -78,7 +148,8 @@ def pack_i32x4_i32x8(lo, hi): class S2RLoader: - """Raw 16-byte LDS reader used to assemble an i32x8 BF16 K64 fragment.""" + """LDS readers used to assemble BF16 K64 fragments.""" + def __init__(self, wave_idx, n_tiles): self.lane_id = fx.thread_idx.x % 64 self.wave_idx = wave_idx @@ -90,4 +161,63 @@ def _vec_load_16bytes(self, lds_src, offset): 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) + return self._vec_load_16bytes( + lds_src, + lds_offset, + ).bitcast(fx.Int32) + + def _ds_read_b64_tr_b16( + self, + lds_src, + byte_offset, + immediate_offset=0, + ): + """Issue one gfx950 ``ds_read_b64_tr_b16`` and return i32x2.""" + if immediate_offset == 0: + asm = "ds_read_b64_tr_b16 $0, $1 offset:0\n" + elif immediate_offset == 0x1000: + asm = "ds_read_b64_tr_b16 $0, $1 offset:4096\n" + else: + raise ValueError( + "ds_read_b64_tr_b16 supports immediate offsets 0 and 0x1000, " + f"got {immediate_offset:#x}" + ) + + 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)], + asm, + "=v,v,~{memory}", + has_side_effects=True, + ) + return Vec( + vector.BitCastOp(raw_type, raw).result, + (2,), + fx.Int32, + ) + + def load_one_transpose_bf16( + self, + lds_src, + first_byte_offset, + second_byte_offset, + immediate_offset=0, + ): + """Return one i32x4 K32 BF16 fragment from two transpose reads.""" + lo = self._ds_read_b64_tr_b16( + lds_src, + first_byte_offset, + immediate_offset, + ) + hi = self._ds_read_b64_tr_b16( + lds_src, + second_byte_offset, + immediate_offset, + ) + return lo.shuffle(hi, [0, 1, 2, 3]) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index ebcdbbca9..4f3fef8aa 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -237,10 +237,9 @@ def _canonicalize_blas_pair( 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 + """Swap TE BLAS operand ownership without changing either tensor layout.""" + del transa, transb + return B_data, A_data def _flatten_rowwise(t: torch.Tensor, name: str) -> torch.Tensor: @@ -275,11 +274,10 @@ def _canonicalize_blas_operands( a_flydsl: [M, K] b_flydsl: [K, N] - The standard conversion is to swap A/B and apply the original transpose - flags to the swapped operands: + Operand ownership is swapped without creating tensor transpose views: - a_flydsl = op(B) - b_flydsl = op(A) + a_flydsl = B + b_flydsl = A """ if transa and transb: raise NotImplementedError( @@ -338,6 +336,112 @@ def _validate_or_allocate_output( return D +def _run_bf16_gemm( + A, + transa, + B, + transb, + D, + *, + output_dtype: torch.dtype, +): + """Dispatch BF16 using the original row-major operand allocations. + + No operand transpose view is created: + + TN: kernel A = TE B [M,K], kernel B = TE A [N,K] + NN: kernel A = TE B [M,K], kernel B = TE A [K,N] + NT: kernel A = TE B [K,M], kernel B = TE A [K,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 torch.bfloat16 inputs, " + 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}" + ) + + dispatch = { + (True, False): "TN", + (False, False): "NN", + (False, True): "NT", + } + try: + layout = dispatch[(bool(transa), bool(transb))] + except KeyError as exc: + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) from exc + + output_shape = _get_gemm_output_shape(A, transa, B, transb) + + # Preserve the original row-major storage. This only collapses leading + # batch dimensions, matching the wrapper's existing regular-GEMM contract. + A_data = _flatten_rowwise(A, "A") + B_data = _flatten_rowwise(B, "B") + + # Kernel ownership is always swapped relative to TE's BLAS arguments. + a_flydsl = B_data + b_flydsl = A_data + + if layout == "TN": + m, k = a_flydsl.shape + n, kb = b_flydsl.shape + expected_a = (m, k) + expected_b = (n, k) + elif layout == "NN": + m, k = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a = (m, k) + expected_b = (k, n) + else: + k, m = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a = (k, m) + expected_b = (k, n) + + if kb != k: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 {layout} received incompatible row-major operands: " + f"a={tuple(a_flydsl.shape)}, b={tuple(b_flydsl.shape)}" + ) + if tuple(a_flydsl.shape) != expected_a or tuple(b_flydsl.shape) != expected_b: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 {layout} physical contract mismatch: " + f"a={tuple(a_flydsl.shape)} expected={expected_a}; " + f"b={tuple(b_flydsl.shape)} expected={expected_b}" + ) + + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL BF16 logical output shape {tuple(output_shape)} " + f"does not match kernel output {(m, n)}" + ) + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=A.device, + backend_name=f"BF16 {layout}", + ) + + bf16_matmul( + a_flydsl, + b_flydsl, + D.view(m, n), + layout=layout, + m=m, + n=n, + k=k, + ) + return D + + def _run_regular_gemm( A, transa, @@ -1125,15 +1229,12 @@ def te_generic_gemm_flydsl( "FlyDSL BF16 supports FP16, BF16, or FP32 output, " f"got {output_dtype}" ) - D = _run_regular_gemm( + D = _run_bf16_gemm( A, transa, B, transb, D, - dtype=torch.bfloat16, - matmul=bf16_matmul, - backend_name="BF16", output_dtype=bf16_output_dtypes[output_dtype], ) return D, None, None, None From e928833ccb8fe8ff8b1cf5fa7f4b588d452071c9 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 29 Jul 2026 18:21:37 +0000 Subject: [PATCH 27/31] add shape specialization for flydsl fp16 gemm backend --- .../pytorch/flydsl_kernels/gemm/fp16_gemm.py | 485 ++++++++++++++---- .../flydsl_kernels/gemm/fp16_gemm_utils.py | 20 + .../flydsl_kernels/gemm/gemm_wrappers.py | 112 +++- 3 files changed, 516 insertions(+), 101 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py index 709f76484..ace8ecda4 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py @@ -2,14 +2,18 @@ # # See LICENSE for license information. -"""FlyDSL FP16 4-wave GEMM kernel for Transformer Engine. +"""FlyDSL FP16 TN/NN/NT 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, -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. +All supported layouts share one source-level kernel generator while compiling +to separate cached binaries: + + TN: A [M,K] normal read, B [N,K] normal read + NN: A [M,K] normal read, B [K,N] transpose read + NT: A [K,M] transpose read, B [K,N] transpose read + +The layout is a Python-only cache key. Global addressing and LDS fragment +reconstruction are selected while building each specialization, so no runtime +layout branch is emitted in the GEMM kernel. This module imports ``flydsl`` at import time and must therefore be imported lazily only after FlyDSL availability has been confirmed. @@ -31,8 +35,9 @@ from .fp16_gemm_utils import ( G2SLoader, S2RLoader, + compute_global_fp16_transpose_swizzle, compute_global_swizzle, - make_bf16_byte_buffer_tensor as make_fp16_byte_buffer_tensor, + make_fp16_byte_buffer_tensor, pack_i32x4_i32x8, swizzle_128, ) @@ -92,12 +97,6 @@ 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.""" @@ -194,13 +193,20 @@ def _xcd_swizzle(num_pid_m, num_pid_n): def _compile_kernel( K: int, output_dtype: torch.dtype, + layout: str, use_xcd_remap: bool = True, ): - """Build the specialized 4-wave kernel for compile-time ``K``. + """Build one compile-time-specialized TN, NN, or NT FP16 kernel. ``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. """ + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported FP16 kernel layout: {layout}") + + a_transpose_read = layout == "NT" + b_transpose_read = layout in ("NN", "NT") + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K NUM_THREADS = 256 WARP_SIZE = 64 @@ -253,11 +259,180 @@ def _compile_kernel( LOAD_PASSES_HALF = LDS_BYTES_HALF // (NUM_THREADS * VEC_BYTES) assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + # Resolve layout-specific addressing and fragment reads before capture. + Q0_SCHED_DSRD = 4 if a_transpose_read else 2 + PREFETCH_SCHED_DSRD = 8 if a_transpose_read else 4 + + if a_transpose_read: + def _a_leading_dim_bytes(c_m): + return c_m * ELEM_BYTES + + def _a_global_base_bytes(k_base, subtile, c_m, bx_m_idx): + return ( + k_base * fx.Index(c_m * ELEM_BYTES) + + (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) + * fx.Index(ELEM_BYTES) + ) + + def _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ): + del load_frag_half_at_byte_base, lane_mod_16 + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + local_m_tile = ( + subtile_m_idx * fx.Index(SUBTILE_M) + + fx.Index(mi * MFMA_M) + - fx.Index(sm * (BLOCK_M // 2)) + ) + return load_transposed_frag_half(lds_a[sm], local_m_tile, half) + else: + def _a_leading_dim_bytes(c_m): + del c_m + return K * ELEM_BYTES + + def _a_global_base_bytes(k_base, subtile, c_m, bx_m_idx): + del c_m + return ( + (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) + * fx.Index(K * ELEM_BYTES) + + k_base * fx.Index(ELEM_BYTES) + ) + + def _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ): + del load_transposed_frag_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)) + return load_frag_half_at_byte_base( + lds_a[sm], + half_row * fx.Index(BLOCK_K * ELEM_BYTES), + half, + ) + + if b_transpose_read: + def _b_leading_dim_bytes(c_n): + return c_n * ELEM_BYTES + + def _b_global_base_bytes(k_base, subtile, c_n, by_n_idx): + return ( + k_base * fx.Index(c_n * ELEM_BYTES) + + (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) + * fx.Index(ELEM_BYTES) + ) + + def _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ): + del load_normal_b_frag, lane_mod_16 + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + local_n_tile = ( + subtile_n_idx * fx.Index(SUBTILE_N) + + fx.Index(ni * MFMA_N) + - fx.Index(sn * (BLOCK_N // 2)) + ) + return load_transposed_frag(lds_b[sn], local_n_tile) + else: + def _b_leading_dim_bytes(c_n): + del c_n + return K * ELEM_BYTES + + def _b_global_base_bytes(k_base, subtile, c_n, by_n_idx): + del c_n + return ( + (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) + * fx.Index(K * ELEM_BYTES) + + k_base * fx.Index(ELEM_BYTES) + ) + + def _load_b_ni( + load_transposed_frag, + load_normal_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ): + del load_transposed_frag + 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_normal_b_frag(lds_b, b_row_addr, sn) + + # Resolve global staging maps before FlyDSL captures ``kernel_gemm``. + # FP16 uses K64, so each transpose-read half-page is two independent + # [K64, X64] slices with 128-byte physical rows. + if a_transpose_read: + def _a_global_offsets(lane, wave_id, c_m): + return compute_global_fp16_transpose_swizzle( + lane, + wave_id, + _a_leading_dim_bytes(c_m), + LOAD_PASSES_HALF, + ) + else: + def _a_global_offsets(lane, wave_id, c_m): + del c_m + return compute_global_swizzle( + lane, + wave_id, + K * ELEM_BYTES, + LOAD_PASSES_HALF, + preshuffled=False, + ) + + if b_transpose_read: + def _b_global_offsets(lane, wave_id, c_n): + return compute_global_fp16_transpose_swizzle( + lane, + wave_id, + _b_leading_dim_bytes(c_n), + LOAD_PASSES_HALF, + ) + else: + def _b_global_offsets(lane, wave_id, c_n): + del c_n + return compute_global_swizzle( + lane, + wave_id, + K * ELEM_BYTES, + LOAD_PASSES_HALF, + preshuffled=False, + ) + @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. + # Preserve the passing TN byte-staging contract exactly. A FP16 K64 + # half-page is 128 rows x 128 bytes = 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] @@ -281,8 +456,9 @@ def kernel_gemm( 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. + # A/B arrive as contiguous uint8 byte views of the original + # row-major FP16 tensors. This preserves the validated 16-byte + # BufferCopyLDS128b path and byte-based address arithmetic. 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)) @@ -313,13 +489,26 @@ def kernel_gemm( 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) + # Offsets are always bytes. TN uses the original 128-byte XOR + # swizzle. NN/NT stage K-major BF16 data as two [K64, X64] slices for + # ds_read_b64_tr_b16; the layout choice was resolved before capture. + gl_off_a = _a_global_offsets(lane, wave_id, c_m) + gl_off_b = _b_global_offsets(lane, wave_id, c_n) + + 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)) @@ -416,7 +605,7 @@ def hot_loop_scheduler_q0_refill_a1_2n(): # reads overlap four independent 8-MFMA K32 groups. for _ in range_constexpr(4): rocdl.sched_vmem(2) - rocdl.sched_dsrd(2) + rocdl.sched_dsrd(Q0_SCHED_DSRD) rocdl.sched_mfma(8) rocdl.sched_barrier(0) @@ -424,18 +613,22 @@ 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_dsrd(PREFETCH_SCHED_DSRD) 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) + global_base = _a_global_base_bytes( + k_base, subtile, c_m, bx_m_idx + ) 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) + global_base = _b_global_base_bytes( + k_base, subtile, c_n, by_n_idx + ) b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) def stage_a_subtile(k_base, subtile, lds_a): @@ -465,7 +658,47 @@ def load_frag_at_byte_base(lds_page, row_byte_base): 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)) + return load_frag_at_byte_base( + lds_b[half], + half_row * fx.Index(BLOCK_K * ELEM_BYTES), + ) + + def load_transposed_frag_half(lds_page, local_x_tile, half): + # FP16 uses v_mfma_f32_16x16x32_f16, not the MXFP8 K128 + # instruction. A 128-X half-page is therefore two independent + # swizzled [K64, X64] FP16 slices. One ds_read_b64_tr_b16 returns + # four FP16 values/lane; two reads form one K32 MFMA fragment. + local_x_i32 = fx.Int32(local_x_tile) + slice_idx = local_x_i32 // fx.Int32(64) + x_in_slice = local_x_i32 % fx.Int32(64) + lane_div16_i32 = fx.Int32(lane_div_16) + lane_in16_i32 = fx.Int32(lane_mod_16) + + source_k = ( + lane_div16_i32 * fx.Int32(8) + + lane_in16_i32 // fx.Int32(4) + ) + source_x_byte = ( + x_in_slice * fx.Int32(ELEM_BYTES) + + (lane_in16_i32 % fx.Int32(4)) * fx.Int32(8) + ) + + physical_k, physical_x = swizzle_128(source_k, source_x_byte) + slice_base = slice_idx * fx.Int32(64 * 128) + base = slice_base + physical_k * fx.Int32(128) + physical_x + other = base ^ fx.Int32(0x220) + immediate_offset = 0 if half == 0 else 0x1000 + return s2r.load_one_transpose_fp16( + lds_page, + base, + other, + immediate_offset=immediate_offset, + ) + + def load_transposed_frag(lds_page, local_x_tile): + x0 = load_transposed_frag_half(lds_page, local_x_tile, 0) + x1 = load_transposed_frag_half(lds_page, local_x_tile, 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 @@ -591,9 +824,15 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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) + return _load_b_ni( + load_transposed_frag, + load_b_frag, + lds_b, + sn, + ni, + reg_subtile_n_idx0, + lane_mod_16, + ) def load_b_subtile_regs(lds_b, sn): return ( @@ -604,11 +843,16 @@ def load_b_subtile_regs(lds_b, sn): ) 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) + return _load_a_half( + load_transposed_frag_half, + load_frag_half_at_byte_base, + lds_a, + sm, + mi, + half, + reg_subtile_m_idx0, + lane_mod_16, + ) def load_a_subtile_mi_regs(lds_a, sm, mi): x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) @@ -1005,16 +1249,17 @@ def launch_gemm( return launch_gemm - @functools.lru_cache(maxsize=None) def _cached_launch( K: int, output_dtype: torch.dtype, + layout: str, use_xcd_remap: bool = True, ): return _compile_kernel( K, output_dtype, + layout, use_xcd_remap=use_xcd_remap, ) @@ -1023,95 +1268,124 @@ def fp16_matmul( a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, + *, + layout: str, + m: int, + n: int, + k: int, stream=None, ): - """TE-facing TN FP16 GEMM adapter. - - Public/backend contract: - a: [M, K] FP16 - b: [K, N] FP16 - 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 - transpose view of contiguous rowwise weight storage, so ``b.T`` is already - contiguous and does not require a physical transpose. - """ + """Launch the wrapper-selected BF16 TN/NN/NT specialization.""" + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported FP16 layout: {layout}") if a.ndim != 2 or b.ndim != 2: raise ValueError( - f"FlyDSL FP16 TN expects rank-2 operands, got A{tuple(a.shape)} " + f"FlyDSL BF16 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}" + "FlyDSL FP16 GEMM expects torch.float16 operands, " + f"got A={a.dtype}, B={b.dtype}" + ) + if not a.is_contiguous() or not b.is_contiguous(): + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 {layout} requires original contiguous row-major " + f"operands, got A stride={tuple(a.stride())}, " + f"B stride={tuple(b.stride())}" ) + + m = int(m) + n = int(n) + k = int(k) + + expected_shapes = { + "TN": ((m, k), (n, k)), + "NN": ((m, k), (k, n)), + "NT": ((k, m), (k, n)), + } + expected_a, expected_b = expected_shapes[layout] + if tuple(a.shape) != expected_a or tuple(b.shape) != expected_b: + raise ValueError( + f"FlyDSL BF16 {layout} physical operands do not match contract: " + f"A{tuple(a.shape)} expected {expected_a}; " + f"B{tuple(b.shape)} expected {expected_b}" + ) + 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, - ): + if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise TypeError( - "FlyDSL FP16 GEMM output dtype must be torch.float16, " - f"torch.bfloat16, or torch.float32, got {c.dtype}" + "FlyDSL FP16 output must be float16, bfloat16, or float32, " + f"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}" + 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 FP16 GEMM requires contiguous output storage") - b_hk = b.transpose(0, 1).contiguous() - doGemm(a, b_hk, c, stream=stream) - + doGemm( + a, + b, + c, + layout=layout, + m=m, + n=n, + k=k, + stream=stream, + ) def doGemm( A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, + *, + layout: str, + m: int, + n: int, + k: int, stream=None, use_xcd_remap: bool = True, ): - """Launch the private K-specialized FP16 core. + """Launch one cached K/output/layout-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. + A and B are passed unchanged from ``gemm_wrappers.py``. Their pointers + reference the original rowwise allocations: + + TN: A backing [M,K], B backing [N,K] + NN: A backing [M,K], B backing [K,N] + NT: A backing [K,M], B backing [K,N] + + NN/NT orientation is implemented by compile-time global addressing and + ``ds_read_b64_tr_b16`` only. """ - 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 in ( - torch.float16, - torch.bfloat16, - torch.float32, - ), ( - "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " - f"got {C.dtype}" - ) + if layout not in ("TN", "NN", "NT"): + raise ValueError(f"Unsupported FP16 layout: {layout}") + + M_runtime = int(m) + N_runtime = int(n) + K_runtime = int(k) + + if A.dtype != torch.float16 or B.dtype != torch.float16: + raise TypeError( + f"BF16 {layout} requires BF16 inputs, got {A.dtype} and {B.dtype}" + ) + if C.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError(f"Unsupported FP16 output dtype: {C.dtype}") + 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}, " @@ -1124,16 +1398,33 @@ def doGemm( 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 tuple(C.shape) != (M_runtime, N_runtime): + raise ValueError( + f"C shape {tuple(C.shape)} != expected {(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), + K_runtime, C.dtype, + layout, bool(use_xcd_remap), ) - launch(A_arg, B_arg, C_arg, M_runtime, N_runtime, stream=stream) + # Preserve the original validated byte-addressed G2L path. These are + # metadata-only dtype/flatten views of the already-contiguous row-major + # tensors selected by gemm_wrappers.py; no transpose or copy is performed. + A_arg = A.view(torch.uint8).view(-1) + B_arg = B.view(torch.uint8).view(-1) + C_arg = C.view(-1) + + 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 index b0629f21a..99a7d5200 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py @@ -38,6 +38,8 @@ def make_bf16_buffer_tensor(arg_bf16): # Keep the exact existing behavior; this is only a symbol alias. make_bf16_byte_buffer_tensor = make_bf16_buffer_tensor +make_fp16_byte_buffer_tensor = make_bf16_byte_buffer_tensor + def compute_global_swizzle( lane_id, @@ -96,6 +98,8 @@ def compute_global_bf16_transpose_swizzle( return offsets +compute_global_fp16_transpose_swizzle = compute_global_bf16_transpose_swizzle + class G2SLoader: """Issue native 16-byte BF16 BufferDesc-to-BF16 LDS copies.""" @@ -221,3 +225,19 @@ def load_one_transpose_bf16( immediate_offset, ) return lo.shuffle(hi, [0, 1, 2, 3]) + + + def load_one_transpose_fp16( + self, + lds_src, + first_byte_offset, + second_byte_offset, + immediate_offset=0, + ): + """Return one i32x4 K32 FP16 fragment from two transpose reads.""" + return self.load_one_transpose_bf16( + lds_src, + first_byte_offset, + second_byte_offset, + immediate_offset, + ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 4f3fef8aa..866b4e354 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -442,6 +442,113 @@ def _run_bf16_gemm( return D + +def _run_fp16_gemm( + A, + transa, + B, + transb, + D, + *, + output_dtype: torch.dtype, +): + """Dispatch FP16 using the original row-major operand allocations. + + No operand transpose view is created: + + TN: kernel A = TE B [M,K], kernel B = TE A [N,K] + NN: kernel A = TE B [M,K], kernel B = TE A [K,N] + NT: kernel A = TE B [K,M], kernel B = TE A [K,N] + """ + 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 torch.float16 inputs, " + 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}" + ) + + dispatch = { + (True, False): "TN", + (False, False): "NN", + (False, True): "NT", + } + try: + layout = dispatch[(bool(transa), bool(transb))] + except KeyError as exc: + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) from exc + + output_shape = _get_gemm_output_shape(A, transa, B, transb) + + # Preserve the original row-major storage. This only collapses leading + # batch dimensions, matching the wrapper's existing regular-GEMM contract. + A_data = _flatten_rowwise(A, "A") + B_data = _flatten_rowwise(B, "B") + + # Kernel ownership is always swapped relative to TE's BLAS arguments. + a_flydsl = B_data + b_flydsl = A_data + + if layout == "TN": + m, k = a_flydsl.shape + n, kb = b_flydsl.shape + expected_a = (m, k) + expected_b = (n, k) + elif layout == "NN": + m, k = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a = (m, k) + expected_b = (k, n) + else: + k, m = a_flydsl.shape + kb, n = b_flydsl.shape + expected_a = (k, m) + expected_b = (k, n) + + if kb != k: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 {layout} received incompatible row-major operands: " + f"a={tuple(a_flydsl.shape)}, b={tuple(b_flydsl.shape)}" + ) + if tuple(a_flydsl.shape) != expected_a or tuple(b_flydsl.shape) != expected_b: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 {layout} physical contract mismatch: " + f"a={tuple(a_flydsl.shape)} expected={expected_a}; " + f"b={tuple(b_flydsl.shape)} expected={expected_b}" + ) + + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL FP16 logical output shape {tuple(output_shape)} " + f"does not match kernel output {(m, n)}" + ) + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=A.device, + backend_name=f"FP16 {layout}", + ) + + fp16_matmul( + a_flydsl, + b_flydsl, + D.view(m, n), + layout=layout, + m=m, + n=n, + k=k, + ) + return D + + def _run_regular_gemm( A, transa, @@ -1251,15 +1358,12 @@ def te_generic_gemm_flydsl( "FlyDSL FP16 supports FP16, BF16, or FP32 output, " f"got {output_dtype}" ) - D = _run_regular_gemm( + D = _run_fp16_gemm( A, transa, B, transb, D, - dtype=torch.float16, - matmul=fp16_matmul, - backend_name="FP16", output_dtype=fp16_output_dtypes[output_dtype], ) return D, None, None, None From 9d6cbfb826780b28c075baeccc610f2a41368945 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 29 Jul 2026 19:16:56 +0000 Subject: [PATCH 28/31] add todo comments --- .../flydsl_kernels/gemm/gemm_wrappers.py | 115 +++++++++++++----- 1 file changed, 85 insertions(+), 30 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 866b4e354..405ba1587 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -19,6 +19,12 @@ from .fp8_gemm import fp8_matmul from .mxfp8_gemm import mxfp8_matmul +# TODO: Some backend-independent GEMM wrapper utilities overlap with the +# Triton GEMM backend (PR#667), including operand classification, logical +# output-shape derivation, and quantized-storage inspection. Once both +# integrations stabilize, factor the genuinely common pieces into a shared +# GEMM wrapper utility module while preserving backend-specific layout and +# storage canonicalization. def _product(shape): """Return the product of dimensions in ``shape``.""" @@ -101,11 +107,13 @@ def _validate_common_epilogue( "FlyDSL GEMM currently supports only alpha=1 and beta=0" ) + # TODO: Add accumulate option if accumulate: raise NotImplementedError( "FlyDSL GEMM accumulation is not implemented" ) + # TODO: Add fused bias and BGRADB epilogues if bias is not None and bias.numel() != 0: raise NotImplementedError( "FlyDSL GEMM bias is not implemented" @@ -549,58 +557,108 @@ def _run_fp16_gemm( return D -def _run_regular_gemm( +def _run_fp32_gemm( A, transa, B, transb, D, - *, - dtype, - matmul, - backend_name, - output_dtype=None, ): - """Run FP16/BF16/FP32 through shared TN/NN/NT shape handling.""" + """Normalize FP32 TN/NN/NT inputs to the current kernel's TN interface. + + The existing FP32 entry point expects ordinary row-major GEMM operands: + + a_tn: [M, K] + b_tn: [K, N] + + TE provides BLAS-shaped operands, so ownership is swapped and only the + operands whose BLAS transpose flags require it are materialized: + + TN: a_tn = B + b_tn = A.T + + NN: a_tn = B + b_tn = A + + NT: a_tn = B.T + b_tn = A + + ``transpose(...).contiguous()`` is therefore used only for the FP32 + operands that are not already in the current TN kernel orientation. + BF16/FP16/FP8/MXFP8 dispatch is unchanged. + """ 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( - f"FlyDSL {backend_name} GEMM expects plain torch.Tensor operands" - ) - if A.dtype != dtype or B.dtype != dtype: - raise TypeError( - f"FlyDSL {backend_name} GEMM requires {dtype} inputs, " + "FlyDSL FP32 GEMM requires torch.float32 inputs, " 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}" ) + if bool(transa) and bool(transb): + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) 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: + A_flat = _flatten_rowwise(A, "A") + B_flat = _flatten_rowwise(B, "B") + + # Standard BLAS-column-major -> row-major conversion: + # swap operands, then apply the original operand transpose flags. + # TODO: Optimize FP32 NN/NT execution. These layouts are currently + # materialized into the TN kernel contract with explicit transpose copies. + if bool(transb): + a_tn = B_flat.transpose(0, 1).contiguous() + else: + a_tn = B_flat + + if bool(transa): + b_tn = A_flat.transpose(0, 1).contiguous() + else: + b_tn = A_flat + + if not a_tn.is_contiguous(): + a_tn = a_tn.contiguous() + if not b_tn.is_contiguous(): + b_tn = b_tn.contiguous() + + if a_tn.ndim != 2 or b_tn.ndim != 2: raise RuntimeError( - f"FlyDSL {backend_name} logical output shape {tuple(output_shape)} " - f"does not match flattened GEMM shape {(m, n)}" + f"FlyDSL FP32 TN normalization produced rank mismatch: " + f"a={tuple(a_tn.shape)}, b={tuple(b_tn.shape)}" + ) + + m, k = a_tn.shape + kb, n = b_tn.shape + if kb != k: + layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 {layout} could not normalize to TN: " + f"a_tn={tuple(a_tn.shape)}, b_tn={tuple(b_tn.shape)}" ) - if output_dtype is None: - output_dtype = dtype + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL FP32 logical output shape {tuple(output_shape)} " + f"does not match normalized TN output {(m, n)}" + ) D = _validate_or_allocate_output( D, shape=output_shape, - dtype=output_dtype, + dtype=torch.float32, device=A.device, - backend_name=backend_name, + backend_name="FP32 via TN core", ) - matmul( - a_flydsl, - b_flydsl, + fp32_matmul( + a_tn, + b_tn, D.view(m, n), ) return D @@ -1374,15 +1432,12 @@ def te_generic_gemm_flydsl( "FlyDSL FP32 currently supports only FP32 output, " f"got {output_dtype}" ) - D = _run_regular_gemm( + D = _run_fp32_gemm( A, transa, B, transb, D, - dtype=torch.float32, - matmul=fp32_matmul, - backend_name="FP32", ) return D, None, None, None @@ -1390,4 +1445,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 + ) From e5f4a0c4a587e5b13b1e3ce0ce308d06a3bf86e6 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 29 Jul 2026 19:20:09 +0000 Subject: [PATCH 29/31] add missing EOLs --- transformer_engine/pytorch/flydsl_kernels/__init__.py | 2 +- transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py | 2 +- transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/__init__.py b/transformer_engine/pytorch/flydsl_kernels/__init__.py index 92fa250e8..c64b988c6 100644 --- a/transformer_engine/pytorch/flydsl_kernels/__init__.py +++ b/transformer_engine/pytorch/flydsl_kernels/__init__.py @@ -1,3 +1,3 @@ from . import gemm -__all__ = ["gemm"] \ No newline at end of file +__all__ = ["gemm"] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py index 5acdce6a2..4eae105ef 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py @@ -10,4 +10,4 @@ __all__ = [ "FlyDSLUnsupportedError", "te_generic_gemm_flydsl", -] \ No newline at end of file +] diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py b/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py index 1ae38569a..b7fc19a23 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py @@ -3,4 +3,4 @@ # 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 + """The GEMM request is valid but unsupported by the available FlyDSL kernels.""" From 0948d9a3e68c26f4af006fd549075e5d1c91594d Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 29 Jul 2026 19:54:06 +0000 Subject: [PATCH 30/31] remove calls to old utility get_tolerances and use dtype_tols instead for flydsl test --- tests/pytorch/test_numerics.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 7685ba1d4..2bc5356e8 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -1434,7 +1434,9 @@ def test_linear_accuracy_flydsl( os.environ.pop("NVTE_FLYDSL_GEMM_WARN_FALLBACK", None) FP8GlobalStateManager.reset() - atol, rtol = get_tolerances(dtype) + tols = dtype_tols(dtype) + atol = tols["atol"] + rtol = tols["rtol"] if fp8: atol = max(atol, 1e-2) From 63c5c4c60b8605b0476bb6ac9b7ea20a9501e37e Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 29 Jul 2026 19:58:23 +0000 Subject: [PATCH 31/31] add gpu arch gating to flyDSL GEMM backend enablement --- transformer_engine/pytorch/cpp_extensions/gemm.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f20f39dd3..0e5b03523 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -494,7 +494,9 @@ def general_gemm( } if not _is_nvfp4_row_scaled_tensor(A) and not _is_nvfp4_row_scaled_tensor(B): - use_gemm_flydsl = IS_HIP_EXTENSION and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0"))) + use_gemm_flydsl = (IS_HIP_EXTENSION + and get_device_compute_capability() == (9, 5) + 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 (