From cf54ce65c7f662ff35423fa2fae8557179731ee3 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 00:33:21 +0000 Subject: [PATCH 01/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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/65] 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 ( From 12d06c712c1c545ea96e0583b1d470a1c264b832 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Mon, 3 Aug 2026 21:12:30 +0000 Subject: [PATCH 32/65] Pin FlyDSL below 0.3 for buffer_ops compatibility --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index d6f7cff46..c12303180 100644 --- a/setup.py +++ b/setup.py @@ -201,7 +201,7 @@ def setup_requirements() -> Tuple[List[str], List[str]]: and "pytorch" in frameworks and bool(int(os.getenv("NVTE_USE_FLYDSL", "0"))) ): - install_reqs.extend(["flydsl"]) + install_reqs.extend(["flydsl>=0.2.4,<0.3"]) # Framework-specific requirements if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): From 10a0813d2f1b5b8ca9147f0c104fb5c483774702 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Fri, 7 Aug 2026 16:56:13 +0000 Subject: [PATCH 33/65] Migrate FlyDSL GEMM kernels off removed buffer_ops API flydsl 0.3.0 removes flydsl.expr.buffer_ops (moved out of FlyDSL), so every GEMM kernel that imported it failed at import. Port all buffer_ops usage to the modern tile/layout copy API and pin flydsl==0.3.0. Conversions: - create_buffer_resource -> make_buffer_tensor - buffer_load -> logical_divide + slice + BufferCopy copy-atom + memref_load_vec - buffer_store -> BufferCopy copy-atom + memref_store_vec + slice Non-obvious details: - Scale-pack kernel receives a 2-D scale tensor, so logical_divide on it walks the layout column-major and a linear slice offset does NOT map to base+off. Rebuild a flat 1-D view (make_view(get_iter(t), make_layout(N,1))) before logical_divide to restore the legacy buffer_load addressing. - The C-store descriptor bias (add_offset on a dynamic Index) segfaults on this build, so fold the per-CTA tile base into each store's linear coordinate instead. - BufferCopy8b lowers to BUFFER_LOAD_UBYTE but is sign-extended when widened; mask with & 0xFF to preserve raw E8M0 scale bytes >= 0x80. - Allocate a fresh rmem register per load/store; a shared register aliases across the many simultaneously-live scales and in-flight epilogue stores. - Pack-kernel unsupported-shape guards now raise FlyDSLUnsupportedError so small shapes fall back to the default backend instead of hard-failing. - Remove dead G2STransposeLoader (its only buffer_ops user in utils). Verified: tests/pytorch/test_numerics.py::test_linear_accuracy_flydsl (48 passed) across small/126m x {None, Float8CurrentScaling, DelayedScaling, MXFP8BlockScaling} x {fp16, bf16, fp32} x bs {1,2}. Co-Authored-By: Claude --- setup.py | 2 +- .../pytorch/flydsl_kernels/gemm/bf16_gemm.py | 28 ++--- .../pytorch/flydsl_kernels/gemm/fp16_gemm.py | 28 ++--- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 27 +++-- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 46 ++++---- .../flydsl_kernels/gemm/fp8_gemm_utils.py | 66 +----------- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 102 ++++++++++-------- 7 files changed, 125 insertions(+), 174 deletions(-) diff --git a/setup.py b/setup.py index c12303180..81ed18830 100644 --- a/setup.py +++ b/setup.py @@ -201,7 +201,7 @@ def setup_requirements() -> Tuple[List[str], List[str]]: and "pytorch" in frameworks and bool(int(os.getenv("NVTE_USE_FLYDSL", "0"))) ): - install_reqs.extend(["flydsl>=0.2.4,<0.3"]) + install_reqs.extend(["flydsl==0.3.0"]) # Framework-specific requirements if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py index cf267d695..a13219eaa 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py @@ -26,7 +26,7 @@ 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 import arith, const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec @@ -516,17 +516,15 @@ def kernel_gemm( 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, + # Per-CTA tile base in C elements, folded into each store's linear + # coordinate below (matching the scale-load addressing on this build; + # add_offset on a dynamic Index is unsupported here). + c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx + gC = fx.rocdl.make_buffer_tensor(C, max_size=True) + c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + c_store_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + output_fx_dtype, ) PIN_ACC_BASE = 0 @@ -795,11 +793,13 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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 + c_idx = c_tile_base_elems + row * fx.Index(c_n) + col 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) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) + fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) # Explicit register coordinates for HK-style four-quadrant mapping. diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py index ace8ecda4..40e6df5be 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py @@ -26,7 +26,7 @@ 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 import arith, const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec @@ -516,17 +516,15 @@ def kernel_gemm( 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, + # Per-CTA tile base in C elements, folded into each store's linear + # coordinate below (matching the scale-load addressing on this build; + # add_offset on a dynamic Index is unsupported here). + c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx + gC = fx.rocdl.make_buffer_tensor(C, max_size=True) + c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + c_store_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + output_fx_dtype, ) PIN_ACC_BASE = 0 @@ -795,11 +793,13 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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 + c_idx = c_tile_base_elems + row * fx.Index(c_n) + col 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) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) + fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) # Explicit register coordinates for HK-style four-quadrant mapping. diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index 6cbd61102..aa0e2c041 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -21,7 +21,7 @@ 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 import arith, const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec @@ -306,18 +306,13 @@ def kernel_gemm( 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, - ) + # Per-CTA tile base in C elements, folded into each store's linear + # coordinate below (matching the scale-load addressing on this build; + # add_offset on a dynamic Index is unsupported here). + c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx + gC = fx.rocdl.make_buffer_tensor(C, max_size=True) + c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + c_store_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) PIN_ACC_BASE = 0 @@ -540,8 +535,10 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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) + c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.memref_store_vec(Vec.filled(1, Vec(acc)[ii], fx.Float32), reg) + fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) # Explicit register coordinates for HK-style four-quadrant mapping. diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 45c7ad66e..f524daf3d 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -21,7 +21,7 @@ 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 import arith, const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec @@ -295,12 +295,18 @@ def kernel_gemm( 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) - ) + a_scale_rsrc = fx.rocdl.make_buffer_tensor(A_scale_inv, max_size=True) + b_scale_rsrc = fx.rocdl.make_buffer_tensor(B_scale_inv, max_size=True) + a_scale_div = fx.logical_divide(a_scale_rsrc, fx.make_layout(1, 1)) + b_scale_div = fx.logical_divide(b_scale_rsrc, fx.make_layout(1, 1)) + scale_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + + def _load_scale(scale_div): + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.copy(scale_ld_atom, fx.slice(scale_div, (None, fx.Int32(0))), reg) + return fx.memref_load_vec(reg)[0] + + output_scale = _load_scale(a_scale_div) * _load_scale(b_scale_div) tx = gpu.thread_id("x") num_blocks_m = c_m // BLOCK_M @@ -341,17 +347,15 @@ def kernel_gemm( 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, + # Per-CTA tile base in C elements, folded into each store's linear + # coordinate below (matching the scale-load addressing on this build; + # add_offset on a dynamic Index is unsupported here). + c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx + gC = fx.rocdl.make_buffer_tensor(C, max_size=True) + c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + c_store_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + output_fx_dtype, ) PIN_ACC_BASE = 0 @@ -549,11 +553,13 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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 + c_idx = c_tile_base_elems + 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) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) + fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) # Explicit register coordinates for HK-style four-quadrant mapping. 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 2c0e0534b..87e80fa5b 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,7 @@ 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, buffer_ops, const_expr, range_constexpr, rocdl +from flydsl.expr import arith, 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 @@ -120,70 +120,6 @@ 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))) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 76e82b5c5..c76fdf2f9 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -29,7 +29,7 @@ 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 import arith, gpu, range_constexpr, rocdl from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec @@ -83,18 +83,18 @@ def _compile_mx32_scale_pack_kernel( and no eager PyTorch shift/index/OR kernels. """ if dim % 64 != 0: - raise ValueError( + raise FlyDSLUnsupportedError( f"Scale outer dimension={dim} must be a multiple of 64" ) if qk % 4 != 0: - raise ValueError( + raise FlyDSLUnsupportedError( 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( + raise FlyDSLUnsupportedError( f"Packed scale words={total_words} must be divisible by " f"{_SCALE_PACK_THREADS}" ) @@ -119,8 +119,22 @@ def _source_offset(source_k32, source_row): @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) + # Address both buffers as flat 1-D arrays so a linear slice coordinate + # maps to base+off, matching the legacy buffer_load semantics. Building + # buffer tensors directly from the 2-D src/dst would make logical_divide + # walk the layout column-major and corrupt every strided access. + src_flat = fx.rocdl.make_buffer_tensor( + fx.Tensor(fx.make_view(fx.get_iter(src), fx.make_layout(dim * qk, 1))), + max_size=True, + ) + dst_flat = fx.rocdl.make_buffer_tensor( + fx.Tensor(fx.make_view(fx.get_iter(dst), fx.make_layout(total_words, 1))), + max_size=True, + ) + src_div = fx.logical_divide(src_flat, fx.make_layout(1, 1)) + dst_div = fx.logical_divide(dst_flat, fx.make_layout(1, 1)) + scale_load_atom = fx.make_copy_atom(fx.rocdl.BufferCopy8b(), fx.Uint8) + scale_store_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) linear = ( fx.Index(fx.block_idx.x) * fx.Index(_SCALE_PACK_THREADS) @@ -140,22 +154,23 @@ def load_scale_byte(group): + 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)) + off = _source_offset(source_k32, source_row) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Uint8) + fx.copy(scale_load_atom, fx.slice(src_div, (None, fx.Int32(off))), reg) + value_u8 = fx.memref_load_vec(reg)[0] + # The byte load widens via sext, so mask to the low 8 bits to + # preserve the raw E8M0 byte for scales >= 0x80 before packing. + return fx.Int32(value_u8) & fx.Int32(0xFF) 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) + reg_i32 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) + fx.memref_store_vec(Vec.filled(1, packed, fx.Int32), reg_i32) + fx.copy(scale_store_atom, reg_i32, fx.slice(dst_div, (None, fx.Int32(linear)))) + @flyc.jit def launch_pack_mx32_scales( @@ -598,8 +613,11 @@ def kernel_gemm( 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) + as_rsrc = fx.rocdl.make_buffer_tensor(As, max_size=True) + bs_rsrc = fx.rocdl.make_buffer_tensor(Bs, max_size=True) + as_div = fx.logical_divide(as_rsrc, fx.make_layout(1, 1)) + bs_div = fx.logical_divide(bs_rsrc, fx.make_layout(1, 1)) + scale_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) tx = gpu.thread_id("x") num_blocks_m = c_m // BLOCK_M @@ -661,17 +679,15 @@ def kernel_gemm( 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, + # Per-CTA tile base in C elements, folded into each store's linear + # coordinate below (matching the scale-load addressing on this build; + # add_offset on a dynamic Index is unsupported here). + c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx + gC = fx.rocdl.make_buffer_tensor(C, max_size=True) + c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + c_store_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + output_fx_dtype, ) PIN_ACC_BASE = 0 @@ -773,22 +789,16 @@ def hot_loop_scheduler_q_prefetch_4n(): 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 + off = k128 * c_m_idx + bx_m_idx + row + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) + fx.copy(scale_ld_atom, fx.slice(as_div, (None, fx.Int32(off))), reg) + return fx.memref_load_vec(reg)[0] 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 + off = k128 * c_n_idx + by_n_idx + row + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) + fx.copy(scale_ld_atom, fx.slice(bs_div, (None, fx.Int32(off))), reg) + return fx.memref_load_vec(reg)[0] def load_a_scale_subtile(k128, sm): subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) @@ -974,11 +984,13 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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 + c_idx = c_tile_base_elems + 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) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) + fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) # Explicit register coordinates for HK-style four-quadrant mapping. From fe3fe6b3e9bd1cb1089a3903d3da9f0087bc71e4 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Fri, 7 Aug 2026 17:26:35 +0000 Subject: [PATCH 34/65] Prefer system ROCm tree for ROCM_PATH in _rocm_init _rocm_init (added by #602 for the TheRock CI migration) set ROCM_PATH to the rocm-sdk devel wheel (get_devel_root()) when unset. FlyDSL 0.3.0's MLIR AMDGPU serialization resolves ld.lld relative to ROCM_PATH and cannot drive the devel-wheel layout, so fresh kernel compiles fail with "lld invocation failed" unless the user manually exports ROCM_PATH. Prefer the system /opt/rocm tree when present (the CI container provides it, and FlyDSL's linker resolves it), falling back to the devel wheel only for wheel-only environments with no system tree. This matches the Dockerfile, which builds with ROCM_PATH=/opt/rocm. Verified: after a clean rebuild with the flydsl cache cleared, the MXFP8 GEMM test links and passes with no ROCM_PATH export. Co-Authored-By: Claude --- build_tools/templates/_rocm_init.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/build_tools/templates/_rocm_init.py b/build_tools/templates/_rocm_init.py index aad9a4336..50c13b454 100644 --- a/build_tools/templates/_rocm_init.py +++ b/build_tools/templates/_rocm_init.py @@ -30,5 +30,12 @@ def initialize() -> None: return if not os.getenv("ROCM_PATH"): - os.environ["ROCM_PATH"] = str(get_devel_root()) + # Prefer the system ROCm tree when present: FlyDSL's MLIR linker + # resolution expects that layout to locate ld.lld. Fall back to the + # rocm-sdk devel wheel for wheel-only environments with no system tree. + _system_rocm = "/opt/rocm" + if os.path.exists(_system_rocm): + os.environ["ROCM_PATH"] = _system_rocm + else: + os.environ["ROCM_PATH"] = str(get_devel_root()) rocm_sdk.initialize_process(preload_shortnames=list(_PRELOAD_LIBS)) From 9aa462144b1c66bc4ae67a78f757fd610637579d Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Mon, 10 Aug 2026 16:08:24 +0000 Subject: [PATCH 35/65] Deduplicate FlyDSL GEMM helpers into gemm_common_utils The dtype-independent GEMM primitives (cdiv/ceildiv, divmod, min, encode_waitcnt, barrier, xcd_swizzle, swizzle_128, pack_i32x4_i32x8) were duplicated across fp16_gemm_utils.py and fp8_gemm_utils.py, and inlined again in each of the five GEMM kernel modules. Consolidate them into a single gemm_common_utils.py source of truth; the per-dtype utils modules re-export them so the kernel files import the same names unchanged. Dtype-specific code (compute_global_swizzle, the G2S/S2R loaders, buffer-tensor makers) stays in the per-dtype modules. This also fixes the ImportError where fp16/bf16/fp32_gemm imported xcd_swizzle and barrier from fp16_gemm_utils, which never defined them. Co-Authored-By: Claude --- .../pytorch/flydsl_kernels/gemm/bf16_gemm.py | 119 ++--------------- .../pytorch/flydsl_kernels/gemm/fp16_gemm.py | 119 ++--------------- .../flydsl_kernels/gemm/fp16_gemm_utils.py | 35 ++--- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 125 ++--------------- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 118 ++-------------- .../flydsl_kernels/gemm/fp8_gemm_utils.py | 55 +++----- .../flydsl_kernels/gemm/gemm_common_utils.py | 126 ++++++++++++++++++ .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 115 +++------------- 8 files changed, 233 insertions(+), 579 deletions(-) create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py index a13219eaa..511efced9 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py @@ -40,6 +40,8 @@ make_bf16_byte_buffer_tensor, pack_i32x4_i32x8, swizzle_128, + xcd_swizzle, + barrier ) @@ -97,99 +99,6 @@ 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, output_dtype: torch.dtype, @@ -469,7 +378,7 @@ def kernel_gemm( 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) + 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) @@ -880,7 +789,7 @@ def hk_one_k_with_refill( ): # 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) + 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, @@ -939,7 +848,7 @@ def hk_one_k_with_refill( # 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) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = pack_frag_halves(a10_x0, a10_x1) @@ -976,7 +885,7 @@ def hk_one_k_with_refill( # 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) + 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) @@ -1011,7 +920,7 @@ def hk_one_k_with_refill( 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) + 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 @@ -1027,7 +936,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = load_a_subtile_mi_regs(cur_a, 1, 0) @@ -1041,7 +950,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): 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) + 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) @@ -1076,7 +985,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): 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) + barrier(vmcnt=0, lgkmcnt=0) a00, a01, a02, a03 = a0_regs b00, b01, b02, b03 = b0_regs @@ -1089,7 +998,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): b13 = load_b_subtile_ni_regs(cur_b, 1, 3) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = load_a_subtile_mi_regs(cur_a, 1, 0) @@ -1098,7 +1007,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): a13 = load_a_subtile_mi_regs(cur_a, 1, 3) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) @@ -1164,13 +1073,13 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): 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) + 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) + 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) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py index 40e6df5be..1f2b72a7e 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py @@ -40,6 +40,8 @@ make_fp16_byte_buffer_tensor, pack_i32x4_i32x8, swizzle_128, + xcd_swizzle, + barrier ) @@ -97,99 +99,6 @@ 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, output_dtype: torch.dtype, @@ -469,7 +378,7 @@ def kernel_gemm( 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) + 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) @@ -880,7 +789,7 @@ def hk_one_k_with_refill( ): # 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) + 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, @@ -939,7 +848,7 @@ def hk_one_k_with_refill( # 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) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = pack_frag_halves(a10_x0, a10_x1) @@ -976,7 +885,7 @@ def hk_one_k_with_refill( # 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) + 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) @@ -1011,7 +920,7 @@ def hk_one_k_with_refill( 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) + 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 @@ -1027,7 +936,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = load_a_subtile_mi_regs(cur_a, 1, 0) @@ -1041,7 +950,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): 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) + 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) @@ -1076,7 +985,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): 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) + barrier(vmcnt=0, lgkmcnt=0) a00, a01, a02, a03 = a0_regs b00, b01, b02, b03 = b0_regs @@ -1089,7 +998,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): b13 = load_b_subtile_ni_regs(cur_b, 1, 3) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = load_a_subtile_mi_regs(cur_a, 1, 0) @@ -1098,7 +1007,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): a13 = load_a_subtile_mi_regs(cur_a, 1, 3) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) @@ -1164,13 +1073,13 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): 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) + 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) + 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) 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 99a7d5200..4b3e4a4bb 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py @@ -9,24 +9,19 @@ from flydsl.expr.typing import Vector as Vec from flydsl.expr.utils.arith import _to_raw as as_mlir_value - -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 +# Dtype-independent primitives live in the shared module; re-export them so the +# GEMM kernels can keep importing them from this per-dtype module unchanged. +from .gemm_common_utils import ( + barrier, + cdiv, + ceildiv, + divmod, + encode_waitcnt, + min, + pack_i32x4_i32x8, + swizzle_128, + xcd_swizzle, +) def make_bf16_buffer_tensor(arg_bf16): @@ -147,10 +142,6 @@ def load_one(self, lds_dst, byte_offset, step): ) -def pack_i32x4_i32x8(lo, hi): - return lo.shuffle(hi, list(range(8))) - - class S2RLoader: """LDS readers used to assemble BF16 K64 fragments.""" diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index aa0e2c041..954de30ee 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -34,6 +34,8 @@ make_bf16_byte_buffer_tensor as make_fp32_byte_buffer_tensor, pack_i32x4_i32x8, swizzle_128, + xcd_swizzle, + barrier ) @@ -90,105 +92,6 @@ 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``. @@ -272,7 +175,7 @@ def kernel_gemm( 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) + 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) @@ -608,7 +511,7 @@ def hk_one_k_with_refill( ): # 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) + 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, @@ -667,7 +570,7 @@ def hk_one_k_with_refill( # 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) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = pack_frag_halves(a10_x0, a10_x1) @@ -704,7 +607,7 @@ def hk_one_k_with_refill( # 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) + 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) @@ -739,7 +642,7 @@ def hk_one_k_with_refill( 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) + 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 @@ -755,7 +658,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = load_a_subtile_mi_regs(cur_a, 1, 0) @@ -769,7 +672,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): 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) + 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) @@ -804,7 +707,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): 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) + barrier(vmcnt=0, lgkmcnt=0) a00, a01, a02, a03 = a0_regs b00, b01, b02, b03 = b0_regs @@ -817,7 +720,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): b13 = load_b_subtile_ni_regs(cur_b, 1, 3) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = load_a_subtile_mi_regs(cur_a, 1, 0) @@ -826,7 +729,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): a13 = load_a_subtile_mi_regs(cur_a, 1, 3) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) @@ -892,13 +795,13 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): 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) + 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) + 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) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index f524daf3d..b28576fc0 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -35,6 +35,8 @@ make_fp8_buffer_tensor, pack_i32x4_i32x8, swizzle_128, + xcd_swizzle, + barrier ) @@ -92,98 +94,6 @@ 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, @@ -313,7 +223,7 @@ def _load_scale(scale_div): 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) + 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) @@ -629,7 +539,7 @@ def hk_one_k_with_refill( ): # 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) + 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, @@ -687,7 +597,7 @@ def hk_one_k_with_refill( # 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) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = pack_frag_halves(a10_x0, a10_x1) @@ -724,7 +634,7 @@ def hk_one_k_with_refill( # 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) + 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) @@ -759,7 +669,7 @@ def hk_one_k_with_refill( 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) + 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 @@ -775,7 +685,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = load_a_subtile_mi_regs(cur_a, 1, 0) @@ -789,7 +699,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): 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) + 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) @@ -824,7 +734,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): 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) + barrier(vmcnt=0, lgkmcnt=0) a00, a01, a02, a03 = a0_regs b00, b01, b02, b03 = b0_regs @@ -837,7 +747,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): b13 = load_b_subtile_ni_regs(cur_b, 1, 3) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = load_a_subtile_mi_regs(cur_a, 1, 0) @@ -846,7 +756,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): a13 = load_a_subtile_mi_regs(cur_a, 1, 3) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) @@ -912,13 +822,13 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): 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) + 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) + 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) 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 87e80fa5b..453ee3a3a 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py @@ -10,21 +10,19 @@ 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): - """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) - +# Dtype-independent primitives live in the shared module; re-export them so the +# GEMM kernels can keep importing them from this per-dtype module unchanged. +from .gemm_common_utils import ( + barrier, + cdiv, + ceildiv, + divmod, + encode_waitcnt, + min, + pack_i32x4_i32x8, + swizzle_128, + xcd_swizzle, +) def preshuffle_b(b_t): """Permute row-major ``B_T`` ``(N, K)`` for ``b_preshuffled=True``.""" @@ -49,13 +47,11 @@ def make_fp8_buffer_tensor(arg_i8, fp8_ir_t): 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 - - +# Returns, for one lane (lane_id, wave_id), a list of n_rounds swizzled flat +# global offsets indexed by DMA pass: offsets[step] = r*K + c where (r,c) = +# swizzle_128(row, col). Each is the static per-thread/per-pass source of one +# 16-byte load; the dynamic K-tile base is added later as soffset. K is the +# global row stride (a_leading_dim / b_leading_dim), not the 128 tile width. def compute_global_swizzle(lane_id, wave_id, K, n_rounds, preshuffled): offsets = [] n_waves = fx.block_dim.x // 64 @@ -120,11 +116,6 @@ def load_one(self, lds_dst, k_offset, 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 @@ -291,16 +282,6 @@ def store(self, c_frag, base_row, base_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)) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py new file mode 100644 index 000000000..287fb2dab --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors +"""Dtype-independent primitives shared by the FlyDSL GEMM kernels. + +These helpers carry no dtype specialization, so the per-dtype utils modules +(``fp16_gemm_utils`` / ``fp8_gemm_utils``) re-export them instead of keeping +private copies. Anything that varies by element type (global-swizzle layouts, +G2S/S2R loaders, buffer-tensor makers) stays in the per-dtype modules. +""" + +import flydsl.expr as fx +from flydsl.expr import arith, rocdl + + +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 min(a, b): + return arith.select(a < b, a, b) + + +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 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), + ) + + +# XOR swizzle over a 128-wide tile: col ^= ((row//2) % 8) * 16, row unchanged. +# Equivalent to CuTe Swizzle (vec=16, perPhase=2, maxPhase=8). +# Callers must keep col < 128 so the XOR stays within the row. Self-inverse, so +# the same call is used on both the LDS store (via compute_global_swizzle) and +# the MFMA-feeding LDS read (S2RLoader). ``col_in_bytes`` is a byte coordinate. +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 pack_i32x4_i32x8(lo, hi): + # Pack two i32x4 as one i32x8 + return lo.shuffle(hi, list(range(8))) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index c76fdf2f9..965bc3de6 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -42,6 +42,8 @@ make_fp8_buffer_tensor, pack_i32x4_i32x8, swizzle_128, + xcd_swizzle, + barrier ) @@ -76,7 +78,7 @@ def _compile_mx32_scale_pack_kernel( stride0: int, stride1: int, ): - """Build one fused raw-E8M0 -> HK-scale packing kernel. + """Build one fused 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 @@ -269,88 +271,6 @@ def pack_mx32_scales_for_hk( ) return packed -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, @@ -623,7 +543,7 @@ def kernel_gemm( 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) + pid_m, pid_n = xcd_swizzle(num_blocks_m, num_blocks_n) bx_m = pid_m * BLOCK_M by_n = pid_n * BLOCK_N @@ -644,6 +564,11 @@ def kernel_gemm( a_leading_dim = _a_leading_dim(c_m) b_leading_dim = _b_leading_dim(c_n) + # gl_off_a/gl_off_b: per-lane lists of LOAD_PASSES_HALF swizzled flat + # global offsets, indexed by DMA pass. gl_off_a[step] is this lane's + # static 16-byte source for pass `step`; G2SLoader adds the dynamic + # K-tile base as soffset. Offsets are pre-swizzled so bytes land in the + # bank-conflict-free LDS slots the MFMA read (S2RLoader) expects. gl_off_a = compute_global_swizzle( lane, wave_id, @@ -1083,7 +1008,7 @@ def hk_one_k_with_refill( # 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) + 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. @@ -1146,7 +1071,7 @@ def hk_one_k_with_refill( # 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) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10 = pack_frag_halves(a10_x0, a10_x1) @@ -1187,7 +1112,7 @@ def hk_one_k_with_refill( # 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) + 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) @@ -1240,7 +1165,7 @@ def hk_one_k_with_refill( 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) + 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 @@ -1256,7 +1181,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs, cur_ mfma_4n(_acc_idx(0, 3, 0), a03, as03, b00, b01, b02, b03, bs00, bs01, bs02, bs03) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) @@ -1270,7 +1195,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs, cur_ 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) + barrier(vmcnt=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) @@ -1323,7 +1248,7 @@ def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs, cur_ 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) + 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 @@ -1336,7 +1261,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) @@ -1345,7 +1270,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) rocdl.sched_barrier(0) - _barrier(lgkmcnt=0) + barrier(lgkmcnt=0) rocdl.sched_barrier(0) a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) @@ -1424,7 +1349,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): 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) + 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. @@ -1438,7 +1363,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): 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) + 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. From 66633b948c0b0a19efab1efcaf4b94834a10ec4e Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Mon, 10 Aug 2026 16:40:31 +0000 Subject: [PATCH 36/65] Extract shared M/N/K block-tiling guard into require_block_tiling The host-side shape validation in every doGemm/do_gemm (M/N/K multiple-of-BLOCK checks plus the minimum-4-K-tiles guard) was duplicated near-identically across all five GEMM kernel modules, differing only by the dtype label and mxfp8's layout interpolation. Consolidate it into a single require_block_tiling helper in gemm_common_utils.py; the error messages are reproduced verbatim. This validation runs before kernel launch and does not touch the hand-tuned kernel_gemm bodies. Co-Authored-By: Claude --- .../pytorch/flydsl_kernels/gemm/bf16_gemm.py | 32 ++++++----------- .../pytorch/flydsl_kernels/gemm/fp16_gemm.py | 32 ++++++----------- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 35 ++++++------------- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 32 ++++++----------- .../flydsl_kernels/gemm/gemm_common_utils.py | 31 ++++++++++++++++ .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 31 ++++++---------- 6 files changed, 81 insertions(+), 112 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py index 511efced9..486a84f7c 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py @@ -32,6 +32,7 @@ # Transformer Engine-local FlyDSL utilities. from .exceptions import FlyDSLUnsupportedError +from .gemm_common_utils import require_block_tiling from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -1285,28 +1286,15 @@ def doGemm( 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}, " - f"got K={K_runtime}" - ) - - num_k_tiles = K_runtime // _BLOCK_K - 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)" - ) + require_block_tiling( + M_runtime, + N_runtime, + K_runtime, + block_m=_BLOCK_M, + block_n=_BLOCK_N, + block_k=_BLOCK_K, + label="BF16 GEMM", + ) if tuple(C.shape) != (M_runtime, N_runtime): raise ValueError( diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py index 1f2b72a7e..315550c5a 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py @@ -32,6 +32,7 @@ # Transformer Engine-local FlyDSL utilities. from .exceptions import FlyDSLUnsupportedError +from .gemm_common_utils import require_block_tiling from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -1285,28 +1286,15 @@ def doGemm( 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}, " - f"got K={K_runtime}" - ) - - num_k_tiles = K_runtime // _BLOCK_K - 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)" - ) + require_block_tiling( + M_runtime, + N_runtime, + K_runtime, + block_m=_BLOCK_M, + block_n=_BLOCK_N, + block_k=_BLOCK_K, + label="FP16 GEMM", + ) if tuple(C.shape) != (M_runtime, N_runtime): raise ValueError( diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index 954de30ee..712ca567b 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -26,7 +26,7 @@ from flydsl.expr.typing import Vector as Vec # Transformer Engine-local FlyDSL utilities. -from .exceptions import FlyDSLUnsupportedError +from .gemm_common_utils import require_block_tiling from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -956,30 +956,15 @@ 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 - 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 - 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)" - ) + require_block_tiling( + M_runtime, + N_runtime, + K_runtime, + block_m=_BLOCK_M, + block_n=_BLOCK_N, + block_k=_BLOCK_K, + label="FP32 GEMM", + ) 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 b28576fc0..1fbd66f22 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -25,7 +25,7 @@ from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec -from .exceptions import FlyDSLUnsupportedError +from .gemm_common_utils import require_block_tiling # Transformer Engine-local FlyDSL utilities. from .fp8_gemm_utils import ( @@ -1040,27 +1040,15 @@ def doGemm( 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)" - ) + require_block_tiling( + M_runtime, + N_runtime, + K_runtime, + block_m=_BLOCK_M, + block_n=_BLOCK_N, + block_k=_BLOCK_K, + label="FP8 GEMM", + ) 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/gemm_common_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py index 287fb2dab..0b3e42ce6 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py @@ -11,6 +11,37 @@ import flydsl.expr as fx from flydsl.expr import arith, rocdl +from .exceptions import FlyDSLUnsupportedError + + +def require_block_tiling(m, n, k, *, block_m, block_n, block_k, label, min_k_tiles=4): + """Validate the host-side M/N/K tiling contract shared by every GEMM core. + + ``label`` is the human-readable kernel identifier used in the error text + (e.g. ``"FP16 GEMM"`` or ``"MXFP8 TN GEMM"``). Raises + ``FlyDSLUnsupportedError`` so callers fall back to the default backend on + unsupported shapes. Returns the number of K tiles. + """ + if m % block_m != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL {label} requires M to be a multiple of {block_m}, got M={m}" + ) + if n % block_n != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL {label} requires N to be a multiple of {block_n}, got N={n}" + ) + if k % block_k != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL {label} requires K to be a multiple of {block_k}, got K={k}" + ) + num_k_tiles = k // block_k + if num_k_tiles < min_k_tiles: + raise FlyDSLUnsupportedError( + f"FlyDSL {label} requires at least {min_k_tiles} K{block_k} tiles, " + f"got K={k} ({num_k_tiles} tiles)" + ) + return num_k_tiles + def cdiv(numer: int, denom: int) -> int: return (numer + denom - 1) // denom diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 965bc3de6..050cc4fe8 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -35,6 +35,7 @@ # Transformer Engine-local FlyDSL utilities. from .exceptions import FlyDSLUnsupportedError +from .gemm_common_utils import require_block_tiling from .fp8_gemm_utils import ( G2SLoader, S2RLoader, @@ -1492,27 +1493,15 @@ def do_gemm( 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 {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 {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 {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 {layout} GEMM requires at least 4 K{_BLOCK_K} " - f"tiles, got K={K_runtime} ({num_k_tiles} tiles)" - ) + require_block_tiling( + M_runtime, + N_runtime, + K_runtime, + block_m=_BLOCK_M, + block_n=_BLOCK_N, + block_k=_BLOCK_K, + label=f"MXFP8 {layout} GEMM", + ) expected_as = (K_runtime // _BLOCK_K, M_runtime) expected_bs = (K_runtime // _BLOCK_K, N_runtime) From 91412fd9c896a009106d2e8575b641ac5d8c512c Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Mon, 10 Aug 2026 17:11:26 +0000 Subject: [PATCH 37/65] Fuse FP16 and BF16 GEMM kernels into half_prec_gemm fp16_gemm.py and bf16_gemm.py were byte-identical apart from the MFMA opcode (v_mfma_f32_16x16x32_f16 vs _bf16), the operand dtype guards, and naming. That opcode is a compile-time f-string, so the two kernels can share one generator without any runtime dtype branch. Merge both into half_prec_gemm.py: _compile_kernel takes an mfma_suffix argument threaded into the opcode string, and _cached_launch keys on it so FP16 and BF16 still compile to separate cached binaries. A shared _half_prec_matmul core parametrized by input_dtype/label backs the thin fp16_matmul/bf16_matmul wrappers, whose public signatures are unchanged. gemm_wrappers.py imports both from the new module; fp16_gemm.py and bf16_gemm.py are deleted. Also drop the dead LDS_SYM_*/LDS_ALIAS_DOMAIN/SCOPE_IDS constants carried over from the copy. Genericize the shared byte-staging helpers in fp16_gemm_utils.py that these kernels (and fp32) reuse: make_bf16_buffer_tensor -> make_byte_buffer_tensor, compute_global_bf16_transpose_swizzle -> compute_global_transpose_swizzle, and S2RLoader.load_one_transpose_bf16 -> load_one_transpose, dropping the redundant fp16/bf16 aliases. These helpers operate on flat byte views and are dtype-independent, so the dtype labels were misleading. The fused _compile_kernel body is source-identical to the originals after the mechanical renames, so codegen is unchanged. Verified imports/compile only; GPU numerics still to be run. Co-Authored-By: Claude --- .../pytorch/flydsl_kernels/gemm/bf16_gemm.py | 1327 ----------------- .../flydsl_kernels/gemm/fp16_gemm_utils.py | 57 +- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 2 +- .../flydsl_kernels/gemm/gemm_wrappers.py | 3 +- .../gemm/{fp16_gemm.py => half_prec_gemm.py} | 189 ++- 5 files changed, 154 insertions(+), 1424 deletions(-) delete mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py rename transformer_engine/pytorch/flydsl_kernels/gemm/{fp16_gemm.py => half_prec_gemm.py} (90%) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py deleted file mode 100644 index 486a84f7c..000000000 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py +++ /dev/null @@ -1,1327 +0,0 @@ -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. - -"""FlyDSL BF16 TN/NN/NT 4-wave GEMM kernel for Transformer Engine. - -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. -""" - -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, 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 .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import require_block_tiling -from .fp16_gemm_utils import ( - G2SLoader, - S2RLoader, - compute_global_bf16_transpose_swizzle, - compute_global_swizzle, - make_bf16_byte_buffer_tensor, - pack_i32x4_i32x8, - swizzle_128, - xcd_swizzle, - barrier -) - - -_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 _compile_kernel( - K: int, - output_dtype: torch.dtype, - layout: str, - use_xcd_remap: bool = True, -): - """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 - - 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 - - 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 - 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 - - # 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: - # 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] - 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 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)) - 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) - - # 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)) - 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) - - # Per-CTA tile base in C elements, folded into each store's linear - # coordinate below (matching the scale-load addressing on this build; - # add_offset on a dynamic Index is unsupported here). - c_tile_base_elems = bx_m_idx * fx.Index(c_n) + by_n_idx - gC = fx.rocdl.make_buffer_tensor(C, max_size=True) - c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) - c_store_atom = fx.make_copy_atom( - fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), - output_fx_dtype, - ) - - 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(Q0_SCHED_DSRD) - 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(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 = _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 = _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): - 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 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 - - 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 = c_tile_base_elems + row * fx.Index(c_n) + col - value = Vec(acc)[ii] - if const_expr(output_dtype != torch.float32): - value = value.to(output_fx_dtype) - reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) - fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) - fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(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): - 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 ( - 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): - 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) - 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, - output_dtype: torch.dtype, - layout: str, - use_xcd_remap: bool = True, -): - return _compile_kernel( - K, - output_dtype, - layout, - use_xcd_remap=use_xcd_remap, - ) - - -def bf16_matmul( - a: torch.Tensor, - b: torch.Tensor, - c: torch.Tensor, - *, - layout: str, - m: int, - n: int, - k: int, - stream=None, -): - """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 expects rank-2 operands, got A{tuple(a.shape)} " - f"and B{tuple(b.shape)}" - ) - if a.dtype != torch.bfloat16 or b.dtype != torch.bfloat16: - raise TypeError( - "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): - raise TypeError( - "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 " - f"{a.device}, {b.device}, and {c.device}" - ) - if not c.is_contiguous(): - raise ValueError("FlyDSL BF16 GEMM requires contiguous output storage") - - 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 one cached K/output/layout-specialized BF16 core. - - 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. - """ - 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}") - - require_block_tiling( - M_runtime, - N_runtime, - K_runtime, - block_m=_BLOCK_M, - block_n=_BLOCK_N, - block_k=_BLOCK_K, - label="BF16 GEMM", - ) - - 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() - - launch = _cached_launch( - K_runtime, - C.dtype, - layout, - bool(use_xcd_remap), - ) - # 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 4b3e4a4bb..d6cc8f221 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py @@ -1,6 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors -"""Byte-staging helpers for the BF16 four-wave GEMM.""" +"""Byte-level staging helpers for the four-wave GEMM kernels. + +These loaders and swizzle helpers operate on flat byte views and carry no dtype +label, so the half-precision (FP16/BF16) and FP32 cores all share them; each +core selects its own MFMA opcode. The module keeps its historical +``fp16_gemm_utils`` filename. +""" import flydsl.expr as fx from flydsl._mlir import ir @@ -24,16 +30,13 @@ ) -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 make_byte_buffer_tensor(arg): + """Create a BufferDesc directly from the wrapper-provided tensor. -make_fp16_byte_buffer_tensor = make_bf16_byte_buffer_tensor + Dtype-independent: the operand is consumed as a flat byte view, so FP16, + BF16, and FP32 all share this maker. + """ + return fx.rocdl.make_buffer_tensor(arg, max_size=False) def compute_global_swizzle( @@ -48,7 +51,7 @@ def compute_global_swizzle( for round in range_constexpr(n_rounds): if const_expr(preshuffled): raise AssertionError( - "BF16 first-pass port does not support preshuffled operands" + "16-bit 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 @@ -57,13 +60,13 @@ def compute_global_swizzle( return offsets -def compute_global_bf16_transpose_swizzle( +def compute_global_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``. + """Offsets for a K-major 16-bit source staged for ``ds_read_b64_tr_b16``. 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 @@ -71,7 +74,7 @@ def compute_global_bf16_transpose_swizzle( 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. + ``source[k, x_base]`` for a contiguous K-major 16-bit matrix. """ offsets = [] n_waves = fx.block_dim.x // 64 @@ -93,10 +96,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.""" + """Issue native 16-byte BufferDesc-to-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) @@ -143,7 +144,7 @@ def load_one(self, lds_dst, byte_offset, step): class S2RLoader: - """LDS readers used to assemble BF16 K64 fragments.""" + """LDS readers used to assemble 16-bit K64 fragments.""" def __init__(self, wave_idx, n_tiles): self.lane_id = fx.thread_idx.x % 64 @@ -197,14 +198,14 @@ def _ds_read_b64_tr_b16( fx.Int32, ) - def load_one_transpose_bf16( + def load_one_transpose( self, lds_src, first_byte_offset, second_byte_offset, immediate_offset=0, ): - """Return one i32x4 K32 BF16 fragment from two transpose reads.""" + """Return one i32x4 K32 16-bit fragment from two transpose reads.""" lo = self._ds_read_b64_tr_b16( lds_src, first_byte_offset, @@ -216,19 +217,3 @@ 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/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index 712ca567b..e003947ef 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -31,7 +31,7 @@ G2SLoader, S2RLoader, compute_global_swizzle, - make_bf16_byte_buffer_tensor as make_fp32_byte_buffer_tensor, + make_byte_buffer_tensor as make_fp32_byte_buffer_tensor, pack_i32x4_i32x8, swizzle_128, xcd_swizzle, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 405ba1587..4cff94af5 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -13,8 +13,7 @@ from .exceptions import FlyDSLUnsupportedError -from .bf16_gemm import bf16_matmul -from .fp16_gemm import fp16_matmul +from .half_prec_gemm import bf16_matmul, fp16_matmul from .fp32_gemm import fp32_matmul from .fp8_gemm import fp8_matmul from .mxfp8_gemm import mxfp8_matmul diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py similarity index 90% rename from transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py rename to transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 315550c5a..444ffdf6e 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -2,10 +2,16 @@ # # See LICENSE for license information. -"""FlyDSL FP16 TN/NN/NT 4-wave GEMM kernel for Transformer Engine. +"""FlyDSL half-precision (FP16/BF16) TN/NN/NT 4-wave GEMM kernel for Transformer Engine. -All supported layouts share one source-level kernel generator while compiling -to separate cached binaries: +FP16 and BF16 share one source-level kernel generator: the algorithm is +identical and only the ``v_mfma_f32_16x16x32_{f16,bf16}`` opcode differs. That +opcode is selected at compile time via ``mfma_suffix``, so each (dtype, K, +output, layout) combination still compiles to its own cached binary with no +runtime dtype branch. + +All supported layouts share the same 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 @@ -36,15 +42,18 @@ from .fp16_gemm_utils import ( G2SLoader, S2RLoader, - compute_global_fp16_transpose_swizzle, + compute_global_transpose_swizzle, compute_global_swizzle, - make_fp16_byte_buffer_tensor, + make_byte_buffer_tensor, pack_i32x4_i32x8, swizzle_128, xcd_swizzle, barrier ) +# FP16 and BF16 differ only in the MFMA opcode suffix. +_MFMA_SUFFIX = {torch.float16: "f16", torch.bfloat16: "bf16"} + _BLOCK_M = 256 _BLOCK_N = 256 @@ -84,13 +93,6 @@ 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 @@ -104,15 +106,19 @@ def _compile_kernel( K: int, output_dtype: torch.dtype, layout: str, + mfma_suffix: str, use_xcd_remap: bool = True, ): - """Build one compile-time-specialized TN, NN, or NT FP16 kernel. + """Build one compile-time-specialized TN, NN, or NT half-precision kernel. + ``mfma_suffix`` selects the ``v_mfma_f32_16x16x32_{f16,bf16}`` opcode. ``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}") + raise ValueError(f"Unsupported half-precision kernel layout: {layout}") + if mfma_suffix not in ("f16", "bf16"): + raise ValueError(f"Unsupported MFMA suffix: {mfma_suffix}") a_transpose_read = layout == "NT" b_transpose_read = layout in ("NN", "NT") @@ -146,7 +152,7 @@ def _compile_kernel( output_fx_dtype = fx.Float32 else: raise TypeError( - "FlyDSL FP16 GEMM output dtype must be torch.float16, " + "FlyDSL half-precision GEMM output dtype must be torch.float16, " f"torch.bfloat16, or torch.float32, got {output_dtype}" ) @@ -299,11 +305,11 @@ def _load_b_ni( 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 + # Half-precision 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( + return compute_global_transpose_swizzle( lane, wave_id, _a_leading_dim_bytes(c_m), @@ -322,7 +328,7 @@ def _a_global_offsets(lane, wave_id, c_m): if b_transpose_read: def _b_global_offsets(lane, wave_id, c_n): - return compute_global_fp16_transpose_swizzle( + return compute_global_transpose_swizzle( lane, wave_id, _b_leading_dim_bytes(c_n), @@ -341,7 +347,7 @@ def _b_global_offsets(lane, wave_id, c_n): @fx.struct class SharedStorage: - # Preserve the passing TN byte-staging contract exactly. A FP16 K64 + # Preserve the passing TN byte-staging contract exactly. A half-precision 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] @@ -367,10 +373,10 @@ def kernel_gemm( lds_b1 = (lds.b1_0, lds.b1_1) # A/B arrive as contiguous uint8 byte views of the original - # row-major FP16 tensors. This preserves the validated 16-byte + # row-major half-precision 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) + gA = make_byte_buffer_tensor(A) + gB = make_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") @@ -400,7 +406,7 @@ def kernel_gemm( lane = tx_i32 % fx.Int32(WARP_SIZE) # 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 + # swizzle. NN/NT stage K-major 16-bit 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) @@ -572,10 +578,10 @@ def load_b_frag(lds_b, local_row, half): ) def load_transposed_frag_half(lds_page, local_x_tile, half): - # FP16 uses v_mfma_f32_16x16x32_f16, not the MXFP8 K128 + # Half-precision uses v_mfma_f32_16x16x32_{f16,bf16}, 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. + # swizzled [K64, X64] half-precision slices. One ds_read_b64_tr_b16 returns + # four half-precision 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) @@ -596,7 +602,7 @@ def load_transposed_frag_half(lds_page, local_x_tile, half): 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( + return s2r.load_one_transpose( lds_page, base, other, @@ -611,9 +617,9 @@ def load_transposed_frag(lds_page, local_x_tile): 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). + def _k32_frag(full_frag, k32): + # A/B 16x64 half-precision wave fragments are i32x8. Each K32 MFMA + # consumes one contiguous i32x4 slice (eight half-precision values/lane). lo = k32 * 4 v = Vec(full_frag) return Vec.from_elements( @@ -621,13 +627,13 @@ def _fp16_k32_frag(full_frag, k32): fx.Int32, ) - def _pinned_fp16_mfma_once(acc_idx, a_k32, b_k32): + def _pinned_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"v_mfma_f32_16x16x32_{mfma_suffix} " f"a[{acc_pin}:{acc_pin + 3}], " f"$0, $1, " f"a[{acc_pin}:{acc_pin + 3}]" @@ -640,12 +646,12 @@ def _pinned_fp16_mfma_once(acc_idx, a_k32, b_k32): ) def pinned_mfma(acc_idx, a_frag, b_frag): - """Accumulate one logical 16x16x64 FP16 product into pinned AGPRs.""" + """Accumulate one logical 16x16x64 half-precision product into pinned AGPRs.""" for k32 in range_constexpr(2): - _pinned_fp16_mfma_once( + _pinned_mfma_once( acc_idx, - _fp16_k32_frag(a_frag, k32), - _fp16_k32_frag(b_frag, k32), + _k32_frag(a_frag, k32), + _k32_frag(b_frag, k32), ) def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): @@ -668,12 +674,12 @@ def mfma_2n_4mi_k32(subtile_id, n_base, k32, a0, a1, a2, a3, b0, b1): 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) + a_k32 = _k32_frag(a_frags[mi], k32) for nj in range_constexpr(2): - _pinned_fp16_mfma_once( + _pinned_mfma_once( _acc_idx(subtile_id, mi, n_base + nj), a_k32, - _fp16_k32_frag(b_frags[nj], k32), + _k32_frag(b_frags[nj], k32), ) def mfma_4n_4mi_k32(subtile_id, k32, a0, a1, a2, a3, b0, b1, b2, b3): @@ -681,12 +687,12 @@ def mfma_4n_4mi_k32(subtile_id, k32, a0, a1, a2, a3, b0, b1, b2, b3): 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) + a_k32 = _k32_frag(a_frags[mi], k32) for ni in range_constexpr(4): - _pinned_fp16_mfma_once( + _pinned_mfma_once( _acc_idx(subtile_id, mi, ni), a_k32, - _fp16_k32_frag(b_frags[ni], k32), + _k32_frag(b_frags[ni], k32), ) def store_acc_vector_for_logical_idx(logical_acc_idx, acc): @@ -1164,17 +1170,19 @@ def _cached_launch( K: int, output_dtype: torch.dtype, layout: str, + mfma_suffix: str, use_xcd_remap: bool = True, ): return _compile_kernel( K, output_dtype, layout, + mfma_suffix, use_xcd_remap=use_xcd_remap, ) -def fp16_matmul( +def _half_prec_matmul( a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, @@ -1183,24 +1191,30 @@ def fp16_matmul( m: int, n: int, k: int, + input_dtype: torch.dtype, + label: str, stream=None, ): - """Launch the wrapper-selected BF16 TN/NN/NT specialization.""" + """Validate operands and launch a half-precision TN/NN/NT specialization. + + ``input_dtype`` is the required FP16/BF16 operand dtype and ``label`` is the + human-readable kernel name used in error messages. + """ if layout not in ("TN", "NN", "NT"): - raise ValueError(f"Unsupported FP16 layout: {layout}") + raise ValueError(f"Unsupported {label} layout: {layout}") if a.ndim != 2 or b.ndim != 2: raise ValueError( - f"FlyDSL BF16 expects rank-2 operands, got A{tuple(a.shape)} " + f"FlyDSL {label} expects rank-2 operands, got A{tuple(a.shape)} " f"and B{tuple(b.shape)}" ) - if a.dtype != torch.float16 or b.dtype != torch.float16: + if a.dtype != input_dtype or b.dtype != input_dtype: raise TypeError( - "FlyDSL FP16 GEMM expects torch.float16 operands, " + f"FlyDSL {label} GEMM expects {input_dtype} 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"FlyDSL {label} {layout} requires original contiguous row-major " f"operands, got A stride={tuple(a.stride())}, " f"B stride={tuple(b.stride())}" ) @@ -1217,7 +1231,7 @@ def fp16_matmul( 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"FlyDSL {label} {layout} physical operands do not match contract: " f"A{tuple(a.shape)} expected {expected_a}; " f"B{tuple(b.shape)} expected {expected_b}" ) @@ -1226,7 +1240,7 @@ def fp16_matmul( 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 FP16 output must be float16, bfloat16, or float32, " + f"FlyDSL {label} output must be float16, bfloat16, or float32, " f"got {c.dtype}" ) if a.device != b.device or a.device != c.device: @@ -1235,7 +1249,7 @@ def fp16_matmul( f"{a.device}, {b.device}, and {c.device}" ) if not c.is_contiguous(): - raise ValueError("FlyDSL FP16 GEMM requires contiguous output storage") + raise ValueError(f"FlyDSL {label} GEMM requires contiguous output storage") doGemm( a, @@ -1245,9 +1259,64 @@ def fp16_matmul( m=m, n=n, k=k, + input_dtype=input_dtype, + label=label, + stream=stream, + ) + + +def fp16_matmul( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + *, + layout: str, + m: int, + n: int, + k: int, + stream=None, +): + """Launch the wrapper-selected FP16 TN/NN/NT specialization.""" + _half_prec_matmul( + a, + b, + c, + layout=layout, + m=m, + n=n, + k=k, + input_dtype=torch.float16, + label="FP16", + stream=stream, + ) + + +def bf16_matmul( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + *, + layout: str, + m: int, + n: int, + k: int, + stream=None, +): + """Launch the wrapper-selected BF16 TN/NN/NT specialization.""" + _half_prec_matmul( + a, + b, + c, + layout=layout, + m=m, + n=n, + k=k, + input_dtype=torch.bfloat16, + label="BF16", stream=stream, ) + def doGemm( A: torch.Tensor, B: torch.Tensor, @@ -1257,10 +1326,12 @@ def doGemm( m: int, n: int, k: int, + input_dtype: torch.dtype, + label: str, stream=None, use_xcd_remap: bool = True, ): - """Launch one cached K/output/layout-specialized FP16 core. + """Launch one cached K/output/layout-specialized half-precision core. A and B are passed unchanged from ``gemm_wrappers.py``. Their pointers reference the original rowwise allocations: @@ -1273,18 +1344,19 @@ def doGemm( ``ds_read_b64_tr_b16`` only. """ if layout not in ("TN", "NN", "NT"): - raise ValueError(f"Unsupported FP16 layout: {layout}") + raise ValueError(f"Unsupported {label} layout: {layout}") M_runtime = int(m) N_runtime = int(n) K_runtime = int(k) - if A.dtype != torch.float16 or B.dtype != torch.float16: + if A.dtype != input_dtype or B.dtype != input_dtype: raise TypeError( - f"BF16 {layout} requires BF16 inputs, got {A.dtype} and {B.dtype}" + f"{label} {layout} requires {input_dtype} inputs, " + f"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}") + raise TypeError(f"Unsupported {label} output dtype: {C.dtype}") require_block_tiling( M_runtime, @@ -1293,7 +1365,7 @@ def doGemm( block_m=_BLOCK_M, block_n=_BLOCK_N, block_k=_BLOCK_K, - label="FP16 GEMM", + label=f"{label} GEMM", ) if tuple(C.shape) != (M_runtime, N_runtime): @@ -1308,6 +1380,7 @@ def doGemm( K_runtime, C.dtype, layout, + _MFMA_SUFFIX[input_dtype], bool(use_xcd_remap), ) # Preserve the original validated byte-addressed G2L path. These are From 1e206caf63fb925348d4e706e6cc0777e91e61fa Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Mon, 10 Aug 2026 20:11:44 +0000 Subject: [PATCH 38/65] Add fused forward BIAS epilogue across all FlyDSL GEMM backends Implement a compile-time-selected epilogue for every FlyDSL GEMM kernel (mxfp8, tensor-wise fp8, fp16/bf16, fp32), adding a per-output-feature bias vector to the matmul result. The epilogue mode is resolved in _compile_kernel and keyed into the launch cache, so BIAS and DEFAULT compile to separate branch-free binaries; the store loop gates the bias add behind const_expr. Kernel details: - Bias is a length-N fp32 vector indexed by the GLOBAL output-feature index (by_n_idx + tile-local col), matching the C-store per-CTA tile-base fold, and broadcast across the M/token rows. The read is hoisted out of the four-row store loop since it is N-invariant. - Applied to the fp32 accumulator before the output-dtype cast; for fp8 it is added after the tensor-wise output scale, matching reference epilogue order. - DEFAULT passes a dummy 1-element bias so the kernel signature is uniform. - The epilogue enum reserves GELU_AUX / GELU_AUX_BIAS (raise NotImplemented) so the store-loop structure and dispatch are ready for a future GELU stage. Wiring: - Thread epilogue/bias through each *_matmul, doGemm, and lru-cached launch. - gemm_wrappers: a shared _resolve_bias() normalizes the TE bias into the kernel's (epilogue, length-N fp32) contract; threaded into all five _run_* dispatchers. Forward bias is now permitted for every backend; the common validator still rejects BGRADB (grad=True with bias). Tests: extend test_gemm.py with dtype-generic bias coverage (regular fp32/fp16/bf16, tensor-wise fp8, mxfp8), each vs a PyTorch reference (with a guard that bias actually changes the output) and vs the native C++ backend. Verified imports/compile and test collection only; GPU numerics still to run. Co-Authored-By: Claude --- tests/pytorch/flydsl_kernels/test_gemm.py | 283 +++++++++++++++++- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 94 +++++- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 87 ++++++ .../flydsl_kernels/gemm/gemm_wrappers.py | 62 +++- .../flydsl_kernels/gemm/half_prec_gemm.py | 91 ++++++ .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 90 +++++- 6 files changed, 689 insertions(+), 18 deletions(-) diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py index 48510880f..f88118ad7 100644 --- a/tests/pytorch/flydsl_kernels/test_gemm.py +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -13,9 +13,11 @@ - same-format and mixed-format MXFP8 - TN / NN / NT layouts - batched multidimensional FP8 flattening +- fused BIAS epilogue across all backends (regular, tensor-wise FP8, MXFP8) -Fused BIAS and BGRADB epilogues are intentionally not included yet because the -FlyDSL GEMM path does not currently support them. +The fused forward BIAS epilogue is implemented for every FlyDSL GEMM backend. +BGRADB (fused bias-gradient, grad=True) is not implemented on any FlyDSL path +yet. Each test compares the FlyDSL path against two independent references: @@ -208,6 +210,31 @@ def call_gemm(A, B, layout, out_dtype, use_flydsl=True): return output +def call_gemm_with_bias(A, B, layout, out_dtype, bias, use_flydsl=True): + """Call ``general_gemm`` with a fused forward BIAS epilogue. + + Bias is a 1-D vector along the output feature axis (the last dim of the + returned ``(*, out_features)`` tensor) and is added to the matmul result. + """ + 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=bias, + quantization_params=None, + gelu=False, + grad=False, + accumulate=False, + ) + + assert gelu_input is None + assert extra_output is None + return output, bias_grad + + def assert_gemm_close(actual, expected, *, atol, rtol): """Compare through FP32 so output narrowing does not hide diagnostics.""" torch.testing.assert_close( @@ -443,6 +470,237 @@ def test_flydsl_vs_cpp_mxfp8(M, K, N, layout, fp8_format): assert_gemm_close(flydsl_out, cpp_out, atol=5e-3, rtol=1e-2) +# ============================================================================== +# Fused BIAS epilogue coverage +# +# Every FlyDSL GEMM backend (regular fp32/fp16/bf16, tensor-wise FP8, MXFP8) +# adds bias along the output-feature (N) axis -- the last dim of the returned +# ``(M, N)`` tensor -- broadcast across the M/token rows. These guard the full +# public plumbing: general_gemm(bias=...) -> te_generic_gemm_flydsl -> +# _run_ -> _matmul(epilogue="BIAS"). +# +# Each backend has a vs-pytorch test (with a guard that bias actually changes +# the output, catching a silent decay to DEFAULT) and a vs-cpp cross-check. +# ============================================================================== + +@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_bias(M, K, N, layout, dtype): + """Regular fp32/fp16/bf16 GEMM with a fused BIAS epilogue vs PyTorch.""" + 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 + + expected_ab = compute_pytorch_reference(A.float(), B.float(), layout) + out_features = expected_ab.shape[-1] + bias = torch.randn(out_features, dtype=dtype, device="cuda") + + output, bias_grad = call_gemm_with_bias( + A, B, layout, out_dtype=dtype, bias=bias, use_flydsl=True, + ) + assert bias_grad is None + + no_bias_out = call_gemm(A, B, layout, out_dtype=dtype, use_flydsl=True) + assert not torch.allclose(output.float(), no_bias_out.float(), atol=1e-4), ( + "FlyDSL output matches the no-bias output; BIAS epilogue appears inactive." + ) + + expected = expected_ab + bias.float() + 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("dtype", REGULAR_DTYPES, ids=["fp32", "fp16", "bf16"]) +def test_flydsl_vs_cpp_regular_bias(M, K, N, layout, dtype): + """Regular BIAS epilogue: FlyDSL must match 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 + + out_features = compute_pytorch_reference(A.float(), B.float(), layout).shape[-1] + bias = torch.randn(out_features, dtype=dtype, device="cuda") + + flydsl_out, _ = call_gemm_with_bias( + A, B, layout, out_dtype=dtype, bias=bias, use_flydsl=True, + ) + cpp_out, _ = call_gemm_with_bias( + A, B, layout, out_dtype=dtype, bias=bias, 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_pytorch_fp8_bias(M, K, N, layout, fp8_format): + """Tensor-wise FP8 GEMM with a fused BIAS epilogue vs PyTorch.""" + 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, + ) + + expected_ab = compute_pytorch_reference(A_deq.float(), B_deq.float(), layout) + out_features = expected_ab.shape[-1] + bias = torch.randn(out_features, dtype=torch.float32, device="cuda") + + output, bias_grad = call_gemm_with_bias( + A_fp8, B_fp8, layout, out_dtype=torch.float32, bias=bias, use_flydsl=True, + ) + assert bias_grad is None + + no_bias_out = call_gemm( + A_fp8, B_fp8, layout, out_dtype=torch.float32, use_flydsl=True, + ) + assert not torch.allclose(output.float(), no_bias_out.float(), atol=1e-4), ( + "FlyDSL FP8 output matches the no-bias output; BIAS epilogue appears inactive." + ) + + expected = expected_ab + bias.float() + assert_gemm_close(output, expected, atol=5e-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_bias(M, K, N, layout, fp8_format): + """Tensor-wise FP8 BIAS epilogue: FlyDSL must match the native C++ backend.""" + 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, + ) + + out_features = compute_pytorch_reference( + A_deq.float(), B_deq.float(), layout, + ).shape[-1] + bias = torch.randn(out_features, dtype=torch.float32, device="cuda") + + flydsl_out, _ = call_gemm_with_bias( + A_fp8, B_fp8, layout, out_dtype=torch.float32, bias=bias, use_flydsl=True, + ) + cpp_out, _ = call_gemm_with_bias( + A_fp8, B_fp8, layout, out_dtype=torch.float32, bias=bias, 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_pytorch_mxfp8_bias(M, K, N, layout, fp8_format): + """MXFP8 forward GEMM with a fused BIAS epilogue vs a PyTorch reference.""" + 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, + ) + + expected_ab = compute_pytorch_reference(A_deq.float(), B_deq.float(), layout) + # Bias is a vector along the output-feature axis (last dim of the output). + out_features = expected_ab.shape[-1] + bias = torch.randn(out_features, dtype=torch.float32, device="cuda") + + output, bias_grad = call_gemm_with_bias( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + bias=bias, + use_flydsl=True, + ) + assert bias_grad is None + + # Bias must actually change the result -- guards against the BIAS epilogue + # silently decaying to DEFAULT and the test passing vacuously. + no_bias_out = call_gemm( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + assert not torch.allclose(output.float(), no_bias_out.float(), atol=1e-4), ( + "FlyDSL MXFP8 output matches the no-bias output; " + "the BIAS epilogue appears inactive." + ) + + expected = expected_ab + bias.float() + 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_cpp_mxfp8_bias(M, K, N, layout, fp8_format): + """MXFP8 forward BIAS epilogue: FlyDSL must match the native C++ backend.""" + 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, + ) + + out_features = compute_pytorch_reference( + A_deq.float(), + B_deq.float(), + layout, + ).shape[-1] + bias = torch.randn(out_features, dtype=torch.float32, device="cuda") + + flydsl_out, _ = call_gemm_with_bias( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + bias=bias, + use_flydsl=True, + ) + cpp_out, _ = call_gemm_with_bias( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + bias=bias, + use_flydsl=False, + ) + + assert_gemm_close(flydsl_out, cpp_out, atol=5e-3, rtol=1e-2) + + # ============================================================================== # Batched multidimensional FP8 coverage # ============================================================================== @@ -538,6 +796,20 @@ def test_flydsl_vs_pytorch_fp8_multidim( "TN", (tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2), ) + test_flydsl_vs_pytorch_regular_bias( + 256, + 512, + 256, + "TN", + torch.bfloat16, + ) + test_flydsl_vs_pytorch_fp8_bias( + 256, + 512, + 256, + "TN", + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), + ) if has_mxfp8_support: test_flydsl_vs_pytorch_mxfp8( @@ -547,5 +819,12 @@ def test_flydsl_vs_pytorch_fp8_multidim( "TN", (tex.DType.kFloat8E5M2, tex.DType.kFloat8E4M3), ) + test_flydsl_vs_pytorch_mxfp8_bias( + 256, + 512, + 256, + "TN", + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), + ) print("All FlyDSL GEMM smoke tests passed!") diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index e003947ef..036ff083e 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -92,12 +92,32 @@ assert LOAD_PASSES_B % 2 == 0 -def _compile_kernel(K: int, use_xcd_remap: bool = True): +def _compile_kernel(K: int, use_xcd_remap: bool = True, epilogue: str = "DEFAULT"): """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. + + ``epilogue`` selects the fused post-GEMM stages, resolved at compile time + so the store loop stays branch-free: + + DEFAULT plain matmul + BIAS + per-output-feature bias vector (indexed by N) + GELU_AUX (reserved) GELU with saved pre-activation aux + GELU_AUX_BIAS (reserved) bias then GELU with saved aux + + Only DEFAULT and BIAS are implemented; the GELU modes are accepted so the + store-loop structure and dispatch signature are already in place. """ + if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): + raise ValueError(f"Unsupported FP32 epilogue: {epilogue}") + if epilogue in ("GELU_AUX", "GELU_AUX_BIAS"): + raise NotImplementedError( + f"FP32 epilogue {epilogue} is reserved but not yet implemented" + ) + has_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") + has_gelu = epilogue in ("GELU_AUX", "GELU_AUX_BIAS") + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K NUM_THREADS = 256 WARP_SIZE = 64 @@ -154,6 +174,7 @@ def kernel_gemm( A: fx.Tensor, B: fx.Tensor, C: fx.Tensor, + Bias: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, ): @@ -217,6 +238,20 @@ def kernel_gemm( c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) c_store_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + # Bias is a length-N fp32 vector indexed by the global output-feature + # (N) coordinate and broadcast across the M/token rows. const_expr folds + # this compile-time flag at trace time so the setup (and load_bias) are + # inlined into the kernel scope with no runtime dispatch branch. + if const_expr(has_bias): + gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) + bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) + bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + + def load_bias(col): + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) + return fx.memref_load_vec(reg)[0] + PIN_ACC_BASE = 0 def _reg_list(prefix, start, end): @@ -436,11 +471,25 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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 + + # Bias depends only on the output-feature (N) coordinate, so read it + # once per column (global index by_n_idx + col) and reuse across the + # four M rows below. + if const_expr(has_bias): + bias_value = load_bias(by_n_idx + col) + for ii in range_constexpr(4): row = row_base + fx.Index(ii) c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + + # Epilogue stages run on the fp32 accumulator; output is fp32 + # so there is no dtype narrowing. GELU will slot in here later. + value = Vec(acc)[ii] + if const_expr(has_bias): + value = value + bias_value + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) - fx.memref_store_vec(Vec.filled(1, Vec(acc)[ii], fx.Float32), reg) + fx.memref_store_vec(Vec.filled(1, value, fx.Float32), reg) fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) @@ -863,6 +912,7 @@ def launch_gemm( A: fx.Tensor, B: fx.Tensor, C: fx.Tensor, + Bias: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, stream: fx.Stream = fx.Stream(None), @@ -873,6 +923,7 @@ def launch_gemm( A, B, C, + Bias, c_m, c_n, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, @@ -881,8 +932,8 @@ 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, use_xcd_remap: bool = True, epilogue: str = "DEFAULT"): + return _compile_kernel(K, use_xcd_remap=use_xcd_remap, epilogue=epilogue) @@ -891,6 +942,9 @@ def fp32_matmul( b: torch.Tensor, c: torch.Tensor, stream=None, + *, + epilogue: str = "DEFAULT", + bias: torch.Tensor = None, ): """TE-facing TN FP32 GEMM adapter. @@ -936,7 +990,7 @@ def fp32_matmul( raise ValueError("FlyDSL FP32 GEMM requires contiguous output storage") b_hk = b.transpose(0, 1).contiguous() - doGemm(a, b_hk, c, stream=stream) + doGemm(a, b_hk, c, stream=stream, epilogue=epilogue, bias=bias) def doGemm( @@ -945,6 +999,8 @@ def doGemm( C: torch.Tensor, stream=None, use_xcd_remap: bool = True, + epilogue: str = "DEFAULT", + bias: torch.Tensor = None, ): """Launch the private K-specialized FP32 core. @@ -966,11 +1022,35 @@ def doGemm( label="FP32 GEMM", ) assert C.shape == (M_runtime, N_runtime) + + if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): + raise ValueError(f"Unsupported FP32 epilogue: {epilogue}") + needs_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") + if needs_bias: + if bias is None: + raise ValueError(f"FP32 epilogue {epilogue} requires a bias tensor") + # Bias is indexed by the output-feature (N) axis and broadcast over M. + if bias.dtype != torch.float32: + raise TypeError(f"FP32 bias must be float32, got {bias.dtype}") + if bias.numel() != N_runtime: + raise ValueError( + f"FP32 bias length {bias.numel()} != N (out_features) {N_runtime}" + ) + if bias.device != A.device: + raise ValueError("bias must be on the same device as A, B, and C") + elif bias is not None: + raise ValueError(f"FP32 epilogue {epilogue} does not accept a bias tensor") + 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) + # DEFAULT keeps the kernel signature uniform with a dummy 1-element bias. + if needs_bias: + Bias_arg = bias.contiguous().view(-1) + else: + Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + launch = _cached_launch(int(K_runtime), bool(use_xcd_remap), epilogue) + launch(A_arg, B_arg, C_arg, Bias_arg, M_runtime, N_runtime, stream=stream) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 1fbd66f22..958eb1961 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -100,12 +100,33 @@ def _compile_kernel( b_fp8_dtype: torch.dtype, output_dtype: torch.dtype, use_xcd_remap: bool = True, + epilogue: str = "DEFAULT", ): """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. + + ``epilogue`` selects the fused post-GEMM stages, resolved at compile time + so the store loop stays branch-free: + + DEFAULT plain matmul + BIAS + per-output-feature bias vector (indexed by N) + GELU_AUX (reserved) GELU with saved pre-activation aux + GELU_AUX_BIAS (reserved) bias then GELU with saved aux + + Only DEFAULT and BIAS are implemented; the GELU modes are accepted so the + store-loop structure and dispatch signature are already in place. """ + if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): + raise ValueError(f"Unsupported FP8 epilogue: {epilogue}") + if epilogue in ("GELU_AUX", "GELU_AUX_BIAS"): + raise NotImplementedError( + f"FP8 epilogue {epilogue} is reserved but not yet implemented" + ) + has_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") + has_gelu = epilogue in ("GELU_AUX", "GELU_AUX_BIAS") + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K fp8_input_types = { @@ -190,6 +211,7 @@ def kernel_gemm( C: fx.Tensor, A_scale_inv: fx.Tensor, B_scale_inv: fx.Tensor, + Bias: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, ): @@ -268,6 +290,20 @@ def _load_scale(scale_div): output_fx_dtype, ) + # Bias is a length-N fp32 vector indexed by the global output-feature + # (N) coordinate and broadcast across the M/token rows. const_expr folds + # this compile-time flag at trace time so the setup (and load_bias) are + # inlined into the kernel scope with no runtime dispatch branch. + if const_expr(has_bias): + gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) + bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) + bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + + def load_bias(col): + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) + return fx.memref_load_vec(reg)[0] + PIN_ACC_BASE = 0 def _reg_list(prefix, start, end): @@ -461,10 +497,24 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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 + + # Bias depends only on the output-feature (N) coordinate, so read it + # once per column (global index by_n_idx + col) and reuse across the + # four M rows below. + if const_expr(has_bias): + bias_value = load_bias(by_n_idx + col) + for ii in range_constexpr(4): row = row_base + fx.Index(ii) c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + + # Epilogue stages run on the fp32 accumulator, in order, before + # the output-dtype narrowing. Bias is added after the tensor-wise + # output scale, matching the reference epilogue ordering. value = Vec(acc)[ii] * output_scale + if const_expr(has_bias): + value = value + bias_value + if output_dtype != torch.float32: value = value.to(output_fx_dtype) reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) @@ -892,6 +942,7 @@ def launch_gemm( C: fx.Tensor, A_scale_inv: fx.Tensor, B_scale_inv: fx.Tensor, + Bias: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, stream: fx.Stream = fx.Stream(None), @@ -904,6 +955,7 @@ def launch_gemm( C, A_scale_inv, B_scale_inv, + Bias, c_m, c_n, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, @@ -918,6 +970,7 @@ def _cached_launch( b_fp8_dtype: torch.dtype, output_dtype: torch.dtype, use_xcd_remap: bool = True, + epilogue: str = "DEFAULT", ): return _compile_kernel( K, @@ -925,6 +978,7 @@ def _cached_launch( b_fp8_dtype, output_dtype, use_xcd_remap=use_xcd_remap, + epilogue=epilogue, ) @@ -936,6 +990,9 @@ def fp8_matmul( b_scale_inv: torch.Tensor, c: torch.Tensor, stream=None, + *, + epilogue: str = "DEFAULT", + bias: torch.Tensor = None, ): """Launch TN tensor-wise FP8 GEMM using final kernel operand order. @@ -1011,6 +1068,8 @@ def fp8_matmul( a_scale_inv, b_scale_inv, stream=stream, + epilogue=epilogue, + bias=bias, ) def doGemm( @@ -1021,6 +1080,8 @@ def doGemm( B_scale_inv: torch.Tensor, stream=None, use_xcd_remap: bool = True, + epilogue: str = "DEFAULT", + bias: torch.Tensor = None, ): """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" M_runtime, K_runtime = A.shape @@ -1054,6 +1115,25 @@ def doGemm( assert C.shape == (M_runtime, N_runtime), ( f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" ) + + if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): + raise ValueError(f"Unsupported FP8 epilogue: {epilogue}") + needs_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") + if needs_bias: + if bias is None: + raise ValueError(f"FP8 epilogue {epilogue} requires a bias tensor") + # Bias is indexed by the output-feature (N) axis and broadcast over M. + if bias.dtype != torch.float32: + raise TypeError(f"FP8 bias must be float32, got {bias.dtype}") + if bias.numel() != N_runtime: + raise ValueError( + f"FP8 bias length {bias.numel()} != N (out_features) {N_runtime}" + ) + if bias.device != A.device: + raise ValueError("bias must be on the same device as A, B, and C") + elif bias is not None: + raise ValueError(f"FP8 epilogue {epilogue} does not accept a bias tensor") + if stream is None: stream = torch.cuda.current_stream() @@ -1062,6 +1142,11 @@ def doGemm( C_arg = C.contiguous().view(-1) A_scale_arg = A_scale_inv.contiguous().view(-1) B_scale_arg = B_scale_inv.contiguous().view(-1) + # DEFAULT keeps the kernel signature uniform with a dummy 1-element bias. + if needs_bias: + Bias_arg = bias.contiguous().view(-1) + else: + Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) launch = _cached_launch( int(K_runtime), @@ -1069,6 +1154,7 @@ def doGemm( B.dtype, C.dtype, bool(use_xcd_remap), + epilogue, ) launch( A_arg, @@ -1076,6 +1162,7 @@ def doGemm( C_arg, A_scale_arg, B_scale_arg, + Bias_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 4cff94af5..3db6a2fa5 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -95,7 +95,12 @@ def _validate_common_epilogue( alpha, beta, ): - """Validate features not yet implemented by the FlyDSL GEMM backend.""" + """Validate features not yet implemented by the FlyDSL GEMM backend. + + Fused forward BIAS is supported by all FlyDSL GEMM backends (mxfp8, fp8, + fp16, bf16, fp32). BGRADB (fused bias gradient, ``grad=True`` with a bias) + is not implemented on any FlyDSL path yet. + """ if quantizer is not None: raise NotImplementedError( "FlyDSL GEMM output quantization is not implemented" @@ -112,11 +117,28 @@ def _validate_common_epilogue( "FlyDSL GEMM accumulation is not implemented" ) - # TODO: Add fused bias and BGRADB epilogues - if bias is not None and bias.numel() != 0: + # Forward BIAS is implemented across all backends; the fused bias-gradient + # (BGRADB) path is not. TODO: add BGRADB to the FlyDSL GEMM backends. + if grad and bias is not None and bias.numel() != 0: raise NotImplementedError( - "FlyDSL GEMM bias is not implemented" + "FlyDSL GEMM fused bias gradient (BGRADB) is not implemented" + ) + + +def _resolve_bias(bias, n): + """Normalize a TE bias into the kernel's ``(epilogue, bias_arg)`` contract. + + FlyDSL kernels index bias by the output-feature (N) axis and expect a + contiguous length-N fp32 vector. TE may hand bias in the output dtype. + Returns ``("DEFAULT", None)`` when no bias is present. + """ + if bias is None or bias.numel() == 0: + return "DEFAULT", None + if bias.numel() != n: + raise ValueError( + f"FlyDSL GEMM bias length {bias.numel()} != N (out_features) {n}" ) + return "BIAS", bias.reshape(-1).to(torch.float32).contiguous() def _classify_input(t): @@ -351,6 +373,7 @@ def _run_bf16_gemm( D, *, output_dtype: torch.dtype, + bias=None, ): """Dispatch BF16 using the original row-major operand allocations. @@ -437,6 +460,7 @@ def _run_bf16_gemm( backend_name=f"BF16 {layout}", ) + epilogue, bias_arg = _resolve_bias(bias, n) bf16_matmul( a_flydsl, b_flydsl, @@ -445,6 +469,8 @@ def _run_bf16_gemm( m=m, n=n, k=k, + epilogue=epilogue, + bias=bias_arg, ) return D @@ -458,6 +484,7 @@ def _run_fp16_gemm( D, *, output_dtype: torch.dtype, + bias=None, ): """Dispatch FP16 using the original row-major operand allocations. @@ -544,6 +571,7 @@ def _run_fp16_gemm( backend_name=f"FP16 {layout}", ) + epilogue, bias_arg = _resolve_bias(bias, n) fp16_matmul( a_flydsl, b_flydsl, @@ -552,6 +580,8 @@ def _run_fp16_gemm( m=m, n=n, k=k, + epilogue=epilogue, + bias=bias_arg, ) return D @@ -562,6 +592,8 @@ def _run_fp32_gemm( B, transb, D, + *, + bias=None, ): """Normalize FP32 TN/NN/NT inputs to the current kernel's TN interface. @@ -655,10 +687,13 @@ def _run_fp32_gemm( backend_name="FP32 via TN core", ) + epilogue, bias_arg = _resolve_bias(bias, n) fp32_matmul( a_tn, b_tn, D.view(m, n), + epilogue=epilogue, + bias=bias_arg, ) return D @@ -825,6 +860,7 @@ def _run_mxfp8( D, *, output_dtype: torch.dtype, + bias=None, ): """Dispatch MXFP8 through exact TN/NN/NT physical contracts. @@ -1028,6 +1064,7 @@ def _run_mxfp8( f"M={m}, N={n}, K={k}" ) + epilogue, bias_arg = _resolve_bias(bias, n) mxfp8_matmul( a_flydsl, a_scale, @@ -1035,6 +1072,8 @@ def _run_mxfp8( b_scale, D.view(m, n), layout=kernel_layout, + epilogue=epilogue, + bias=bias_arg, ) return D @@ -1108,6 +1147,7 @@ def _run_fp8( D, *, output_dtype: torch.dtype, + bias=None, ): """Normalize tensor-wise FP8 storage and invoke the common FP8 core.""" supported_fp8_dtypes = ( @@ -1230,12 +1270,15 @@ def _run_fp8( _fp8_debug(f"derived M={m}, N={n}, K={k}") _fp8_tensor_debug("output/D", D) + epilogue, bias_arg = _resolve_bias(bias, n) matmul( a_flydsl, a_scale, b_flydsl, b_scale, D.view(m, n), + epilogue=epilogue, + bias=bias_arg, ) return D @@ -1295,6 +1338,9 @@ def te_generic_gemm_flydsl( "FlyDSL GEMM does not support transa=True, transb=True (TT)" ) + a_kind, _ = _classify_input(A) + b_kind, _ = _classify_input(B) + _validate_common_epilogue( quantizer=quantizer, bias=bias, @@ -1305,9 +1351,6 @@ def te_generic_gemm_flydsl( beta=beta, ) - a_kind, _ = _classify_input(A) - b_kind, _ = _classify_input(B) - if a_kind == "mxfp8" or b_kind == "mxfp8": # Validate both are MXFP8 if a_kind != b_kind: @@ -1340,6 +1383,7 @@ def te_generic_gemm_flydsl( transb, D, output_dtype=mxfp8_output_dtypes[output_dtype], + bias=bias, ) return D, None, None, None @@ -1368,6 +1412,7 @@ def te_generic_gemm_flydsl( transb, D, output_dtype=fp8_output_dtypes[output_dtype], + bias=bias, ) return D, None, None, None @@ -1400,6 +1445,7 @@ def te_generic_gemm_flydsl( transb, D, output_dtype=bf16_output_dtypes[output_dtype], + bias=bias, ) return D, None, None, None @@ -1422,6 +1468,7 @@ def te_generic_gemm_flydsl( transb, D, output_dtype=fp16_output_dtypes[output_dtype], + bias=bias, ) return D, None, None, None @@ -1437,6 +1484,7 @@ def te_generic_gemm_flydsl( B, transb, D, + bias=bias, ) return D, None, None, None diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 444ffdf6e..8a4bec15f 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -108,17 +108,37 @@ def _compile_kernel( layout: str, mfma_suffix: str, use_xcd_remap: bool = True, + epilogue: str = "DEFAULT", ): """Build one compile-time-specialized TN, NN, or NT half-precision kernel. ``mfma_suffix`` selects the ``v_mfma_f32_16x16x32_{f16,bf16}`` opcode. ``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. + + ``epilogue`` selects the fused post-GEMM stages, resolved at compile time + so the store loop stays branch-free: + + DEFAULT plain matmul + BIAS + per-output-feature bias vector (indexed by N) + GELU_AUX (reserved) GELU with saved pre-activation aux + GELU_AUX_BIAS (reserved) bias then GELU with saved aux + + Only DEFAULT and BIAS are implemented; the GELU modes are accepted so the + store-loop structure and dispatch signature are already in place. """ if layout not in ("TN", "NN", "NT"): raise ValueError(f"Unsupported half-precision kernel layout: {layout}") if mfma_suffix not in ("f16", "bf16"): raise ValueError(f"Unsupported MFMA suffix: {mfma_suffix}") + if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): + raise ValueError(f"Unsupported half-precision epilogue: {epilogue}") + if epilogue in ("GELU_AUX", "GELU_AUX_BIAS"): + raise NotImplementedError( + f"half-precision epilogue {epilogue} is reserved but not yet implemented" + ) + has_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") + has_gelu = epilogue in ("GELU_AUX", "GELU_AUX_BIAS") a_transpose_read = layout == "NT" b_transpose_read = layout in ("NN", "NT") @@ -363,6 +383,7 @@ def kernel_gemm( A: fx.Tensor, B: fx.Tensor, C: fx.Tensor, + Bias: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, ): @@ -443,6 +464,20 @@ def kernel_gemm( output_fx_dtype, ) + # Bias is a length-N fp32 vector indexed by the global output-feature + # (N) coordinate and broadcast across the M/token rows. const_expr folds + # this compile-time flag at trace time so the setup (and load_bias) are + # inlined into the kernel scope with no runtime dispatch branch. + if const_expr(has_bias): + gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) + bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) + bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + + def load_bias(col): + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) + return fx.memref_load_vec(reg)[0] + PIN_ACC_BASE = 0 def _reg_list(prefix, start, end): @@ -707,10 +742,23 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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 + + # Bias depends only on the output-feature (N) coordinate, so read it + # once per column (global index by_n_idx + col) and reuse across the + # four M rows below. + if const_expr(has_bias): + bias_value = load_bias(by_n_idx + col) + for ii in range_constexpr(4): row = row_base + fx.Index(ii) c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + + # Epilogue stages run on the fp32 accumulator, in order, before + # the output-dtype narrowing. GELU will slot in here later. value = Vec(acc)[ii] + if const_expr(has_bias): + value = value + bias_value + if const_expr(output_dtype != torch.float32): value = value.to(output_fx_dtype) reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) @@ -1148,6 +1196,7 @@ def launch_gemm( A: fx.Tensor, B: fx.Tensor, C: fx.Tensor, + Bias: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, stream: fx.Stream = fx.Stream(None), @@ -1158,6 +1207,7 @@ def launch_gemm( A, B, C, + Bias, c_m, c_n, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, @@ -1172,6 +1222,7 @@ def _cached_launch( layout: str, mfma_suffix: str, use_xcd_remap: bool = True, + epilogue: str = "DEFAULT", ): return _compile_kernel( K, @@ -1179,6 +1230,7 @@ def _cached_launch( layout, mfma_suffix, use_xcd_remap=use_xcd_remap, + epilogue=epilogue, ) @@ -1194,6 +1246,8 @@ def _half_prec_matmul( input_dtype: torch.dtype, label: str, stream=None, + epilogue: str = "DEFAULT", + bias: torch.Tensor = None, ): """Validate operands and launch a half-precision TN/NN/NT specialization. @@ -1262,6 +1316,8 @@ def _half_prec_matmul( input_dtype=input_dtype, label=label, stream=stream, + epilogue=epilogue, + bias=bias, ) @@ -1275,6 +1331,8 @@ def fp16_matmul( n: int, k: int, stream=None, + epilogue: str = "DEFAULT", + bias: torch.Tensor = None, ): """Launch the wrapper-selected FP16 TN/NN/NT specialization.""" _half_prec_matmul( @@ -1288,6 +1346,8 @@ def fp16_matmul( input_dtype=torch.float16, label="FP16", stream=stream, + epilogue=epilogue, + bias=bias, ) @@ -1301,6 +1361,8 @@ def bf16_matmul( n: int, k: int, stream=None, + epilogue: str = "DEFAULT", + bias: torch.Tensor = None, ): """Launch the wrapper-selected BF16 TN/NN/NT specialization.""" _half_prec_matmul( @@ -1314,6 +1376,8 @@ def bf16_matmul( input_dtype=torch.bfloat16, label="BF16", stream=stream, + epilogue=epilogue, + bias=bias, ) @@ -1330,6 +1394,8 @@ def doGemm( label: str, stream=None, use_xcd_remap: bool = True, + epilogue: str = "DEFAULT", + bias: torch.Tensor = None, ): """Launch one cached K/output/layout-specialized half-precision core. @@ -1373,6 +1439,24 @@ def doGemm( f"C shape {tuple(C.shape)} != expected {(M_runtime, N_runtime)}" ) + if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): + raise ValueError(f"Unsupported {label} epilogue: {epilogue}") + needs_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") + if needs_bias: + if bias is None: + raise ValueError(f"{label} epilogue {epilogue} requires a bias tensor") + # Bias is indexed by the output-feature (N) axis and broadcast over M. + if bias.dtype != torch.float32: + raise TypeError(f"{label} bias must be float32, got {bias.dtype}") + if bias.numel() != N_runtime: + raise ValueError( + f"{label} bias length {bias.numel()} != N (out_features) {N_runtime}" + ) + if bias.device != A.device: + raise ValueError("bias must be on the same device as A, B, and C") + elif bias is not None: + raise ValueError(f"{label} epilogue {epilogue} does not accept a bias tensor") + if stream is None: stream = torch.cuda.current_stream() @@ -1382,6 +1466,7 @@ def doGemm( layout, _MFMA_SUFFIX[input_dtype], bool(use_xcd_remap), + epilogue, ) # Preserve the original validated byte-addressed G2L path. These are # metadata-only dtype/flatten views of the already-contiguous row-major @@ -1389,11 +1474,17 @@ def doGemm( A_arg = A.view(torch.uint8).view(-1) B_arg = B.view(torch.uint8).view(-1) C_arg = C.view(-1) + # DEFAULT keeps the kernel signature uniform with a dummy 1-element bias. + if needs_bias: + Bias_arg = bias.contiguous().view(-1) + else: + Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) launch( A_arg, B_arg, C_arg, + Bias_arg, M_runtime, N_runtime, stream=stream, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 050cc4fe8..98c6db724 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -29,7 +29,7 @@ import flydsl.compiler as flyc import flydsl.expr as fx from flydsl._mlir.dialects import llvm -from flydsl.expr import arith, gpu, range_constexpr, rocdl +from flydsl.expr import arith, const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec @@ -279,15 +279,35 @@ def _compile_kernel( b_fp8_dtype: torch.dtype, output_dtype: torch.dtype, layout: str, + epilogue: str = "DEFAULT", ): """Build one compile-time-specialized TN, NN, or NT kernel. ``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. + + ``epilogue`` selects the fused post-GEMM stages, resolved at compile time + so the store loop stays branch-free: + + DEFAULT plain matmul + BIAS + per-output-feature bias vector (indexed by N) + GELU_AUX (reserved) GELU with saved pre-activation aux + GELU_AUX_BIAS (reserved) bias then GELU with saved aux + + Only DEFAULT and BIAS are implemented; the GELU modes are accepted so the + store-loop structure and dispatch signature are already in place. """ if layout not in ("TN", "NN", "NT"): raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") + if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): + raise ValueError(f"Unsupported MXFP8 epilogue: {epilogue}") + if epilogue in ("GELU_AUX", "GELU_AUX_BIAS"): + raise NotImplementedError( + f"MXFP8 epilogue {epilogue} is reserved but not yet implemented" + ) + has_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") + has_gelu = epilogue in ("GELU_AUX", "GELU_AUX_BIAS") a_transpose_read = layout == "NT" b_transpose_read = layout in ("NN", "NT") @@ -520,7 +540,7 @@ class SharedStorage: @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 + A: fx.Tensor, As: fx.Tensor, B: fx.Tensor, Bs: fx.Tensor, C: fx.Tensor, Bias: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32 ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -616,6 +636,22 @@ def kernel_gemm( output_fx_dtype, ) + # Bias is a length-N fp32 vector indexed by the output-feature (N) + # coordinate and broadcast across the M/token rows. Only staged when + # the epilogue requests it; DEFAULT receives a dummy 1-element tensor. + # const_expr folds this compile-time flag at trace time so the setup + # (and load_bias) are inlined into the kernel scope with no runtime + # dispatch branch. + if const_expr(has_bias): + gBias = fx.rocdl.make_buffer_tensor(Bias, max_size=True) + bias_div = fx.logical_divide(gBias, fx.make_layout(1, 1)) + bias_ld_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + + def load_bias(col): + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) + return fx.memref_load_vec(reg)[0] + PIN_ACC_BASE = 0 def _reg_list(prefix, start, end): @@ -908,10 +944,25 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): 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 + + # Bias depends only on the output-feature (N) coordinate, so read it + # once per column and reuse across the four M rows below. Index by + # the GLOBAL feature (by_n_idx + col), matching the C store's + # c_tile_base_elems fold; the tile-local col alone would reread + # bias[0..BLOCK_N) for every N-block. + if const_expr(has_bias): + bias_value = load_bias(by_n_idx + col) + for ii in range_constexpr(4): row = row_base + fx.Index(ii) c_idx = c_tile_base_elems + row * fx.Index(c_n) + col + + # Epilogue stages run on the fp32 accumulator, in order, before + # the output-dtype narrowing. GELU will slot in here later. value = Vec(acc)[ii] + if const_expr(has_bias): + value = value + bias_value + if output_dtype != torch.float32: value = value.to(output_fx_dtype) reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) @@ -1446,6 +1497,7 @@ def launch_gemm( B: fx.Tensor, Bs: fx.Tensor, C: fx.Tensor, + Bias: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, stream: fx.Stream = fx.Stream(None), @@ -1458,6 +1510,7 @@ def launch_gemm( B, Bs, C, + Bias, c_m, c_n, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, @@ -1474,6 +1527,8 @@ def do_gemm( stream=None, *, layout: str = "TN", + epilogue: str = "DEFAULT", + bias: torch.Tensor = None, ): """Launch one cached compile-time MXFP8 layout specialization.""" if layout == "TN": @@ -1522,6 +1577,24 @@ def do_gemm( 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 epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): + raise ValueError(f"Unsupported MXFP8 epilogue: {epilogue}") + needs_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") + if needs_bias: + if bias is None: + raise ValueError(f"MXFP8 epilogue {epilogue} requires a bias tensor") + # Bias is indexed by the output-feature (N) axis and broadcast over M. + if bias.dtype != torch.float32: + raise TypeError(f"MXFP8 bias must be float32, got {bias.dtype}") + if bias.numel() != N_runtime: + raise ValueError( + f"MXFP8 bias length {bias.numel()} != N (out_features) {N_runtime}" + ) + if bias.device != A.device: + raise ValueError("bias must be on the same device as A, B, and C") + elif bias is not None: + raise ValueError(f"MXFP8 epilogue {epilogue} does not accept a bias tensor") + if stream is None: stream = torch.cuda.current_stream() @@ -1531,6 +1604,11 @@ def do_gemm( As_arg = As.contiguous().view(-1) Bs_arg = Bs.contiguous().view(-1) C_arg = C.contiguous().view(-1) + # DEFAULT keeps the kernel signature uniform with a dummy 1-element bias. + if needs_bias: + Bias_arg = bias.contiguous().view(-1) + else: + Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) _cached_launch( K_runtime, @@ -1538,12 +1616,14 @@ def do_gemm( B.dtype, C.dtype, layout, + epilogue, )( A_arg, As_arg, B_arg, Bs_arg, C_arg, + Bias_arg, M_runtime, N_runtime, stream=stream, @@ -1557,6 +1637,7 @@ def _cached_launch( b_fp8_dtype: torch.dtype, output_dtype: torch.dtype, layout: str, + epilogue: str = "DEFAULT", ): """Cache independent TN/NN/NT binaries with no runtime layout argument.""" return _compile_kernel( @@ -1565,6 +1646,7 @@ def _cached_launch( b_fp8_dtype, output_dtype, layout, + epilogue, ) @@ -1604,6 +1686,8 @@ def mxfp8_matmul( stream=None, *, layout: str = "TN", + epilogue: str = "DEFAULT", + bias: torch.Tensor = None, ): """Normalize scale orientation and launch a compile-time layout binary. @@ -1727,6 +1811,8 @@ def mxfp8_matmul( D.view(m, n), layout=layout, stream=stream, + epilogue=epilogue, + bias=bias, ) return D From b7606218cc7fb25844bf6666ca498e3c0de4561a Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Mon, 10 Aug 2026 21:34:21 +0000 Subject: [PATCH 39/65] Add fused GELU_AUX / GELU_AUX_BIAS epilogue to MXFP8 GEMM Implement forward GELU for the MXFP8 FlyDSL kernel. Unlike a plain fused activation, GELU here always saves the pre-activation values (A@B, or A@B + bias for GELU_AUX_BIAS) to a second M x N output so the backward pass can recompute the GELU gradient -- matching the hipkittens epilogue where the aux store and the activation are inseparable. Kernel (mxfp8_gemm.py): - _compile_kernel gains the GELU_AUX / GELU_AUX_BIAS epilogue modes (were reserved). In the store loop, when has_gelu, save the post-bias pre-activation to the Aux output, then apply tanh-approx GELU to the C value. The activation is inlined from scalar math (the FlyDSL preshuffle reference's overflow-safe exp(-2|y|) form), so there is no GELU-primitive dependency. - kernel_gemm / launch_gemm gain an Aux tensor argument; do_gemm validates the caller-allocated aux (M x N, C dtype) and passes it through, filled in place. DEFAULT/BIAS keep the signature uniform with a dummy 1-element aux buffer. Wrapper (gemm_wrappers.py): - _resolve_epilogue combines bias + gelu into DEFAULT / BIAS / GELU_AUX / GELU_AUX_BIAS. _run_mxfp8 allocates the aux, launches in place, and returns it; the dispatcher hands it back in general_gemm's gelu_input (3rd) slot, matching how layernorm_mlp unpacks the fused-FC1 result. - GELU is MXFP8-only for now: the fp8/regular dispatch paths reject gelu, and _validate_common_epilogue rejects DGELU (gelu + grad). Tests: mxfp8 GELU_AUX and GELU_AUX_BIAS vs PyTorch, asserting both the aux (pre-activation) and the output (tanh-approx GELU of it), with a guard that GELU actually changes the result. Verified imports/compile and test collection only; GPU numerics still to run. Co-Authored-By: Claude --- tests/pytorch/flydsl_kernels/test_gemm.py | 148 +++++++++++++++++- .../flydsl_kernels/gemm/gemm_wrappers.py | 50 +++++- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 94 +++++++++-- 3 files changed, 271 insertions(+), 21 deletions(-) diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py index f88118ad7..6d19938b1 100644 --- a/tests/pytorch/flydsl_kernels/test_gemm.py +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -14,10 +14,12 @@ - TN / NN / NT layouts - batched multidimensional FP8 flattening - fused BIAS epilogue across all backends (regular, tensor-wise FP8, MXFP8) +- fused GELU_AUX / GELU_AUX_BIAS epilogue for MXFP8 The fused forward BIAS epilogue is implemented for every FlyDSL GEMM backend. -BGRADB (fused bias-gradient, grad=True) is not implemented on any FlyDSL path -yet. +Fused forward GELU (GELU_AUX, saving the pre-activation aux) is implemented for +MXFP8 only. BGRADB (fused bias-gradient) and DGELU (fused GELU gradient, +grad=True) are not implemented on any FlyDSL path yet. Each test compares the FlyDSL path against two independent references: @@ -235,6 +237,37 @@ def call_gemm_with_bias(A, B, layout, out_dtype, bias, use_flydsl=True): return output, bias_grad +def call_gemm_with_gelu(A, B, layout, out_dtype, bias=None, use_flydsl=True): + """Call ``general_gemm`` with a fused forward GELU (GELU_AUX) epilogue. + + Returns ``(output, gelu_input)`` where ``output`` is ``gelu(A@B[+bias])`` + and ``gelu_input`` is the saved pre-activation (``A@B[+bias]``) that the + backward pass consumes. + """ + 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=bias, + quantization_params=None, + gelu=True, + grad=False, + accumulate=False, + ) + + assert bias_grad is None + assert extra_output is None + return output, gelu_input + + +def gelu_tanh_ref(x): + """tanh-approx GELU reference (matches the kernel and torch approximate='tanh').""" + return torch.nn.functional.gelu(x, approximate="tanh") + + def assert_gemm_close(actual, expected, *, atol, rtol): """Compare through FP32 so output narrowing does not hide diagnostics.""" torch.testing.assert_close( @@ -701,6 +734,110 @@ def test_flydsl_vs_cpp_mxfp8_bias(M, K, N, layout, fp8_format): assert_gemm_close(flydsl_out, cpp_out, atol=5e-3, rtol=1e-2) +# ============================================================================== +# Fused GELU epilogue coverage (MXFP8) +# +# GELU_AUX applies tanh-approx GELU to the C output while saving the +# pre-activation (A@B, or A@B+bias for GELU_AUX_BIAS) to a second aux output +# for the backward pass. general_gemm returns the aux in the third tuple slot +# (gelu_input). Currently implemented for the MXFP8 backend only. +# ============================================================================== + +@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_gelu(M, K, N, layout, fp8_format): + """MXFP8 forward GELU_AUX vs PyTorch: check both output and saved aux.""" + 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, + ) + + pre_act = compute_pytorch_reference(A_deq.float(), B_deq.float(), layout) + + output, gelu_input = call_gemm_with_gelu( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + assert gelu_input is not None, "GELU_AUX did not return the pre-activation aux." + + # GELU must actually change the result vs the plain matmul output. + no_gelu_out = call_gemm( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + assert not torch.allclose(output.float(), no_gelu_out.float(), atol=1e-4), ( + "FlyDSL MXFP8 output matches the no-GELU output; " + "the GELU epilogue appears inactive." + ) + + # Aux is the pre-activation (A@B); output is gelu(A@B). + assert_gemm_close(gelu_input, pre_act, atol=5e-3, rtol=1e-2) + assert_gemm_close(output, gelu_tanh_ref(pre_act), 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_gelu_bias(M, K, N, layout, fp8_format): + """MXFP8 forward GELU_AUX_BIAS vs PyTorch: bias folded before GELU, aux saved.""" + 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, + ) + + ab = compute_pytorch_reference(A_deq.float(), B_deq.float(), layout) + out_features = ab.shape[-1] + bias = torch.randn(out_features, dtype=torch.float32, device="cuda") + pre_act = ab + bias.float() + + output, gelu_input = call_gemm_with_gelu( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + bias=bias, + use_flydsl=True, + ) + assert gelu_input is not None, "GELU_AUX_BIAS did not return the pre-activation aux." + + # Aux is the post-bias pre-activation (A@B + bias); output is gelu of it. + assert_gemm_close(gelu_input, pre_act, atol=5e-3, rtol=1e-2) + assert_gemm_close(output, gelu_tanh_ref(pre_act), atol=5e-3, rtol=1e-2) + + # ============================================================================== # Batched multidimensional FP8 coverage # ============================================================================== @@ -826,5 +963,12 @@ def test_flydsl_vs_pytorch_fp8_multidim( "TN", (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), ) + test_flydsl_vs_pytorch_mxfp8_gelu( + 256, + 512, + 256, + "TN", + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), + ) print("All FlyDSL GEMM smoke tests passed!") diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 3db6a2fa5..c230bdd80 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -98,8 +98,9 @@ def _validate_common_epilogue( """Validate features not yet implemented by the FlyDSL GEMM backend. Fused forward BIAS is supported by all FlyDSL GEMM backends (mxfp8, fp8, - fp16, bf16, fp32). BGRADB (fused bias gradient, ``grad=True`` with a bias) - is not implemented on any FlyDSL path yet. + fp16, bf16, fp32). Fused forward GELU_AUX is implemented for MXFP8 only + (the per-backend ``_run_*`` reject it elsewhere). BGRADB (fused bias + gradient) and DGELU (fused GELU gradient) are not implemented anywhere yet. """ if quantizer is not None: raise NotImplementedError( @@ -124,6 +125,13 @@ def _validate_common_epilogue( "FlyDSL GEMM fused bias gradient (BGRADB) is not implemented" ) + # Fused forward GELU (GELU_AUX) is supported; the backward fused GELU + # gradient (DGELU) is not. TODO: add DGELU to the FlyDSL GEMM backends. + if gelu and grad: + raise NotImplementedError( + "FlyDSL GEMM fused GELU gradient (DGELU) is not implemented" + ) + def _resolve_bias(bias, n): """Normalize a TE bias into the kernel's ``(epilogue, bias_arg)`` contract. @@ -141,6 +149,21 @@ def _resolve_bias(bias, n): return "BIAS", bias.reshape(-1).to(torch.float32).contiguous() +def _resolve_epilogue(bias, n, gelu=False): + """Resolve the ``(epilogue, bias_arg)`` contract including fused GELU. + + Combines the bias vector (see :func:`_resolve_bias`) with an optional + forward GELU into one of DEFAULT / BIAS / GELU_AUX / GELU_AUX_BIAS. The + GELU epilogues additionally produce a pre-activation aux output, which the + caller allocates and passes to the kernel. + """ + bias_epilogue, bias_arg = _resolve_bias(bias, n) + has_bias = bias_epilogue == "BIAS" + if gelu: + return ("GELU_AUX_BIAS" if has_bias else "GELU_AUX"), bias_arg + return bias_epilogue, bias_arg + + def _classify_input(t): """Classify a GEMM operand for the FlyDSL backend.""" try: @@ -861,6 +884,7 @@ def _run_mxfp8( *, output_dtype: torch.dtype, bias=None, + gelu=False, ): """Dispatch MXFP8 through exact TN/NN/NT physical contracts. @@ -1064,7 +1088,11 @@ def _run_mxfp8( f"M={m}, N={n}, K={k}" ) - epilogue, bias_arg = _resolve_bias(bias, n) + epilogue, bias_arg = _resolve_epilogue(bias, n, gelu=gelu) + # GELU_AUX modes emit a pre-activation aux output (M x N, output dtype) + # for the backward pass; the wrapper allocates it and the kernel fills it + # in place. Return it so the dispatcher can hand it back as gelu_input. + aux = torch.empty_like(D) if gelu else None mxfp8_matmul( a_flydsl, a_scale, @@ -1074,8 +1102,9 @@ def _run_mxfp8( layout=kernel_layout, epilogue=epilogue, bias=bias_arg, + aux=aux.view(m, n) if aux is not None else None, ) - return D + return D, aux def _select_fp8_storage_for_layout(A, transa, B, transb): @@ -1376,7 +1405,7 @@ def te_generic_gemm_flydsl( f"got {output_dtype}" ) - D = _run_mxfp8( + D, gelu_input = _run_mxfp8( A, transa, B, @@ -1384,8 +1413,17 @@ def te_generic_gemm_flydsl( D, output_dtype=mxfp8_output_dtypes[output_dtype], bias=bias, + gelu=gelu, + ) + return D, None, gelu_input, None + + # Fused forward GELU is implemented for MXFP8 only so far; the tensor-wise + # FP8 and regular (fp16/bf16/fp32) paths do not support it yet. + # TODO: extend GELU_AUX to the other FlyDSL GEMM backends. + if gelu: + raise NotImplementedError( + "FlyDSL GEMM fused GELU is currently implemented for MXFP8 only" ) - return D, None, None, None if a_kind == "fp8" or b_kind == "fp8": if a_kind != b_kind: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 98c6db724..2a78f7f44 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -29,7 +29,7 @@ import flydsl.compiler as flyc import flydsl.expr as fx from flydsl._mlir.dialects import llvm -from flydsl.expr import arith, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import arith, const_expr, gpu, math, range_constexpr, rocdl from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec @@ -292,20 +292,16 @@ def _compile_kernel( DEFAULT plain matmul BIAS + per-output-feature bias vector (indexed by N) - GELU_AUX (reserved) GELU with saved pre-activation aux - GELU_AUX_BIAS (reserved) bias then GELU with saved aux + GELU_AUX GELU(A@B), saving the pre-activation to the Aux output + GELU_AUX_BIAS GELU(A@B + bias), saving the pre-activation to Aux - Only DEFAULT and BIAS are implemented; the GELU modes are accepted so the - store-loop structure and dispatch signature are already in place. + The GELU modes write a second M x N output (pre-activation, tanh-approx + GELU applied to C) for the backward pass. """ if layout not in ("TN", "NN", "NT"): raise ValueError(f"Unsupported MXFP8 kernel layout: {layout}") if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): raise ValueError(f"Unsupported MXFP8 epilogue: {epilogue}") - if epilogue in ("GELU_AUX", "GELU_AUX_BIAS"): - raise NotImplementedError( - f"MXFP8 epilogue {epilogue} is reserved but not yet implemented" - ) has_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") has_gelu = epilogue in ("GELU_AUX", "GELU_AUX_BIAS") @@ -540,7 +536,7 @@ class SharedStorage: @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, Bias: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32 + A: fx.Tensor, As: fx.Tensor, B: fx.Tensor, Bs: fx.Tensor, C: fx.Tensor, Bias: fx.Tensor, Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32 ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -652,6 +648,34 @@ def load_bias(col): fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) return fx.memref_load_vec(reg)[0] + # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N + # output so the backward pass can recompute the GELU gradient. Same tile + # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. + if const_expr(has_gelu): + gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) + aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) + aux_store_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + output_fx_dtype, + ) + + def gelu_tanh(x): + # tanh-approx GELU (matches PyTorch approximate='tanh' and the + # FlyDSL preshuffle reference), expressed through a non-positive + # exponent so exp() cannot overflow: + # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) + half_f32 = fx.Float32(0.5) + one_f32 = fx.Float32(1.0) + zero_f32 = fx.Float32(0.0) + two_f32 = fx.Float32(2.0) + x3 = x * x * x + y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) + abs_y = fx.Float32(y).maximumf(zero_f32 - y) + e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) + denom = one_f32 + e_neg2abs + numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) + return half_f32 * x * (numerator * (one_f32 / denom)) + PIN_ACC_BASE = 0 def _reg_list(prefix, start, end): @@ -958,11 +982,24 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): c_idx = c_tile_base_elems + row * fx.Index(c_n) + col # Epilogue stages run on the fp32 accumulator, in order, before - # the output-dtype narrowing. GELU will slot in here later. + # the output-dtype narrowing: + # value = acc [+ bias] (pre-activation) + # GELU_AUX: save pre-activation to Aux, then value = gelu(value) value = Vec(acc)[ii] if const_expr(has_bias): value = value + bias_value + if const_expr(has_gelu): + # Save the pre-activation (post-bias) value for backward, + # then apply GELU to the C output. + aux_val = value + if output_dtype != torch.float32: + aux_val = aux_val.to(output_fx_dtype) + aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) + fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) + value = gelu_tanh(value) + if output_dtype != torch.float32: value = value.to(output_fx_dtype) reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) @@ -1498,6 +1535,7 @@ def launch_gemm( Bs: fx.Tensor, C: fx.Tensor, Bias: fx.Tensor, + Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, stream: fx.Stream = fx.Stream(None), @@ -1511,6 +1549,7 @@ def launch_gemm( Bs, C, Bias, + Aux, c_m, c_n, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, @@ -1529,8 +1568,13 @@ def do_gemm( layout: str = "TN", epilogue: str = "DEFAULT", bias: torch.Tensor = None, + aux: torch.Tensor = None, ): - """Launch one cached compile-time MXFP8 layout specialization.""" + """Launch one cached compile-time MXFP8 layout specialization. + + The GELU epilogues require a caller-allocated ``aux`` output (M x N, same + dtype as C), filled in place with the pre-activation values for backward. + """ if layout == "TN": M_runtime, K_runtime = A.shape N_runtime, Kb_runtime = B.shape @@ -1595,6 +1639,23 @@ def do_gemm( elif bias is not None: raise ValueError(f"MXFP8 epilogue {epilogue} does not accept a bias tensor") + needs_aux = epilogue in ("GELU_AUX", "GELU_AUX_BIAS") + if needs_aux: + # Pre-activation output for the backward pass: caller-allocated M x N, + # same dtype as C, filled in place through its buffer descriptor. + if aux is None: + raise ValueError(f"MXFP8 epilogue {epilogue} requires an aux output tensor") + if tuple(aux.shape) != (M_runtime, N_runtime): + raise ValueError( + f"MXFP8 aux shape {tuple(aux.shape)} != {(M_runtime, N_runtime)}" + ) + if aux.dtype != C.dtype: + raise TypeError(f"MXFP8 aux dtype {aux.dtype} != C dtype {C.dtype}") + if aux.device != A.device: + raise ValueError("aux must be on the same device as A, B, and C") + elif aux is not None: + raise ValueError(f"MXFP8 epilogue {epilogue} does not accept an aux tensor") + if stream is None: stream = torch.cuda.current_stream() @@ -1604,11 +1665,15 @@ def do_gemm( As_arg = As.contiguous().view(-1) Bs_arg = Bs.contiguous().view(-1) C_arg = C.contiguous().view(-1) - # DEFAULT keeps the kernel signature uniform with a dummy 1-element bias. + # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + if needs_aux: + Aux_arg = aux.view(-1) + else: + Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) _cached_launch( K_runtime, @@ -1624,6 +1689,7 @@ def do_gemm( Bs_arg, C_arg, Bias_arg, + Aux_arg, M_runtime, N_runtime, stream=stream, @@ -1688,6 +1754,7 @@ def mxfp8_matmul( layout: str = "TN", epilogue: str = "DEFAULT", bias: torch.Tensor = None, + aux: torch.Tensor = None, ): """Normalize scale orientation and launch a compile-time layout binary. @@ -1813,6 +1880,7 @@ def mxfp8_matmul( stream=stream, epilogue=epilogue, bias=bias, + aux=aux.view(m, n) if aux is not None else None, ) return D From 87b1bb03ee9740beea1418796426f04620be048f Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 00:44:04 +0000 Subject: [PATCH 40/65] Extend fused GELU_AUX / GELU_AUX_BIAS epilogue to all FlyDSL GEMM backends Following the MXFP8 implementation, add forward GELU (with saved pre-activation aux) to the half-precision (fp16/bf16), tensor-wise FP8, and FP32 kernels so fused GELU is now uniform across every FlyDSL GEMM backend. Kernels (half_prec_gemm, fp8_gemm, fp32_gemm): - _compile_kernel accepts the GELU_AUX / GELU_AUX_BIAS modes (were reserved). In the store loop, when has_gelu, save the post-bias pre-activation to the Aux output, then apply the inlined tanh-approx GELU to the C value. All gated by const_expr so DEFAULT/BIAS stay branch-free. - kernel_gemm / launch_gemm gain an Aux argument; doGemm validates the caller-allocated aux (M x N, C dtype; fp32 for the fp32 kernel) and passes it through, filled in place. DEFAULT/BIAS keep the signature uniform with a dummy 1-element aux. Per-backend ordering preserved: fp8 saves the aux after the tensor-wise output scale; fp32 has no output cast. Wrapper (gemm_wrappers.py): - Drop the MXFP8-only GELU guard. _run_bf16/_run_fp16/_run_fp32/_run_fp8 now take gelu, allocate the aux, launch in place via _resolve_epilogue, and return (D, aux); the dispatcher hands the aux back in general_gemm's gelu_input slot for every backend. DGELU (gelu + grad) stays rejected. Tests: dtype-generic GELU_AUX coverage for regular (fp32/fp16/bf16) and tensor-wise FP8, mirroring the MXFP8 tests (output == gelu(pre_act), aux == pre_act, guard that GELU changes the output). Verified imports/compile and test collection only; GPU numerics still to run. Co-Authored-By: Claude --- tests/pytorch/flydsl_kernels/test_gemm.py | 87 ++++++++++++++++-- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 81 ++++++++++++++--- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 84 ++++++++++++++--- .../flydsl_kernels/gemm/gemm_wrappers.py | 56 +++++++----- .../flydsl_kernels/gemm/half_prec_gemm.py | 89 ++++++++++++++++--- 5 files changed, 331 insertions(+), 66 deletions(-) diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py index 6d19938b1..2605af856 100644 --- a/tests/pytorch/flydsl_kernels/test_gemm.py +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -14,12 +14,12 @@ - TN / NN / NT layouts - batched multidimensional FP8 flattening - fused BIAS epilogue across all backends (regular, tensor-wise FP8, MXFP8) -- fused GELU_AUX / GELU_AUX_BIAS epilogue for MXFP8 +- fused GELU_AUX / GELU_AUX_BIAS epilogue across all backends -The fused forward BIAS epilogue is implemented for every FlyDSL GEMM backend. -Fused forward GELU (GELU_AUX, saving the pre-activation aux) is implemented for -MXFP8 only. BGRADB (fused bias-gradient) and DGELU (fused GELU gradient, -grad=True) are not implemented on any FlyDSL path yet. +The fused forward BIAS and GELU (GELU_AUX, saving the pre-activation aux) +epilogues are implemented for every FlyDSL GEMM backend. BGRADB (fused +bias-gradient) and DGELU (fused GELU gradient, grad=True) are not implemented +on any FlyDSL path yet. Each test compares the FlyDSL path against two independent references: @@ -735,14 +735,73 @@ def test_flydsl_vs_cpp_mxfp8_bias(M, K, N, layout, fp8_format): # ============================================================================== -# Fused GELU epilogue coverage (MXFP8) +# Fused GELU epilogue coverage # # GELU_AUX applies tanh-approx GELU to the C output while saving the # pre-activation (A@B, or A@B+bias for GELU_AUX_BIAS) to a second aux output # for the backward pass. general_gemm returns the aux in the third tuple slot -# (gelu_input). Currently implemented for the MXFP8 backend only. +# (gelu_input). Implemented across all FlyDSL backends (regular fp32/fp16/bf16, +# tensor-wise FP8, MXFP8). Each test checks output == gelu(pre_act) and +# aux == pre_act, with a guard that GELU actually changes the output. # ============================================================================== +@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_gelu(M, K, N, layout, dtype): + """Regular fp32/fp16/bf16 forward GELU_AUX vs PyTorch.""" + 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 + + pre_act = compute_pytorch_reference(A.float(), B.float(), layout) + + output, gelu_input = call_gemm_with_gelu( + A, B, layout, out_dtype=dtype, use_flydsl=True, + ) + assert gelu_input is not None, "GELU_AUX did not return the pre-activation aux." + + no_gelu_out = call_gemm(A, B, layout, out_dtype=dtype, use_flydsl=True) + assert not torch.allclose(output.float(), no_gelu_out.float(), atol=1e-4), ( + "FlyDSL output matches the no-GELU output; GELU epilogue appears inactive." + ) + + assert_gemm_close(gelu_input, pre_act, atol=1e-3, rtol=1e-2) + assert_gemm_close(output, gelu_tanh_ref(pre_act), 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_gelu(M, K, N, layout, fp8_format): + """Tensor-wise FP8 forward GELU_AUX vs PyTorch.""" + 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, + ) + + pre_act = compute_pytorch_reference(A_deq.float(), B_deq.float(), layout) + + output, gelu_input = call_gemm_with_gelu( + A_fp8, B_fp8, layout, out_dtype=torch.float32, use_flydsl=True, + ) + assert gelu_input is not None, "GELU_AUX did not return the pre-activation aux." + + no_gelu_out = call_gemm( + A_fp8, B_fp8, layout, out_dtype=torch.float32, use_flydsl=True, + ) + assert not torch.allclose(output.float(), no_gelu_out.float(), atol=1e-4), ( + "FlyDSL FP8 output matches the no-GELU output; GELU epilogue appears inactive." + ) + + assert_gemm_close(gelu_input, pre_act, atol=5e-3, rtol=1e-2) + assert_gemm_close(output, gelu_tanh_ref(pre_act), atol=5e-3, rtol=1e-2) + + @requires_mxfp8_support @pytest.mark.parametrize("M, K, N", MXFP8_SHAPES) @pytest.mark.parametrize("layout", LAYOUTS) @@ -947,6 +1006,20 @@ def test_flydsl_vs_pytorch_fp8_multidim( "TN", (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), ) + test_flydsl_vs_pytorch_regular_gelu( + 256, + 512, + 256, + "TN", + torch.bfloat16, + ) + test_flydsl_vs_pytorch_fp8_gelu( + 256, + 512, + 256, + "TN", + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), + ) if has_mxfp8_support: test_flydsl_vs_pytorch_mxfp8( diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index 036ff083e..afbff4585 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -21,7 +21,7 @@ import flydsl.compiler as flyc import flydsl.expr as fx from flydsl._mlir.dialects import llvm -from flydsl.expr import arith, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import arith, const_expr, gpu, math, range_constexpr, rocdl from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec @@ -103,18 +103,14 @@ def _compile_kernel(K: int, use_xcd_remap: bool = True, epilogue: str = "DEFAULT DEFAULT plain matmul BIAS + per-output-feature bias vector (indexed by N) - GELU_AUX (reserved) GELU with saved pre-activation aux - GELU_AUX_BIAS (reserved) bias then GELU with saved aux + GELU_AUX GELU(A@B), saving the pre-activation to the Aux output + GELU_AUX_BIAS GELU(A@B + bias), saving the pre-activation to Aux - Only DEFAULT and BIAS are implemented; the GELU modes are accepted so the - store-loop structure and dispatch signature are already in place. + The GELU modes write a second M x N output (pre-activation, tanh-approx + GELU applied to C) for the backward pass. """ if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): raise ValueError(f"Unsupported FP32 epilogue: {epilogue}") - if epilogue in ("GELU_AUX", "GELU_AUX_BIAS"): - raise NotImplementedError( - f"FP32 epilogue {epilogue} is reserved but not yet implemented" - ) has_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") has_gelu = epilogue in ("GELU_AUX", "GELU_AUX_BIAS") @@ -175,6 +171,7 @@ def kernel_gemm( B: fx.Tensor, C: fx.Tensor, Bias: fx.Tensor, + Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, ): @@ -252,6 +249,31 @@ def load_bias(col): fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) return fx.memref_load_vec(reg)[0] + # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N + # fp32 output so the backward pass can recompute the GELU gradient. Same + # tile base as C; DEFAULT gets a dummy 1-element tensor. + if const_expr(has_gelu): + gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) + aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) + aux_store_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + + def gelu_tanh(x): + # tanh-approx GELU (matches PyTorch approximate='tanh' and the + # FlyDSL preshuffle reference), expressed through a non-positive + # exponent so exp() cannot overflow: + # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) + half_f32 = fx.Float32(0.5) + one_f32 = fx.Float32(1.0) + zero_f32 = fx.Float32(0.0) + two_f32 = fx.Float32(2.0) + x3 = x * x * x + y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) + abs_y = fx.Float32(y).maximumf(zero_f32 - y) + e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) + denom = one_f32 + e_neg2abs + numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) + return half_f32 * x * (numerator * (one_f32 / denom)) + PIN_ACC_BASE = 0 def _reg_list(prefix, start, end): @@ -483,11 +505,19 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): c_idx = c_tile_base_elems + row * fx.Index(c_n) + col # Epilogue stages run on the fp32 accumulator; output is fp32 - # so there is no dtype narrowing. GELU will slot in here later. + # so there is no dtype narrowing: + # value = acc [+ bias] (pre-activation) + # GELU_AUX: save pre-activation to Aux, then value = gelu(value) value = Vec(acc)[ii] if const_expr(has_bias): value = value + bias_value + if const_expr(has_gelu): + aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + fx.memref_store_vec(Vec.filled(1, value, fx.Float32), aux_reg) + fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) + value = gelu_tanh(value) + reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) fx.memref_store_vec(Vec.filled(1, value, fx.Float32), reg) fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(c_idx)))) @@ -913,6 +943,7 @@ def launch_gemm( B: fx.Tensor, C: fx.Tensor, Bias: fx.Tensor, + Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, stream: fx.Stream = fx.Stream(None), @@ -924,6 +955,7 @@ def launch_gemm( B, C, Bias, + Aux, c_m, c_n, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, @@ -945,6 +977,7 @@ def fp32_matmul( *, epilogue: str = "DEFAULT", bias: torch.Tensor = None, + aux: torch.Tensor = None, ): """TE-facing TN FP32 GEMM adapter. @@ -990,7 +1023,7 @@ def fp32_matmul( raise ValueError("FlyDSL FP32 GEMM requires contiguous output storage") b_hk = b.transpose(0, 1).contiguous() - doGemm(a, b_hk, c, stream=stream, epilogue=epilogue, bias=bias) + doGemm(a, b_hk, c, stream=stream, epilogue=epilogue, bias=bias, aux=aux) def doGemm( @@ -1001,6 +1034,7 @@ def doGemm( use_xcd_remap: bool = True, epilogue: str = "DEFAULT", bias: torch.Tensor = None, + aux: torch.Tensor = None, ): """Launch the private K-specialized FP32 core. @@ -1041,16 +1075,37 @@ def doGemm( elif bias is not None: raise ValueError(f"FP32 epilogue {epilogue} does not accept a bias tensor") + needs_aux = epilogue in ("GELU_AUX", "GELU_AUX_BIAS") + if needs_aux: + # Pre-activation output for the backward pass: caller-allocated M x N + # fp32, filled in place through its buffer descriptor. + if aux is None: + raise ValueError(f"FP32 epilogue {epilogue} requires an aux output tensor") + if tuple(aux.shape) != (M_runtime, N_runtime): + raise ValueError( + f"FP32 aux shape {tuple(aux.shape)} != {(M_runtime, N_runtime)}" + ) + if aux.dtype != torch.float32: + raise TypeError(f"FP32 aux must be float32, got {aux.dtype}") + if aux.device != A.device: + raise ValueError("aux must be on the same device as A, B, and C") + elif aux is not None: + raise ValueError(f"FP32 epilogue {epilogue} does not accept an aux tensor") + 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) - # DEFAULT keeps the kernel signature uniform with a dummy 1-element bias. + # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + if needs_aux: + Aux_arg = aux.contiguous().view(-1) + else: + Aux_arg = torch.zeros(1, dtype=torch.float32, device=A.device) launch = _cached_launch(int(K_runtime), bool(use_xcd_remap), epilogue) - launch(A_arg, B_arg, C_arg, Bias_arg, M_runtime, N_runtime, stream=stream) + launch(A_arg, B_arg, C_arg, Bias_arg, Aux_arg, M_runtime, N_runtime, stream=stream) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 958eb1961..8080af817 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -21,7 +21,7 @@ import flydsl.compiler as flyc import flydsl.expr as fx from flydsl._mlir.dialects import llvm -from flydsl.expr import arith, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import arith, const_expr, gpu, math, range_constexpr, rocdl from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec @@ -112,18 +112,14 @@ def _compile_kernel( DEFAULT plain matmul BIAS + per-output-feature bias vector (indexed by N) - GELU_AUX (reserved) GELU with saved pre-activation aux - GELU_AUX_BIAS (reserved) bias then GELU with saved aux + GELU_AUX GELU(A@B), saving the pre-activation to the Aux output + GELU_AUX_BIAS GELU(A@B + bias), saving the pre-activation to Aux - Only DEFAULT and BIAS are implemented; the GELU modes are accepted so the - store-loop structure and dispatch signature are already in place. + The GELU modes write a second M x N output (pre-activation, tanh-approx + GELU applied to C) for the backward pass. """ if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): raise ValueError(f"Unsupported FP8 epilogue: {epilogue}") - if epilogue in ("GELU_AUX", "GELU_AUX_BIAS"): - raise NotImplementedError( - f"FP8 epilogue {epilogue} is reserved but not yet implemented" - ) has_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") has_gelu = epilogue in ("GELU_AUX", "GELU_AUX_BIAS") @@ -212,6 +208,7 @@ def kernel_gemm( A_scale_inv: fx.Tensor, B_scale_inv: fx.Tensor, Bias: fx.Tensor, + Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, ): @@ -304,6 +301,34 @@ def load_bias(col): fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) return fx.memref_load_vec(reg)[0] + # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N + # output so the backward pass can recompute the GELU gradient. Same tile + # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. + if const_expr(has_gelu): + gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) + aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) + aux_store_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + output_fx_dtype, + ) + + def gelu_tanh(x): + # tanh-approx GELU (matches PyTorch approximate='tanh' and the + # FlyDSL preshuffle reference), expressed through a non-positive + # exponent so exp() cannot overflow: + # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) + half_f32 = fx.Float32(0.5) + one_f32 = fx.Float32(1.0) + zero_f32 = fx.Float32(0.0) + two_f32 = fx.Float32(2.0) + x3 = x * x * x + y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) + abs_y = fx.Float32(y).maximumf(zero_f32 - y) + e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) + denom = one_f32 + e_neg2abs + numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) + return half_f32 * x * (numerator * (one_f32 / denom)) + PIN_ACC_BASE = 0 def _reg_list(prefix, start, end): @@ -510,11 +535,21 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): # Epilogue stages run on the fp32 accumulator, in order, before # the output-dtype narrowing. Bias is added after the tensor-wise - # output scale, matching the reference epilogue ordering. + # output scale, matching the reference epilogue ordering; GELU_AUX + # then saves the pre-activation and applies GELU to the C value. value = Vec(acc)[ii] * output_scale if const_expr(has_bias): value = value + bias_value + if const_expr(has_gelu): + aux_val = value + if output_dtype != torch.float32: + aux_val = aux_val.to(output_fx_dtype) + aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) + fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) + value = gelu_tanh(value) + if output_dtype != torch.float32: value = value.to(output_fx_dtype) reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) @@ -943,6 +978,7 @@ def launch_gemm( A_scale_inv: fx.Tensor, B_scale_inv: fx.Tensor, Bias: fx.Tensor, + Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, stream: fx.Stream = fx.Stream(None), @@ -956,6 +992,7 @@ def launch_gemm( A_scale_inv, B_scale_inv, Bias, + Aux, c_m, c_n, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, @@ -993,6 +1030,7 @@ def fp8_matmul( *, epilogue: str = "DEFAULT", bias: torch.Tensor = None, + aux: torch.Tensor = None, ): """Launch TN tensor-wise FP8 GEMM using final kernel operand order. @@ -1070,6 +1108,7 @@ def fp8_matmul( stream=stream, epilogue=epilogue, bias=bias, + aux=aux, ) def doGemm( @@ -1082,6 +1121,7 @@ def doGemm( use_xcd_remap: bool = True, epilogue: str = "DEFAULT", bias: torch.Tensor = None, + aux: torch.Tensor = None, ): """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" M_runtime, K_runtime = A.shape @@ -1134,6 +1174,23 @@ def doGemm( elif bias is not None: raise ValueError(f"FP8 epilogue {epilogue} does not accept a bias tensor") + needs_aux = epilogue in ("GELU_AUX", "GELU_AUX_BIAS") + if needs_aux: + # Pre-activation output for the backward pass: caller-allocated M x N, + # same dtype as C, filled in place through its buffer descriptor. + if aux is None: + raise ValueError(f"FP8 epilogue {epilogue} requires an aux output tensor") + if tuple(aux.shape) != (M_runtime, N_runtime): + raise ValueError( + f"FP8 aux shape {tuple(aux.shape)} != {(M_runtime, N_runtime)}" + ) + if aux.dtype != C.dtype: + raise TypeError(f"FP8 aux dtype {aux.dtype} != C dtype {C.dtype}") + if aux.device != A.device: + raise ValueError("aux must be on the same device as A, B, and C") + elif aux is not None: + raise ValueError(f"FP8 epilogue {epilogue} does not accept an aux tensor") + if stream is None: stream = torch.cuda.current_stream() @@ -1142,11 +1199,15 @@ def doGemm( C_arg = C.contiguous().view(-1) A_scale_arg = A_scale_inv.contiguous().view(-1) B_scale_arg = B_scale_inv.contiguous().view(-1) - # DEFAULT keeps the kernel signature uniform with a dummy 1-element bias. + # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + if needs_aux: + Aux_arg = aux.view(-1) + else: + Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) launch = _cached_launch( int(K_runtime), @@ -1163,6 +1224,7 @@ def doGemm( A_scale_arg, B_scale_arg, Bias_arg, + Aux_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 c230bdd80..11685c81e 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -397,6 +397,7 @@ def _run_bf16_gemm( *, output_dtype: torch.dtype, bias=None, + gelu=False, ): """Dispatch BF16 using the original row-major operand allocations. @@ -483,7 +484,8 @@ def _run_bf16_gemm( backend_name=f"BF16 {layout}", ) - epilogue, bias_arg = _resolve_bias(bias, n) + epilogue, bias_arg = _resolve_epilogue(bias, n, gelu=gelu) + aux = torch.empty_like(D) if gelu else None bf16_matmul( a_flydsl, b_flydsl, @@ -494,8 +496,9 @@ def _run_bf16_gemm( k=k, epilogue=epilogue, bias=bias_arg, + aux=aux.view(m, n) if aux is not None else None, ) - return D + return D, aux @@ -508,6 +511,7 @@ def _run_fp16_gemm( *, output_dtype: torch.dtype, bias=None, + gelu=False, ): """Dispatch FP16 using the original row-major operand allocations. @@ -594,7 +598,8 @@ def _run_fp16_gemm( backend_name=f"FP16 {layout}", ) - epilogue, bias_arg = _resolve_bias(bias, n) + epilogue, bias_arg = _resolve_epilogue(bias, n, gelu=gelu) + aux = torch.empty_like(D) if gelu else None fp16_matmul( a_flydsl, b_flydsl, @@ -605,8 +610,9 @@ def _run_fp16_gemm( k=k, epilogue=epilogue, bias=bias_arg, + aux=aux.view(m, n) if aux is not None else None, ) - return D + return D, aux def _run_fp32_gemm( @@ -617,6 +623,7 @@ def _run_fp32_gemm( D, *, bias=None, + gelu=False, ): """Normalize FP32 TN/NN/NT inputs to the current kernel's TN interface. @@ -710,15 +717,17 @@ def _run_fp32_gemm( backend_name="FP32 via TN core", ) - epilogue, bias_arg = _resolve_bias(bias, n) + epilogue, bias_arg = _resolve_epilogue(bias, n, gelu=gelu) + aux = torch.empty_like(D) if gelu else None fp32_matmul( a_tn, b_tn, D.view(m, n), epilogue=epilogue, bias=bias_arg, + aux=aux.view(m, n) if aux is not None else None, ) - return D + return D, aux @@ -1177,6 +1186,7 @@ def _run_fp8( *, output_dtype: torch.dtype, bias=None, + gelu=False, ): """Normalize tensor-wise FP8 storage and invoke the common FP8 core.""" supported_fp8_dtypes = ( @@ -1299,7 +1309,8 @@ def _run_fp8( _fp8_debug(f"derived M={m}, N={n}, K={k}") _fp8_tensor_debug("output/D", D) - epilogue, bias_arg = _resolve_bias(bias, n) + epilogue, bias_arg = _resolve_epilogue(bias, n, gelu=gelu) + aux = torch.empty_like(D) if gelu else None matmul( a_flydsl, a_scale, @@ -1308,8 +1319,9 @@ def _run_fp8( D.view(m, n), epilogue=epilogue, bias=bias_arg, + aux=aux.view(m, n) if aux is not None else None, ) - return D + return D, aux def te_generic_gemm_flydsl( @@ -1417,14 +1429,6 @@ def te_generic_gemm_flydsl( ) return D, None, gelu_input, None - # Fused forward GELU is implemented for MXFP8 only so far; the tensor-wise - # FP8 and regular (fp16/bf16/fp32) paths do not support it yet. - # TODO: extend GELU_AUX to the other FlyDSL GEMM backends. - if gelu: - raise NotImplementedError( - "FlyDSL GEMM fused GELU is currently implemented for MXFP8 only" - ) - if a_kind == "fp8" or b_kind == "fp8": if a_kind != b_kind: raise ValueError( @@ -1443,7 +1447,7 @@ def te_generic_gemm_flydsl( f"got {output_dtype}" ) - D = _run_fp8( + D, gelu_input = _run_fp8( A, transa, B, @@ -1451,8 +1455,9 @@ def te_generic_gemm_flydsl( D, output_dtype=fp8_output_dtypes[output_dtype], bias=bias, + gelu=gelu, ) - return D, None, None, None + return D, None, gelu_input, None if a_kind != "regular" or b_kind != "regular": raise TypeError( @@ -1476,7 +1481,7 @@ def te_generic_gemm_flydsl( "FlyDSL BF16 supports FP16, BF16, or FP32 output, " f"got {output_dtype}" ) - D = _run_bf16_gemm( + D, gelu_input = _run_bf16_gemm( A, transa, B, @@ -1484,8 +1489,9 @@ def te_generic_gemm_flydsl( D, output_dtype=bf16_output_dtypes[output_dtype], bias=bias, + gelu=gelu, ) - return D, None, None, None + return D, None, gelu_input, None if A.dtype == torch.float16 and B.dtype == torch.float16: fp16_output_dtypes = { @@ -1499,7 +1505,7 @@ def te_generic_gemm_flydsl( "FlyDSL FP16 supports FP16, BF16, or FP32 output, " f"got {output_dtype}" ) - D = _run_fp16_gemm( + D, gelu_input = _run_fp16_gemm( A, transa, B, @@ -1507,8 +1513,9 @@ def te_generic_gemm_flydsl( D, output_dtype=fp16_output_dtypes[output_dtype], bias=bias, + gelu=gelu, ) - return D, None, None, None + return D, None, gelu_input, None if A.dtype == torch.float32 and B.dtype == torch.float32: if output_dtype not in (None, tex.DType.kFloat32): @@ -1516,15 +1523,16 @@ def te_generic_gemm_flydsl( "FlyDSL FP32 currently supports only FP32 output, " f"got {output_dtype}" ) - D = _run_fp32_gemm( + D, gelu_input = _run_fp32_gemm( A, transa, B, transb, D, bias=bias, + gelu=gelu, ) - return D, None, None, None + return D, None, gelu_input, None raise NotImplementedError( "FlyDSL GEMM currently supports only MXFP8, tensor-wise E4M3 FP8, " diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 8a4bec15f..089c544e5 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -32,7 +32,7 @@ import flydsl.compiler as flyc import flydsl.expr as fx from flydsl._mlir.dialects import llvm -from flydsl.expr import arith, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import arith, const_expr, gpu, math, range_constexpr, rocdl from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec @@ -121,11 +121,11 @@ def _compile_kernel( DEFAULT plain matmul BIAS + per-output-feature bias vector (indexed by N) - GELU_AUX (reserved) GELU with saved pre-activation aux - GELU_AUX_BIAS (reserved) bias then GELU with saved aux + GELU_AUX GELU(A@B), saving the pre-activation to the Aux output + GELU_AUX_BIAS GELU(A@B + bias), saving the pre-activation to Aux - Only DEFAULT and BIAS are implemented; the GELU modes are accepted so the - store-loop structure and dispatch signature are already in place. + The GELU modes write a second M x N output (pre-activation, tanh-approx + GELU applied to C) for the backward pass. """ if layout not in ("TN", "NN", "NT"): raise ValueError(f"Unsupported half-precision kernel layout: {layout}") @@ -133,10 +133,6 @@ def _compile_kernel( raise ValueError(f"Unsupported MFMA suffix: {mfma_suffix}") if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): raise ValueError(f"Unsupported half-precision epilogue: {epilogue}") - if epilogue in ("GELU_AUX", "GELU_AUX_BIAS"): - raise NotImplementedError( - f"half-precision epilogue {epilogue} is reserved but not yet implemented" - ) has_bias = epilogue in ("BIAS", "GELU_AUX_BIAS") has_gelu = epilogue in ("GELU_AUX", "GELU_AUX_BIAS") @@ -384,6 +380,7 @@ def kernel_gemm( B: fx.Tensor, C: fx.Tensor, Bias: fx.Tensor, + Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, ): @@ -478,6 +475,34 @@ def load_bias(col): fx.copy(bias_ld_atom, fx.slice(bias_div, (None, fx.Int32(col))), reg) return fx.memref_load_vec(reg)[0] + # GELU_AUX saves the pre-activation value (A@B[+bias]) to a second M x N + # output so the backward pass can recompute the GELU gradient. Same tile + # base / store atom shape as C; DEFAULT gets a dummy 1-element tensor. + if const_expr(has_gelu): + gAux = fx.rocdl.make_buffer_tensor(Aux, max_size=True) + aux_div = fx.logical_divide(gAux, fx.make_layout(1, 1)) + aux_store_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy32b() if output_element_bytes == 4 else fx.rocdl.BufferCopy16b(), + output_fx_dtype, + ) + + def gelu_tanh(x): + # tanh-approx GELU (matches PyTorch approximate='tanh' and the + # FlyDSL preshuffle reference), expressed through a non-positive + # exponent so exp() cannot overflow: + # 0.5*x*(1 + tanh(y)), y = sqrt(2/pi)*(x + 0.044715*x^3) + half_f32 = fx.Float32(0.5) + one_f32 = fx.Float32(1.0) + zero_f32 = fx.Float32(0.0) + two_f32 = fx.Float32(2.0) + x3 = x * x * x + y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3) + abs_y = fx.Float32(y).maximumf(zero_f32 - y) + e_neg2abs = math.exp(fx.Float32(-2.0) * abs_y) + denom = one_f32 + e_neg2abs + numerator = (y > zero_f32).select(two_f32, two_f32 * e_neg2abs) + return half_f32 * x * (numerator * (one_f32 / denom)) + PIN_ACC_BASE = 0 def _reg_list(prefix, start, end): @@ -754,11 +779,22 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): c_idx = c_tile_base_elems + row * fx.Index(c_n) + col # Epilogue stages run on the fp32 accumulator, in order, before - # the output-dtype narrowing. GELU will slot in here later. + # the output-dtype narrowing: + # value = acc [+ bias] (pre-activation) + # GELU_AUX: save pre-activation to Aux, then value = gelu(value) value = Vec(acc)[ii] if const_expr(has_bias): value = value + bias_value + if const_expr(has_gelu): + aux_val = value + if const_expr(output_dtype != torch.float32): + aux_val = aux_val.to(output_fx_dtype) + aux_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) + fx.memref_store_vec(Vec.filled(1, aux_val, output_fx_dtype), aux_reg) + fx.copy(aux_store_atom, aux_reg, fx.slice(aux_div, (None, fx.Int32(c_idx)))) + value = gelu_tanh(value) + if const_expr(output_dtype != torch.float32): value = value.to(output_fx_dtype) reg = fx.make_rmem_tensor(fx.make_layout(1, 1), output_fx_dtype) @@ -1197,6 +1233,7 @@ def launch_gemm( B: fx.Tensor, C: fx.Tensor, Bias: fx.Tensor, + Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32, stream: fx.Stream = fx.Stream(None), @@ -1208,6 +1245,7 @@ def launch_gemm( B, C, Bias, + Aux, c_m, c_n, value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, @@ -1248,6 +1286,7 @@ def _half_prec_matmul( stream=None, epilogue: str = "DEFAULT", bias: torch.Tensor = None, + aux: torch.Tensor = None, ): """Validate operands and launch a half-precision TN/NN/NT specialization. @@ -1318,6 +1357,7 @@ def _half_prec_matmul( stream=stream, epilogue=epilogue, bias=bias, + aux=aux, ) @@ -1333,6 +1373,7 @@ def fp16_matmul( stream=None, epilogue: str = "DEFAULT", bias: torch.Tensor = None, + aux: torch.Tensor = None, ): """Launch the wrapper-selected FP16 TN/NN/NT specialization.""" _half_prec_matmul( @@ -1348,6 +1389,7 @@ def fp16_matmul( stream=stream, epilogue=epilogue, bias=bias, + aux=aux, ) @@ -1363,6 +1405,7 @@ def bf16_matmul( stream=None, epilogue: str = "DEFAULT", bias: torch.Tensor = None, + aux: torch.Tensor = None, ): """Launch the wrapper-selected BF16 TN/NN/NT specialization.""" _half_prec_matmul( @@ -1378,6 +1421,7 @@ def bf16_matmul( stream=stream, epilogue=epilogue, bias=bias, + aux=aux, ) @@ -1396,6 +1440,7 @@ def doGemm( use_xcd_remap: bool = True, epilogue: str = "DEFAULT", bias: torch.Tensor = None, + aux: torch.Tensor = None, ): """Launch one cached K/output/layout-specialized half-precision core. @@ -1457,6 +1502,23 @@ def doGemm( elif bias is not None: raise ValueError(f"{label} epilogue {epilogue} does not accept a bias tensor") + needs_aux = epilogue in ("GELU_AUX", "GELU_AUX_BIAS") + if needs_aux: + # Pre-activation output for the backward pass: caller-allocated M x N, + # same dtype as C, filled in place through its buffer descriptor. + if aux is None: + raise ValueError(f"{label} epilogue {epilogue} requires an aux output tensor") + if tuple(aux.shape) != (M_runtime, N_runtime): + raise ValueError( + f"{label} aux shape {tuple(aux.shape)} != {(M_runtime, N_runtime)}" + ) + if aux.dtype != C.dtype: + raise TypeError(f"{label} aux dtype {aux.dtype} != C dtype {C.dtype}") + if aux.device != A.device: + raise ValueError("aux must be on the same device as A, B, and C") + elif aux is not None: + raise ValueError(f"{label} epilogue {epilogue} does not accept an aux tensor") + if stream is None: stream = torch.cuda.current_stream() @@ -1474,17 +1536,22 @@ def doGemm( A_arg = A.view(torch.uint8).view(-1) B_arg = B.view(torch.uint8).view(-1) C_arg = C.view(-1) - # DEFAULT keeps the kernel signature uniform with a dummy 1-element bias. + # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. if needs_bias: Bias_arg = bias.contiguous().view(-1) else: Bias_arg = torch.zeros(1, dtype=torch.float32, device=A.device) + if needs_aux: + Aux_arg = aux.view(-1) + else: + Aux_arg = torch.zeros(1, dtype=C.dtype, device=A.device) launch( A_arg, B_arg, C_arg, Bias_arg, + Aux_arg, M_runtime, N_runtime, stream=stream, From 3b7bdafbaa97e6ac42ebf4436fceb5f40b8380a3 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 01:07:02 +0000 Subject: [PATCH 41/65] Raise FlyDSLUnsupportedError so unsupported GEMM configs fall back general_gemm only catches FlyDSLUnsupportedError to fall back to the C++ backend; NotImplementedError is not a subclass, so any config the FlyDSL dispatcher rejected with NotImplementedError crashed instead of falling back. This was reachable in production: layernorm_mlp fc2_dgrad passes gelu=True, grad=True (DGELU) to general_gemm, which the validator rejects. Under NVTE_USE_FLYDSL=1 on gfx950 that took down the backward pass rather than degrading to the working C++ path. Convert every unsupported-config raise in te_generic_gemm_flydsl / _validate_common_epilogue to FlyDSLUnsupportedError: quantizer, alpha/beta, accumulate, BGRADB (bias+grad), DGELU (gelu+grad), TT layout, and the unsupported input/output dtype guards. These are valid GEMM requests FlyDSL just cannot serve, so falling back is correct. Also fix the now-stale "GELU_AUX is MXFP8 only" docstring. Drop the dead _canonicalize_blas_operands / _canonicalize_blas_pair helpers (never called; the former held the last stray NotImplementedError). Co-Authored-By: Claude --- .../flydsl_kernels/gemm/gemm_wrappers.py | 92 ++++--------------- 1 file changed, 19 insertions(+), 73 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 11685c81e..dcca82fa7 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -97,38 +97,41 @@ def _validate_common_epilogue( ): """Validate features not yet implemented by the FlyDSL GEMM backend. - Fused forward BIAS is supported by all FlyDSL GEMM backends (mxfp8, fp8, - fp16, bf16, fp32). Fused forward GELU_AUX is implemented for MXFP8 only - (the per-backend ``_run_*`` reject it elsewhere). BGRADB (fused bias - gradient) and DGELU (fused GELU gradient) are not implemented anywhere yet. + Fused forward BIAS and forward GELU_AUX are supported by all FlyDSL GEMM + backends (mxfp8, fp8, fp16, bf16, fp32). BGRADB (fused bias gradient) and + DGELU (fused GELU gradient) are not implemented anywhere yet. + + Unsupported configurations raise ``FlyDSLUnsupportedError`` so + ``general_gemm`` transparently falls back to the C++ backend rather than + crashing -- these are valid GEMM requests, just not ones FlyDSL can serve. """ if quantizer is not None: - raise NotImplementedError( + raise FlyDSLUnsupportedError( "FlyDSL GEMM output quantization is not implemented" ) if float(alpha) != 1.0 or float(beta) != 0.0: - raise NotImplementedError( + raise FlyDSLUnsupportedError( "FlyDSL GEMM currently supports only alpha=1 and beta=0" ) # TODO: Add accumulate option if accumulate: - raise NotImplementedError( + raise FlyDSLUnsupportedError( "FlyDSL GEMM accumulation is not implemented" ) # Forward BIAS is implemented across all backends; the fused bias-gradient # (BGRADB) path is not. TODO: add BGRADB to the FlyDSL GEMM backends. if grad and bias is not None and bias.numel() != 0: - raise NotImplementedError( + raise FlyDSLUnsupportedError( "FlyDSL GEMM fused bias gradient (BGRADB) is not implemented" ) # Fused forward GELU (GELU_AUX) is supported; the backward fused GELU # gradient (DGELU) is not. TODO: add DGELU to the FlyDSL GEMM backends. if gelu and grad: - raise NotImplementedError( + raise FlyDSLUnsupportedError( "FlyDSL GEMM fused GELU gradient (DGELU) is not implemented" ) @@ -283,17 +286,6 @@ def _fp8_scale_debug(name: str, scale: torch.Tensor) -> None: ) -def _canonicalize_blas_pair( - A_data: torch.Tensor, - transa: bool, - B_data: torch.Tensor, - transb: bool, -): - """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: """Flatten all leading dimensions while preserving the final dimension.""" if t.ndim < 2: @@ -312,52 +304,6 @@ def _flatten_columnwise(t: torch.Tensor, name: str) -> torch.Tensor: return t.reshape(t.shape[0], -1) -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. - - TE's generic GEMM interface follows BLAS column-major interpretation. - FlyDSL kernels consume ordinary row-major matrices: - - a_flydsl: [M, K] - b_flydsl: [K, N] - - Operand ownership is swapped without creating tensor transpose views: - - a_flydsl = B - b_flydsl = A - """ - if transa and transb: - raise NotImplementedError( - "FlyDSL GEMM does not support transa=True, transb=True (TT)" - ) - - 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, - transb, - ) - - 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 _validate_or_allocate_output( D, *, @@ -1375,7 +1321,7 @@ def te_generic_gemm_flydsl( del bulk_overlap if transa and transb: - raise NotImplementedError( + raise FlyDSLUnsupportedError( "FlyDSL GEMM does not support transa=True, transb=True (TT)" ) @@ -1412,7 +1358,7 @@ def te_generic_gemm_flydsl( tex.DType.kFloat32: torch.float32, } if output_dtype not in mxfp8_output_dtypes: - raise NotImplementedError( + raise FlyDSLUnsupportedError( "FlyDSL MXFP8 supports FP16, BF16, or FP32 output, " f"got {output_dtype}" ) @@ -1442,7 +1388,7 @@ def te_generic_gemm_flydsl( tex.DType.kFloat32: torch.float32, } if output_dtype not in fp8_output_dtypes: - raise NotImplementedError( + raise FlyDSLUnsupportedError( "FlyDSL tensor-wise FP8 supports FP16, BF16, or FP32 output, " f"got {output_dtype}" ) @@ -1477,7 +1423,7 @@ def te_generic_gemm_flydsl( tex.DType.kFloat32: torch.float32, } if output_dtype not in bf16_output_dtypes: - raise NotImplementedError( + raise FlyDSLUnsupportedError( "FlyDSL BF16 supports FP16, BF16, or FP32 output, " f"got {output_dtype}" ) @@ -1501,7 +1447,7 @@ def te_generic_gemm_flydsl( tex.DType.kFloat32: torch.float32, } if output_dtype not in fp16_output_dtypes: - raise NotImplementedError( + raise FlyDSLUnsupportedError( "FlyDSL FP16 supports FP16, BF16, or FP32 output, " f"got {output_dtype}" ) @@ -1519,7 +1465,7 @@ def te_generic_gemm_flydsl( if A.dtype == torch.float32 and B.dtype == torch.float32: if output_dtype not in (None, tex.DType.kFloat32): - raise NotImplementedError( + raise FlyDSLUnsupportedError( "FlyDSL FP32 currently supports only FP32 output, " f"got {output_dtype}" ) @@ -1534,7 +1480,7 @@ def te_generic_gemm_flydsl( ) return D, None, gelu_input, None - raise NotImplementedError( + raise FlyDSLUnsupportedError( "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 c131f20f765195081ff44a173d2f90598444274f Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 01:21:30 +0000 Subject: [PATCH 42/65] Reject comm+GEMM overlap instead of silently dropping it te_generic_gemm_flydsl del'd comm_overlap / comm_type / extra_output / bulk_overlap and returned extra_output=None. But extra_output is the reduce-scatter destination that Linear (linear.py) passes for tensor-parallel comm+GEMM overlap: general_gemm assigns the None return back over the caller's tensor, so the reduce-scatter silently never happens -- wrong results or a downstream None deref, not an error. Replace the four drops with a guard that raises FlyDSLUnsupportedError when any overlap argument is set, so general_gemm falls back to the C++ backend that actually performs the collective. workspace / workspaceSize / bias_type / use_split_accumulator remain dropped -- they are hints or unused given the other guards. Co-Authored-By: Claude --- .../flydsl_kernels/gemm/gemm_wrappers.py | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index dcca82fa7..19c82fd14 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -1310,15 +1310,29 @@ def te_generic_gemm_flydsl( - FP16 input with FP16, BF16, or FP32 output - FP32 input with FP32 output """ + # These are hints or unused given the epilogue/dtype guards below: bias_type + # is inferred, gelu_in is only for DGELU (rejected), and FlyDSL has no + # separate workspace or split-accumulator path. 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 + + # Comm+GEMM overlap (tensor-parallel) is not implemented. These carry real + # side effects -- extra_output is the reduce-scatter destination -- so they + # must be rejected, not dropped: silently discarding them would no-op the + # collective and return wrong results. Raise FlyDSLUnsupportedError so + # general_gemm falls back to the C++ backend. + if ( + comm_overlap is not None + or comm_type is not None + or extra_output is not None + or bulk_overlap + ): + raise FlyDSLUnsupportedError( + "FlyDSL GEMM does not support comm+GEMM overlap" + ) if transa and transb: raise FlyDSLUnsupportedError( From 343b8c3074cc7044c18b88c05359fde86b25e6de Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 01:32:11 +0000 Subject: [PATCH 43/65] Reject oversized GEMM operands that overflow the int32 launch signature FlyDSL packs flat operand byte views into an int32 launch signature, so any single operand >= 2 GiB overflows the 'i' pack format and aborts the run with struct.error at launch time. struct.error is not a FlyDSLUnsupportedError, so it escapes general_gemm's fallback and hard-fails -- these shapes cannot be skipped. This hits real large-vocab / large-MLP layers at MBS=4 (e.g. Llama-3.1-405B M=32768,K=53248: A is 3.49 GB; Qwen2.5-72B C is 3.87 GB). Add a shared require_launch_size helper in gemm_common_utils.py that raises FlyDSLUnsupportedError for any operand exceeding 2**31 - 1 bytes, and call it in every doGemm right after the block-tiling guard (mxfp8 checks A/B/As/Bs/C; fp8, half-precision, and fp32 check A/B/C). Oversized shapes now fall back to the C++ backend instead of aborting. Co-Authored-By: Claude --- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 3 ++- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 3 ++- .../flydsl_kernels/gemm/gemm_common_utils.py | 26 +++++++++++++++++++ .../flydsl_kernels/gemm/half_prec_gemm.py | 3 ++- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 6 ++++- 5 files changed, 37 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index afbff4585..c9258515f 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -26,7 +26,7 @@ from flydsl.expr.typing import Vector as Vec # Transformer Engine-local FlyDSL utilities. -from .gemm_common_utils import require_block_tiling +from .gemm_common_utils import require_block_tiling, require_launch_size from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -1055,6 +1055,7 @@ def doGemm( block_k=_BLOCK_K, label="FP32 GEMM", ) + require_launch_size("FP32 GEMM", ("A", A), ("B", B), ("C", C)) assert C.shape == (M_runtime, N_runtime) if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index 8080af817..d43718ca7 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -25,7 +25,7 @@ from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec -from .gemm_common_utils import require_block_tiling +from .gemm_common_utils import require_block_tiling, require_launch_size # Transformer Engine-local FlyDSL utilities. from .fp8_gemm_utils import ( @@ -1150,6 +1150,7 @@ def doGemm( block_k=_BLOCK_K, label="FP8 GEMM", ) + require_launch_size("FP8 GEMM", ("A", A), ("B", B), ("C", C)) 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/gemm_common_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py index 0b3e42ce6..ce3c75e35 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py @@ -43,6 +43,32 @@ def require_block_tiling(m, n, k, *, block_m, block_n, block_k, label, min_k_til return num_k_tiles +# FlyDSL packs flat operand byte views into an int32 launch signature, so any +# single operand of >= 2 GiB overflows the 'i' pack format at launch. Guard for +# it up front so oversized shapes fall back to the default backend instead of +# aborting with an unrecoverable struct.error. +_MAX_LAUNCH_BYTES = 2**31 - 1 + + +def require_launch_size(label, *tensors): + """Reject operands whose byte size exceeds the int32 launch-argument limit. + + ``tensors`` is an iterable of ``(name, tensor)`` pairs. Raises + ``FlyDSLUnsupportedError`` (so ``general_gemm`` falls back to the C++ + backend) for the first operand at or above ``_MAX_LAUNCH_BYTES``; skips + ``None`` tensors (e.g. an absent bias/aux). + """ + for name, t in tensors: + if t is None: + continue + nbytes = t.numel() * t.element_size() + if nbytes > _MAX_LAUNCH_BYTES: + raise FlyDSLUnsupportedError( + f"FlyDSL {label} operand {name} is {nbytes} bytes, which exceeds " + f"the int32 launch-argument limit of {_MAX_LAUNCH_BYTES}" + ) + + def cdiv(numer: int, denom: int) -> int: return (numer + denom - 1) // denom diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 089c544e5..9b29274a8 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -38,7 +38,7 @@ # Transformer Engine-local FlyDSL utilities. from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import require_block_tiling +from .gemm_common_utils import require_block_tiling, require_launch_size from .fp16_gemm_utils import ( G2SLoader, S2RLoader, @@ -1478,6 +1478,7 @@ def doGemm( block_k=_BLOCK_K, label=f"{label} GEMM", ) + require_launch_size(f"{label} GEMM", ("A", A), ("B", B), ("C", C)) if tuple(C.shape) != (M_runtime, N_runtime): raise ValueError( diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 2a78f7f44..75bfdcc2d 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -35,7 +35,7 @@ # Transformer Engine-local FlyDSL utilities. from .exceptions import FlyDSLUnsupportedError -from .gemm_common_utils import require_block_tiling +from .gemm_common_utils import require_block_tiling, require_launch_size from .fp8_gemm_utils import ( G2SLoader, S2RLoader, @@ -1601,6 +1601,10 @@ def do_gemm( block_k=_BLOCK_K, label=f"MXFP8 {layout} GEMM", ) + require_launch_size( + f"MXFP8 {layout} GEMM", + ("A", A), ("B", B), ("As", As), ("Bs", Bs), ("C", C), + ) expected_as = (K_runtime // _BLOCK_K, M_runtime) expected_bs = (K_runtime // _BLOCK_K, N_runtime) From 4402b562a7e62d611c82cfbdd19b1c5bee181b96 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 02:07:01 +0000 Subject: [PATCH 44/65] Address FlyDSL GEMM test review: copyright, CI enrollment, drop smoke block - Copyright header uses a single current year (2026), not a range, since the file is new in this PR. - Enroll flydsl_kernels/test_gemm.py in ci/pytorch.sh, gated on the gfx950 MXFP8 availability check (the FlyDSL kernels currently require gfx950, so on other archs general_gemm won't select FlyDSL and the suite would silently exercise the C++ backend). The suite was previously never run in CI. - Remove the __main__ smoke-test block: it duplicated pytest, bypassed the cleanup_env fixture, and leaked NVTE_USE_FLYDSL=1 into the process. The broader "unify with the Triton test into one backend-parametrized suite plus a single backend-selection env" refactor is deferred until #667 lands. Co-Authored-By: Claude --- ci/pytorch.sh | 4 ++ tests/pytorch/flydsl_kernels/test_gemm.py | 76 +---------------------- 2 files changed, 5 insertions(+), 75 deletions(-) diff --git a/ci/pytorch.sh b/ci/pytorch.sh index 38cef08ce..a3e42ebf7 100755 --- a/ci/pytorch.sh +++ b/ci/pytorch.sh @@ -97,6 +97,10 @@ run_test_config(){ run_default_fa 1 triton_kernels/test_utils.py NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa 1 triton_kernels/test_norms.py NVTE_ROCM_ENABLE_MXFP8=1 NVTE_TEST_TRITON_AUTOTUNE=1 run_default_fa_lbl "autotune" 3 triton_kernels/test_norms.py + # The FlyDSL GEMM kernels currently require gfx950; on other archs + # general_gemm won't select FlyDSL, so gate on the same gfx950 check the + # MXFP8 tests use to avoid silently exercising the C++ backend instead. + check_mxfp8_supported && NVTE_USE_FLYDSL=1 NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa_lbl "flydsl" 1 flydsl_kernels/test_gemm.py run_default_fa 1 test_parallel_cross_entropy.py NVTE_USE_DEQUANTIZE_TRITON=1 NVTE_USE_CAST_TRANSPOSE_TRITON=1 NVTE_USE_RMSNORM_TRITON=1 NVTE_USE_LAYERNORM_TRITON=1 run_default_fa_lbl "triton" 3 test_numerics.py NVTE_USE_CAST_TRANSPOSE_TRITON=1 NVTE_USE_RMSNORM_TRITON=1 run_default_fa_lbl "triton" 1 test_fusible_ops.py diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py index 2605af856..eb4a5665a 100644 --- a/tests/pytorch/flydsl_kernels/test_gemm.py +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # License for AMD contributions = MIT. See LICENSE for more information @@ -971,77 +971,3 @@ def test_flydsl_vs_pytorch_fp8_multidim( 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), - ) - test_flydsl_vs_pytorch_regular_bias( - 256, - 512, - 256, - "TN", - torch.bfloat16, - ) - test_flydsl_vs_pytorch_fp8_bias( - 256, - 512, - 256, - "TN", - (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), - ) - test_flydsl_vs_pytorch_regular_gelu( - 256, - 512, - 256, - "TN", - torch.bfloat16, - ) - test_flydsl_vs_pytorch_fp8_gelu( - 256, - 512, - 256, - "TN", - (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), - ) - - if has_mxfp8_support: - test_flydsl_vs_pytorch_mxfp8( - 256, - 512, - 256, - "TN", - (tex.DType.kFloat8E5M2, tex.DType.kFloat8E4M3), - ) - test_flydsl_vs_pytorch_mxfp8_bias( - 256, - 512, - 256, - "TN", - (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), - ) - test_flydsl_vs_pytorch_mxfp8_gelu( - 256, - 512, - 256, - "TN", - (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), - ) - - print("All FlyDSL GEMM smoke tests passed!") From 1a0ab36b1abb5d179d3cd4e118f14dfb1cfaade8 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 02:28:44 +0000 Subject: [PATCH 45/65] Address review: __init__ header, DSL divmod, import fallback, formatting Batch of quick review fixes: - flydsl_kernels/__init__.py: add the AMD copyright header and module docstring (was headerless, tripping pylint C0114) and drop the eager 'from . import gemm' so touching the package no longer pulls in the whole FlyDSL stack. The consumer's 'from ..flydsl_kernels.gemm import ...' keeps the lazy-import property. - half_prec/fp8/fp32 kernels: import the DSL-safe divmod from the utils layer so the use_xcd_remap=False branch (pid = divmod(block_idx, ...)) no longer calls the Python builtin, which rejects DSL scalars. - cpp_extensions/gemm.py: move the lazy flydsl import inside the try and catch ImportError alongside FlyDSLUnsupportedError, so NVTE_USE_FLYDSL=1 on a wheel built without flydsl degrades to the default backend instead of a bare ImportError. Also reflow the use_gemm_flydsl condition per black. - setup.py: collapse the FlyDSL install-req guard to one line and use append. - Run black (repo config) over the FlyDSL kernel package and test file; strip trailing whitespace. Shared files (setup.py, gemm.py) are touched only on the changed lines, not wholesale reformatted. Co-Authored-By: Claude --- setup.py | 8 +- tests/pytorch/flydsl_kernels/test_gemm.py | 136 +++++++--- .../pytorch/cpp_extensions/gemm.py | 25 +- .../pytorch/flydsl_kernels/__init__.py | 7 +- .../pytorch/flydsl_kernels/gemm/exceptions.py | 1 + .../flydsl_kernels/gemm/fp16_gemm_utils.py | 10 +- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 68 +++-- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 95 +++---- .../flydsl_kernels/gemm/fp8_gemm_utils.py | 23 +- .../flydsl_kernels/gemm/gemm_wrappers.py | 194 +++++--------- .../flydsl_kernels/gemm/half_prec_gemm.py | 136 +++++----- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 236 +++++++----------- 12 files changed, 426 insertions(+), 513 deletions(-) diff --git a/setup.py b/setup.py index 81ed18830..b7dd1a29c 100644 --- a/setup.py +++ b/setup.py @@ -196,12 +196,8 @@ 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==0.3.0"]) + if rocm_build() and "pytorch" in frameworks and bool(int(os.getenv("NVTE_USE_FLYDSL", "0"))): + install_reqs.append("flydsl==0.3.0") # Framework-specific requirements if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py index eb4a5665a..9326de034 100644 --- a/tests/pytorch/flydsl_kernels/test_gemm.py +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -96,6 +96,7 @@ # --- Fixtures ----------------------------------------------------------------- + @pytest.fixture(autouse=True) def cleanup_env(): """Save and restore FlyDSL-related environment variables between tests.""" @@ -117,6 +118,7 @@ def cleanup_env(): # --- Helpers ------------------------------------------------------------------ + def get_shapes(layout, M, K, N): """Return the A/B storage shapes used by TE's public GEMM tests.""" if layout == "TN": @@ -283,6 +285,7 @@ def assert_gemm_close(actual, expected, *, atol, rtol): # Approach 1: FlyDSL vs PyTorch torch.matmul reference # ============================================================================== + @pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) @pytest.mark.parametrize("layout", LAYOUTS) @pytest.mark.parametrize( @@ -390,6 +393,7 @@ def test_flydsl_vs_pytorch_mxfp8(M, K, N, layout, fp8_format): # 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( @@ -516,6 +520,7 @@ def test_flydsl_vs_cpp_mxfp8(M, K, N, layout, fp8_format): # the output, catching a silent decay to DEFAULT) and a vs-cpp cross-check. # ============================================================================== + @pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) @pytest.mark.parametrize("layout", LAYOUTS) @pytest.mark.parametrize("dtype", REGULAR_DTYPES, ids=["fp32", "fp16", "bf16"]) @@ -532,14 +537,19 @@ def test_flydsl_vs_pytorch_regular_bias(M, K, N, layout, dtype): bias = torch.randn(out_features, dtype=dtype, device="cuda") output, bias_grad = call_gemm_with_bias( - A, B, layout, out_dtype=dtype, bias=bias, use_flydsl=True, + A, + B, + layout, + out_dtype=dtype, + bias=bias, + use_flydsl=True, ) assert bias_grad is None no_bias_out = call_gemm(A, B, layout, out_dtype=dtype, use_flydsl=True) - assert not torch.allclose(output.float(), no_bias_out.float(), atol=1e-4), ( - "FlyDSL output matches the no-bias output; BIAS epilogue appears inactive." - ) + assert not torch.allclose( + output.float(), no_bias_out.float(), atol=1e-4 + ), "FlyDSL output matches the no-bias output; BIAS epilogue appears inactive." expected = expected_ab + bias.float() assert_gemm_close(output, expected, atol=1e-3, rtol=1e-2) @@ -560,10 +570,20 @@ def test_flydsl_vs_cpp_regular_bias(M, K, N, layout, dtype): bias = torch.randn(out_features, dtype=dtype, device="cuda") flydsl_out, _ = call_gemm_with_bias( - A, B, layout, out_dtype=dtype, bias=bias, use_flydsl=True, + A, + B, + layout, + out_dtype=dtype, + bias=bias, + use_flydsl=True, ) cpp_out, _ = call_gemm_with_bias( - A, B, layout, out_dtype=dtype, bias=bias, use_flydsl=False, + A, + B, + layout, + out_dtype=dtype, + bias=bias, + use_flydsl=False, ) assert_gemm_close(flydsl_out, cpp_out, atol=1e-3, rtol=1e-2) @@ -578,7 +598,12 @@ def test_flydsl_vs_pytorch_fp8_bias(M, K, N, layout, fp8_format): 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, + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, ) expected_ab = compute_pytorch_reference(A_deq.float(), B_deq.float(), layout) @@ -586,16 +611,25 @@ def test_flydsl_vs_pytorch_fp8_bias(M, K, N, layout, fp8_format): bias = torch.randn(out_features, dtype=torch.float32, device="cuda") output, bias_grad = call_gemm_with_bias( - A_fp8, B_fp8, layout, out_dtype=torch.float32, bias=bias, use_flydsl=True, + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + bias=bias, + use_flydsl=True, ) assert bias_grad is None no_bias_out = call_gemm( - A_fp8, B_fp8, layout, out_dtype=torch.float32, use_flydsl=True, - ) - assert not torch.allclose(output.float(), no_bias_out.float(), atol=1e-4), ( - "FlyDSL FP8 output matches the no-bias output; BIAS epilogue appears inactive." + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, ) + assert not torch.allclose( + output.float(), no_bias_out.float(), atol=1e-4 + ), "FlyDSL FP8 output matches the no-bias output; BIAS epilogue appears inactive." expected = expected_ab + bias.float() assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) @@ -610,19 +644,36 @@ def test_flydsl_vs_cpp_fp8_bias(M, K, N, layout, fp8_format): 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, + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, ) out_features = compute_pytorch_reference( - A_deq.float(), B_deq.float(), layout, + A_deq.float(), + B_deq.float(), + layout, ).shape[-1] bias = torch.randn(out_features, dtype=torch.float32, device="cuda") flydsl_out, _ = call_gemm_with_bias( - A_fp8, B_fp8, layout, out_dtype=torch.float32, bias=bias, use_flydsl=True, + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + bias=bias, + use_flydsl=True, ) cpp_out, _ = call_gemm_with_bias( - A_fp8, B_fp8, layout, out_dtype=torch.float32, bias=bias, use_flydsl=False, + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + bias=bias, + use_flydsl=False, ) assert_gemm_close(flydsl_out, cpp_out, atol=5e-3, rtol=1e-2) @@ -675,10 +726,9 @@ def test_flydsl_vs_pytorch_mxfp8_bias(M, K, N, layout, fp8_format): out_dtype=torch.float32, use_flydsl=True, ) - assert not torch.allclose(output.float(), no_bias_out.float(), atol=1e-4), ( - "FlyDSL MXFP8 output matches the no-bias output; " - "the BIAS epilogue appears inactive." - ) + assert not torch.allclose( + output.float(), no_bias_out.float(), atol=1e-4 + ), "FlyDSL MXFP8 output matches the no-bias output; the BIAS epilogue appears inactive." expected = expected_ab + bias.float() assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) @@ -745,6 +795,7 @@ def test_flydsl_vs_cpp_mxfp8_bias(M, K, N, layout, fp8_format): # aux == pre_act, with a guard that GELU actually changes the output. # ============================================================================== + @pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) @pytest.mark.parametrize("layout", LAYOUTS) @pytest.mark.parametrize("dtype", REGULAR_DTYPES, ids=["fp32", "fp16", "bf16"]) @@ -759,14 +810,18 @@ def test_flydsl_vs_pytorch_regular_gelu(M, K, N, layout, dtype): pre_act = compute_pytorch_reference(A.float(), B.float(), layout) output, gelu_input = call_gemm_with_gelu( - A, B, layout, out_dtype=dtype, use_flydsl=True, + A, + B, + layout, + out_dtype=dtype, + use_flydsl=True, ) assert gelu_input is not None, "GELU_AUX did not return the pre-activation aux." no_gelu_out = call_gemm(A, B, layout, out_dtype=dtype, use_flydsl=True) - assert not torch.allclose(output.float(), no_gelu_out.float(), atol=1e-4), ( - "FlyDSL output matches the no-GELU output; GELU epilogue appears inactive." - ) + assert not torch.allclose( + output.float(), no_gelu_out.float(), atol=1e-4 + ), "FlyDSL output matches the no-GELU output; GELU epilogue appears inactive." assert_gemm_close(gelu_input, pre_act, atol=1e-3, rtol=1e-2) assert_gemm_close(output, gelu_tanh_ref(pre_act), atol=1e-3, rtol=1e-2) @@ -781,22 +836,35 @@ def test_flydsl_vs_pytorch_fp8_gelu(M, K, N, layout, fp8_format): 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, + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, ) pre_act = compute_pytorch_reference(A_deq.float(), B_deq.float(), layout) output, gelu_input = call_gemm_with_gelu( - A_fp8, B_fp8, layout, out_dtype=torch.float32, use_flydsl=True, + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, ) assert gelu_input is not None, "GELU_AUX did not return the pre-activation aux." no_gelu_out = call_gemm( - A_fp8, B_fp8, layout, out_dtype=torch.float32, use_flydsl=True, - ) - assert not torch.allclose(output.float(), no_gelu_out.float(), atol=1e-4), ( - "FlyDSL FP8 output matches the no-GELU output; GELU epilogue appears inactive." + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, ) + assert not torch.allclose( + output.float(), no_gelu_out.float(), atol=1e-4 + ), "FlyDSL FP8 output matches the no-GELU output; GELU epilogue appears inactive." assert_gemm_close(gelu_input, pre_act, atol=5e-3, rtol=1e-2) assert_gemm_close(output, gelu_tanh_ref(pre_act), atol=5e-3, rtol=1e-2) @@ -844,10 +912,9 @@ def test_flydsl_vs_pytorch_mxfp8_gelu(M, K, N, layout, fp8_format): out_dtype=torch.float32, use_flydsl=True, ) - assert not torch.allclose(output.float(), no_gelu_out.float(), atol=1e-4), ( - "FlyDSL MXFP8 output matches the no-GELU output; " - "the GELU epilogue appears inactive." - ) + assert not torch.allclose( + output.float(), no_gelu_out.float(), atol=1e-4 + ), "FlyDSL MXFP8 output matches the no-GELU output; the GELU epilogue appears inactive." # Aux is the pre-activation (A@B); output is gelu(A@B). assert_gemm_close(gelu_input, pre_act, atol=5e-3, rtol=1e-2) @@ -901,6 +968,7 @@ def test_flydsl_vs_pytorch_mxfp8_gelu_bias(M, K, N, layout, fp8_format): # Batched multidimensional FP8 coverage # ============================================================================== + @pytest.mark.parametrize( "batch_size, M, K, N", [ diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 0e5b03523..e8a62fe56 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -494,22 +494,27 @@ 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 get_device_compute_capability() == (9, 5) - 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 ( - FlyDSLUnsupportedError, - te_generic_gemm_flydsl, - ) - try: + # Lazy import keeps FlyDSL off the normal Transformer Engine + # import path. It is done inside the try so a wheel built without + # flydsl (NVTE_USE_FLYDSL unset at build time) degrades to the + # default backend instead of raising a bare ImportError. + from ..flydsl_kernels.gemm import ( + FlyDSLUnsupportedError, + te_generic_gemm_flydsl, + ) + out, bias_grad, gelu_input, extra_output = te_generic_gemm_flydsl( *args, **kwargs, ) - except FlyDSLUnsupportedError as exc: + except (FlyDSLUnsupportedError, ImportError) as exc: warn_fallback = os.environ.get( "NVTE_FLYDSL_GEMM_WARN_FALLBACK", "0", diff --git a/transformer_engine/pytorch/flydsl_kernels/__init__.py b/transformer_engine/pytorch/flydsl_kernels/__init__.py index c64b988c6..0057be1e2 100644 --- a/transformer_engine/pytorch/flydsl_kernels/__init__.py +++ b/transformer_engine/pytorch/flydsl_kernels/__init__.py @@ -1,3 +1,4 @@ -from . import gemm - -__all__ = ["gemm"] +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +"""FlyDSL kernels for ROCm gfx950 GEMM replacement.""" diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py b/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py index b7fc19a23..c4a4018c1 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py @@ -2,5 +2,6 @@ # # See LICENSE for license information. + class FlyDSLUnsupportedError(RuntimeError): """The GEMM request is valid but unsupported by the available FlyDSL kernels.""" 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 d6cc8f221..76628bdbd 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py @@ -50,9 +50,7 @@ def compute_global_swizzle( n_waves = fx.block_dim.x // 64 for round in range_constexpr(n_rounds): if const_expr(preshuffled): - raise AssertionError( - "16-bit first-pass port does not support preshuffled operands" - ) + raise AssertionError("16-bit 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) @@ -88,11 +86,7 @@ def compute_global_transpose_swizzle( # 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 - ) + offsets.append(logical_k * leading_dim_bytes + slice_idx * 64 * 2 + logical_x_bytes) return offsets diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index c9258515f..92b538026 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -31,11 +31,12 @@ G2SLoader, S2RLoader, compute_global_swizzle, + divmod, make_byte_buffer_tensor as make_fp32_byte_buffer_tensor, pack_i32x4_i32x8, swizzle_128, xcd_swizzle, - barrier + barrier, ) @@ -51,7 +52,7 @@ WARP_SIZE = 64 NUM_WAVES = NUM_THREADS // WARP_SIZE -SUBTILE_M = 64 +SUBTILE_M = 64 SUBTILE_N = 64 MFMA_M = 16 @@ -144,7 +145,9 @@ def _compile_kernel(K: int, use_xcd_remap: bool = True, epilogue: str = "DEFAULT 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" + 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 @@ -216,8 +219,12 @@ def kernel_gemm( # 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) + 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) @@ -365,11 +372,15 @@ def hot_loop_scheduler_q_prefetch_4n(): 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 = (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) + 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): @@ -402,7 +413,9 @@ def load_b_frag(lds_b, local_row, half): 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 + 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 @@ -415,16 +428,11 @@ def _pinned_fp32_mfma_once(acc_idx, a_k4, b_k4): 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}}}" - ), + "v_mfma_f32_16x16x4_f32 " + f"a[{acc_pin}:{acc_pin + 3}], " + "$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]", + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", has_side_effects=True, ) @@ -522,7 +530,6 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): fx.memref_store_vec(Vec.filled(1, value, fx.Float32), reg) fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(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: @@ -818,7 +825,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): # # Finalize accumulators in their own physical AGPR slots, but delay # each AGPR read/store until several independent final MFMAs have - # been issued. + # been issued. # # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, # MFMA 4, drain 1, MFMA 5, drain 2, ... @@ -936,7 +943,6 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): ) hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) - @flyc.jit def launch_gemm( A: fx.Tensor, @@ -963,12 +969,12 @@ def launch_gemm( return launch_gemm + @functools.lru_cache(maxsize=None) def _cached_launch(K: int, use_xcd_remap: bool = True, epilogue: str = "DEFAULT"): return _compile_kernel(K, use_xcd_remap=use_xcd_remap, epilogue=epilogue) - def fp32_matmul( a: torch.Tensor, b: torch.Tensor, @@ -993,16 +999,13 @@ def fp32_matmul( """ 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)}" + f"FlyDSL FP32 TN expects rank-2 operands, got A{tuple(a.shape)} 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)}" - ) + 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, " @@ -1016,8 +1019,7 @@ def fp32_matmul( ) 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}" + 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 FP32 GEMM requires contiguous output storage") @@ -1068,9 +1070,7 @@ def doGemm( if bias.dtype != torch.float32: raise TypeError(f"FP32 bias must be float32, got {bias.dtype}") if bias.numel() != N_runtime: - raise ValueError( - f"FP32 bias length {bias.numel()} != N (out_features) {N_runtime}" - ) + raise ValueError(f"FP32 bias length {bias.numel()} != N (out_features) {N_runtime}") if bias.device != A.device: raise ValueError("bias must be on the same device as A, B, and C") elif bias is not None: @@ -1083,9 +1083,7 @@ def doGemm( if aux is None: raise ValueError(f"FP32 epilogue {epilogue} requires an aux output tensor") if tuple(aux.shape) != (M_runtime, N_runtime): - raise ValueError( - f"FP32 aux shape {tuple(aux.shape)} != {(M_runtime, N_runtime)}" - ) + raise ValueError(f"FP32 aux shape {tuple(aux.shape)} != {(M_runtime, N_runtime)}") if aux.dtype != torch.float32: raise TypeError(f"FP32 aux must be float32, got {aux.dtype}") if aux.device != A.device: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index d43718ca7..be2a8fc9d 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -32,15 +32,15 @@ G2SLoader, S2RLoader, compute_global_swizzle, + divmod, make_fp8_buffer_tensor, pack_i32x4_i32x8, swizzle_128, xcd_swizzle, - barrier + barrier, ) - _BLOCK_M = 256 _BLOCK_N = 256 _BLOCK_K = 128 @@ -53,7 +53,7 @@ WARP_SIZE = 64 NUM_WAVES = NUM_THREADS // WARP_SIZE -SUBTILE_M = 64 +SUBTILE_M = 64 SUBTILE_N = 64 MFMA_M = 16 @@ -149,8 +149,7 @@ def _compile_kernel( output_fx_dtype = fx.Float32 else: raise TypeError( - "FlyDSL FP8 supports only float16, bfloat16, and float32 " - f"outputs, got {output_dtype}" + f"FlyDSL FP8 supports only float16, bfloat16, and float32 outputs, got {output_dtype}" ) NUM_THREADS = 256 WARP_SIZE = 64 @@ -181,7 +180,9 @@ def _compile_kernel( 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" + 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) @@ -451,7 +452,9 @@ def load_b_frag(lds_b, local_row, half): 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 + 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.""" @@ -462,17 +465,12 @@ def pinned_mfma(acc_idx, a_frag, b_frag): 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}}}" - ), + "v_mfma_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + "$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}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", has_side_effects=True, ) @@ -486,17 +484,12 @@ def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): 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}}}" - ), + "v_mfma_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + "$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}}},~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}", has_side_effects=True, ) @@ -556,7 +549,6 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(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: @@ -851,7 +843,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): # # Finalize accumulators in their own physical AGPR slots, but delay # each AGPR read/store until several independent final MFMAs have - # been issued. + # been issued. # # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, # MFMA 4, drain 1, MFMA 5, drain 2, ... @@ -969,7 +961,6 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): ) hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) - @flyc.jit def launch_gemm( A: fx.Tensor, @@ -1000,6 +991,7 @@ def launch_gemm( return launch_gemm + @functools.lru_cache(maxsize=None) def _cached_launch( K: int, @@ -1019,7 +1011,6 @@ def _cached_launch( ) - def fp8_matmul( a: torch.Tensor, a_scale_inv: torch.Tensor, @@ -1050,8 +1041,7 @@ def fp8_matmul( 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)}" + f"FlyDSL FP8 TN expects rank-2 operands, got A{tuple(a.shape)} and B{tuple(b.shape)}" ) supported_fp8_dtypes = ( @@ -1060,16 +1050,13 @@ def fp8_matmul( ) 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}" + f"FlyDSL FP8 GEMM expects E4M3 or E5M2 payloads, 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)}" - ) + 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), @@ -1087,17 +1074,14 @@ def fp8_matmul( 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}" + f"FlyDSL FP8 supports only float16, bfloat16, and float32 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" - ) + raise ValueError("A, B, inverse scales, and C must be on the same device") doGemm( a, @@ -1111,6 +1095,7 @@ def fp8_matmul( aux=aux, ) + def doGemm( A: torch.Tensor, B: torch.Tensor, @@ -1136,10 +1121,7 @@ def doGemm( torch.float16, torch.bfloat16, torch.float32, - ), ( - "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " - f"got {C.dtype}" - ) + ), f"C dtype must be torch.float16, torch.bfloat16, or torch.float32, got {C.dtype}" assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" require_block_tiling( M_runtime, @@ -1153,9 +1135,10 @@ def doGemm( require_launch_size("FP8 GEMM", ("A", A), ("B", B), ("C", C)) 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})" - ) + assert C.shape == ( + M_runtime, + N_runtime, + ), f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): raise ValueError(f"Unsupported FP8 epilogue: {epilogue}") @@ -1167,9 +1150,7 @@ def doGemm( if bias.dtype != torch.float32: raise TypeError(f"FP8 bias must be float32, got {bias.dtype}") if bias.numel() != N_runtime: - raise ValueError( - f"FP8 bias length {bias.numel()} != N (out_features) {N_runtime}" - ) + raise ValueError(f"FP8 bias length {bias.numel()} != N (out_features) {N_runtime}") if bias.device != A.device: raise ValueError("bias must be on the same device as A, B, and C") elif bias is not None: @@ -1182,9 +1163,7 @@ def doGemm( if aux is None: raise ValueError(f"FP8 epilogue {epilogue} requires an aux output tensor") if tuple(aux.shape) != (M_runtime, N_runtime): - raise ValueError( - f"FP8 aux shape {tuple(aux.shape)} != {(M_runtime, N_runtime)}" - ) + raise ValueError(f"FP8 aux shape {tuple(aux.shape)} != {(M_runtime, N_runtime)}") if aux.dtype != C.dtype: raise TypeError(f"FP8 aux dtype {aux.dtype} != C dtype {C.dtype}") if aux.device != A.device: 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 453ee3a3a..1c75b6be7 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py @@ -24,6 +24,7 @@ xcd_swizzle, ) + def preshuffle_b(b_t): """Permute row-major ``B_T`` ``(N, K)`` for ``b_preshuffled=True``.""" n, k = b_t.shape[-2:] @@ -60,7 +61,11 @@ def compute_global_swizzle(lane_id, wave_id, K, n_rounds, 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) + (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) @@ -264,10 +269,12 @@ def _store_bf16(self, value_bf16, 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) + 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) + 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 @@ -316,7 +323,10 @@ def call(self, a, b, c, *, set_prio=True): a_frags = [self._make_operand_frag(a[idx]) for idx in range_constexpr(self.n_tiles_a)] b_frags = [self._make_operand_frag(b[idx]) for idx in range_constexpr(self.n_tiles_b)] - c_frags = [self._make_accum_frag(c[idx]) for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b)] + 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): @@ -326,7 +336,10 @@ def call(self, a, b, c, *, set_prio=True): 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)] + 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 diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 19c82fd14..42f700387 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -25,6 +25,7 @@ # GEMM wrapper utility module while preserving backend-specific layout and # storage canonicalization. + def _product(shape): """Return the product of dimensions in ``shape``.""" result = 1 @@ -69,22 +70,15 @@ def reinterpret_as_fp8_tensor( use_ocp_fp8 = capability == (9, 5) if dtype == tex.DType.kFloat8E4M3: - torch_dtype = ( - torch.float8_e4m3fn - if use_ocp_fp8 - else torch.float8_e4m3fnuz - ) + 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 - ) + 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, @@ -106,34 +100,24 @@ def _validate_common_epilogue( crashing -- these are valid GEMM requests, just not ones FlyDSL can serve. """ if quantizer is not None: - raise FlyDSLUnsupportedError( - "FlyDSL GEMM output quantization is not implemented" - ) + raise FlyDSLUnsupportedError("FlyDSL GEMM output quantization is not implemented") if float(alpha) != 1.0 or float(beta) != 0.0: - raise FlyDSLUnsupportedError( - "FlyDSL GEMM currently supports only alpha=1 and beta=0" - ) + raise FlyDSLUnsupportedError("FlyDSL GEMM currently supports only alpha=1 and beta=0") # TODO: Add accumulate option if accumulate: - raise FlyDSLUnsupportedError( - "FlyDSL GEMM accumulation is not implemented" - ) + raise FlyDSLUnsupportedError("FlyDSL GEMM accumulation is not implemented") # Forward BIAS is implemented across all backends; the fused bias-gradient # (BGRADB) path is not. TODO: add BGRADB to the FlyDSL GEMM backends. if grad and bias is not None and bias.numel() != 0: - raise FlyDSLUnsupportedError( - "FlyDSL GEMM fused bias gradient (BGRADB) is not implemented" - ) + raise FlyDSLUnsupportedError("FlyDSL GEMM fused bias gradient (BGRADB) is not implemented") # Fused forward GELU (GELU_AUX) is supported; the backward fused GELU # gradient (DGELU) is not. TODO: add DGELU to the FlyDSL GEMM backends. if gelu and grad: - raise FlyDSLUnsupportedError( - "FlyDSL GEMM fused GELU gradient (DGELU) is not implemented" - ) + raise FlyDSLUnsupportedError("FlyDSL GEMM fused GELU gradient (DGELU) is not implemented") def _resolve_bias(bias, n): @@ -146,9 +130,7 @@ def _resolve_bias(bias, n): if bias is None or bias.numel() == 0: return "DEFAULT", None if bias.numel() != n: - raise ValueError( - f"FlyDSL GEMM bias length {bias.numel()} != N (out_features) {n}" - ) + raise ValueError(f"FlyDSL GEMM bias length {bias.numel()} != N (out_features) {n}") return "BIAS", bias.reshape(-1).to(torch.float32).contiguous() @@ -174,6 +156,7 @@ def _classify_input(t): from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import ( Float8TensorStorage, ) + if isinstance(t, (Float8Tensor, Float8TensorStorage)): return "fp8", t except ImportError: @@ -184,6 +167,7 @@ def _classify_input(t): from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import ( MXFP8TensorStorage, ) + if isinstance(t, (MXFP8Tensor, MXFP8TensorStorage)): return "mxfp8", t except ImportError: @@ -193,12 +177,13 @@ def _classify_input(t): from transformer_engine.pytorch.quantized_tensor import ( QuantizedTensorStorage, ) + if isinstance(t, QuantizedTensorStorage): raise ValueError( - f"The FlyDSL GEMM backend does not support " + "The FlyDSL GEMM backend does not support " f"{type(t).__name__}. Only Float8Tensor / " - f"Float8TensorStorage and MXFP8Tensor / " - f"MXFP8TensorStorage are implemented." + "Float8TensorStorage and MXFP8Tensor / " + "MXFP8TensorStorage are implemented." ) except ImportError: pass @@ -209,17 +194,13 @@ 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 FlyDSLUnsupportedError( - 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, tex.DType.kFloat8E5M2, ): - raise TypeError( - f"{name} has unsupported TE FP8 dtype metadata: {fp8_dtype}" - ) + 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. @@ -289,18 +270,14 @@ def _fp8_scale_debug(name: str, scale: torch.Tensor) -> None: 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"FlyDSL GEMM expects {name} to have rank >= 2, got {tuple(t.shape)}" - ) + raise ValueError(f"FlyDSL GEMM expects {name} to have rank >= 2, got {tuple(t.shape)}") 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)}" - ) + raise ValueError(f"FlyDSL GEMM expects {name} to have rank >= 2, got {tuple(t.shape)}") return t.reshape(t.shape[0], -1) @@ -316,21 +293,13 @@ def _validate_or_allocate_output( return torch.empty(shape, dtype=dtype, device=device) if tuple(D.shape) != tuple(shape): - raise ValueError( - f"D shape {tuple(D.shape)} does not match expected {tuple(shape)}" - ) + raise ValueError(f"D shape {tuple(D.shape)} does not match expected {tuple(shape)}") if D.dtype != dtype: - raise TypeError( - f"FlyDSL {backend_name} requires {dtype} output, got {D.dtype}" - ) + raise TypeError(f"FlyDSL {backend_name} requires {dtype} output, got {D.dtype}") if D.device != device: - raise ValueError( - f"D must be on {device}, got {D.device}" - ) + raise ValueError(f"D must be on {device}, got {D.device}") if not D.is_contiguous(): - raise ValueError( - f"FlyDSL {backend_name} requires contiguous output storage" - ) + raise ValueError(f"FlyDSL {backend_name} requires contiguous output storage") return D @@ -357,13 +326,10 @@ def _run_bf16_gemm( 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}" + f"FlyDSL BF16 GEMM requires torch.bfloat16 inputs, 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}" - ) + raise ValueError(f"A and B must be on the same device, got {A.device} and {B.device}") dispatch = { (True, False): "TN", @@ -447,7 +413,6 @@ def _run_bf16_gemm( return D, aux - def _run_fp16_gemm( A, transa, @@ -471,13 +436,10 @@ def _run_fp16_gemm( 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}" + f"FlyDSL FP16 GEMM requires torch.float16 inputs, 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}" - ) + raise ValueError(f"A and B must be on the same device, got {A.device} and {B.device}") dispatch = { (True, False): "TN", @@ -598,17 +560,12 @@ def _run_fp32_gemm( 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 torch.float32 inputs, " - f"got A={A.dtype} and B={B.dtype}" + f"FlyDSL FP32 GEMM requires torch.float32 inputs, 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}" - ) + 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)" - ) + raise FlyDSLUnsupportedError("FlyDSL GEMM does not support transa=True, transb=True (TT)") output_shape = _get_gemm_output_shape(A, transa, B, transb) @@ -636,7 +593,7 @@ def _run_fp32_gemm( if a_tn.ndim != 2 or b_tn.ndim != 2: raise RuntimeError( - f"FlyDSL FP32 TN normalization produced rank mismatch: " + "FlyDSL FP32 TN normalization produced rank mismatch: " f"a={tuple(a_tn.shape)}, b={tuple(b_tn.shape)}" ) @@ -676,14 +633,11 @@ def _run_fp32_gemm( return D, aux - def _get_fp8_rowwise_payload(t, name): """Return TE's existing rowwise ``_data`` payload without copying.""" data = getattr(t, "_data", None) if data is None: - raise FlyDSLUnsupportedError( - f"FlyDSL FP8 requires existing {name} rowwise (_data) storage" - ) + raise FlyDSLUnsupportedError(f"FlyDSL FP8 requires existing {name} rowwise (_data) storage") return _reinterpret_fp8_payload( data, getattr(t, "_fp8_dtype", None), @@ -765,9 +719,7 @@ def _select_mxfp8_data_and_scale( ) if data is None or scale is None: - raise RuntimeError( - f"{name} does not contain required {orientation} MXFP8 data and scales" - ) + raise RuntimeError(f"{name} does not contain required {orientation} MXFP8 data and scales") _mxfp8_debug( f"{name} selected data shape={tuple(data.shape)}, " @@ -789,11 +741,10 @@ def _mxfp8_logical_shape(t, name: str) -> torch.Size: 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" - ) + 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, @@ -809,10 +760,7 @@ def _flatten_mxfp8_scale( [K/32, ...] -> [K/32, outer] """ if t.ndim < 2: - raise ValueError( - f"FlyDSL MXFP8 expects {name} scale rank >= 2, " - f"got {tuple(t.shape)}" - ) + raise ValueError(f"FlyDSL MXFP8 expects {name} scale rank >= 2, got {tuple(t.shape)}") original_shape = tuple(t.shape) if source_colwise: @@ -865,10 +813,7 @@ def _run_mxfp8( tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2, ) - if ( - a_fp8_dtype not in supported_fp8_dtypes - or b_fp8_dtype not in supported_fp8_dtypes - ): + if a_fp8_dtype not in supported_fp8_dtypes or b_fp8_dtype not in supported_fp8_dtypes: raise FlyDSLUnsupportedError( "FlyDSL MXFP8 supports E4M3 and E5M2 independently for A/B; " f"got A={a_fp8_dtype} and B={b_fp8_dtype}" @@ -993,9 +938,7 @@ def _run_mxfp8( ) if k % 32 != 0: - raise ValueError( - f"K={k} must be divisible by MXFP8 scale group size 32" - ) + raise ValueError(f"K={k} must be divisible by MXFP8 scale group size 32") if tuple(a_scale.shape) != expected_a_scale: raise ValueError( @@ -1109,9 +1052,7 @@ def _select_fp8_storage_for_layout(A, transa, B, transb): B_data = _flatten_columnwise(B_payload, B_storage) else: - raise FlyDSLUnsupportedError( - "FlyDSL GEMM does not support transa=True, transb=True (TT)" - ) + raise FlyDSLUnsupportedError("FlyDSL GEMM does not support transa=True, transb=True (TT)") return ( A_data, @@ -1141,10 +1082,7 @@ def _run_fp8( ) 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 - ): + if a_fp8_dtype not in supported_fp8_dtypes or b_fp8_dtype not in supported_fp8_dtypes: raise FlyDSLUnsupportedError( "FlyDSL FP8 supports E4M3 and E5M2 independently for A/B; " f"got A={a_fp8_dtype} and B={b_fp8_dtype}" @@ -1330,14 +1268,10 @@ def te_generic_gemm_flydsl( or extra_output is not None or bulk_overlap ): - raise FlyDSLUnsupportedError( - "FlyDSL GEMM does not support comm+GEMM overlap" - ) + raise FlyDSLUnsupportedError("FlyDSL GEMM does not support comm+GEMM overlap") if transa and transb: - raise FlyDSLUnsupportedError( - "FlyDSL GEMM does not support transa=True, transb=True (TT)" - ) + raise FlyDSLUnsupportedError("FlyDSL GEMM does not support transa=True, transb=True (TT)") a_kind, _ = _classify_input(A) b_kind, _ = _classify_input(B) @@ -1353,16 +1287,20 @@ def te_generic_gemm_flydsl( ) if a_kind == "mxfp8" or b_kind == "mxfp8": - # Validate both are MXFP8 + # Validate both are MXFP8 if a_kind != b_kind: - raise ValueError( - "Mixed MXFP8 and non-MXFP8 FlyDSL GEMM inputs are not supported" - ) + 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: + 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: + 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") mxfp8_output_dtypes = { @@ -1373,8 +1311,7 @@ def te_generic_gemm_flydsl( } if output_dtype not in mxfp8_output_dtypes: raise FlyDSLUnsupportedError( - "FlyDSL MXFP8 supports FP16, BF16, or FP32 output, " - f"got {output_dtype}" + f"FlyDSL MXFP8 supports FP16, BF16, or FP32 output, got {output_dtype}" ) D, gelu_input = _run_mxfp8( @@ -1391,9 +1328,7 @@ def te_generic_gemm_flydsl( 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" - ) + raise ValueError("Mixed regular FP8 and non-FP8 FlyDSL GEMM inputs are not supported") fp8_output_dtypes = { None: torch.float16, @@ -1403,8 +1338,7 @@ def te_generic_gemm_flydsl( } if output_dtype not in fp8_output_dtypes: raise FlyDSLUnsupportedError( - "FlyDSL tensor-wise FP8 supports FP16, BF16, or FP32 output, " - f"got {output_dtype}" + f"FlyDSL tensor-wise FP8 supports FP16, BF16, or FP32 output, got {output_dtype}" ) D, gelu_input = _run_fp8( @@ -1421,13 +1355,10 @@ def te_generic_gemm_flydsl( if a_kind != "regular" or b_kind != "regular": raise TypeError( - "Unsupported FlyDSL GEMM operand types: " - f"{type(A).__name__} and {type(B).__name__}" + f"Unsupported FlyDSL GEMM operand types: {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" - ) + raise TypeError("FlyDSL regular GEMM expects plain torch.Tensor operands") if A.dtype == torch.bfloat16 and B.dtype == torch.bfloat16: bf16_output_dtypes = { @@ -1438,8 +1369,7 @@ def te_generic_gemm_flydsl( } if output_dtype not in bf16_output_dtypes: raise FlyDSLUnsupportedError( - "FlyDSL BF16 supports FP16, BF16, or FP32 output, " - f"got {output_dtype}" + f"FlyDSL BF16 supports FP16, BF16, or FP32 output, got {output_dtype}" ) D, gelu_input = _run_bf16_gemm( A, @@ -1462,8 +1392,7 @@ def te_generic_gemm_flydsl( } if output_dtype not in fp16_output_dtypes: raise FlyDSLUnsupportedError( - "FlyDSL FP16 supports FP16, BF16, or FP32 output, " - f"got {output_dtype}" + f"FlyDSL FP16 supports FP16, BF16, or FP32 output, got {output_dtype}" ) D, gelu_input = _run_fp16_gemm( A, @@ -1480,8 +1409,7 @@ def te_generic_gemm_flydsl( if A.dtype == torch.float32 and B.dtype == torch.float32: if output_dtype not in (None, tex.DType.kFloat32): raise FlyDSLUnsupportedError( - "FlyDSL FP32 currently supports only FP32 output, " - f"got {output_dtype}" + f"FlyDSL FP32 currently supports only FP32 output, got {output_dtype}" ) D, gelu_input = _run_fp32_gemm( A, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 9b29274a8..06ba10e19 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -44,11 +44,12 @@ S2RLoader, compute_global_transpose_swizzle, compute_global_swizzle, + divmod, make_byte_buffer_tensor, pack_i32x4_i32x8, swizzle_128, xcd_swizzle, - barrier + barrier, ) # FP16 and BF16 differ only in the MFMA opcode suffix. @@ -68,7 +69,7 @@ WARP_SIZE = 64 NUM_WAVES = NUM_THREADS // WARP_SIZE -SUBTILE_M = 64 +SUBTILE_M = 64 SUBTILE_N = 64 MFMA_M = 16 @@ -184,7 +185,9 @@ def _compile_kernel( 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" + 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 @@ -196,15 +199,14 @@ def _compile_kernel( 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) - ) + 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, @@ -224,18 +226,18 @@ def _load_a_half( - 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) - ) + 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, @@ -249,11 +251,7 @@ def _load_a_half( ): 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 - ) + 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], @@ -262,15 +260,14 @@ def _load_a_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) - ) + 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, @@ -289,18 +286,18 @@ def _load_b_ni( - 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) - ) + 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, @@ -313,17 +310,14 @@ def _load_b_ni( ): 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 - ) + 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``. # Half-precision 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_transpose_swizzle( lane, @@ -331,7 +325,9 @@ def _a_global_offsets(lane, wave_id, c_m): _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( @@ -343,6 +339,7 @@ def _a_global_offsets(lane, wave_id, c_m): ) if b_transpose_read: + def _b_global_offsets(lane, wave_id, c_n): return compute_global_transpose_swizzle( lane, @@ -350,7 +347,9 @@ def _b_global_offsets(lane, wave_id, c_n): _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( @@ -594,15 +593,11 @@ def hot_loop_scheduler_q_prefetch_4n(): 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 = _a_global_base_bytes( - k_base, subtile, c_m, bx_m_idx - ) + 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 = _b_global_base_bytes( - k_base, subtile, c_n, by_n_idx - ) + 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): @@ -648,14 +643,10 @@ def load_transposed_frag_half(lds_page, local_x_tile, half): 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) - ) + 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) @@ -675,7 +666,9 @@ def load_transposed_frag(lds_page, local_x_tile): 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 + return ( + subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + ) def _k32_frag(full_frag, k32): # A/B 16x64 half-precision wave fragments are i32x8. Each K32 MFMA @@ -692,16 +685,11 @@ def _pinned_mfma_once(acc_idx, a_k32, b_k32): llvm.InlineAsmOp( None, [arith._to_raw(a_k32), arith._to_raw(b_k32)], - ( - f"v_mfma_f32_16x16x32_{mfma_suffix} " - 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}}}" - ), + f"v_mfma_f32_16x16x32_{mfma_suffix} " + f"a[{acc_pin}:{acc_pin + 3}], " + "$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]", + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}", has_side_effects=True, ) @@ -801,7 +789,6 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(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: @@ -1108,7 +1095,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): # # Finalize accumulators in their own physical AGPR slots, but delay # each AGPR read/store until several independent final MFMAs have - # been issued. + # been issued. # # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, # MFMA 4, drain 1, MFMA 5, drain 2, ... @@ -1226,7 +1213,6 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): ) hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) - @flyc.jit def launch_gemm( A: fx.Tensor, @@ -1253,6 +1239,7 @@ def launch_gemm( return launch_gemm + @functools.lru_cache(maxsize=None) def _cached_launch( K: int, @@ -1297,13 +1284,11 @@ def _half_prec_matmul( raise ValueError(f"Unsupported {label} layout: {layout}") if a.ndim != 2 or b.ndim != 2: raise ValueError( - f"FlyDSL {label} expects rank-2 operands, got A{tuple(a.shape)} " - f"and B{tuple(b.shape)}" + f"FlyDSL {label} expects rank-2 operands, got A{tuple(a.shape)} and B{tuple(b.shape)}" ) if a.dtype != input_dtype or b.dtype != input_dtype: raise TypeError( - f"FlyDSL {label} GEMM expects {input_dtype} operands, " - f"got A={a.dtype}, B={b.dtype}" + f"FlyDSL {label} GEMM expects {input_dtype} operands, got A={a.dtype}, B={b.dtype}" ) if not a.is_contiguous() or not b.is_contiguous(): raise FlyDSLUnsupportedError( @@ -1333,13 +1318,11 @@ def _half_prec_matmul( raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise TypeError( - f"FlyDSL {label} output must be float16, bfloat16, or float32, " - f"got {c.dtype}" + f"FlyDSL {label} output must be float16, bfloat16, or float32, 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}" + 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(f"FlyDSL {label} GEMM requires contiguous output storage") @@ -1463,8 +1446,7 @@ def doGemm( if A.dtype != input_dtype or B.dtype != input_dtype: raise TypeError( - f"{label} {layout} requires {input_dtype} inputs, " - f"got {A.dtype} and {B.dtype}" + f"{label} {layout} requires {input_dtype} inputs, got {A.dtype} and {B.dtype}" ) if C.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise TypeError(f"Unsupported {label} output dtype: {C.dtype}") @@ -1481,9 +1463,7 @@ def doGemm( require_launch_size(f"{label} GEMM", ("A", A), ("B", B), ("C", C)) if tuple(C.shape) != (M_runtime, N_runtime): - raise ValueError( - f"C shape {tuple(C.shape)} != expected {(M_runtime, N_runtime)}" - ) + raise ValueError(f"C shape {tuple(C.shape)} != expected {(M_runtime, N_runtime)}") if epilogue not in ("DEFAULT", "BIAS", "GELU_AUX", "GELU_AUX_BIAS"): raise ValueError(f"Unsupported {label} epilogue: {epilogue}") @@ -1495,9 +1475,7 @@ def doGemm( if bias.dtype != torch.float32: raise TypeError(f"{label} bias must be float32, got {bias.dtype}") if bias.numel() != N_runtime: - raise ValueError( - f"{label} bias length {bias.numel()} != N (out_features) {N_runtime}" - ) + raise ValueError(f"{label} bias length {bias.numel()} != N (out_features) {N_runtime}") if bias.device != A.device: raise ValueError("bias must be on the same device as A, B, and C") elif bias is not None: @@ -1510,9 +1488,7 @@ def doGemm( if aux is None: raise ValueError(f"{label} epilogue {epilogue} requires an aux output tensor") if tuple(aux.shape) != (M_runtime, N_runtime): - raise ValueError( - f"{label} aux shape {tuple(aux.shape)} != {(M_runtime, N_runtime)}" - ) + raise ValueError(f"{label} aux shape {tuple(aux.shape)} != {(M_runtime, N_runtime)}") if aux.dtype != C.dtype: raise TypeError(f"{label} aux dtype {aux.dtype} != C dtype {C.dtype}") if aux.device != A.device: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 75bfdcc2d..b75384fdc 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -44,7 +44,7 @@ pack_i32x4_i32x8, swizzle_128, xcd_swizzle, - barrier + barrier, ) @@ -86,39 +86,31 @@ def _compile_mx32_scale_pack_kernel( and no eager PyTorch shift/index/OR kernels. """ if dim % 64 != 0: - raise FlyDSLUnsupportedError( - f"Scale outer dimension={dim} must be a multiple of 64" - ) + raise FlyDSLUnsupportedError(f"Scale outer dimension={dim} must be a multiple of 64") if qk % 4 != 0: - raise FlyDSLUnsupportedError( - f"Scale K/32 dimension={qk} must be divisible by 4" - ) + raise FlyDSLUnsupportedError(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 FlyDSLUnsupportedError( - f"Packed scale words={total_words} must be divisible by " - f"{_SCALE_PACK_THREADS}" + f"Packed scale words={total_words} must be divisible by {_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: + 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) - ) + return source_k32 * fx.Index(stride0) + source_row * fx.Index(stride1) + 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) - ) + 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): @@ -139,9 +131,8 @@ def kernel_pack_mx32_scales(src: fx.Tensor, dst: fx.Tensor): scale_load_atom = fx.make_copy_atom(fx.rocdl.BufferCopy8b(), fx.Uint8) scale_store_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) - linear = ( - fx.Index(fx.block_idx.x) * fx.Index(_SCALE_PACK_THREADS) - + fx.Index(gpu.thread_id("x")) + 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) @@ -152,11 +143,7 @@ def kernel_pack_mx32_scales(src: fx.Tensor, dst: fx.Tensor): 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 - ) + source_row = tile * fx.Index(64) + fx.Index(group * 16) + row_within_16 off = _source_offset(source_k32, source_row) reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Uint8) fx.copy(scale_load_atom, fx.slice(src_div, (None, fx.Int32(off))), reg) @@ -174,7 +161,6 @@ def load_scale_byte(group): fx.memref_store_vec(Vec.filled(1, packed, fx.Int32), reg_i32) fx.copy(scale_store_atom, reg_i32, fx.slice(dst_div, (None, fx.Int32(linear)))) - @flyc.jit def launch_pack_mx32_scales( src: fx.Tensor, @@ -199,9 +185,7 @@ def _cached_mx32_scale_pack_launch( stride1: int, ): """Cache orientation-and-stride-specialized fused pack binaries.""" - return _compile_mx32_scale_pack_kernel( - dim, qk, source_colwise, stride0, stride1 - ) + return _compile_mx32_scale_pack_kernel(dim, qk, source_colwise, stride0, stride1) def pack_mx32_scales_for_hk( @@ -220,21 +204,13 @@ def pack_mx32_scales_for_hk( * ``[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}" - ) + 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 {tuple(scales_u8.shape)}" - ) + raise ValueError(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()}" - ) + raise ValueError(f"MXFP8 scale packing requires positive strides, got {scales_u8.stride()}") if source_colwise: qk, dim = scales_u8.shape @@ -242,13 +218,9 @@ def pack_mx32_scales_for_hk( dim, qk = scales_u8.shape if qk % 4 != 0: - raise ValueError( - f"Scale K/32 dimension={qk} must be divisible by 4" - ) + 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" - ) + raise ValueError(f"Scale outer dimension={dim} must be a multiple of 64") packed = torch.empty( (qk // 4, dim), @@ -334,8 +306,7 @@ def _compile_kernel( output_fx_dtype = fx.Float32 else: raise TypeError( - "FlyDSL MXFP8 supports only float16, bfloat16, and float32 " - f"outputs, got {output_dtype}" + f"FlyDSL MXFP8 supports only float16, bfloat16, and float32 outputs, got {output_dtype}" ) NUM_THREADS = 256 @@ -368,7 +339,9 @@ def _compile_kernel( 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" + 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) @@ -381,15 +354,12 @@ def _compile_kernel( 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)) - ) + return k_base * fx.Index(c_m) + bx_m_idx + fx.Index(subtile * (BLOCK_M // 2)) def _load_a_half( load_transposed_frag_half, @@ -412,18 +382,16 @@ def _load_a_half( 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 - ) + return (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) + k_base def _load_a_half( load_transposed_frag_half, @@ -437,11 +405,7 @@ def _load_a_half( ): 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 - ) + 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( @@ -451,15 +415,12 @@ def _load_a_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)) - ) + return k_base * fx.Index(c_n) + by_n_idx + fx.Index(subtile * (BLOCK_N // 2)) def _load_b_ni( load_transposed_frag, @@ -478,18 +439,16 @@ def _load_b_ni( - 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 - ) + return (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + k_base def _load_b_ni( load_transposed_frag, @@ -502,21 +461,20 @@ def _load_b_ni( ): 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 - ) + 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) @@ -536,7 +494,15 @@ class SharedStorage: @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, Bias: fx.Tensor, Aux: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32 + A: fx.Tensor, + As: fx.Tensor, + B: fx.Tensor, + Bs: fx.Tensor, + C: fx.Tensor, + Bias: fx.Tensor, + Aux: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, ): lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_a0 = (lds.a0_0, lds.a0_1) @@ -858,14 +824,8 @@ 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) - ) + 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 @@ -885,7 +845,9 @@ def load_transposed_frag(lds_page, local_x_tile): 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 + 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. @@ -901,17 +863,15 @@ def pinned_mfma(acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): _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}}}"), + "v_mfma_scale_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + "$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}], " + "$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, ) @@ -929,17 +889,15 @@ def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag, a_scale, b_scale, m _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}}}"), + "v_mfma_scale_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + "$0, $1, " + f"a[{old_pin}:{old_pin + 3}], " + "$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, ) @@ -1006,7 +964,6 @@ def store_acc_vector_for_logical_idx(logical_acc_idx, acc): fx.memref_store_vec(Vec.filled(1, value, output_fx_dtype), reg) fx.copy(c_store_atom, reg, fx.slice(c_div, (None, fx.Int32(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: @@ -1201,7 +1158,10 @@ def hk_one_k_with_refill( # 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) + 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) @@ -1253,7 +1213,9 @@ def hk_one_k_with_refill( 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): + 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 @@ -1371,7 +1333,7 @@ def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): # # Finalize accumulators in their own physical AGPR slots, but delay # each AGPR read/store until several independent final MFMAs have - # been issued. + # been issued. # # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, # MFMA 4, drain 1, MFMA 5, drain 2, ... @@ -1557,6 +1519,7 @@ def launch_gemm( return launch_gemm + def do_gemm( A: torch.Tensor, As: torch.Tensor, @@ -1603,7 +1566,11 @@ def do_gemm( ) require_launch_size( f"MXFP8 {layout} GEMM", - ("A", A), ("B", B), ("As", As), ("Bs", Bs), ("C", C), + ("A", A), + ("B", B), + ("As", As), + ("Bs", Bs), + ("C", C), ) expected_as = (K_runtime // _BLOCK_K, M_runtime) @@ -1612,13 +1579,13 @@ def do_gemm( assert Bs.dtype == torch.int32, f"Bs dtype {Bs.dtype} != torch.int32 packed scales" 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)}" - ) + 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}" + f"C dtype must be torch.float16, torch.bfloat16, or torch.float32, got {C.dtype}" ) tensors = (A, As, B, Bs, C) @@ -1635,9 +1602,7 @@ def do_gemm( if bias.dtype != torch.float32: raise TypeError(f"MXFP8 bias must be float32, got {bias.dtype}") if bias.numel() != N_runtime: - raise ValueError( - f"MXFP8 bias length {bias.numel()} != N (out_features) {N_runtime}" - ) + raise ValueError(f"MXFP8 bias length {bias.numel()} != N (out_features) {N_runtime}") if bias.device != A.device: raise ValueError("bias must be on the same device as A, B, and C") elif bias is not None: @@ -1650,9 +1615,7 @@ def do_gemm( if aux is None: raise ValueError(f"MXFP8 epilogue {epilogue} requires an aux output tensor") if tuple(aux.shape) != (M_runtime, N_runtime): - raise ValueError( - f"MXFP8 aux shape {tuple(aux.shape)} != {(M_runtime, N_runtime)}" - ) + raise ValueError(f"MXFP8 aux shape {tuple(aux.shape)} != {(M_runtime, N_runtime)}") if aux.dtype != C.dtype: raise TypeError(f"MXFP8 aux dtype {aux.dtype} != C dtype {C.dtype}") if aux.device != A.device: @@ -1741,10 +1704,7 @@ def _validate_common_payloads( 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}" - ) + raise TypeError(f"FlyDSL MXFP8 output must be float16, bfloat16, or float32, got {D.dtype}") def mxfp8_matmul( @@ -1789,16 +1749,12 @@ def mxfp8_matmul( if kb != k: raise ValueError( - f"Incompatible MXFP8 {layout} operands: " - f"A{tuple(a.shape)} and B{tuple(b.shape)}" + f"Incompatible MXFP8 {layout} operands: A{tuple(a.shape)} and B{tuple(b.shape)}" ) if tuple(D.shape) != (m, n): 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}" - ) + raise ValueError(f"K={k} must be divisible by MXFP8 scale group size {SCALE_GROUP_SIZE}") if layout == "NT": expected_a_scale = (k // SCALE_GROUP_SIZE, m) @@ -1812,13 +1768,11 @@ def mxfp8_matmul( if tuple(a_scale.shape) != expected_a_scale: raise ValueError( - f"a_scale shape {tuple(a_scale.shape)} != expected " - f"{expected_a_scale} for {layout}" + f"a_scale shape {tuple(a_scale.shape)} != expected {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} for {layout}" + f"b_scale shape {tuple(b_scale.shape)} != expected {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") From 77a069dd9f244a463bfd1e14e918d55caa3543ec Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 13:46:26 +0000 Subject: [PATCH 46/65] Make test_linear_accuracy_flydsl detect silent FlyDSL fallback The test only skipped on non-HIP, but FlyDSL dispatch is gfx950-gated in cpp_extensions/gemm.py -- so on other archs NVTE_USE_FLYDSL=1 was a no-op and the test compared the native backend against itself. It also never checked that the FlyDSL path was actually taken: every GEMM could raise FlyDSLUnsupportedError and fall back, and the test would still pass. - Gate on gfx950 (get_device_compute_capability() == (9, 5)) and pytest.importorskip("flydsl"), so the test skips honestly on unsupported arch / a build without flydsl instead of passing vacuously or erroring. - Capture [FLYDSL WARNING] fallback notices (NVTE_FLYDSL_GEMM_WARN_FALLBACK is already set) around the FlyDSL forward+backward and fail if a config the PR claims to support silently fell back. "small" is the designated fallback case (K=128 is below the 4-K-tile minimum) and is exempt. Co-Authored-By: Claude --- tests/pytorch/test_numerics.py | 39 ++++++++++++++++++++++++++++------ 1 file changed, 33 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 2bc5356e8..f4020f182 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -6,6 +6,7 @@ import math import os +import warnings from typing import Dict, List, Tuple, Optional import pytest @@ -1351,12 +1352,23 @@ def test_linear_accuracy_flydsl( ): """Compare FlyDSL and native TE Linear forward, dgrad, and wgrad.""" - if not IS_HIP_EXTENSION: - pytest.skip("FlyDSL GEMM is only supported on HIP.") + # FlyDSL GEMM dispatch is gated on gfx950 in cpp_extensions/gemm.py. On any + # other arch NVTE_USE_FLYDSL=1 is a no-op and the FlyDSL path would just be + # the native backend compared against itself, so skip rather than pass + # vacuously. + if not IS_HIP_EXTENSION or get_device_compute_capability() != (9, 5): + pytest.skip("FlyDSL GEMM is only supported on gfx950.") + # flydsl is only installed when NVTE_USE_FLYDSL=1 at build time; without it + # the lazy import in general_gemm raises, so skip instead of erroring. + pytest.importorskip("flydsl", reason="FlyDSL package is not installed.") fp8 = fp8_recipe is not None config = model_configs[model] + # "small" is the designated fallback case (not a config this PR claims to + # support); every other model is expected to run on FlyDSL. + expect_fallback = model == "small" + if isinstance(fp8_recipe, recipe.MXFP8BlockScaling): if not mxfp8_available: pytest.skip(reason_for_no_mxfp8) @@ -1423,17 +1435,32 @@ def test_linear_accuracy_flydsl( reset_rng_states() FP8GlobalStateManager.reset() - with autocast(enabled=fp8, recipe=fp8_recipe): - out_flydsl = linear_flydsl(inp_flydsl) + # Capture the [FLYDSL WARNING] fallback notices emitted by + # cpp_extensions/gemm.py so we can tell whether FlyDSL actually ran. + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with autocast(enabled=fp8, recipe=fp8_recipe): + out_flydsl = linear_flydsl(inp_flydsl) - out_flydsl.sum().backward() - torch.cuda.synchronize() + out_flydsl.sum().backward() + torch.cuda.synchronize() + + fell_back = any("[FLYDSL WARNING]" in str(w.message) for w in caught) finally: os.environ.pop("NVTE_USE_FLYDSL", None) os.environ.pop("NVTE_FLYDSL_GEMM_WARN_FALLBACK", None) FP8GlobalStateManager.reset() + # A silent fallback on a supported config means FlyDSL never ran, so the + # correctness asserts below would pass vacuously (native vs native). + if not expect_fallback and fell_back: + pytest.fail( + "FlyDSL GEMM unexpectedly fell back to the native backend for " + f"model={model}, dtype={dtype}, fp8_recipe={fp8_recipe}; " + "the FlyDSL path was not exercised." + ) + tols = dtype_tols(dtype) atol = tols["atol"] rtol = tols["rtol"] From 86b0e2ab4a2b896fbe633249e073d9ab9ee0fa1a Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 13:51:04 +0000 Subject: [PATCH 47/65] Gate flydsl_kernels/test_gemm on gfx950 and detect silent fallback The suite gated only the MXFP8 tests on gfx950, but all FlyDSL dispatch is gfx950-gated in cpp_extensions/gemm.py. On other archs the regular/FP8 tests ran hipBLASLt (not FlyDSL) and the vs-cpp cases became native-vs-native tautologies. - Skip the whole module unless on gfx950, via a pytestmark, and pytest.importorskip("flydsl"). Compute capability is read behind a helper guarded by torch.cuda.is_available() so importing the module never initialises CUDA at collection time (which would error a CPU-only runner instead of skipping). - Route every general_gemm call through a helper that enables NVTE_FLYDSL_GEMM_WARN_FALLBACK and captures [FLYDSL WARNING] notices; when FlyDSL is requested and a fallback is detected, fail. All shapes here are tile-aligned supported configs, so any fallback means FlyDSL did not run. Co-Authored-By: Claude --- tests/pytorch/flydsl_kernels/test_gemm.py | 141 +++++++++++++++++----- 1 file changed, 109 insertions(+), 32 deletions(-) diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py index 9326de034..3399d3059 100644 --- a/tests/pytorch/flydsl_kernels/test_gemm.py +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -33,6 +33,7 @@ """ import os +import warnings import pytest import torch @@ -49,10 +50,35 @@ # --- Feature detection -------------------------------------------------------- -major, minor = torch.cuda.get_device_capability() + +def _device_capability(): + """Compute capability as (major, minor), or None on a CPU-only box. + + Deferred behind a helper so importing this module never initialises CUDA at + collection time (which would error the whole module on a CPU-only runner + instead of skipping it). + """ + if not torch.cuda.is_available(): + return None + return torch.cuda.get_device_capability() + + +# All FlyDSL GEMM dispatch is gated on gfx950 in cpp_extensions/gemm.py, not +# just MXFP8. On any other arch general_gemm silently runs the C++ backend, so +# every test here would either exercise hipBLASLt or (for the vs-cpp cases) be +# a tautology. Skip the whole module unless we are on gfx950 with FlyDSL +# installed (flydsl is only present when NVTE_USE_FLYDSL=1 at build time). +_CAP = _device_capability() +has_flydsl_support = _CAP == (9, 5) +pytestmark = pytest.mark.skipif( + not has_flydsl_support, + reason="FlyDSL GEMM dispatch requires gfx950", +) +if has_flydsl_support: + pytest.importorskip("flydsl", reason="FlyDSL package is not installed") # The current FlyDSL MXFP8 implementation uses the gfx950 fp8-scaled MFMA. -has_mxfp8_support = major == 9 and minor >= 5 +has_mxfp8_support = has_flydsl_support requires_mxfp8_support = pytest.mark.skipif( not has_mxfp8_support, @@ -60,6 +86,34 @@ ) +# --- FlyDSL fallback detection ------------------------------------------------ + +_FLYDSL_FALLBACK_TAG = "[FLYDSL WARNING]" + + +def _run_capturing_fallback(fn): + """Run ``fn`` with FlyDSL fallback warnings enabled and captured. + + Returns ``(result, fell_back)`` where ``fell_back`` is True if the dispatch + emitted a ``[FLYDSL WARNING]`` fallback for any GEMM in ``fn``. The env var + makes cpp_extensions/gemm.py warn (instead of silently) when a FlyDSL GEMM + is unsupported and it falls back to the default backend. + """ + old = os.environ.get("NVTE_FLYDSL_GEMM_WARN_FALLBACK") + os.environ["NVTE_FLYDSL_GEMM_WARN_FALLBACK"] = "1" + try: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result = fn() + fell_back = any(_FLYDSL_FALLBACK_TAG in str(w.message) for w in caught) + return result, fell_back + finally: + if old is None: + os.environ.pop("NVTE_FLYDSL_GEMM_WARN_FALLBACK", None) + else: + os.environ["NVTE_FLYDSL_GEMM_WARN_FALLBACK"] = old + + # --- Test parameters ---------------------------------------------------------- # The current FlyDSL kernels have no M/N edge masks and specialize K in K128 @@ -192,21 +246,38 @@ def create_mxfp8_tensors( ) +def _assert_flydsl_ran(use_flydsl, fell_back): + """Fail if FlyDSL was requested for a supported config but silently fell back. + + Every shape in this suite is tile-aligned and a config the PR claims to + support, so a fallback means FlyDSL did not actually run and the comparison + below would be vacuous (native vs native). + """ + if use_flydsl and fell_back: + pytest.fail( + "FlyDSL GEMM unexpectedly fell back to the native backend; " + "the FlyDSL path was not exercised." + ) + + 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, + (output, bias_grad, gelu_input, extra_output), fell_back = _run_capturing_fallback( + lambda: general_gemm( + A=A, + B=B, + out_dtype=out_dtype, + layout=layout, + bias=None, + quantization_params=None, + gelu=False, + grad=False, + accumulate=False, + ) ) + _assert_flydsl_ran(use_flydsl, fell_back) assert bias_grad is None assert gelu_input is None @@ -222,17 +293,20 @@ def call_gemm_with_bias(A, B, layout, out_dtype, bias, use_flydsl=True): """ 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=bias, - quantization_params=None, - gelu=False, - grad=False, - accumulate=False, + (output, bias_grad, gelu_input, extra_output), fell_back = _run_capturing_fallback( + lambda: general_gemm( + A=A, + B=B, + out_dtype=out_dtype, + layout=layout, + bias=bias, + quantization_params=None, + gelu=False, + grad=False, + accumulate=False, + ) ) + _assert_flydsl_ran(use_flydsl, fell_back) assert gelu_input is None assert extra_output is None @@ -248,17 +322,20 @@ def call_gemm_with_gelu(A, B, layout, out_dtype, bias=None, use_flydsl=True): """ 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=bias, - quantization_params=None, - gelu=True, - grad=False, - accumulate=False, + (output, bias_grad, gelu_input, extra_output), fell_back = _run_capturing_fallback( + lambda: general_gemm( + A=A, + B=B, + out_dtype=out_dtype, + layout=layout, + bias=bias, + quantization_params=None, + gelu=True, + grad=False, + accumulate=False, + ) ) + _assert_flydsl_ran(use_flydsl, fell_back) assert bias_grad is None assert extra_output is None From b4612f01b21dbbdf6ec0b5128b9fd2150ee13f49 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 13:56:00 +0000 Subject: [PATCH 48/65] Fix flydsl test get_shapes to reference N for NN/NT + add M!=N shape get_shapes ignored N for the NN and NT layouts (NN returned (M,K),(K,M); NT returned (M,K),(M,K)), so those layouts only ever produced square outputs and the M/N operand-swap contract in gemm_wrappers was never exercised asymmetrically -- a transposed-M/N bug would still pass. This matched only because every shape happened to have N == M. Adopt the corrected per-layout shapes from the sibling Triton suite (NN -> (K,M),(N,K); NT -> (K,M),(K,N)), which are consistent with the shared general_gemm (N, M) output contract already encoded in compute_pytorch_reference, and add a (512, 512, 1024) M != N shape to FLYDSL_SHAPES and MXFP8_SHAPES so the asymmetry is actually covered. Co-Authored-By: Claude --- tests/pytorch/flydsl_kernels/test_gemm.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py index 3399d3059..4f3d0bb13 100644 --- a/tests/pytorch/flydsl_kernels/test_gemm.py +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -122,11 +122,13 @@ def _run_capturing_fallback(fn): (512, 512, 512), (512, 1024, 512), (1024, 512, 1024), + (512, 512, 1024), # M != N: exercises the operand-swap contract asymmetrically ] MXFP8_SHAPES = [ (512, 512, 512), (512, 1024, 512), + (512, 512, 1024), # M != N ] LAYOUTS = ["TN", "NN", "NT"] @@ -174,13 +176,19 @@ def cleanup_env(): def get_shapes(layout, M, K, N): - """Return the A/B storage shapes used by TE's public GEMM tests.""" - if layout == "TN": + """Return the (A, B) storage shapes for a given TE ``general_gemm`` layout. + + Every layout produces a logical ``(N, M)`` output (see + ``compute_pytorch_reference``), so each must reference both M and N -- + otherwise NN/NT silently collapse to square outputs and never exercise the + M/N operand-swap contract asymmetrically. + """ + if layout == "TN": # A transposed, B not transposed return (M, K), (N, K) - if layout == "NN": - return (M, K), (K, M) - if layout == "NT": - return (M, K), (M, K) + if layout == "NN": # neither transposed + return (K, M), (N, K) + if layout == "NT": # A not transposed, B transposed + return (K, M), (K, N) raise ValueError(f"Unsupported layout: {layout}") From 5c14856970c280d3bd9ca287401eceea3286f245 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 14:10:16 +0000 Subject: [PATCH 49/65] Fix rank mismatch in multidim FP8 test reference test_flydsl_vs_pytorch_fp8_multidim compared the (batch, N, batch*M) output that general_gemm returns (it keeps B's leading dims for a non-transposed B) against a flattened (batch*N, batch*M) reference, so torch.testing.assert_close raised on the shape difference before checking values. The two are the same fully-flattened matmul, so reshape the reference to output.shape. No value change -- just matching the rank the wrapper produces; verified all 8 parametrizations pass on gfx950. Co-Authored-By: Claude --- tests/pytorch/flydsl_kernels/test_gemm.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py index 4f3d0bb13..b05b0b895 100644 --- a/tests/pytorch/flydsl_kernels/test_gemm.py +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -1119,8 +1119,11 @@ def test_flydsl_vs_pytorch_fp8_multidim( use_flydsl=True, ) + # general_gemm keeps B's leading dims for a non-transposed B, so the output + # is (batch, N, batch*M) while the flattened reference is (batch*N, batch*M). + # The values are the same fully-flattened matmul; reshape to match rank. A_flat = A_fp8.dequantize().reshape(-1, K) B_flat = B_fp8.dequantize().reshape(-1, K) - expected = torch.matmul(B_flat, A_flat.T) + expected = torch.matmul(B_flat, A_flat.T).reshape(output.shape) assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) From 68d7b1792148e492a6451b45803a35111221dec8 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 14:23:49 +0000 Subject: [PATCH 50/65] Add AMD copyright + provenance to adapted FlyDSL GEMM utils fp8_gemm_utils.py and fp16_gemm_utils.py were adapted from the FlyDSL project's GEMM utility helpers, not vendored verbatim: the cdiv/ceildiv import was repointed off FlyDSL's kernels.common.utils onto TE's shared gemm_common_utils, and symbols were added/removed during the dedup refactor. Per the repo's header rules for AMD-modified files, add the AMD copyright line above the existing Apache-2.0 / FlyDSL Project Contributors line and state the provenance explicitly, rather than relocating to 3rdparty/. Co-Authored-By: Claude --- .../pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py | 4 ++++ .../pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py | 4 ++++ 2 files changed, 8 insertions(+) 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 76628bdbd..fb93e9c01 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py @@ -1,5 +1,9 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors +# +# Adapted by AMD from the FlyDSL project's GEMM utility helpers. """Byte-level staging helpers for the four-wave GEMM kernels. These loaders and swizzle helpers operate on flat byte views and carry no dtype 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 1c75b6be7..9562d7086 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py @@ -1,5 +1,9 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors +# +# Adapted by AMD from the FlyDSL project's GEMM utility helpers. import flydsl.expr as fx from flydsl._mlir import ir From 4d0448a21898fe80ab82b55cc2c8802cb40dc981 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 15:56:45 +0000 Subject: [PATCH 51/65] Remove dead symbols from FlyDSL GEMM kernels Prune helpers carried over from the upstream FlyDSL utils but never used by the TE kernels: - fp8_gemm_utils.py: preshuffle_b, compute_global_linear_128x128, StoreC, Mfma16x16x128 (all definition-only), and the now-unused arith / T imports. - fp8_gemm.py and fp32_gemm.py: the LDS_SYM_* / LDS_ALIAS_DOMAIN / SCOPE_IDS module constants, which were defined but never referenced. (wait_barrier, G2STransposeLoader, make_fp32_inputs and the unused _divmod were already removed in earlier refactors.) Pure deletion, no behavior change; imports and test collection unaffected. Co-Authored-By: Claude --- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 7 - .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 7 - .../flydsl_kernels/gemm/fp8_gemm_utils.py | 147 +----------------- 3 files changed, 1 insertion(+), 160 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index 92b538026..b3c4dd2eb 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -77,13 +77,6 @@ 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 diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index be2a8fc9d..f4fff8fbb 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -78,13 +78,6 @@ 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 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 9562d7086..5df836b70 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py @@ -9,8 +9,7 @@ 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 T +from flydsl.expr import 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 @@ -29,13 +28,6 @@ ) -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 @@ -79,23 +71,6 @@ 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) @@ -229,123 +204,3 @@ def load_one_transpose( immediate_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): - 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)) - - -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)]) From 064c1e0382ddefa6f1a2d2baede2e95dd2e4a3c2 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 16:22:50 +0000 Subject: [PATCH 52/65] Remove redundant FP32 double-transpose; core is TN-only The FP32 wrapper transposed operand B into [K, N] and fp32_matmul then transposed it back to [N, K] before launching the core, a redundant round-trip on the TN hot path. Drop the internal transpose so the core consumes B as [N, K] directly, matching every other dtype backend, and move all NN/NT layout normalization into the wrapper. Co-Authored-By: Claude --- .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 20 ++++++------- .../flydsl_kernels/gemm/gemm_wrappers.py | 30 ++++++++++--------- 2 files changed, 25 insertions(+), 25 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index b3c4dd2eb..c53863314 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -7,8 +7,8 @@ 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. +shaped [M, N]. The public ``fp32_matmul`` entry point forwards these operands +to the core directly; NN/NT layout normalization is handled by the wrapper. This module imports ``flydsl`` at import time and must therefore be imported lazily only after FlyDSL availability has been confirmed. @@ -980,15 +980,14 @@ def fp32_matmul( ): """TE-facing TN FP32 GEMM adapter. - Public/backend contract: + Public/backend contract (matching every other dtype backend): a: [M, K] FP32 - b: [K, N] FP32 + b: [N, K] 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. + The optimized core streams both operands with K contiguous and consumes B + as [N, K] directly, so this adapter performs no internal transpose. The + wrapper is responsible for normalizing NN/NT layouts into this contract. """ if a.ndim != 2 or b.ndim != 2: raise ValueError( @@ -996,7 +995,7 @@ def fp32_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)}") if a.dtype != torch.float32 or b.dtype != torch.float32: @@ -1017,8 +1016,7 @@ def fp32_matmul( 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, epilogue=epilogue, bias=bias, aux=aux) + doGemm(a, b, c, stream=stream, epilogue=epilogue, bias=bias, aux=aux) def doGemm( diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 42f700387..83e51203d 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -535,26 +535,27 @@ def _run_fp32_gemm( ): """Normalize FP32 TN/NN/NT inputs to the current kernel's TN interface. - The existing FP32 entry point expects ordinary row-major GEMM operands: + The private FP32 core is TN-only and, like every other dtype backend, + consumes its operands as: a_tn: [M, K] - b_tn: [K, N] + b_tn: [N, K] 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 + b_tn = A NN: a_tn = B - b_tn = A + b_tn = A.T NT: a_tn = B.T - b_tn = A + b_tn = A.T ``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. + operands that are not already in the TN kernel orientation. The core + performs no internal transpose. 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") @@ -572,19 +573,20 @@ def _run_fp32_gemm( 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. + # Swap TE's BLAS operands into the core's TN ownership and materialize the + # NN/NT layouts into the [M, K] x [N, K] kernel contract with explicit + # transpose copies here in the wrapper. The core (like every other dtype + # backend) consumes b as [N, K] directly and performs no internal transpose. + # TODO: Optimize FP32 NN/NT execution. 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 + else: + b_tn = A_flat.transpose(0, 1).contiguous() if not a_tn.is_contiguous(): a_tn = a_tn.contiguous() @@ -598,7 +600,7 @@ def _run_fp32_gemm( ) m, k = a_tn.shape - kb, n = b_tn.shape + n, kb = b_tn.shape if kb != k: layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" raise FlyDSLUnsupportedError( From 6f78bb73f68e23710680f3526e59994aef87e62b Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 16:39:30 +0000 Subject: [PATCH 53/65] Make MXFP8 shape rejections fall back instead of crashing MXFP8 scale-packing and matmul rejected unaligned shapes with plain ValueError, which escapes general_gemm's FlyDSLUnsupportedError-only fallback and crashes the GEMM. Convert the shape-dependent rejections (qk % 4, dim % 64, k % scale-group-size) to FlyDSLUnsupportedError so an untileable MXFP8 GEMM degrades to the C++ backend. Genuine caller contract violations (dtype, device, stride, scale-shape mismatch) keep raising ValueError/TypeError. Co-Authored-By: Claude --- .../pytorch/flydsl_kernels/gemm/gemm_wrappers.py | 2 +- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 83e51203d..df39b5367 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -940,7 +940,7 @@ def _run_mxfp8( ) if k % 32 != 0: - raise ValueError(f"K={k} must be divisible by MXFP8 scale group size 32") + raise FlyDSLUnsupportedError(f"K={k} must be divisible by MXFP8 scale group size 32") if tuple(a_scale.shape) != expected_a_scale: raise ValueError( diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index b75384fdc..6220fdabc 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -218,9 +218,9 @@ def pack_mx32_scales_for_hk( dim, qk = scales_u8.shape if qk % 4 != 0: - raise ValueError(f"Scale K/32 dimension={qk} must be divisible by 4") + raise FlyDSLUnsupportedError(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") + raise FlyDSLUnsupportedError(f"Scale outer dimension={dim} must be a multiple of 64") packed = torch.empty( (qk // 4, dim), @@ -1754,7 +1754,9 @@ def mxfp8_matmul( if tuple(D.shape) != (m, n): 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 {SCALE_GROUP_SIZE}") + raise FlyDSLUnsupportedError( + f"K={k} must be divisible by MXFP8 scale group size {SCALE_GROUP_SIZE}" + ) if layout == "NT": expected_a_scale = (k // SCALE_GROUP_SIZE, m) From 0711978c657cd52f38b5476897ad0f4268bbbed4 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 16:50:47 +0000 Subject: [PATCH 54/65] Harden FlyDSL GEMM fallback, test isolation, and lint Address three review items: - Output-format rejections (non-contiguous output, unsupported output dtype) now raise FlyDSLUnsupportedError across all backends so an unservable-but-valid GEMM falls back to the C++ backend instead of crashing, matching the existing non-contiguous A/B behavior. MXFP8 do_gemm gains an explicit output-contiguity guard so a non-contiguous D no longer silently receives a throwaway .contiguous() copy. - Add tests/pytorch/flydsl_kernels/__init__.py so test_gemm.py gets a package-qualified module name and no longer collides under pytest's prepend import mode with triton_kernels/test_gemm.py. - Silence pylint W0622 on the DSL-safe divmod/min helpers that intentionally shadow builtins. Co-Authored-By: Claude --- tests/pytorch/flydsl_kernels/__init__.py | 10 ++++++++++ .../pytorch/flydsl_kernels/gemm/fp32_gemm.py | 5 +++-- .../pytorch/flydsl_kernels/gemm/fp8_gemm.py | 5 +++-- .../pytorch/flydsl_kernels/gemm/gemm_common_utils.py | 4 ++-- .../pytorch/flydsl_kernels/gemm/gemm_wrappers.py | 6 ++++-- .../pytorch/flydsl_kernels/gemm/half_prec_gemm.py | 4 ++-- .../pytorch/flydsl_kernels/gemm/mxfp8_gemm.py | 12 ++++++++++-- 7 files changed, 34 insertions(+), 12 deletions(-) create mode 100644 tests/pytorch/flydsl_kernels/__init__.py diff --git a/tests/pytorch/flydsl_kernels/__init__.py b/tests/pytorch/flydsl_kernels/__init__.py new file mode 100644 index 000000000..14d651bc8 --- /dev/null +++ b/tests/pytorch/flydsl_kernels/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# License for AMD contributions = MIT. See LICENSE for more information + +"""FlyDSL GEMM test package. + +Gives ``test_gemm.py`` a package-qualified module name so it does not collide +under pytest's default prepend import mode with the identically named +``tests/pytorch/triton_kernels/test_gemm.py``. +""" diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py index c53863314..1b7f9250e 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 .gemm_common_utils import require_block_tiling, require_launch_size from .fp16_gemm_utils import ( G2SLoader, @@ -1006,7 +1007,7 @@ def fp32_matmul( if tuple(c.shape) != (m, n): raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") if c.dtype != torch.float32: - raise TypeError( + raise FlyDSLUnsupportedError( f"The current FlyDSL FP32 kernel stores torch.float32 output, got {c.dtype}" ) if a.device != b.device or a.device != c.device: @@ -1014,7 +1015,7 @@ def fp32_matmul( 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 FP32 GEMM requires contiguous output storage") + raise FlyDSLUnsupportedError("FlyDSL FP32 GEMM requires contiguous output storage") doGemm(a, b, c, stream=stream, epilogue=epilogue, bias=bias, aux=aux) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py index f4fff8fbb..7a9427f95 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -25,6 +25,7 @@ from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec +from .exceptions import FlyDSLUnsupportedError from .gemm_common_utils import require_block_tiling, require_launch_size # Transformer Engine-local FlyDSL utilities. @@ -1066,11 +1067,11 @@ def fp8_matmul( 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( + raise FlyDSLUnsupportedError( f"FlyDSL FP8 supports only float16, bfloat16, and float32 outputs, got {c.dtype}" ) if not c.is_contiguous(): - raise ValueError("FlyDSL FP8 requires contiguous output storage") + raise FlyDSLUnsupportedError("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:]): diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py index ce3c75e35..da8ceb476 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py @@ -76,7 +76,7 @@ def cdiv(numer: int, denom: int) -> int: ceildiv = cdiv -def divmod(a, b): +def divmod(a, b): # pylint: disable=redefined-builtin """Integer divmod that works on DSL values (e.g. ``Int32``). The builtin ``divmod`` rejects DSL scalar types, so this uses the overloaded @@ -85,7 +85,7 @@ def divmod(a, b): return (a // b, a % b) -def min(a, b): +def min(a, b): # pylint: disable=redefined-builtin return arith.select(a < b, a, b) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index df39b5367..6e5691426 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -295,11 +295,13 @@ def _validate_or_allocate_output( if tuple(D.shape) != tuple(shape): raise ValueError(f"D shape {tuple(D.shape)} does not match expected {tuple(shape)}") if D.dtype != dtype: - raise TypeError(f"FlyDSL {backend_name} requires {dtype} output, got {D.dtype}") + raise FlyDSLUnsupportedError( + f"FlyDSL {backend_name} requires {dtype} output, got {D.dtype}" + ) if D.device != device: raise ValueError(f"D must be on {device}, got {D.device}") if not D.is_contiguous(): - raise ValueError(f"FlyDSL {backend_name} requires contiguous output storage") + raise FlyDSLUnsupportedError(f"FlyDSL {backend_name} requires contiguous output storage") return D diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py index 06ba10e19..36ed68d5b 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/half_prec_gemm.py @@ -1325,7 +1325,7 @@ def _half_prec_matmul( 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(f"FlyDSL {label} GEMM requires contiguous output storage") + raise FlyDSLUnsupportedError(f"FlyDSL {label} GEMM requires contiguous output storage") doGemm( a, @@ -1449,7 +1449,7 @@ def doGemm( f"{label} {layout} requires {input_dtype} inputs, got {A.dtype} and {B.dtype}" ) if C.dtype not in (torch.float16, torch.bfloat16, torch.float32): - raise TypeError(f"Unsupported {label} output dtype: {C.dtype}") + raise FlyDSLUnsupportedError(f"Unsupported {label} output dtype: {C.dtype}") require_block_tiling( M_runtime, diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py index 6220fdabc..375f11d5a 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -1626,12 +1626,18 @@ def do_gemm( if stream is None: stream = torch.cuda.current_stream() + # A non-contiguous output cannot receive the kernel's writes: the flat + # descriptor below would target a throwaway ``.contiguous()`` copy. Reject + # it here so ``general_gemm`` falls back rather than silently dropping C. + if not C.is_contiguous(): + raise FlyDSLUnsupportedError("FlyDSL MXFP8 requires contiguous output storage") + # 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) + C_arg = C.view(-1) # DEFAULT keeps the kernel signature uniform with dummy 1-element buffers. if needs_bias: Bias_arg = bias.contiguous().view(-1) @@ -1704,7 +1710,9 @@ def _validate_common_payloads( 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(f"FlyDSL MXFP8 output must be float16, bfloat16, or float32, got {D.dtype}") + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 output must be float16, bfloat16, or float32, got {D.dtype}" + ) def mxfp8_matmul( From 0b89f0a5f2ab25bd052a09a6bcba2742a837d4cd Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Tue, 11 Aug 2026 18:09:02 +0000 Subject: [PATCH 55/65] Add MXFP8 fallback regression test Verify that an MXFP8 shape FlyDSL cannot tile (M not a multiple of the 256-wide M tile) degrades to the C++ backend via general_gemm rather than crashing, and that the fallen-back result is numerically correct. Uses an M-tiling mismatch rather than an untileable K, since the C++ MXFP8 backend shares FlyDSL's K%128 constraint and could not serve a K-based rejection. Co-Authored-By: Claude --- tests/pytorch/flydsl_kernels/test_gemm.py | 62 +++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py index b05b0b895..8f0229fb1 100644 --- a/tests/pytorch/flydsl_kernels/test_gemm.py +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -474,6 +474,68 @@ def test_flydsl_vs_pytorch_mxfp8(M, K, N, layout, fp8_format): assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) +def test_flydsl_mxfp8_unsupported_shape_falls_back(): + """An MXFP8 shape FlyDSL cannot tile must fall back, not crash. + + M=128 is not a multiple of the 256-wide kernel M tile, so FlyDSL rejects it + with ``FlyDSLUnsupportedError`` -- the only type ``general_gemm`` catches. + Dispatch must degrade to the C++ backend rather than propagating an + uncatchable ValueError/RuntimeError, and the fallen-back result must still + be correct. + + An M/N-tiling mismatch is used rather than an untileable K: the C++ MXFP8 + backend shares FlyDSL's K%128 constraint, so a K-based rejection would also + fail C++ and could never produce a comparable result. K=512 here keeps the + scale packing and K-tile count valid, isolating the M-tiling rejection. + """ + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = "1" + os.environ["NVTE_USE_FLYDSL"] = "1" + torch.manual_seed(42) + + M, K, N = 128, 512, 256 + fp8_dtype = tex.DType.kFloat8E4M3 + A_mxfp8, B_mxfp8, A_deq, B_deq = create_mxfp8_tensors( + M, + K, + N, + "TN", + fp8_dtype, + fp8_dtype, + ) + + old_warn = os.environ.get("NVTE_FLYDSL_GEMM_WARN_FALLBACK") + os.environ["NVTE_FLYDSL_GEMM_WARN_FALLBACK"] = "1" + try: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + output, _, _, _ = general_gemm( + A=A_mxfp8, + B=B_mxfp8, + out_dtype=torch.float32, + layout="TN", + bias=None, + quantization_params=None, + gelu=False, + grad=False, + accumulate=False, + ) + finally: + if old_warn is None: + os.environ.pop("NVTE_FLYDSL_GEMM_WARN_FALLBACK", None) + else: + os.environ["NVTE_FLYDSL_GEMM_WARN_FALLBACK"] = old_warn + + fallback_msgs = [str(w.message) for w in caught if _FLYDSL_FALLBACK_TAG in str(w.message)] + for msg in fallback_msgs: + print(msg) # visible with ``pytest -s`` + assert fallback_msgs, ( + "an untileable MXFP8 shape should fall back to the C++ backend, " + "but no [FLYDSL WARNING] fallback was emitted" + ) + expected = compute_pytorch_reference(A_deq.float(), B_deq.float(), "TN") + assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) + + # ============================================================================== # Approach 2: FlyDSL vs native C++ ``generic_gemm`` reference # ============================================================================== From d8481b6ff6bddadd6d73c0d5c4a5905db536212b Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 12 Aug 2026 13:56:40 +0000 Subject: [PATCH 56/65] Gate FlyDSL CI suite on gfx950, detected once Two fixes to the FlyDSL test invocation in run_test_config: - run_test_config runs once per fused-attn backend, so gating on check_mxfp8_supported re-spawned the availability subprocess up to 5 times. Detect FlyDSL support once and cache it in _FLYDSL_SUPPORTED. - FlyDSL covers fp32/fp16/bf16 and tensor-wise FP8 GEMMs, not just MXFP8. Gate on gfx950 -- the arch FlyDSL dispatch actually requires -- instead of MXFP8 availability, so the non-MXFP8 cases run too. NVTE_ROCM_ENABLE_MXFP8=1 stays on the invocation so the MXFP8 cases still run where supported. Co-Authored-By: Claude --- ci/pytorch.sh | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/ci/pytorch.sh b/ci/pytorch.sh index a3e42ebf7..98de11265 100755 --- a/ci/pytorch.sh +++ b/ci/pytorch.sh @@ -49,6 +49,24 @@ check_mxfp8_supported() { fi } +# FlyDSL GEMM dispatch requires gfx950 (see cpp_extensions/gemm.py); on other +# archs general_gemm silently runs the C++ backend, so the FlyDSL suite would +# exercise the wrong path. Detect once and cache, since run_test_config runs +# per fused-attn backend. Not MXFP8-gated: the suite also covers fp32/16/8. +_FLYDSL_SUPPORTED="" +check_flydsl_supported() { + if [ -z "$_FLYDSL_SUPPORTED" ]; then + _cap=$(python -c "${PYTHON_TE_IMPORT}; from transformer_engine.pytorch.utils import get_device_compute_capability; print(get_device_compute_capability())" 2>/dev/null) + if [ "$_cap" = "(9, 5)" ]; then + _FLYDSL_SUPPORTED="yes" + else + _FLYDSL_SUPPORTED="no" + echo "FlyDSL GEMM requires gfx950, skipping FlyDSL tests" >&2 + fi + fi + [ "$_FLYDSL_SUPPORTED" = "yes" ] +} + run_test_config(){ echo ==== Run with Fused attention backend: $_fus_attn ==== #_WORKERS_COUNT=$TEST_WORKERS @@ -97,10 +115,10 @@ run_test_config(){ run_default_fa 1 triton_kernels/test_utils.py NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa 1 triton_kernels/test_norms.py NVTE_ROCM_ENABLE_MXFP8=1 NVTE_TEST_TRITON_AUTOTUNE=1 run_default_fa_lbl "autotune" 3 triton_kernels/test_norms.py - # The FlyDSL GEMM kernels currently require gfx950; on other archs - # general_gemm won't select FlyDSL, so gate on the same gfx950 check the - # MXFP8 tests use to avoid silently exercising the C++ backend instead. - check_mxfp8_supported && NVTE_USE_FLYDSL=1 NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa_lbl "flydsl" 1 flydsl_kernels/test_gemm.py + # FlyDSL covers fp32/16/8 and MXFP8 GEMMs; NVTE_ROCM_ENABLE_MXFP8=1 lets the + # MXFP8 cases run where supported. Gate on gfx950 (the arch FlyDSL dispatch + # requires), not MXFP8 availability, so the non-MXFP8 cases still run. + check_flydsl_supported && NVTE_USE_FLYDSL=1 NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa_lbl "flydsl" 1 flydsl_kernels/test_gemm.py run_default_fa 1 test_parallel_cross_entropy.py NVTE_USE_DEQUANTIZE_TRITON=1 NVTE_USE_CAST_TRANSPOSE_TRITON=1 NVTE_USE_RMSNORM_TRITON=1 NVTE_USE_LAYERNORM_TRITON=1 run_default_fa_lbl "triton" 3 test_numerics.py NVTE_USE_CAST_TRANSPOSE_TRITON=1 NVTE_USE_RMSNORM_TRITON=1 run_default_fa_lbl "triton" 1 test_fusible_ops.py From b391dca19b5a529c431514f82414f516176a85f9 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 12 Aug 2026 14:13:40 +0000 Subject: [PATCH 57/65] Simplify FP8 payload reinterpret to gfx950 OCP-only FlyDSL dispatch is gated to gfx950 (CDNA4), which has no FNUZ hardware -- FNUZ is gfx94x-only. The fnuz torch dtypes therefore cannot occur on this path, so drop the dead capability/fnuz branch and map the TE FP8 dtypes directly to the OCP torch dtypes (e4m3fn / e5m2). Remove the now-unused get_device_compute_capability import. TODO: add the FNUZ mapping when gfx942 support lands in a future PR. Co-Authored-By: Claude --- .../flydsl_kernels/gemm/gemm_wrappers.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 6e5691426..4a1dbb09a 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -9,8 +9,6 @@ import torch import transformer_engine_torch as tex -from transformer_engine.pytorch.utils import get_device_compute_capability - from .exceptions import FlyDSLUnsupportedError from .half_prec_gemm import bf16_matmul, fp16_matmul @@ -63,16 +61,17 @@ 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) + """View TE's uint8 payload as the native torch FP8 dtype. + FlyDSL dispatch is gated to gfx950 (CDNA4), which supports only OCP FP8 + (e4m3fn / e5m2); the hardware has no FNUZ format. FNUZ is gfx94x-only, so + the fnuz torch dtypes cannot occur on this path. TODO: add the FNUZ mapping + when gfx942 support is added to the FlyDSL GEMM backends in a future PR. + """ if dtype == tex.DType.kFloat8E4M3: - torch_dtype = torch.float8_e4m3fn if use_ocp_fp8 else torch.float8_e4m3fnuz + torch_dtype = torch.float8_e4m3fn elif dtype == tex.DType.kFloat8E5M2: - torch_dtype = torch.float8_e5m2 if use_ocp_fp8 else torch.float8_e5m2fnuz + torch_dtype = torch.float8_e5m2 else: raise TypeError(f"Unsupported TE FP8 dtype: {dtype}") From 9d549bdee5c72d9302b88c0a3428c9c8f16492b4 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 12 Aug 2026 14:56:21 +0000 Subject: [PATCH 58/65] Drop FlyDSL from install_requires; make it user-installed Per maintainer review: the install_requires entry only takes effect for source builds with NVTE_USE_FLYDSL=1 set at build time, so it never reaches users of the release wheels it was meant to help. Treat FlyDSL like FlashAttention -- an optional backend the user installs themselves. The runtime already degrades gracefully: general_gemm catches ImportError alongside FlyDSLUnsupportedError and falls back to the C++ backend with a warning when flydsl is absent. Co-Authored-By: Claude --- setup.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/setup.py b/setup.py index b7dd1a29c..385c9eef9 100644 --- a/setup.py +++ b/setup.py @@ -195,10 +195,6 @@ 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.append("flydsl==0.3.0") - # Framework-specific requirements if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): if "pytorch" in frameworks: From 3d276b73c0ffde84c38b470dc24709b55f0719b0 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 12 Aug 2026 15:03:36 +0000 Subject: [PATCH 59/65] Warn when NVTE_USE_FLYDSL=1 but flydsl is not installed Now that flydsl is user-installed rather than a declared dependency, a user can set NVTE_USE_FLYDSL=1 without the package present. Previously this fell back to the C++ backend silently unless the opt-in fallback-warning flag was set. Split the dispatch except so a missing-package ImportError always warns once (pointing at `pip install flydsl`) regardless of the flag, while the FlyDSLUnsupportedError path keeps its existing opt-in warning for shapes and configs FlyDSL genuinely cannot serve. Co-Authored-By: Claude --- .../pytorch/cpp_extensions/gemm.py | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index e8a62fe56..ea19a1be5 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -335,6 +335,11 @@ def _nvfp4_row_scaled_gemm_inputs( ) +# Warn only once when NVTE_USE_FLYDSL=1 but the flydsl package is missing, so a +# misconfigured run is surfaced without spamming the per-GEMM hot path. +_flydsl_import_warned = False + + def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -514,7 +519,24 @@ def general_gemm( *args, **kwargs, ) - except (FlyDSLUnsupportedError, ImportError) as exc: + except ImportError as exc: + # NVTE_USE_FLYDSL=1 was requested but the flydsl package is not + # installed. This is a misconfiguration, not an unsupported GEMM + # config, so always warn (once) regardless of the opt-in fallback + # flag before degrading to the default backend. + global _flydsl_import_warned + if not _flydsl_import_warned: + _flydsl_import_warned = True + warnings.warn( + "[FLYDSL WARNING]: NVTE_USE_FLYDSL=1 but the flydsl package " + "is not installed; falling back to the default backend. " + f"Install it (e.g. `pip install flydsl`) to enable it. Reason: {exc}", + UserWarning, + stacklevel=2, + ) + + out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) + except FlyDSLUnsupportedError as exc: warn_fallback = os.environ.get( "NVTE_FLYDSL_GEMM_WARN_FALLBACK", "0", From 3296331cd8800a508bf074ab495f2fb8b460bd97 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 12 Aug 2026 15:07:48 +0000 Subject: [PATCH 60/65] Add AMD provenance header to gemm_common_utils This module carried only the verbatim FlyDSL SPDX/copyright line, but it is AMD-adapted TE integration code (require_block_tiling/require_launch_size raise the TE-local FlyDSLUnsupportedError to drive general_gemm fallback, plus the int32-launch guard and gfx950 helpers). Add the AMD copyright and "Adapted by AMD" provenance line to match fp16_gemm_utils.py/fp8_gemm_utils.py. Co-Authored-By: Claude --- .../pytorch/flydsl_kernels/gemm/gemm_common_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py index da8ceb476..899374203 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_common_utils.py @@ -1,5 +1,9 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors +# +# Adapted by AMD from the FlyDSL project's GEMM utility helpers. """Dtype-independent primitives shared by the FlyDSL GEMM kernels. These helpers carry no dtype specialization, so the per-dtype utils modules From 563f6f7969fe0dc351383f1c1d0d95628238f484 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 12 Aug 2026 15:12:22 +0000 Subject: [PATCH 61/65] Probe system ROCm roots in canonical order in _rocm_init Per review: match the ROCM_PATH resolution order used elsewhere in TE (common/__init__.py, build_tools/utils.py::rocm_path). Honor an explicit ROCM_PATH, then probe /opt/rocm/core and /opt/rocm, then fall back to the rocm-sdk devel root. The previous code only checked bare /opt/rocm and skipped the /opt/rocm/core layout that the ROCm wheels/Dockerfile use. Co-Authored-By: Claude --- build_tools/templates/_rocm_init.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/build_tools/templates/_rocm_init.py b/build_tools/templates/_rocm_init.py index 50c13b454..a58b8fb50 100644 --- a/build_tools/templates/_rocm_init.py +++ b/build_tools/templates/_rocm_init.py @@ -30,12 +30,15 @@ def initialize() -> None: return if not os.getenv("ROCM_PATH"): - # Prefer the system ROCm tree when present: FlyDSL's MLIR linker - # resolution expects that layout to locate ld.lld. Fall back to the - # rocm-sdk devel wheel for wheel-only environments with no system tree. - _system_rocm = "/opt/rocm" - if os.path.exists(_system_rocm): - os.environ["ROCM_PATH"] = _system_rocm + # Prefer a system ROCm tree when present: FlyDSL's MLIR linker resolution + # expects that layout to locate ld.lld. Probe the standard system roots in + # the same order as the rest of TE (see common/__init__.py and + # build_tools/utils.py::rocm_path), then fall back to the rocm-sdk devel + # wheel for wheel-only environments with no system tree. + for _candidate in ("/opt/rocm/core", "/opt/rocm"): + if os.path.exists(_candidate): + os.environ["ROCM_PATH"] = _candidate + break else: os.environ["ROCM_PATH"] = str(get_devel_root()) rocm_sdk.initialize_process(preload_shortnames=list(_PRELOAD_LIBS)) From f10a4165a3c14359e001e8af40808d2442c2f658 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 12 Aug 2026 15:36:44 +0000 Subject: [PATCH 62/65] Require flydsl >= 0.3.0 and document the FlyDSL GEMM backend Since flydsl is user-installed rather than a declared dependency, the version expectation lived nowhere in the repo. Add _MIN_FLYDSL = (0, 3) in flydsl_kernels/gemm/__init__.py as the single source of truth and check it at import time; a missing or too-old package raises ImportError, which general_gemm already catches to warn once and fall back to the default backend. Generalize that warning to cover the too-old case. Document the backend in README.rst: gfx950 is the targeted architecture for now (gfx942 planned), requires flydsl >= 0.3.0 (install it yourself, like flash-attention), enabled via NVTE_USE_FLYDSL=1. Co-Authored-By: Claude --- README.rst | 17 +++++++++++ .../pytorch/cpp_extensions/gemm.py | 14 ++++++---- .../pytorch/flydsl_kernels/gemm/__init__.py | 28 +++++++++++++++++++ 3 files changed, 53 insertions(+), 6 deletions(-) diff --git a/README.rst b/README.rst index 98ba14f55..5fbfd1c5a 100644 --- a/README.rst +++ b/README.rst @@ -224,6 +224,23 @@ These backends are not enabled by default. The following environment variables c When none are set, TE uses multi-stream dispatch (one hipBLASLt GEMM per expert). +FlyDSL GEMM Backend on ROCm +^^^^^^^^^^^^^^^^^^^^^^^^^^^ +ROCm TE provides an optional FlyDSL GEMM backend for dense (non-grouped) GEMMs, covering BF16, FP16, FP32, tensor-wise FP8, and MXFP8. + +Support matrix: + +* **Architecture** -- gfx950 (CDNA4) is the targeted architecture for now. On any other architecture the backend is not selected and TE uses its default GEMM path. gfx942 support is planned for a future release. +* **FlyDSL package** -- requires ``flydsl >= 0.3.0``. FlyDSL is not a declared dependency of TE (similar to flash-attention); install it yourself with ``pip install flydsl``. + +The backend is off by default and enabled via an environment variable: + +* ``NVTE_USE_FLYDSL=1`` -- dispatch dense GEMMs through FlyDSL when running on gfx950. Requires ``flydsl`` to be installed. +* ``NVTE_FLYDSL_GEMM_WARN_FALLBACK=1`` -- emit a warning whenever a GEMM that FlyDSL cannot serve (unsupported shape/config) falls back to the default backend. Off by default. + +If ``NVTE_USE_FLYDSL=1`` is set but ``flydsl`` is missing or older than ``0.3.0``, TE warns once and falls back to the default GEMM backend. Configurations FlyDSL does not support (e.g. shapes that are not tile-aligned) also fall back transparently. + + Fused Attention Backends on ROCm ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Currently ROCm TE supports two backends, AOTriton and CK, for fused attention. diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index ea19a1be5..b88acdae8 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -520,17 +520,19 @@ def general_gemm( **kwargs, ) except ImportError as exc: - # NVTE_USE_FLYDSL=1 was requested but the flydsl package is not - # installed. This is a misconfiguration, not an unsupported GEMM - # config, so always warn (once) regardless of the opt-in fallback - # flag before degrading to the default backend. + # NVTE_USE_FLYDSL=1 was requested but the flydsl package is + # missing or too old (see flydsl_kernels.gemm._MIN_FLYDSL). This + # is a misconfiguration, not an unsupported GEMM config, so always + # warn (once) regardless of the opt-in fallback flag before + # degrading to the default backend. global _flydsl_import_warned if not _flydsl_import_warned: _flydsl_import_warned = True warnings.warn( "[FLYDSL WARNING]: NVTE_USE_FLYDSL=1 but the flydsl package " - "is not installed; falling back to the default backend. " - f"Install it (e.g. `pip install flydsl`) to enable it. Reason: {exc}", + "is unavailable; falling back to the default backend. " + f"Install a supported version (e.g. `pip install flydsl`) to " + f"enable it. Reason: {exc}", UserWarning, stacklevel=2, ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py index 4eae105ef..baa600be7 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py @@ -4,6 +4,34 @@ """FlyDSL GEMM kernels (dense, non-grouped) for BF16/FP16/FP32/FP8/MXFP8.""" +# Minimum supported flydsl release. Single source of truth for the version +# requirement, since flydsl is user-installed (not a declared dependency). A +# too-old (or missing) package raises ImportError here, which general_gemm +# catches to warn once and fall back to the default backend. +_MIN_FLYDSL = (0, 3) + + +def _check_flydsl_version() -> None: + from importlib.metadata import version + + # PackageNotFoundError subclasses ImportError, so a missing package is + # handled by the same fallback path as a too-old one. + installed = version("flydsl") + try: + major_minor = tuple(int(part) for part in installed.split(".")[:2]) + except ValueError: + # Unparseable version string: the package imported, so let it proceed + # rather than falsely block a valid install. + return + if major_minor < _MIN_FLYDSL: + raise ImportError( + f"flydsl {installed} is installed but the FlyDSL GEMM backend requires " + f">= {_MIN_FLYDSL[0]}.{_MIN_FLYDSL[1]}" + ) + + +_check_flydsl_version() + from .exceptions import FlyDSLUnsupportedError from .gemm_wrappers import te_generic_gemm_flydsl From e118c9692da24e949e9faf9149260844701e91a7 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 12 Aug 2026 19:41:03 +0000 Subject: [PATCH 63/65] Drop redundant CI arch gate; rely on test self-skip The FlyDSL suite already self-gates: a module-level skipif handles non-gfx950 and importorskip handles a missing flydsl package, so on unsupported archs it collects-and-skips (exit 0). The shell-side check_flydsl_supported gate in ci/pytorch.sh duplicated that gfx950 detection, so remove it and invoke the suite directly. Also correct a stale comment in test_gemm.py: flydsl is now a user-installed package, not present only when NVTE_USE_FLYDSL=1 at build time. Co-Authored-By: Claude --- ci/pytorch.sh | 25 ++++------------------- tests/pytorch/flydsl_kernels/test_gemm.py | 4 ++-- 2 files changed, 6 insertions(+), 23 deletions(-) diff --git a/ci/pytorch.sh b/ci/pytorch.sh index 98de11265..2685b99c4 100755 --- a/ci/pytorch.sh +++ b/ci/pytorch.sh @@ -49,24 +49,6 @@ check_mxfp8_supported() { fi } -# FlyDSL GEMM dispatch requires gfx950 (see cpp_extensions/gemm.py); on other -# archs general_gemm silently runs the C++ backend, so the FlyDSL suite would -# exercise the wrong path. Detect once and cache, since run_test_config runs -# per fused-attn backend. Not MXFP8-gated: the suite also covers fp32/16/8. -_FLYDSL_SUPPORTED="" -check_flydsl_supported() { - if [ -z "$_FLYDSL_SUPPORTED" ]; then - _cap=$(python -c "${PYTHON_TE_IMPORT}; from transformer_engine.pytorch.utils import get_device_compute_capability; print(get_device_compute_capability())" 2>/dev/null) - if [ "$_cap" = "(9, 5)" ]; then - _FLYDSL_SUPPORTED="yes" - else - _FLYDSL_SUPPORTED="no" - echo "FlyDSL GEMM requires gfx950, skipping FlyDSL tests" >&2 - fi - fi - [ "$_FLYDSL_SUPPORTED" = "yes" ] -} - run_test_config(){ echo ==== Run with Fused attention backend: $_fus_attn ==== #_WORKERS_COUNT=$TEST_WORKERS @@ -116,9 +98,10 @@ run_test_config(){ NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa 1 triton_kernels/test_norms.py NVTE_ROCM_ENABLE_MXFP8=1 NVTE_TEST_TRITON_AUTOTUNE=1 run_default_fa_lbl "autotune" 3 triton_kernels/test_norms.py # FlyDSL covers fp32/16/8 and MXFP8 GEMMs; NVTE_ROCM_ENABLE_MXFP8=1 lets the - # MXFP8 cases run where supported. Gate on gfx950 (the arch FlyDSL dispatch - # requires), not MXFP8 availability, so the non-MXFP8 cases still run. - check_flydsl_supported && NVTE_USE_FLYDSL=1 NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa_lbl "flydsl" 1 flydsl_kernels/test_gemm.py + # MXFP8 cases run where supported. The suite self-gates on gfx950 and a + # present flydsl package (module-level skipif + importorskip in test_gemm.py), + # so it collects-and-skips on unsupported archs without a shell-side check. + NVTE_USE_FLYDSL=1 NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa_lbl "flydsl" 1 flydsl_kernels/test_gemm.py run_default_fa 1 test_parallel_cross_entropy.py NVTE_USE_DEQUANTIZE_TRITON=1 NVTE_USE_CAST_TRANSPOSE_TRITON=1 NVTE_USE_RMSNORM_TRITON=1 NVTE_USE_LAYERNORM_TRITON=1 run_default_fa_lbl "triton" 3 test_numerics.py NVTE_USE_CAST_TRANSPOSE_TRITON=1 NVTE_USE_RMSNORM_TRITON=1 run_default_fa_lbl "triton" 1 test_fusible_ops.py diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py index 8f0229fb1..36e96186e 100644 --- a/tests/pytorch/flydsl_kernels/test_gemm.py +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -66,8 +66,8 @@ def _device_capability(): # All FlyDSL GEMM dispatch is gated on gfx950 in cpp_extensions/gemm.py, not # just MXFP8. On any other arch general_gemm silently runs the C++ backend, so # every test here would either exercise hipBLASLt or (for the vs-cpp cases) be -# a tautology. Skip the whole module unless we are on gfx950 with FlyDSL -# installed (flydsl is only present when NVTE_USE_FLYDSL=1 at build time). +# a tautology. Skip the whole module unless we are on gfx950 with the +# user-installed flydsl package present (importorskip below). _CAP = _device_capability() has_flydsl_support = _CAP == (9, 5) pytestmark = pytest.mark.skipif( From 777a53d808b55655be8f4ea8da6e1693d8c4ef02 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 12 Aug 2026 20:15:37 +0000 Subject: [PATCH 64/65] Add GELU_AUX_BIAS test coverage for regular and fp8 backends The GELU_AUX_BIAS epilogue (fused GELU over A@B + bias, saving the pre-activation aux) is implemented in every FlyDSL GEMM backend, but only the mxfp8 path had a test. Add the matching regular (fp32/fp16/bf16) and tensor-wise fp8 tests so every dtype variant that supports the epilogue is exercised, mirroring test_flydsl_vs_pytorch_mxfp8_gelu_bias. Co-Authored-By: Claude --- tests/pytorch/flydsl_kernels/test_gemm.py | 68 +++++++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py index 36e96186e..253b8097a 100644 --- a/tests/pytorch/flydsl_kernels/test_gemm.py +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -974,6 +974,37 @@ def test_flydsl_vs_pytorch_regular_gelu(M, K, N, layout, dtype): assert_gemm_close(output, gelu_tanh_ref(pre_act), atol=1e-3, rtol=1e-2) +@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_gelu_bias(M, K, N, layout, dtype): + """Regular fp32/fp16/bf16 forward GELU_AUX_BIAS vs PyTorch: bias folded before GELU, aux saved.""" + 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 + + ab = compute_pytorch_reference(A.float(), B.float(), layout) + out_features = ab.shape[-1] + bias = torch.randn(out_features, dtype=torch.float32, device="cuda") + pre_act = ab + bias.float() + + output, gelu_input = call_gemm_with_gelu( + A, + B, + layout, + out_dtype=dtype, + bias=bias, + use_flydsl=True, + ) + assert gelu_input is not None, "GELU_AUX_BIAS did not return the pre-activation aux." + + # Aux is the post-bias pre-activation (A@B + bias); output is gelu of it. + assert_gemm_close(gelu_input, pre_act, atol=1e-3, rtol=1e-2) + assert_gemm_close(output, gelu_tanh_ref(pre_act), 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) @@ -1017,6 +1048,43 @@ def test_flydsl_vs_pytorch_fp8_gelu(M, K, N, layout, fp8_format): assert_gemm_close(output, gelu_tanh_ref(pre_act), atol=5e-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_gelu_bias(M, K, N, layout, fp8_format): + """Tensor-wise FP8 forward GELU_AUX_BIAS vs PyTorch: bias folded before GELU, aux saved.""" + 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, + ) + + ab = compute_pytorch_reference(A_deq.float(), B_deq.float(), layout) + out_features = ab.shape[-1] + bias = torch.randn(out_features, dtype=torch.float32, device="cuda") + pre_act = ab + bias.float() + + output, gelu_input = call_gemm_with_gelu( + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + bias=bias, + use_flydsl=True, + ) + assert gelu_input is not None, "GELU_AUX_BIAS did not return the pre-activation aux." + + # Aux is the post-bias pre-activation (A@B + bias); output is gelu of it. + assert_gemm_close(gelu_input, pre_act, atol=5e-3, rtol=1e-2) + assert_gemm_close(output, gelu_tanh_ref(pre_act), atol=5e-3, rtol=1e-2) + + @requires_mxfp8_support @pytest.mark.parametrize("M, K, N", MXFP8_SHAPES) @pytest.mark.parametrize("layout", LAYOUTS) From fd65884215a5b6665ebf0871e8480dfadd310cac Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 12 Aug 2026 20:19:45 +0000 Subject: [PATCH 65/65] Add epilogue-inactivity guard to GELU_AUX_BIAS tests Mirror the plain-GELU tests: compare the fused output against a no-epilogue call_gemm and assert they differ, so a GELU_AUX_BIAS epilogue that silently reverted to DEFAULT would fail instead of passing vacuously. Added to the regular, fp8, and mxfp8 gelu_bias tests (the mxfp8 one predated the guard). Co-Authored-By: Claude --- tests/pytorch/flydsl_kernels/test_gemm.py | 27 +++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py index 253b8097a..94dd26f88 100644 --- a/tests/pytorch/flydsl_kernels/test_gemm.py +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -1000,6 +1000,11 @@ def test_flydsl_vs_pytorch_regular_gelu_bias(M, K, N, layout, dtype): ) assert gelu_input is not None, "GELU_AUX_BIAS did not return the pre-activation aux." + no_gelu_out = call_gemm(A, B, layout, out_dtype=dtype, use_flydsl=True) + assert not torch.allclose( + output.float(), no_gelu_out.float(), atol=1e-4 + ), "FlyDSL output matches the no-epilogue output; GELU_AUX_BIAS appears inactive." + # Aux is the post-bias pre-activation (A@B + bias); output is gelu of it. assert_gemm_close(gelu_input, pre_act, atol=1e-3, rtol=1e-2) assert_gemm_close(output, gelu_tanh_ref(pre_act), atol=1e-3, rtol=1e-2) @@ -1080,6 +1085,17 @@ def test_flydsl_vs_pytorch_fp8_gelu_bias(M, K, N, layout, fp8_format): ) assert gelu_input is not None, "GELU_AUX_BIAS did not return the pre-activation aux." + no_gelu_out = call_gemm( + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + assert not torch.allclose( + output.float(), no_gelu_out.float(), atol=1e-4 + ), "FlyDSL FP8 output matches the no-epilogue output; GELU_AUX_BIAS appears inactive." + # Aux is the post-bias pre-activation (A@B + bias); output is gelu of it. assert_gemm_close(gelu_input, pre_act, atol=5e-3, rtol=1e-2) assert_gemm_close(output, gelu_tanh_ref(pre_act), atol=5e-3, rtol=1e-2) @@ -1174,6 +1190,17 @@ def test_flydsl_vs_pytorch_mxfp8_gelu_bias(M, K, N, layout, fp8_format): ) assert gelu_input is not None, "GELU_AUX_BIAS did not return the pre-activation aux." + no_gelu_out = call_gemm( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + assert not torch.allclose( + output.float(), no_gelu_out.float(), atol=1e-4 + ), "FlyDSL MXFP8 output matches the no-epilogue output; GELU_AUX_BIAS appears inactive." + # Aux is the post-bias pre-activation (A@B + bias); output is gelu of it. assert_gemm_close(gelu_input, pre_act, atol=5e-3, rtol=1e-2) assert_gemm_close(output, gelu_tanh_ref(pre_act), atol=5e-3, rtol=1e-2)