From cf54ce65c7f662ff35423fa2fae8557179731ee3 Mon Sep 17 00:00:00 2001 From: Aristotle Martin Date: Wed, 22 Jul 2026 00:33:21 +0000 Subject: [PATCH 01/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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/43] 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 f429253d8efe126d12560cc511e6591a21335077 Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Wed, 5 Aug 2026 18:25:30 +0000 Subject: [PATCH 33/43] Add FlyDSL permute-free MoE grouped GEMM kernels and integrate into GroupedLinear Introduces FlyDSL-based permute-free MoE grouped GEMM support for ROCm: - New tensor shim helpers for FlyDSL pointer/resource handling. - Shared permute-free GEMM utilities and a dense 4-quadrant pipelined MMA loop for bf16. - Forward, dgrad, and wgrad kernels for permute-free gather/route-read grouped GEMMs. - PyTorch integration exposing the permute-free path in `GroupedLinear` behind an environment flag. --- .../pytorch/flydsl_kernels/gemm/bf16_gemm.py | 192 +++ .../flydsl_kernels/gemm/pf_gemm_utils.py | 480 +++++++ .../permute_free_grouped_gemm/pf_dgrad.py | 378 +++++ .../permute_free_grouped_gemm/pf_fwd.py | 390 +++++ .../permute_free_grouped_gemm/pf_wgrad.py | 640 +++++++++ .../pytorch/flydsl_kernels/tensor_shim.py | 36 + .../pytorch/module/grouped_linear.py | 295 +++- transformer_engine/pytorch/moe/__init__.py | 45 + transformer_engine/pytorch/moe/moe_routing.py | 138 ++ .../pytorch/moe/permute_free_grouped_gemm.py | 1258 +++++++++++++++++ .../pytorch/moe/pf_fwd_wrapper.py | 461 ++++++ .../pytorch/moe/pf_helper_kernels.py | 825 +++++++++++ .../pytorch/moe/pf_wgrad_wrapper.py | 227 +++ 13 files changed, 5339 insertions(+), 26 deletions(-) create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/pf_gemm_utils.py create mode 100644 transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py create mode 100644 transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py create mode 100644 transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py create mode 100644 transformer_engine/pytorch/flydsl_kernels/tensor_shim.py create mode 100644 transformer_engine/pytorch/moe/__init__.py create mode 100644 transformer_engine/pytorch/moe/moe_routing.py create mode 100644 transformer_engine/pytorch/moe/permute_free_grouped_gemm.py create mode 100644 transformer_engine/pytorch/moe/pf_fwd_wrapper.py create mode 100644 transformer_engine/pytorch/moe/pf_helper_kernels.py create mode 100644 transformer_engine/pytorch/moe/pf_wgrad_wrapper.py diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py index cf267d695..7764803a4 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py @@ -25,6 +25,7 @@ import flydsl.compiler as flyc import flydsl.expr as fx +from flydsl.compiler.ast_rewriter import ASTRewriter 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 @@ -41,6 +42,7 @@ pack_i32x4_i32x8, swizzle_128, ) +from .fp8_gemm_utils import wait_barrier _BLOCK_M = 256 @@ -1428,3 +1430,193 @@ def doGemm( N_runtime, stream=stream, ) + + +@ASTRewriter.transform +def dense_mma_pipeline_bf16( + lds, + a_g2s, + b_g2s, + a_s2r, + b_s2r, + mfma, + store_c, + A0_gl_offset, + A1_gl_offset, + B0_gl_offset, + B1_gl_offset, + a_k_step, + b_k_step, + block_m, + block_n, + wave_m, + wave_n, + K, + BLOCK_M, + BLOCK_N, + nt_vmcnt, + a_g2s_hi=None, +): + """Shared 4-quadrant pipelined MMA loop + store epilogue for the fixed-K bf16 tile (NT/NN/TN). + + ``a_g2s_hi`` (optional): a second A global->LDS loader used ONLY for the upper LDS + half-tile (rows [LDS_BLOCK_M, BLOCK_M)). Defaults to ``a_g2s`` (contiguous dense/grouped + GEMM, where both halves share one swizzle). A gathering GEMM passes a distinct loader whose + per-lane offsets are redirected through the gather index for the upper rows. + """ + a_g2s_hi = a_g2s if a_g2s_hi is None else a_g2s_hi + K_ITERS = K // BLOCK_K + assert K_ITERS >= 2, f"K_ITERS={K_ITERS} too small; need K >= {2 * BLOCK_K}" + N_TILES_A = BLOCK_M // 128 + N_TILES_B = BLOCK_N // 256 + N_ACCUMS = N_TILES_A * N_TILES_B + LDS_BLOCK_M = BLOCK_M // 2 + LDS_BLOCK_N = BLOCK_N // 2 + N_LDS_STEPS_A = LDS_BLOCK_M // 64 + N_LDS_STEPS_B = LDS_BLOCK_N // 64 + + a_cur0 = lds.A_lds_cur_0 + a_cur1 = lds.A_lds_cur_1 + a_next0 = lds.A_lds_next_0 + a_next1 = lds.A_lds_next_1 + b_cur0 = lds.B_lds_cur_0 + b_cur1 = lds.B_lds_cur_1 + b_next0 = lds.B_lds_next_0 + b_next1 = lds.B_lds_next_1 + + c00_frag = [mfma.zero_value] * N_ACCUMS + c01_frag = [mfma.zero_value] * N_ACCUMS + c10_frag = [mfma.zero_value] * N_ACCUMS + c11_frag = [mfma.zero_value] * N_ACCUMS + + b_g2s.load(b_cur0, B0_gl_offset + 0 * b_k_step) + a_g2s.load(a_cur0, A0_gl_offset + 0 * a_k_step) + b_g2s.load(b_cur1, B1_gl_offset + 0 * b_k_step) + a_g2s_hi.load(a_cur1, A1_gl_offset + 0 * a_k_step) + + if wave_m == 1: + rocdl.s_barrier() + wait_barrier(N_LDS_STEPS_A + N_LDS_STEPS_B) + + b_g2s.load(b_next0, B0_gl_offset + 1 * b_k_step) + a_g2s.load(a_next0, A0_gl_offset + 1 * a_k_step) + b_g2s.load(b_next1, B1_gl_offset + 1 * b_k_step) + + wait_barrier(N_LDS_STEPS_A + 2 * N_LDS_STEPS_B) + + for k in range_constexpr(K_ITERS - 2): + b0_frag = b_s2r.load(b_cur0) + a0_frag = a_s2r.load(a_cur0) + a_g2s_hi.load(a_next1, A1_gl_offset + (k + 1) * a_k_step) + rocdl.s_barrier() + + rocdl.s_setprio(1) + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + b1_frag = b_s2r.load(b_cur1) + b_g2s.load(b_cur0, B0_gl_offset + (k + 2) * b_k_step) + rocdl.s_barrier() + + rocdl.s_setprio(1) + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + a1_frag = a_s2r.load(a_cur1) + a_g2s.load(a_cur0, A0_gl_offset + (k + 2) * a_k_step) + rocdl.s_barrier() + + rocdl.s_setprio(1) + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + b_g2s.load(b_cur1, B1_gl_offset + (k + 2) * b_k_step) + wait_barrier(2 * N_LDS_STEPS_A + N_LDS_STEPS_B) + + rocdl.s_setprio(1) + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + if const_expr(nt_vmcnt >= 0): + llvm.inline_asm( + res=None, + operands_=[], + asm_string=f"s_waitcnt vmcnt({nt_vmcnt})", + constraints="", + has_side_effects=True, + ) + a_cur0, a_next0 = a_next0, a_cur0 + a_cur1, a_next1 = a_next1, a_cur1 + b_cur0, b_next0 = b_next0, b_cur0 + b_cur1, b_next1 = b_next1, b_cur1 + + k = K_ITERS - 2 + b0_frag = b_s2r.load(b_cur0) + a0_frag = a_s2r.load(a_cur0) + rocdl.s_barrier() + rocdl.s_setprio(1) + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + b1_frag = b_s2r.load(b_cur1) + rocdl.s_barrier() + rocdl.s_setprio(1) + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + a1_frag = a_s2r.load(a_cur1) + rocdl.s_barrier() + rocdl.s_setprio(1) + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + b0_frag = b_s2r.load(b_next0) + a_g2s_hi.load(a_next1, A1_gl_offset + (k + 1) * a_k_step) + rocdl.s_barrier() + rocdl.s_setprio(1) + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + a_cur0, a_next0 = a_next0, a_cur0 + a_cur1, a_next1 = a_next1, a_cur1 + b_cur0, b_next0 = b_next0, b_cur0 + b_cur1, b_next1 = b_next1, b_cur1 + + a0_frag = a_s2r.load(a_cur0) + wait_barrier(0) + rocdl.s_setprio(1) + c00_frag = mfma.call(a0_frag, b0_frag, c00_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + b1_frag = b_s2r.load(b_cur1) + rocdl.s_barrier() + rocdl.s_setprio(1) + c01_frag = mfma.call(a0_frag, b1_frag, c01_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + a1_frag = a_s2r.load(a_cur1) + rocdl.s_barrier() + rocdl.s_setprio(1) + c10_frag = mfma.call(a1_frag, b0_frag, c10_frag) + c11_frag = mfma.call(a1_frag, b1_frag, c11_frag) + rocdl.s_setprio(0) + rocdl.s_barrier() + + wave_n_offset = wave_n * (N_TILES_B * 32) + wave_m_offset = wave_m * (N_TILES_A * 32) + base_row = block_m * BLOCK_M + wave_m_offset + base_col = block_n * BLOCK_N + wave_n_offset + store_c.store(c00_frag, base_row + 0, base_col + 0) + store_c.store(c01_frag, base_row + 0, base_col + LDS_BLOCK_N) + store_c.store(c10_frag, base_row + LDS_BLOCK_M, base_col + 0) + store_c.store(c11_frag, base_row + LDS_BLOCK_M, base_col + LDS_BLOCK_N) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/pf_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/pf_gemm_utils.py new file mode 100644 index 000000000..1bc109412 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/pf_gemm_utils.py @@ -0,0 +1,480 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL helpers for permute-free MoE grouped GEMM (Mega 8-wave / 32x32x16). + +Shared building blocks for ``pf_fwd.py`` and ``pf_dgrad.py``. Dense 4-wave BF16 +utilities live in ``fp16_gemm_utils.py``; the pipelined MMA loop lives in +``bf16_gemm.py`` as ``dense_mma_pipeline_bf16``. +""" + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as fly_dialect +from flydsl._mlir.dialects import llvm as _llvm +from flydsl.expr import buffer_ops +from flydsl.expr.arith import _to_raw as _raw +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec +from flydsl.expr.utils.arith import ArithValue +from flydsl.expr import arith, range_constexpr +from flydsl.expr.buffer_ops import _unwrap_value, buffer_store, create_buffer_resource + +from .bf16_gemm import BLOCK_K, dense_mma_pipeline_bf16 +from .fp16_gemm_utils import G2SLoader, ceildiv, make_bf16_buffer_tensor, swizzle_128 + + +def _inttoptr_lds(byte_addr): + """Integer byte address -> !llvm.ptr<3> (LDS). Parsed per call: the type is + bound to the current MLIRContext and cannot be cached across compiles.""" + return _llvm.inttoptr(ir.Type.parse("!llvm.ptr<3>"), _raw(fx.Int64(byte_addr))) + + +_gep = buffer_ops.get_element_ptr + + +def _lds_ptr_from_i32(addr_i32, byte_offset=0): + """Build an LDS pointer (ptr<3>) from an i32 byte address + optional static offset.""" + ptr = _inttoptr_lds(ArithValue(addr_i32).extui(T.i64)) + if byte_offset != 0: + ptr = _gep(ptr, static_byte_offset=byte_offset) + return ptr + + +def _packed_ds_read_tr16(base_ptr, byte_offsets): + n = len(byte_offsets) + v2i32 = ir.VectorType.get([2], ir.IntegerType.get_signless(32)) + struct_t = _llvm.StructType.get_literal([v2i32] * n) + asm = "\n".join(f"ds_read_b64_tr_b16 ${k}, ${n} offset:{byte_offsets[k]}" for k in range(n)) + constraints = ",".join(["=&v"] * n + ["v"] + ["~{memory}"]) + op = _llvm.InlineAsmOp( + res=struct_t, + operands_=[_raw(base_ptr)], + asm_string=asm, + constraints=constraints, + has_side_effects=True, + ) + return [Vec(_llvm.extractvalue(v2i32, op.result, [k])).bitcast(fx.BFloat16) for k in range(n)] + + +class _S2RLoaderBase: + """Shared ctor for LDS->register operand loaders: caches the per-lane id, + this wave's tile index, and the tile count.""" + + 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 + + +class _S2RLoaderBf16(_S2RLoaderBase): + """Shared skeleton for the bf16 operand loaders (cf. _MfmaBf16): n_tiles output + tiles, each a list of k-sub fragments. Subclasses supply the per-sub offset table + and _tile(), which holds the LDS address math -- transposed ds_read_tr_b16 or + swizzled buffer load (too different to share beyond this loop).""" + + def load(self, lds_src): + return [self._tile(lds_src, i) for i in range_constexpr(self.n_tiles)] + + +def _read_tr16_sub(base_i32, sub16, row_off): + """One tr16 sub-block (512 elems/block): packed double-read, the pair 128 bytes + apart, at sub16*512 + row_off, then assembled.""" + ptr = _lds_ptr_from_i32(base_i32 + (sub16 * 512 + row_off) * 2) + r0, r1 = _packed_ds_read_tr16(ptr, [0, 128]) + return r0.shuffle(r1, list(range(8))) + + +class S2RLoaderTrBf16(_S2RLoaderBf16): + """mfma_f32_32x32x16 operand via ds_read_tr_b16 transpose. Like S2RLoaderTr's + _K_BASE, _SUB lists the tr16 sub-block of each inst_k=16 mfma step (consecutive + here); its length is both the sub count and the per-tile block stride.""" + + _SUB = (0, 1, 2, 3) + + def _tile(self, lds_src, i): + m, kblk = self.lane_id % 32, self.lane_id // 32 + row_off = (m // 16) * 256 + kblk * 128 + (m % 16) * 4 + base_i32 = fx.Int32(fx.ptrtoint(lds_src.ptr)) + sub0 = (self.wave_idx * self.n_tiles + i) * len(self._SUB) + return [ + _read_tr16_sub(base_i32, sub0 + self._SUB[c], row_off) + for c in range_constexpr(len(self._SUB)) + ] + + +def _load8_bf16(lds_src, byte_off): + i8 = fx.recast_iter(fx.Uint8, lds_src.ptr) + p = fx.add_offset(i8, fx.make_int_tuple(byte_off)) + v = fx.make_view(p, fx.make_layout(16, 1)).load() + return v.bitcast(fx.BFloat16) + + +class S2RLoaderBf16(_S2RLoaderBf16): + """mfma_f32_32x32x16 operand (swizzled, non-transposed). Mirroring S2RLoaderTr, + _K_BASE lists the K-column (elems) of each inst_k=16 sub of a 32-row tile; its + length is the sub count -- no BLOCK_K needed.""" + + _K_BASE = (0, 16, 32, 48) + + def _tile(self, lds_src, i): + m, kblk = self.lane_id % 32, self.lane_id // 32 + row = self.wave_idx * (self.n_tiles * 32) + i * 32 + m + subs = [] + for c in range_constexpr(len(self._K_BASE)): + col_byte = (self._K_BASE[c] + kblk * 8) * 2 + _, cs = swizzle_128(row, col_byte) + subs.append(_load8_bf16(lds_src, row * 128 + cs)) + return subs + + +class _MfmaBf16: + """Grouped bf16 mfma: accumulate n_tiles_a x n_tiles_b output tiles. The k-sub + count is taken from each operand's fragment list (len(a[i])), so the atom's + (m, n, inst_k) is the only shape this class needs -- no BLOCK_K coupling.""" + + def __init__(self, n_tiles_a, n_tiles_b, m, n, inst_k): + self.atom = fx.make_mma_atom(fx.rocdl.MFMA(m, n, inst_k, fx.BFloat16)) + acc_len = m * n // 64 # f32 accum lanes per wave + self.accum_type = Vec.make_type(acc_len, fx.Float32) + self.zero_value = Vec.filled(acc_len, 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 call(self, a, b, c): + assert len(a) == self.n_tiles_a + assert len(b) == self.n_tiles_b + for i in range_constexpr(self.n_tiles_a): + for j in range_constexpr(self.n_tiles_b): + acc = c[self.idx(i, j)] + for ks in range_constexpr(len(a[i])): + acc = fly_dialect.mma_atom_call_ssa( + [self.accum_type], self.atom, a[i][ks], b[j][ks], acc + ) + c[self.idx(i, j)] = acc + return c + + +class Mfma32x32x16(_MfmaBf16): + def __init__(self, n_tiles_a, n_tiles_b): + super().__init__(n_tiles_a, n_tiles_b, 32, 32, 16) + + +class StoreCBf16: + def __init__(self, C, c_rows, c_cols, out_ty, cache_modifier=0): + self.c_rows = c_rows + self.c_cols = c_cols + self.lane_id = fx.thread_idx.x % 64 + self.out_ty = out_ty + self.cache_modifier = cache_modifier + c_nbytes = c_rows * c_cols * 2 + gC = fx.rocdl.make_buffer_tensor(C, max_size=False, num_records_bytes=c_nbytes) + self.c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + self.out_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), out_ty) + self.reg_out_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), out_ty) + self.c_rsrc = ( + create_buffer_resource(C, max_size=False, num_records_bytes=c_nbytes) + if cache_modifier + else None + ) + self.oob = fx.Int32(c_rows * c_cols) # out-of-bounds sink index + + def _store_masked(self, value, c_index, valid): + """Store one element to c_index (masked to the OOB sink when invalid).""" + idx = arith.select(valid, c_index, self.oob) + val = value.to(self.out_ty) + if self.cache_modifier: + buffer_store(val, self.c_rsrc, fx.Int32(idx), cache_modifier=self.cache_modifier) + else: + fx.memref_store_vec(Vec.filled(1, val, self.out_ty), self.reg_out_1) + fx.copy(self.out_atom_1, self.reg_out_1, fx.slice(self.c_div, (None, fx.Int32(idx)))) + + def store(self, c_frag, base_row, base_col): + n = self.lane_id % 32 + m_hi = (self.lane_id // 32) * 4 + col = base_col + n + col_valid = col < self.c_cols + for ti in range_constexpr(len(c_frag)): + acc = Vec(c_frag[ti]) + for r in range_constexpr(16): + row = base_row + ti * 32 + (r // 4) * 8 + m_hi + (r % 4) + self._store_masked(acc[r], row * self.c_cols + col, col_valid) + + def store16(self, c_frag, base_row, base_col): + n = self.lane_id % 16 + m_hi = (self.lane_id // 16) * 4 + col = base_col + n + col_valid = col < self.c_cols + for ti in range_constexpr(len(c_frag)): + acc = Vec(c_frag[ti]) + for r in range_constexpr(4): + row = base_row + ti * 16 + m_hi + r + self._store_masked(acc[r], row * self.c_cols + col, col_valid) + + def store_trans16(self, c_frag, group_idx, base_m, base_n, out_m, out_n): + n = self.lane_id % 16 + m_hi = (self.lane_id // 16) * 4 + glob_n = base_n + n + n_valid = glob_n < out_n + row_base = (group_idx * out_n + glob_n) * out_m + for ti in range_constexpr(len(c_frag)): + acc = Vec(c_frag[ti]) + for r in range_constexpr(4): + m = base_m + ti * 16 + m_hi + r + self._store_masked(acc[r], row_base + m, n_valid) + + +def compute_global_swizzle_bf16(lane_id, wave_id, K, n_rounds): + offsets = [] + n_waves = fx.block_dim.x // 64 + for r in range_constexpr(n_rounds): + row = lane_id // 8 + wave_id * 8 + r * (n_waves * 8) + col_byte = (lane_id % 8) * 16 + _, c = swizzle_128(row, col_byte) + offsets.append(row * K + c // 2) + return offsets + + +def compute_global_gather_swizzle_bf16(lane_id, wave_id, K, n_rounds, sorted_res, sorted_row_base): + """Per-lane global A offsets for a *gathering* grouped GEMM. + + Identical to :func:`compute_global_swizzle_bf16` except the tile row is redirected through a + gather index: the source row for tile row ``row`` is ``SORTED_IDS[sorted_row_base + row]`` + instead of the contiguous pool row. The bank swizzle (``c``) is still keyed on the *tile* row + so the LDS destination layout (and therefore the S2R transpose-read) is unchanged -- only the + global fetch address is redirected. The index loads happen once (K-invariant), so they are + amortized over the whole K-loop. + """ + offsets = [] + n_waves = fx.block_dim.x // 64 + for r in range_constexpr(n_rounds): + row = lane_id // 8 + wave_id * 8 + r * (n_waves * 8) + col_byte = (lane_id % 8) * 16 + _, c = swizzle_128(row, col_byte) + src_row = buffer_ops.buffer_load( + sorted_res, sorted_row_base + row, vec_width=1, dtype=T.i32 + ) + offsets.append(src_row * K + c // 2) + return offsets + + +def compute_global_swizzle_nn_bf16(lane_id, wave_id, c_n, n_steps): + offsets = [] + n_waves = fx.block_dim.x // 64 + kk = (lane_id % 32) // 2 + g = lane_id // 32 + n_in = g * 16 + (lane_id % 2) * 8 + for step in range_constexpr(n_steps): + idx = wave_id + step * n_waves + n_tile = idx // 4 + ks = idx % 4 + offsets.append((ks * 16 + kk) * c_n + n_tile * 32 + n_in) + return offsets + + +def make_value_attrs(waves_per_eu, agpr_alloc, fwg): + """Kernel value_attrs. agpr_alloc: 0 = compiler default; N>0 = force exactly + N AGPRs ("N,N"); -N = allow up to N ("0,N").""" + d = {"rocdl.waves_per_eu": waves_per_eu, "rocdl.flat_work_group_size": fwg} + if agpr_alloc != 0: + if agpr_alloc < 0: + alloc = f"0,{-agpr_alloc}" + else: + alloc = f"{agpr_alloc},{agpr_alloc}" + d["passthrough"] = [ + ["amdgpu-agpr-alloc", alloc], + ["amdgpu-mfma-vgpr-form", "false"], + ] + return d + + +def xcd_remap_pid(pid, total_pids, num_xcd): + """Remap the tile id so same-XCD workgroups gather into one contiguous + block, keeping each XCD's L2 reuse within that XCD. Bijection over + [0, total_pids); identity when num_xcd <= 1.""" + if num_xcd <= 1: + return pid + per_xcd = total_pids // num_xcd # floor + rem = total_pids - per_xcd * num_xcd + xcd = pid % num_xcd + local = pid // num_xcd + offset = xcd * per_xcd + arith.select(xcd < rem, xcd, rem) + return offset + local + + +def _i64(v): + # widen an i32 runtime value to i64 (avoids overflow in worst-case base offsets) + return ArithValue(arith.extsi(fx.T.i64(), _unwrap_value(v)), signed=True) + + +def _make_shared_storage(BLOCK_M, BLOCK_N): + a_lds_size = (BLOCK_M // 2) * BLOCK_K + b_lds_size = (BLOCK_N // 2) * BLOCK_K + + @fx.struct + class SharedStorage: + A_lds_cur_0: fx.Array[fx.BFloat16, a_lds_size, 16] + A_lds_cur_1: fx.Array[fx.BFloat16, a_lds_size, 16] + A_lds_next_0: fx.Array[fx.BFloat16, a_lds_size, 16] + A_lds_next_1: fx.Array[fx.BFloat16, a_lds_size, 16] + B_lds_cur_0: fx.Array[fx.BFloat16, b_lds_size, 16] + B_lds_cur_1: fx.Array[fx.BFloat16, b_lds_size, 16] + B_lds_next_0: fx.Array[fx.BFloat16, b_lds_size, 16] + B_lds_next_1: fx.Array[fx.BFloat16, b_lds_size, 16] + + return SharedStorage + + +def _gemm_bf16_nn_tn_tile_impl( + A, + B, + C, + c_m, + c_n, + lds, + block_m, + block_n, + *, + a_transpose, + K, + BLOCK_M, + BLOCK_N, + n_blocks=None, + GROUP_M=1, + num_xcd=8, + out_fp16=False, + nt_vmcnt=3, + b_group_base=None, + c_cache_modifier=0, +): + assert BLOCK_M >= 128 and BLOCK_N >= 256 and BLOCK_M % 128 == 0 and BLOCK_N % 256 == 0 + assert K % BLOCK_K == 0, f"bf16 NN/TN needs K % {BLOCK_K} == 0 (got K={K})" + N_TILES_A = BLOCK_M // 128 + N_TILES_B = BLOCK_N // 256 + LDS_BLOCK_M = BLOCK_M // 2 + LDS_BLOCK_N = BLOCK_N // 2 + N_LDS_STEPS_A = LDS_BLOCK_M // 64 + N_LDS_STEPS_B = LDS_BLOCK_N // 64 + N_LDS_ROUNDS = max(N_LDS_STEPS_A, N_LDS_STEPS_B) + + lane_id = fx.thread_idx.x % 64 + wave_id = fx.thread_idx.x // 64 + wave_m = wave_id // 4 + wave_n = wave_id % 4 + + if block_m is None: + num_pid_m = ceildiv(c_m, BLOCK_M) + pid = xcd_remap_pid(fx.block_idx.x, num_pid_m * n_blocks, num_xcd) + num_pid_in_group = GROUP_M * n_blocks + group_id = pid // num_pid_in_group + pid_in_group = pid % num_pid_in_group + first_pid_m = group_id * GROUP_M + remaining_m = num_pid_m - first_pid_m + group_size_m = arith.select(remaining_m < GROUP_M, remaining_m, fx.Int32(GROUP_M)) + block_m = first_pid_m + (pid_in_group % group_size_m) + block_n = pid_in_group // group_size_m + + if a_transpose: + A0_gl_offset = block_m * BLOCK_M + 0 + A1_gl_offset = block_m * BLOCK_M + LDS_BLOCK_M + a_k_step = BLOCK_K * c_m + else: + A0_gl_offset = (block_m * BLOCK_M) * K + A1_gl_offset = (block_m * BLOCK_M + LDS_BLOCK_M) * K + a_k_step = BLOCK_K + B0_gl_offset = block_n * BLOCK_N + 0 + B1_gl_offset = block_n * BLOCK_N + LDS_BLOCK_N + b_k_step = BLOCK_K * c_n + if b_group_base is not None: + B0_gl_offset = B0_gl_offset + b_group_base + B1_gl_offset = B1_gl_offset + b_group_base + + gA = make_bf16_buffer_tensor(A) + gB = make_bf16_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)) + if a_transpose: + gl_off_a = compute_global_swizzle_nn_bf16(lane_id, wave_id, c_m, N_LDS_STEPS_A) + else: + gl_off_a = compute_global_swizzle_bf16(lane_id, wave_id, K, N_LDS_ROUNDS) + gl_off_b = compute_global_swizzle_nn_bf16(lane_id, wave_id, c_n, N_LDS_STEPS_B) + + mfma = Mfma32x32x16(N_TILES_A, N_TILES_B) + a_g2s = G2SLoader(a_div, gl_off_a, N_LDS_STEPS_A, fx.BFloat16.ir_type, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, N_LDS_STEPS_B, fx.BFloat16.ir_type, wave_id) + a_s2r = S2RLoaderTrBf16(wave_m, N_TILES_A) if a_transpose else S2RLoaderBf16(wave_m, N_TILES_A) + b_s2r = S2RLoaderTrBf16(wave_n, N_TILES_B) + _out_ty = fx.Float16 if out_fp16 else fx.BFloat16 + store_c = StoreCBf16(C, c_m, c_n, _out_ty, cache_modifier=c_cache_modifier) + + dense_mma_pipeline_bf16( + lds, + a_g2s, + b_g2s, + a_s2r, + b_s2r, + mfma, + store_c, + A0_gl_offset, + A1_gl_offset, + B0_gl_offset, + B1_gl_offset, + a_k_step, + b_k_step, + block_m, + block_n, + wave_m, + wave_n, + K, + BLOCK_M, + BLOCK_N, + nt_vmcnt, + ) + + +def gemm_bf16_nn_tile( + A, + B, + C, + c_m, + c_n, + lds, + block_m=None, + block_n=None, + *, + K, + BLOCK_M, + BLOCK_N, + n_blocks=None, + GROUP_M=1, + num_xcd=8, + out_fp16=False, + nt_vmcnt=3, + b_group_base=None, + c_cache_modifier=0, +): + _gemm_bf16_nn_tn_tile_impl( + A, + B, + C, + c_m, + c_n, + lds, + block_m, + block_n, + a_transpose=False, + K=K, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + n_blocks=n_blocks, + GROUP_M=GROUP_M, + num_xcd=num_xcd, + out_fp16=out_fp16, + nt_vmcnt=nt_vmcnt, + b_group_base=b_group_base, + c_cache_modifier=c_cache_modifier, + ) diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py new file mode 100644 index 000000000..ddf48cd3d --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Permute-free MoE data-gradient (dgrad) grouped-GEMM (v3): MegaMOE's fast bf16 NN GEMM. + +Backward companion to ``pf_fwd``. dgrad contracts the incoming grad against the +weight over the output-feature axis (NN layout):: + + dXrow[s, k] = sum_n grad[src(s), n] * W[e][n, k] + +There are two flavours, matching the TE permute-free contract (``index_a_by_route_pos``): + +* **route-read** (FC1 dgrad, ``gather=False``): ``grad`` is already compact route order + ``[em_max, N]``; row ``s`` is read directly. ``grad`` and ``dXrow`` share the row count. +* **gather** (FC2 dgrad, ``gather=True``): ``grad`` is token-space ``[num_recv, N]`` and each + route slot ``s`` gathers ``grad[SORTED[s]]`` (sentinel ``SORTED[s] == num_recv`` -> 0), writing + the compact route-order ``dXrow[em_max, K]``. Mirrors the forward NT gather, one row map + redirecting the two LDS A half-tiles, only on the NN tile. + +The weight is bit-identical to the forward ``[E, N, K]`` (forward reads it NT, dgrad NN). + +Contract: + * ``grad_y`` [rows, N] bf16 incoming grad (rows = em_max route-read / num_recv gather) + * ``weight`` [E, N, K] bf16 per-expert weights (shared with forward) + * ``dx`` [em_max, K] bf16 compact expert-major input grad per slot (in place) + * ``expert_ids`` [num_m_blocks] i32 expert id per BLOCK_M slot block + * ``num_tile_blocks`` [1] i32 real (non-padding) BLOCK_M block count (device) + * ``sorted_slot_ids`` [em_max] i32 gather index (gather=True); unused for route-read +""" + +from __future__ import annotations + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import arith +from flydsl.expr.buffer_ops import ( + buffer_load, + create_buffer_resource, + extract_base_index, +) +from flydsl.expr.typing import AddressSpace, PointerType + +from ..gemm.bf16_gemm import BLOCK_K, dense_mma_pipeline_bf16 +from ..gemm.fp16_gemm_utils import G2SLoader, ceildiv, make_bf16_buffer_tensor +from ..gemm.pf_gemm_utils import ( + Mfma32x32x16, + S2RLoaderBf16, + S2RLoaderTrBf16, + StoreCBf16, + _i64, + _make_shared_storage, + compute_global_gather_swizzle_bf16, + compute_global_swizzle_nn_bf16, + gemm_bf16_nn_tile, + make_value_attrs, + xcd_remap_pid, +) + +__all__ = ["compile_grouped_gemm_dgrad_bf16", "grouped_gemm_dgrad_bf16"] + + +def gemm_bf16_nn_gather_tile( + A, # flat [num_recv, N] grad buffer (gather source) + B, # weight [E, N, K] flat + C, # dx tile (already rebased to this row block) + c_n, + lds, + sorted_res, + sorted_row_base, + block_n, + *, + Kc, # contraction dim (forward intermediate feature N) + BLOCK_M, + BLOCK_N, + out_fp16=False, + nt_vmcnt=3, + b_group_base, +): + """One NN dgrad tile with the A (grad) rows *gathered* via ``sorted_slot_ids``. + + Mirrors the forward :func:`pf_fwd.gemm_bf16_nt_gather_tile` (same two-loader row + redirection, A base 0, store block_m 0) but on the NN B path (transpose-read B, ``b_k_step = + BLOCK_K * c_n``), contracting over ``Kc`` (= N). + """ + assert BLOCK_M >= 128 and BLOCK_N >= 256 and BLOCK_M % 128 == 0 and BLOCK_N % 256 == 0 + assert Kc % BLOCK_K == 0, f"NN gather needs N % {BLOCK_K} == 0 (got N={Kc})" + N_TILES_A = BLOCK_M // 128 + N_TILES_B = BLOCK_N // 256 + LDS_BLOCK_M = BLOCK_M // 2 + LDS_BLOCK_N = BLOCK_N // 2 + N_LDS_STEPS_A = LDS_BLOCK_M // 64 + N_LDS_STEPS_B = LDS_BLOCK_N // 64 + N_LDS_ROUNDS = max(N_LDS_STEPS_A, N_LDS_STEPS_B) + + lane_id = fx.thread_idx.x % 64 + wave_id = fx.thread_idx.x // 64 + wave_m = wave_id // 4 + wave_n = wave_id % 4 + + A0_gl_offset = fx.Int32(0) + A1_gl_offset = fx.Int32(0) + B0_gl_offset = block_n * BLOCK_N + b_group_base + B1_gl_offset = block_n * BLOCK_N + LDS_BLOCK_N + b_group_base + + gA = make_bf16_buffer_tensor(A) + gB = make_bf16_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)) + + gl_off_a_lo = compute_global_gather_swizzle_bf16( + lane_id, wave_id, Kc, N_LDS_ROUNDS, sorted_res, sorted_row_base + ) + gl_off_a_hi = compute_global_gather_swizzle_bf16( + lane_id, wave_id, Kc, N_LDS_ROUNDS, sorted_res, sorted_row_base + fx.Int32(LDS_BLOCK_M) + ) + gl_off_b = compute_global_swizzle_nn_bf16(lane_id, wave_id, c_n, N_LDS_STEPS_B) + + mfma = Mfma32x32x16(N_TILES_A, N_TILES_B) + a_g2s = G2SLoader(a_div, gl_off_a_lo, N_LDS_STEPS_A, fx.BFloat16.ir_type, wave_id) + a_g2s_hi = G2SLoader(a_div, gl_off_a_hi, N_LDS_STEPS_A, fx.BFloat16.ir_type, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, N_LDS_STEPS_B, fx.BFloat16.ir_type, wave_id) + a_s2r = S2RLoaderBf16(wave_m, N_TILES_A) + b_s2r = S2RLoaderTrBf16(wave_n, N_TILES_B) + _out_ty = fx.Float16 if out_fp16 else fx.BFloat16 + store_c = StoreCBf16(C, fx.Int32(BLOCK_M), c_n, _out_ty) + + dense_mma_pipeline_bf16( + lds, + a_g2s, + b_g2s, + a_s2r, + b_s2r, + mfma, + store_c, + A0_gl_offset, + A1_gl_offset, + B0_gl_offset, + B1_gl_offset, + BLOCK_K, # a_k_step (contraction rides soffset) + BLOCK_K * c_n, # b_k_step (NN: B is [K, c_n] row-major) + fx.Int32(0), # store block_m: C is already rebased to this tile + block_n, + wave_m, + wave_n, + Kc, + BLOCK_M, + BLOCK_N, + nt_vmcnt, + a_g2s_hi=a_g2s_hi, + ) + + +def _dgrad_gather_body( + GY_flat, GY_tile, WEIGHT, DX_tile, lds, sorted_res, sorted_row_base, block_n, gbase, + *, N, Kout, BLOCK_M, BLOCK_N, out_fp16, nt_vmcnt, +): + """FC2 dgrad tile: gather the token-space grad rows via ``sorted_slot_ids`` (NN gather).""" + gemm_bf16_nn_gather_tile( + GY_flat, WEIGHT, DX_tile, fx.Int32(Kout), lds, sorted_res, sorted_row_base, block_n, + Kc=N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, out_fp16=out_fp16, nt_vmcnt=nt_vmcnt, + b_group_base=gbase, + ) + + +def _dgrad_routeread_body( + GY_flat, GY_tile, WEIGHT, DX_tile, lds, sorted_res, sorted_row_base, block_n, gbase, + *, N, Kout, BLOCK_M, BLOCK_N, out_fp16, nt_vmcnt, +): + """FC1 dgrad tile: read the compact route grad directly (plain Mega NN tile).""" + gemm_bf16_nn_tile( + GY_tile, WEIGHT, DX_tile, fx.Int32(BLOCK_M), fx.Int32(Kout), lds, fx.Int32(0), block_n, + K=N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, out_fp16=out_fp16, nt_vmcnt=nt_vmcnt, + b_group_base=gbase, + ) + + +@functools.lru_cache(maxsize=256) +def compile_grouped_gemm_dgrad_bf16( + N, # contraction dim (forward intermediate feature) + Kout, # output cols (forward hidden input feature) + BLOCK_M=256, + BLOCK_N=256, + GROUP_M=4, + num_xcd=1, + nt_vmcnt=3, + waves_per_eu=2, + agpr_alloc=0, + out_fp16=False, + gather=False, +): + """Compile (cached) the grouped BF16 NN dgrad launcher for one ``(N, Kout, tile)`` combo. + + ``gather=False`` (FC1 dgrad): ``grad`` is compact route order, read per tile (plain Mega NN + tile). ``gather=True`` (FC2 dgrad): ``grad`` is token-space, gathered via ``SORTED`` into the + compact route output (NN gather tile). Both rebase the C tile in i64 to survive worst-case + pools; the grid front-loads via XCD swizzle over the real tile range. + """ + SharedStorage = _make_shared_storage(BLOCK_M, BLOCK_N) + # Compile-time tile selector (plain Python; resolved before the AST rewriter so the kernel + # body stays branch-free -- a device ``if gather`` would be lowered to real control flow). + tile_body = _dgrad_gather_body if gather else _dgrad_routeread_body + + @flyc.kernel(known_block_size=[512, 1, 1]) + def grouped_gemm_dgrad_k( + GRAD_Y: fx.Tensor, + WEIGHT: fx.Tensor, + DX: fx.Tensor, + TILE_TO_GROUP: fx.Tensor, + NUM_TILE_BLOCKS: fx.Int32, + SORTED: fx.Tensor, + A_ELEMS: fx.Int32, + c_m: fx.Int32, + ): + n_blocks = ceildiv(fx.Int32(Kout), BLOCK_N) + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + group_res = create_buffer_resource(TILE_TO_GROUP, max_size=True) + sorted_res = create_buffer_resource(SORTED, max_size=True) + # Real BLOCK_M tile count as a scalar (host-known, capture-safe -- no per-call device + # tensor, which a HIP graph capture forbids). Matches the v2 int contract. + real_tiles = NUM_TILE_BLOCKS + real_grid = real_tiles * n_blocks + + pool_ptr_ty = PointerType.get( + elem_ty=fx.BFloat16.ir_type, address_space=AddressSpace.Global, alignment=16 + ) + gy_base = fx.arith.ArithValue( + arith.index_cast(fx.T.i64(), extract_base_index(GRAD_Y)), signed=True + ) + dx_base = fx.arith.ArithValue( + arith.index_cast(fx.T.i64(), extract_base_index(DX)), signed=True + ) + # Flat 1D grad view (gather source; bounded to A_ELEMS so the sentinel row reads 0). + # Built unconditionally: the route-read tile simply ignores it. + GY_flat = fx.make_view(fx.inttoptr(pool_ptr_ty, gy_base), fx.make_layout(A_ELEMS, 1)) + + def _emit(): + pid = xcd_remap_pid(fx.block_idx.x, real_grid, num_xcd) + num_pid_m = real_tiles + num_pid_in_group = GROUP_M * n_blocks + group_id = pid // num_pid_in_group + pid_in_group = pid % num_pid_in_group + first_pid_m = group_id * GROUP_M + remaining_m = num_pid_m - first_pid_m + group_size_m = arith.select(remaining_m < GROUP_M, remaining_m, fx.Int32(GROUP_M)) + block_m = first_pid_m + (pid_in_group % group_size_m) + block_n = pid_in_group // group_size_m + g_idx = buffer_load(group_res, block_m, vec_width=1, dtype=fx.T.i32()) + # PF padding blocks mark expert_ids=-1; skip the full Mega pipeline (v2 parity). + if g_idx >= fx.Int32(0): + gbase = g_idx * fx.Int32(N) * fx.Int32(Kout) + sorted_row_base = block_m * fx.Int32(BLOCK_M) + + dx_byte_off = _i64(block_m * fx.Int32(BLOCK_M)) * _i64(fx.Int32(Kout)) * fx.Int64(2) + DX_tile = fx.make_view( + fx.inttoptr(pool_ptr_ty, dx_base + dx_byte_off), + fx.make_layout(fx.Int32(BLOCK_M) * fx.Int32(Kout), 1), + ) + # Route-read grad tile (rebased); ignored by the gather tile. + gy_byte_off = _i64(block_m * fx.Int32(BLOCK_M)) * _i64(fx.Int32(N)) * fx.Int64(2) + GY_tile = fx.make_view( + fx.inttoptr(pool_ptr_ty, gy_base + gy_byte_off), + fx.make_layout(fx.Int32(BLOCK_M) * fx.Int32(N), 1), + ) + + tile_body( + GY_flat, GY_tile, WEIGHT, DX_tile, lds, sorted_res, sorted_row_base, block_n, + gbase, N=N, Kout=Kout, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, out_fp16=out_fp16, + nt_vmcnt=nt_vmcnt, + ) + + if fx.block_idx.x < real_grid: + _emit() + + @flyc.jit + def launch(GRAD_Y, WEIGHT, DX, TILE_TO_GROUP, NUM_TILE_BLOCKS: fx.Int32, SORTED, A_ELEMS: fx.Int32, c_m: fx.Int32, stream: fx.Stream): + grid_x = ceildiv(c_m, BLOCK_M) * ceildiv(fx.Int32(Kout), BLOCK_N) + grouped_gemm_dgrad_k( + GRAD_Y, + WEIGHT, + DX, + TILE_TO_GROUP, + NUM_TILE_BLOCKS, + SORTED, + A_ELEMS, + c_m, + value_attrs=make_value_attrs(waves_per_eu, agpr_alloc, "512,512"), + ).launch(grid=(grid_x, 1, 1), block=(512, 1, 1), stream=stream) + + return launch + + +def grouped_gemm_dgrad_bf16( + grad_y, # [rows, N] bf16 incoming grad (rows = em_max route-read / num_recv gather) + weight, # [E, N, K] bf16 per-expert weights (shared with forward) + dx, # [em_max, K] bf16 compact expert-major input grad per slot (in place) + expert_ids, # [num_m_blocks] i32 + num_tile_blocks, # int real BLOCK_M block count (host scalar; capture-safe) + sorted_slot_ids=None, # [em_max] i32 gather index (required when gather=True) + *, + gather=False, + BLOCK_M=256, + BLOCK_N=256, + GROUP_M=4, + num_xcd=1, + nt_vmcnt=3, + waves_per_eu=2, + agpr_alloc=0, +): + """Host entry: grouped bf16 NN dgrad ``dXrow[s] = grad[src(s)] @ W[expert]``. + + ``gather=False`` reads ``grad`` at the route row (FC1 dgrad, ``grad`` rows == ``dx`` rows). + ``gather=True`` gathers ``grad[SORTED[s]]`` from a token-space buffer (FC2 dgrad); ``dx`` is + written over its full padded ``[em_max, K]`` extent (pad rows carry dead values). + """ + assert grad_y.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16 + assert dx.dtype == torch.bfloat16 + E, N, K = weight.shape + c_m = int(dx.shape[0]) + assert grad_y.shape[1] == N, f"grad_y N={grad_y.shape[1]} != weight N={N}" + assert dx.shape[1] == K, f"dx K={dx.shape[1]} != weight K={K}" + if not gather: + assert grad_y.shape[0] == c_m, ( + f"route-read dgrad expects grad_y rows == dx rows ({grad_y.shape[0]} != {c_m}); " + "did you mean gather=True (token-space grad)?" + ) + + grad_y = grad_y.contiguous() + if gather: + a_elems = int(grad_y.numel()) + if a_elems >= (1 << 31): + raise ValueError( + f"grad_y has {a_elems} elems (>= 2^31); the flat-view bound is int32. " + "Large-A support needs a per-lane i64 SRD rebase (TODO)." + ) + else: + # route-read rebases GRAD_Y per tile; the flat gather view (A_ELEMS) is unused. + a_elems = int(BLOCK_M) * int(N) + expert_ids_i32 = expert_ids.to(torch.int32) + if gather: + assert sorted_slot_ids is not None, "gather=True dgrad requires sorted_slot_ids" + sorted_arg = sorted_slot_ids.to(torch.int32) + else: + # route-read reads GRAD_Y at the route row; SORTED is unread. Reuse a live tensor to + # avoid a capture-illegal per-call allocation. + sorted_arg = expert_ids_i32 + + weight_flat = weight.reshape(E * N, K) + if not weight_flat.is_contiguous(): + weight_flat = weight_flat.contiguous() + weight_flat = weight_flat.view(-1) + launch = compile_grouped_gemm_dgrad_bf16( + N=N, + Kout=K, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + GROUP_M=GROUP_M, + num_xcd=num_xcd, + nt_vmcnt=int(nt_vmcnt), + waves_per_eu=int(waves_per_eu), + agpr_alloc=int(agpr_alloc), + gather=bool(gather), + ) + launch( + grad_y, + weight_flat, + dx, + expert_ids_i32, + int(num_tile_blocks), + sorted_arg, + a_elems, + c_m, + stream=torch.cuda.current_stream(), + ) + return dx diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py new file mode 100644 index 000000000..c70590cda --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Permute-free MoE forward *gather* grouped-GEMM (v3): MegaMOE's fast bf16 GEMM + PF gather. + +This is a port of the MegaMOE grouped bf16 GEMM (32x32x16 MFMA, 8-wave / 512-thread +workgroup, deep distance-2 3-buffer DMA ring, XCD swizzle with ``GROUP_M`` front-loading) +into Transformer Engine FlyDSL, with a *single* change: the A operand is fetched +through a per-row gather index instead of a contiguous pool. + +Motivation. The v2 permute-free kernel (16x16x32, gather-in-K-loop) trails MegaMOE's dense +grouped GEMM by ~13%. Swapping v2's MFMA atom to 32x32x16 or deepening ``block_k`` both +regressed, so the lever is not the atom or LDS depth -- it is the gather structure vs. Mega's +pipeline. v3 tests the opposite direction: keep Mega's pipeline *verbatim* and only redirect the +A fetch through ``sorted_slot_ids``, so we pay Mega's throughput while keeping PF's memory model +(no pre-permutation, gather-on-demand). + +Contract (mirrors MegaMOE ``grouped_gemm_bf16_only`` + aiter PF routing metadata): + * ``A`` [num_recv, K] bf16 received-token activations, UNPERMUTED (gather source) + * ``B`` [E, N, K] bf16 per-expert weights, NT (contiguous inner K) + * ``C`` [em_max, N] bf16 compact expert-major output (block-padded, in place) + * ``sorted_slot_ids`` [em_max] i32 received-token row per padded slot (sentinel = num_recv) + * ``expert_ids`` [num_m_blocks] i32 expert id per ``BLOCK_M`` output block (padding tail: + ``-1``; those blocks early-exit like v2) + * ``num_tile_blocks`` [1] i32 real (non-padding) ``BLOCK_M`` block count (device) + +Only the plain (non-activation) forward GEMM is ported here; the fused gated epilogue lives in +v2 and can be layered on later once the GEMM-throughput parity is confirmed. +""" + +from __future__ import annotations + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import arith, range_constexpr +from flydsl.expr.buffer_ops import ( + buffer_load, + create_buffer_resource, + extract_base_index, +) +from flydsl.expr.typing import AddressSpace, PointerType + +from ..gemm.bf16_gemm import BLOCK_K, dense_mma_pipeline_bf16 +from ..gemm.fp16_gemm_utils import G2SLoader, ceildiv, make_bf16_buffer_tensor, swizzle_128 +from ..gemm.pf_gemm_utils import ( + Mfma32x32x16, + S2RLoaderBf16, + StoreCBf16, + _i64, + _make_shared_storage, + compute_global_gather_swizzle_bf16, + compute_global_swizzle_bf16, + make_value_attrs, + xcd_remap_pid, +) + +__all__ = ["compile_grouped_gemm_gather_bf16", "grouped_gemm_gather_bf16"] + + +def compute_global_identity_swizzle_bf16(lane_id, wave_id, K, n_rounds, sorted_row_base): + """Per-lane global A offsets reading the *route row directly* (identity gather). + + Same flat-buffer pipeline as :func:`compute_global_gather_swizzle_bf16` but the source row is + ``sorted_row_base + row`` computed arithmetically instead of loaded from a gather table. This + is the FC2 route-read (``index_a_by_route_pos``) path: it needs no ``sorted_slot_ids`` tensor, + so it stays inside a HIP graph capture (no per-call identity allocation) while reusing the + proven whole-buffer / base-0 gather tile (avoiding the per-tile A-rebase multi-block hazard). + """ + offsets = [] + n_waves = fx.block_dim.x // 64 + for r in range_constexpr(n_rounds): + row = lane_id // 8 + wave_id * 8 + r * (n_waves * 8) + col_byte = (lane_id % 8) * 16 + _, c = swizzle_128(row, col_byte) + offsets.append((sorted_row_base + row) * K + c // 2) + return offsets + + +def gemm_bf16_nt_gather_tile( + A, + B_T, + C, + c_m, + c_n, + lds, + sorted_res, + sorted_row_base, + block_n, + *, + K, + BLOCK_M, + BLOCK_N, + out_fp16=False, + nt_vmcnt=3, + b_group_base, + gather=True, +): + """One NT tile of the grouped GEMM with the A rows *gathered* via ``sorted_slot_ids``. + + Identical to MegaMOE's ``gemm_bf16_nt_tile`` except: + * the two LDS A half-tiles use *gather* swizzles (rows redirected through + ``sorted_slot_ids[sorted_row_base + tile_row]``), so the A base offset is 0 and only the + K-step rides the load ``soffset``; + * ``C`` is already rebased to this tile's row block by the caller, so the store block_m is 0. + """ + assert BLOCK_M >= 128 and BLOCK_N >= 256 and BLOCK_M % 128 == 0 and BLOCK_N % 256 == 0 + assert K % BLOCK_K == 0, f"bf16 NT gather needs K % {BLOCK_K} == 0 (got K={K})" + N_TILES_A = BLOCK_M // 128 + N_TILES_B = BLOCK_N // 256 + LDS_BLOCK_M = BLOCK_M // 2 + LDS_BLOCK_N = BLOCK_N // 2 + N_LDS_STEPS_A = LDS_BLOCK_M // 64 + N_LDS_STEPS_B = LDS_BLOCK_N // 64 + N_LDS_ROUNDS = max(N_LDS_STEPS_A, N_LDS_STEPS_B) + + lane_id = fx.thread_idx.x % 64 + wave_id = fx.thread_idx.x // 64 + wave_m = wave_id // 4 + wave_n = wave_id % 4 + + # A rows carried by the gather offsets -> global base is 0 (K rides soffset only). + A0_gl_offset = fx.Int32(0) + A1_gl_offset = fx.Int32(0) + B0_gl_offset = (block_n * BLOCK_N) * K + B1_gl_offset = (block_n * BLOCK_N + LDS_BLOCK_N) * K + if b_group_base is not None: + B0_gl_offset = B0_gl_offset + b_group_base + B1_gl_offset = B1_gl_offset + b_group_base + + # ``A`` MUST be a flat 1D buffer view (built by the caller): a linear gather offset + # (src_row*K + col) indexes row-major elements. A raw 2D tensor's logical_divide/slice + # indexes the outer (row) dim, so a flat offset runs off the end -> garbage. + gA = make_bf16_buffer_tensor(A) + gB = make_bf16_buffer_tensor(B_T) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + + # Two gather swizzles: the lo half covers tile rows [0, LDS_BLOCK_M), the hi half + # [LDS_BLOCK_M, BLOCK_M); each redirects its tile row through sorted_slot_ids. The two + # halves gather independent rows, so they need distinct loaders (a_g2s / a_g2s_hi). + if gather: + gl_off_a_lo = compute_global_gather_swizzle_bf16( + lane_id, wave_id, K, N_LDS_ROUNDS, sorted_res, sorted_row_base + ) + gl_off_a_hi = compute_global_gather_swizzle_bf16( + lane_id, wave_id, K, N_LDS_ROUNDS, sorted_res, sorted_row_base + fx.Int32(LDS_BLOCK_M) + ) + else: + # FC2 route-read: identity (row = route position), no sorted_slot_ids load. + gl_off_a_lo = compute_global_identity_swizzle_bf16( + lane_id, wave_id, K, N_LDS_ROUNDS, sorted_row_base + ) + gl_off_a_hi = compute_global_identity_swizzle_bf16( + lane_id, wave_id, K, N_LDS_ROUNDS, sorted_row_base + fx.Int32(LDS_BLOCK_M) + ) + gl_off_b = compute_global_swizzle_bf16(lane_id, wave_id, K, N_LDS_ROUNDS) + + mfma = Mfma32x32x16(N_TILES_A, N_TILES_B) + a_g2s = G2SLoader(a_div, gl_off_a_lo, N_LDS_STEPS_A, fx.BFloat16.ir_type, wave_id) + a_g2s_hi = G2SLoader(a_div, gl_off_a_hi, N_LDS_STEPS_A, fx.BFloat16.ir_type, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, N_LDS_STEPS_B, fx.BFloat16.ir_type, wave_id) + a_s2r = S2RLoaderBf16(wave_m, N_TILES_A) + b_s2r = S2RLoaderBf16(wave_n, N_TILES_B) + _out_ty = fx.Float16 if out_fp16 else fx.BFloat16 + store_c = StoreCBf16(C, c_m, c_n, _out_ty) + + dense_mma_pipeline_bf16( + lds, + a_g2s, + b_g2s, + a_s2r, + b_s2r, + mfma, + store_c, + A0_gl_offset, + A1_gl_offset, + B0_gl_offset, + B1_gl_offset, + BLOCK_K, + BLOCK_K, + fx.Int32(0), # store block_m: C is already rebased to this tile + block_n, + wave_m, + wave_n, + K, + BLOCK_M, + BLOCK_N, + nt_vmcnt, + a_g2s_hi=a_g2s_hi, + ) + + +@functools.lru_cache(maxsize=256) +def compile_grouped_gemm_gather_bf16( + K, + BLOCK_M=256, + BLOCK_N=256, + GROUP_M=4, + num_xcd=1, + nt_vmcnt=3, # gfx950 G2S LDS hazard: vmcnt>=4 races (nondeterministic); 3 is det + waves_per_eu=2, + agpr_alloc=0, + out_fp16=False, + gather=True, +): + """Compile (cached) the gathering grouped BF16 GEMM launcher for one ``(K, tile)`` combo. + + Grid is over-launched to the padded output pool; each block early-exits past the real tile + range (``num_tile_blocks``) or when ``expert_ids[block_m] < 0`` (PF padding tail). Mirrors + MegaMOE's ``compile_grouped_gemm_bf16`` (NT) exactly apart from the gather A fetch, the extra + ``SORTED`` argument, and the padding-block guard. Returns the flyc launch callable. + + The FC2 route-read (``index_a_by_route_pos``) case reuses this same gathering kernel with an + identity ``SORTED`` (``SORTED[s] = s``), so no separate no-gather tile is needed. + """ + SharedStorage = _make_shared_storage(BLOCK_M, BLOCK_N) + + @flyc.kernel(known_block_size=[512, 1, 1]) + def grouped_gemm_gather_k( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + TILE_TO_GROUP: fx.Tensor, + NUM_TILE_BLOCKS: fx.Int32, + SORTED: fx.Tensor, + A_ELEMS: fx.Int32, + c_m: fx.Int32, + c_n: fx.Int32, + ): + n_blocks = ceildiv(c_n, BLOCK_N) + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + group_res = create_buffer_resource(TILE_TO_GROUP, max_size=True) + sorted_res = create_buffer_resource(SORTED, max_size=True) + # Real (non-padding) BLOCK_M tile count as a scalar (host-known, capture-safe -- avoids a + # per-call device tensor that a HIP graph capture forbids). Matches the v2 int contract. + real_tiles = NUM_TILE_BLOCKS + # XCD-swizzle over the REAL tile range only (front-loaded); swizzling the full padded + # pool scatters real tiles -> ~2x slower. + real_grid = real_tiles * n_blocks + + pool_ptr_ty = PointerType.get( + elem_ty=fx.BFloat16.ir_type, address_space=AddressSpace.Global, alignment=16 + ) + # Flat 1D view of the whole received-token buffer, bounded to A_ELEMS elements. The + # gather offsets index this row-major buffer directly; the buffer resource clamps the + # padding sentinel (src_row == num_recv) to an OOB read of 0. (int32 element count: + # holds up to 2^31 elems; larger A needs a per-lane i64 rebase like Mega's fp8 path.) + a_base = fx.arith.ArithValue(arith.index_cast(fx.T.i64(), extract_base_index(A)), signed=True) + A_flat = fx.make_view(fx.inttoptr(pool_ptr_ty, a_base), fx.make_layout(A_ELEMS, 1)) + + def _emit(): + pid = xcd_remap_pid(fx.block_idx.x, real_grid, num_xcd) + num_pid_m = real_tiles + num_pid_in_group = GROUP_M * n_blocks + group_id = pid // num_pid_in_group + pid_in_group = pid % num_pid_in_group + first_pid_m = group_id * GROUP_M + remaining_m = num_pid_m - first_pid_m + group_size_m = arith.select(remaining_m < GROUP_M, remaining_m, fx.Int32(GROUP_M)) + block_m = first_pid_m + (pid_in_group % group_size_m) + block_n = pid_in_group // group_size_m + g_idx = buffer_load(group_res, block_m, vec_width=1, dtype=fx.T.i32()) + # PF padding blocks mark expert_ids=-1; skip the full Mega pipeline (v2 parity). + if g_idx >= fx.Int32(0): + gbase = g_idx * fx.Int32(K) * c_n + # Worst-case pool (cap*N > 2^31): rebase C per tile in int64, int32 in-resource + # offset. Mirrors the fused nt path. + c_byte_off = _i64(block_m * fx.Int32(BLOCK_M)) * _i64(c_n) * fx.Int64(2) + c_base = fx.arith.ArithValue(arith.index_cast(fx.T.i64(), extract_base_index(C)), signed=True) + C_tile = fx.make_view( + fx.inttoptr(pool_ptr_ty, c_base + c_byte_off), + fx.make_layout(fx.Int32(BLOCK_M) * c_n, 1), + ) + sorted_row_base = block_m * fx.Int32(BLOCK_M) + gemm_bf16_nt_gather_tile( + A_flat, + B, + C_tile, + fx.Int32(BLOCK_M), + c_n, + lds, + sorted_res, + sorted_row_base, + block_n, + K=K, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + out_fp16=out_fp16, + nt_vmcnt=nt_vmcnt, + b_group_base=gbase, + gather=gather, + ) + + if fx.block_idx.x < real_grid: + _emit() + + @flyc.jit + def launch(A, B, C, TILE_TO_GROUP, NUM_TILE_BLOCKS: fx.Int32, SORTED, A_ELEMS: fx.Int32, c_m: fx.Int32, c_n: fx.Int32, stream: fx.Stream): + grid_x = ceildiv(c_m, BLOCK_M) * ceildiv(c_n, BLOCK_N) + grouped_gemm_gather_k( + A, + B, + C, + TILE_TO_GROUP, + NUM_TILE_BLOCKS, + SORTED, + A_ELEMS, + c_m, + c_n, + value_attrs=make_value_attrs(waves_per_eu, agpr_alloc, "512,512"), + ).launch(grid=(grid_x, 1, 1), block=(512, 1, 1), stream=stream) + + return launch + + +def grouped_gemm_gather_bf16( + A, # [num_recv, K] bf16 received-token activations (UNPERMUTED gather source) + weight, # [E, N, K] bf16 per-expert B (NT) + output, # [em_max, N] bf16 compact expert-major C (in place) + expert_ids, # [num_m_blocks] i32 expert per BLOCK_M output block + num_tile_blocks, # int real BLOCK_M block count (host scalar; capture-safe) + sorted_slot_ids, # [em_max] i32 received-token row per padded slot (sentinel = num_recv) + *, + BLOCK_M=256, + BLOCK_N=256, + GROUP_M=4, + num_xcd=1, + nt_vmcnt=3, + waves_per_eu=2, + agpr_alloc=0, + gather=True, +): + """Host entry: grouped bf16 NT GEMM. With ``gather=True`` (FC1) ``C[pos] = A[SORTED[pos]] + @ B[expert]^T``; with ``gather=False`` (FC2 route-read) ``A`` is the compact ``[em_max, K]`` + pool read at the route row (``sorted_slot_ids`` unused, may be a dummy). + + ``output`` is written in place over its full padded ``[em_max, N]`` extent (padding rows carry + dead values, ignored by downstream stages keyed on the same routing metadata). ``c_m`` is the + padded slot count (``output.shape[0]``); the grid self-bounds to ``num_tile_blocks``. + """ + assert A.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16 + assert output.dtype == torch.bfloat16 + c_m = int(output.shape[0]) + E, N, K = weight.shape + assert K == A.shape[1], f"weight K={K} != A K={A.shape[1]}" + assert output.shape[1] == N, f"output N={output.shape[1]} != weight N={N}" + + A = A.contiguous() + a_elems = int(A.numel()) + if a_elems >= (1 << 31): + raise ValueError( + f"A has {a_elems} elems (>= 2^31); the flat-view gather bound is int32. " + "Large-A support needs a per-lane i64 SRD rebase (TODO)." + ) + weight_flat = weight.reshape(E * N, K).contiguous().view(-1) + expert_ids_i32 = expert_ids.to(torch.int32) + if gather: + sorted_arg = sorted_slot_ids.to(torch.int32) + else: + # FC2 route-read: the kernel synthesizes the identity index on-device, so SORTED is + # unread. Reuse a live tensor as the (unused) arg to avoid a capture-illegal allocation. + sorted_arg = expert_ids_i32 + launch = compile_grouped_gemm_gather_bf16( + K=K, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + GROUP_M=GROUP_M, + num_xcd=num_xcd, + nt_vmcnt=int(nt_vmcnt), + waves_per_eu=int(waves_per_eu), + agpr_alloc=int(agpr_alloc), + gather=bool(gather), + ) + launch( + A, + weight_flat, + output, + expert_ids_i32, + int(num_tile_blocks), + sorted_arg, + a_elems, + c_m, + N, + stream=torch.cuda.current_stream(), + ) + return output diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py new file mode 100644 index 000000000..487ff4ba3 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py @@ -0,0 +1,640 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Permute-free MoE weight-gradient (wgrad) grouped GEMM in FlyDSL. + +Contract: the gradient operand is the *compact* ``[num_routes, N]`` buffer +and ``SORTED`` maps each padded route slot to a received-token row:: + + dW[e][n, k] = sum_{routed slot s of e, valid} grad[route_start[e] + local(s), n] + * x[SORTED[s], k] + +where ``local(s)`` is the slot's offset within expert ``e``'s block-padded range and +padding slots (``SORTED[s] == num_recv_tokens``) are masked to zero. The grad walk is a +plain contiguous row scan (no ``SORTED`` indirection); ``SORTED`` is only consulted for +the ``x`` gather token and the padding mask. + +The contraction tile is staged through LDS and transposed on-read: + + 1. Coalesced fill: each 32-slot contraction step loads ``grad[slot, n_feat]`` and + ``x[token(slot), k_feat]`` into LDS as ``[slot(row), feature(col)]`` tiles with + wide vector loads along the contiguous feature axis. + 2. Hardware transpose-read: ``ds_read_tr16_b64`` reads the ``[slot, feature]`` tile + transposed into the MFMA A/B fragment layout ``[feature, slot]`` -- so the token + slot becomes the matrix-core contraction axis with no VGPR shuffle and no strided + global gather. + +A workgroup of ``warps_n x warps_k`` warps computes one ``block_n x block_k`` output +tile. All warps **cooperatively fill** the shared LDS contraction tile once per step +(amortizing the global gather), then each warp owns a ``(block_n/warps_n) x +(block_k/warps_k)`` sub-tile of 16x16 MFMA atoms, transpose-reading its own feature +columns out of the shared tile. ``warps_n = warps_k = 1`` reduces to the single-warp v2. + +Fixed configuration (no runtime toggles): + - bf16 inputs, bf16 output (FC1: compact grad + token-gathered ``x``) + - optional ``accumulate``: overwrite (default) or read-modify-write into ``dW`` + - DMA + XOR chunk swizzle fill, 3-stage LDS pipeline, LLVM DMA alias scopes +""" + +from __future__ import annotations + +import functools + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, const_expr, gpu, ptrtoint, range_constexpr, rocdl, vector +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch +from flydsl.utils.smem_allocator import SmemAllocator + +from ..tensor_shim import ptr_rsrc + +__all__ = ["compile_moe_wgrad_v2", "WGRAD_BLOCK_M"] + +# MFMA atom dims (16x16x32 bf16). +WMMA_M = 16 # output rows -> N (grad feature) +WMMA_N = 16 # output cols -> K (x feature) +WMMA_K = 32 # contraction -> token-slot tile +C_FRAG = 4 # f32 values per lane for the C (dW) fragment +WARP_SIZE = 64 + +# Coalesced fill vector width (bf16 elems). 8 bf16 = 16 B = one global_load_dwordx4. +FILL_V = 8 + +WGRAD_BLOCK_M = 32 # contraction (slot) step; matches the align block_size +NUM_BUF = 3 # triple-buffered DMA pipeline (pipe_stages == 3) + + +@functools.lru_cache(maxsize=None) +def compile_moe_wgrad_v2( + *, + block_n: int = 64, + block_k: int = 64, + warps_n: int = 1, + warps_k: int = 1, + accumulate: bool = False, +): + if block_n % (warps_n * WMMA_M) != 0 or block_k % (warps_k * WMMA_N) != 0: + raise ValueError("block_n/block_k must be multiples of warps_*16") + + n_threads = warps_n * warps_k * WARP_SIZE + if (WGRAD_BLOCK_M * block_n) % (n_threads * FILL_V) != 0: + raise ValueError("32*block_n must be a multiple of n_threads*FILL_V") + if (WGRAD_BLOCK_M * block_k) % (n_threads * FILL_V) != 0: + raise ValueError("32*block_k must be a multiple of n_threads*FILL_V") + + gpu_arch = get_rocm_arch() + WN = block_n // warps_n # per-warp grad-feature span + WK = block_k // warps_k # per-warp x-feature span + M_STEPS = WN // WMMA_M # grad-feature atoms per warp (MFMA-M) + N_STEPS = WK // WMMA_N # x-feature atoms per warp (MFMA-N) + NACC = M_STEPS * N_STEPS + + # DMA fill writes each lane's 16B contiguously into LDS (no per-row pad possible), + # so the swizzle path uses an un-padded stride and breaks bank conflicts with an + # XOR chunk-swizzle instead (see the fill + transpose-read below). + SG = block_n # grad LDS row stride (bf16 elems) + SX = block_k # x LDS row stride + # Swizzle granule = FILL_V bf16 (one 16B DMA unit). The XOR maps the feature chunk + # index with the slot so consecutive contraction slots land on distinct banks. + CPR_G_SWZ = block_n // FILL_V # feature chunks per grad row + CPR_X_SWZ = block_k // FILL_V # feature chunks per x row + + G_TILE_ELEMS = WGRAD_BLOCK_M * SG + X_TILE_ELEMS = WGRAD_BLOCK_M * SX + G_FILLS = (WGRAD_BLOCK_M * block_n) // (n_threads * FILL_V) + X_FILLS = (WGRAD_BLOCK_M * block_k) // (n_threads * FILL_V) + + KERNEL_NAME = ( + f"moe_wgrad_routelist_bf16_{block_n}x{block_k}_w{warps_n}x{warps_k}" + f"{'_acc' if accumulate else ''}_dsz_s3_dsa_v2" + ) + + # LDS allocation: NUM_BUF-buffered grad tile + x tile (2 bytes/bf16). + allocator = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem") + g_lds_off = allocator._align(allocator.ptr, 16) + allocator.ptr = g_lds_off + G_TILE_ELEMS * 2 * NUM_BUF + x_lds_off = allocator._align(allocator.ptr, 16) + allocator.ptr = x_lds_off + X_TILE_ELEMS * 2 * NUM_BUF + # The DMA path backs LDS with a raw llvm addrspace(3) global (buffer_load_lds needs a + # real global for M0; the memref allocator does not work). Same offsets/size. + LDS_TOTAL_BYTES = allocator._align(allocator.ptr, 128) + LDS_SYM = KERNEL_NAME + "_lds" + + @flyc.kernel(known_block_size=[n_threads, 1, 1]) + def wgrad_kernel( + dW: fx.Pointer, # [E, N, K] bf16 output + X: fx.Pointer, # [num_recv_tokens, K] bf16 (received-token activations) + GRAD: fx.Pointer, # [num_routes, N] bf16 (compact per-route gradient) + SORTED: fx.Pointer, # [padded] i32 received-token row per route slot (sentinel = num_recv_tokens) + BLOCK_START: fx.Pointer, # [E] i32 (block units) + BLOCKS_PER_EXPERT: fx.Pointer, # [E] i32 + ROUTE_START: fx.Pointer, # [E] i32 (compact first-route index = cumsum(counts) - counts) + N: fx.Int32, + K: fx.Int32, + num_recv_tokens: fx.Int32, + ): + bf16 = T.bf16 + c0 = arith.constant(0, index=True) + + dW_rsrc = ptr_rsrc(dW) + x_rsrc = ptr_rsrc(X) + grad_rsrc = ptr_rsrc(GRAD) + sorted_rsrc = ptr_rsrc(SORTED) + bstart_rsrc = ptr_rsrc(BLOCK_START) + bpe_rsrc = ptr_rsrc(BLOCKS_PER_EXPERT) + rstart_rsrc = ptr_rsrc(ROUTE_START) + + # DMA path: raw addrspace(3) global backs LDS; reads + DMA both GEP off it. + smem_raw_ptr = _llvm.mlir_addressof(ir.Type.parse("!llvm.ptr<3>"), LDS_SYM) + + tid = fx.Int32(gpu.thread_id("x")) + n_tile = fx.Int32(gpu.block_id("x")) # along N (grad feature) + k_tile = fx.Int32(gpu.block_id("y")) # along K (x feature) + expert = fx.Int32(gpu.block_id("z")) # expert id + + wid = tid // WARP_SIZE + lane = tid % WARP_SIZE + wn_id = wid // warps_k # warp row (grad feature) + wk_id = wid % warps_k # warp col (x feature) + lane_n = lane % 16 # MFMA C col (k_feat) + lane_m_base = lane // 16 # 0..3 + tr_k_group = (lane % 16) // 4 # 0..3 + tr_col_sub = lane % 4 # 0..3 + + warp_n_base = wn_id * fx.Int32(WN) # grad-feature col base of this warp + warp_k_base = wk_id * fx.Int32(WK) # x-feature col base of this warp + + # One alias scope per (operand, ring slot): the grad tile and the x tile each own + # NUM_BUF disjoint LDS byte ranges. A fill of slot r and a read of slot r share a + # scope so their RAW dependency survives; every other pair (different slot, or the + # other operand) is marked noalias, which is what lets SIInsertWaitcnts keep older + # in-flight fills streaming instead of draining vmcnt to 0 before each transpose read. + _ALIAS_DOMAIN = '#llvm.alias_scope_domain' + _SCOPE_IDS = tuple( + [f"g{r}" for r in range(NUM_BUF)] + [f"x{r}" for r in range(NUM_BUF)] + ) + + def _scope_attr(ids): + inner = ", ".join( + f'#llvm.alias_scope' + for sid in ids + ) + return ir.Attribute.parse(f"[{inner}]") + + _MY_SCOPE = {sid: _scope_attr((sid,)) for sid in _SCOPE_IDS} + _NOALIAS_SCOPE = { + sid: _scope_attr(tuple(o for o in _SCOPE_IDS if o != sid)) + for sid in _SCOPE_IDS + } + + def _g_sid(slot): + return f"g{slot % NUM_BUF}" + + def _x_sid(slot): + return f"x{slot % NUM_BUF}" + + def _scope_kw(sid): + """alias/noalias metadata kwargs for one (operand, ring slot), or {}.""" + if sid is None: + return {} + return { + "alias_scopes": _MY_SCOPE[sid], + "noalias_scopes": _NOALIAS_SCOPE[sid], + } + + N_idx = arith.index_cast(T.index, N) + K_idx = arith.index_cast(T.index, K) + nrecv_idx = arith.index_cast(T.index, num_recv_tokens) + + # DMA fill cannot mask padding slots in registers, so bound the token-gathered + # ``x`` operand's resource to its real [num_recv, K] extent: the sentinel token + # (== num_recv) and any pipeline overrun then read out-of-bounds -> hardware + # returns 0. A zeroed gathered column makes the padding slot's outer product + # ``grad (x) 0 == 0``, so the paired contiguous-walk grad operand may safely read + # garbage (clamped to row 0) for those slots. + x_addr_i64 = arith.index_cast(T.i64, ptrtoint(X)) + x_nrec_bytes = nrecv_idx * K_idx * arith.index(2) + x_rsrc = buffer_ops.create_buffer_resource_from_addr( + x_addr_i64, num_records_bytes=x_nrec_bytes + ) + + n_block_base = n_tile * block_n + k_block_base = k_tile * block_k + n_base_idx = arith.index_cast(T.index, n_block_base) + k_base_idx = arith.index_cast(T.index, k_block_base) + + # Per-expert routed-slot range. ``base_slot`` is the block-padded slot offset into + # ``SORTED`` (holds the received-token row for the ``x`` gather); ``route_start_e`` + # is the compact first-route index into the ``[num_routes, N]`` grad buffer. + bstart = buffer_load_i32(bstart_rsrc, expert) + nblocks = buffer_load_i32(bpe_rsrc, expert) + rstart = buffer_load_i32(rstart_rsrc, expert) + base_slot = arith.index_cast(T.index, bstart) * arith.index(WGRAD_BLOCK_M) + num_slots = arith.index_cast(T.index, nblocks) * arith.index(WGRAD_BLOCK_M) + route_start_e_idx = arith.index_cast(T.index, rstart) + + CPR_G = block_n // FILL_V + CPR_X = block_k // FILL_V + + def _tile_idx(i, cpr): + chunk = tid + fx.Int32(i * n_threads) + slot_i32 = chunk // fx.Int32(cpr) + feat_i32 = (chunk % fx.Int32(cpr)) * fx.Int32(FILL_V) + slot_idx = arith.index_cast(T.index, slot_i32) + feat_idx = arith.index_cast(T.index, feat_i32) + return slot_idx, feat_idx, slot_i32, feat_i32 + + # ``load_slot_ids`` issues the (indirect) ``sorted`` index loads. These are + # prefetched *two* steps ahead and carried across the loop as iter_args. + def load_slot_ids(s_base_idx): + g_ids = [ + buffer_load_i32_idx(sorted_rsrc, base_slot + s_base_idx + _tile_idx(i, CPR_G)[0]) + for i in range_constexpr(G_FILLS) + ] + x_ids = [ + buffer_load_i32_idx(sorted_rsrc, base_slot + s_base_idx + _tile_idx(i, CPR_X)[0]) + for i in range_constexpr(X_FILLS) + ] + return g_ids, x_ids + + def _dma_one(rsrc, ids, slot_base_idx, cpr, feat_base_idx, dim_idx, lds_off, + buf_byte, n_fills, clamp_row, sid=None): + # Issue the global->LDS DMA for one operand: each lane streams FILL_V bf16 + # from ``global[row, swizzled_feat]`` straight into contiguous LDS (no VGPR + # staging). The LDS destination base is wave-uniform (readfirstlane); the + # hardware spreads lane L to base + L*16B, reconstructing the contiguous + # ``phys_linear * FILL_V`` layout the swizzled transpose read expects. + for i in range_constexpr(n_fills): + phys = tid + fx.Int32(i * n_threads) + slot = phys // fx.Int32(cpr) + chunk = phys % fx.Int32(cpr) + slot_idx = arith.index_cast(T.index, slot) + token = arith.index_cast(T.index, ids[i]) + in_range = arith.cmpi( + arith.CmpIPredicate.ult, slot_base_idx + slot_idx, num_slots + ) + valid = arith.andi( + in_range, arith.cmpi(arith.CmpIPredicate.ult, token, nrecv_idx) + ) + if const_expr(clamp_row): + # grad: contiguous route walk; clamp overrun to row 0 for fault safety + # (its padding contribution is cancelled by the zeroed x column). + row_idx = valid.select(route_start_e_idx + slot_base_idx + slot_idx, c0) + else: + # x: gather by received-token; sentinel/OOB row -> hardware 0. + row_idx = token + swz = slot & fx.Int32(cpr - 1) + glob_feat = (chunk ^ swz) * fx.Int32(FILL_V) + glob_feat_idx = arith.index_cast(T.index, glob_feat) + voff_elem = row_idx * dim_idx + feat_base_idx + glob_feat_idx + voff_byte = arith.index_cast(T.i32, voff_elem * arith.index(2)) + lds_perlane = fx.Int32(lds_off) + buf_byte + phys * fx.Int32(FILL_V * 2) + lds_base = rocdl.readfirstlane(T.i32, lds_perlane) + rocdl.raw_ptr_buffer_load_lds( + rsrc, _gep_lds(smem_raw_ptr, lds_base), fx.Int32(FILL_V * 2), + voff_byte, fx.Int32(0), fx.Int32(0), fx.Int32(1), + **_scope_kw(sid), + ) + + def dma_fill(g_ids, x_ids, slot_base_idx, g_buf_byte, x_buf_byte, wbuf=None): + # FC1: token-gather ``x`` (clamp_row=False); compact grad route walk (clamp_row=True). + # ``wbuf`` is the Python ring-slot index this tile is being staged into (for alias scopes). + g_sid = _g_sid(wbuf) if wbuf is not None else None + x_sid = _x_sid(wbuf) if wbuf is not None else None + _dma_one( + grad_rsrc, g_ids, slot_base_idx, CPR_G_SWZ, n_base_idx, N_idx, + g_lds_off, g_buf_byte, G_FILLS, clamp_row=True, sid=g_sid, + ) + _dma_one( + x_rsrc, x_ids, slot_base_idx, CPR_X_SWZ, k_base_idx, K_idx, + x_lds_off, x_buf_byte, X_FILLS, clamp_row=False, sid=x_sid, + ) + + def _dma_barrier(keep=0): + # DMA lands on vmcnt (global load); drain it before the workgroup barrier so + # all waves observe the freshly-staged LDS tile. ``keep`` leaves that many + # vmem ops in flight (graduated wait) -- for the 3-stage pipeline this keeps the + # just-issued tile's DMA streaming across the barrier so it overlaps the next + # iteration's MFMA too (only the tile read next is fully drained). lgkmcnt(0) + # retires this wave's ds_reads before the buffer is recycled NUM_BUF steps on. + asm = f"s_waitcnt vmcnt({keep}) lgkmcnt(0)\ns_barrier" + _llvm.InlineAsmOp( + res=None, operands_=[], asm_string=asm, + constraints="", has_side_effects=True, is_align_stack=False, + ) + + def compute(accs, g_buf_byte, x_buf_byte, dma_prefetch=None, slot=None): + # A fragments are read one-per-mi to keep VGPR pressure low. B fragments are + # loop-invariant across mi, fetched once and interleaved with the mi==0 MFMAs + # so their ds_read latency overlaps compute. The MFMA region runs at raised + # priority so the matrix pipe stays fed while reads are in flight. + g_read_sid = _g_sid(slot) if slot is not None else None + x_read_sid = _x_sid(slot) if slot is not None else None + + def read_a(mi): + return _tr_read_frag_swz( + smem_raw_ptr, g_lds_off, SG, CPR_G_SWZ, warp_n_base, mi * WMMA_M, + lane_m_base, tr_k_group, tr_col_sub, g_buf_byte, + alias_kw=_scope_kw(g_read_sid), + ) + + def read_b(nj): + return _tr_read_frag_swz( + smem_raw_ptr, x_lds_off, SX, CPR_X_SWZ, warp_k_base, nj * WMMA_N, + lane_m_base, tr_k_group, tr_col_sub, x_buf_byte, + alias_kw=_scope_kw(x_read_sid), + ) + + new_accs = [None] * NACC + b_frags = [None] * N_STEPS + + if const_expr(dma_prefetch is not None): + # Burst *all* current-tile transpose reads first, then issue the next-tile + # DMA. Keeping the reads ahead of the DMA in program order stops the compiler + # from planting an ``s_waitcnt vmcnt(0)`` in front of the first ds_read; + # the DMA then streams under the MFMA burst instead. + a_frags = [read_a(mi) for mi in range_constexpr(M_STEPS)] + for nj in range_constexpr(N_STEPS): + b_frags[nj] = read_b(nj) + rocdl.sched_barrier(0) + dma_prefetch() + rocdl.sched_barrier(0) + rocdl.s_setprio(1) + for mi in range_constexpr(M_STEPS): + for nj in range_constexpr(N_STEPS): + idx = mi * N_STEPS + nj + new_accs[idx] = rocdl.mfma_f32_16x16x32_bf16( + T.vec(C_FRAG, T.f32), + [a_frags[mi], b_frags[nj], accs[idx], 0, 0, 0], + ) + rocdl.sched_barrier(0) + rocdl.s_setprio(0) + return new_accs + + a_next = read_a(0) + b_frags[0] = read_b(0) # first B needed for the very first MFMA + rocdl.s_setprio(1) + for mi in range_constexpr(M_STEPS): + a_cur = a_next + if const_expr(mi + 1 < M_STEPS): + a_next = read_a(mi + 1) # prefetch next A while MFMA-ing current + for nj in range_constexpr(N_STEPS): + # Prefetch the next B fragment during the first mi only; its ds_read + # then overlaps this step's MFMA and all later mi reuse the resident frag. + if const_expr(mi == 0 and nj + 1 < N_STEPS): + b_frags[nj + 1] = read_b(nj + 1) + idx = mi * N_STEPS + nj + new_accs[idx] = rocdl.mfma_f32_16x16x32_bf16( + T.vec(C_FRAG, T.f32), [a_cur, b_frags[nj], accs[idx], 0, 0, 0] + ) + rocdl.sched_barrier(0) + rocdl.s_setprio(0) + return new_accs + + acc_init = [ + arith.constant_vector(0.0, T.vec(C_FRAG, T.f32)) + for _ in range(NACC) + ] + acc_iter_args = NACC + + # NUM_BUF-stage ping-pong pipeline: MFMA the current LDS buffer while the next + # step's global gather is in flight. Prefetch distance D == NUM_BUF - 1: stage D + # tiles up front and keep the most-recently-issued tile's DMA in flight across the + # barrier (graduated ``vmcnt``) so it overlaps two MFMA steps instead of one. + D = NUM_BUF - 1 + PER_TILE_DMA = G_FILLS + X_FILLS + # Graduated ``vmcnt`` kept in flight across each sub-tile's barrier: leaves the + # just-issued tile's DMA streaming (only the tile read next is fully drained). + KEEP = (NUM_BUF - 2) * PER_TILE_DMA + acc_ty = T.vec(C_FRAG, T.f32) + + # Static NUM_BUF-buffer ring, distance-D, slot-loop unrolled by NUM_BUF. Sub-tile j + # always reads buffer j and prefetches tile (base+j+D) into buffer (j+D) % NUM_BUF -- + # both Python constants, so every LDS byte offset is compile-time constant and the + # backend can prove read(buf j) never aliases the in-flight DMA write(buf (j+D)). + def g_byte(j): + return fx.Int32((j % NUM_BUF) * G_TILE_ELEMS * 2) + + def x_byte(j): + return fx.Int32((j % NUM_BUF) * X_TILE_ELEMS * 2) + + ID_SET = G_FILLS + X_FILLS + + # Slot ids are carried NUM_BUF-ahead so their ~500-cycle SORTED-load latency is + # retired before ``dma_fill`` consumes them. + def pack_idsets(sets): + out = [] + for g_ids, x_ids in sets: + out += g_ids + x_ids + return out + + def unpack_idsets(args, base): + sets = [] + for s in range_constexpr(NUM_BUF): + off = base + s * ID_SET + g_ids = [args[off + i] for i in range(G_FILLS)] + x_ids = [args[off + G_FILLS + i] for i in range(X_FILLS)] + sets.append((g_ids, x_ids)) + return sets + + def mk_prefetch(g_ids, x_ids, pf_slot, wbuf): + return lambda: dma_fill( + g_ids, x_ids, pf_slot, g_byte(wbuf), x_byte(wbuf), wbuf=wbuf + ) + + # Prologue: stage tiles 0..D-1 into buffers 0..D-1 (distance-D), then drain fully. + for j in range_constexpr(D): + j_slot = arith.index(j * WMMA_K) + gidj, xidj = load_slot_ids(j_slot) + dma_fill(gidj, xidj, j_slot, g_byte(j), x_byte(j), wbuf=j) + _dma_barrier() + + # Initial carried id-sets: iteration 0 prefetches tiles D..D+NUM_BUF-1. + init_sets = [ + load_slot_ids(arith.index((D + j) * WMMA_K)) + for j in range_constexpr(NUM_BUF) + ] + + ntiles = fx.Int32(arith.index_cast(T.i32, num_slots)) // fx.Int32(WMMA_K) + group_end = arith.index_cast( + T.index, (ntiles // fx.Int32(NUM_BUF)) * fx.Int32(NUM_BUF * WMMA_K) + ) + step_big = arith.index(NUM_BUF * WMMA_K) + + loop = scf.ForOp( + c0, group_end, step_big, iter_args=acc_init + pack_idsets(init_sets) + ) + with ir.InsertionPoint(loop.body): + s_base = loop.induction_variable + accs = [loop.body.arguments[1 + i] for i in range(NACC)] + cur_sets = unpack_idsets(loop.body.arguments, 1 + acc_iter_args) + + # sub-tile j: read buffer j (tile s_base+j), prefetch tile s_base+j+D into + # buffer (j+D) % NUM_BUF using its carried id-set. + for j in range_constexpr(NUM_BUF): + g_ids_j, x_ids_j = cur_sets[j] + pf_slot = s_base + arith.index((j + D) * WMMA_K) + accs = compute( + accs, g_byte(j), x_byte(j), + dma_prefetch=mk_prefetch(g_ids_j, x_ids_j, pf_slot, j + D), + slot=j, + ) + _dma_barrier(KEEP) + + nxt_sets = [ + load_slot_ids(s_base + arith.index((NUM_BUF + D + j) * WMMA_K)) + for j in range_constexpr(NUM_BUF) + ] + scf.YieldOp(accs + pack_idsets(nxt_sets)) + + accs = [loop.results[i] for i in range(NACC)] + + # Tail: rem in 0..NUM_BUF-1 leftover tiles, already prefetched into buffers + # 0..rem-1 by the last group. Drain any in-flight DMA, then consume read-only. + _dma_barrier() + rem = ntiles - (ntiles // fx.Int32(NUM_BUF)) * fx.Int32(NUM_BUF) + + def _tail(j, accs): + if const_expr(j >= NUM_BUF - 1): + return accs + has_j = arith.cmpi(arith.CmpIPredicate.ugt, rem, fx.Int32(j)) + tif = scf.IfOp(has_j, results_=[acc_ty] * NACC, has_else=True) + with ir.InsertionPoint(tif.then_block): + a2 = compute(accs, g_byte(j), x_byte(j), dma_prefetch=None, slot=j) + scf.YieldOp(_tail(j + 1, a2)) + with ir.InsertionPoint(tif.else_block): + scf.YieldOp(accs) + return [tif.results[i] for i in range(NACC)] + + accs = _tail(0, accs) + + # Epilogue: C[m=n_feat, n=k_feat], lane holds 4 rows. Each dW element is owned by + # exactly one workgroup (grid = N x K x E over disjoint output tiles), so when + # ``accumulate`` is set the read-modify-write into the destination is race-free. + E_NK_row = arith.index_cast(T.index, expert) * N_idx * K_idx + for mi in range_constexpr(M_STEPS): + for nj in range_constexpr(N_STEPS): + acc_idx = mi * N_STEPS + nj + acc = accs[acc_idx] + c_n = k_block_base + warp_k_base + fx.Int32(nj * WMMA_N) + lane_n + c_n_idx = arith.index_cast(T.index, c_n) + k_ok = arith.cmpi(arith.CmpIPredicate.ult, c_n_idx, K_idx) + for ii in range_constexpr(C_FRAG): + n_out = ( + n_block_base + warp_n_base + fx.Int32(mi * WMMA_M) + + lane_m_base * C_FRAG + fx.Int32(ii) + ) + n_out_idx = arith.index_cast(T.index, n_out) + in_bounds = arith.andi( + arith.cmpi(arith.CmpIPredicate.ult, n_out_idx, N_idx), k_ok + ) + store_if = scf.IfOp(in_bounds, results_=[], has_else=False) + with ir.InsertionPoint(store_if.then_block): + val = vector.extract( + acc, static_position=[ii], dynamic_position=[] + ) # f32 accumulator + out_off = E_NK_row + n_out_idx * K_idx + c_n_idx + if const_expr(accumulate): + prev = buffer_ops.buffer_load( + dW_rsrc, out_off, vec_width=1, dtype=bf16 + ) + val = arith.addf(val, arith.extf(T.f32, prev)) + store_val = arith.truncf(bf16, val) + buffer_ops.buffer_store(store_val, dW_rsrc, out_off) + scf.YieldOp([]) + + @flyc.jit + def launch_wgrad( + dW: fx.Pointer, + X: fx.Pointer, + GRAD: fx.Pointer, + SORTED: fx.Pointer, + BLOCK_START: fx.Pointer, + BLOCKS_PER_EXPERT: fx.Pointer, + ROUTE_START: fx.Pointer, + N: fx.Int32, + K: fx.Int32, + num_recv_tokens: fx.Int32, + num_experts: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + _llvm.GlobalOp( + global_type=ir.Type.parse(f"!llvm.array<{LDS_TOTAL_BYTES} x i8>"), + sym_name=LDS_SYM, + linkage=ir.Attribute.parse("#llvm.linkage"), + addr_space=3, + alignment=1024, + ) + gx = (N + block_n - 1) // block_n + gy = (K + block_k - 1) // block_k + gz = num_experts + wgrad_kernel._func.__name__ = KERNEL_NAME + wgrad_kernel( + dW, X, GRAD, SORTED, BLOCK_START, BLOCKS_PER_EXPERT, ROUTE_START, + N, K, num_recv_tokens, + ).launch(grid=(gx, gy, gz), block=(n_threads, 1, 1), stream=stream) + + return launch_wgrad + + +def _gep_lds(base_ptr, byte_i32): + """GEP an LDS !llvm.ptr<3> by a runtime i8 byte offset off a real LDS base pointer. + + ``buffer_load_lds`` derives its M0 write base from the LDS pointer, which the backend + only lowers correctly when the pointer is a GEP off a genuine addrspace(3) global. + We route the reads through the same base so alias analysis ties the DMA writes to the + transpose reads. + """ + return _llvm.getelementptr( + ir.Type.parse("!llvm.ptr<3>"), rocdl._to_ir(base_ptr), [rocdl._to_ir(byte_i32)], + [-(2 ** 31)], T.i8, None, + ) + + +def _tr_read_frag_swz( + smem_base, lds_off, stride, cpr, warp_col_base, col_const, + lane_m_base, tr_k_group, tr_col_sub, buf_byte, alias_kw=None, +): + """Swizzled transpose-read for the DMA fill (un-padded LDS). + + The DMA writes each contraction tile contiguously (stride == feature span), so bank + conflicts on the transpose read are broken by an XOR *chunk* swizzle: the physical + feature chunk of logical ``(slot, feat)`` is ``(feat // FILL_V) XOR (slot & (cpr - 1))`` + -- the *same* map the fill applies to the global gather column, so the read lands on + exactly the element the DMA staged. + """ + col_run = warp_col_base + tr_col_sub * fx.Int32(4) + feat_log = col_run + fx.Int32(col_const) # logical feature column + chunk = feat_log // fx.Int32(FILL_V) + within = feat_log % fx.Int32(FILL_V) + row_lo = lane_m_base * fx.Int32(8) + tr_k_group # contraction slot (0..31) + + def _read(slot): + swz = slot & fx.Int32(cpr - 1) + phys_chunk = chunk ^ swz + phys_feat = phys_chunk * fx.Int32(FILL_V) + within + elem = slot * fx.Int32(stride) + phys_feat + byte = elem * fx.Int32(2) + fx.Int32(lds_off) + buf_byte + raw = rocdl.ds_read_tr16_b64( + T.vec(4, T.bf16), _gep_lds(smem_base, byte), **(alias_kw or {}) + ).result + return fx.Vector(raw, (4,), fx.BFloat16) + + lo = _read(row_lo) + hi = _read(row_lo + fx.Int32(4)) + return lo.shuffle(hi, [0, 1, 2, 3, 4, 5, 6, 7]) + + +def buffer_load_i32(rsrc, off_i32): + return buffer_ops.buffer_load(rsrc, off_i32, vec_width=1, dtype=T.i32) + + +def buffer_load_i32_idx(rsrc, off_idx): + return buffer_ops.buffer_load(rsrc, off_idx, vec_width=1, dtype=T.i32) diff --git a/transformer_engine/pytorch/flydsl_kernels/tensor_shim.py b/transformer_engine/pytorch/flydsl_kernels/tensor_shim.py new file mode 100644 index 000000000..ccea9cb4d --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/tensor_shim.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import torch +import flydsl.compiler as flyc +from flydsl.expr import arith, buffer_ops, ptrtoint +from flydsl.expr.typing import T + + +def ptr_rsrc(ptr): + """Convert an fx.Pointer kernel arg to a buffer resource for buffer_load/store.""" + addr_i64 = arith.index_cast(T.i64, ptrtoint(ptr)) + return buffer_ops.create_buffer_resource_from_addr(addr_i64) + + +def ptr_arg(t: torch.Tensor): + """Wrap a torch.Tensor as an fx.Pointer (PointerJitArg) for kernel launch.""" + import flydsl.expr as fx + + type_name = type(t).__name__ + module_name = type(t).__module__ + if type_name == "FakeTensor" or "fake_tensor" in module_name: + return flyc.from_c_void_p(fx.Uint8, 0) + return flyc.from_c_void_p(fx.Uint8, t.data_ptr()) + + +def _run_compiled(exe, *args): + """First call: ``flyc.compile(exe, *args)`` compiles **and** executes the kernel. + Subsequent calls: fast dispatch via the cached ``CompiledFunction``. + """ + cf = getattr(exe, "_cf", None) + if cf is None: + cf = flyc.compile(exe, *args) + exe._cf = cf + else: + cf(*args) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index f534da5c3..4cf5fecee 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -70,6 +70,11 @@ if IS_HIP_EXTENSION: from transformer_engine.pytorch.triton_kernels.grouped_gemm import general_grouped_gemm_triton + from transformer_engine.pytorch.moe import ( + is_permute_free_grouped_gemm_enabled, + permute_free_grouped_gemm_forward, + permute_free_grouped_gemm_backward, + ) import os __all__ = ["GroupedLinear"] @@ -405,6 +410,7 @@ def forward( ctx, inp: torch.Tensor, m_splits: torch.Tensor, + dispatched_probs: Optional[torch.Tensor], non_tensor_args: Tuple, *weights_and_biases, ) -> Tuple[torch.Tensor, list]: @@ -437,7 +443,12 @@ def forward( m_splits_tensor, actual_m_splits, unpad_output, + routing_metadata, + grouped_weight_param, ) = non_tensor_args + # The gated-activation fusion hint rides on the routing metadata (single source of + # truth); ``None`` on every non-permute-free / non-FC1 path. + perm_free_activation = getattr(routing_metadata, "activation", None) if fp8: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override else: @@ -448,11 +459,34 @@ def forward( # Check if Triton kernel should be used use_grouped_gemm_triton = IS_HIP_EXTENSION and os.getenv("NVTE_USE_GROUPED_GEMM_TRITON", "0") == "1" and not fp8 and not fuse_wgrad_accumulation + # Permute-free gather GEMM (fwd + dgrad gather, wgrad on default path). + # Training is supported: fp8 / bias are excluded. ``fuse_wgrad_accumulation`` is only + # supported together with a single grouped weight param, where the permute-free backward + # accumulates wgrad directly into the grouped param's ``main_grad``. + use_perm_free_grouped_gemm = ( + IS_HIP_EXTENSION + and is_permute_free_grouped_gemm_enabled() + and routing_metadata is not None + and not fp8 + and (not fuse_wgrad_accumulation or grouped_weight_param is not None) + and activation_dtype == torch.bfloat16 + and not use_bias + ) + if use_perm_free_grouped_gemm and use_grouped_gemm_triton: + raise RuntimeError( + "NVTE_PERMUTE_FREE_GROUPED_GEMM and NVTE_USE_GROUPED_GEMM_TRITON cannot both be enabled." + ) + num_gemms = len(m_splits) weights = weights_and_biases[:num_gemms] biases = weights_and_biases[num_gemms:] device = inp.device - weight_requires_grad = weights[0].requires_grad + # Grouped weights expose detached per-expert views; the grouped param carries requires_grad. + weight_requires_grad = ( + grouped_weight_param.requires_grad + if grouped_weight_param is not None + else weights[0].requires_grad + ) # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): @@ -503,7 +537,20 @@ def forward( # Initialize input tensors in_features = weights[0].size(-1) - if inp.size(-1) != in_features: + perm_free_route_space = ( + getattr(routing_metadata, "route_space", False) if routing_metadata is not None else False + ) + # FC2 with a fused gated prologue consumes FC1's raw 2F [gate|up] buffer (width 2F). + expect_in_features = ( + 2 * in_features + if ( + use_perm_free_grouped_gemm + and perm_free_route_space + and perm_free_activation is not None + ) + else in_features + ) + if inp.size(-1) != expect_in_features: raise ValueError( f"Input tensor (shape={tuple(inp.size())}) is not compatible with " f"weight tensor (shape={tuple(weights[0].size())})" @@ -565,6 +612,8 @@ def forward( ) elif use_grouped_gemm_triton: inputmats = [cast_if_needed(inp_view, activation_dtype)] + elif use_perm_free_grouped_gemm: + inputmats = [cast_if_needed(inp_view, activation_dtype)] else: inputmats = torch.split(cast_if_needed(inp_view, activation_dtype), m_splits) @@ -597,12 +646,15 @@ def forward( if fp8 and activation_dtype == torch.float32: bias_dtype = torch.bfloat16 # FP8 GEMM only supports BF16/FP16 bias biases = [cast_if_needed(bias, bias_dtype) for bias in biases] if use_bias else biases - # Initialize output tensor - out = torch.empty( - [sum(m_splits), weights_fp8[0].size(0)], - dtype=activation_dtype, - device=device, - ) + # Initialize output tensor. The permute-free path allocates its own worst-case padded + # [T * min(topk, E), out_features] output inside permute_free_grouped_gemm_bf16 (valid rows are + # the compact route range [0, num_routes); the tail is inert zero padding). + if not use_perm_free_grouped_gemm: + out = torch.empty( + [sum(m_splits), weights_fp8[0].size(0)], + dtype=activation_dtype, + device=device, + ) # Choose whether to use split accumulator use_split_accumulator = _2X_ACC_FPROP @@ -612,25 +664,39 @@ def forward( use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator # Perform GEMM - if use_grouped_gemm_triton: + if use_perm_free_grouped_gemm: + perm_free_route_space = getattr(routing_metadata, "route_space", False) + # FC1 emits raw 2F [gate|up]; the ``activation`` hint on the metadata is consumed on + # FC2, which applies the gated activation in a standalone pass and then runs a plain + # GEMM (the fused-prologue path regressed throughput). Route probs ride with FC2 too. + pf_result = permute_free_grouped_gemm_forward( + inputmats[0], + weights_fp8, + routing_metadata, + activation=perm_free_activation if perm_free_route_space else None, + dispatched_probs=dispatched_probs if perm_free_route_space else None, + ) + out = pf_result.out + elif use_grouped_gemm_triton: general_grouped_gemm_func = general_grouped_gemm_triton kwargs = {"m_splits_tensor": m_splits_tensor} else: general_grouped_gemm_func = general_grouped_gemm kwargs = {} - general_grouped_gemm_func( - weights_fp8, - inputmats, - [out], - output_quantizers, - activation_dtype, - single_output=True, - m_splits=m_splits, - bias=biases, - use_bias=use_bias, - use_split_accumulator=use_split_accumulator, - **kwargs, - ) + if not use_perm_free_grouped_gemm: + general_grouped_gemm_func( + weights_fp8, + inputmats, + [out], + output_quantizers, + activation_dtype, + single_output=True, + m_splits=m_splits, + bias=biases, + use_bias=use_bias, + use_split_accumulator=use_split_accumulator, + **kwargs, + ) output_unpadded = False @@ -681,21 +747,51 @@ def forward( if backward_override == "high_precision" and inp.requires_grad else [None] * num_gemms ) + # Permute-free FC2 backward needs the route probs when the forward fused them. tensors_to_save, tensor_objects = prepare_for_saving( *inputmats, *weights_fp8, *saved_weights, *biases, + dispatched_probs, ) ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects + ctx.perm_free_fc2_activation = ( + perm_free_activation + if ( + use_perm_free_grouped_gemm + and getattr(routing_metadata, "route_space", False) + ) + else None + ) ctx.grad_input_quantizers = grad_input_quantizers ctx.grad_output_quantizers = grad_output_quantizers ctx.grad_weight_quantizers = grad_weight_quantizers - ctx.weights_requires_grad = weights[0].requires_grad - if fuse_wgrad_accumulation and ctx.weights_requires_grad: + ctx.weights_requires_grad = weight_requires_grad + # Permute-free backward routes the whole [E, out, in] wgrad straight to the grouped + # param, rather than through the detached per-expert views (which carry no grad edge). + # With fuse_wgrad_accumulation it accumulates into ``main_grad``; otherwise it + # accumulates into the grouped param's autograd ``.grad`` (standard accumulation). + ctx.perm_free_grouped = ( + use_perm_free_grouped_gemm + and grouped_weight_param is not None + and ctx.weights_requires_grad + ) + ctx.grouped_fuse_wgrad = fuse_wgrad_accumulation + if ctx.perm_free_grouped: + ctx.grouped_weight_ref = weakref.ref(grouped_weight_param) + if ctx.grouped_fuse_wgrad: + ctx.grouped_overwrite_main_grad = getattr( + grouped_weight_param, "overwrite_main_grad", False + ) + if hasattr(grouped_weight_param, "__fsdp_param__"): + ctx.grouped_main_grad_func = grouped_weight_param.get_main_grad + else: + ctx.grouped_main_grad_func = lambda: grouped_weight_param.main_grad + elif fuse_wgrad_accumulation and ctx.weights_requires_grad: # Keep weakrefs to weights to preserve attributes like main_grad # when we need to modify the weight python objects ctx.origin_weight_refs = [weakref.ref(w) for w in weights] @@ -742,6 +838,8 @@ def forward( ctx.input_quantizers = input_quantizers ctx.use_grouped_gemm_triton = use_grouped_gemm_triton ctx.num_input_tensors = len(inputmats) + ctx.use_perm_free_grouped_gemm = use_perm_free_grouped_gemm + ctx.routing_metadata = routing_metadata if use_perm_free_grouped_gemm else None # backward overrides if backward_override is not None: @@ -756,7 +854,10 @@ def forward( ctx.grad_output_quantizers = [None] * num_gemms ctx.reduce_and_update_bwd_fp8_tensors = False - # [*, in_features] -> [*, out_features] except first dimension changes for SP + # [*, in_features] -> [*, out_features], or worst-case padded [T * min(topk, E), out_features] + # (permute-free route-list path; valid rows are the compact range [0, num_routes)). + if use_perm_free_grouped_gemm: + return out, new_workspaces return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @staticmethod @@ -961,6 +1062,98 @@ def backward( weights = saved_tensors[num_inputs: num_inputs + N] saved_weights = saved_tensors[num_inputs + N : num_inputs + 2 * N] biases = saved_tensors[num_inputs + 2 * N : num_inputs + 3 * N] + dispatched_probs = saved_tensors[num_inputs + 3 * N] + + # Permute-free gather GEMM: both dgrad and wgrad gather along the + # contraction axis inside a single Triton kernel. + if getattr(ctx, "use_perm_free_grouped_gemm", False): + # Single grouped weight: the GEMM only ever sees detached per-expert views, so + # the whole [E, out, in] wgrad goes straight to the grouped param (the positional + # autograd return is dead). Point the kernel at the destination accumulator up + # front so it folds the wgrad in directly -- no scratch dW and no separate + # add/copy. Both FC1 and FC2 (route_space, via swap_gather) emit [E, out, in] + # directly into the buffer. + grouped = getattr(ctx, "perm_free_grouped", False) and ctx.weights_requires_grad + wgrad_out = None + wgrad_accumulate = False + if grouped: + grouped_weight = ctx.grouped_weight_ref() + assert ( + grouped_weight is not None + ), "grouped weight was removed before its wgrad could be applied" + gw_shape = tuple(grouped_weight.shape) + if ctx.grouped_fuse_wgrad: + wgrad_out = ctx.grouped_main_grad_func().view(gw_shape) + wgrad_accumulate = not ctx.grouped_overwrite_main_grad + else: + # Accumulate into the grouped param's autograd ``.grad``; allocate it on + # the first backward (overwrite) and add in place thereafter so grad + # accumulation across microbatches still works. + if grouped_weight.grad is None: + grouped_weight.grad = torch.empty( + gw_shape, dtype=grouped_weight.dtype, device=grad_output.device + ) + wgrad_accumulate = False + else: + wgrad_accumulate = True + wgrad_out = grouped_weight.grad + + # The wrapper decides FC1 vs FC2 dgrad/wgrad from the routing metadata. + pf_result = permute_free_grouped_gemm_backward( + grad_output, + routing=ctx.routing_metadata, + weights=weights, + num_gemms=ctx.num_gemms, + hidden_states=inputmats[0], + requires_dgrad=ctx.requires_dgrad, + requires_wgrad=ctx.weights_requires_grad, + dispatched_probs=dispatched_probs, + fc2_activation=getattr(ctx, "perm_free_fc2_activation", None), + wgrad_out=wgrad_out, + wgrad_accumulate=wgrad_accumulate, + ) + if getattr(ctx, "perm_free_grouped", False) and pf_result.wgrad_stacked is not None: + grouped_weight = ctx.grouped_weight_ref() + assert ( + grouped_weight is not None + ), "grouped weight was removed before its wgrad could be applied" + if pf_result.wgrad_applied: + # FC1: the kernel already folded the wgrad into main_grad / .grad. + if ctx.grouped_fuse_wgrad and hasattr( + grouped_weight, "grad_added_to_main_grad" + ): + grouped_weight.grad_added_to_main_grad = True + else: + # FC2 (transpose) or no direct kernel: sink the returned stacked wgrad. + dW = pf_result.wgrad_stacked + if ctx.grouped_fuse_wgrad: + main_grad = ctx.grouped_main_grad_func().view(dW.shape) + if ctx.grouped_overwrite_main_grad: + main_grad.copy_(dW) + else: + main_grad.add_(dW) + if hasattr(grouped_weight, "grad_added_to_main_grad"): + grouped_weight.grad_added_to_main_grad = True + elif grouped_weight.grad is None: + grouped_weight.grad = dW + else: + grouped_weight.grad.add_(dW) + wgrad_list = [None] * ctx.num_gemms + else: + # Fall back to returning positional per-expert wgrad (separate leaf params). + # Split the [E, out, in] gradient into zero-copy per-expert views. + wgrad_list = ( + list(pf_result.wgrad_stacked) + if pf_result.wgrad_stacked is not None + else [None] * ctx.num_gemms + ) + return ( + pf_result.dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, + pf_result.grad_probs, + None, + *wgrad_list, + *([None] * ctx.num_gemms), + ) # Restore from weakrefs to get original weight python objects # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) @@ -1270,6 +1463,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): return ( dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, None, # m_splits + None, # dispatched_probs None, # non_tensor_args *wgrad_list, *grad_biases, @@ -1346,6 +1540,19 @@ class GroupedLinear(TransformerEngineBaseModule): GroupedLinear doesn't really handle the TP communications inside. The ``tp_size`` and ``parallel_mode`` are used to determine the shapes of weights and biases. The TP communication should be handled in the dispatch and combine stages of MoE models. + + Permute-free MoE (ROCm, bf16) + ----------------------------- + When ``NVTE_PERMUTE_FREE_GROUPED_GEMM=1``, pass a ``permute_free_metadata`` + (:class:`PermuteFreeMetadata`, carrying the boolean ``routing_map`` + ``[num_recv_tokens, num_local_experts]`` + a ``route_space`` direction) instead of + permuting activations before this module. The caller must skip ``moe_permute``. FC1 + (``route_space=False``) takes ``[num_recv_tokens, in_features]`` and produces the + worst-case padded ``[T * min(topk, E), out_features]`` route buffer (valid rows are the compact + route range ``[0, num_routes)``; the tail is inert zero padding); FC2 + (``route_space=True``) takes the route-ordered ``[T * min(topk, E), in_features]`` and fuses the + scatter back to token order, returning ``[num_recv_tokens, out_features]``. Requires + ``bias=False`` and bf16. The router-weight combine happens upstream (at the activation). """ def __init__( @@ -1767,6 +1974,8 @@ def forward( m_splits_tensor: Optional[torch.Tensor] = None, actual_m_splits: Optional[List[int]] = None, unpad_output: bool = False, + permute_free_metadata: Optional["PermuteFreeMetadata"] = None, + dispatched_probs: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: """ Apply the linear transformation to the input. @@ -1798,6 +2007,30 @@ def forward( When True, unpad the GEMM output from sum(m_splits) to sum(actual_m_splits) rows before returning. Used by the ROCm fused-pad-cast-transpose path; ignored on CUDA. + permute_free_metadata : PermuteFreeMetadata, optional + Route-list routing metadata (carrying the boolean ``routing_map`` + ``[num_recv_tokens, num_local_experts]`` plus a ``route_space`` + direction). When set with ``NVTE_PERMUTE_FREE_GROUPED_GEMM=1``, run + the route-list GEMM on unpermuted bf16 activations. ``route_space= + False`` (FC1) gathers per expert into the worst-case padded + ``[T * min(topk, E), out_features]`` route buffer (valid rows are the compact + range ``[0, num_routes)``; the tail is inert zero padding); + ``route_space=True`` (FC2) reads route-ordered input and fuses the + scatter back to token order, returning ``[num_recv_tokens, + out_features]``. TE builds/caches the expert-sorted alignment buffers + on the metadata. The gated-activation fusion hint + (``permute_free_metadata.activation`` = ``"silu"`` / ``"gelu"``) rides on + this object: when set on the FC1 direction (``route_space=False``) it + fuses the **gated** activation into the GEMM epilogue (weight output dim + is the gate+up width ``2F``, laid out as ``[gate | up]``; the returned + buffer is the ``F``-wide ``act(gate) * up`` and the separate activation + pass is skipped). Ignored on other paths. + dispatched_probs : torch.Tensor, optional + ``[num_recv_tokens, num_local_experts]`` gating probabilities. When given + with a fused ``activation``, each route's ``prob[token, expert]`` is + multiplied into the activation in-kernel (skipping the separate route-prob + pass); its gradient is returned to the router through autograd. Must be a + leaf/differentiable tensor for training. """ debug = self.is_debug_iter() is_grad_enabled = torch.is_grad_enabled() @@ -1836,6 +2069,12 @@ def forward( # Preprocess input tensor if isinstance(inp, QuantizedTensorStorage): raise TypeError("GroupedLinear doesn't support input tensor in FP8.") + + # The permute-free path is driven entirely by the PermuteFreeMetadata (which carries + # the boolean routing_map + the route_space direction). It is built once by the + # caller and shared across FC1/FC2 to avoid a duplicate align build; TE builds/caches + # the align buffers on the object. + routing_metadata = permute_free_metadata inp = self.prepare_forward(inp, num_gemms=self.num_gemms) try: @@ -1897,9 +2136,13 @@ def forward( m_splits_tensor, actual_m_splits, unpad_output, + routing_metadata, + # Grouped weight param (single_grouped_weight): permute-free backward accumulates + # wgrad into its ``main_grad`` instead of returning it through the detached views. + getattr(self, "weight", None) if self.single_grouped_weight else None, ) out, new_workspaces = linear_fn( - *autograd_ctx, inp, m_splits, non_tensor_args, *weight_tensors, *bias_tensors + *autograd_ctx, inp, m_splits, dispatched_probs, non_tensor_args, *weight_tensors, *bias_tensors ) if cache_weight: diff --git a/transformer_engine/pytorch/moe/__init__.py b/transformer_engine/pytorch/moe/__init__.py new file mode 100644 index 000000000..fee15bb56 --- /dev/null +++ b/transformer_engine/pytorch/moe/__init__.py @@ -0,0 +1,45 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Permute-free MoE grouped GEMM (FlyDSL) integration for PyTorch.""" + +from .permute_free_grouped_gemm import ( + MoERoutingMetadata, + PermuteFreeBackwardResult, + PermuteFreeForwardResult, + PermuteFreeMetadata, + get_default_moe_kernel_config, + is_permute_free_grouped_gemm_enabled, + permute_free_grouped_gemm_backward, + permute_free_grouped_gemm_bf16, + permute_free_grouped_gemm_bf16_dgrad, + permute_free_grouped_gemm_bf16_fc2, + permute_free_grouped_gemm_bf16_fc2_dgrad, + permute_free_grouped_gemm_bf16_fc2_wgrad, + permute_free_grouped_gemm_bf16_wgrad, + permute_free_grouped_gemm_forward, + permute_free_gated_act_bwd, + permute_free_gated_act_recompute, + prepare_moe_align, +) + +__all__ = [ + "MoERoutingMetadata", + "PermuteFreeBackwardResult", + "PermuteFreeForwardResult", + "PermuteFreeMetadata", + "get_default_moe_kernel_config", + "is_permute_free_grouped_gemm_enabled", + "permute_free_grouped_gemm_backward", + "permute_free_grouped_gemm_bf16", + "permute_free_grouped_gemm_bf16_dgrad", + "permute_free_grouped_gemm_bf16_fc2", + "permute_free_grouped_gemm_bf16_fc2_dgrad", + "permute_free_grouped_gemm_bf16_fc2_wgrad", + "permute_free_grouped_gemm_bf16_wgrad", + "permute_free_grouped_gemm_forward", + "permute_free_gated_act_bwd", + "permute_free_gated_act_recompute", + "prepare_moe_align", +] diff --git a/transformer_engine/pytorch/moe/moe_routing.py b/transformer_engine/pytorch/moe/moe_routing.py new file mode 100644 index 000000000..62bc19cc3 --- /dev/null +++ b/transformer_engine/pytorch/moe/moe_routing.py @@ -0,0 +1,138 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""MoE routing metadata for permute-free grouped GEMM.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch + + +@dataclass +class MoERoutingMetadata: + """Routing tensors for the route-list gather-in-GEMM MoE path. + + Parameters + ---------- + routing_map: + Boolean mask ``[num_recv_tokens, num_local_experts]``, True where a received token + feeds a local expert. ``num_routes = routing_map.sum()``. + num_experts: + Local expert count on this rank. Optional -- defaults to ``routing_map.size(1)``. + topk: + Upstream router top-k -- the host-known maximum number of local experts any received + token can feed. Optional. When provided, ``prepare_moe_align`` tightens the (still + sync-free) static over-allocation of the block-padded route buffers from the dense + ``num_recv_tokens * num_experts`` bound down to + ``num_recv_tokens * min(topk, num_experts)``, shrinking ``em_max`` (and the zero-init + it costs) by up to ``num_experts / topk``. Leave ``None`` to keep the dense bound. + + The remaining fields are lazily populated by ``prepare_moe_align`` and cached for + FC1 fwd/dgrad reuse (all in the expert-sorted, block-padded route-list layout): + + The align buffers are built sync-free and over-allocated to static, shape-derived + upper bounds; the real extents are carried as device scalars so no host sync is + needed to construct them. + + sorted_slot_ids: + ``[T * min(topk, E)]`` received-token row to gather for each block-padded position (sentinel + ``num_recv_tokens`` for padding and for the over-allocated tail). + expert_ids: + ``[blocks_max]`` local expert owning each ``BLOCK_SIZE_M`` block (``-1`` past the + real block count; those blocks are never visited). + num_tokens_post_padded: + ``[1]`` device scalar = real ``em`` (block-padded route count). Bounds the kernel. + block_start: + ``[num_experts]`` per-expert first block index (block units). + route_start: + ``[num_experts]`` per-expert first *compact* route index (``cumsum(counts) - counts``). + Maps a block-padded position to its compact output row. + route_to_token: + ``[routes_max]`` received-token row for each compact route (first ``num_routes`` + entries valid); used by the dgrad scatter-add back to ``[num_recv_tokens, K]``. + token_routes / token_route_count: + Inverse of ``route_to_token`` (token -> its compact route positions), built sync-free + for the contention-free gather-combine that replaces the atomic scatter in the token + combine (FC2 fwd) and the FC1 input-grad reduction (FC1 dgrad). ``token_routes`` is + ``[num_recv_tokens, min(topk, num_experts)]`` (int32); for token ``t`` the first + ``token_route_count[t]`` entries are its route positions (expert-ascending), the rest + are unused padding. + block_size_m: + ``BLOCK_SIZE_M`` used to build the fwd/dgrad align buffers. + wgrad_*: + Separate block-``CONTRACT_M`` align buffers for the route-list wgrad kernel. + route_counts / route_within: + Cached block-size-independent scan (per-expert counts and within-expert ranks) + shared by the fwd/dgrad and wgrad align builds. + """ + + routing_map: torch.Tensor + num_experts: Optional[int] = None + topk: Optional[int] = None + sorted_slot_ids: Optional[torch.Tensor] = None + expert_ids: Optional[torch.Tensor] = None + num_tokens_post_padded: Optional[torch.Tensor] = None + block_start: Optional[torch.Tensor] = None + route_start: Optional[torch.Tensor] = None + route_to_token: Optional[torch.Tensor] = None + token_routes: Optional[torch.Tensor] = None + token_route_count: Optional[torch.Tensor] = None + block_size_m: Optional[int] = None + wgrad_sorted_slot_ids: Optional[torch.Tensor] = None + wgrad_block_start: Optional[torch.Tensor] = None + wgrad_blocks_per_expert: Optional[torch.Tensor] = None + wgrad_block_size: Optional[int] = None + # Block-size-independent scan (per-expert counts + within-expert ranks), shared by the + # fwd/dgrad and wgrad align builds so it is computed once per routing map. + route_counts: Optional[torch.Tensor] = None + route_within: Optional[torch.Tensor] = None + + def __post_init__(self): + # num_experts is redundant with the routing map width (one column per local + # expert), so the caller can pass just ``routing_map``. + if self.num_experts is None: + self.num_experts = int(self.routing_map.size(1)) + + @property + def num_recv_tokens(self) -> int: + """Number of received tokens (rows of ``routing_map`` / the activation buffer).""" + return int(self.routing_map.size(0)) + + +@dataclass +class PermuteFreeMetadata(MoERoutingMetadata): + """Routing metadata for the permute-free grouped GEMM, tagged with a direction. + + Extends :class:`MoERoutingMetadata` + + - ``route_space=False`` (FC1): the input lives in **received-token order** + ``[num_recv_tokens, in]``. The forward *gathers* per expert (``index_a_by_route_pos= + False``) into the compact/padded ``[T * min(topk, E), out]`` route buffer; the dgrad combines + the input gradient back to token rows (contention-free gather-combine). + - ``route_space=True`` (FC2): the input is already in **route order** + ``[T * min(topk, E), in]`` (FC1's output). The forward reads by route position + (``index_a_by_route_pos=True``) and combines each token's routes back to + ``[num_recv_tokens, out]`` (contention-free gather-combine); the dgrad gathers + the token-space grad back into the compact route buffer. + + The align buffers are identical for both directions, so a single built metadata can be + reused for FC1 and FC2 (e.g. via ``dataclasses.replace(meta, route_space=True)``), + avoiding a duplicate align build. + + Fusion hint (optional): + + activation: + Gated activation to fuse into the FC2 GEMM prologue -- ``"silu"`` or ``"gelu"``. + ``None`` leaves the activation to the caller (no fusion). FC1 emits raw ``2F``; + this hint is consumed on the FC2 direction (``route_space=True``). + + (The per-route gating probabilities are *not* carried here: they need a gradient, so they + are passed as a separate autograd tensor argument to the module rather than as metadata.) + """ + + route_space: bool = False + activation: Optional[str] = None \ No newline at end of file diff --git a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py new file mode 100644 index 000000000..199825e76 --- /dev/null +++ b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py @@ -0,0 +1,1258 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Permute-free route-list grouped GEMM for MoE (bf16). + +- TE builds one expert-sorted ``sorted_slot_ids`` (received-token row per route) plus a + compact-output map (``route_start``/``block_start``), then runs the gather-GEMM. +- FC1 fwd output is worst-case padded ``[T * min(topk, E), out_features]`` in expert order (valid rows + are the compact route range ``[0, num_routes)``, tail is inert zero padding); dgrad returns + ``dA = [num_recv_tokens, in_features]`` (scatter-add of the per-route gradients). +""" + +from __future__ import annotations + +import os +import warnings +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple + +import torch + +from .moe_routing import MoERoutingMetadata, PermuteFreeMetadata +from .pf_helper_kernels import ( + fused_gated_act_prob_bwd, + fused_gated_act_prob_fwd, + route_gather_combine, + route_list_align, + route_list_scan, +) + +__all__ = [ + "MoERoutingMetadata", + "PermuteFreeMetadata", + "permute_free_grouped_gemm_bf16", + "permute_free_grouped_gemm_bf16_dgrad", + "permute_free_grouped_gemm_bf16_wgrad", + "permute_free_grouped_gemm_forward", + "permute_free_grouped_gemm_backward", + "PermuteFreeForwardResult", + "PermuteFreeBackwardResult", + "prepare_moe_align", + "get_default_moe_kernel_config", + "is_permute_free_grouped_gemm_enabled", +] + +_WGRAD_CONTRACT_M = 32 + +# Max forward align/kernel ``block_m`` for fwd/dgrad FlyDSL Permute-free Grouped GEMM. +_FLYDSL_GATED_BLOCK_M = 256 +_FLYDSL_FWD_BLOCK_M = 256 +_FLYDSL_FWD_LARGE_TIER = 128 + + +def _get_flydsl_fwd(): + """Return ``(flydsl_moe_fwd_autotuned, flydsl_moe_fwd_supported)``.""" + from .pf_fwd_wrapper import flydsl_moe_fwd_autotuned, flydsl_moe_fwd_supported + + return flydsl_moe_fwd_autotuned, flydsl_moe_fwd_supported + + +def _expert_per_route(routing: MoERoutingMetadata, routes_max: int) -> torch.Tensor: + """Per-route local expert id ``[routes_max]`` from the route-list block metadata.""" + if ( + routing.expert_ids is None + or routing.block_size_m is None + or routing.block_size_m <= 0 + ): + raise ValueError("_expert_per_route requires prepared routing align buffers.") + block_m = int(routing.block_size_m) + pos = torch.arange(routes_max, device=routing.expert_ids.device, dtype=torch.int64) + block_idx = (pos // block_m).clamp(max=routing.expert_ids.numel() - 1) + return routing.expert_ids[block_idx.to(torch.int64)].to(torch.int32) + +def _env_int(name: str, default: int) -> int: + v = os.environ.get(name) + try: + return int(v) if v is not None and v.strip() != "" else default + except ValueError: + return default + + +def _get_flydsl_wgrad(): + """Return the FC1 FlyDSL wgrad launcher.""" + from .pf_wgrad_wrapper import flydsl_moe_wgrad_autotuned + + return flydsl_moe_wgrad_autotuned + + +def _pf_moe_fwd( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + routing: MoERoutingMetadata, + *, + num_recv_tokens: int, + config: Dict[str, Any], + index_a_by_route_pos: bool = False, + activation: Optional[str] = None, + dispatched_probs: Optional[torch.Tensor] = None, + preact_out: Optional[torch.Tensor] = None, + gated_a: bool = False, +) -> None: + """Run the route-list gather-GEMM (forward or dgrad) via FlyDSL.""" + autotuned, supported = _get_flydsl_fwd() + block_m = int(config["BLOCK_SIZE_M"]) + if not supported(A, B, block_m=block_m): + raise RuntimeError( + "FlyDSL grouped GEMM does not support these operands " + f"(block_m={block_m}, A={tuple(A.shape)}, B={tuple(B.shape)})." + ) + autotuned( + A, + B, + C, + routing.sorted_slot_ids, + routing.expert_ids, + routing.block_start, + routing.route_start, + num_recv_tokens=num_recv_tokens, + block_m=block_m, + index_a_by_route_pos=index_a_by_route_pos, + activation=activation, + dispatched_probs=dispatched_probs, + preact_out=preact_out, + gated_a=gated_a, + ) + + +def _fwd_align_block_size_m( + A: torch.Tensor, + B: torch.Tensor, + *, + gated: bool, + default_block_m: int, + num_tokens: int, + gated_a: bool = False, +) -> int: + """Pick the forward align/kernel ``block_m``, favoring the backend that will run. + + ``block_m`` is baked into ``prepare_moe_align`` (it sets the row padding) and is *shared* by a + layer's FC1 fwd, FC1 dgrad and FC2 fwd (they reuse ``routing.block_size_m``), so it is chosen + once here. When the FlyDSL backend will run and the workload is on the large-token tier + (``num_tokens >= _FLYDSL_FWD_LARGE_TIER``), prefer :data:`_FLYDSL_FWD_BLOCK_M` (256): it lifts + the non-gated FC1 fwd, FC1 dgrad, and ``gated_a`` FC2 onto the faster ``256x32`` tile. The + gated FC1 epilogue stages a ``2F`` [gate|up] B-tile that overflows LDS at 256, so it drops back + to :data:`_FLYDSL_GATED_BLOCK_M` (128). Small/medium tiers keep the token-count default. + """ + if gated: + # The gated FC1 epilogue stages a 2F [gate|up] B-tile: even the LDS-valid 256 tiles are + # the slow tiny ones, so the measured-best policy is a hard cap at _FLYDSL_GATED_BLOCK_M + # (~1.35x at 128 vs ~1.06x at 256). Never bump; only allow the picker to confirm/keep a + # value <= 128 (dropping the default when it exceeds the cap). + cap = min(default_block_m, _FLYDSL_GATED_BLOCK_M) + candidates = {c for c in (_FLYDSL_GATED_BLOCK_M, default_block_m) if c <= cap} + else: + # Non-gated FC1 / dgrad / gated_a FC2: offer the default and any smaller floor (so the + # picker can drop to a smaller LDS-valid block_m), plus the 256 bump on the large-token + # tier, where it is a measured win and the extra E_local*128 align padding is a negligible + # fraction of the routed work. + candidates = {c for c in (_FLYDSL_GATED_BLOCK_M, default_block_m) if c <= default_block_m} + if num_tokens >= _FLYDSL_FWD_LARGE_TIER: + candidates.add(_FLYDSL_FWD_BLOCK_M) + + # Prefer the in-tree FlyDSL block_m picker when available. + try: + from .pf_fwd_wrapper import flydsl_moe_fwd_pick_block_m + except Exception: # pylint: disable=broad-except + flydsl_moe_fwd_pick_block_m = None + if flydsl_moe_fwd_pick_block_m is not None: + picked = flydsl_moe_fwd_pick_block_m( + A, B, gated=gated, gated_a=gated_a, candidates=tuple(candidates) + ) + if picked is not None: + return picked + + # Legacy fallback: cap the gated fwd to 128, else default. + if not gated or default_block_m <= _FLYDSL_GATED_BLOCK_M: + return default_block_m + _, supported = _get_flydsl_fwd() + if supported(A, B, block_m=_FLYDSL_GATED_BLOCK_M): + return _FLYDSL_GATED_BLOCK_M + return default_block_m + + +def is_permute_free_grouped_gemm_enabled() -> bool: + from torch.utils.cpp_extension import IS_HIP_EXTENSION + + return IS_HIP_EXTENSION and os.getenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "0") == "1" + + +def get_default_moe_kernel_config(num_tokens: int) -> Dict[str, Any]: + """Return align ``block_m`` tile config for the route-list bf16 MoE path.""" + # aiter's tuned config is an optional accelerator. + try: + from aiter.ops.triton.utils.moe_config_utils import get_optimal_moe_config + + return get_optimal_moe_config(torch.bfloat16, M=num_tokens) + except Exception: # pylint: disable=broad-except + pass + + if num_tokens <= 32: + block_m = 16 + elif num_tokens <= 96: + block_m = 32 + elif num_tokens <= 512: + block_m = 64 + else: + block_m = 128 + + block_n = 64 if num_tokens <= 64 else 128 + block_k = 128 if num_tokens <= 64 else 64 + group_m = 16 if num_tokens // max(block_m, 1) > 128 else 1 + + return { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": group_m, + "num_warps": 4 if num_tokens <= 128 else 8, + "num_stages": 2, + } + + +def moe_align_route_list( + routing_map: torch.Tensor, + *, + num_experts: int, + block_size: int, + scan=None, + topk: Optional[int] = None, + build_inverse_map: bool = False, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + Optional[torch.Tensor], + Optional[torch.Tensor], +]: + """Build the expert-sorted route-list buffers from a boolean routing map. + + Parameters + ---------- + routing_map: + Boolean ``[num_recv_tokens, num_experts]``; True where a received token feeds a + local expert. + num_experts: + Local expert count. + block_size: + ``BLOCK_SIZE_M`` (fwd/dgrad) or ``CONTRACT_M`` (wgrad) the layout is padded to. + scan: + Optional ``(counts, within)`` from :func:`route_list_scan`, shared across block + sizes so the block-independent scan is only paid once per routing map. + topk: + Host-known upper bound (the router top-k) on the number of experts any token routes + to. When provided, tightens the static over-allocation from ``T * num_experts`` to + ``T * min(topk, num_experts)`` (still sync-free). + + Returns + ------- + ``(sorted_slot_ids, expert_ids, num_tokens_post_padded, block_start, blocks_per_expert, + route_start, route_to_token, token_routes, token_route_count)``. Index tensors are + ``int32``; ``num_tokens_post_padded`` (``[1]``) is a device scalar. The last two are the + token->routes inverse map, ``None`` unless ``build_inverse_map``. + """ + return route_list_align( + routing_map, + num_experts=num_experts, + block_size=block_size, + scan=scan, + topk=topk, + build_inverse_map=build_inverse_map, + ) + + +def _ensure_route_scan(metadata: MoERoutingMetadata): + """Compute + cache the block-independent scan (counts, within) once per routing map.""" + if metadata.route_counts is None or metadata.route_within is None: + metadata.route_counts, metadata.route_within = route_list_scan( + metadata.routing_map, num_experts=metadata.num_experts + ) + return metadata.route_counts, metadata.route_within + + +def prepare_moe_align(metadata: MoERoutingMetadata, block_m: int) -> MoERoutingMetadata: + """Build and cache the fwd/dgrad route-list align buffers on ``metadata`` (sync-free).""" + if ( + metadata.sorted_slot_ids is not None + and metadata.block_size_m == block_m + and metadata.token_routes is not None + ): + return metadata + + counts, within = _ensure_route_scan(metadata) + + ( + sorted_slot_ids, + expert_ids, + num_tokens_post_padded, + block_start, + _blocks_per_expert, + route_start, + route_to_token, + token_routes, + token_route_count, + ) = moe_align_route_list( + metadata.routing_map, + num_experts=metadata.num_experts, + block_size=block_m, + scan=(counts, within), + topk=metadata.topk, + build_inverse_map=True, + ) + metadata.sorted_slot_ids = sorted_slot_ids # [T * min(topk, E)] int32: block-padded slot -> token + metadata.expert_ids = expert_ids # [blocks_max] int32: expert owning each block (-1 past end) + metadata.num_tokens_post_padded = num_tokens_post_padded # [1] int32 device scalar: padded extent + metadata.block_start = block_start # [E] int32: first block index of each expert (block units) + metadata.route_start = route_start # [E] int32: first compact route index of each expert + metadata.route_to_token = route_to_token # [routes_max] int32: compact route -> token + metadata.block_size_m = block_m # int: BLOCK_SIZE_M the layout is padded to + metadata.token_routes = token_routes # [T, min(topk, E)] int32: token -> its compact route ids + metadata.token_route_count = token_route_count # [T] int32: number of routes per token + return metadata + + +def _prepare_wgrad_align( + metadata: MoERoutingMetadata, contract_m: int +) -> MoERoutingMetadata: + """Build and cache the block-``contract_m`` align buffers for the wgrad kernel.""" + if ( + metadata.wgrad_sorted_slot_ids is not None + and metadata.wgrad_block_size == contract_m + ): + return metadata + + ( + sorted_slot_ids, + _expert_ids, + _num_tokens_post_padded, + block_start, + blocks_per_expert, + route_start, + _route_to_token, + _token_routes, + _token_route_count, + ) = moe_align_route_list( + metadata.routing_map, + num_experts=metadata.num_experts, + block_size=contract_m, + scan=_ensure_route_scan(metadata), + topk=metadata.topk, + ) + metadata.wgrad_sorted_slot_ids = sorted_slot_ids + metadata.wgrad_block_start = block_start + metadata.wgrad_blocks_per_expert = blocks_per_expert + metadata.wgrad_block_size = contract_m + # route_start is block-size-independent; keep whichever is already cached. + if metadata.route_start is None: + metadata.route_start = route_start + return metadata + + +def _try_view_grouped(weights: list[torch.Tensor]) -> Optional[torch.Tensor]: + """Return a zero-copy ``[E, out, in]`` view when the per-expert tensors are contiguous, + uniformly-strided, sequential slices of **one shared storage** (e.g. a single grouped weight + buffer). + + Returns ``None`` when the layout is not a single contiguous block, so the caller can fall + back to ``torch.stack``. Note that sequential ``data_ptr``s are not sufficient: separate + parameters can be allocated back-to-back yet own distinct storages, so ``as_strided`` from + the first tensor would overrun its storage. We therefore require a single shared storage. + """ + w0 = weights[0] + if not w0.is_contiguous(): + return None + out, in_ = w0.shape + stride = out * in_ + itemsize = w0.element_size() + base_ptr = w0.data_ptr() + storage_ptr = w0.untyped_storage().data_ptr() + for i, w in enumerate(weights): + if ( + w.shape != w0.shape + or w.dtype != w0.dtype + or not w.is_contiguous() + or w.untyped_storage().data_ptr() != storage_ptr + or w.data_ptr() != base_ptr + i * stride * itemsize + ): + return None + # Guard: the shared storage must actually span all E experts from w0's offset. + needed = (w0.storage_offset() + len(weights) * stride) * itemsize + if w0.untyped_storage().nbytes() < needed: + return None + # All experts form one contiguous [E, out, in] block starting at w0's storage offset. + return torch.as_strided(w0, (len(weights), out, in_), (stride, in_, 1)) + + +def _stack_expert_weights(weights: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: + if isinstance(weights, torch.Tensor): + if weights.dim() != 3: + raise ValueError( + f"Stacked expert weights must be 3D [num_experts, out, in], got {weights.shape}." + ) + return weights + if not weights: + raise ValueError("At least one expert weight tensor is required.") + grouped = _try_view_grouped(weights) + return grouped if grouped is not None else torch.stack(weights, dim=0) + + +def permute_free_grouped_gemm_bf16( + hidden_states: torch.Tensor, + weights: torch.Tensor | list[torch.Tensor], + routing: MoERoutingMetadata, + *, + config: Optional[Dict[str, Any]] = None, + activation: Optional[str] = None, + dispatched_probs: Optional[torch.Tensor] = None, + return_preact: bool = False, +) -> torch.Tensor: + """Route-list gather-in-GEMM (FC1 forward) for bf16 MoE. + + Computes ``C[route] = A[route_to_token[route]] @ W[e]^T`` for each expert-sorted route, + writing directly into a compact ``[num_routes, out_features]`` output. + + Parameters + ---------- + hidden_states: + Received activations ``[num_recv_tokens, in_features]``, bf16, contiguous. + weights: + Expert weights ``[num_experts, out_features, in_features]`` or list of ``[out, in]``. + routing: + ``MoERoutingMetadata`` carrying (or able to build) the route-list align buffers. + activation: + When set (``"silu"`` / ``"gelu"``), fuses the **gated** activation into the GEMM + epilogue: ``out_features`` is the gate+up width (``2F``), the output is the ``F``-wide + activated buffer ``act(gate) * up``, and the separate activation kernel is skipped. The + weight output dim must be laid out as ``[gate | up]``. + dispatched_probs: + Optional ``[num_recv_tokens, num_experts]`` gating probabilities. When given (with a + fused ``activation``), each route's ``prob[token, expert]`` is multiplied into the + activation in-kernel, skipping the separate route-prob-apply pass. + return_preact: + When set (only valid with a fused ``activation``), also allocate and fill a + ``[T * min(topk, E), 2F]`` buffer with the raw ``[gate | up]`` pre-activation and return it + alongside the activated output. The backward needs it to reconstruct the 2F GEMM-output + gradient; keep ``False`` for inference to skip the extra store. + + Returns + ------- + torch.Tensor + Worst-case padded ``[T * min(topk, E), out_features]`` (or ``[T * min(topk, E), F]`` with a fused + ``activation``), bf16 (expert-contiguous). The valid rows are the compact route range + ``[0, num_routes)``; the tail ``[num_routes, em_max)`` is inert, *uninitialized* + padding. Consumers MUST read only the compact range, using the routing metadata to + locate each expert's rows: + + - ``route_start[e] = counts[0] + ... + counts[e-1]`` (``= cumsum(counts) - counts``): + its starting offset in the packed output. + - ``num_tokens_post_padded = sum_e ceil(counts[e] / block_size) * block_size``: the + block-padded route count (each expert's row count rounded up to ``block_size``). + """ + if hidden_states.dtype != torch.bfloat16: + raise TypeError( + f"permute_free_grouped_gemm_bf16 requires bf16 input, got {hidden_states.dtype}." + ) + if not hidden_states.is_contiguous(): + hidden_states = hidden_states.contiguous() + + weights_stacked = _stack_expert_weights(weights) + if weights_stacked.stride(-1) != 1: + weights_stacked = weights_stacked.contiguous() + + num_recv_tokens, in_features = hidden_states.shape + num_experts, out_features, in_k = weights_stacked.shape + if in_k != in_features: + raise ValueError( + f"Weight in_features ({in_k}) does not match hidden_states ({in_features})." + ) + if num_experts != routing.num_experts: + raise ValueError( + f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." + ) + + gated_act = activation is not None + if gated_act and out_features % 2 != 0: + raise ValueError( + f"Gated activation requires an even (gate+up) out_features, got {out_features}." + ) + # Fused gated activation halves the stored width (2F -> F). + stored_features = out_features // 2 if gated_act else out_features + + # Copy so the backend-aware block_m override never mutates the caller's config dict. + kernel_config = dict(config or get_default_moe_kernel_config(num_recv_tokens)) + block_size_m = _fwd_align_block_size_m( + hidden_states, + weights_stacked, + gated=gated_act, + default_block_m=int(kernel_config["BLOCK_SIZE_M"]), + num_tokens=num_recv_tokens, + ) + kernel_config["BLOCK_SIZE_M"] = block_size_m + + routing = prepare_moe_align(routing, block_size_m) + + # Worst-case (sync-free) allocation: size the output to the block-padded upper bound + # em_max = sorted_slot_ids.shape[0], which is derived purely from shapes (num_recv_tokens, + # num_experts, block_size) and is always >= num_routes. This avoids the device->host sync + # (.item()) that a compact [num_routes, N] allocation would require. + em_max = routing.sorted_slot_ids.shape[0] + output = torch.empty( + (em_max, stored_features), + dtype=torch.bfloat16, + device=hidden_states.device, + ) + + if return_preact and not gated_act: + raise ValueError("return_preact requires a fused activation.") + preact = ( + torch.empty((em_max, out_features), dtype=torch.bfloat16, device=hidden_states.device) + if return_preact + else None + ) + + _pf_moe_fwd( + hidden_states, + weights_stacked, + output, + routing, + num_recv_tokens=num_recv_tokens, + config=kernel_config, + index_a_by_route_pos=False, + activation=activation, + dispatched_probs=dispatched_probs, + preact_out=preact, + ) + if return_preact: + return output, preact + return output + + +def permute_free_gated_act_bwd( + grad_output: torch.Tensor, + preact: torch.Tensor, + routing: MoERoutingMetadata, + *, + activation: str, + dispatched_probs: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Activation (+ route-prob) backward for the fused permute-free FC1 epilogue. + + Reconstructs the raw ``2F`` GEMM-output gradient ``dpre = [d_gate | d_up]`` (and, when the + forward fused the route prob, the prob gradient) from the saved ``[T * min(topk, E), 2F]`` + pre-activation. The returned ``dpre`` feeds the *unchanged* route-list dgrad / wgrad, so the + FC1 backward stays symmetric with the non-activation path; the caller owns the autograd + boundary and routes ``dprob`` back to ``dispatched_probs``. + + Parameters + ---------- + grad_output: + ``[T * min(topk, E), F]`` grad wrt the fused FC1 output (route/padded layout). + preact: + ``[T * min(topk, E), 2F]`` raw ``[gate | up]`` pre-activation saved by the forward. + dispatched_probs: + ``[num_recv_tokens, E]`` gating probs (or ``None`` for a silu/gelu-only fusion). + + Returns + ------- + ``(dpre, dprob)`` -- ``dpre`` is ``[T * min(topk, E), 2F]`` bf16; ``dprob`` matches + ``dispatched_probs`` (or ``None`` when no probs were fused). + """ + routes_max = int(routing.route_to_token.shape[0]) + token = routing.route_to_token.to(torch.int32) + expert = _expert_per_route(routing, routes_max) + # The route buffers are statically sized to the worst case ``routes_max = T * topk``, but + # only the dense head ``[0, num_routes)`` is real (under EP this can be ~topk*E_local/E + # smaller). ``num_tokens_post_padded`` is a device scalar >= num_routes (block-padded), so + # it bounds the kernel to the real routes -- sync-free -- and lets the tail exit early. + return fused_gated_act_prob_bwd( + grad_output.contiguous(), + preact, + token, + expert, + num_recv_tokens=routing.num_recv_tokens, + activation=activation, + dispatched_probs=dispatched_probs, + num_routes_bound=routing.num_tokens_post_padded, + ) + + +def permute_free_gated_act_recompute( + preact: torch.Tensor, + routing: MoERoutingMetadata, + *, + activation: str, + dispatched_probs: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Rebuild the fused FC1 activation ``act(gate) * up * prob`` from the saved pre-activation. + + Lets the backward checkpoint only the ``2F`` pre-activation (``[T * min(topk, E), 2F]``) and + reconstruct the ``F``-wide activation just-in-time for the FC2 wgrad, feeding the + *unchanged* full-speed wgrad kernel a transient buffer (freed right after) instead of + persisting the activation across the fwd/bwd boundary. Uses the same per-route routing + arrays as :func:`permute_free_gated_act_bwd`. + + Returns + ------- + torch.Tensor + ``act``, shape ``[T * min(topk, E), F]``, bf16 (route/padded layout). + """ + routes_max = int(routing.route_to_token.shape[0]) + token = routing.route_to_token.to(torch.int32) + expert = _expert_per_route(routing, routes_max) + return fused_gated_act_prob_fwd( + preact, + token, + expert, + num_recv_tokens=routing.num_recv_tokens, + activation=activation, + dispatched_probs=dispatched_probs, + num_routes_bound=routing.num_tokens_post_padded, + ) + + +def permute_free_grouped_gemm_bf16_dgrad( + grad_output: torch.Tensor, + weights: torch.Tensor | list[torch.Tensor], + routing: MoERoutingMetadata, + *, + config: Optional[Dict[str, Any]] = None, +) -> torch.Tensor: + """Route-list gather-in-GEMM dgrad (FC1 backward wrt input). + + ``grad_output`` is the compact per-route gradient ``[num_routes, out_features]``. The + kernel contracts over ``out_features`` to produce per-route ``dX = [num_routes, in]``, + which is then scatter-added back onto the received tokens. + + Returns + ------- + torch.Tensor + ``dA``, shape ``[num_recv_tokens, in_features]``, bf16. + """ + if grad_output.dtype != torch.bfloat16: + raise TypeError( + f"permute_free_grouped_gemm_bf16_dgrad requires bf16 grad, got {grad_output.dtype}." + ) + if grad_output.dim() != 2: + raise ValueError( + f"grad_output must be compact [num_routes, out_features], got {grad_output.shape}." + ) + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + weights_stacked = _stack_expert_weights(weights) + num_experts, out_features, in_features = weights_stacked.shape + if grad_output.shape[-1] != out_features: + raise ValueError( + f"grad_output out_features ({grad_output.shape[-1]}) " + f"does not match weights ({out_features})." + ) + if num_experts != routing.num_experts: + raise ValueError( + f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." + ) + + # Contract over out_features via the [E, in, out] transposed *view* (stride relabel only). + if weights_stacked.stride(-1) != 1: + weights_stacked = weights_stacked.contiguous() + weights_t = weights_stacked.transpose(1, 2) + + if routing.sorted_slot_ids is None or routing.block_size_m is None: + fwd_config = config or get_default_moe_kernel_config(routing.num_recv_tokens) + routing = prepare_moe_align(routing, int(fwd_config["BLOCK_SIZE_M"])) + block_size_m = int(routing.block_size_m) + + dgrad_config = get_default_moe_kernel_config(int(grad_output.shape[0])) + dgrad_config = {**dgrad_config, "BLOCK_SIZE_M": block_size_m} + + # Two-stage dgrad reduction (contention-free): (1) plain compact store of each per-route + # dX[route] = grad[route] @ W1[e] into an [T * min(topk, E), in] bf16 buffer (coalesced, no atomics), + # then (2) a token-parallel gather-combine summing each token's route rows. This replaces + # the fused atomic scatter-to-token, whose per-token contention dominated the FC1 dgrad. + em_max = routing.sorted_slot_ids.shape[0] + compact = torch.empty( + (em_max, in_features), + dtype=torch.bfloat16, + device=grad_output.device, + ) + _pf_moe_fwd( + grad_output, + weights_t, + compact, + routing, + num_recv_tokens=routing.num_recv_tokens, + config=dgrad_config, + index_a_by_route_pos=True, + ) + return route_gather_combine( + compact, + routing.token_routes, + routing.token_route_count, + routing.num_recv_tokens, + out_dtype=torch.bfloat16, + ) + + +def permute_free_grouped_gemm_bf16_wgrad( + hidden_states: torch.Tensor, + grad_output: torch.Tensor, + weights_shape, + routing: MoERoutingMetadata, + *, + config: Optional[Dict[str, Any]] = None, + out: Optional[torch.Tensor] = None, + accumulate: bool = False, + swap_gather: bool = False, +) -> torch.Tensor: + """Route-list fused weight-gradient (FC1 backward wrt weights), FlyDSL-only. + + Computes ``dW[e] = sum_{route in e} grad[route]^T @ A[route_to_token[route]]`` by + gathering the activation operand (received-token row) and reading the compact grad row + in the FlyDSL kernel. + + Parameters + ---------- + hidden_states: + Received activations ``[num_recv_tokens, in_features]``, bf16. + grad_output: + Compact per-route gradient ``[num_routes, out_features]``, bf16. + weights_shape: + ``(num_experts, out_features, in_features)``. + out: + Optional ``[E, out, in]`` bf16 destination to fold the wgrad into directly + (``+=`` if ``accumulate`` else ``=``). + accumulate: + With ``out``: add into it rather than overwrite. + swap_gather: + FC2 wgrad mode (token-gather on grad). Not implemented in the current FlyDSL kernel. + + Returns + ------- + torch.Tensor + The wgrad ``[num_experts, out_features, in_features]`` -- ``out`` when supplied, else a + freshly allocated bf16 tensor. + """ + if hidden_states.dtype != torch.bfloat16 or grad_output.dtype != torch.bfloat16: + raise TypeError("permute_free_grouped_gemm_bf16_wgrad requires bf16 inputs.") + if grad_output.dim() != 2: + raise ValueError( + f"grad_output must be compact [num_routes, out_features], got {grad_output.shape}." + ) + + num_experts, out_features, in_features = (int(v) for v in weights_shape) + if grad_output.shape[-1] != out_features: + raise ValueError( + f"grad_output out_features ({grad_output.shape[-1]}) " + f"does not match weights ({out_features})." + ) + if hidden_states.shape[-1] != in_features: + raise ValueError( + f"hidden_states in_features ({hidden_states.shape[-1]}) " + f"does not match weights ({in_features})." + ) + if num_experts != routing.num_experts: + raise ValueError( + f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." + ) + + x = hidden_states + if not x.is_contiguous(): + x = x.contiguous() + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + contract_m = _WGRAD_CONTRACT_M + routing = _prepare_wgrad_align(routing, contract_m) + + if swap_gather: + raise RuntimeError( + "swap_gather wgrad (FC2 direct [E, H, F] layout) is not implemented in the " + "current FlyDSL wgrad kernel." + ) + + flydsl_wgrad = _get_flydsl_wgrad() + + if out is not None: + assert tuple(out.shape) == (num_experts, out_features, in_features), ( + f"wgrad out shape {tuple(out.shape)} != {(num_experts, out_features, in_features)}" + ) + if out.dtype != torch.bfloat16: + raise TypeError( + f"FlyDSL wgrad requires bf16 out, got {out.dtype}." + ) + flydsl_wgrad( + x, + grad_output, + out, + routing.wgrad_sorted_slot_ids, + routing.wgrad_block_start, + routing.wgrad_blocks_per_expert, + routing.route_start, + num_recv_tokens=routing.num_recv_tokens, + accumulate=bool(accumulate), + ) + return out + + dW = torch.zeros( + (num_experts, out_features, in_features), + dtype=torch.bfloat16, + device=hidden_states.device, + ) + flydsl_wgrad( + x, + grad_output, + dW, + routing.wgrad_sorted_slot_ids, + routing.wgrad_block_start, + routing.wgrad_blocks_per_expert, + routing.route_start, + num_recv_tokens=routing.num_recv_tokens, + accumulate=bool(accumulate), + ) + return dW + + +def permute_free_grouped_gemm_bf16_fc2( + fc2_input: torch.Tensor, + weights: torch.Tensor | list[torch.Tensor], + routing: MoERoutingMetadata, + *, + config: Optional[Dict[str, Any]] = None, + activation: Optional[str] = None, + dispatched_probs: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Route-list FC2 forward with gather-combine to token order (bf16 MoE). + + ``fc2_input`` is the route-ordered FC1 output. When ``activation`` is set (``"silu"`` / + ``"gelu"``) it is the raw ``2F`` ``[gate | up]`` pre-activation and the kernel fuses + ``act(gate) * up [* prob]`` on the ``A`` operand before the GEMM; otherwise it is the + ``F``-wide activated buffer (legacy path). For each route the kernel reads its row + (``index_a_by_route_pos=True``) and computes the compact per-route GEMM; a separate + contention-free gather-combine pass sums each token's routes into its output row. + + Parameters + ---------- + fc2_input: + Route-ordered activations ``[T * min(topk, E), in_features]`` (``F``) or, when + ``activation`` is set, the raw ``2F`` ``[gate | up]`` pre-activation from FC1. + weights: + Expert weights ``[num_experts, out_features, in_features]`` (W2) or list of ``[out, + in]``. + activation: + When set, fuse ``act(gate) * up`` (+ optional ``dispatched_probs``) on the ``A`` + operand (FC2 prologue). Requires ``fc2_input.shape[-1] == 2 * in_features``. + dispatched_probs: + Optional ``[num_recv_tokens, num_experts]`` gating probabilities (fused in the FC2 + prologue when ``activation`` is set). + routing: + ``MoERoutingMetadata`` carrying (or able to build) the route-list align buffers. + + Returns + ------- + torch.Tensor + ``[num_recv_tokens, out_features]``, bf16 (token order); ready for the cross-rank + combine. No separate unpermute is needed. + """ + if fc2_input.dtype != torch.bfloat16: + raise TypeError( + f"permute_free_grouped_gemm_bf16_fc2 requires bf16 input, got {fc2_input.dtype}." + ) + if not fc2_input.is_contiguous(): + fc2_input = fc2_input.contiguous() + + weights_stacked = _stack_expert_weights(weights) + if weights_stacked.stride(-1) != 1: + weights_stacked = weights_stacked.contiguous() + + num_experts, out_features, in_features = weights_stacked.shape + gated_a = activation is not None + if gated_a: + if fc2_input.shape[-1] != 2 * in_features: + raise ValueError( + f"gated FC2 expects 2F preact input ({2 * in_features} cols), " + f"got {fc2_input.shape[-1]}." + ) + elif fc2_input.shape[-1] != in_features: + raise ValueError( + f"fc2_input in_features ({fc2_input.shape[-1]}) does not match weights " + f"({in_features})." + ) + if num_experts != routing.num_experts: + raise ValueError( + f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." + ) + + if routing.sorted_slot_ids is None or routing.block_size_m is None: + fwd_config = config or get_default_moe_kernel_config(routing.num_recv_tokens) + routing = prepare_moe_align(routing, int(fwd_config["BLOCK_SIZE_M"])) + block_size_m = int(routing.block_size_m) + kernel_config = get_default_moe_kernel_config(int(fc2_input.shape[0])) + kernel_config = {**kernel_config, "BLOCK_SIZE_M": block_size_m} + + # Two-stage combine (contention-free): (1) plain compact store of each per-route result + # y[route] = fc2_input[route] @ W2[e] into an [T * min(topk, E), out] bf16 buffer (coalesced, no + # atomics), then (2) a token-parallel gather-combine that sums each token's route rows. + # This replaces the fused atomic scatter-to-token, whose per-token atomic contention made + # the FC2 forward ~2x the FC1 forward. Gather has no contention, so the combine is ~free. + em_max = routing.sorted_slot_ids.shape[0] + compact = torch.empty( + (em_max, out_features), + dtype=torch.bfloat16, + device=fc2_input.device, + ) + _pf_moe_fwd( + fc2_input, + weights_stacked, + compact, + routing, + num_recv_tokens=routing.num_recv_tokens, + config=kernel_config, + index_a_by_route_pos=True, + activation=activation, + dispatched_probs=dispatched_probs, + gated_a=gated_a, + ) + return route_gather_combine( + compact, + routing.token_routes, + routing.token_route_count, + routing.num_recv_tokens, + out_dtype=torch.bfloat16, + ) + + +def permute_free_grouped_gemm_bf16_fc2_dgrad( + grad_output: torch.Tensor, + weights: torch.Tensor | list[torch.Tensor], + routing: MoERoutingMetadata, + *, + config: Optional[Dict[str, Any]] = None, +) -> torch.Tensor: + """FC2 dgrad: gather the token-space grad back into the compact route buffer. + + ``grad_output`` is the token-space gradient ``[num_recv_tokens, out_features]`` (the grad + of FC2's fused-scatter forward output). For each route the kernel gathers its received + token's grad row (``index_a_by_route_pos=False``) and computes ``grad[token] @ W2[e]`` + (contracting over ``out_features``), writing the compact per-route + ``d(fc2_input) = [T * min(topk, E), in_features]``. + + Returns + ------- + torch.Tensor + ``[T * min(topk, E), in_features]``, bf16 (route order; valid rows are ``[0, num_routes)``). + """ + if grad_output.dtype != torch.bfloat16: + raise TypeError( + f"permute_free_grouped_gemm_bf16_fc2_dgrad requires bf16 grad, got " + f"{grad_output.dtype}." + ) + if grad_output.dim() != 2: + raise ValueError( + f"grad_output must be [num_recv_tokens, out_features], got {grad_output.shape}." + ) + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + weights_stacked = _stack_expert_weights(weights) + num_experts, out_features, in_features = weights_stacked.shape + if grad_output.shape[-1] != out_features: + raise ValueError( + f"grad_output out_features ({grad_output.shape[-1]}) does not match weights " + f"({out_features})." + ) + if num_experts != routing.num_experts: + raise ValueError( + f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." + ) + + # Contract over out_features via the [E, in, out] transposed view (stride relabel only). + if weights_stacked.stride(-1) != 1: + weights_stacked = weights_stacked.contiguous() + weights_t = weights_stacked.transpose(1, 2) + + if routing.sorted_slot_ids is None or routing.block_size_m is None: + fwd_config = config or get_default_moe_kernel_config(routing.num_recv_tokens) + routing = prepare_moe_align(routing, int(fwd_config["BLOCK_SIZE_M"])) + block_size_m = int(routing.block_size_m) + dgrad_config = get_default_moe_kernel_config(int(grad_output.shape[0])) + dgrad_config = {**dgrad_config, "BLOCK_SIZE_M": block_size_m} + + # Padded route-order output; the gather-GEMM writes only the compact [0, num_routes) + # range and never visits the tail (bounded by num_tokens_post_padded, sentinel-masked). + # Left uninitialized (torch.empty) to skip the em_max*N zero-init; consumers read only the + # compact range via the routing metadata, so the garbage tail is never observed. + em_max = routing.sorted_slot_ids.shape[0] + dgrad = torch.empty( + (em_max, in_features), + dtype=torch.bfloat16, + device=grad_output.device, + ) + _pf_moe_fwd( + grad_output, + weights_t, + dgrad, + routing, + num_recv_tokens=grad_output.shape[0], + config=dgrad_config, + index_a_by_route_pos=False, + ) + return dgrad + + +def permute_free_grouped_gemm_bf16_fc2_wgrad( + fc2_input: torch.Tensor, + grad_output: torch.Tensor, + weights_shape, + routing: MoERoutingMetadata, + *, + config: Optional[Dict[str, Any]] = None, + out: Optional[torch.Tensor] = None, + accumulate: bool = False, + preact: Optional[torch.Tensor] = None, + dispatched_probs: Optional[torch.Tensor] = None, + activation: Optional[str] = None, +) -> torch.Tensor: + """FC2 wgrad: ``dW2[e] = sum_{route in e} grad[token(route)]^T @ fc2_input[route]``. + + The two operands play mirror roles to FC1: ``grad_output`` is token-space and must be + token-gathered, ``fc2_input`` is route-ordered and read contiguously. Which one the kernel + gathers decides the output orientation: + + * **FC2 wgrad** uses operand swap + transpose: compute ``[E, F, H]`` via the FC1 FlyDSL + kernel, then ``transpose(1, 2)`` to ``[E, H, F]``. + + Parameters + ---------- + fc2_input: + Route-ordered activations ``[T * min(topk, E), in_features(F)]``, bf16. + grad_output: + Token-space gradient ``[num_recv_tokens, out_features(H)]``, bf16. + weights_shape: + ``(num_experts, out_features(H), in_features(F))`` -- the W2 shape. + out / accumulate: + Optional ``[E, H, F]`` destination; ``+=`` if ``accumulate`` else ``=``. Returned when given. + preact / dispatched_probs / activation: + Recompute-from-preact mode. When ``preact`` (``[T * min(topk, E), 2F]`` raw ``[gate | up]``) + is given, ``fc2_input`` is ignored and the ``F``-wide activation ``act(gate)*up*prob`` is + re-materialised here into a transient buffer (freed after the wgrad), so the backward can + checkpoint only the ``2F`` preact instead of the ``F``-wide activation. ``activation`` picks + the nonlinearity (``"silu"``/``"gelu"``); ``dispatched_probs`` (``[num_recv, E]``) is the + fused route-prob table (omit when the forward did not fuse probs). The wgrad kernel itself + is unchanged and runs at full stored-act speed. + """ + num_experts, out_features, in_features = (int(v) for v in weights_shape) + + # Recompute-from-preact: rebuild the transient activation for the (unchanged) wgrad. + if preact is not None: + if activation is None: + raise ValueError("permute_free FC2 wgrad recompute requires an activation.") + fc2_input = permute_free_gated_act_recompute( + preact, routing, activation=activation, dispatched_probs=dispatched_probs + ) + + # FC2 wgrad: operand swap + transpose via the FC1 FlyDSL kernel. + dW_t = permute_free_grouped_gemm_bf16_wgrad( + grad_output, + fc2_input, + (num_experts, in_features, out_features), + routing, + config=config, + ) + dW = dW_t.transpose(1, 2).contiguous() # [E, H, F] + if out is None: + return dW + if accumulate: + out.add_(dW) + else: + out.copy_(dW) + return out + + +# --------------------------------------------------------------------------- +# Direction-aware dispatch: pick the FC1/FC2 fwd/bwd kernels from the routing +# metadata so callers (e.g. GroupedLinear) don't have to branch on route_space / +# activation themselves. +# --------------------------------------------------------------------------- +@dataclass +class PermuteFreeForwardResult: + """Output of :func:`permute_free_grouped_gemm_forward`. + + ``preact`` is the raw ``2F`` pre-activation saved for the FC1 gated-activation backward; + it is ``None`` on every other path (FC2, or FC1 without a fused activation / without a + backward). + """ + + out: torch.Tensor + preact: Optional[torch.Tensor] = None + + +@dataclass +class PermuteFreeBackwardResult: + """Gradients from :func:`permute_free_grouped_gemm_backward`. + + ``dgrad`` / ``wgrad_stacked`` are ``None`` when the corresponding gradient was not requested. + ``wgrad_stacked`` is the weight gradient as a single ``[E, out, in]`` tensor (the natural + kernel output): accumulate it directly into a grouped param's ``main_grad``, or split it into + per-expert views (``list(wgrad_stacked)``) for the positional autograd return. ``grad_probs`` + is the route-prob gradient for the FC1 fused-prob path (``None`` otherwise). + """ + + dgrad: Optional[torch.Tensor] = None + grad_probs: Optional[torch.Tensor] = None + wgrad_stacked: Optional[torch.Tensor] = None + # True when the wgrad was written directly into the caller-provided ``wgrad_out`` (FC1 + # path), so the caller must not re-apply it. False when ``wgrad_stacked`` is a fresh tensor + # the caller still needs to sink (FC2 -- its transpose precludes an in-place accumulate -- + # or the plain positional-return path). + wgrad_applied: bool = False + + +def permute_free_grouped_gemm_forward( + hidden_states: torch.Tensor, + weights: torch.Tensor | list[torch.Tensor], + routing: MoERoutingMetadata, + *, + activation: Optional[str] = None, + dispatched_probs: Optional[torch.Tensor] = None, +) -> PermuteFreeForwardResult: + """Dispatch the permute-free grouped-GEMM forward from the routing direction + fusion hints. + + ``routing.route_space`` selects the direction: + + - ``True`` (FC2): route-ordered ``2F`` pre-activation (FC1 output). The gated activation + ``act(gate)*up[*prob]`` is applied in a standalone pass into an ``F``-wide transient, then + a plain gather-GEMM + gather-combine -> ``[num_recv, out]``. (Fusing the activation into + the FC2 GEMM prologue regressed throughput, so FC2 is kept as a plain DMA GEMM mirroring + FC1.) The ``F``-wide transient is freed after FC2; only the ``2F`` pre-activation is + checkpointed for backward (which recomputes the ``F`` activation just-in-time). + - ``False`` (FC1): gather-in-GEMM -> padded ``[T * min(topk, E), 2F]`` raw ``[gate | up]`` + pre-activation (no activation fusion here; the ``activation`` hint is consumed on FC2). + """ + if getattr(routing, "route_space", False): + fc2_input = hidden_states + if activation is not None: + # Standalone gated-activation pass on the raw 2F [gate|up] pre-activation, producing + # the F-wide FC2 operand. This is the same route-wise kernel the backward uses to + # recompute the activation, so the caller can keep checkpointing only the 2F + # pre-activation (this F-wide buffer is transient and freed after the FC2 GEMM). + fc2_input = permute_free_gated_act_recompute( + hidden_states, + routing, + activation=activation, + dispatched_probs=dispatched_probs, + ) + return PermuteFreeForwardResult( + permute_free_grouped_gemm_bf16_fc2(fc2_input, weights, routing) + ) + return PermuteFreeForwardResult( + permute_free_grouped_gemm_bf16(hidden_states, weights, routing) + ) + + +def permute_free_grouped_gemm_backward( + grad_output: torch.Tensor, + *, + routing: MoERoutingMetadata, + weights: list[torch.Tensor], + num_gemms: int, + hidden_states: Optional[torch.Tensor] = None, + requires_dgrad: bool = False, + requires_wgrad: bool = False, + fc1_activation: Optional[str] = None, + preact: Optional[torch.Tensor] = None, + dispatched_probs: Optional[torch.Tensor] = None, + fc2_activation: Optional[str] = None, + wgrad_out: Optional[torch.Tensor] = None, + wgrad_accumulate: bool = False, +) -> PermuteFreeBackwardResult: + """Dispatch the permute-free grouped-GEMM backward (mirror of the forward dispatch). + + ``routing.route_space`` picks the FC2 vs FC1 dgrad/wgrad kernels. On the FC1 + gated-activation path (``fc1_activation`` set) the raw ``2F`` GEMM-output gradient (and the + route-prob gradient) is first reconstructed from the saved ``preact`` and fed to the + *unchanged* dgrad/wgrad. ``hidden_states`` (the forward input) is only needed for wgrad. + + On the FC2 path (``route_space=True``), when ``fc2_activation`` is set together with + ``preact`` the wgrad rebuilds its F-wide input (``act(gate)*up*prob``) in-flight from the FC1 + pre-activation instead of consuming a stored ``hidden_states`` -- the FC1->FC2 recompute + handoff, letting the FC2 forward skip saving its F-wide input. + + ``wgrad_out`` (optional ``[E, out, in]`` accumulator) folds the wgrad straight into the + caller's buffer (``+=`` if ``wgrad_accumulate`` else ``=``) instead of returning a fresh + tensor. FC2 applies via transpose after the FlyDSL kernel. When used, ``result.wgrad_applied`` + is ``True`` and ``result.wgrad_stacked`` is that same buffer. + """ + grad_output = grad_output.contiguous() + route_space = getattr(routing, "route_space", False) + dgrad = None + grad_probs = None + wgrad_stacked = None + wgrad_applied = False + + if route_space: + # FC2: grad is token-space [num_recv, out]. dgrad gathers back to the compact route + # buffer [T * min(topk, E), F] (GEMM operand), then the gated-activation backward maps + # dL/dF -> dL/d(2F) (+ dprob) when the forward fused the FC2 prologue. + if requires_dgrad: + weights_stacked = _stack_expert_weights(weights) # [E, H, F], zero-copy when grouped + dgrad_f = permute_free_grouped_gemm_bf16_fc2_dgrad( + grad_output, weights_stacked, routing + ) + if fc2_activation is not None and hidden_states is not None: + dgrad, grad_probs = permute_free_gated_act_bwd( + dgrad_f, + hidden_states, + routing, + activation=fc2_activation, + dispatched_probs=dispatched_probs, + ) + if not (dispatched_probs is not None and dispatched_probs.requires_grad): + grad_probs = None + else: + dgrad = dgrad_f + if requires_wgrad: + weights_shape = (num_gemms, weights[0].size(0), weights[0].size(1)) + recompute = fc2_activation is not None and hidden_states is not None + dW = permute_free_grouped_gemm_bf16_fc2_wgrad( + hidden_states, grad_output, weights_shape, routing, + out=wgrad_out, accumulate=wgrad_accumulate, + preact=hidden_states if recompute else None, + dispatched_probs=dispatched_probs if recompute else None, + activation=fc2_activation if recompute else None, + ) # [E, H, F] + wgrad_stacked = dW + wgrad_applied = wgrad_out is not None + else: + # FC1: grad is the padded [T * min(topk, E), 2F] route buffer (from FC2 act-bwd); dgrad + # scatters the input grad back to received-token rows. No FC1 activation backward. + if requires_dgrad: + weights_stacked = _stack_expert_weights(weights) # [E, N, H], zero-copy when grouped + dgrad = permute_free_grouped_gemm_bf16_dgrad(grad_output, weights_stacked, routing) + if requires_wgrad: + weights_shape = (num_gemms, weights[0].size(0), weights[0].size(1)) + dW = permute_free_grouped_gemm_bf16_wgrad( + hidden_states, grad_output, weights_shape, routing, + out=wgrad_out, accumulate=wgrad_accumulate, + ) # [E, N, H] + wgrad_stacked = dW + wgrad_applied = wgrad_out is not None + + return PermuteFreeBackwardResult( + dgrad=dgrad, grad_probs=grad_probs, wgrad_stacked=wgrad_stacked, + wgrad_applied=wgrad_applied, + ) diff --git a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py new file mode 100644 index 000000000..b9ef557f2 --- /dev/null +++ b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py @@ -0,0 +1,461 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""FlyDSL permute-free MoE route-list forward gather-GEMM op wrapper. + +Mirrors the route-list forward gather-GEMM contract so the same routing +metadata (``sorted_slot_ids`` / ``expert_ids`` / ``block_start`` / ``route_start``) can be +reused verbatim. Writes the compact ``[em_max, WIDTH_N]`` route output in place. + +Forward gather-GEMM is FlyDSL-only. This module retains Triton kernels for routing +metadata construction, gated-activation recompute/bwd, and the token-order gather-combine +pass that follows the compact route-list GEMM outputs. +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch + +__all__ = [ + "flydsl_moe_fwd", + "flydsl_moe_fwd_autotuned", + "flydsl_moe_fwd_supported", + "flydsl_moe_fwd_pick_block_m", +] + +# Activation ids for fused-epilogue kernels (v2, not yet in-tree). +ACT_SILU = 0 +ACT_GELU = 1 +_ACT_IDS = {"silu": ACT_SILU, "gelu": ACT_GELU} +_WMMA = 16 +_FILL_V = 8 +_WARP = 64 +_LDS_PAD = 8 +_LDS_LIMIT = 163840 # gfx950 per-workgroup LDS (160 KB) + +_WARP_CHOICES = [1, 2, 4, 8, 16] + + +def _env_flag(name: str, default: bool) -> bool: + v = os.environ.get(name) + if v is None: + return default + return v.strip().lower() not in ("0", "false", "no", "off", "") + + +def _v3_enabled() -> bool: + """Whether to route plain GEMMs through the in-tree v3 (MegaMOE-ported) kernels.""" + return _env_flag("AITER_MOE_FLYDSL_V3", True) + + +# MegaMOE's hand-tuned bf16 grouped-GEMM tile geometry (offline-swept in primus-turbo's +# bench_mega_moe: BLOCK_M/BLOCK_N=256, GROUP_M=4 fwd / GROUP_M=8 FC1 NN dgrad, num_xcd=1, +# nt_vmcnt=3). TE's FlyDSL fwd align already bumps non-gated FC1 / dgrad / plain FC2 to +# block_size_m=256; v3 pins the same Mega M-tile rather than trusting the wrapper arg. +_V3_BLOCK_M = 256 +_V3_BLOCK_N = 256 +_V3_GROUP_M = 4 +_V3_DGRAD_FC1_GROUP_M = 8 # bench_mega_moe grouped_gemm_combine L1-dgrad (NN) sweep +_V3_NUM_XCD = 1 + + +def _run_v3_fwd( + A, B, C, sorted_slot_ids, expert_ids, *, num_recv_tokens, block_m, + transpose_b, index_a_by_route_pos, gated, gated_a, mul_prob, save_preact, +): + """Dispatch a plain fwd/dgrad GEMM to the v3 (MegaMOE-ported) kernels. + + Covers the three non-fused paths: FC1 gather (``index_a_by_route_pos=False``), FC2 + route-read (``index_a_by_route_pos=True``) and dgrad (``transpose_b``). Fused epilogues + (gated activation / route-prob / pre-activation save) have no v3 equivalent and raise. + v3 uses MegaMOE's fixed tile geometry (32x32x16, ``BLOCK_N=256``, ``GROUP_M=4``, + ``num_xcd=1``), so the wrapper's ``block_n``/``block_k``/warp args are ignored here. + """ + if gated or gated_a or mul_prob or save_preact: + raise RuntimeError( + "Fused permute-free forward (gated activation / dispatched_probs / preact_out) " + "requires the v2 FlyDSL kernel (moe_fwd_flydsl_v2), which is not in-tree yet. " + "Use standalone gated-act + plain v3 GEMM (GroupedLinear FC2 path), or port v2." + ) + block_m = int(block_m) + em_max = int(sorted_slot_ids.shape[0]) + if em_max % block_m != 0: + raise ValueError( + f"v3 expects em_max ({em_max}) divisible by block_m ({block_m}); " + "check routing align block_size_m." + ) + num_tile_blocks = em_max // block_m + expert_ids_i32 = expert_ids.to(torch.int32) + if expert_ids_i32.numel() != num_tile_blocks: + raise ValueError( + f"v3 expects expert_ids length {num_tile_blocks} (em_max={em_max}, BLOCK_M={block_m}), " + f"got {expert_ids_i32.numel()}; routing align block_size_m must match block_m={block_m}" + ) + + if transpose_b: + # dgrad: v3 is a native NN GEMM contracting the incoming grad against the weight over + # the output-feature axis. The wrapper's B is the transposed-weight *view* [E, out, in] + # (stride relabel); transpose(1,2) recovers the [E, N=out, K=in] forward weight v3 + # dgrad reads NN. FC1 dgrad (index_a_by_route_pos=True) is compact route-read (grad rows + # == dx rows); FC2 dgrad (index_a_by_route_pos=False) gathers the token-space grad + # [num_recv, N] into the compact route output [em_max, K]. + from ..flydsl_kernels.permute_free_grouped_gemm.pf_dgrad import grouped_gemm_dgrad_bf16 + + # TE passes a stride-relabelled transpose *view* (``stride(2) != 1``); undo the view to + # recover the forward ``[E, N, K]`` storage without copying (``transpose`` twice == id). + weight = B.transpose(1, 2) if B.stride(2) != 1 else B + # Real M-tile count for the dgrad grid: derive from the padded pool shape (host-known, + # graph-capture safe). Padding tail blocks (expert_ids=-1) early-exit in-kernel, same as + # the forward v3 path. Do not count active blocks on-device ((expert_ids>=0).sum().item()) + # -- that syncs the GPU and breaks HIP/CUDA graph capture. + dgrad_group_m = _V3_DGRAD_FC1_GROUP_M if index_a_by_route_pos else _V3_GROUP_M + grouped_gemm_dgrad_bf16( + A, weight, C, expert_ids_i32, num_tile_blocks, sorted_slot_ids, + gather=not index_a_by_route_pos, + BLOCK_M=block_m, BLOCK_N=_V3_BLOCK_N, GROUP_M=dgrad_group_m, num_xcd=_V3_NUM_XCD, + ) + return + + from ..flydsl_kernels.permute_free_grouped_gemm.pf_fwd import grouped_gemm_gather_bf16 + + grouped_gemm_gather_bf16( + A, B, C, expert_ids_i32, num_tile_blocks, sorted_slot_ids, + gather=not index_a_by_route_pos, + BLOCK_M=block_m, BLOCK_N=_V3_BLOCK_N, GROUP_M=_V3_GROUP_M, num_xcd=_V3_NUM_XCD, + ) + + +def _mfma_dim(transpose_b: bool) -> int: + """MFMA output-tile edge: forward GEMM uses the 32x32x16 atom, dgrad the 16x16x32 atom. + + Kept in sync with the kernel's ``MOE_FWD_MFMA32`` escape hatch so the autotuner sweeps the + tile-divisibility that the compiled atom actually requires. + """ + if transpose_b: + return 16 + return 32 if _env_flag("MOE_FWD_MFMA32", False) else 16 + + +def _warp_valid(block_m, block_n, block_k, wm, wn, transpose_b=False): + n_threads = wm * wn * _WARP + wmma = _mfma_dim(transpose_b) + if block_m % (wm * wmma) or block_n % (wn * wmma): + return False + if (block_m * block_k) % (n_threads * _FILL_V): + return False + if (block_n * block_k) % (n_threads * _FILL_V): + return False + return True + + +def _pick_warps(block_m: int, block_n: int, block_k: int, transpose_b=False): + """Pick (warps_m, warps_n) balancing per-warp MFMA tile (M_STEPS x N_STEPS) vs occupancy. + + Prefers keeping the per-warp atom counts moderate (fewer accumulators -> more waves) while + landing a 256-512 thread workgroup, which measured fastest across the Qwen MoE shapes. + """ + wmma = _mfma_dim(transpose_b) + best = None + for wm in _WARP_CHOICES: + for wn in _WARP_CHOICES: + if not _warp_valid(block_m, block_n, block_k, wm, wn, transpose_b): + continue + n_threads = wm * wn * _WARP + if n_threads > 512: + continue + m_steps = block_m // (wm * wmma) + n_steps = block_n // (wn * wmma) + # Favor a small, balanced per-warp atom footprint (fewer accumulators -> more + # waves), then a 256-512 thread workgroup. Ties broken toward |m_steps-n_steps| + # small (balanced reuse of A and B fragments). + score = ( + m_steps * n_steps, + abs(m_steps - n_steps), + 0 if 256 <= n_threads <= 512 else 1, + n_threads, + ) + if best is None or score < best[0]: + best = (score, (wm, wn)) + return best[1] if best is not None else None + + +def _fwd_buffering(): + """(num_buffers, lds_pad) for the production DMA+swizzle fill path. + + Both default on (opt out with ``MOE_FWD_DMA=0`` / ``MOE_FWD_SWZ=0``). The DMA path runs + a distance-2, 3-buffer ring; the register fallback keeps 2-buffer ping/pong. + """ + use_dma = _env_flag("MOE_FWD_DMA", True) + swz = _env_flag("MOE_FWD_SWZ", True) + pad = 0 if (use_dma or swz) else _LDS_PAD + return (3 if use_dma else 2), pad + + +def _lds_bytes(block_m, block_n, block_k, gated, transpose_b=False, gated_a=False): + n_bt = 2 if gated else 1 + nbuf, pad = _fwd_buffering() + a_tile = block_m * (block_k + pad) + # Hybrid gated_a DMA path stages the raw ``up`` half in a parallel A tile (2x A LDS). + hybrid_a = gated_a and _env_flag("MOE_FWD_DMA", True) and _env_flag("MOE_FWD_GATEDA_DMA", False) + a_tiles = 2 if hybrid_a else 1 + # dgrad stages B as [k, n] (row stride = block_n+pad); fwd as [n, k] (block_k+pad). + b_tile = block_k * (block_n + pad) if transpose_b else block_n * (block_k + pad) + return (a_tiles * a_tile + n_bt * b_tile) * nbuf * 2 + + +def _default_block_n(block_m, block_k, gated, transpose_b=False): + """Widest N tile (128 then 64) that fits LDS -- wider N raises arithmetic intensity.""" + for bn in (128, 64): + if _lds_bytes(block_m, bn, block_k, gated, transpose_b) <= _LDS_LIMIT and bn % _mfma_dim(transpose_b) == 0: + return bn + return 64 + + +def flydsl_moe_fwd_supported( + A: torch.Tensor, + B: torch.Tensor, + *, + block_m: int, + block_n: int = 64, + block_k: int = 64, +) -> bool: + """Whether the in-tree FlyDSL fwd/dgrad kernels can handle these operands.""" + if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: + return False + if A.stride(1) != 1: + return False + if B.stride(2) != 1 and B.stride(1) != 1: + return False + em_max = None # only known at launch for C; skip em_max % block_m precheck here + _ = (block_n, block_k, em_max) + return True + + +def flydsl_moe_fwd( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + sorted_slot_ids: torch.Tensor, + expert_ids: torch.Tensor, + block_start: torch.Tensor, + route_start: torch.Tensor, + *, + num_recv_tokens: int, + block_m: int, + block_n: Optional[int] = None, + block_k: int = 64, + warps_m: Optional[int] = None, + warps_n: Optional[int] = None, + index_a_by_route_pos: bool = False, + activation: Optional[str] = None, + dispatched_probs: Optional[torch.Tensor] = None, + preact_out: Optional[torch.Tensor] = None, + gated_a: bool = False, +) -> None: + """Route-list gather-GEMM forward, writing the compact ``C[em_max, WIDTH_N]`` in place. + + ``A`` is ``[*, K]`` (received-token acts, gathered by ``sorted_slot_ids`` when + ``index_a_by_route_pos=False``, else read at the compact route row). ``B`` is + ``[num_experts, N_OUT, K]`` (contiguous inner ``K``). With ``activation`` set the fused + **gated** epilogue (FC1, ``gated_a=False``) applies ``act(gate) * up`` over ``N_OUT = 2F`` + into the ``F``-wide ``C``; ``dispatched_probs`` multiplies the per-route prob after the + activation, and ``preact_out`` (``[em_max, 2F]``) saves the raw ``[gate | up]`` + pre-activation. With ``gated_a=True`` (FC2) ``A`` is the raw ``[gate | up]`` pre-activation + (width ``2F``), the prologue applies ``act(gate) * up [* prob]`` into an ``F``-wide tile, + and the GEMM contracts over ``K = F`` against ``B[e, H, F]``. + """ + assert A.dtype == B.dtype == C.dtype == torch.bfloat16 + assert A.stride(1) == 1, "A must be contiguous along the contraction (K)" + # B is [E, N_OUT, K]: contiguous along K (fwd) or along N (dgrad transposed-weight view). + transpose_b = B.stride(2) != 1 + if transpose_b: + assert B.stride(1) == 1, "transposed B must be contiguous along N (dgrad view)" + assert activation is None, "dgrad (transposed B) does not support fused activation" + + gated = activation is not None and not gated_a + if gated_a: + if activation is None: + raise ValueError("gated_a requires activation ('silu' or 'gelu')") + if not index_a_by_route_pos: + raise ValueError("gated_a requires index_a_by_route_pos=True") + + # Plain GEMM: in-tree v3 kernels (pf_fwd / pf_dgrad). Fused epilogues need v2 (not ported). + _run_v3_fwd( + A, B, C, sorted_slot_ids, expert_ids, + num_recv_tokens=num_recv_tokens, block_m=block_m, + transpose_b=transpose_b, index_a_by_route_pos=index_a_by_route_pos, + gated=gated, gated_a=gated_a, + mul_prob=dispatched_probs is not None, save_preact=preact_out is not None, + ) + + +# Tile/warp configs the autotuner sweeps: (block_n, block_k, warps_m, warps_n). Fill path +# (DMA+swizzle, 3-buffer) is fixed at the env defaults above -- only tile geometry is tuned. +# +# Two shapes matter for the wide-M (block_m=256) MoE cases, where every block_k=128 and +# 256x64 entry below is filtered out by the LDS limit, leaving only 128x64 candidates: +# * a square-ish w4x4 warp grid, which spreads the cooperative A-fill (2 fills) / B-fill +# (1 fill) DMA and the LDS fragment reads more evenly than the tall w8x2 layout; +# * block_n=256 with block_k=32, which fits LDS at 3 buffers and raises arithmetic +# intensity to block_m*block_n/(2*(block_m+block_n)) = 64 MAC/byte vs 42.7 at block_n=128 +# (the ratio is independent of block_k, so widening N is what pays). +# Measured on FC1 no-act (qwen235b, block_m=256), idle machine, alias scopes on: +# 256x32 w4x4 ~1628us, 128x64 w4x4 ~1694us, 128x64 w8x2 ~1726us. +# +# An exhaustive sweep of the valid space (all block_n in 64..512 x block_k in 32..256 x warp +# grids >=4 waves; benchmarks/microbenchmarks/sweep_fc1_configs.py) found nothing better, so +# the list below is not missing a winner. Two directions are dead ends and are deliberately +# absent: block_n>=384 costs 4-22x (per-wave accumulator spill plus 120-144KB LDS pinning +# occupancy to 1 workgroup), and trading block_m down to reach block_k=128 -- 4x fewer +# barriers, which the ATT trace makes look attractive -- costs 59% (2581us at block_m=128 +# bk=128) because arithmetic intensity falls faster than barrier count. Note block_k>=128 does +# not fit LDS at all once NUM_BUF>=3, which the kernel enforces. +# +# Do not add 256x384x32 w1x4 or 512x32 w1x4 / w2x2: they abort the backend outright +# ("Bad machine code: Virtual register defs don't dominate all uses"), which would take the +# autotuner down with them rather than being skipped. Pre-existing, unrelated to alias scopes. +_FWD_TUNE_CONFIGS = [ + (64, 64, 2, 2), + (128, 64, 2, 2), + (128, 64, 4, 2), + (128, 64, 2, 4), + (128, 64, 4, 4), + (128, 64, 8, 2), + (256, 32, 4, 4), + (256, 32, 2, 4), + (256, 32, 4, 2), + (256, 64, 4, 2), + (128, 128, 2, 2), + (128, 128, 4, 2), + (64, 128, 2, 2), + (256, 64, 2, 4), +] + +# Winner cache: {shape/mode key -> (block_n, block_k, warps_m, warps_n)}. +_FWD_CACHE: dict = {} + + +def _valid_config(block_m, bn, bk, wm, wn, K, gated, gated_a=False, transpose_b=False): + if K % bk != 0 or bn % _mfma_dim(transpose_b) != 0: + return False + if not _warp_valid(block_m, bn, bk, wm, wn, transpose_b): + return False + return _lds_bytes(block_m, bn, bk, gated, gated_a=gated_a) <= _LDS_LIMIT + + +def flydsl_moe_fwd_pick_block_m( + A: torch.Tensor, + B: torch.Tensor, + *, + gated: bool = False, + gated_a: bool = False, + candidates=(256, 128), +) -> Optional[int]: + """Largest ``block_m`` in ``candidates`` the FlyDSL fwd can actually run for these operands + and epilogue mode, or ``None`` if the operands are unsupported at every candidate. + + "Can run" == bf16 operands contiguous along the contraction AND at least one autotuner tile + in :data:`_FWD_TUNE_CONFIGS` fits the per-workgroup LDS budget for the epilogue. The gated FC1 + epilogue stages a ``2F`` ``[gate|up]`` B-tile, so it only fits ``block_m <= 128``; the non-gated + FC1 fwd and the ``gated_a`` FC2 prologue fit ``block_m = 256``, which lifts the shared + fwd/dgrad/FC2 align onto the faster ``256x32`` MegaMOE-like tile. Callers should include their + token-count default among ``candidates`` (the picker only walks high->low over what is passed), + so a small-token workload is never padded up beyond what the caller offered. + """ + if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: + return None + if A.stride(1) != 1: # A must be contiguous along the contraction (K) + return None + if B.stride(2) != 1 and B.stride(1) != 1: # B contiguous along K (fwd) or N (dgrad view) + return None + K = int(B.shape[2]) + for block_m in sorted({int(c) for c in candidates}, reverse=True): + if any( + _valid_config(block_m, bn, bk, wm, wn, K, gated, gated_a) + for (bn, bk, wm, wn) in _FWD_TUNE_CONFIGS + ): + return block_m + return None + + +def flydsl_moe_fwd_autotuned( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + sorted_slot_ids: torch.Tensor, + expert_ids: torch.Tensor, + block_start: torch.Tensor, + route_start: torch.Tensor, + *, + num_recv_tokens: int, + block_m: int, + block_k: int = 64, + index_a_by_route_pos: bool = False, + activation: Optional[str] = None, + dispatched_probs: Optional[torch.Tensor] = None, + preact_out: Optional[torch.Tensor] = None, + gated_a: bool = False, + warmup: int = 3, + iters: int = 10, +) -> None: + """Shape-autotuned :func:`flydsl_moe_fwd`. + + On the first call for a given (block_m, GEMM shape, epilogue mode) the valid subset of + ``_FWD_TUNE_CONFIGS`` is benchmarked and the fastest ``(block_n, block_k, warps_m, warps_n)`` + is cached. The production DMA+swizzle fill path is always used; only tile geometry is swept. + """ + gated = activation is not None and not gated_a + N_OUT, K = int(B.shape[1]), int(B.shape[2]) + width_n = int(C.shape[1]) + key = ( + int(block_m), N_OUT, K, width_n, bool(gated), bool(gated_a), + activation, dispatched_probs is not None, preact_out is not None, + bool(index_a_by_route_pos), + ) + + def _launch(bn, bk, wm, wn): + flydsl_moe_fwd( + A, B, C, sorted_slot_ids, expert_ids, block_start, route_start, + num_recv_tokens=num_recv_tokens, block_m=block_m, block_n=bn, block_k=bk, + warps_m=wm, warps_n=wn, index_a_by_route_pos=index_a_by_route_pos, + activation=activation, dispatched_probs=dispatched_probs, preact_out=preact_out, + gated_a=gated_a, + ) + + best = _FWD_CACHE.get(key) + if best is None: + candidates = [ + (bn, bk, wm, wn) + for (bn, bk, wm, wn) in _FWD_TUNE_CONFIGS + if _valid_config(block_m, bn, bk, wm, wn, K, gated, gated_a) + ] + if not candidates: + _launch(None, block_k, None, None) # heuristic fallback + return + best_t = None + for cfg in candidates: + try: + for _ in range(warmup): + _launch(*cfg) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + start.record() + for _ in range(iters): + _launch(*cfg) + end.record() + torch.cuda.synchronize() + t = start.elapsed_time(end) / iters + except Exception: # noqa: BLE001 -- skip configs that fail to compile/run + continue + if best_t is None or t < best_t: + best_t, best = t, cfg + if best is None: + _launch(None, block_k, None, None) + return + _FWD_CACHE[key] = best + + _launch(*best) diff --git a/transformer_engine/pytorch/moe/pf_helper_kernels.py b/transformer_engine/pytorch/moe/pf_helper_kernels.py new file mode 100644 index 000000000..c01984654 --- /dev/null +++ b/transformer_engine/pytorch/moe/pf_helper_kernels.py @@ -0,0 +1,825 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Route-list MoE Triton helpers (bf16): align/scan, gated-act, gather-combine. + +Grouped GEMM for FC1/FC2 forward, backward (dgrad), and wgrad is FlyDSL-only +(``pf_fwd_wrapper``, ``pf_wgrad_wrapper``). This module retains Triton kernels for +routing metadata construction, gated-activation recompute/bwd, and the token-order +gather-combine pass that follows the compact route-list GEMM outputs. +""" + +from __future__ import annotations + +from typing import Any, Dict, Optional + +import torch +import triton +import triton.language as tl + + +# --- Fused gated-activation epilogue helpers (exp2-based, no libdevice) --- +@triton.jit +def _tanh(x): + # tanh via exp2 (1.44269504089 = 1/ln2): tanh(x) = 2*sigmoid(2x) - 1. + return 2.0 / (1.0 + tl.exp2(-2.0 * x * 1.44269504089)) - 1.0 + + +@triton.jit +def _silu(x): + # x * sigmoid(x), sigmoid via exp2. + return x / (1.0 + tl.exp2(-(x * 1.44269504089))) + + +@triton.jit +def _gelu_tanh(x): + # tanh approximation of GELU (matches PyTorch approximate='tanh'). + inner = 0.7978845608028654 * (x + 0.044715 * x * x * x) + return 0.5 * x * (1.0 + _tanh(inner)) + + +# Activation selector for the fused epilogue. Keep the set small and explicit; add a helper +# above and an entry here to extend. Passed to the kernel as a compile-time int so each +# activation specializes to its own instance (no runtime branch in the hot path). +ACT_SILU: tl.constexpr = 0 +ACT_GELU: tl.constexpr = 1 +_ACT_IDS = {"silu": 0, "gelu": 1} + + +@triton.jit +def _apply_activation(x, ACTIVATION: tl.constexpr): + if ACTIVATION == 0: + return _silu(x) + else: + return _gelu_tanh(x) + + +# --- Activation derivatives (for the fused gated-activation backward) --- +@triton.jit +def _silu_grad(x): + # d/dx [x*sigmoid(x)] = sigmoid(x) * (1 + x*(1 - sigmoid(x))). + s = 1.0 / (1.0 + tl.exp2(-(x * 1.44269504089))) + return s + x * s * (1.0 - s) + + +@triton.jit +def _gelu_tanh_grad(x): + # d/dx [0.5*x*(1+tanh(inner))], inner = c*(x + 0.044715 x^3), c = sqrt(2/pi). + inner = 0.7978845608028654 * (x + 0.044715 * x * x * x) + t = _tanh(inner) + dinner = 0.7978845608028654 * (1.0 + 3.0 * 0.044715 * x * x) + return 0.5 * (1.0 + t) + 0.5 * x * (1.0 - t * t) * dinner + + +@triton.jit +def _apply_activation_grad(x, ACTIVATION: tl.constexpr): + if ACTIVATION == 0: + return _silu_grad(x) + else: + return _gelu_tanh_grad(x) + + +# --- Fused activation + derivative (for the gated-activation backward) --- +# The backward needs *both* act(x) and act'(x) on the same input. Computing them via +# the separate helpers above evaluates the transcendental twice (sigmoid for silu, +# tanh for gelu). These return the pair from a single transcendental -- the backward +# kernel is VALU-bound on exp2, so sharing it is the dominant win. +@triton.jit +def _silu_and_grad(x): + # s = sigmoid(x); silu = x*s; d/dx silu = s + x*s*(1-s). + s = 1.0 / (1.0 + tl.exp2(-(x * 1.44269504089))) + act = x * s + return act, s + act * (1.0 - s) + + +@triton.jit +def _gelu_tanh_and_grad(x): + inner = 0.7978845608028654 * (x + 0.044715 * x * x * x) + t = _tanh(inner) + dinner = 0.7978845608028654 * (1.0 + 3.0 * 0.044715 * x * x) + return 0.5 * x * (1.0 + t), 0.5 * (1.0 + t) + 0.5 * x * (1.0 - t * t) * dinner + + +@triton.jit +def _apply_activation_and_grad(x, ACTIVATION: tl.constexpr): + if ACTIVATION == 0: + return _silu_and_grad(x) + else: + return _gelu_tanh_and_grad(x) + + +def _gated_act_bwd_autotune_configs() -> list: + """Tile/warp configs for the gated-act-bwd autotuner. + + Each program owns a ``BLOCK_M`` x ``BLOCK_H`` route/feature tile. The kernel is + latency-bound (the per-route work is tiny and the memory unit is never stalled), so + the sweep brackets the route-tile height ``BLOCK_M`` (more routes/block = more loads + in flight to hide latency) against the feature width ``BLOCK_H`` and the warp count, + trading memory-level parallelism against VGPR/occupancy. + """ + return [ + triton.Config({"BLOCK_M": bm, "BLOCK_H": bh}, num_warps=w) + for bm, bh, w in ( + # ``num_warps`` such that BLOCK_H/(warps*64) >= 8 gives 128-bit (dwordx4) + # global loads/stores; the (1,1024,2) / (2,1024,2) points below hit that, + # while the wider-warp points keep dwordx2 -- the tuner picks per shape. + (1, 1024, 4), + (1, 1024, 2), + (2, 1024, 2), + (2, 1024, 4), + (2, 512, 4), + (4, 1024, 8), + (8, 512, 4), + (16, 256, 4), + (16, 512, 8), + (32, 256, 8), + (32, 128, 4), + (8, 256, 4), + (4, 512, 4), + (64, 128, 8), + ) + ] + + +@triton.autotune(configs=_gated_act_bwd_autotune_configs(), key=["F", "HAS_PROBS"]) +@triton.jit +def _gated_act_prob_bwd_kernel( + grad_out_ptr, # [T * min(topk, E), F] grad wrt the fused FC1 output (= act(g)*u*prob) + preact_ptr, # [T * min(topk, E), 2F] raw pre-activation [gate | up] + probs_ptr, # [num_recv_tokens, E] + token_ptr, # [routes_max] route -> received-token row + expert_ptr, # [routes_max] route -> local expert + dpre_ptr, # [T * min(topk, E), 2F] out: grad wrt the raw 2F GEMM output + grad_probs_ptr, # [num_recv_tokens, E] out (fp32) + nbound_ptr, # [1] int32 device scalar: dynamic upper bound on compact routes + num_recv_tokens, + F, + stride_gom, + stride_goh, + stride_prem, + stride_preh, + stride_pm, + stride_pe, + stride_dprem, + stride_dpreh, + stride_gpm, + stride_gpe, + ACTIVATION: tl.constexpr, + HAS_PROBS: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """Backward of the fused gated-activation (+ optional route-prob) FC1 epilogue. + + Each program owns a ``BLOCK_M`` x ``BLOCK_H`` tile of routes x gate features. Given + ``grad_out`` (F), the saved pre-activation ``[gate | up]`` (2F) and the per-route prob + ``p``, it emits the 2F grad wrt the raw GEMM output (``dpre = [d_gate | d_up]``) and, + when probs are fused, reduces ``dp = sum_f grad_out * act(gate) * up`` per route and + scatters it to ``grad_dispatched_probs[token, expert]``. + + ``act(gate)`` and ``act'(gate)`` are produced together from a single transcendental + (the kernel is VALU-heavy on ``exp2``). Each ``(token, expert)`` cell is written by + exactly one route, so the prob-grad scatter is a plain store (no atomics). Padded / + over-allocated routes carry the ``token == num_recv_tokens`` sentinel and are masked + out of the prob load/store; their ``dpre`` rows are inert (ignored downstream). + + The compact route buffers are statically over-allocated to the worst-case + ``routes_max = T * topk`` (sync-free shape bound), but the real routes occupy only the + dense head ``[0, num_routes)`` -- under expert parallelism this can be ~topk*E_local/E + times smaller. ``nbound_ptr`` carries the actual (block-padded) route extent as a device + scalar so tail programs beyond it exit before touching HBM, instead of grinding through + the padding at full memory bandwidth. + """ + pid = tl.program_id(axis=0) + num_routes_bound = tl.load(nbound_ptr) + if pid * BLOCK_M >= num_routes_bound: + return + r_offs = pid * BLOCK_M + tl.arange(0, BLOCK_M) + r_mask = r_offs < num_routes_bound + token = tl.load(token_ptr + r_offs, mask=r_mask, other=num_recv_tokens).to(tl.int64) + valid = token < num_recv_tokens + expert = tl.load(expert_ptr + r_offs, mask=r_mask, other=0).to(tl.int64) + if HAS_PROBS: + prob = tl.load( + probs_ptr + token * stride_pm + expert * stride_pe, mask=valid, other=0.0 + ).to(tl.float32) + + dp_acc = tl.zeros((BLOCK_M,), dtype=tl.float32) + for h0 in range(0, F, BLOCK_H): + offs = h0 + tl.arange(0, BLOCK_H) + hmask = offs < F + m = r_mask[:, None] & hmask[None, :] + prow = r_offs[:, None] * stride_prem + g = tl.load( + preact_ptr + prow + offs[None, :] * stride_preh, mask=m, other=0.0 + ).to(tl.float32) + u = tl.load( + preact_ptr + prow + (offs[None, :] + F) * stride_preh, mask=m, other=0.0 + ).to(tl.float32) + go = tl.load( + grad_out_ptr + r_offs[:, None] * stride_gom + offs[None, :] * stride_goh, + mask=m, + other=0.0, + ).to(tl.float32) + + act_g, dact_g = _apply_activation_and_grad(g, ACTIVATION) + if HAS_PROBS: + dp_acc += tl.sum(go * act_g * u, axis=1) + da = go * prob[:, None] + else: + da = go + d_up = da * act_g + d_gate = da * u * dact_g + drow = r_offs[:, None] * stride_dprem + tl.store( + dpre_ptr + drow + offs[None, :] * stride_dpreh, + d_gate.to(dpre_ptr.dtype.element_ty), + mask=m, + ) + tl.store( + dpre_ptr + drow + (offs[None, :] + F) * stride_dpreh, + d_up.to(dpre_ptr.dtype.element_ty), + mask=m, + ) + + if HAS_PROBS: + tl.store( + grad_probs_ptr + token * stride_gpm + expert * stride_gpe, + dp_acc, + mask=valid, + ) + + +def fused_gated_act_prob_bwd( + grad_out: torch.Tensor, + preact: torch.Tensor, + token: torch.Tensor, + expert: torch.Tensor, + *, + num_recv_tokens: int, + activation: str, + dispatched_probs: Optional[torch.Tensor] = None, + grad_probs_shape: Optional[torch.Size] = None, + num_routes_bound: Optional[torch.Tensor] = None, +): + """Backward of the fused gated-activation (+ route-prob) FC1 epilogue. + + Parameters + ---------- + grad_out: + ``[T * min(topk, E), F]`` grad wrt the fused FC1 output (route/padded layout). + preact: + ``[T * min(topk, E), 2F]`` raw pre-activation ``[gate | up]`` saved by the forward. + token, expert: + ``[routes_max]`` per-route received-token row / local-expert id (int32). + dispatched_probs: + ``[num_recv_tokens, E]`` gating probs, or ``None`` if the forward did not fuse the + route-prob multiply. When given, its gradient is returned. + num_routes_bound: + Optional ``[1]`` int32 device scalar giving a (block-padded) upper bound on the number + of *real* compact routes. The route buffers are statically sized to the worst case + ``routes_max = T * topk``, but under expert parallelism only a small dense head is + populated; passing the actual extent (e.g. ``num_tokens_post_padded`` from the routing + metadata) lets tail programs exit early instead of streaming the padding through HBM. + When ``None`` the full static ``routes_max`` is used (no early exit). + + Returns + ------- + (dpre, grad_probs) + ``dpre`` is ``[T * min(topk, E), 2F]`` (bf16), grad wrt the raw GEMM output; ``grad_probs`` is + ``[num_recv_tokens, E]`` (matching ``dispatched_probs.dtype``) or ``None``. + """ + em_max, F = grad_out.shape + if preact.shape[0] != em_max or preact.shape[1] != 2 * F: + raise ValueError( + f"preact must be [T * min(topk, E), 2F]={ (em_max, 2 * F) }, got {tuple(preact.shape)}." + ) + routes_max = int(token.shape[0]) + has_probs = dispatched_probs is not None + + dpre = torch.empty((em_max, 2 * F), dtype=torch.bfloat16, device=grad_out.device) + if has_probs: + grad_probs = torch.zeros(dispatched_probs.shape, dtype=torch.float32, device=grad_out.device) + probs_ptr = dispatched_probs + stride_pm, stride_pe = dispatched_probs.stride(0), dispatched_probs.stride(1) + stride_gpm, stride_gpe = grad_probs.stride(0), grad_probs.stride(1) + else: + grad_probs = None + probs_ptr = grad_out # unused + stride_pm = stride_pe = stride_gpm = stride_gpe = 0 + + act_id = _ACT_IDS[activation] + if num_routes_bound is None: + nbound = torch.tensor([routes_max], dtype=torch.int32, device=grad_out.device) + else: + nbound = num_routes_bound + grid = lambda meta: (triton.cdiv(routes_max, meta["BLOCK_M"]),) # noqa: E731 + _gated_act_prob_bwd_kernel[grid]( + grad_out, + preact, + probs_ptr, + token, + expert, + dpre, + grad_probs if has_probs else grad_out, + nbound, + num_recv_tokens, + F, + grad_out.stride(0), + grad_out.stride(1), + preact.stride(0), + preact.stride(1), + stride_pm, + stride_pe, + dpre.stride(0), + dpre.stride(1), + stride_gpm, + stride_gpe, + ACTIVATION=act_id, + HAS_PROBS=has_probs, + ) + if has_probs: + grad_probs = grad_probs.to(dispatched_probs.dtype) + return dpre, grad_probs + + +@triton.autotune(configs=_gated_act_bwd_autotune_configs(), key=["F", "HAS_PROBS"]) +@triton.jit +def _gated_act_prob_fwd_kernel( + preact_ptr, # [routes_max, 2F] raw pre-activation [gate | up] + probs_ptr, # [num_recv_tokens, E] + token_ptr, # [routes_max] route -> received-token row + expert_ptr, # [routes_max] route -> local expert + act_ptr, # [routes_max, F] out: act(gate) * up * prob + nbound_ptr, # [1] int32 device scalar: dynamic upper bound on compact routes + num_recv_tokens, + F, + stride_prem, + stride_preh, + stride_pm, + stride_pe, + stride_am, + stride_ah, + ACTIVATION: tl.constexpr, + HAS_PROBS: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """Recompute the fused gated-activation FC1 output ``act(gate) * up * prob`` from preact. + + This is the *forward* counterpart of :func:`_gated_act_prob_bwd_kernel`: given the saved + ``2F`` pre-activation ``[gate | up]`` and the per-route prob it re-materialises the + ``F``-wide activation that the forward fused into the FC1 epilogue, so the backward can + checkpoint only the ``2F`` preact (never the ``F``-wide act) and rebuild it just-in-time + for the FC2 wgrad. Padded / over-allocated routes carry the ``token == num_recv_tokens`` + sentinel and get ``prob == 0`` (their act rows are ignored downstream). ``nbound_ptr`` + bounds tail programs to the real compact route extent (sync-free, EP-friendly early exit). + """ + pid = tl.program_id(axis=0) + num_routes_bound = tl.load(nbound_ptr) + if pid * BLOCK_M >= num_routes_bound: + return + r_offs = pid * BLOCK_M + tl.arange(0, BLOCK_M) + r_mask = r_offs < num_routes_bound + if HAS_PROBS: + token = tl.load(token_ptr + r_offs, mask=r_mask, other=num_recv_tokens).to(tl.int64) + valid = token < num_recv_tokens + expert = tl.load(expert_ptr + r_offs, mask=r_mask, other=0).to(tl.int64) + prob = tl.load( + probs_ptr + token * stride_pm + expert * stride_pe, mask=valid, other=0.0 + ).to(tl.float32) + + for h0 in range(0, F, BLOCK_H): + offs = h0 + tl.arange(0, BLOCK_H) + hmask = offs < F + m = r_mask[:, None] & hmask[None, :] + prow = r_offs[:, None] * stride_prem + g = tl.load( + preact_ptr + prow + offs[None, :] * stride_preh, mask=m, other=0.0 + ).to(tl.float32) + u = tl.load( + preact_ptr + prow + (offs[None, :] + F) * stride_preh, mask=m, other=0.0 + ).to(tl.float32) + act = _apply_activation(g, ACTIVATION) * u + if HAS_PROBS: + act = act * prob[:, None] + tl.store( + act_ptr + r_offs[:, None] * stride_am + offs[None, :] * stride_ah, + act.to(act_ptr.dtype.element_ty), + mask=m, + ) + + +def fused_gated_act_prob_fwd( + preact: torch.Tensor, + token: torch.Tensor, + expert: torch.Tensor, + *, + num_recv_tokens: int, + activation: str, + dispatched_probs: Optional[torch.Tensor] = None, + num_routes_bound: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Recompute the fused gated-activation FC1 output from the saved pre-activation. + + Re-materialises ``act = act(gate) * up * prob`` (``[routes_max, F]``, bf16) from the + ``[routes_max, 2F]`` preact ``[gate | up]`` -- the inverse of checkpointing the ``F``-wide + activation. Feeding this transient buffer to the *unchanged* FC2 wgrad keeps that kernel at + full (stored-act) speed while only the ``2F`` preact is persisted across the fwd/bwd + boundary. See :func:`fused_gated_act_prob_bwd` for the argument contract (token / expert / + ``num_routes_bound`` are the same per-route routing arrays). + """ + routes_max, two_f = preact.shape + if two_f % 2 != 0: + raise ValueError(f"preact must be [routes_max, 2F], got a {two_f}-wide last dim.") + F = two_f // 2 + has_probs = dispatched_probs is not None + + act = torch.empty((routes_max, F), dtype=torch.bfloat16, device=preact.device) + if has_probs: + probs_ptr = dispatched_probs + stride_pm, stride_pe = dispatched_probs.stride(0), dispatched_probs.stride(1) + else: + probs_ptr = preact # unused + stride_pm = stride_pe = 0 + + act_id = _ACT_IDS[activation] + if num_routes_bound is None: + nbound = torch.tensor([routes_max], dtype=torch.int32, device=preact.device) + else: + nbound = num_routes_bound + grid = lambda meta: (triton.cdiv(routes_max, meta["BLOCK_M"]),) # noqa: E731 + _gated_act_prob_fwd_kernel[grid]( + preact, + probs_ptr, + token, + expert, + act, + nbound, + num_recv_tokens, + F, + preact.stride(0), + preact.stride(1), + stride_pm, + stride_pe, + act.stride(0), + act.stride(1), + ACTIVATION=act_id, + HAS_PROBS=has_probs, + ) + return act + + + +@triton.jit +def _gather_combine_kernel( + src_ptr, # compact [T * min(topk, E), N] (route order; valid rows [0, num_routes)) + token_routes_ptr, # [T, MAXK] int32: compact route positions per token + token_count_ptr, # [T] int32: number of routes for each token + out_ptr, # [T, N] out + N, + stride_sm, + stride_sn, + stride_om, + stride_on, + MAXK: tl.constexpr, + BLOCK_N: tl.constexpr, + compute_type: tl.constexpr, +): + t = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + cnt = tl.load(token_count_ptr + t) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + n_mask = offs_n < N + acc = tl.zeros((BLOCK_N,), dtype=tl.float32) + # Sum the token's route rows (expert-ascending -> deterministic). Columns >= cnt are + # unused padding and skipped, so no padded/garbage row is ever gathered. + for j in range(0, MAXK): + if j < cnt: + r = tl.load(token_routes_ptr + t * MAXK + j).to(tl.int64) + v = tl.load(src_ptr + r * stride_sm + offs_n * stride_sn, mask=n_mask, other=0.0) + acc += v.to(tl.float32) + tl.store(out_ptr + t * stride_om + offs_n * stride_on, acc.to(compute_type), mask=n_mask) + + +def route_gather_combine( + src: torch.Tensor, + token_routes: torch.Tensor, + token_route_count: torch.Tensor, + num_recv_tokens: int, + *, + out_dtype: torch.dtype = torch.bfloat16, + block_n: Optional[int] = None, +) -> torch.Tensor: + """Combine per-route rows into per-token rows via a contention-free gather. + + Parameters + ---------- + src: + Compact per-route buffer ``[T * min(topk, E), N]`` (route order; only ``[0, num_routes)`` valid). + token_routes / token_route_count: + Token->routes inverse map from the align place kernel (``[T, MAXK]`` route positions + and the per-token route count). + num_recv_tokens: + Number of output token rows ``T``. + block_n: + N-tile width (``BLOCK_N``). ``None`` (default) picks ``min(next_pow2(N), 4096)``: wide + tiles issue larger contiguous per-route loads and cut redundant index loads. + + Returns + ------- + torch.Tensor + ``[num_recv_tokens, N]`` in ``out_dtype`` -- each row is the fp32 sum of its token's + route rows, cast once. No atomics, no host sync. + """ + if src.dim() != 2: + raise ValueError(f"src must be [T * min(topk, E), N], got {tuple(src.shape)}.") + if not src.is_contiguous(): + src = src.contiguous() + n = src.shape[1] + if block_n is None: + block_n = min(triton.next_power_of_2(n), 4096) + maxk = int(token_routes.shape[1]) + out = torch.empty((num_recv_tokens, n), dtype=out_dtype, device=src.device) + compute_type = tl.bfloat16 if out_dtype == torch.bfloat16 else tl.float32 + grid = (num_recv_tokens, triton.cdiv(n, block_n)) + _gather_combine_kernel[grid]( + src, + token_routes, + token_route_count, + out, + n, + src.stride(0), + src.stride(1), + out.stride(0), + out.stride(1), + MAXK=maxk, + BLOCK_N=block_n, + compute_type=compute_type, + ) + return out + + +@triton.jit +def _counts_within_kernel( + routing_map_ptr, # [T, E] (bool/int8), True where token t feeds local expert e + within_ptr, # [E, T] int32 out: exclusive within-expert rank of each routed cell + counts_ptr, # [E] int32 out: routed-token count per expert + T, + stride_t, + stride_e, + BLOCK_T: tl.constexpr, +): + e = tl.program_id(axis=0) + offs = tl.arange(0, BLOCK_T) + mask = offs < T + vals = tl.load(routing_map_ptr + offs * stride_t + e * stride_e, mask=mask, other=0).to( + tl.int32 + ) + # Exclusive prefix sum over tokens => within-expert rank; total => per-expert count. + incl = tl.cumsum(vals, axis=0) + excl = incl - vals + tl.store(within_ptr + e * T + offs, excl, mask=mask) + tl.store(counts_ptr + e, tl.sum(vals, axis=0)) + + +@triton.jit +def _expert_meta_kernel( + counts_ptr, # [E] int32 + blocks_per_expert_ptr, # [E] int32 out + block_start_ptr, # [E] int32 out (block units) + route_start_ptr, # [E] int32 out (compact route units) + expert_ids_ptr, # [blocks_max] int32 out (expert owning each block, -1 past the end) + ntpp_ptr, # [1] int32 out: block-padded token extent + E, + blocks_max, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_E: tl.constexpr, + BLOCK_B: tl.constexpr, +): + offs_e = tl.arange(0, BLOCK_E) + mask_e = offs_e < E + counts = tl.load(counts_ptr + offs_e, mask=mask_e, other=0) + blocks_per_expert = (counts + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + cblocks = tl.cumsum(blocks_per_expert, axis=0) # inclusive prefix over experts + # block_start is over padded block token counts + block_start = cblocks - blocks_per_expert + # route_start is over raw token counts + route_start = tl.cumsum(counts, axis=0) - counts + total_blocks = tl.max(tl.where(mask_e, cblocks, 0), axis=0) + + pid = tl.program_id(axis=0) + if pid == 0: + tl.store(blocks_per_expert_ptr + offs_e, blocks_per_expert, mask=mask_e) + tl.store(block_start_ptr + offs_e, block_start, mask=mask_e) + tl.store(route_start_ptr + offs_e, route_start, mask=mask_e) + tl.store(ntpp_ptr, total_blocks * BLOCK_SIZE_M) + + # expert_ids[b] = #{e : cblocks[e] <= b} (== searchsorted(cblocks, b, right=True)), + # then -1 for blocks past the real extent. + offs_b = pid * BLOCK_B + tl.arange(0, BLOCK_B) + mask_b = offs_b < blocks_max + cblocks_valid = tl.where(mask_e, cblocks, (1 << 30)) + le = (cblocks_valid[None, :] <= offs_b[:, None]).to(tl.int32) + expert_ids = tl.sum(le, axis=1) + expert_ids = tl.where(offs_b < total_blocks, expert_ids, -1) + tl.store(expert_ids_ptr + offs_b, expert_ids, mask=mask_b) + + +@triton.jit +def _route_list_place_kernel( + routing_map_ptr, # [T, E] (bool/int8), True where token t feeds local expert e + within_ptr, # [E, T] int32, exclusive within-expert rank of each routed cell + block_start_ptr, # [E] int32: first block index of each expert (block units) + route_start_ptr, # [E] int32: first compact route index of each expert + sorted_slot_ids_ptr, # [T * min(topk, E)] int32, sentinel-init (T) + route_to_token_ptr, # [routes_max] int32, sentinel-init (T) + token_routes_ptr, # [T, MAXK] int32 out: token->route positions (only if BUILD_INVERSE) + token_count_ptr, # [T] int32 out: routes per token (only if BUILD_INVERSE) + T, + E, + stride_t, + stride_e, + BLOCK_SIZE_M: tl.constexpr, + BUILD_INVERSE: tl.constexpr, + MAXK: tl.constexpr, +): + # One row of the map per program; place the (few) routed cells deterministically. Because + # the same per-token expert scan already computes each routed cell's compact route position + # ``pos = route_start[e] + within[e, t]``, the token->routes inverse map (used by the + # contention-free gather-combine) is emitted here too when ``BUILD_INVERSE`` -- folding what + # was a second per-token kernel launch into this one (expert-ascending, no atomics). + t = tl.program_id(axis=0) + if t >= T: + return + j = tl.zeros((), dtype=tl.int32) + for e in range(0, E): + is_routed = tl.load(routing_map_ptr + t * stride_t + e * stride_e) + if is_routed != 0: + w = tl.load(within_ptr + e * T + t) + bs = tl.load(block_start_ptr + e) + rs = tl.load(route_start_ptr + e) + pos = rs + w + tl.store(sorted_slot_ids_ptr + bs * BLOCK_SIZE_M + w, t) + tl.store(route_to_token_ptr + pos, t) + if BUILD_INVERSE: + tl.store(token_routes_ptr + t * MAXK + j, pos) + j += 1 + if BUILD_INVERSE: + tl.store(token_count_ptr + t, j) + + +def route_list_scan( + routing_map: torch.Tensor, + *, + num_experts: int, +): + """Block-size-independent scan: per-expert token counts + within-expert ranks. + + Returned ``(counts [E], within [E, T])`` can be reused across multiple block sizes + (e.g. the fwd/dgrad ``BLOCK_SIZE_M`` and the wgrad ``CONTRACT_M``), so the scan is + only paid once per routing map. Pass them to :func:`route_list_align` via ``scan=``. + """ + device = routing_map.device + T = int(routing_map.size(0)) + E = int(num_experts) + within = torch.empty((E, T), dtype=torch.int32, device=device) + counts = torch.empty((E,), dtype=torch.int32, device=device) + _counts_within_kernel[(E,)]( + routing_map, + within, + counts, + T, + routing_map.stride(0), + routing_map.stride(1), + BLOCK_T=triton.next_power_of_2(max(T, 1)), + ) + return counts, within + + +def route_list_align( + routing_map: torch.Tensor, + *, + num_experts: int, + block_size: int, + scan=None, + topk: int | None = None, + build_inverse_map: bool = False, +): + """Sync-free fused build of the route-list align buffers. + + Parameters + ---------- + scan: + Optional ``(counts, within)`` from :func:`route_list_scan` for this ``routing_map``. + When supplied the block-independent scan kernel is skipped (shared across block + sizes); otherwise it is computed here. + topk: + Host-known upper bound on the number of experts any token routes to (the router + top-k). When provided, the static over-allocation bound is tightened from the dense + ``T * num_experts`` to ``T * min(topk, num_experts)`` -- still + sync-free, but shrinking the padded buffers by ``num_experts / topk``. + build_inverse_map: + When True, the place kernel also emits the token->routes inverse map (used by the + contention-free gather-combine in FC2 fwd / FC1 dgrad) in the same launch, so no + separate inverse-map kernel is needed. Block-independent, so build it on the fwd align + only. + + Returns + ------- + ``(sorted_slot_ids, expert_ids, num_tokens_post_padded, block_start, blocks_per_expert, + route_start, route_to_token, token_routes, token_route_count)`` -- index tensors + ``int32``; ``num_tokens_post_padded`` (``[1]``) is a device scalar. ``token_routes`` + (``[T, min(topk, E)]``) and ``token_route_count`` (``[T]``) are ``None`` unless + ``build_inverse_map``. + """ + if routing_map.dtype != torch.bool: + routing_map = routing_map.bool() + routing_map = routing_map.contiguous() + device = routing_map.device + T = int(routing_map.size(0)) + E = int(num_experts) + + # Static (sync-free) upper bounds from shapes only. Each token routes to at most + # ``min(topk, E)`` experts, so that tightens the dense ``T * E`` bound. + max_per_token = E if topk is None else min(int(topk), E) + routes_max = T * max_per_token + blocks_max = (routes_max + block_size - 1) // block_size + E + em_max = blocks_max * block_size + + # Per-expert count + exclusive within-expert rank (one program per expert). Reused + # across block sizes when the caller passes a precomputed scan. + if scan is None: + counts, within = route_list_scan(routing_map, num_experts=E) + else: + counts, within = scan + + # Per-expert placement metadata + per-block expert ids (single launch). + blocks_per_expert = torch.empty((E,), dtype=torch.int32, device=device) + block_start = torch.empty((E,), dtype=torch.int32, device=device) + route_start = torch.empty((E,), dtype=torch.int32, device=device) + expert_ids = torch.empty((blocks_max,), dtype=torch.int32, device=device) + num_tokens_post_padded = torch.empty((1,), dtype=torch.int32, device=device) + block_b = 256 + _expert_meta_kernel[(triton.cdiv(blocks_max, block_b),)]( + counts, + blocks_per_expert, + block_start, + route_start, + expert_ids, + num_tokens_post_padded, + E, + blocks_max, + BLOCK_SIZE_M=block_size, + BLOCK_E=triton.next_power_of_2(max(E, 1)), + BLOCK_B=block_b, + ) + + # Optional token->routes inverse map, emitted by the same place kernel. Width is the + # tightened per-token bound; zero-init so any unused tail column is a safe (in-range) + # index (the gather masks columns >= count anyway). + if build_inverse_map: + maxk = max(max_per_token, 1) + token_routes = torch.zeros((T, maxk), dtype=torch.int32, device=device) + token_route_count = torch.empty((T,), dtype=torch.int32, device=device) + else: + maxk = 1 + token_routes = torch.empty((1,), dtype=torch.int32, device=device) # unused stub + token_route_count = token_routes + + # Scatter each routed cell into its deterministic (block-padded / compact) slot. Both + # sentinel buffers are carved from a single fill (one launch); the place kernel then + # overwrites the routed slots in each contiguous view. + sentinel = torch.full((em_max + routes_max,), T, dtype=torch.int32, device=device) + sorted_slot_ids = sentinel[:em_max] + route_to_token = sentinel[em_max:] + _route_list_place_kernel[(T,)]( + routing_map, + within, + block_start, + route_start, + sorted_slot_ids, + route_to_token, + token_routes, + token_route_count, + T, + E, + routing_map.stride(0), + routing_map.stride(1), + BLOCK_SIZE_M=block_size, + BUILD_INVERSE=build_inverse_map, + MAXK=maxk, + ) + + return ( + sorted_slot_ids, + expert_ids, + num_tokens_post_padded, + block_start, + blocks_per_expert, + route_start, + route_to_token, + token_routes if build_inverse_map else None, + token_route_count if build_inverse_map else None, + ) diff --git a/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py b/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py new file mode 100644 index 000000000..a2302078a --- /dev/null +++ b/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""FlyDSL permute-free MoE weight-gradient (wgrad) op wrapper. + +Mirrors the route-list wgrad contract so the same routing metadata +(``sorted_slot_ids`` holding the received-token row per slot, plus ``block_start`` / +``blocks_per_expert`` / ``route_start``) can be reused verbatim. + +Fixed kernel configuration (matches ``pf_wgrad.py``): FC1 route-list wgrad with compact +``grad`` + token-gathered ``x``, bf16 into ``dw`` (overwrite or accumulate), DMA + XOR +chunk swizzle fill, 3-stage LDS pipeline. Only the workgroup tile geometry is selectable +/ autotuned. +""" + +from __future__ import annotations + +import torch + +from ..flydsl_kernels.permute_free_grouped_gemm.pf_wgrad import ( + WGRAD_BLOCK_M, + compile_moe_wgrad_v2, +) +from ..flydsl_kernels.tensor_shim import _run_compiled, ptr_arg + +__all__ = ["flydsl_moe_wgrad", "flydsl_moe_wgrad_autotuned", "WGRAD_BLOCK_M"] + + +def flydsl_moe_wgrad( + x: torch.Tensor, + grad: torch.Tensor, + dw: torch.Tensor, + sorted_slot_ids: torch.Tensor, + block_start: torch.Tensor, + blocks_per_expert: torch.Tensor, + route_start: torch.Tensor, + *, + num_recv_tokens: int, + block_n: int = 128, + block_k: int = 128, + warps_n: int = 2, + warps_k: int = 2, + accumulate: bool = False, +) -> None: + """Compute FC1 grouped wgrad ``grad[route]^T @ x[token(route)]`` into ``dw``, per expert. + + See the permute-free wgrad API for the argument contract: ``x`` is + ``[num_recv_tokens, K]`` (gathered by received-token row), ``grad`` is the compact + ``[num_routes, N]`` per-route gradient, ``sorted_slot_ids`` maps each block-padded + route slot to its received-token row (sentinel ``num_recv_tokens`` for padding), and + ``route_start[e]`` is the compact first-route index of expert ``e``. + + ``block_n``/``block_k`` and ``warps_n``/``warps_k`` select the workgroup tile; the + defaults (``128x128`` over ``2x2`` warps) are a strong general config on CDNA4. + """ + num_experts, N, K = dw.shape + + assert x.dtype == grad.dtype == torch.bfloat16 + assert dw.dtype == torch.bfloat16 + assert x.is_contiguous() and grad.is_contiguous() and dw.is_contiguous() + + exe = compile_moe_wgrad_v2( + block_n=int(block_n), + block_k=int(block_k), + warps_n=int(warps_n), + warps_k=int(warps_k), + accumulate=bool(accumulate), + ) + + _run_compiled( + exe, + ptr_arg(dw), + ptr_arg(x), + ptr_arg(grad), + ptr_arg(sorted_slot_ids), + ptr_arg(block_start), + ptr_arg(blocks_per_expert), + ptr_arg(route_start), + int(N), + int(K), + int(num_recv_tokens), + int(num_experts), + torch.cuda.current_stream(), + ) + + +# Tile configs the autotuner sweeps (block_n, block_k, warps_n, warps_k). +_AUTOTUNE_TILES = [ + (64, 64, 1, 1), + (128, 64, 1, 1), + (128, 128, 2, 2), + (128, 128, 4, 2), + (256, 128, 4, 2), + (256, 256, 4, 4), + (128, 256, 2, 4), + (256, 128, 2, 2), + (256, 256, 2, 4), + (256, 256, 4, 2), +] + +_wgrad_autotuner = None + + +def _wgrad_run( + dw, + x, + grad, + sorted_slot_ids, + block_start, + blocks_per_expert, + route_start, + N, + K, + num_recv_tokens, + num_experts, + block_n=128, + block_k=128, + warps_n=2, + warps_k=2, + accumulate=False, +): + """Dispatch target for the FlyDSL autotuner: compile (lru-cached) + launch one tile.""" + exe = compile_moe_wgrad_v2( + block_n=int(block_n), + block_k=int(block_k), + warps_n=int(warps_n), + warps_k=int(warps_k), + accumulate=bool(accumulate), + ) + _run_compiled( + exe, + ptr_arg(dw), + ptr_arg(x), + ptr_arg(grad), + ptr_arg(sorted_slot_ids), + ptr_arg(block_start), + ptr_arg(blocks_per_expert), + ptr_arg(route_start), + int(N), + int(K), + int(num_recv_tokens), + int(num_experts), + torch.cuda.current_stream(), + ) + + +def _get_autotuner(warmup=10, rep=30): + """Build the shape-keyed Autotuner lazily (one instance, disk-cached results).""" + global _wgrad_autotuner + if _wgrad_autotuner is None: + from flydsl.autotune import Autotuner, Config + + configs = [ + Config(block_n=bn, block_k=bk, warps_n=wn, warps_k=wk) + for (bn, bk, wn, wk) in _AUTOTUNE_TILES + ] + _wgrad_autotuner = Autotuner( + _wgrad_run, + configs, + key=["x", "N", "num_experts"], + warmup=warmup, + rep=rep, + ) + return _wgrad_autotuner + + +def _select_wgrad_config( + x, grad, sorted_slot_ids, block_start, blocks_per_expert, route_start, + N, K, num_recv_tokens, num_experts, +): + """Return the autotuned ``(block_n, block_k, warps_n, warps_k)`` for this problem.""" + tuner = _get_autotuner() + scratch = torch.empty(num_experts, N, K, device=x.device, dtype=torch.bfloat16) + args = ( + scratch, x, grad, sorted_slot_ids, block_start, blocks_per_expert, route_start, + int(N), int(K), int(num_recv_tokens), int(num_experts), + ) + key = tuner._make_key(args, {}) + if key not in tuner.cache: + tuner(*args) + cfg = tuner.cache[key].kwargs + return cfg["block_n"], cfg["block_k"], cfg["warps_n"], cfg["warps_k"] + + +def flydsl_moe_wgrad_autotuned( + x: torch.Tensor, + grad: torch.Tensor, + dw: torch.Tensor, + sorted_slot_ids: torch.Tensor, + block_start: torch.Tensor, + blocks_per_expert: torch.Tensor, + route_start: torch.Tensor, + *, + num_recv_tokens: int, + accumulate: bool = False, +) -> None: + """Shape-autotuned variant of :func:`flydsl_moe_wgrad`. + + First call for a given ``(x.shape, N, num_experts)`` benchmarks every tile in + ``_AUTOTUNE_TILES`` and caches the fastest (in-memory + on disk under + ``~/.flydsl/autotune/``). + """ + num_experts, N, K = dw.shape + + assert x.dtype == grad.dtype == torch.bfloat16 + assert dw.dtype == torch.bfloat16 + assert x.is_contiguous() and grad.is_contiguous() and dw.is_contiguous() + + block_n, block_k, warps_n, warps_k = _select_wgrad_config( + x, grad, sorted_slot_ids, block_start, blocks_per_expert, route_start, + N, K, num_recv_tokens, num_experts, + ) + flydsl_moe_wgrad( + x, + grad, + dw, + sorted_slot_ids, + block_start, + blocks_per_expert, + route_start, + num_recv_tokens=int(num_recv_tokens), + block_n=block_n, + block_k=block_k, + warps_n=warps_n, + warps_k=warps_k, + accumulate=bool(accumulate), + ) From 2383b10bbeb7407615929ac39fc2f612ac00862b Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Wed, 5 Aug 2026 19:38:49 +0000 Subject: [PATCH 34/43] Use block-padded slot layout throughout permute-free MoE grouped GEMM Replaces the compact [num_routes, ...] route indexing with a block-padded [em_max, ...] slot canonical layout across forward, dgrad, wgrad, activation, and routing metadata. This removes route_start/route_to_token from MoERoutingMetadata, updates the FlyDSL kernels and wrappers to index by padded slot, and fixes the GroupedLinear permute-free return tuple and reshape. Also adds a new benchmark comparing permute-free vs permute+grouped GEMM backends, unit tests for the route-list kernels, and Qwen3-235B/optional DSV3-GateUP grouped GEMM test cases. --- .../microbenchmarks/benchmark_grouped_gemm.py | 13 +- .../benchmark_perm_free_grouped_gemm.py | 546 ++++++++++++++ .../pytorch/test_perm_free_grouped_linear.py | 711 ++++++++++++++++++ .../permute_free_grouped_gemm/pf_wgrad.py | 39 +- .../pytorch/module/grouped_linear.py | 9 +- transformer_engine/pytorch/moe/moe_routing.py | 19 +- .../pytorch/moe/permute_free_grouped_gemm.py | 114 +-- .../pytorch/moe/pf_fwd_wrapper.py | 8 +- .../pytorch/moe/pf_helper_kernels.py | 56 +- .../pytorch/moe/pf_wgrad_wrapper.py | 33 +- 10 files changed, 1409 insertions(+), 139 deletions(-) create mode 100644 benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py create mode 100644 tests/pytorch/test_perm_free_grouped_linear.py diff --git a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_grouped_gemm.py index a576568a0..27f486900 100755 --- a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py +++ b/benchmarks/microbenchmarks/benchmark_grouped_gemm.py @@ -71,11 +71,18 @@ def _generate_moe_test_cases( return test_cases -def generate_deepseekv3_test_cases(): - # DSV3-GateUP hangs on some hardware; only benchmark DSV3-Down. +def generate_deepseekv3_test_cases(include_gateup: bool = False): + # DSV3-GateUP hangs on some hardware; only benchmark DSV3-Down by default. return _generate_moe_test_cases( "DSV3", n_routed_experts=256, moe_intermediate_size=2048, hidden_size=7168, - skip_shapes=["GateUP"], + skip_shapes=None if include_gateup else ["GateUP"], + ) + + +def generate_qwen3_235b_test_cases(): + # Qwen3-235B-A22B: 128 routed experts, top-8, moe_intermediate_size 1536, hidden 4096. + return _generate_moe_test_cases( + "Qwen3-235B", n_routed_experts=128, moe_intermediate_size=1536, hidden_size=4096 ) diff --git a/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py new file mode 100644 index 000000000..7cd991a97 --- /dev/null +++ b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py @@ -0,0 +1,546 @@ +#!/usr/bin/env python +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""Compare permute-free gather-GEMM vs permute + grouped GEMM backends on ROCm. + +Backends: + - permute_free: TE permute-free (FlyDSL) route-list gather-in-GEMM (align + GEMM), no permute + - hipblaslt: TE multistream hipBLASLt grouped GEMM after moe_permute + - ck: TE CK grouped GEMM after moe_permute (NVTE_USE_CUTLASS_GROUPED_GEMM=1) + - triton: AITER Triton grouped GEMM after moe_permute + +Run from this directory:: + + python benchmark_perm_free_grouped_gemm.py + python benchmark_perm_free_grouped_gemm.py --quick --csv +""" + +from __future__ import annotations + +import os +from typing import Callable, Dict, List, Tuple + +import torch +from torch.utils.cpp_extension import IS_HIP_EXTENSION + +from benchmark_grouped_gemm import ( + EP_SIZE_LIST, + GROUPED_GEMM_M_SIZE_LIST, + generate_deepseekv2_lite_test_cases, + generate_deepseekv2_test_cases, + generate_deepseekv3_test_cases, + generate_grok_v2_test_cases, + generate_qwen3_235b_test_cases, +) +from utils import compute_tflops, make_metric_record, make_parser, run_benchmarks, time_func + +DEFAULT_TOPK = 8 +BACKENDS = ("permute_free", "hipblaslt", "ck", "triton") +# ``permute_free_act`` is an opt-in, GateUP-only backend: the permute-free FC1 gather-GEMM +# with the gated SiLU activation (``silu(gate) * up``) fused into the epilogue. It is not in +# the default sweep (it only applies to gate+up shapes); request it via ``--backends``. +KNOWN_BACKENDS = BACKENDS + ("permute_free_act",) + +# Training phases benchmarked (and reported) separately in --train mode. +PHASES = ("fwd", "dgrad", "wgrad") + +# Short backend labels for the per-phase --train metric names. +_BACKEND_SHORT = { + "permute_free": "PermuteFree", + "permute_free_act": "PermuteFreeAct", + "hipblaslt": "hipBLASLt", + "ck": "CK", + "triton": "Triton", +} + + +def _require_rocm(): + if not IS_HIP_EXTENSION or not torch.cuda.is_available(): + raise RuntimeError("This benchmark requires ROCm (HIP extension) and a CUDA device.") + + +def _make_routing(num_tokens: int, num_experts: int, topk: int, device: str, seed: int): + from transformer_engine.pytorch.moe import MoERoutingMetadata + + gen = torch.Generator(device=device) + gen.manual_seed(seed) + logits = torch.randn(num_tokens, num_experts, device=device, generator=gen) + probs = torch.softmax(logits, dim=-1) + topk_weights, topk_ids = torch.topk(probs, k=topk, dim=-1) + topk_ids = topk_ids.to(torch.int32) + topk_weights = topk_weights.to(torch.float32) + # MoERoutingMetadata now takes a boolean routing_map as its primary input; build it + # from the sampled topk_ids so the permute-free and permute backends see identical + # routing. Passing ``topk`` tightens the block-padded over-allocation (em_max). topk_ids/ + # topk_weights are attached for the permute (moe_permute) backends. + routing_map = torch.zeros(num_tokens, num_experts, dtype=torch.bool, device=device) + routing_map.scatter_(1, topk_ids.to(torch.long), True) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts, topk=topk) + routing.topk_ids = topk_ids + routing.topk_weights = topk_weights + return routing + + +def _m_splits_from_topk(topk_ids: torch.Tensor, num_experts: int) -> List[int]: + counts = torch.bincount(topk_ids.reshape(-1).long(), minlength=num_experts) + return [int(v) for v in counts.tolist()] + + +def _permute_hidden( + hidden: torch.Tensor, topk_ids: torch.Tensor, num_tokens: int, topk: int +) -> torch.Tensor: + from transformer_engine.pytorch import moe_permute + + num_out_tokens = num_tokens * topk + permuted, _ = moe_permute(hidden, topk_ids, num_out_tokens, map_type="index") + return permuted + + +def _grouped_gemm_hip( + permuted: torch.Tensor, + weights: torch.Tensor, + m_splits: List[int], + *, + use_ck: bool, +) -> torch.Tensor: + from transformer_engine.pytorch.cpp_extensions import general_grouped_gemm + + prev = os.environ.get("NVTE_USE_CUTLASS_GROUPED_GEMM") + os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = "1" if use_ck else "0" + try: + b = len(m_splits) + n = int(weights.shape[1]) + sum_m = sum(m_splits) + xs = list(torch.split(permuted.view(sum_m, -1), m_splits)) + weight_list = [weights[i] for i in range(b)] + out = torch.empty((sum_m, n), device=permuted.device, dtype=permuted.dtype) + general_grouped_gemm( + A=weight_list, + B=xs, + out=[out], + quantization_params=[None] * b, + out_dtype=permuted.dtype, + single_output=True, + m_splits=m_splits, + use_bias=False, + bias=None, + layout="TN", + ) + return out + finally: + if prev is None: + os.environ.pop("NVTE_USE_CUTLASS_GROUPED_GEMM", None) + else: + os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = prev + + +def _grouped_gemm_triton( + permuted: torch.Tensor, + weights: torch.Tensor, + m_splits: List[int], +) -> torch.Tensor: + from transformer_engine.pytorch.triton_kernels.grouped_gemm import general_grouped_gemm_triton + + b = len(m_splits) + n = int(weights.shape[1]) + sum_m = sum(m_splits) + xs = list(torch.split(permuted.view(sum_m, -1), m_splits)) + weight_list = [weights[i] for i in range(b)] + out = torch.empty((sum_m, n), device=permuted.device, dtype=permuted.dtype) + general_grouped_gemm_triton( + A=weight_list, + B=xs, + out=[out], + quantization_params=[None] * b, + out_dtype=permuted.dtype, + single_output=True, + m_splits=m_splits, + use_bias=False, + bias=None, + layout="TN", + ) + return out + + +def _permute_free_gemm( + hidden: torch.Tensor, + weights: torch.Tensor, + routing, +) -> torch.Tensor: + from transformer_engine.pytorch.moe import ( + permute_free_grouped_gemm_bf16, + ) + + return permute_free_grouped_gemm_bf16(hidden, weights, routing) + + +def _permute_free_act_gemm( + hidden: torch.Tensor, + weights: torch.Tensor, + routing, +) -> torch.Tensor: + """Permute-free FC1 gather-GEMM with the gated SiLU activation fused into the epilogue. + + ``weights`` is the gate+up projection ``[E, 2F, K]``; the kernel emits the F-wide + activated buffer ``silu(gate) * up`` directly, skipping the separate activation pass. + Only valid for gate+up (GateUP) shapes -- ``out_features`` must be even. + """ + from transformer_engine.pytorch.moe import ( + permute_free_grouped_gemm_bf16, + ) + + return permute_free_grouped_gemm_bf16(hidden, weights, routing, activation="silu") + + +def _build_backend_fns( + hidden: torch.Tensor, + weights: torch.Tensor, + routing, + m_splits: List[int], + num_tokens: int, + topk: int, + is_gated: bool = False, +) -> Dict[str, Callable[[], torch.Tensor]]: + fns: Dict[str, Callable[[], torch.Tensor]] = {} + + def hipblaslt_fn(): + permuted = _permute_hidden(hidden, routing.topk_ids, num_tokens, topk) + return _grouped_gemm_hip(permuted, weights, m_splits, use_ck=False) + + def ck_fn(): + permuted = _permute_hidden(hidden, routing.topk_ids, num_tokens, topk) + return _grouped_gemm_hip(permuted, weights, m_splits, use_ck=True) + + def triton_fn(): + permuted = _permute_hidden(hidden, routing.topk_ids, num_tokens, topk) + return _grouped_gemm_triton(permuted, weights, m_splits) + + fns["hipblaslt"] = hipblaslt_fn + fns["ck"] = ck_fn + fns["triton"] = triton_fn + fns["permute_free"] = lambda: _permute_free_gemm(hidden, weights, routing) + # GateUP-only: permute-free FC1 with the fused gated SiLU epilogue. + if is_gated: + fns["permute_free_act"] = lambda: _permute_free_act_gemm(hidden, weights, routing) + return fns + + +def _build_train_phase_fns( + hidden: torch.Tensor, + weights: torch.Tensor, + routing, + m_splits: List[int], + num_tokens: int, + topk: int, + dtype: torch.dtype, +) -> Dict[str, Dict[str, Callable[[], torch.Tensor]]]: + """Per-phase (fwd / dgrad / wgrad) closures for each backend. + + Returns ``{backend: {phase: fn}}`` so each phase can be timed independently. + + - permute_free: gather fwd + gather dgrad + fused wgrad (FlyDSL / Triton), matching the + ``_GroupedLinear`` permute-free autograd path. Every kernel gathers from token-major + tensors, so no standalone permute/unpermute is charged. + - hipblaslt/ck: the traditional permuted grouped GEMMs. The token permute is charged to + ``fwd`` and the dgrad unpermute to ``dgrad`` (both produce a token-major-equivalent + result); ``wgrad`` reuses the operands already permuted in fwd, so it is GEMM-only -- + mirroring a real training iteration where the permute is paid once. (Triton GG is + excluded from --train: its grouped backend does not expose the NN/NT grad layouts.) + """ + from transformer_engine.pytorch import moe_permute, moe_unpermute + from transformer_engine.pytorch.cpp_extensions import general_grouped_gemm + from transformer_engine.pytorch.moe import ( + permute_free_grouped_gemm_bf16, + permute_free_grouped_gemm_bf16_dgrad, + permute_free_grouped_gemm_bf16_wgrad, + ) + + device = hidden.device + b = len(m_splits) + n = int(weights.shape[1]) + k = int(weights.shape[2]) + total_m = sum(m_splits) + weight_list = [weights[i] for i in range(b)] + weights_shape = (b, n, k) + + # Run the permute-free forward once (outside timing) to finalize the block-padded align on + # ``routing`` (the fwd enforces a v3 block-size floor, which fixes em_max) and to learn the + # padded slot extent. The dgrad/wgrad now consume the block-padded ``[em_max, out_features]`` + # slot gradient -- i.e. the gradient of this fwd output -- not a compact ``[num_routes]`` one. + em_max = int(permute_free_grouped_gemm_bf16(hidden, weights, routing).shape[0]) + + # Upstream gradients. ``grad_out_pf`` is the block-padded slot grad for the permute-free + # dgrad/wgrad; the traditional path keeps its compact expert-sorted ``[total_m, n]`` grad. + # Allocated once, outside the timed region. + grad_out_pf = torch.randn(em_max, n, dtype=dtype, device=device) + grad_out_perm = torch.randn(total_m, n, dtype=dtype, device=device) + grad_splits = list(torch.split(grad_out_perm, m_splits)) + + # Pre-permuted activations reused by the traditional dgrad/wgrad phases. The permute + # itself is (re)charged inside the traditional fwd phase below. + permuted_hidden, row_id_map = moe_permute( + hidden, routing.topk_ids, total_m, map_type="index" + ) + xs_perm = list(torch.split(permuted_hidden.view(total_m, -1), m_splits)) + + # --- permute-free phases (gather-in-GEMM; token-major in/out) --- + def pf_fwd(): + return permute_free_grouped_gemm_bf16(hidden, weights, routing) + + def pf_dgrad(): + return permute_free_grouped_gemm_bf16_dgrad(grad_out_pf, weights, routing) + + def pf_wgrad(): + return permute_free_grouped_gemm_bf16_wgrad( + hidden, grad_out_pf, weights_shape, routing + ) + + # --- traditional (permute + grouped GEMM) phases --- + def _with_ck_env(use_ck: bool, fn: Callable[[], torch.Tensor]): + def run(): + prev = os.environ.get("NVTE_USE_CUTLASS_GROUPED_GEMM") + os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = "1" if use_ck else "0" + try: + return fn() + finally: + if prev is None: + os.environ.pop("NVTE_USE_CUTLASS_GROUPED_GEMM", None) + else: + os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = prev + + return run + + def _trad_fwd(): + permuted, _ = moe_permute(hidden, routing.topk_ids, total_m, map_type="index") + xs = list(torch.split(permuted.view(total_m, -1), m_splits)) + out = torch.empty((total_m, n), device=device, dtype=dtype) + general_grouped_gemm( + A=weight_list, B=xs, out=[out], quantization_params=[None] * b, + out_dtype=dtype, single_output=True, m_splits=m_splits, + use_bias=False, bias=None, layout="TN", + ) + return out + + def _trad_dgrad(): + dx_buf = torch.empty((total_m, k), device=device, dtype=dtype) + dxs = list(torch.split(dx_buf, m_splits)) + general_grouped_gemm( + A=weight_list, B=grad_splits, out=dxs, quantization_params=[None] * b, + out_dtype=dtype, single_output=False, m_splits=m_splits, grad=False, + use_bias=False, bias=None, layout="NN", + ) + return moe_unpermute(dx_buf, row_id_map, map_type="index") + + def _trad_wgrad(): + dw = torch.empty((b, n, k), device=device, dtype=dtype) + dws = [dw[i] for i in range(b)] + general_grouped_gemm( + A=xs_perm, B=grad_splits, out=dws, quantization_params=[None] * b, + out_dtype=dtype, single_output=False, m_splits=m_splits, grad=False, + use_bias=False, bias=None, layout="NT", + ) + return dw + + def _trad_phase_fns(use_ck: bool) -> Dict[str, Callable[[], torch.Tensor]]: + return { + "fwd": _with_ck_env(use_ck, _trad_fwd), + "dgrad": _with_ck_env(use_ck, _trad_dgrad), + "wgrad": _with_ck_env(use_ck, _trad_wgrad), + } + + return { + "permute_free": {"fwd": pf_fwd, "dgrad": pf_dgrad, "wgrad": pf_wgrad}, + "hipblaslt": _trad_phase_fns(use_ck=False), + "ck": _trad_phase_fns(use_ck=True), + } + + +def bench_moe_grouped_gemm_backends( + Case: str, + B: int, + M: int, + N: int, + K: int, + dtype: torch.dtype, + topk: int = DEFAULT_TOPK, + seed: int = 0, + backends: Tuple[str, ...] = BACKENDS, + train: bool = False, +): + _require_rocm() + device = "cuda" + num_tokens = M + num_experts = B + + effective_topk = min(topk, num_experts) + if effective_topk < 1: + raise ValueError(f"Need at least one expert (B={num_experts}).") + + hidden = torch.randn(num_tokens, K, dtype=dtype, device=device) + weights = torch.randn(num_experts, N, K, dtype=dtype, device=device) + routing = _make_routing(num_tokens, num_experts, effective_topk, device, seed) + m_splits = _m_splits_from_topk(routing.topk_ids, num_experts) + + total_m = num_tokens * effective_topk + + if train: + # Per-phase analysis: time fwd / dgrad / wgrad independently for each backend. + phase_fns = _build_train_phase_fns( + hidden, weights, routing, m_splits, num_tokens, effective_topk, dtype + ) + # Each of fwd / dgrad / wgrad performs a ~2*total_m*N*K FLOP GEMM. + phase_flops = 2 * total_m * N * K + + # Warmup every selected phase (also builds the permute-free align buffers so the + # dgrad/wgrad timings reflect cache reuse, as in a real iteration). + for name in backends: + if name not in phase_fns: + continue + for phase in PHASES: + phase_fns[name][phase]() + torch.cuda.synchronize() + + # Emit records phase-major so the printout groups fwd, then dgrad, then wgrad. + records = [] + for phase in PHASES: + for name in backends: + if name not in phase_fns: + continue + ms, measurement = time_func(phase_fns[name][phase]) + records.append( + make_metric_record( + f"{phase} \u00b7 {_BACKEND_SHORT[name]}", + ms, + "TFLOPS", + compute_tflops(phase_flops, ms), + measurement=measurement, + ) + ) + return records + + # GateUP shapes carry a gate+up projection (out_features == 2F, even), which is the only + # case where the fused gated-activation permute-free backend applies. + is_gated = Case.endswith("GateUP") + backend_fns = _build_backend_fns( + hidden, weights, routing, m_splits, num_tokens, effective_topk, is_gated=is_gated + ) + fwd_flops = 2 * total_m * N * K + + # Permute-free backends emit a padded/route-major (or F-wide, for the activation fusion) + # buffer, so their output shape differs from the permuted [total_m, N] result. + _permute_free_names = ("permute_free", "permute_free_act") + + # Warmup + correctness spot-check (permute-free vs hipblaslt layout differs; skip allclose) + for name in backends: + if name not in backend_fns: + continue + out = backend_fns[name]() + torch.cuda.synchronize() + if name not in _permute_free_names: + assert out.shape == (total_m, N), f"{name}: bad shape {out.shape}" + + records = [] + for name in backends: + if name not in backend_fns: + continue + label = { + "permute_free": "PermuteFree Gather-GEMM", + "permute_free_act": "PermuteFree Gather-GEMM + SiLU", + "hipblaslt": "Permute+hipBLASLt Grouped GEMM", + "ck": "Permute+CK Grouped GEMM", + "triton": "Permute+AITER Triton Grouped GEMM", + }[name] + ms, measurement = time_func(backend_fns[name]) + records.append( + make_metric_record( + label, + ms, + "TFLOPS", + compute_tflops(fwd_flops, ms), + measurement=measurement, + ) + ) + return records + + +def _filter_test_cases(test_cases, quick: bool): + if not quick: + return test_cases + allowed_m = {512, 1024} + filtered = [c for c in test_cases if c["M"] in allowed_m and c["Case"].endswith("-Down")] + return filtered[:4] if filtered else test_cases[:2] + + +def main(): + base = make_parser(description="Benchmark MoE grouped GEMM backends on ROCm.") + base.add_argument("--quick", action="store_true", help="Run a small subset of shapes.") + base.add_argument( + "--train", + action="store_true", + help="Benchmark a full training iteration (fwd + dgrad + wgrad) instead of fwd only.", + ) + base.add_argument("--topk", type=int, default=DEFAULT_TOPK, help="MoE top-k routing.") + base.add_argument( + "--backends", + type=str, + default=",".join(BACKENDS), + help=f"Comma-separated backends (default: {','.join(BACKENDS)}).", + ) + base.add_argument( + "--case-prefix", + type=str, + default=None, + help="Only run cases whose Case name starts with this prefix (e.g. DSV3).", + ) + base.add_argument( + "--include-dsv3-gateup", + action="store_true", + help="Also benchmark DSV3-GateUP (skipped by default; known to run on gfx950).", + ) + args = base.parse_args() + + _require_rocm() + + test_cases = ( + generate_deepseekv2_lite_test_cases() + + generate_deepseekv2_test_cases() + + generate_deepseekv3_test_cases(include_gateup=args.include_dsv3_gateup) + + generate_grok_v2_test_cases() + + generate_qwen3_235b_test_cases() + ) + test_cases = _filter_test_cases(test_cases, args.quick) + if args.case_prefix: + prefix = args.case_prefix + test_cases = [ + c + for c in test_cases + if c["Case"].startswith(prefix) + and not (prefix == "DSV2-" and c["Case"].startswith("DSV2-Lite")) + ] + for case in test_cases: + case["topk"] = min(args.topk, case["B"]) + case["seed"] = 42 + + selected_backends = tuple(b.strip() for b in args.backends.split(",") if b.strip()) + for b in selected_backends: + if b not in KNOWN_BACKENDS: + raise ValueError(f"Unknown backend {b!r}, expected one of {KNOWN_BACKENDS}") + + def bench_fn(**case): + return bench_moe_grouped_gemm_backends( + **case, backends=selected_backends, train=args.train + ) + + run_benchmarks( + test_cases=test_cases, + bench_fn=bench_fn, + param_columns=["Case", "B", "M", "N", "K", "dtype", "topk"], + args=args, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/pytorch/test_perm_free_grouped_linear.py b/tests/pytorch/test_perm_free_grouped_linear.py new file mode 100644 index 000000000..36088af9f --- /dev/null +++ b/tests/pytorch/test_perm_free_grouped_linear.py @@ -0,0 +1,711 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for the route-list permute-free bf16 MoE gather-GEMM kernels.""" + +import pytest +import torch +from torch.utils.cpp_extension import IS_HIP_EXTENSION + +from transformer_engine.pytorch import GroupedLinear +from transformer_engine.pytorch.moe import ( + MoERoutingMetadata, + PermuteFreeMetadata, + get_default_moe_kernel_config, + permute_free_grouped_gemm_backward, + permute_free_grouped_gemm_bf16, + permute_free_grouped_gemm_bf16_dgrad, + permute_free_grouped_gemm_bf16_fc2_wgrad, + permute_free_grouped_gemm_bf16_wgrad, + prepare_moe_align, +) +from transformer_engine.pytorch.moe.permute_free_grouped_gemm import ( + _FLYDSL_FWD_BLOCK_M, + _expert_per_route, +) + +pytestmark = pytest.mark.skipif( + not (IS_HIP_EXTENSION and torch.cuda.is_available()), + reason="Permute-free grouped GEMM tests require ROCm and CUDA device.", +) + + +@pytest.fixture(autouse=True) +def _cuda_sync_after_test(): + yield + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def _rel_l2(actual: torch.Tensor, ref: torch.Tensor) -> float: + return ((actual.float() - ref.float()).norm() / ref.float().norm().clamp_min(1e-8)).item() + + +def _random_routing_map(num_recv_tokens, num_experts, max_hits, device, seed): + """Boolean [num_recv_tokens, num_experts] with 1..max_hits local experts per token.""" + gen = torch.Generator(device=device).manual_seed(seed) + routing_map = torch.zeros(num_recv_tokens, num_experts, dtype=torch.bool, device=device) + for t in range(num_recv_tokens): + k = int(torch.randint(1, max_hits + 1, (1,), device=device, generator=gen).item()) + experts = torch.randperm(num_experts, device=device, generator=gen)[:k] + routing_map[t, experts] = True + # Guarantee every expert owns at least one route so per-expert refs are exercised. + for e in range(num_experts): + if not routing_map[:, e].any(): + routing_map[int(e) % num_recv_tokens, e] = True + return routing_map + + +def _compact_route_order(routing_map): + """Expert-sorted (token, expert) route lists, matching ``moe_align_route_list``.""" + tok, exp = routing_map.nonzero(as_tuple=True) + order = torch.argsort(exp, stable=True) + return tok[order].to(torch.int64), exp[order].to(torch.int64) + +# --------------------------------------------------------------------------- +# Route-list gather-GEMM: fwd / dgrad / wgrad +# --------------------------------------------------------------------------- +def test_route_list_fwd(): + torch.manual_seed(7) + num_recv_tokens, in_features, out_features = 128, 128, 256 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=1) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + + hidden = torch.randn(num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16) + weights = torch.randn( + num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16 + ) + + out_full = permute_free_grouped_gemm_bf16(hidden, weights, routing) + + route_to_token, route_expert = _compact_route_order(routing_map) + num_routes = route_to_token.numel() + # Block-padded canonical layout: the output is [em_max, out] in expert-sorted, block-padded + # slot order. Padded slot s holds hidden[sorted_slot_ids[s]] @ W[expert(s)]^T for the valid + # slots (sorted_slot_ids[s] < num_recv_tokens); padding slots carry inert dead values and + # are dropped only at the final padded->token gather-combine. + em_max = int(routing.sorted_slot_ids.shape[0]) + assert out_full.shape == (em_max, out_features) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + ref = torch.einsum( + "rk,rnk->rn", + hidden[tok].float(), + weights[slot_expert.to(torch.int64)].float(), + ) + assert int(valid.sum().item()) == num_routes + assert _rel_l2(out_full[valid], ref[valid]) < 2e-2 + + +def test_route_list_dgrad(): + torch.manual_seed(11) + num_recv_tokens, in_features, out_features = 128, 96, 128 + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=2) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + num_routes = int(routing_map.sum().item()) + # Block-padded canonical: the incoming route grad lives in the same [em_max] block-padded + # slot order as the FC1 forward output (the FC2 dgrad hands it back in this layout). Prepare + # the align at the forward block size so the standalone dgrad reuses the same slot layout. + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + + weights = torch.randn( + num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16 + ) + grad = torch.zeros(em_max, out_features, device="cuda", dtype=torch.bfloat16) + grad[valid] = torch.randn( + int(valid.sum().item()), out_features, device="cuda", dtype=torch.bfloat16 + ) + + dA = permute_free_grouped_gemm_bf16_dgrad(grad, weights, routing) + assert dA.shape == (num_recv_tokens, in_features) + + # dA[t] = sum_{padded slots s with token==t, valid} grad[s] @ W[expert(s)] + per_slot = torch.einsum( + "rn,rnk->rk", grad.float(), weights[slot_expert.to(torch.int64)].float() + ) + per_slot = per_slot * valid[:, None].float() + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + ref = torch.zeros(num_recv_tokens, in_features, device="cuda", dtype=torch.float32) + ref.index_add_(0, tok, per_slot) + assert int(valid.sum().item()) == num_routes + assert _rel_l2(dA, ref) < 2e-2 + + +def test_route_list_wgrad(): + torch.manual_seed(17) + num_recv_tokens, in_features, out_features = 128, 96, 128 + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=3) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + num_routes = int(routing_map.sum().item()) + # Block-padded canonical: the incoming route grad lives in the [em_max] block-padded slot + # order (same layout the forward writes / FC2 dgrad hands back). Prepare the forward align so + # the wgrad reads grad at block_start[e]*block_size_m + within-rank. + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + + hidden = torch.randn(num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16) + weights_shape = (num_experts, out_features, in_features) + grad = torch.zeros(em_max, out_features, device="cuda", dtype=torch.bfloat16) + grad[valid] = torch.randn( + int(valid.sum().item()), out_features, device="cuda", dtype=torch.bfloat16 + ) + + dW = permute_free_grouped_gemm_bf16_wgrad(hidden, grad, weights_shape, routing) + assert dW.shape == weights_shape + + # dW[e] = sum_{valid slots s with expert==e} outer(grad[s], hidden[slot_token[s]]) + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + per_slot = torch.einsum("rn,rk->rnk", grad.float(), hidden[tok].float()) + per_slot = per_slot * valid[:, None, None].float() + ref = torch.zeros(num_experts, out_features, in_features, device="cuda", dtype=torch.float32) + # Past-end padding slots carry expert id -1; clamp for the scatter (their rows are zeroed). + ref.index_add_(0, slot_expert.to(torch.int64).clamp_min(0), per_slot) + assert int(valid.sum().item()) == num_routes + assert _rel_l2(dW, ref) < 2e-2 + + +def test_route_list_fc2_wgrad(): + """FC2 wgrad via operand-swap + transpose; ``out``/``accumulate`` fold in bf16 or fp32.""" + torch.manual_seed(19) + num_recv_tokens, in_features, out_features = 128, 96, 128 # F=in, H=out (W2 is [E, H, F]) + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=4) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + num_routes = int(routing_map.sum().item()) + # Block-padded canonical: fc2_input (the FC1 output) lives in the [em_max] block-padded slot + # order; grad_output stays token-space and is gathered internally. + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + + fc2_input = torch.zeros(em_max, in_features, device="cuda", dtype=torch.bfloat16) + fc2_input[valid] = torch.randn( + int(valid.sum().item()), in_features, device="cuda", dtype=torch.bfloat16 + ) + grad_output = torch.randn(num_recv_tokens, out_features, device="cuda", dtype=torch.bfloat16) + weights_shape = (num_experts, out_features, in_features) # [E, H, F] + + # dW2[e] = sum_{valid slots s with expert==e} outer(grad_output[slot_token[s]], fc2_input[s]) + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + per_slot = torch.einsum("rh,rf->rhf", grad_output[tok].float(), fc2_input.float()) + per_slot = per_slot * valid[:, None, None].float() + ref = torch.zeros(num_experts, out_features, in_features, device="cuda", dtype=torch.float32) + # Past-end padding slots carry expert id -1; clamp for the scatter (their rows are zeroed). + ref.index_add_(0, slot_expert.to(torch.int64).clamp_min(0), per_slot) + + # Fresh output (no transpose needed downstream). + dW2 = permute_free_grouped_gemm_bf16_fc2_wgrad(fc2_input, grad_output, weights_shape, routing) + assert dW2.shape == weights_shape + assert _rel_l2(dW2, ref) < 2e-2 + + # Direct fp32 accumulate into an existing buffer: overwrite then add == 2x. + out = torch.zeros(weights_shape, device="cuda", dtype=torch.float32) + permute_free_grouped_gemm_bf16_fc2_wgrad( + fc2_input, grad_output, weights_shape, routing, out=out, accumulate=False + ) + assert _rel_l2(out, ref) < 2e-2 + permute_free_grouped_gemm_bf16_fc2_wgrad( + fc2_input, grad_output, weights_shape, routing, out=out, accumulate=True + ) + assert _rel_l2(out, 2.0 * ref) < 2e-2 + + +def test_route_list_fc2_wgrad_recompute_from_preact(): + """FC2 wgrad recompute-from-preact: rebuild act = act(gate)*up*prob from the saved 2F + pre-activation into a transient buffer, feed the unchanged wgrad, and match a stored-act + reference (the backward then checkpoints only the 2F preact, never the F-wide act).""" + from transformer_engine.pytorch.moe import ( + get_default_moe_kernel_config, + prepare_moe_align, + ) + + torch.manual_seed(29) + num_recv_tokens, in_features, out_features = 128, 96, 128 # F=in, H=out (W2 is [E, H, F]) + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=6) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + + # Block-padded canonical: preact / act / grad all live in the [em_max] block-padded slot + # order (padding slots have token sentinel >= num_recv_tokens and are dropped by the kernel). + em_max = int(routing.sorted_slot_ids.shape[0]) + num_routes = int(routing_map.sum().item()) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + preact = torch.zeros(em_max, 2 * in_features, device="cuda", dtype=torch.bfloat16) + preact[valid] = torch.randn( + int(valid.sum().item()), 2 * in_features, device="cuda", dtype=torch.bfloat16 + ) + grad_output = torch.randn(num_recv_tokens, out_features, device="cuda", dtype=torch.bfloat16) + probs = torch.rand(num_recv_tokens, num_experts, device="cuda", dtype=torch.float32) + weights_shape = (num_experts, out_features, in_features) # [E, H, F] + + # Reference act in the block-padded slot layout using the kernel's own slot arrays. + g = preact[:, :in_features].float() + u = preact[:, in_features:].float() + p = probs[tok, exp] + act_ref = torch.nn.functional.silu(g) * u * p[:, None] + act_ref = act_ref * valid[:, None].float() + act_stored = act_ref.to(torch.bfloat16) + + # Stored-act baseline: feed the materialized activation directly. + dW2_stored = permute_free_grouped_gemm_bf16_fc2_wgrad( + act_stored, grad_output, weights_shape, routing + ) + # Recompute path: feed preact + probs, rebuild act in-flight. + dW2_rc = permute_free_grouped_gemm_bf16_fc2_wgrad( + None, grad_output, weights_shape, routing, + preact=preact, dispatched_probs=probs, activation="silu", + ) + + # fp32 reference dW2[e] = sum_{valid slots s: exp==e} outer(grad_output[tok_s], act_ref[s]). + ref = torch.zeros(num_experts, out_features, in_features, device="cuda", dtype=torch.float32) + per_route = torch.einsum("rh,rf->rhf", grad_output[tok].float(), act_ref) + ref.index_add_(0, exp, per_route) + + assert _rel_l2(dW2_stored, ref) < 2e-2 + assert _rel_l2(dW2_rc, ref) < 2e-2 + # Recompute should also track the stored-act path closely (same GEMM, act rebuilt). + assert _rel_l2(dW2_rc, dW2_stored) < 1e-2 + + +def test_fc2_backward_dispatch_recompute_matches_stored(): + """FC2 backward dispatch: wgrad from the saved 2F preact (+ fc2_activation) matches the + legacy stored-F activation path.""" + from transformer_engine.pytorch.moe import ( + get_default_moe_kernel_config, + permute_free_grouped_gemm_backward, + prepare_moe_align, + ) + + torch.manual_seed(31) + num_recv_tokens, in_features, out_features = 128, 96, 128 # F=in, H=out (W2 is [E, H, F]) + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=8) + routing = PermuteFreeMetadata( + routing_map=routing_map, num_experts=num_experts, route_space=True, activation="silu" + ) + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + + # Block-padded canonical: preact / act live in the [em_max] block-padded slot order. + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + preact = torch.zeros(em_max, 2 * in_features, device="cuda", dtype=torch.bfloat16) + preact[valid] = torch.randn( + int(valid.sum().item()), 2 * in_features, device="cuda", dtype=torch.bfloat16 + ) + grad_output = torch.randn(num_recv_tokens, out_features, device="cuda", dtype=torch.bfloat16) + probs = torch.rand(num_recv_tokens, num_experts, device="cuda", dtype=torch.float32) + weights = [ + torch.randn(out_features, in_features, device="cuda", dtype=torch.bfloat16) + for _ in range(num_experts) + ] + + g = preact[:, :in_features].float() + u = preact[:, in_features:].float() + p = probs[tok, exp] + act_stored = (torch.nn.functional.silu(g) * u * p[:, None] * valid[:, None].float()).to( + torch.bfloat16 + ) + + stored = permute_free_grouped_gemm_backward( + grad_output, + routing=routing, + weights=weights, + num_gemms=num_experts, + hidden_states=act_stored, + requires_wgrad=True, + ) + recompute = permute_free_grouped_gemm_backward( + grad_output, + routing=routing, + weights=weights, + num_gemms=num_experts, + hidden_states=preact, + requires_wgrad=True, + dispatched_probs=probs, + fc2_activation="silu", + ) + assert recompute.wgrad_stacked.shape == (num_experts, out_features, in_features) + assert _rel_l2(recompute.wgrad_stacked, stored.wgrad_stacked) < 1e-2 + + +def test_fc1_fc2_gated_pipeline(monkeypatch): + """End-to-end FC1 raw 2F -> FC2 fused activation: forward outputs and backward grads match a + PyTorch reference through the permute-free GroupedLinear modules.""" + import dataclasses + + from transformer_engine.pytorch.moe import ( + get_default_moe_kernel_config, + prepare_moe_align, + ) + + monkeypatch.setenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "1") + torch.manual_seed(37) + # FFN width is a multiple of the FlyDSL block_k (64) so the fused FC2 gated_a prologue + # runs on FlyDSL (the default backend) rather than falling back to Triton. + num_recv_tokens, hidden, ffn = 128, 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=9) + fc1_meta = PermuteFreeMetadata( + routing_map=routing_map, num_experts=num_experts, activation="silu" + ) + # Prepare at the v3-compatible forward block_m (>= 128, multiple of 128); the token-count + # default (64 at 128 tokens) is below the v3 gather/dgrad tile minimum. FC1 and FC2 share the + # same block-padded slot layout, so the derived fc2_meta inherits this align. + prepare_moe_align(fc1_meta, _FLYDSL_FWD_BLOCK_M) + fc2_meta = dataclasses.replace(fc1_meta, route_space=True) + num_routes = int(routing_map.sum().item()) + m_splits = [num_recv_tokens // num_experts] * num_experts + m_splits[-1] += num_recv_tokens - sum(m_splits) + + inp = torch.randn(num_recv_tokens, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + probs = torch.rand(num_recv_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True) + + fc1 = GroupedLinear( + num_experts, hidden, 2 * ffn, bias=False, params_dtype=torch.bfloat16, device="cuda" + ) + fc2 = GroupedLinear( + num_experts, ffn, hidden, bias=False, params_dtype=torch.bfloat16, device="cuda" + ) + torch.manual_seed(38) + with torch.no_grad(): + for i in range(num_experts): + getattr(fc1, f"weight{i}").normal_(0, 0.05) + getattr(fc2, f"weight{i}").normal_(0, 0.05) + + preact = fc1(inp, m_splits, permute_free_metadata=fc1_meta) + assert preact.shape[1] == 2 * ffn + preact.retain_grad() # non-leaf: keep its grad so we can sanity-check the FC2->FC1 edge + out = fc2(preact, m_splits, permute_free_metadata=fc2_meta, dispatched_probs=probs) + + # Reference: FC1 raw 2F [gate|up] -> silu(gate)*up*prob -> FC2, summed over each token's + # experts (the fused FC2 prologue must reproduce this token-space output). Computed on the + # CPU so no GPU (re)allocation happens between the forward and the backward -- the FlyDSL + # FC2 dgrad autotune is sensitive to heap layout shifts (a latent OOB read faults only when + # a large device alloc/free reshuffles the pool between fwd and bwd). + with torch.no_grad(): + w1 = torch.stack([getattr(fc1, f"weight{i}").cpu() for i in range(num_experts)]).float() + w2 = torch.stack([getattr(fc2, f"weight{i}").cpu() for i in range(num_experts)]).float() + pre_ref = torch.einsum("th,enh->ten", inp.detach().cpu().float(), w1) # [T, E, 2F] + gate, up = pre_ref[..., :ffn], pre_ref[..., ffn:] + act = torch.nn.functional.silu(gate) * up * probs.detach().cpu().float()[..., None] + y = torch.einsum("tef,ehf->teh", act, w2) # [T, E, H] + ref = (y * routing_map.cpu().float()[..., None]).sum(dim=1) # [T, H] + rel = (out.detach().cpu().float() - ref).norm() / ref.norm().clamp_min(1e-8) + assert rel < 3e-2, rel.item() + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None + assert probs.grad is not None + for i in range(num_experts): + assert getattr(fc1, f"weight{i}").grad is not None + assert getattr(fc2, f"weight{i}").grad is not None + assert getattr(fc2, f"weight{i}").grad.abs().sum() > 0 + # Sanity: compact route head got a non-zero FC1 output grad (FC2 act-bwd -> FC1). + assert preact.grad[:num_routes].abs().sum() > 0 + + # Numerical check on probs.grad (locks the gated-act backward's per-slot (token, expert) + # mapping): with loss = out.sum(), dL/dprob[t, e] = routing_map[t, e] * . A wrong expert-per-slot map (e.g. compact route index misread as a + # block-padded slot) silently corrupts this even though probs.grad stays non-zero. + with torch.no_grad(): + act_noprob = torch.nn.functional.silu(gate) * up # [T, E, F] (CPU float) + w2sum = w2.sum(dim=1) # [E, F] = sum over output features + dprob_ref = routing_map.cpu().float() * torch.einsum("tef,ef->te", act_noprob, w2sum) + rel_p = (probs.grad.detach().cpu() - dprob_ref).norm() / dprob_ref.norm().clamp_min(1e-8) + assert rel_p < 3e-2, rel_p.item() + + +def test_route_list_fwd_dgrad_wgrad_consistency(): + """fwd + dgrad + wgrad against a single autograd reference in compact route space.""" + torch.manual_seed(23) + # The v3 gather/dgrad tiles need both contraction dims (in for fwd, out for dgrad) to be a + # multiple of the FlyDSL block_k (64) and >= 128 (K_ITERS >= 2), so use 128-wide features. + num_recv_tokens, in_features, out_features = 96, 128, 128 + num_experts, max_hits = 6, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=5) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + route_to_token, route_expert = _compact_route_order(routing_map) + num_routes = route_to_token.numel() + # Prepare the fwd/dgrad align (block-padded slot layout) up front so we can map the compact + # per-route grad into padded-slot order for the block-padded dgrad. + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + valid = routing.sorted_slot_ids < num_recv_tokens + assert int(valid.sum().item()) == num_routes + + hidden = torch.randn(num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16) + weights = torch.randn( + num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16 + ) + + # Forward output is block-padded [em_max]; the valid padded slots are expert-ascending, the + # same order as the compact route list, so out_full[valid] lines up with the compact ref. + out_full = permute_free_grouped_gemm_bf16(hidden, weights, routing) + out = out_full[valid] + # One per-route gradient in the block-padded [em_max] slot layout that both the dgrad and the + # wgrad now route-read (padded slot = block_start[e]*block_size_m + within-rank). The valid + # slots are expert-ascending, matching the compact autograd reference order. ``grad`` keeps a + # compact [num_routes] copy only to seed the compact reference backward. + grad = torch.randn(num_routes, out_features, device="cuda", dtype=torch.bfloat16) + grad_bp = torch.zeros(em_max, out_features, device="cuda", dtype=torch.bfloat16) + grad_bp[valid] = grad + dA = permute_free_grouped_gemm_bf16_dgrad(grad_bp, weights, routing) + dW = permute_free_grouped_gemm_bf16_wgrad(hidden, grad_bp, weights.shape, routing) + + # Autograd reference in compact route space. + ref_hidden = hidden.float().clone().requires_grad_(True) + ref_w = weights.float().clone().requires_grad_(True) + ref_out = torch.einsum( + "rk,rnk->rn", ref_hidden[route_to_token], ref_w[route_expert] + ) + assert _rel_l2(out, ref_out) < 2e-2 + ref_out.backward(grad.float()) + + assert _rel_l2(dA, ref_hidden.grad) < 2e-2 + assert _rel_l2(dW, ref_w.grad) < 2e-2 + + +@pytest.mark.skip(reason="apply_route_probs not ported to transformer_engine.pytorch.moe yet") +def test_apply_route_probs_fwd_bwd(): + """Fused per-route prob apply (gather+multiply) vs an autograd advanced-index reference.""" + from transformer_engine.pytorch.moe import ( + get_default_moe_kernel_config, + prepare_moe_align, + ) + # apply_route_probs was not ported; kept for when route_prob helper lands in moe/. + _ = get_default_moe_kernel_config, prepare_moe_align + + torch.manual_seed(31) + num_recv_tokens, hidden_dim = 128, 192 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=9) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + cfg = get_default_moe_kernel_config(num_recv_tokens) + prepare_moe_align(routing, int(cfg["BLOCK_SIZE_M"])) + + em_max = int(routing.sorted_slot_ids.shape[0]) + num_routes = int(routing_map.sum().item()) + act = torch.randn(em_max, hidden_dim, device="cuda", dtype=torch.bfloat16) + probs = torch.rand(num_recv_tokens, num_experts, device="cuda", dtype=torch.float32) + + a1 = act.clone().requires_grad_(True) + p1 = probs.clone().requires_grad_(True) + out = apply_route_probs(a1, p1, routing) + + # Reference: differentiable advanced index over the block-padded slot order. + tok = routing.sorted_slot_ids.to(torch.int64).clamp(0, num_recv_tokens - 1) + exp = _expert_per_route(routing, em_max).to(torch.int64).clamp_min(0) + valid = routing.sorted_slot_ids < num_recv_tokens + a2 = act.clone().requires_grad_(True) + p2 = probs.clone().requires_grad_(True) + pr = torch.where(valid, p2[tok, exp], torch.zeros_like(p2[tok, exp])) + ref = a2 * pr[:, None] + + assert _rel_l2(out[valid], ref[valid]) < 2e-2 + + g = torch.randn(em_max, hidden_dim, device="cuda", dtype=torch.bfloat16) + g[num_routes:] = 0 + out.backward(g) + ref.backward(g) + assert _rel_l2(a1.grad[:num_routes], a2.grad[:num_routes]) < 2e-2 + assert _rel_l2(p1.grad, p2.grad) < 2e-2 + + +# --------------------------------------------------------------------------- +# Module-level weight-gradient accumulation: grouped weight (-> main_grad) vs. +# separate per-expert weights (-> autograd .grad). +# --------------------------------------------------------------------------- +def _make_grouped_linear(num_gemms, in_features, out_features, *, grouped, fuse_wgrad): + mod = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + fuse_wgrad_accumulation=fuse_wgrad, + single_grouped_weight=grouped, + ) + return mod + + +def test_grouped_weight_main_grad_matches_ungrouped_grad(monkeypatch): + """Permute-free wgrad lands in the grouped param's ``main_grad`` and matches the + ungrouped autograd ``.grad`` expert-for-expert (same kernel dW, different sink).""" + monkeypatch.setenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "1") + # single_grouped_weight requires this gate, else the module falls back to per-expert params. + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + torch.manual_seed(41) + + # The v3 fwd/dgrad tiles need both contraction dims (in for fwd, out for dgrad) >= 128 and a + # multiple of block_k (64); use 128-wide features. + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=13) + num_routes = int(routing_map.sum().item()) + m_splits = [num_recv_tokens // num_experts] * num_experts + + # inp requires grad (as in real training, where it comes from a prior layer): for the + # grouped path the detached per-expert views do not drive autograd, so the activation is + # what makes the Function output require grad and triggers the backward. + inp = torch.randn( + num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + # Shared initial weights for both modules (so the kernel dW is identical). + W = torch.randn(num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16) + # Shared upstream gradient over the compact route range. + g = torch.randn(num_routes, out_features, device="cuda", dtype=torch.bfloat16) + + # --- ungrouped: separate per-expert params, wgrad via autograd .grad --- + mod_a = _make_grouped_linear( + num_experts, in_features, out_features, grouped=False, fuse_wgrad=False + ) + with torch.no_grad(): + for i in range(num_experts): + getattr(mod_a, f"weight{i}").copy_(W[i]) + routing_a = PermuteFreeMetadata(routing_map=routing_map, num_experts=num_experts) + out_a = mod_a(inp, m_splits, permute_free_metadata=routing_a) + # Block-padded canonical: the module output is [em_max, out] in padded slot order; the valid + # (expert-ascending) slots line up with the compact upstream grad g. Slicing [:num_routes] + # would cut across the padding gaps, so gather the valid slots instead. + valid_a = routing_a.sorted_slot_ids < num_recv_tokens + (out_a[valid_a].float() * g.float()).sum().backward() + grad_a = torch.stack( + [getattr(mod_a, f"weight{i}").grad.float() for i in range(num_experts)], dim=0 + ) + + # --- grouped: single grouped param, wgrad accumulated into main_grad --- + mod_b = _make_grouped_linear( + num_experts, in_features, out_features, grouped=True, fuse_wgrad=True + ) + with torch.no_grad(): + for i, view in enumerate(mod_b._get_weight_tensors()): + view.copy_(W[i]) + mod_b.weight.main_grad = torch.zeros( + num_experts, out_features, in_features, device="cuda", dtype=torch.float32 + ) + mod_b.weight.grad_added_to_main_grad = False + routing_b = PermuteFreeMetadata(routing_map=routing_map, num_experts=num_experts) + out_b = mod_b(inp, m_splits, permute_free_metadata=routing_b) + valid_b = routing_b.sorted_slot_ids < num_recv_tokens + (out_b[valid_b].float() * g.float()).sum().backward() + + # The grouped param collects grad in main_grad, not in .grad. + assert mod_b.weight.grad is None + assert getattr(mod_b.weight, "grad_added_to_main_grad", False) is True + main_grad = mod_b.weight.main_grad.view(num_experts, out_features, in_features).float() + + # Both paths fold the same bf16 kernel dW into their sink (grouped -> fp32 main_grad, + # ungrouped -> bf16 .grad), so they match to within bf16 rounding (~1e-3 relative). + assert _rel_l2(main_grad, grad_a) < 5e-3 + # And every expert actually received a non-zero gradient (guards against frozen experts). + for e in range(num_experts): + assert main_grad[e].abs().sum() > 0, f"expert {e} has a zero main_grad (frozen)." + + +def test_grouped_weight_nonfused_grad_matches_ungrouped(monkeypatch): + """single_grouped_weight without fuse_wgrad_accumulation routes the [E, out, in] wgrad into + the grouped param's autograd ``.grad`` (the GEMM only sees detached per-expert views), and it + must match the ungrouped per-expert ``.grad`` expert-for-expert and accumulate across + backwards.""" + monkeypatch.setenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "1") + # single_grouped_weight requires this gate, else the module falls back to per-expert params. + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + torch.manual_seed(43) + + # 128-wide features to satisfy the v3 fwd/dgrad tile (K>=128, multiple of block_k=64). + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 3 + m_splits = [num_recv_tokens // num_experts] * num_experts + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=17) + num_routes = int(routing_map.sum().item()) + # inp drives autograd for the grouped path (detached views carry no grad edge). + inp = torch.randn( + num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + W = torch.randn(num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16) + g = torch.randn(num_routes, out_features, device="cuda", dtype=torch.bfloat16) + + # --- ungrouped reference: separate per-expert params, wgrad via autograd .grad --- + mod_a = _make_grouped_linear( + num_experts, in_features, out_features, grouped=False, fuse_wgrad=False + ) + with torch.no_grad(): + for i in range(num_experts): + getattr(mod_a, f"weight{i}").copy_(W[i]) + routing_a = PermuteFreeMetadata(routing_map=routing_map, num_experts=num_experts) + out_a = mod_a(inp, m_splits, permute_free_metadata=routing_a) + # Block-padded canonical: gather the valid (expert-ascending) slots rather than slicing + # [:num_routes], which would cut across the padding gaps. + valid_a = routing_a.sorted_slot_ids < num_recv_tokens + (out_a[valid_a].float() * g.float()).sum().backward() + grad_a = torch.stack( + [getattr(mod_a, f"weight{i}").grad.float() for i in range(num_experts)], dim=0 + ) + + # --- grouped, no fusion: wgrad accumulates into the grouped param's .grad --- + mod_b = _make_grouped_linear( + num_experts, in_features, out_features, grouped=True, fuse_wgrad=False + ) + with torch.no_grad(): + for i, view in enumerate(mod_b._get_weight_tensors()): + view.copy_(W[i]) + routing_b = PermuteFreeMetadata(routing_map=routing_map, num_experts=num_experts) + out_b = mod_b(inp, m_splits, permute_free_metadata=routing_b) + valid_b = routing_b.sorted_slot_ids < num_recv_tokens + (out_b[valid_b].float() * g.float()).sum().backward() + + # Grad lands in the grouped param's .grad (not main_grad), shape [E, out, in]. + assert getattr(mod_b.weight, "grad_added_to_main_grad", False) is False + assert mod_b.weight.grad is not None + assert tuple(mod_b.weight.grad.shape) == (num_experts, out_features, in_features) + grad_b = mod_b.weight.grad.view(num_experts, out_features, in_features).float() + assert _rel_l2(grad_b, grad_a) < 1e-3 + for e in range(num_experts): + assert grad_b[e].abs().sum() > 0, f"expert {e} has a zero grad (frozen)." + + # A second backward without zero_grad accumulates in place (grad-accumulation semantics). + out_b2 = mod_b(inp, m_splits, permute_free_metadata=routing_b) + (out_b2[valid_b].float() * g.float()).sum().backward() + grad_b2 = mod_b.weight.grad.view(num_experts, out_features, in_features).float() + assert _rel_l2(grad_b2, 2.0 * grad_a) < 1e-3 diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py index 487ff4ba3..07f48a3a1 100644 --- a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py @@ -3,16 +3,17 @@ """Permute-free MoE weight-gradient (wgrad) grouped GEMM in FlyDSL. -Contract: the gradient operand is the *compact* ``[num_routes, N]`` buffer -and ``SORTED`` maps each padded route slot to a received-token row:: +Contract: the gradient operand is the *block-padded* ``[em_max, N]`` buffer +and ``SORTED`` maps each padded slot to a received-token row:: - dW[e][n, k] = sum_{routed slot s of e, valid} grad[route_start[e] + local(s), n] + dW[e][n, k] = sum_{routed slot s of e, valid} grad[grad_base[e] + local(s), n] * x[SORTED[s], k] -where ``local(s)`` is the slot's offset within expert ``e``'s block-padded range and -padding slots (``SORTED[s] == num_recv_tokens``) are masked to zero. The grad walk is a -plain contiguous row scan (no ``SORTED`` indirection); ``SORTED`` is only consulted for -the ``x`` gather token and the padding mask. +where ``grad_base[e] = block_start[e] * block_size_m`` is the expert's block-padded first-slot +row, ``local(s)`` is the slot's offset within expert ``e``'s block-padded range, and padding +slots (``SORTED[s] == num_recv_tokens``) are masked to zero. The grad walk is a plain +contiguous row scan (no ``SORTED`` indirection); ``SORTED`` is only consulted for the ``x`` +gather token and the padding mask. The contraction tile is staged through LDS and transposed on-read: @@ -129,11 +130,11 @@ def compile_moe_wgrad_v2( def wgrad_kernel( dW: fx.Pointer, # [E, N, K] bf16 output X: fx.Pointer, # [num_recv_tokens, K] bf16 (received-token activations) - GRAD: fx.Pointer, # [num_routes, N] bf16 (compact per-route gradient) - SORTED: fx.Pointer, # [padded] i32 received-token row per route slot (sentinel = num_recv_tokens) + GRAD: fx.Pointer, # [em_max, N] bf16 (block-padded per-slot gradient) + SORTED: fx.Pointer, # [padded] i32 received-token row per slot (sentinel = num_recv_tokens) BLOCK_START: fx.Pointer, # [E] i32 (block units) BLOCKS_PER_EXPERT: fx.Pointer, # [E] i32 - ROUTE_START: fx.Pointer, # [E] i32 (compact first-route index = cumsum(counts) - counts) + GRAD_BASE: fx.Pointer, # [E] i32 (block-padded first-slot row = block_start[e] * block_size_m) N: fx.Int32, K: fx.Int32, num_recv_tokens: fx.Int32, @@ -147,7 +148,7 @@ def wgrad_kernel( sorted_rsrc = ptr_rsrc(SORTED) bstart_rsrc = ptr_rsrc(BLOCK_START) bpe_rsrc = ptr_rsrc(BLOCKS_PER_EXPERT) - rstart_rsrc = ptr_rsrc(ROUTE_START) + gbase_rsrc = ptr_rsrc(GRAD_BASE) # DMA path: raw addrspace(3) global backs LDS; reads + DMA both GEP off it. smem_raw_ptr = _llvm.mlir_addressof(ir.Type.parse("!llvm.ptr<3>"), LDS_SYM) @@ -228,15 +229,15 @@ def _scope_kw(sid): n_base_idx = arith.index_cast(T.index, n_block_base) k_base_idx = arith.index_cast(T.index, k_block_base) - # Per-expert routed-slot range. ``base_slot`` is the block-padded slot offset into - # ``SORTED`` (holds the received-token row for the ``x`` gather); ``route_start_e`` - # is the compact first-route index into the ``[num_routes, N]`` grad buffer. + # Per-expert routed-slot range. ``base_slot`` is the wgrad-align slot offset into + # ``SORTED`` (holds the received-token row for the ``x`` gather); ``grad_base_idx`` is + # expert ``e``'s block-padded first-slot row into the ``[em_max, N]`` grad buffer. bstart = buffer_load_i32(bstart_rsrc, expert) nblocks = buffer_load_i32(bpe_rsrc, expert) - rstart = buffer_load_i32(rstart_rsrc, expert) + gbase = buffer_load_i32(gbase_rsrc, expert) base_slot = arith.index_cast(T.index, bstart) * arith.index(WGRAD_BLOCK_M) num_slots = arith.index_cast(T.index, nblocks) * arith.index(WGRAD_BLOCK_M) - route_start_e_idx = arith.index_cast(T.index, rstart) + grad_base_idx = arith.index_cast(T.index, gbase) CPR_G = block_n // FILL_V CPR_X = block_k // FILL_V @@ -284,7 +285,7 @@ def _dma_one(rsrc, ids, slot_base_idx, cpr, feat_base_idx, dim_idx, lds_off, if const_expr(clamp_row): # grad: contiguous route walk; clamp overrun to row 0 for fault safety # (its padding contribution is cancelled by the zeroed x column). - row_idx = valid.select(route_start_e_idx + slot_base_idx + slot_idx, c0) + row_idx = valid.select(grad_base_idx + slot_base_idx + slot_idx, c0) else: # x: gather by received-token; sentinel/OOB row -> hardware 0. row_idx = token @@ -556,7 +557,7 @@ def launch_wgrad( SORTED: fx.Pointer, BLOCK_START: fx.Pointer, BLOCKS_PER_EXPERT: fx.Pointer, - ROUTE_START: fx.Pointer, + GRAD_BASE: fx.Pointer, N: fx.Int32, K: fx.Int32, num_recv_tokens: fx.Int32, @@ -577,7 +578,7 @@ def launch_wgrad( gz = num_experts wgrad_kernel._func.__name__ = KERNEL_NAME wgrad_kernel( - dW, X, GRAD, SORTED, BLOCK_START, BLOCKS_PER_EXPERT, ROUTE_START, + dW, X, GRAD, SORTED, BLOCK_START, BLOCKS_PER_EXPERT, GRAD_BASE, N, K, num_recv_tokens, ).launch(grid=(gx, gy, gz), block=(n_threads, 1, 1), stream=stream) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 4cf5fecee..2a24d1ca9 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -593,7 +593,9 @@ def forward( # Convert splits to list of ints for compatibility with split functions m_splits = m_splits.tolist() - inp_view = inp.reshape(-1, in_features) + # Use ``expect_in_features`` (2F on the gated FC2 route-space path, else in_features) so a + # raw 2F [gate|up] preact is not mangled into F-wide rows by the reshape. + inp_view = inp.reshape(-1, expect_in_features) inputmats: list if fp8 and not debug: # Disable bulk allocation when CPU offloading is active: offloading skips small @@ -1149,8 +1151,9 @@ def backward( ) return ( pf_result.dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, - pf_result.grad_probs, - None, + None, # m_splits + pf_result.grad_probs, # dispatched_probs + None, # non_tensor_args *wgrad_list, *([None] * ctx.num_gemms), ) diff --git a/transformer_engine/pytorch/moe/moe_routing.py b/transformer_engine/pytorch/moe/moe_routing.py index 62bc19cc3..6b85f86b4 100644 --- a/transformer_engine/pytorch/moe/moe_routing.py +++ b/transformer_engine/pytorch/moe/moe_routing.py @@ -47,19 +47,14 @@ class MoERoutingMetadata: num_tokens_post_padded: ``[1]`` device scalar = real ``em`` (block-padded route count). Bounds the kernel. block_start: - ``[num_experts]`` per-expert first block index (block units). - route_start: - ``[num_experts]`` per-expert first *compact* route index (``cumsum(counts) - counts``). - Maps a block-padded position to its compact output row. - route_to_token: - ``[routes_max]`` received-token row for each compact route (first ``num_routes`` - entries valid); used by the dgrad scatter-add back to ``[num_recv_tokens, K]``. + ``[num_experts]`` per-expert first block index (block units). Expert ``e``'s block-padded + slots start at ``block_start[e] * block_size_m``. token_routes / token_route_count: - Inverse of ``route_to_token`` (token -> its compact route positions), built sync-free - for the contention-free gather-combine that replaces the atomic scatter in the token - combine (FC2 fwd) and the FC1 input-grad reduction (FC1 dgrad). ``token_routes`` is + Token -> its block-padded slot positions, built sync-free for the contention-free + gather-combine that replaces the atomic scatter in the token combine (FC2 fwd) and the + FC1 input-grad reduction (FC1 dgrad). ``token_routes`` is ``[num_recv_tokens, min(topk, num_experts)]`` (int32); for token ``t`` the first - ``token_route_count[t]`` entries are its route positions (expert-ascending), the rest + ``token_route_count[t]`` entries are its padded slots (expert-ascending), the rest are unused padding. block_size_m: ``BLOCK_SIZE_M`` used to build the fwd/dgrad align buffers. @@ -77,8 +72,6 @@ class MoERoutingMetadata: expert_ids: Optional[torch.Tensor] = None num_tokens_post_padded: Optional[torch.Tensor] = None block_start: Optional[torch.Tensor] = None - route_start: Optional[torch.Tensor] = None - route_to_token: Optional[torch.Tensor] = None token_routes: Optional[torch.Tensor] = None token_route_count: Optional[torch.Tensor] = None block_size_m: Optional[int] = None diff --git a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py index 199825e76..e7b66a4be 100644 --- a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py +++ b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py @@ -4,8 +4,8 @@ """Permute-free route-list grouped GEMM for MoE (bf16). -- TE builds one expert-sorted ``sorted_slot_ids`` (received-token row per route) plus a - compact-output map (``route_start``/``block_start``), then runs the gather-GEMM. +- TE builds one expert-sorted ``sorted_slot_ids`` (received-token row per block-padded slot) + plus the per-expert ``block_start`` (block units), then runs the gather-GEMM. - FC1 fwd output is worst-case padded ``[T * min(topk, E), out_features]`` in expert order (valid rows are the compact route range ``[0, num_routes)``, tail is inert zero padding); dgrad returns ``dA = [num_recv_tokens, in_features]`` (scatter-add of the per-route gradients). @@ -116,7 +116,6 @@ def _pf_moe_fwd( routing.sorted_slot_ids, routing.expert_ids, routing.block_start, - routing.route_start, num_recv_tokens=num_recv_tokens, block_m=block_m, index_a_by_route_pos=index_a_by_route_pos, @@ -161,6 +160,15 @@ def _fwd_align_block_size_m( candidates = {c for c in (_FLYDSL_GATED_BLOCK_M, default_block_m) if c <= default_block_m} if num_tokens >= _FLYDSL_FWD_LARGE_TIER: candidates.add(_FLYDSL_FWD_BLOCK_M) + # The v3 gather/dgrad tile requires block_m >= 128 (128x256 MFMA minimum), so on a + # small-token tier where the token-count default is < 128 we must still floor the align + # block_m to 128 rather than emit a sub-tile the kernel cannot run. + from .pf_fwd_wrapper import _v3_enabled + + if _v3_enabled(): + candidates = {c for c in candidates if c >= _FLYDSL_GATED_BLOCK_M} or { + _FLYDSL_GATED_BLOCK_M + } # Prefer the in-tree FlyDSL block_m picker when available. try: @@ -263,9 +271,9 @@ def moe_align_route_list( Returns ------- ``(sorted_slot_ids, expert_ids, num_tokens_post_padded, block_start, blocks_per_expert, - route_start, route_to_token, token_routes, token_route_count)``. Index tensors are - ``int32``; ``num_tokens_post_padded`` (``[1]``) is a device scalar. The last two are the - token->routes inverse map, ``None`` unless ``build_inverse_map``. + token_routes, token_route_count)``. Index tensors are ``int32``; + ``num_tokens_post_padded`` (``[1]``) is a device scalar. The last two are the token->routes + inverse map, ``None`` unless ``build_inverse_map``. """ return route_list_align( routing_map, @@ -303,8 +311,6 @@ def prepare_moe_align(metadata: MoERoutingMetadata, block_m: int) -> MoERoutingM num_tokens_post_padded, block_start, _blocks_per_expert, - route_start, - route_to_token, token_routes, token_route_count, ) = moe_align_route_list( @@ -319,8 +325,6 @@ def prepare_moe_align(metadata: MoERoutingMetadata, block_m: int) -> MoERoutingM metadata.expert_ids = expert_ids # [blocks_max] int32: expert owning each block (-1 past end) metadata.num_tokens_post_padded = num_tokens_post_padded # [1] int32 device scalar: padded extent metadata.block_start = block_start # [E] int32: first block index of each expert (block units) - metadata.route_start = route_start # [E] int32: first compact route index of each expert - metadata.route_to_token = route_to_token # [routes_max] int32: compact route -> token metadata.block_size_m = block_m # int: BLOCK_SIZE_M the layout is padded to metadata.token_routes = token_routes # [T, min(topk, E)] int32: token -> its compact route ids metadata.token_route_count = token_route_count # [T] int32: number of routes per token @@ -343,8 +347,6 @@ def _prepare_wgrad_align( _num_tokens_post_padded, block_start, blocks_per_expert, - route_start, - _route_to_token, _token_routes, _token_route_count, ) = moe_align_route_list( @@ -358,9 +360,6 @@ def _prepare_wgrad_align( metadata.wgrad_block_start = block_start metadata.wgrad_blocks_per_expert = blocks_per_expert metadata.wgrad_block_size = contract_m - # route_start is block-size-independent; keep whichever is already cached. - if metadata.route_start is None: - metadata.route_start = route_start return metadata @@ -424,8 +423,8 @@ def permute_free_grouped_gemm_bf16( ) -> torch.Tensor: """Route-list gather-in-GEMM (FC1 forward) for bf16 MoE. - Computes ``C[route] = A[route_to_token[route]] @ W[e]^T`` for each expert-sorted route, - writing directly into a compact ``[num_routes, out_features]`` output. + Computes ``C[slot] = A[sorted_slot_ids[slot]] @ W[expert(slot)]^T`` for each block-padded + slot, writing directly into the ``[em_max, out_features]`` block-padded output. Parameters ---------- @@ -453,14 +452,13 @@ def permute_free_grouped_gemm_bf16( Returns ------- torch.Tensor - Worst-case padded ``[T * min(topk, E), out_features]`` (or ``[T * min(topk, E), F]`` with a fused - ``activation``), bf16 (expert-contiguous). The valid rows are the compact route range - ``[0, num_routes)``; the tail ``[num_routes, em_max)`` is inert, *uninitialized* - padding. Consumers MUST read only the compact range, using the routing metadata to - locate each expert's rows: - - - ``route_start[e] = counts[0] + ... + counts[e-1]`` (``= cumsum(counts) - counts``): - its starting offset in the packed output. + Block-padded ``[em_max, out_features]`` (or ``[em_max, F]`` with a fused ``activation``), + bf16 (expert-contiguous, each expert padded to ``block_size``). The valid rows are the + block-padded slots whose ``sorted_slot_ids[slot] < num_recv_tokens``; the padding slots + carry inert dead values. Consumers locate each expert's rows via the routing metadata: + + - ``block_start[e]`` (block units): expert ``e``'s rows occupy padded slots + ``[block_start[e] * block_size, ...)``, with within-rank offset matching each route. - ``num_tokens_post_padded = sum_e ceil(counts[e] / block_size) * block_size``: the block-padded route count (each expert's row count rounded up to ``block_size``). """ @@ -573,13 +571,16 @@ def permute_free_gated_act_bwd( ``(dpre, dprob)`` -- ``dpre`` is ``[T * min(topk, E), 2F]`` bf16; ``dprob`` matches ``dispatched_probs`` (or ``None`` when no probs were fused). """ - routes_max = int(routing.route_to_token.shape[0]) - token = routing.route_to_token.to(torch.int32) - expert = _expert_per_route(routing, routes_max) - # The route buffers are statically sized to the worst case ``routes_max = T * topk``, but - # only the dense head ``[0, num_routes)`` is real (under EP this can be ~topk*E_local/E - # smaller). ``num_tokens_post_padded`` is a device scalar >= num_routes (block-padded), so - # it bounds the kernel to the real routes -- sync-free -- and lets the tail exit early. + # Block-padded canonical layout: grad_out (FC2 dgrad) and preact both live in the [em_max] + # padded slot order, and the kernel indexes grad_out/preact/dpre by the same row as + # token/expert -- so we must feed the padded slot arrays (sorted_slot_ids + slot expert), + # not the compact route arrays, or the padded valid slots beyond routes_max never get a dpre + # row (and each route pairs a correct token with the wrong padded grad row). + em_max = int(routing.sorted_slot_ids.shape[0]) + token = routing.sorted_slot_ids.to(torch.int32) + expert = _expert_per_route(routing, em_max) + # ``num_tokens_post_padded`` is a device scalar bounding the real padded extent, so the tail + # programs exit early (sync-free) instead of streaming the padding through HBM. return fused_gated_act_prob_bwd( grad_output.contiguous(), preact, @@ -610,11 +611,15 @@ def permute_free_gated_act_recompute( Returns ------- torch.Tensor - ``act``, shape ``[T * min(topk, E), F]``, bf16 (route/padded layout). + ``act``, shape ``[em_max, F]``, bf16 (block-padded slot layout, matching the FC1 + forward output and the layout the wgrad route-reads). """ - routes_max = int(routing.route_to_token.shape[0]) - token = routing.route_to_token.to(torch.int32) - expert = _expert_per_route(routing, routes_max) + # Block-padded canonical layout: emit one row per padded slot (keyed by sorted_slot_ids) so + # the rebuilt activation lines up with the [em_max] grad the wgrad route-reads. Padding slots + # (token sentinel >= num_recv_tokens) are masked to zero by the kernel. + em_max = int(routing.sorted_slot_ids.shape[0]) + token = routing.sorted_slot_ids.to(torch.int32) + expert = _expert_per_route(routing, em_max) return fused_gated_act_prob_fwd( preact, token, @@ -721,9 +726,9 @@ def permute_free_grouped_gemm_bf16_wgrad( ) -> torch.Tensor: """Route-list fused weight-gradient (FC1 backward wrt weights), FlyDSL-only. - Computes ``dW[e] = sum_{route in e} grad[route]^T @ A[route_to_token[route]]`` by - gathering the activation operand (received-token row) and reading the compact grad row - in the FlyDSL kernel. + Computes ``dW[e] = sum_{slot in e} grad[slot]^T @ A[sorted_slot_ids[slot]]`` by gathering + the activation operand (received-token row) and reading the block-padded grad row (base + ``block_start[e] * block_size_m`` + within-rank) in the FlyDSL kernel. Parameters ---------- @@ -777,7 +782,15 @@ def permute_free_grouped_gemm_bf16_wgrad( grad_output = grad_output.contiguous() contract_m = _WGRAD_CONTRACT_M + # The grad operand lives in the forward's block-padded [em_max] slot layout (its row for + # route (e, w) is block_start[e]*block_size_m + w). Ensure the forward align is present so we + # can pass that padded per-expert base to the kernel; routes are contiguous within an expert, + # so base + within-rank indexes the correct padded grad row. + if routing.block_start is None or routing.block_size_m is None: + fwd_config = config or get_default_moe_kernel_config(routing.num_recv_tokens) + routing = prepare_moe_align(routing, int(fwd_config["BLOCK_SIZE_M"])) routing = _prepare_wgrad_align(routing, contract_m) + grad_base = (routing.block_start.to(torch.int64) * int(routing.block_size_m)).to(torch.int32) if swap_gather: raise RuntimeError( @@ -802,7 +815,7 @@ def permute_free_grouped_gemm_bf16_wgrad( routing.wgrad_sorted_slot_ids, routing.wgrad_block_start, routing.wgrad_blocks_per_expert, - routing.route_start, + grad_base, num_recv_tokens=routing.num_recv_tokens, accumulate=bool(accumulate), ) @@ -820,7 +833,7 @@ def permute_free_grouped_gemm_bf16_wgrad( routing.wgrad_sorted_slot_ids, routing.wgrad_block_start, routing.wgrad_blocks_per_expert, - routing.route_start, + grad_base, num_recv_tokens=routing.num_recv_tokens, accumulate=bool(accumulate), ) @@ -1245,10 +1258,21 @@ def permute_free_grouped_gemm_backward( dgrad = permute_free_grouped_gemm_bf16_dgrad(grad_output, weights_stacked, routing) if requires_wgrad: weights_shape = (num_gemms, weights[0].size(0), weights[0].size(1)) - dW = permute_free_grouped_gemm_bf16_wgrad( - hidden_states, grad_output, weights_shape, routing, - out=wgrad_out, accumulate=wgrad_accumulate, - ) # [E, N, H] + if wgrad_out is not None and wgrad_out.dtype != torch.bfloat16: + # fp32 main_grad sink: the bf16 kernel can't write it directly, so emit the bf16 + # dW and fold it into the fp32 buffer here (mirrors the FC2 wgrad out handling). + dW = permute_free_grouped_gemm_bf16_wgrad( + hidden_states, grad_output, weights_shape, routing, + ) # [E, N, H] + if wgrad_accumulate: + wgrad_out.add_(dW) + else: + wgrad_out.copy_(dW) + else: + dW = permute_free_grouped_gemm_bf16_wgrad( + hidden_states, grad_output, weights_shape, routing, + out=wgrad_out, accumulate=wgrad_accumulate, + ) # [E, N, H] wgrad_stacked = dW wgrad_applied = wgrad_out is not None diff --git a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py index b9ef557f2..bb1fa968b 100644 --- a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py +++ b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py @@ -4,8 +4,8 @@ """FlyDSL permute-free MoE route-list forward gather-GEMM op wrapper. Mirrors the route-list forward gather-GEMM contract so the same routing -metadata (``sorted_slot_ids`` / ``expert_ids`` / ``block_start`` / ``route_start``) can be -reused verbatim. Writes the compact ``[em_max, WIDTH_N]`` route output in place. +metadata (``sorted_slot_ids`` / ``expert_ids`` / ``block_start``) can be reused verbatim. +Writes the block-padded ``[em_max, WIDTH_N]`` slot output in place. Forward gather-GEMM is FlyDSL-only. This module retains Triton kernels for routing metadata construction, gated-activation recompute/bwd, and the token-order gather-combine @@ -241,7 +241,6 @@ def flydsl_moe_fwd( sorted_slot_ids: torch.Tensor, expert_ids: torch.Tensor, block_start: torch.Tensor, - route_start: torch.Tensor, *, num_recv_tokens: int, block_m: int, @@ -388,7 +387,6 @@ def flydsl_moe_fwd_autotuned( sorted_slot_ids: torch.Tensor, expert_ids: torch.Tensor, block_start: torch.Tensor, - route_start: torch.Tensor, *, num_recv_tokens: int, block_m: int, @@ -418,7 +416,7 @@ def flydsl_moe_fwd_autotuned( def _launch(bn, bk, wm, wn): flydsl_moe_fwd( - A, B, C, sorted_slot_ids, expert_ids, block_start, route_start, + A, B, C, sorted_slot_ids, expert_ids, block_start, num_recv_tokens=num_recv_tokens, block_m=block_m, block_n=bn, block_k=bk, warps_m=wm, warps_n=wn, index_a_by_route_pos=index_a_by_route_pos, activation=activation, dispatched_probs=dispatched_probs, preact_out=preact_out, diff --git a/transformer_engine/pytorch/moe/pf_helper_kernels.py b/transformer_engine/pytorch/moe/pf_helper_kernels.py index c01984654..6a106c1d1 100644 --- a/transformer_engine/pytorch/moe/pf_helper_kernels.py +++ b/transformer_engine/pytorch/moe/pf_helper_kernels.py @@ -474,8 +474,8 @@ def fused_gated_act_prob_fwd( @triton.jit def _gather_combine_kernel( - src_ptr, # compact [T * min(topk, E), N] (route order; valid rows [0, num_routes)) - token_routes_ptr, # [T, MAXK] int32: compact route positions per token + src_ptr, # block-padded [em_max, N] (padded slot order; padding rows carry dead values) + token_routes_ptr, # [T, MAXK] int32: block-padded slot positions per token token_count_ptr, # [T] int32: number of routes for each token out_ptr, # [T, N] out N, @@ -494,7 +494,8 @@ def _gather_combine_kernel( n_mask = offs_n < N acc = tl.zeros((BLOCK_N,), dtype=tl.float32) # Sum the token's route rows (expert-ascending -> deterministic). Columns >= cnt are - # unused padding and skipped, so no padded/garbage row is ever gathered. + # unused padding and skipped, so no padded/garbage slot is ever gathered (this masked + # padded->token reduction is the sole output-side masking point in the block-padded path). for j in range(0, MAXK): if j < cnt: r = tl.load(token_routes_ptr + t * MAXK + j).to(tl.int64) @@ -589,7 +590,6 @@ def _expert_meta_kernel( counts_ptr, # [E] int32 blocks_per_expert_ptr, # [E] int32 out block_start_ptr, # [E] int32 out (block units) - route_start_ptr, # [E] int32 out (compact route units) expert_ids_ptr, # [blocks_max] int32 out (expert owning each block, -1 past the end) ntpp_ptr, # [1] int32 out: block-padded token extent E, @@ -605,15 +605,12 @@ def _expert_meta_kernel( cblocks = tl.cumsum(blocks_per_expert, axis=0) # inclusive prefix over experts # block_start is over padded block token counts block_start = cblocks - blocks_per_expert - # route_start is over raw token counts - route_start = tl.cumsum(counts, axis=0) - counts total_blocks = tl.max(tl.where(mask_e, cblocks, 0), axis=0) pid = tl.program_id(axis=0) if pid == 0: tl.store(blocks_per_expert_ptr + offs_e, blocks_per_expert, mask=mask_e) tl.store(block_start_ptr + offs_e, block_start, mask=mask_e) - tl.store(route_start_ptr + offs_e, route_start, mask=mask_e) tl.store(ntpp_ptr, total_blocks * BLOCK_SIZE_M) # expert_ids[b] = #{e : cblocks[e] <= b} (== searchsorted(cblocks, b, right=True)), @@ -632,9 +629,7 @@ def _route_list_place_kernel( routing_map_ptr, # [T, E] (bool/int8), True where token t feeds local expert e within_ptr, # [E, T] int32, exclusive within-expert rank of each routed cell block_start_ptr, # [E] int32: first block index of each expert (block units) - route_start_ptr, # [E] int32: first compact route index of each expert sorted_slot_ids_ptr, # [T * min(topk, E)] int32, sentinel-init (T) - route_to_token_ptr, # [routes_max] int32, sentinel-init (T) token_routes_ptr, # [T, MAXK] int32 out: token->route positions (only if BUILD_INVERSE) token_count_ptr, # [T] int32 out: routes per token (only if BUILD_INVERSE) T, @@ -645,11 +640,10 @@ def _route_list_place_kernel( BUILD_INVERSE: tl.constexpr, MAXK: tl.constexpr, ): - # One row of the map per program; place the (few) routed cells deterministically. Because - # the same per-token expert scan already computes each routed cell's compact route position - # ``pos = route_start[e] + within[e, t]``, the token->routes inverse map (used by the - # contention-free gather-combine) is emitted here too when ``BUILD_INVERSE`` -- folding what - # was a second per-token kernel launch into this one (expert-ascending, no atomics). + # One row of the map per program; place the (few) routed cells deterministically into their + # block-padded slot ``bs*BLOCK_SIZE_M + within[e, t]``. The token->routes inverse map (used by + # the contention-free gather-combine) is emitted here too when ``BUILD_INVERSE`` -- folding + # what was a second per-token kernel launch into this one (expert-ascending, no atomics). t = tl.program_id(axis=0) if t >= T: return @@ -659,12 +653,14 @@ def _route_list_place_kernel( if is_routed != 0: w = tl.load(within_ptr + e * T + t) bs = tl.load(block_start_ptr + e) - rs = tl.load(route_start_ptr + e) - pos = rs + w - tl.store(sorted_slot_ids_ptr + bs * BLOCK_SIZE_M + w, t) - tl.store(route_to_token_ptr + pos, t) + slot = bs * BLOCK_SIZE_M + w + tl.store(sorted_slot_ids_ptr + slot, t) if BUILD_INVERSE: - tl.store(token_routes_ptr + t * MAXK + j, pos) + # Block-padded canonical layout: every intermediate route tensor (the GEMM + # output, the FC2 route-read input, the activation) is indexed by the padded + # slot ``bs*BLOCK_SIZE_M + w``. The token->routes inverse map therefore stores the + # padded slot so the final padded->token gather-combine reads the correct rows. + tl.store(token_routes_ptr + t * MAXK + j, slot) j += 1 if BUILD_INVERSE: tl.store(token_count_ptr + t, j) @@ -729,10 +725,9 @@ def route_list_align( Returns ------- ``(sorted_slot_ids, expert_ids, num_tokens_post_padded, block_start, blocks_per_expert, - route_start, route_to_token, token_routes, token_route_count)`` -- index tensors - ``int32``; ``num_tokens_post_padded`` (``[1]``) is a device scalar. ``token_routes`` - (``[T, min(topk, E)]``) and ``token_route_count`` (``[T]``) are ``None`` unless - ``build_inverse_map``. + token_routes, token_route_count)`` -- index tensors ``int32``; ``num_tokens_post_padded`` + (``[1]``) is a device scalar. ``token_routes`` (``[T, min(topk, E)]``) and + ``token_route_count`` (``[T]``) are ``None`` unless ``build_inverse_map``. """ if routing_map.dtype != torch.bool: routing_map = routing_map.bool() @@ -758,7 +753,6 @@ def route_list_align( # Per-expert placement metadata + per-block expert ids (single launch). blocks_per_expert = torch.empty((E,), dtype=torch.int32, device=device) block_start = torch.empty((E,), dtype=torch.int32, device=device) - route_start = torch.empty((E,), dtype=torch.int32, device=device) expert_ids = torch.empty((blocks_max,), dtype=torch.int32, device=device) num_tokens_post_padded = torch.empty((1,), dtype=torch.int32, device=device) block_b = 256 @@ -766,7 +760,6 @@ def route_list_align( counts, blocks_per_expert, block_start, - route_start, expert_ids, num_tokens_post_padded, E, @@ -788,19 +781,14 @@ def route_list_align( token_routes = torch.empty((1,), dtype=torch.int32, device=device) # unused stub token_route_count = token_routes - # Scatter each routed cell into its deterministic (block-padded / compact) slot. Both - # sentinel buffers are carved from a single fill (one launch); the place kernel then - # overwrites the routed slots in each contiguous view. - sentinel = torch.full((em_max + routes_max,), T, dtype=torch.int32, device=device) - sorted_slot_ids = sentinel[:em_max] - route_to_token = sentinel[em_max:] + # Scatter each routed cell into its deterministic block-padded slot. The sentinel-init + # buffer is filled once (one launch); the place kernel then overwrites the routed slots. + sorted_slot_ids = torch.full((em_max,), T, dtype=torch.int32, device=device) _route_list_place_kernel[(T,)]( routing_map, within, block_start, - route_start, sorted_slot_ids, - route_to_token, token_routes, token_route_count, T, @@ -818,8 +806,6 @@ def route_list_align( num_tokens_post_padded, block_start, blocks_per_expert, - route_start, - route_to_token, token_routes if build_inverse_map else None, token_route_count if build_inverse_map else None, ) diff --git a/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py b/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py index a2302078a..102559179 100644 --- a/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py +++ b/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py @@ -5,9 +5,10 @@ Mirrors the route-list wgrad contract so the same routing metadata (``sorted_slot_ids`` holding the received-token row per slot, plus ``block_start`` / -``blocks_per_expert`` / ``route_start``) can be reused verbatim. +``blocks_per_expert`` and the per-expert block-padded grad base ``grad_base``) can be reused +verbatim. -Fixed kernel configuration (matches ``pf_wgrad.py``): FC1 route-list wgrad with compact +Fixed kernel configuration (matches ``pf_wgrad.py``): FC1 route-list wgrad with block-padded ``grad`` + token-gathered ``x``, bf16 into ``dw`` (overwrite or accumulate), DMA + XOR chunk swizzle fill, 3-stage LDS pipeline. Only the workgroup tile geometry is selectable / autotuned. @@ -33,7 +34,7 @@ def flydsl_moe_wgrad( sorted_slot_ids: torch.Tensor, block_start: torch.Tensor, blocks_per_expert: torch.Tensor, - route_start: torch.Tensor, + grad_base: torch.Tensor, *, num_recv_tokens: int, block_n: int = 128, @@ -42,13 +43,13 @@ def flydsl_moe_wgrad( warps_k: int = 2, accumulate: bool = False, ) -> None: - """Compute FC1 grouped wgrad ``grad[route]^T @ x[token(route)]`` into ``dw``, per expert. + """Compute FC1 grouped wgrad ``grad[slot]^T @ x[token(slot)]`` into ``dw``, per expert. See the permute-free wgrad API for the argument contract: ``x`` is - ``[num_recv_tokens, K]`` (gathered by received-token row), ``grad`` is the compact - ``[num_routes, N]`` per-route gradient, ``sorted_slot_ids`` maps each block-padded - route slot to its received-token row (sentinel ``num_recv_tokens`` for padding), and - ``route_start[e]`` is the compact first-route index of expert ``e``. + ``[num_recv_tokens, K]`` (gathered by received-token row), ``grad`` is the block-padded + ``[em_max, N]`` per-slot gradient, ``sorted_slot_ids`` maps each block-padded slot to its + received-token row (sentinel ``num_recv_tokens`` for padding), and ``grad_base[e]`` is the + block-padded first-slot row of expert ``e`` (``block_start[e] * block_size_m``). ``block_n``/``block_k`` and ``warps_n``/``warps_k`` select the workgroup tile; the defaults (``128x128`` over ``2x2`` warps) are a strong general config on CDNA4. @@ -75,7 +76,7 @@ def flydsl_moe_wgrad( ptr_arg(sorted_slot_ids), ptr_arg(block_start), ptr_arg(blocks_per_expert), - ptr_arg(route_start), + ptr_arg(grad_base), int(N), int(K), int(num_recv_tokens), @@ -108,7 +109,7 @@ def _wgrad_run( sorted_slot_ids, block_start, blocks_per_expert, - route_start, + grad_base, N, K, num_recv_tokens, @@ -135,7 +136,7 @@ def _wgrad_run( ptr_arg(sorted_slot_ids), ptr_arg(block_start), ptr_arg(blocks_per_expert), - ptr_arg(route_start), + ptr_arg(grad_base), int(N), int(K), int(num_recv_tokens), @@ -165,14 +166,14 @@ def _get_autotuner(warmup=10, rep=30): def _select_wgrad_config( - x, grad, sorted_slot_ids, block_start, blocks_per_expert, route_start, + x, grad, sorted_slot_ids, block_start, blocks_per_expert, grad_base, N, K, num_recv_tokens, num_experts, ): """Return the autotuned ``(block_n, block_k, warps_n, warps_k)`` for this problem.""" tuner = _get_autotuner() scratch = torch.empty(num_experts, N, K, device=x.device, dtype=torch.bfloat16) args = ( - scratch, x, grad, sorted_slot_ids, block_start, blocks_per_expert, route_start, + scratch, x, grad, sorted_slot_ids, block_start, blocks_per_expert, grad_base, int(N), int(K), int(num_recv_tokens), int(num_experts), ) key = tuner._make_key(args, {}) @@ -189,7 +190,7 @@ def flydsl_moe_wgrad_autotuned( sorted_slot_ids: torch.Tensor, block_start: torch.Tensor, blocks_per_expert: torch.Tensor, - route_start: torch.Tensor, + grad_base: torch.Tensor, *, num_recv_tokens: int, accumulate: bool = False, @@ -207,7 +208,7 @@ def flydsl_moe_wgrad_autotuned( assert x.is_contiguous() and grad.is_contiguous() and dw.is_contiguous() block_n, block_k, warps_n, warps_k = _select_wgrad_config( - x, grad, sorted_slot_ids, block_start, blocks_per_expert, route_start, + x, grad, sorted_slot_ids, block_start, blocks_per_expert, grad_base, N, K, num_recv_tokens, num_experts, ) flydsl_moe_wgrad( @@ -217,7 +218,7 @@ def flydsl_moe_wgrad_autotuned( sorted_slot_ids, block_start, blocks_per_expert, - route_start, + grad_base, num_recv_tokens=int(num_recv_tokens), block_n=block_n, block_k=block_k, From 0406870a90b88aafb218c33553b2615701dd0697 Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Wed, 5 Aug 2026 22:08:23 +0000 Subject: [PATCH 35/43] Extend FlyDSL permute-free wgrad to support FC2 swap-gather and fp32 accumulation Adds `swap_gather` and `out_dtype` options to the FlyDSL wgrad kernel so it can compute the FC2 weight gradient directly in the native `[E, H, F]` layout (token-gathering `grad_output`, contiguous walk over block-padded `fc2_input`) without a transpose. Also enables fp32 `dw` output, allowing fp32 `main_grad` sinks to accumulate in-kernel instead of writing a bf16 scratch buffer and folding separately. Updates the PyTorch wrappers and `GroupedLinear` backward path to use these new kernel variants. --- .../microbenchmarks/benchmark_grouped_gemm.py | 6 +- .../permute_free_grouped_gemm/pf_wgrad.py | 62 +++++++++++---- .../pytorch/moe/permute_free_grouped_gemm.py | 75 ++++++++++--------- .../pytorch/moe/pf_wgrad_wrapper.py | 20 ++++- 4 files changed, 107 insertions(+), 56 deletions(-) diff --git a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_grouped_gemm.py index 27f486900..249b64ae6 100755 --- a/benchmarks/microbenchmarks/benchmark_grouped_gemm.py +++ b/benchmarks/microbenchmarks/benchmark_grouped_gemm.py @@ -71,11 +71,11 @@ def _generate_moe_test_cases( return test_cases -def generate_deepseekv3_test_cases(include_gateup: bool = False): - # DSV3-GateUP hangs on some hardware; only benchmark DSV3-Down by default. +def generate_deepseekv3_test_cases(): + # DSV3-GateUP hangs on some hardware; only benchmark DSV3-Down. return _generate_moe_test_cases( "DSV3", n_routed_experts=256, moe_intermediate_size=2048, hidden_size=7168, - skip_shapes=None if include_gateup else ["GateUP"], + skip_shapes=["GateUP"], ) diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py index 07f48a3a1..6c81fd6d9 100644 --- a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py @@ -78,7 +78,21 @@ def compile_moe_wgrad_v2( warps_n: int = 1, warps_k: int = 1, accumulate: bool = False, + swap_gather: bool = False, + out_dtype: str = "bf16", ): + # ``out_dtype`` selects the ``dW`` element type: ``"bf16"`` (default) truncates the f32 MFMA + # accumulator on store; ``"fp32"`` stores/reads-modify-writes the f32 accumulator directly. + # The latter lets an fp32 ``main_grad`` sink accumulate in-kernel (no bf16 scratch + fold). + if out_dtype not in ("bf16", "fp32"): + raise ValueError(f"moe_wgrad out_dtype must be 'bf16' or 'fp32', got {out_dtype!r}") + # ``swap_gather`` mirrors the two operands' gather/contiguous roles for the FC2 wgrad: + # FC1 walks ``GRAD`` (the [em_max, N] block-padded gradient) contiguously and token-gathers + # ``X`` (the [num_recv, K] activation) -> ``dW[E, N, K]``. FC2 instead token-gathers ``GRAD`` + # (the [num_recv, N] token-space dL/dFC2out) and walks ``X`` (the [em_max, K] block-padded + # activation) contiguously, so the same kernel emits the native ``dW2[E, H, F]`` layout with + # no transpose. Only the per-operand ``clamp_row`` and which operand's resource is bounded to + # ``[num_recv, feat]`` differ. if block_n % (warps_n * WMMA_M) != 0 or block_k % (warps_k * WMMA_N) != 0: raise ValueError("block_n/block_k must be multiples of warps_*16") @@ -110,9 +124,11 @@ def compile_moe_wgrad_v2( G_FILLS = (WGRAD_BLOCK_M * block_n) // (n_threads * FILL_V) X_FILLS = (WGRAD_BLOCK_M * block_k) // (n_threads * FILL_V) + out_is_f32 = out_dtype == "fp32" KERNEL_NAME = ( f"moe_wgrad_routelist_bf16_{block_n}x{block_k}_w{warps_n}x{warps_k}" - f"{'_acc' if accumulate else ''}_dsz_s3_dsa_v2" + f"{'_o32' if out_is_f32 else ''}{'_acc' if accumulate else ''}" + f"{'_sg' if swap_gather else ''}_dsz_s3_dsa_v2" ) # LDS allocation: NUM_BUF-buffered grad tile + x tile (2 bytes/bf16). @@ -212,17 +228,25 @@ def _scope_kw(sid): K_idx = arith.index_cast(T.index, K) nrecv_idx = arith.index_cast(T.index, num_recv_tokens) - # DMA fill cannot mask padding slots in registers, so bound the token-gathered - # ``x`` operand's resource to its real [num_recv, K] extent: the sentinel token + # DMA fill cannot mask padding slots in registers, so bound the *token-gathered* + # operand's resource to its real [num_recv, feat] extent: the sentinel token # (== num_recv) and any pipeline overrun then read out-of-bounds -> hardware # returns 0. A zeroed gathered column makes the padding slot's outer product - # ``grad (x) 0 == 0``, so the paired contiguous-walk grad operand may safely read - # garbage (clamped to row 0) for those slots. - x_addr_i64 = arith.index_cast(T.i64, ptrtoint(X)) - x_nrec_bytes = nrecv_idx * K_idx * arith.index(2) - x_rsrc = buffer_ops.create_buffer_resource_from_addr( - x_addr_i64, num_records_bytes=x_nrec_bytes - ) + # ``gathered (x) 0 == 0``, so the paired contiguous-walk operand may safely read + # garbage (clamped to row 0) for those slots. FC1 gathers ``x`` (stride K); FC2 + # (``swap_gather``) gathers ``grad`` (stride N). + if swap_gather: + grad_addr_i64 = arith.index_cast(T.i64, ptrtoint(GRAD)) + grad_nrec_bytes = nrecv_idx * N_idx * arith.index(2) + grad_rsrc = buffer_ops.create_buffer_resource_from_addr( + grad_addr_i64, num_records_bytes=grad_nrec_bytes + ) + else: + x_addr_i64 = arith.index_cast(T.i64, ptrtoint(X)) + x_nrec_bytes = nrecv_idx * K_idx * arith.index(2) + x_rsrc = buffer_ops.create_buffer_resource_from_addr( + x_addr_i64, num_records_bytes=x_nrec_bytes + ) n_block_base = n_tile * block_n k_block_base = k_tile * block_k @@ -303,17 +327,19 @@ def _dma_one(rsrc, ids, slot_base_idx, cpr, feat_base_idx, dim_idx, lds_off, ) def dma_fill(g_ids, x_ids, slot_base_idx, g_buf_byte, x_buf_byte, wbuf=None): - # FC1: token-gather ``x`` (clamp_row=False); compact grad route walk (clamp_row=True). + # FC1 (swap_gather=False): block-padded grad route walk (clamp_row=True) + token-gather + # ``x`` (clamp_row=False). FC2 (swap_gather=True) flips both: token-gather ``grad`` + # (clamp_row=False) + block-padded activation walk on ``x`` (clamp_row=True). # ``wbuf`` is the Python ring-slot index this tile is being staged into (for alias scopes). g_sid = _g_sid(wbuf) if wbuf is not None else None x_sid = _x_sid(wbuf) if wbuf is not None else None _dma_one( grad_rsrc, g_ids, slot_base_idx, CPR_G_SWZ, n_base_idx, N_idx, - g_lds_off, g_buf_byte, G_FILLS, clamp_row=True, sid=g_sid, + g_lds_off, g_buf_byte, G_FILLS, clamp_row=not swap_gather, sid=g_sid, ) _dma_one( x_rsrc, x_ids, slot_base_idx, CPR_X_SWZ, k_base_idx, K_idx, - x_lds_off, x_buf_byte, X_FILLS, clamp_row=False, sid=x_sid, + x_lds_off, x_buf_byte, X_FILLS, clamp_row=swap_gather, sid=x_sid, ) def _dma_barrier(keep=0): @@ -517,6 +543,9 @@ def _tail(j, accs): # Epilogue: C[m=n_feat, n=k_feat], lane holds 4 rows. Each dW element is owned by # exactly one workgroup (grid = N x K x E over disjoint output tiles), so when # ``accumulate`` is set the read-modify-write into the destination is race-free. + # ``out_ty`` is the ``dW`` element type: bf16 truncates the f32 accumulator on store; + # fp32 stores it directly (and accumulates without the bf16 round-trip). + out_ty = T.f32 if out_is_f32 else bf16 E_NK_row = arith.index_cast(T.index, expert) * N_idx * K_idx for mi in range_constexpr(M_STEPS): for nj in range_constexpr(N_STEPS): @@ -542,10 +571,11 @@ def _tail(j, accs): out_off = E_NK_row + n_out_idx * K_idx + c_n_idx if const_expr(accumulate): prev = buffer_ops.buffer_load( - dW_rsrc, out_off, vec_width=1, dtype=bf16 + dW_rsrc, out_off, vec_width=1, dtype=out_ty ) - val = arith.addf(val, arith.extf(T.f32, prev)) - store_val = arith.truncf(bf16, val) + prev_f32 = prev if const_expr(out_is_f32) else arith.extf(T.f32, prev) + val = arith.addf(val, prev_f32) + store_val = val if const_expr(out_is_f32) else arith.truncf(bf16, val) buffer_ops.buffer_store(store_val, dW_rsrc, out_off) scf.YieldOp([]) diff --git a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py index e7b66a4be..6b6ffb2ee 100644 --- a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py +++ b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py @@ -739,12 +739,15 @@ def permute_free_grouped_gemm_bf16_wgrad( weights_shape: ``(num_experts, out_features, in_features)``. out: - Optional ``[E, out, in]`` bf16 destination to fold the wgrad into directly - (``+=`` if ``accumulate`` else ``=``). + Optional ``[E, out, in]`` destination (bf16 or fp32) to fold the wgrad into directly + (``+=`` if ``accumulate`` else ``=``). An fp32 sink (e.g. ``main_grad``) accumulates + in-kernel via the fp32-store kernel variant -- no bf16 scratch + separate fold. accumulate: With ``out``: add into it rather than overwrite. swap_gather: - FC2 wgrad mode (token-gather on grad). Not implemented in the current FlyDSL kernel. + FC2 wgrad mode. Gathers ``grad_output`` (token-space ``[num_recv, out_features]``) and + walks ``hidden_states`` (block-padded ``[em_max, in_features]``) contiguously, emitting the + native ``[E, out_features, in_features]`` layout directly (no transpose). Returns ------- @@ -792,21 +795,15 @@ def permute_free_grouped_gemm_bf16_wgrad( routing = _prepare_wgrad_align(routing, contract_m) grad_base = (routing.block_start.to(torch.int64) * int(routing.block_size_m)).to(torch.int32) - if swap_gather: - raise RuntimeError( - "swap_gather wgrad (FC2 direct [E, H, F] layout) is not implemented in the " - "current FlyDSL wgrad kernel." - ) - flydsl_wgrad = _get_flydsl_wgrad() if out is not None: assert tuple(out.shape) == (num_experts, out_features, in_features), ( f"wgrad out shape {tuple(out.shape)} != {(num_experts, out_features, in_features)}" ) - if out.dtype != torch.bfloat16: + if out.dtype not in (torch.bfloat16, torch.float32): raise TypeError( - f"FlyDSL wgrad requires bf16 out, got {out.dtype}." + f"FlyDSL wgrad requires bf16 or fp32 out, got {out.dtype}." ) flydsl_wgrad( x, @@ -818,6 +815,7 @@ def permute_free_grouped_gemm_bf16_wgrad( grad_base, num_recv_tokens=routing.num_recv_tokens, accumulate=bool(accumulate), + swap_gather=bool(swap_gather), ) return out @@ -836,6 +834,7 @@ def permute_free_grouped_gemm_bf16_wgrad( grad_base, num_recv_tokens=routing.num_recv_tokens, accumulate=bool(accumulate), + swap_gather=bool(swap_gather), ) return dW @@ -1046,8 +1045,9 @@ def permute_free_grouped_gemm_bf16_fc2_wgrad( token-gathered, ``fc2_input`` is route-ordered and read contiguously. Which one the kernel gathers decides the output orientation: - * **FC2 wgrad** uses operand swap + transpose: compute ``[E, F, H]`` via the FC1 FlyDSL - kernel, then ``transpose(1, 2)`` to ``[E, H, F]``. + * **FC2 wgrad** uses ``swap_gather``: the FlyDSL kernel token-gathers ``grad_output`` and + walks the block-padded ``fc2_input`` contiguously, writing the native ``[E, H, F]`` layout + directly (no transpose, and ``out``/``accumulate`` fold straight into a bf16 destination). Parameters ---------- @@ -1078,17 +1078,33 @@ def permute_free_grouped_gemm_bf16_fc2_wgrad( preact, routing, activation=activation, dispatched_probs=dispatched_probs ) - # FC2 wgrad: operand swap + transpose via the FC1 FlyDSL kernel. - dW_t = permute_free_grouped_gemm_bf16_wgrad( - grad_output, + # FC2 wgrad via swap_gather: gather grad_output [num_recv, H] (N=H) and walk the block-padded + # fc2_input [em_max, F] (K=F) contiguously -> native [E, H, F], no transpose. + weights_shape = (num_experts, out_features, in_features) # [E, H, F] + if out is not None and out.dtype in (torch.bfloat16, torch.float32): + # bf16 or fp32 sink: the kernel writes/accumulates straight into it (fp32 via the + # fp32-store variant), so no scratch dW and no separate fold pass. + return permute_free_grouped_gemm_bf16_wgrad( + fc2_input, + grad_output, + weights_shape, + routing, + config=config, + out=out, + accumulate=accumulate, + swap_gather=True, + ) + dW = permute_free_grouped_gemm_bf16_wgrad( fc2_input, - (num_experts, in_features, out_features), + grad_output, + weights_shape, routing, config=config, - ) - dW = dW_t.transpose(1, 2).contiguous() # [E, H, F] + swap_gather=True, + ) # [E, H, F] bf16 if out is None: return dW + # Exotic sink dtype the kernel can't write directly: fold the bf16 result here. if accumulate: out.add_(dW) else: @@ -1258,21 +1274,12 @@ def permute_free_grouped_gemm_backward( dgrad = permute_free_grouped_gemm_bf16_dgrad(grad_output, weights_stacked, routing) if requires_wgrad: weights_shape = (num_gemms, weights[0].size(0), weights[0].size(1)) - if wgrad_out is not None and wgrad_out.dtype != torch.bfloat16: - # fp32 main_grad sink: the bf16 kernel can't write it directly, so emit the bf16 - # dW and fold it into the fp32 buffer here (mirrors the FC2 wgrad out handling). - dW = permute_free_grouped_gemm_bf16_wgrad( - hidden_states, grad_output, weights_shape, routing, - ) # [E, N, H] - if wgrad_accumulate: - wgrad_out.add_(dW) - else: - wgrad_out.copy_(dW) - else: - dW = permute_free_grouped_gemm_bf16_wgrad( - hidden_states, grad_output, weights_shape, routing, - out=wgrad_out, accumulate=wgrad_accumulate, - ) # [E, N, H] + # bf16 or fp32 ``main_grad`` sink both accumulate in-kernel (fp32 via the fp32-store + # variant); no scratch dW + separate fold. + dW = permute_free_grouped_gemm_bf16_wgrad( + hidden_states, grad_output, weights_shape, routing, + out=wgrad_out, accumulate=wgrad_accumulate, + ) # [E, N, H] wgrad_stacked = dW wgrad_applied = wgrad_out is not None diff --git a/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py b/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py index 102559179..9758cc5fd 100644 --- a/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py +++ b/transformer_engine/pytorch/moe/pf_wgrad_wrapper.py @@ -42,6 +42,7 @@ def flydsl_moe_wgrad( warps_n: int = 2, warps_k: int = 2, accumulate: bool = False, + swap_gather: bool = False, ) -> None: """Compute FC1 grouped wgrad ``grad[slot]^T @ x[token(slot)]`` into ``dw``, per expert. @@ -51,21 +52,28 @@ def flydsl_moe_wgrad( received-token row (sentinel ``num_recv_tokens`` for padding), and ``grad_base[e]`` is the block-padded first-slot row of expert ``e`` (``block_start[e] * block_size_m``). + ``swap_gather`` selects the FC2 wgrad variant: ``grad`` is the token-space ``[num_recv, N]`` + operand (gathered) and ``x`` is the block-padded ``[em_max, K]`` operand (contiguous walk), + emitting the native ``dw[E, N, K]`` orientation with no transpose. + ``block_n``/``block_k`` and ``warps_n``/``warps_k`` select the workgroup tile; the defaults (``128x128`` over ``2x2`` warps) are a strong general config on CDNA4. """ num_experts, N, K = dw.shape assert x.dtype == grad.dtype == torch.bfloat16 - assert dw.dtype == torch.bfloat16 + assert dw.dtype in (torch.bfloat16, torch.float32) assert x.is_contiguous() and grad.is_contiguous() and dw.is_contiguous() + out_dtype = "fp32" if dw.dtype == torch.float32 else "bf16" exe = compile_moe_wgrad_v2( block_n=int(block_n), block_k=int(block_k), warps_n=int(warps_n), warps_k=int(warps_k), accumulate=bool(accumulate), + swap_gather=bool(swap_gather), + out_dtype=out_dtype, ) _run_compiled( @@ -119,6 +127,7 @@ def _wgrad_run( warps_n=2, warps_k=2, accumulate=False, + swap_gather=False, ): """Dispatch target for the FlyDSL autotuner: compile (lru-cached) + launch one tile.""" exe = compile_moe_wgrad_v2( @@ -127,6 +136,7 @@ def _wgrad_run( warps_n=int(warps_n), warps_k=int(warps_k), accumulate=bool(accumulate), + swap_gather=bool(swap_gather), ) _run_compiled( exe, @@ -194,17 +204,20 @@ def flydsl_moe_wgrad_autotuned( *, num_recv_tokens: int, accumulate: bool = False, + swap_gather: bool = False, ) -> None: """Shape-autotuned variant of :func:`flydsl_moe_wgrad`. First call for a given ``(x.shape, N, num_experts)`` benchmarks every tile in ``_AUTOTUNE_TILES`` and caches the fastest (in-memory + on disk under - ``~/.flydsl/autotune/``). + ``~/.flydsl/autotune/``). The tile geometry is independent of ``swap_gather`` (the two + variants share the same memory-access shape), so the sweep runs on the default variant and + the fastest tile is reused for the ``swap_gather`` launch. """ num_experts, N, K = dw.shape assert x.dtype == grad.dtype == torch.bfloat16 - assert dw.dtype == torch.bfloat16 + assert dw.dtype in (torch.bfloat16, torch.float32) assert x.is_contiguous() and grad.is_contiguous() and dw.is_contiguous() block_n, block_k, warps_n, warps_k = _select_wgrad_config( @@ -225,4 +238,5 @@ def flydsl_moe_wgrad_autotuned( warps_n=warps_n, warps_k=warps_k, accumulate=bool(accumulate), + swap_gather=bool(swap_gather), ) From 5e694a2ad1ec5aeefb90477c40dfb30f73fd636c Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Wed, 5 Aug 2026 22:27:33 +0000 Subject: [PATCH 36/43] Decouple gated activation from FlyDSL permute-free grouped GEMM and add FC2 path Removes fused activation, dispatched_probs, and preact_out handling from the FlyDSL forward wrapper so all plain gather/route-read GEMMs use the v3 MegaMOE-ported kernel. Gated SiLU activation and route-prob scaling now use standalone Triton helpers (permute_free_gated_act_recompute/bwd). Adds permute-free FC2 forward and dgrad variants (permute_free_grouped_gemm_bf16_fc2 and _fc2_dgrad), simplifies block-size selection for the plain-GEMM path, and removes the obsolete get_default_moe_kernel_config helper. Updates the benchmark and unit tests to exercise the standalone activation and FC2 kernels. --- .../benchmark_perm_free_grouped_gemm.py | 10 +- .../pytorch/test_perm_free_grouped_linear.py | 271 +++++++++++++++-- transformer_engine/pytorch/moe/__init__.py | 2 - .../pytorch/moe/permute_free_grouped_gemm.py | 278 +++++------------- .../pytorch/moe/pf_fwd_wrapper.py | 106 ++----- 5 files changed, 354 insertions(+), 313 deletions(-) diff --git a/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py index 7cd991a97..e954c7575 100644 --- a/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py +++ b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py @@ -182,17 +182,17 @@ def _permute_free_act_gemm( weights: torch.Tensor, routing, ) -> torch.Tensor: - """Permute-free FC1 gather-GEMM with the gated SiLU activation fused into the epilogue. + """Permute-free FC1 gather-GEMM + standalone SiLU gated activation (no in-kernel fusion). - ``weights`` is the gate+up projection ``[E, 2F, K]``; the kernel emits the F-wide - activated buffer ``silu(gate) * up`` directly, skipping the separate activation pass. - Only valid for gate+up (GateUP) shapes -- ``out_features`` must be even. + ``weights`` is the gate+up projection ``[E, 2F, K]``; returns the F-wide activated buffer. """ from transformer_engine.pytorch.moe import ( + permute_free_gated_act_recompute, permute_free_grouped_gemm_bf16, ) - return permute_free_grouped_gemm_bf16(hidden, weights, routing, activation="silu") + preact = permute_free_grouped_gemm_bf16(hidden, weights, routing) + return permute_free_gated_act_recompute(preact, routing, activation="silu") def _build_backend_fns( diff --git a/tests/pytorch/test_perm_free_grouped_linear.py b/tests/pytorch/test_perm_free_grouped_linear.py index 36088af9f..ed715bfff 100644 --- a/tests/pytorch/test_perm_free_grouped_linear.py +++ b/tests/pytorch/test_perm_free_grouped_linear.py @@ -13,12 +13,16 @@ from transformer_engine.pytorch.moe import ( MoERoutingMetadata, PermuteFreeMetadata, - get_default_moe_kernel_config, + is_permute_free_grouped_gemm_enabled, + permute_free_gated_act_bwd, permute_free_grouped_gemm_backward, permute_free_grouped_gemm_bf16, permute_free_grouped_gemm_bf16_dgrad, + permute_free_grouped_gemm_bf16_fc2, + permute_free_grouped_gemm_bf16_fc2_dgrad, permute_free_grouped_gemm_bf16_fc2_wgrad, permute_free_grouped_gemm_bf16_wgrad, + permute_free_grouped_gemm_forward, prepare_moe_align, ) from transformer_engine.pytorch.moe.permute_free_grouped_gemm import ( @@ -234,10 +238,7 @@ def test_route_list_fc2_wgrad_recompute_from_preact(): """FC2 wgrad recompute-from-preact: rebuild act = act(gate)*up*prob from the saved 2F pre-activation into a transient buffer, feed the unchanged wgrad, and match a stored-act reference (the backward then checkpoints only the 2F preact, never the F-wide act).""" - from transformer_engine.pytorch.moe import ( - get_default_moe_kernel_config, - prepare_moe_align, - ) + from transformer_engine.pytorch.moe import prepare_moe_align torch.manual_seed(29) num_recv_tokens, in_features, out_features = 128, 96, 128 # F=in, H=out (W2 is [E, H, F]) @@ -298,7 +299,6 @@ def test_fc2_backward_dispatch_recompute_matches_stored(): """FC2 backward dispatch: wgrad from the saved 2F preact (+ fc2_activation) matches the legacy stored-F activation path.""" from transformer_engine.pytorch.moe import ( - get_default_moe_kernel_config, permute_free_grouped_gemm_backward, prepare_moe_align, ) @@ -366,10 +366,7 @@ def test_fc1_fc2_gated_pipeline(monkeypatch): PyTorch reference through the permute-free GroupedLinear modules.""" import dataclasses - from transformer_engine.pytorch.moe import ( - get_default_moe_kernel_config, - prepare_moe_align, - ) + from transformer_engine.pytorch.moe import prepare_moe_align monkeypatch.setenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "1") torch.manual_seed(37) @@ -502,15 +499,256 @@ def test_route_list_fwd_dgrad_wgrad_consistency(): assert _rel_l2(dW, ref_w.grad) < 2e-2 +def test_route_list_fc2_fwd(): + """FC2 forward (F-wide input): per-route GEMM + token gather-combine.""" + torch.manual_seed(51) + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=21) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + fc2_input = torch.zeros(em_max, in_features, device="cuda", dtype=torch.bfloat16) + fc2_input[valid] = torch.randn( + int(valid.sum().item()), in_features, device="cuda", dtype=torch.bfloat16 + ) + weights = torch.randn( + num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16 + ) + + out = permute_free_grouped_gemm_bf16_fc2(fc2_input, weights, routing) + assert out.shape == (num_recv_tokens, out_features) + + per_slot = torch.einsum("rf,rhf->rh", fc2_input.float(), weights[exp].float()) + per_slot = per_slot * valid[:, None].float() + ref = torch.zeros(num_recv_tokens, out_features, device="cuda", dtype=torch.float32) + ref.index_add_(0, tok, per_slot) + assert _rel_l2(out, ref) < 2e-2 + + +def test_route_list_fc2_fwd_standalone_act(): + """FC2 forward dispatch: standalone gated-act recompute + plain GEMM (no in-kernel fusion).""" + torch.manual_seed(53) + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=23) + routing = PermuteFreeMetadata( + routing_map=routing_map, num_experts=num_experts, route_space=True, activation="silu" + ) + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + preact = torch.zeros(em_max, 2 * in_features, device="cuda", dtype=torch.bfloat16) + preact[valid] = torch.randn( + int(valid.sum().item()), 2 * in_features, device="cuda", dtype=torch.bfloat16 + ) + probs = torch.rand(num_recv_tokens, num_experts, device="cuda", dtype=torch.float32) + weights = torch.randn( + num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16 + ) + + out = permute_free_grouped_gemm_forward( + preact, weights, routing, activation="silu", dispatched_probs=probs + ).out + assert out.shape == (num_recv_tokens, out_features) + + g = preact[:, :in_features].float() + u = preact[:, in_features:].float() + act = torch.nn.functional.silu(g) * u * probs[tok, exp][:, None] * valid[:, None].float() + per_slot = torch.einsum("rf,rhf->rh", act, weights[exp].float()) * valid[:, None].float() + ref = torch.zeros(num_recv_tokens, out_features, device="cuda", dtype=torch.float32) + ref.index_add_(0, tok, per_slot) + assert _rel_l2(out, ref) < 2e-2 + + +def test_route_list_fc2_dgrad(): + """FC2 dgrad: token-space grad gathered per route into block-padded ``[em_max, F]``.""" + torch.manual_seed(57) + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=25) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + grad_output = torch.randn(num_recv_tokens, out_features, device="cuda", dtype=torch.bfloat16) + weights = torch.randn( + num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16 + ) + + dgrad = permute_free_grouped_gemm_bf16_fc2_dgrad(grad_output, weights, routing) + assert dgrad.shape == (em_max, in_features) + + per_slot = torch.einsum("rh,rhf->rf", grad_output[tok].float(), weights[exp].float()) + assert _rel_l2(dgrad[valid], per_slot[valid]) < 2e-2 + + +def test_route_list_gated_act_bwd(): + """Standalone gated-activation backward: dpre + dprob vs autograd reference.""" + torch.manual_seed(59) + num_recv_tokens, in_features = 128, 128 + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=27) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + gate = torch.randn(em_max, in_features, device="cuda", dtype=torch.float32) + up = torch.randn(em_max, in_features, device="cuda", dtype=torch.float32) + gate = (gate * valid[:, None].float()).requires_grad_(True) + up = (up * valid[:, None].float()).requires_grad_(True) + probs = torch.rand( + num_recv_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True + ) + grad_out = torch.zeros(em_max, in_features, device="cuda", dtype=torch.bfloat16) + grad_out[valid] = torch.randn( + int(valid.sum().item()), in_features, device="cuda", dtype=torch.bfloat16 + ) + + pr = probs[tok, exp] + act = torch.nn.functional.silu(gate) * up * pr[:, None] * valid[:, None].float() + (act * grad_out.float()).sum().backward() + dpre_ref = torch.cat([gate.grad, up.grad], dim=1) + + preact = torch.cat([gate.detach(), up.detach()], dim=1).to(torch.bfloat16) + dpre, dprob = permute_free_gated_act_bwd( + grad_out, preact, routing, activation="silu", dispatched_probs=probs.detach() + ) + assert dpre.shape == (em_max, 2 * in_features) + assert dprob.shape == (num_recv_tokens, num_experts) + assert _rel_l2(dpre[valid], dpre_ref[valid]) < 2e-2 + assert _rel_l2(dprob, probs.grad) < 2e-2 + + +def test_route_list_wgrad_out_accumulate(): + """FC1 wgrad in-kernel fold into bf16/fp32 ``out`` (overwrite + accumulate).""" + torch.manual_seed(61) + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=29) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + + hidden = torch.randn(num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16) + weights_shape = (num_experts, out_features, in_features) + grad = torch.zeros(em_max, out_features, device="cuda", dtype=torch.bfloat16) + grad[valid] = torch.randn( + int(valid.sum().item()), out_features, device="cuda", dtype=torch.bfloat16 + ) + + per_slot = torch.einsum("rn,rk->rnk", grad.float(), hidden[tok].float()) + per_slot = per_slot * valid[:, None, None].float() + ref = torch.zeros(num_experts, out_features, in_features, device="cuda", dtype=torch.float32) + ref.index_add_(0, slot_expert.to(torch.int64).clamp_min(0), per_slot) + + for dtype in (torch.bfloat16, torch.float32): + out = torch.zeros(weights_shape, device="cuda", dtype=dtype) + ret = permute_free_grouped_gemm_bf16_wgrad( + hidden, grad, weights_shape, routing, out=out, accumulate=False + ) + assert ret is out + assert _rel_l2(out, ref) < 2e-2 + permute_free_grouped_gemm_bf16_wgrad( + hidden, grad, weights_shape, routing, out=out, accumulate=True + ) + assert _rel_l2(out, 2.0 * ref) < 2e-2 + + +def test_permute_free_forward_dispatch(): + """``permute_free_grouped_gemm_forward`` routes FC1 plain GEMM and FC2 standalone-act paths.""" + torch.manual_seed(63) + num_recv_tokens, hidden_dim, ffn = 128, 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=31) + + fc1_routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + fc1_routing = prepare_moe_align(fc1_routing, _FLYDSL_FWD_BLOCK_M) + hidden = torch.randn(num_recv_tokens, hidden_dim, device="cuda", dtype=torch.bfloat16) + w1 = torch.randn(num_experts, 2 * ffn, hidden_dim, device="cuda", dtype=torch.bfloat16) + res_fc1 = permute_free_grouped_gemm_forward(hidden, w1, fc1_routing) + direct = permute_free_grouped_gemm_bf16(hidden, w1, fc1_routing) + valid_fc1 = fc1_routing.sorted_slot_ids < num_recv_tokens + assert _rel_l2(res_fc1.out[valid_fc1], direct[valid_fc1]) < 1e-3 + + fc2_routing = PermuteFreeMetadata( + routing_map=routing_map, num_experts=num_experts, route_space=True, activation="silu" + ) + prepare_moe_align(fc2_routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(fc2_routing.sorted_slot_ids.shape[0]) + slot_token = fc2_routing.sorted_slot_ids + slot_expert = _expert_per_route(fc2_routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + preact = torch.zeros(em_max, 2 * ffn, device="cuda", dtype=torch.bfloat16) + preact[valid] = torch.randn( + int(valid.sum().item()), 2 * ffn, device="cuda", dtype=torch.bfloat16 + ) + probs = torch.rand(num_recv_tokens, num_experts, device="cuda", dtype=torch.float32) + w2 = torch.randn(num_experts, hidden_dim, ffn, device="cuda", dtype=torch.bfloat16) + out = permute_free_grouped_gemm_forward( + preact, w2, fc2_routing, activation="silu", dispatched_probs=probs + ).out + assert out.shape == (num_recv_tokens, hidden_dim) + + g = preact[:, :ffn].float() + u = preact[:, ffn:].float() + act = torch.nn.functional.silu(g) * u * probs[tok, exp][:, None] * valid[:, None].float() + per_slot = torch.einsum("rf,rhf->rh", act, w2[exp].float()) * valid[:, None].float() + ref = torch.zeros(num_recv_tokens, hidden_dim, device="cuda", dtype=torch.float32) + ref.index_add_(0, tok, per_slot) + assert _rel_l2(out, ref) < 2e-2 + + +def test_is_permute_free_grouped_gemm_enabled(monkeypatch): + """Env gate: True only on ROCm with ``NVTE_PERMUTE_FREE_GROUPED_GEMM=1``.""" + monkeypatch.setenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "1") + assert is_permute_free_grouped_gemm_enabled() is True + monkeypatch.setenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "0") + assert is_permute_free_grouped_gemm_enabled() is False + monkeypatch.delenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", raising=False) + assert is_permute_free_grouped_gemm_enabled() is False + + @pytest.mark.skip(reason="apply_route_probs not ported to transformer_engine.pytorch.moe yet") def test_apply_route_probs_fwd_bwd(): """Fused per-route prob apply (gather+multiply) vs an autograd advanced-index reference.""" - from transformer_engine.pytorch.moe import ( - get_default_moe_kernel_config, - prepare_moe_align, - ) + from transformer_engine.pytorch.moe import prepare_moe_align # apply_route_probs was not ported; kept for when route_prob helper lands in moe/. - _ = get_default_moe_kernel_config, prepare_moe_align + _ = prepare_moe_align torch.manual_seed(31) num_recv_tokens, hidden_dim = 128, 192 @@ -518,8 +756,7 @@ def test_apply_route_probs_fwd_bwd(): routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=9) routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) - cfg = get_default_moe_kernel_config(num_recv_tokens) - prepare_moe_align(routing, int(cfg["BLOCK_SIZE_M"])) + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) em_max = int(routing.sorted_slot_ids.shape[0]) num_routes = int(routing_map.sum().item()) diff --git a/transformer_engine/pytorch/moe/__init__.py b/transformer_engine/pytorch/moe/__init__.py index fee15bb56..ef5a96fb0 100644 --- a/transformer_engine/pytorch/moe/__init__.py +++ b/transformer_engine/pytorch/moe/__init__.py @@ -9,7 +9,6 @@ PermuteFreeBackwardResult, PermuteFreeForwardResult, PermuteFreeMetadata, - get_default_moe_kernel_config, is_permute_free_grouped_gemm_enabled, permute_free_grouped_gemm_backward, permute_free_grouped_gemm_bf16, @@ -29,7 +28,6 @@ "PermuteFreeBackwardResult", "PermuteFreeForwardResult", "PermuteFreeMetadata", - "get_default_moe_kernel_config", "is_permute_free_grouped_gemm_enabled", "permute_free_grouped_gemm_backward", "permute_free_grouped_gemm_bf16", diff --git a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py index 6b6ffb2ee..6f9b480e1 100644 --- a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py +++ b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py @@ -14,9 +14,8 @@ from __future__ import annotations import os -import warnings from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple +from typing import List, Optional, Tuple import torch @@ -40,14 +39,13 @@ "PermuteFreeForwardResult", "PermuteFreeBackwardResult", "prepare_moe_align", - "get_default_moe_kernel_config", "is_permute_free_grouped_gemm_enabled", ] _WGRAD_CONTRACT_M = 32 -# Max forward align/kernel ``block_m`` for fwd/dgrad FlyDSL Permute-free Grouped GEMM. -_FLYDSL_GATED_BLOCK_M = 256 +# Minimum v3 gather/dgrad block_m (128x256 MFMA floor). +_FLYDSL_MIN_BLOCK_M = 128 _FLYDSL_FWD_BLOCK_M = 256 _FLYDSL_FWD_LARGE_TIER = 128 @@ -94,16 +92,12 @@ def _pf_moe_fwd( routing: MoERoutingMetadata, *, num_recv_tokens: int, - config: Dict[str, Any], + block_m: int, index_a_by_route_pos: bool = False, - activation: Optional[str] = None, - dispatched_probs: Optional[torch.Tensor] = None, - preact_out: Optional[torch.Tensor] = None, - gated_a: bool = False, ) -> None: - """Run the route-list gather-GEMM (forward or dgrad) via FlyDSL.""" + """Run the route-list gather-GEMM (forward or dgrad) via FlyDSL autotuning.""" autotuned, supported = _get_flydsl_fwd() - block_m = int(config["BLOCK_SIZE_M"]) + block_m = int(block_m) if not supported(A, B, block_m=block_m): raise RuntimeError( "FlyDSL grouped GEMM does not support these operands " @@ -119,10 +113,6 @@ def _pf_moe_fwd( num_recv_tokens=num_recv_tokens, block_m=block_m, index_a_by_route_pos=index_a_by_route_pos, - activation=activation, - dispatched_probs=dispatched_probs, - preact_out=preact_out, - gated_a=gated_a, ) @@ -130,10 +120,7 @@ def _fwd_align_block_size_m( A: torch.Tensor, B: torch.Tensor, *, - gated: bool, - default_block_m: int, num_tokens: int, - gated_a: bool = False, ) -> int: """Pick the forward align/kernel ``block_m``, favoring the backend that will run. @@ -141,34 +128,25 @@ def _fwd_align_block_size_m( layer's FC1 fwd, FC1 dgrad and FC2 fwd (they reuse ``routing.block_size_m``), so it is chosen once here. When the FlyDSL backend will run and the workload is on the large-token tier (``num_tokens >= _FLYDSL_FWD_LARGE_TIER``), prefer :data:`_FLYDSL_FWD_BLOCK_M` (256): it lifts - the non-gated FC1 fwd, FC1 dgrad, and ``gated_a`` FC2 onto the faster ``256x32`` tile. The - gated FC1 epilogue stages a ``2F`` [gate|up] B-tile that overflows LDS at 256, so it drops back - to :data:`_FLYDSL_GATED_BLOCK_M` (128). Small/medium tiers keep the token-count default. + FC1 fwd, FC1 dgrad, and FC2 onto the faster ``256x32`` tile. Small/medium tiers default to + :data:`_FLYDSL_MIN_BLOCK_M` (128), the v3 gather/dgrad floor. """ - if gated: - # The gated FC1 epilogue stages a 2F [gate|up] B-tile: even the LDS-valid 256 tiles are - # the slow tiny ones, so the measured-best policy is a hard cap at _FLYDSL_GATED_BLOCK_M - # (~1.35x at 128 vs ~1.06x at 256). Never bump; only allow the picker to confirm/keep a - # value <= 128 (dropping the default when it exceeds the cap). - cap = min(default_block_m, _FLYDSL_GATED_BLOCK_M) - candidates = {c for c in (_FLYDSL_GATED_BLOCK_M, default_block_m) if c <= cap} - else: - # Non-gated FC1 / dgrad / gated_a FC2: offer the default and any smaller floor (so the - # picker can drop to a smaller LDS-valid block_m), plus the 256 bump on the large-token - # tier, where it is a measured win and the extra E_local*128 align padding is a negligible - # fraction of the routed work. - candidates = {c for c in (_FLYDSL_GATED_BLOCK_M, default_block_m) if c <= default_block_m} - if num_tokens >= _FLYDSL_FWD_LARGE_TIER: - candidates.add(_FLYDSL_FWD_BLOCK_M) - # The v3 gather/dgrad tile requires block_m >= 128 (128x256 MFMA minimum), so on a - # small-token tier where the token-count default is < 128 we must still floor the align - # block_m to 128 rather than emit a sub-tile the kernel cannot run. - from .pf_fwd_wrapper import _v3_enabled - - if _v3_enabled(): - candidates = {c for c in candidates if c >= _FLYDSL_GATED_BLOCK_M} or { - _FLYDSL_GATED_BLOCK_M - } + default_block_m = ( + _FLYDSL_FWD_BLOCK_M if num_tokens >= _FLYDSL_FWD_LARGE_TIER else _FLYDSL_MIN_BLOCK_M + ) + # Offer the default and any smaller floor (so the picker can drop to a smaller LDS-valid + # block_m), plus the 256 bump on the large-token tier, where it is a measured win and the + # extra E_local*128 align padding is a negligible fraction of the routed work. + candidates = {c for c in (_FLYDSL_MIN_BLOCK_M, default_block_m) if c <= default_block_m} + if num_tokens >= _FLYDSL_FWD_LARGE_TIER: + candidates.add(_FLYDSL_FWD_BLOCK_M) + from .pf_fwd_wrapper import _v3_enabled + + # The v3 gather/dgrad tile requires block_m >= 128 (128x256 MFMA minimum), so on a + # small-token tier where the default would be < 128 we must still floor the align + # block_m to 128 rather than emit a sub-tile the kernel cannot run. + if _v3_enabled(): + candidates = {c for c in candidates if c >= _FLYDSL_MIN_BLOCK_M} or {_FLYDSL_MIN_BLOCK_M} # Prefer the in-tree FlyDSL block_m picker when available. try: @@ -176,60 +154,31 @@ def _fwd_align_block_size_m( except Exception: # pylint: disable=broad-except flydsl_moe_fwd_pick_block_m = None if flydsl_moe_fwd_pick_block_m is not None: - picked = flydsl_moe_fwd_pick_block_m( - A, B, gated=gated, gated_a=gated_a, candidates=tuple(candidates) - ) + picked = flydsl_moe_fwd_pick_block_m(A, B, candidates=tuple(candidates)) if picked is not None: return picked - - # Legacy fallback: cap the gated fwd to 128, else default. - if not gated or default_block_m <= _FLYDSL_GATED_BLOCK_M: - return default_block_m - _, supported = _get_flydsl_fwd() - if supported(A, B, block_m=_FLYDSL_GATED_BLOCK_M): - return _FLYDSL_GATED_BLOCK_M return default_block_m +def _ensure_fwd_align( + routing: MoERoutingMetadata, + A: torch.Tensor, + B: torch.Tensor, +) -> int: + """Return ``routing.block_size_m``, building fwd align buffers when missing.""" + if routing.sorted_slot_ids is not None and routing.block_size_m is not None: + return int(routing.block_size_m) + block_m = _fwd_align_block_size_m(A, B, num_tokens=routing.num_recv_tokens) + prepare_moe_align(routing, block_m) + return block_m + + def is_permute_free_grouped_gemm_enabled() -> bool: from torch.utils.cpp_extension import IS_HIP_EXTENSION return IS_HIP_EXTENSION and os.getenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "0") == "1" -def get_default_moe_kernel_config(num_tokens: int) -> Dict[str, Any]: - """Return align ``block_m`` tile config for the route-list bf16 MoE path.""" - # aiter's tuned config is an optional accelerator. - try: - from aiter.ops.triton.utils.moe_config_utils import get_optimal_moe_config - - return get_optimal_moe_config(torch.bfloat16, M=num_tokens) - except Exception: # pylint: disable=broad-except - pass - - if num_tokens <= 32: - block_m = 16 - elif num_tokens <= 96: - block_m = 32 - elif num_tokens <= 512: - block_m = 64 - else: - block_m = 128 - - block_n = 64 if num_tokens <= 64 else 128 - block_k = 128 if num_tokens <= 64 else 64 - group_m = 16 if num_tokens // max(block_m, 1) > 128 else 1 - - return { - "BLOCK_SIZE_M": block_m, - "BLOCK_SIZE_N": block_n, - "BLOCK_SIZE_K": block_k, - "GROUP_SIZE_M": group_m, - "num_warps": 4 if num_tokens <= 128 else 8, - "num_stages": 2, - } - - def moe_align_route_list( routing_map: torch.Tensor, *, @@ -415,16 +364,13 @@ def permute_free_grouped_gemm_bf16( hidden_states: torch.Tensor, weights: torch.Tensor | list[torch.Tensor], routing: MoERoutingMetadata, - *, - config: Optional[Dict[str, Any]] = None, - activation: Optional[str] = None, - dispatched_probs: Optional[torch.Tensor] = None, - return_preact: bool = False, ) -> torch.Tensor: """Route-list gather-in-GEMM (FC1 forward) for bf16 MoE. Computes ``C[slot] = A[sorted_slot_ids[slot]] @ W[expert(slot)]^T`` for each block-padded - slot, writing directly into the ``[em_max, out_features]`` block-padded output. + slot, writing directly into the ``[em_max, out_features]`` block-padded output. Gated + activation is **not** fused here; FC1 emits the raw ``2F`` ``[gate | up]`` pre-activation + and the standalone gated-act helpers apply ``act(gate) * up [* prob]`` on FC2. Parameters ---------- @@ -434,28 +380,14 @@ def permute_free_grouped_gemm_bf16( Expert weights ``[num_experts, out_features, in_features]`` or list of ``[out, in]``. routing: ``MoERoutingMetadata`` carrying (or able to build) the route-list align buffers. - activation: - When set (``"silu"`` / ``"gelu"``), fuses the **gated** activation into the GEMM - epilogue: ``out_features`` is the gate+up width (``2F``), the output is the ``F``-wide - activated buffer ``act(gate) * up``, and the separate activation kernel is skipped. The - weight output dim must be laid out as ``[gate | up]``. - dispatched_probs: - Optional ``[num_recv_tokens, num_experts]`` gating probabilities. When given (with a - fused ``activation``), each route's ``prob[token, expert]`` is multiplied into the - activation in-kernel, skipping the separate route-prob-apply pass. - return_preact: - When set (only valid with a fused ``activation``), also allocate and fill a - ``[T * min(topk, E), 2F]`` buffer with the raw ``[gate | up]`` pre-activation and return it - alongside the activated output. The backward needs it to reconstruct the 2F GEMM-output - gradient; keep ``False`` for inference to skip the extra store. Returns ------- torch.Tensor - Block-padded ``[em_max, out_features]`` (or ``[em_max, F]`` with a fused ``activation``), - bf16 (expert-contiguous, each expert padded to ``block_size``). The valid rows are the - block-padded slots whose ``sorted_slot_ids[slot] < num_recv_tokens``; the padding slots - carry inert dead values. Consumers locate each expert's rows via the routing metadata: + Block-padded ``[em_max, out_features]``, bf16 (expert-contiguous, each expert padded to + ``block_size``). The valid rows are the block-padded slots whose + ``sorted_slot_ids[slot] < num_recv_tokens``; the padding slots carry inert dead values. + Consumers locate each expert's rows via the routing metadata: - ``block_start[e]`` (block units): expert ``e``'s rows occupy padded slots ``[block_start[e] * block_size, ...)``, with within-rank offset matching each route. @@ -484,26 +416,7 @@ def permute_free_grouped_gemm_bf16( f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." ) - gated_act = activation is not None - if gated_act and out_features % 2 != 0: - raise ValueError( - f"Gated activation requires an even (gate+up) out_features, got {out_features}." - ) - # Fused gated activation halves the stored width (2F -> F). - stored_features = out_features // 2 if gated_act else out_features - - # Copy so the backend-aware block_m override never mutates the caller's config dict. - kernel_config = dict(config or get_default_moe_kernel_config(num_recv_tokens)) - block_size_m = _fwd_align_block_size_m( - hidden_states, - weights_stacked, - gated=gated_act, - default_block_m=int(kernel_config["BLOCK_SIZE_M"]), - num_tokens=num_recv_tokens, - ) - kernel_config["BLOCK_SIZE_M"] = block_size_m - - routing = prepare_moe_align(routing, block_size_m) + block_size_m = _ensure_fwd_align(routing, hidden_states, weights_stacked) # Worst-case (sync-free) allocation: size the output to the block-padded upper bound # em_max = sorted_slot_ids.shape[0], which is derived purely from shapes (num_recv_tokens, @@ -511,33 +424,20 @@ def permute_free_grouped_gemm_bf16( # (.item()) that a compact [num_routes, N] allocation would require. em_max = routing.sorted_slot_ids.shape[0] output = torch.empty( - (em_max, stored_features), + (em_max, out_features), dtype=torch.bfloat16, device=hidden_states.device, ) - if return_preact and not gated_act: - raise ValueError("return_preact requires a fused activation.") - preact = ( - torch.empty((em_max, out_features), dtype=torch.bfloat16, device=hidden_states.device) - if return_preact - else None - ) - _pf_moe_fwd( hidden_states, weights_stacked, output, routing, num_recv_tokens=num_recv_tokens, - config=kernel_config, + block_m=block_size_m, index_a_by_route_pos=False, - activation=activation, - dispatched_probs=dispatched_probs, - preact_out=preact, ) - if return_preact: - return output, preact return output @@ -635,8 +535,6 @@ def permute_free_grouped_gemm_bf16_dgrad( grad_output: torch.Tensor, weights: torch.Tensor | list[torch.Tensor], routing: MoERoutingMetadata, - *, - config: Optional[Dict[str, Any]] = None, ) -> torch.Tensor: """Route-list gather-in-GEMM dgrad (FC1 backward wrt input). @@ -677,13 +575,7 @@ def permute_free_grouped_gemm_bf16_dgrad( weights_stacked = weights_stacked.contiguous() weights_t = weights_stacked.transpose(1, 2) - if routing.sorted_slot_ids is None or routing.block_size_m is None: - fwd_config = config or get_default_moe_kernel_config(routing.num_recv_tokens) - routing = prepare_moe_align(routing, int(fwd_config["BLOCK_SIZE_M"])) - block_size_m = int(routing.block_size_m) - - dgrad_config = get_default_moe_kernel_config(int(grad_output.shape[0])) - dgrad_config = {**dgrad_config, "BLOCK_SIZE_M": block_size_m} + block_size_m = _ensure_fwd_align(routing, grad_output, weights_t) # Two-stage dgrad reduction (contention-free): (1) plain compact store of each per-route # dX[route] = grad[route] @ W1[e] into an [T * min(topk, E), in] bf16 buffer (coalesced, no atomics), @@ -701,7 +593,7 @@ def permute_free_grouped_gemm_bf16_dgrad( compact, routing, num_recv_tokens=routing.num_recv_tokens, - config=dgrad_config, + block_m=block_size_m, index_a_by_route_pos=True, ) return route_gather_combine( @@ -719,7 +611,6 @@ def permute_free_grouped_gemm_bf16_wgrad( weights_shape, routing: MoERoutingMetadata, *, - config: Optional[Dict[str, Any]] = None, out: Optional[torch.Tensor] = None, accumulate: bool = False, swap_gather: bool = False, @@ -790,8 +681,11 @@ def permute_free_grouped_gemm_bf16_wgrad( # can pass that padded per-expert base to the kernel; routes are contiguous within an expert, # so base + within-rank indexes the correct padded grad row. if routing.block_start is None or routing.block_size_m is None: - fwd_config = config or get_default_moe_kernel_config(routing.num_recv_tokens) - routing = prepare_moe_align(routing, int(fwd_config["BLOCK_SIZE_M"])) + stub_w = torch.empty( + num_experts, out_features, in_features, + device=hidden_states.device, dtype=torch.bfloat16, + ) + _ensure_fwd_align(routing, hidden_states, stub_w) routing = _prepare_wgrad_align(routing, contract_m) grad_base = (routing.block_start.to(torch.int64) * int(routing.block_size_m)).to(torch.int32) @@ -843,34 +737,21 @@ def permute_free_grouped_gemm_bf16_fc2( fc2_input: torch.Tensor, weights: torch.Tensor | list[torch.Tensor], routing: MoERoutingMetadata, - *, - config: Optional[Dict[str, Any]] = None, - activation: Optional[str] = None, - dispatched_probs: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Route-list FC2 forward with gather-combine to token order (bf16 MoE). - ``fc2_input`` is the route-ordered FC1 output. When ``activation`` is set (``"silu"`` / - ``"gelu"``) it is the raw ``2F`` ``[gate | up]`` pre-activation and the kernel fuses - ``act(gate) * up [* prob]`` on the ``A`` operand before the GEMM; otherwise it is the - ``F``-wide activated buffer (legacy path). For each route the kernel reads its row - (``index_a_by_route_pos=True``) and computes the compact per-route GEMM; a separate - contention-free gather-combine pass sums each token's routes into its output row. + ``fc2_input`` is the route-ordered ``F``-wide FC1 activation (or a transient buffer rebuilt + from the saved ``2F`` pre-activation via :func:`permute_free_gated_act_recompute`). For each + route the kernel reads its row (``index_a_by_route_pos=True``) and computes the compact + per-route GEMM; a separate contention-free gather-combine pass sums each token's routes + into its output row. Parameters ---------- fc2_input: - Route-ordered activations ``[T * min(topk, E), in_features]`` (``F``) or, when - ``activation`` is set, the raw ``2F`` ``[gate | up]`` pre-activation from FC1. + Route-ordered activations ``[em_max, in_features(F)]``, bf16. weights: - Expert weights ``[num_experts, out_features, in_features]`` (W2) or list of ``[out, - in]``. - activation: - When set, fuse ``act(gate) * up`` (+ optional ``dispatched_probs``) on the ``A`` - operand (FC2 prologue). Requires ``fc2_input.shape[-1] == 2 * in_features``. - dispatched_probs: - Optional ``[num_recv_tokens, num_experts]`` gating probabilities (fused in the FC2 - prologue when ``activation`` is set). + Expert weights ``[num_experts, out_features(H), in_features(F)]`` (W2) or list of ``[H, F]``. routing: ``MoERoutingMetadata`` carrying (or able to build) the route-list align buffers. @@ -892,14 +773,7 @@ def permute_free_grouped_gemm_bf16_fc2( weights_stacked = weights_stacked.contiguous() num_experts, out_features, in_features = weights_stacked.shape - gated_a = activation is not None - if gated_a: - if fc2_input.shape[-1] != 2 * in_features: - raise ValueError( - f"gated FC2 expects 2F preact input ({2 * in_features} cols), " - f"got {fc2_input.shape[-1]}." - ) - elif fc2_input.shape[-1] != in_features: + if fc2_input.shape[-1] != in_features: raise ValueError( f"fc2_input in_features ({fc2_input.shape[-1]}) does not match weights " f"({in_features})." @@ -909,12 +783,7 @@ def permute_free_grouped_gemm_bf16_fc2( f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." ) - if routing.sorted_slot_ids is None or routing.block_size_m is None: - fwd_config = config or get_default_moe_kernel_config(routing.num_recv_tokens) - routing = prepare_moe_align(routing, int(fwd_config["BLOCK_SIZE_M"])) - block_size_m = int(routing.block_size_m) - kernel_config = get_default_moe_kernel_config(int(fc2_input.shape[0])) - kernel_config = {**kernel_config, "BLOCK_SIZE_M": block_size_m} + block_size_m = _ensure_fwd_align(routing, fc2_input, weights_stacked) # Two-stage combine (contention-free): (1) plain compact store of each per-route result # y[route] = fc2_input[route] @ W2[e] into an [T * min(topk, E), out] bf16 buffer (coalesced, no @@ -933,11 +802,8 @@ def permute_free_grouped_gemm_bf16_fc2( compact, routing, num_recv_tokens=routing.num_recv_tokens, - config=kernel_config, + block_m=block_size_m, index_a_by_route_pos=True, - activation=activation, - dispatched_probs=dispatched_probs, - gated_a=gated_a, ) return route_gather_combine( compact, @@ -952,8 +818,6 @@ def permute_free_grouped_gemm_bf16_fc2_dgrad( grad_output: torch.Tensor, weights: torch.Tensor | list[torch.Tensor], routing: MoERoutingMetadata, - *, - config: Optional[Dict[str, Any]] = None, ) -> torch.Tensor: """FC2 dgrad: gather the token-space grad back into the compact route buffer. @@ -997,12 +861,7 @@ def permute_free_grouped_gemm_bf16_fc2_dgrad( weights_stacked = weights_stacked.contiguous() weights_t = weights_stacked.transpose(1, 2) - if routing.sorted_slot_ids is None or routing.block_size_m is None: - fwd_config = config or get_default_moe_kernel_config(routing.num_recv_tokens) - routing = prepare_moe_align(routing, int(fwd_config["BLOCK_SIZE_M"])) - block_size_m = int(routing.block_size_m) - dgrad_config = get_default_moe_kernel_config(int(grad_output.shape[0])) - dgrad_config = {**dgrad_config, "BLOCK_SIZE_M": block_size_m} + block_size_m = _ensure_fwd_align(routing, grad_output, weights_t) # Padded route-order output; the gather-GEMM writes only the compact [0, num_routes) # range and never visits the tail (bounded by num_tokens_post_padded, sentinel-masked). @@ -1020,7 +879,7 @@ def permute_free_grouped_gemm_bf16_fc2_dgrad( dgrad, routing, num_recv_tokens=grad_output.shape[0], - config=dgrad_config, + block_m=block_size_m, index_a_by_route_pos=False, ) return dgrad @@ -1032,7 +891,6 @@ def permute_free_grouped_gemm_bf16_fc2_wgrad( weights_shape, routing: MoERoutingMetadata, *, - config: Optional[Dict[str, Any]] = None, out: Optional[torch.Tensor] = None, accumulate: bool = False, preact: Optional[torch.Tensor] = None, @@ -1089,7 +947,6 @@ def permute_free_grouped_gemm_bf16_fc2_wgrad( grad_output, weights_shape, routing, - config=config, out=out, accumulate=accumulate, swap_gather=True, @@ -1099,7 +956,6 @@ def permute_free_grouped_gemm_bf16_fc2_wgrad( grad_output, weights_shape, routing, - config=config, swap_gather=True, ) # [E, H, F] bf16 if out is None: diff --git a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py index bb1fa968b..3189e46ec 100644 --- a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py +++ b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py @@ -7,9 +7,9 @@ metadata (``sorted_slot_ids`` / ``expert_ids`` / ``block_start``) can be reused verbatim. Writes the block-padded ``[em_max, WIDTH_N]`` slot output in place. -Forward gather-GEMM is FlyDSL-only. This module retains Triton kernels for routing -metadata construction, gated-activation recompute/bwd, and the token-order gather-combine -pass that follows the compact route-list GEMM outputs. +Forward gather-GEMM is FlyDSL-only (v3 MegaMOE-ported plain GEMM). Gated activation and +route-prob apply live in standalone Triton helpers (:mod:`pf_helper_kernels`), not in the +GEMM kernel. """ from __future__ import annotations @@ -26,10 +26,6 @@ "flydsl_moe_fwd_pick_block_m", ] -# Activation ids for fused-epilogue kernels (v2, not yet in-tree). -ACT_SILU = 0 -ACT_GELU = 1 -_ACT_IDS = {"silu": ACT_SILU, "gelu": ACT_GELU} _WMMA = 16 _FILL_V = 8 _WARP = 64 @@ -53,8 +49,8 @@ def _v3_enabled() -> bool: # MegaMOE's hand-tuned bf16 grouped-GEMM tile geometry (offline-swept in primus-turbo's # bench_mega_moe: BLOCK_M/BLOCK_N=256, GROUP_M=4 fwd / GROUP_M=8 FC1 NN dgrad, num_xcd=1, -# nt_vmcnt=3). TE's FlyDSL fwd align already bumps non-gated FC1 / dgrad / plain FC2 to -# block_size_m=256; v3 pins the same Mega M-tile rather than trusting the wrapper arg. +# nt_vmcnt=3). TE's FlyDSL fwd align already bumps FC1 / dgrad / FC2 to block_size_m=256; +# v3 pins the same Mega M-tile rather than trusting the wrapper arg. _V3_BLOCK_M = 256 _V3_BLOCK_N = 256 _V3_GROUP_M = 4 @@ -64,22 +60,15 @@ def _v3_enabled() -> bool: def _run_v3_fwd( A, B, C, sorted_slot_ids, expert_ids, *, num_recv_tokens, block_m, - transpose_b, index_a_by_route_pos, gated, gated_a, mul_prob, save_preact, + transpose_b, index_a_by_route_pos, ): """Dispatch a plain fwd/dgrad GEMM to the v3 (MegaMOE-ported) kernels. - Covers the three non-fused paths: FC1 gather (``index_a_by_route_pos=False``), FC2 - route-read (``index_a_by_route_pos=True``) and dgrad (``transpose_b``). Fused epilogues - (gated activation / route-prob / pre-activation save) have no v3 equivalent and raise. - v3 uses MegaMOE's fixed tile geometry (32x32x16, ``BLOCK_N=256``, ``GROUP_M=4``, - ``num_xcd=1``), so the wrapper's ``block_n``/``block_k``/warp args are ignored here. + Covers the three paths: FC1 gather (``index_a_by_route_pos=False``), FC2 route-read + (``index_a_by_route_pos=True``) and dgrad (``transpose_b``). v3 uses MegaMOE's fixed tile + geometry (32x32x16, ``BLOCK_N=256``, ``GROUP_M=4``, ``num_xcd=1``), so the wrapper's + ``block_n``/``block_k``/warp args are ignored here. """ - if gated or gated_a or mul_prob or save_preact: - raise RuntimeError( - "Fused permute-free forward (gated activation / dispatched_probs / preact_out) " - "requires the v2 FlyDSL kernel (moe_fwd_flydsl_v2), which is not in-tree yet. " - "Use standalone gated-act + plain v3 GEMM (GroupedLinear FC2 path), or port v2." - ) block_m = int(block_m) em_max = int(sorted_slot_ids.shape[0]) if em_max % block_m != 0: @@ -194,22 +183,18 @@ def _fwd_buffering(): return (3 if use_dma else 2), pad -def _lds_bytes(block_m, block_n, block_k, gated, transpose_b=False, gated_a=False): - n_bt = 2 if gated else 1 +def _lds_bytes(block_m, block_n, block_k, transpose_b=False): nbuf, pad = _fwd_buffering() a_tile = block_m * (block_k + pad) - # Hybrid gated_a DMA path stages the raw ``up`` half in a parallel A tile (2x A LDS). - hybrid_a = gated_a and _env_flag("MOE_FWD_DMA", True) and _env_flag("MOE_FWD_GATEDA_DMA", False) - a_tiles = 2 if hybrid_a else 1 # dgrad stages B as [k, n] (row stride = block_n+pad); fwd as [n, k] (block_k+pad). b_tile = block_k * (block_n + pad) if transpose_b else block_n * (block_k + pad) - return (a_tiles * a_tile + n_bt * b_tile) * nbuf * 2 + return (a_tile + b_tile) * nbuf * 2 -def _default_block_n(block_m, block_k, gated, transpose_b=False): +def _default_block_n(block_m, block_k, transpose_b=False): """Widest N tile (128 then 64) that fits LDS -- wider N raises arithmetic intensity.""" for bn in (128, 64): - if _lds_bytes(block_m, bn, block_k, gated, transpose_b) <= _LDS_LIMIT and bn % _mfma_dim(transpose_b) == 0: + if _lds_bytes(block_m, bn, block_k, transpose_b) <= _LDS_LIMIT and bn % _mfma_dim(transpose_b) == 0: return bn return 64 @@ -249,22 +234,13 @@ def flydsl_moe_fwd( warps_m: Optional[int] = None, warps_n: Optional[int] = None, index_a_by_route_pos: bool = False, - activation: Optional[str] = None, - dispatched_probs: Optional[torch.Tensor] = None, - preact_out: Optional[torch.Tensor] = None, - gated_a: bool = False, ) -> None: """Route-list gather-GEMM forward, writing the compact ``C[em_max, WIDTH_N]`` in place. ``A`` is ``[*, K]`` (received-token acts, gathered by ``sorted_slot_ids`` when ``index_a_by_route_pos=False``, else read at the compact route row). ``B`` is - ``[num_experts, N_OUT, K]`` (contiguous inner ``K``). With ``activation`` set the fused - **gated** epilogue (FC1, ``gated_a=False``) applies ``act(gate) * up`` over ``N_OUT = 2F`` - into the ``F``-wide ``C``; ``dispatched_probs`` multiplies the per-route prob after the - activation, and ``preact_out`` (``[em_max, 2F]``) saves the raw ``[gate | up]`` - pre-activation. With ``gated_a=True`` (FC2) ``A`` is the raw ``[gate | up]`` pre-activation - (width ``2F``), the prologue applies ``act(gate) * up [* prob]`` into an ``F``-wide tile, - and the GEMM contracts over ``K = F`` against ``B[e, H, F]``. + ``[num_experts, N_OUT, K]`` (contiguous inner ``K``). Gated activation and route-prob + apply are **not** fused here; use the standalone helpers in :mod:`pf_helper_kernels`. """ assert A.dtype == B.dtype == C.dtype == torch.bfloat16 assert A.stride(1) == 1, "A must be contiguous along the contraction (K)" @@ -272,22 +248,12 @@ def flydsl_moe_fwd( transpose_b = B.stride(2) != 1 if transpose_b: assert B.stride(1) == 1, "transposed B must be contiguous along N (dgrad view)" - assert activation is None, "dgrad (transposed B) does not support fused activation" - - gated = activation is not None and not gated_a - if gated_a: - if activation is None: - raise ValueError("gated_a requires activation ('silu' or 'gelu')") - if not index_a_by_route_pos: - raise ValueError("gated_a requires index_a_by_route_pos=True") - # Plain GEMM: in-tree v3 kernels (pf_fwd / pf_dgrad). Fused epilogues need v2 (not ported). + # Plain GEMM: in-tree v3 kernels (pf_fwd / pf_dgrad). _run_v3_fwd( A, B, C, sorted_slot_ids, expert_ids, num_recv_tokens=num_recv_tokens, block_m=block_m, transpose_b=transpose_b, index_a_by_route_pos=index_a_by_route_pos, - gated=gated, gated_a=gated_a, - mul_prob=dispatched_probs is not None, save_preact=preact_out is not None, ) @@ -337,32 +303,27 @@ def flydsl_moe_fwd( _FWD_CACHE: dict = {} -def _valid_config(block_m, bn, bk, wm, wn, K, gated, gated_a=False, transpose_b=False): +def _valid_config(block_m, bn, bk, wm, wn, K, transpose_b=False): if K % bk != 0 or bn % _mfma_dim(transpose_b) != 0: return False if not _warp_valid(block_m, bn, bk, wm, wn, transpose_b): return False - return _lds_bytes(block_m, bn, bk, gated, gated_a=gated_a) <= _LDS_LIMIT + return _lds_bytes(block_m, bn, bk, transpose_b) <= _LDS_LIMIT def flydsl_moe_fwd_pick_block_m( A: torch.Tensor, B: torch.Tensor, *, - gated: bool = False, - gated_a: bool = False, candidates=(256, 128), ) -> Optional[int]: - """Largest ``block_m`` in ``candidates`` the FlyDSL fwd can actually run for these operands - and epilogue mode, or ``None`` if the operands are unsupported at every candidate. + """Largest ``block_m`` in ``candidates`` the FlyDSL fwd can actually run for these operands, + or ``None`` if the operands are unsupported at every candidate. "Can run" == bf16 operands contiguous along the contraction AND at least one autotuner tile - in :data:`_FWD_TUNE_CONFIGS` fits the per-workgroup LDS budget for the epilogue. The gated FC1 - epilogue stages a ``2F`` ``[gate|up]`` B-tile, so it only fits ``block_m <= 128``; the non-gated - FC1 fwd and the ``gated_a`` FC2 prologue fit ``block_m = 256``, which lifts the shared - fwd/dgrad/FC2 align onto the faster ``256x32`` MegaMOE-like tile. Callers should include their - token-count default among ``candidates`` (the picker only walks high->low over what is passed), - so a small-token workload is never padded up beyond what the caller offered. + in :data:`_FWD_TUNE_CONFIGS` fits the per-workgroup LDS budget. Callers should include their + default among ``candidates`` (the picker only walks high->low over what is passed), so a + small-token workload is never padded up beyond what the caller offered. """ if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: return None @@ -373,7 +334,7 @@ def flydsl_moe_fwd_pick_block_m( K = int(B.shape[2]) for block_m in sorted({int(c) for c in candidates}, reverse=True): if any( - _valid_config(block_m, bn, bk, wm, wn, K, gated, gated_a) + _valid_config(block_m, bn, bk, wm, wn, K) for (bn, bk, wm, wn) in _FWD_TUNE_CONFIGS ): return block_m @@ -392,35 +353,24 @@ def flydsl_moe_fwd_autotuned( block_m: int, block_k: int = 64, index_a_by_route_pos: bool = False, - activation: Optional[str] = None, - dispatched_probs: Optional[torch.Tensor] = None, - preact_out: Optional[torch.Tensor] = None, - gated_a: bool = False, warmup: int = 3, iters: int = 10, ) -> None: """Shape-autotuned :func:`flydsl_moe_fwd`. - On the first call for a given (block_m, GEMM shape, epilogue mode) the valid subset of + On the first call for a given (block_m, GEMM shape) the valid subset of ``_FWD_TUNE_CONFIGS`` is benchmarked and the fastest ``(block_n, block_k, warps_m, warps_n)`` is cached. The production DMA+swizzle fill path is always used; only tile geometry is swept. """ - gated = activation is not None and not gated_a N_OUT, K = int(B.shape[1]), int(B.shape[2]) width_n = int(C.shape[1]) - key = ( - int(block_m), N_OUT, K, width_n, bool(gated), bool(gated_a), - activation, dispatched_probs is not None, preact_out is not None, - bool(index_a_by_route_pos), - ) + key = (int(block_m), N_OUT, K, width_n, bool(index_a_by_route_pos)) def _launch(bn, bk, wm, wn): flydsl_moe_fwd( A, B, C, sorted_slot_ids, expert_ids, block_start, num_recv_tokens=num_recv_tokens, block_m=block_m, block_n=bn, block_k=bk, warps_m=wm, warps_n=wn, index_a_by_route_pos=index_a_by_route_pos, - activation=activation, dispatched_probs=dispatched_probs, preact_out=preact_out, - gated_a=gated_a, ) best = _FWD_CACHE.get(key) @@ -428,7 +378,7 @@ def _launch(bn, bk, wm, wn): candidates = [ (bn, bk, wm, wn) for (bn, bk, wm, wn) in _FWD_TUNE_CONFIGS - if _valid_config(block_m, bn, bk, wm, wn, K, gated, gated_a) + if _valid_config(block_m, bn, bk, wm, wn, K) ] if not candidates: _launch(None, block_k, None, None) # heuristic fallback From 8cb39912c19c852ff01d6c1b7304417c6cdc7731 Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Thu, 6 Aug 2026 00:40:08 +0000 Subject: [PATCH 37/43] Cache per-slot expert ids in MoERoutingMetadata and add permute-free backward tests Adds a `slot_expert_ids` field to `MoERoutingMetadata` that `prepare_moe_align` populates by broadcasting block-level `expert_ids` to the padded slot layout. `_expert_per_route` now returns the cached tensor when available, avoiding redundant recomputation for the standalone gated-activation kernels. Also adds unit tests covering the permute-free grouped GEMM backward paths: FC1 dgrad/wgrad, FC2 dgrad with `grad_probs`, and the GeLU variant of the route-list gated-activation backward. --- .../pytorch/test_perm_free_grouped_linear.py | 197 ++++++++++++++++++ transformer_engine/pytorch/moe/moe_routing.py | 5 + .../pytorch/moe/permute_free_grouped_gemm.py | 35 +++- 3 files changed, 233 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/test_perm_free_grouped_linear.py b/tests/pytorch/test_perm_free_grouped_linear.py index ed715bfff..3e4b691f5 100644 --- a/tests/pytorch/test_perm_free_grouped_linear.py +++ b/tests/pytorch/test_perm_free_grouped_linear.py @@ -743,6 +743,203 @@ def test_is_permute_free_grouped_gemm_enabled(monkeypatch): assert is_permute_free_grouped_gemm_enabled() is False +def test_slot_expert_ids_cache(): + """``prepare_moe_align`` caches per-slot expert ids; ``_expert_per_route`` reuses the cache.""" + torch.manual_seed(67) + num_recv_tokens, num_experts, max_hits = 128, 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=33) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + + # Not built until the align runs. + assert routing.slot_expert_ids is None + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + + # Populated at align time with the right shape/dtype. + assert routing.slot_expert_ids is not None + assert routing.slot_expert_ids.shape == (em_max,) + assert routing.slot_expert_ids.dtype == torch.int32 + + # Matches a manual block-broadcast of ``expert_ids`` (expert_ids[slot // block_m]). + block_m = int(routing.block_size_m) + pos = torch.arange(em_max, device="cuda", dtype=torch.int64) + block_idx = (pos // block_m).clamp(max=routing.expert_ids.numel() - 1) + manual = routing.expert_ids[block_idx].to(torch.int32) + assert torch.equal(routing.slot_expert_ids, manual) + + # ``_expert_per_route`` returns the *cached* tensor (no redundant recompute). + assert _expert_per_route(routing, em_max) is routing.slot_expert_ids + + # Valid slots carry the expert-sorted compact route experts. + _, route_expert = _compact_route_order(routing_map) + valid = routing.sorted_slot_ids < num_recv_tokens + assert torch.equal(routing.slot_expert_ids[valid].to(torch.int64), route_expert) + + # Rebuild-on-early-return: drop the cache but keep the align, then re-prepare with the same + # block_m (hits the cached-align early return) -- it must repopulate rather than stay None. + routing.slot_expert_ids = None + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + assert routing.slot_expert_ids is not None + assert torch.equal(routing.slot_expert_ids, manual) + + +def test_permute_free_backward_fc1_dispatch(): + """``permute_free_grouped_gemm_backward`` FC1 path (``route_space=False``): dgrad + wgrad match + the standalone kernels, and ``wgrad_out`` folds in place (``wgrad_applied``).""" + torch.manual_seed(71) + num_recv_tokens, in_features, out_features = 128, 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=35) + routing = PermuteFreeMetadata(routing_map=routing_map, num_experts=num_experts) # FC1 + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + valid = routing.sorted_slot_ids < num_recv_tokens + + hidden = torch.randn(num_recv_tokens, in_features, device="cuda", dtype=torch.bfloat16) + weights = [ + torch.randn(out_features, in_features, device="cuda", dtype=torch.bfloat16) + for _ in range(num_experts) + ] + W = torch.stack(weights, dim=0) + grad = torch.zeros(em_max, out_features, device="cuda", dtype=torch.bfloat16) + grad[valid] = torch.randn( + int(valid.sum().item()), out_features, device="cuda", dtype=torch.bfloat16 + ) + + res = permute_free_grouped_gemm_backward( + grad, routing=routing, weights=weights, num_gemms=num_experts, + hidden_states=hidden, requires_dgrad=True, requires_wgrad=True, + ) + dgrad_ref = permute_free_grouped_gemm_bf16_dgrad(grad, W, routing) + wgrad_ref = permute_free_grouped_gemm_bf16_wgrad( + hidden, grad, (num_experts, out_features, in_features), routing + ) + assert res.dgrad.shape == (num_recv_tokens, in_features) + assert res.wgrad_stacked.shape == (num_experts, out_features, in_features) + assert res.wgrad_applied is False + assert res.grad_probs is None + assert _rel_l2(res.dgrad, dgrad_ref) < 1e-3 + assert _rel_l2(res.wgrad_stacked, wgrad_ref) < 1e-3 + + # wgrad_out fold (fp32 sink, overwrite): wgrad_applied True and the buffer is returned. + wgrad_out = torch.zeros( + num_experts, out_features, in_features, device="cuda", dtype=torch.float32 + ) + res2 = permute_free_grouped_gemm_backward( + grad, routing=routing, weights=weights, num_gemms=num_experts, + hidden_states=hidden, requires_wgrad=True, + wgrad_out=wgrad_out, wgrad_accumulate=False, + ) + assert res2.wgrad_applied is True + assert res2.wgrad_stacked is wgrad_out + assert res2.dgrad is None + assert _rel_l2(wgrad_out, wgrad_ref) < 5e-3 + + +def test_permute_free_backward_fc2_dgrad_probs(): + """FC2 backward dispatch dgrad path: dgrad + grad_probs match the standalone act-bwd, and + grad_probs is suppressed when ``dispatched_probs`` does not require grad.""" + torch.manual_seed(73) + num_recv_tokens, in_features, out_features = 128, 128, 128 # F=in, H=out + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=37) + routing = PermuteFreeMetadata( + routing_map=routing_map, num_experts=num_experts, route_space=True, activation="silu" + ) + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + valid = routing.sorted_slot_ids < num_recv_tokens + + preact = torch.zeros(em_max, 2 * in_features, device="cuda", dtype=torch.bfloat16) + preact[valid] = torch.randn( + int(valid.sum().item()), 2 * in_features, device="cuda", dtype=torch.bfloat16 + ) + grad_output = torch.randn(num_recv_tokens, out_features, device="cuda", dtype=torch.bfloat16) + weights = [ + torch.randn(out_features, in_features, device="cuda", dtype=torch.bfloat16) + for _ in range(num_experts) + ] + W = torch.stack(weights, dim=0) + + probs = torch.rand( + num_recv_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True + ) + res = permute_free_grouped_gemm_backward( + grad_output, routing=routing, weights=weights, num_gemms=num_experts, + hidden_states=preact, requires_dgrad=True, + fc2_activation="silu", dispatched_probs=probs, + ) + # Reference: fc2_dgrad -> standalone gated act-bwd. + dgrad_f = permute_free_grouped_gemm_bf16_fc2_dgrad(grad_output, W, routing) + dpre_ref, dprob_ref = permute_free_gated_act_bwd( + dgrad_f, preact, routing, activation="silu", dispatched_probs=probs.detach() + ) + assert res.dgrad.shape == (em_max, 2 * in_features) + assert res.grad_probs is not None + assert res.grad_probs.shape == (num_recv_tokens, num_experts) + assert _rel_l2(res.dgrad[valid], dpre_ref[valid]) < 2e-2 + assert _rel_l2(res.grad_probs, dprob_ref) < 2e-2 + + # No prob grad requested -> grad_probs suppressed. + res_nop = permute_free_grouped_gemm_backward( + grad_output, routing=routing, weights=weights, num_gemms=num_experts, + hidden_states=preact, requires_dgrad=True, + fc2_activation="silu", dispatched_probs=probs.detach(), + ) + assert res_nop.grad_probs is None + assert _rel_l2(res_nop.dgrad[valid], dpre_ref[valid]) < 2e-2 + + +def test_route_list_gated_act_bwd_gelu(): + """Gelu (tanh) variant of the standalone gated-activation backward (ACT_GELU path).""" + torch.manual_seed(79) + num_recv_tokens, in_features = 128, 128 + num_experts, max_hits = 8, 3 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=39) + routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) + routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + em_max = int(routing.sorted_slot_ids.shape[0]) + slot_token = routing.sorted_slot_ids + slot_expert = _expert_per_route(routing, em_max) + valid = slot_token < num_recv_tokens + tok = slot_token.to(torch.int64).clamp_(0, num_recv_tokens - 1) + exp = slot_expert.to(torch.int64).clamp_min(0) + + gate = torch.randn(em_max, in_features, device="cuda", dtype=torch.float32) + up = torch.randn(em_max, in_features, device="cuda", dtype=torch.float32) + gate = (gate * valid[:, None].float()).requires_grad_(True) + up = (up * valid[:, None].float()).requires_grad_(True) + probs = torch.rand( + num_recv_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True + ) + grad_out = torch.zeros(em_max, in_features, device="cuda", dtype=torch.bfloat16) + grad_out[valid] = torch.randn( + int(valid.sum().item()), in_features, device="cuda", dtype=torch.bfloat16 + ) + + pr = probs[tok, exp] + act = ( + torch.nn.functional.gelu(gate, approximate="tanh") + * up + * pr[:, None] + * valid[:, None].float() + ) + (act * grad_out.float()).sum().backward() + dpre_ref = torch.cat([gate.grad, up.grad], dim=1) + + preact = torch.cat([gate.detach(), up.detach()], dim=1).to(torch.bfloat16) + dpre, dprob = permute_free_gated_act_bwd( + grad_out, preact, routing, activation="gelu", dispatched_probs=probs.detach() + ) + assert dpre.shape == (em_max, 2 * in_features) + assert _rel_l2(dpre[valid], dpre_ref[valid]) < 3e-2 + assert _rel_l2(dprob, probs.grad) < 3e-2 + + @pytest.mark.skip(reason="apply_route_probs not ported to transformer_engine.pytorch.moe yet") def test_apply_route_probs_fwd_bwd(): """Fused per-route prob apply (gather+multiply) vs an autograd advanced-index reference.""" diff --git a/transformer_engine/pytorch/moe/moe_routing.py b/transformer_engine/pytorch/moe/moe_routing.py index 6b85f86b4..79217e252 100644 --- a/transformer_engine/pytorch/moe/moe_routing.py +++ b/transformer_engine/pytorch/moe/moe_routing.py @@ -44,6 +44,10 @@ class MoERoutingMetadata: expert_ids: ``[blocks_max]`` local expert owning each ``BLOCK_SIZE_M`` block (``-1`` past the real block count; those blocks are never visited). + slot_expert_ids: + ``[T * min(topk, E)]`` per-slot local expert id, derived from ``expert_ids`` and + ``block_size_m`` (``expert_ids[slot // block_size_m]``). Populated by + ``prepare_moe_align`` for the standalone gated-activation kernels. num_tokens_post_padded: ``[1]`` device scalar = real ``em`` (block-padded route count). Bounds the kernel. block_start: @@ -70,6 +74,7 @@ class MoERoutingMetadata: topk: Optional[int] = None sorted_slot_ids: Optional[torch.Tensor] = None expert_ids: Optional[torch.Tensor] = None + slot_expert_ids: Optional[torch.Tensor] = None num_tokens_post_padded: Optional[torch.Tensor] = None block_start: Optional[torch.Tensor] = None token_routes: Optional[torch.Tensor] = None diff --git a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py index 6f9b480e1..2e7f532f6 100644 --- a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py +++ b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py @@ -57,6 +57,18 @@ def _get_flydsl_fwd(): return flydsl_moe_fwd_autotuned, flydsl_moe_fwd_supported +def _expand_expert_ids_per_slot( + expert_ids: torch.Tensor, + block_m: int, + routes_max: int, +) -> torch.Tensor: + """Broadcast block-level ``expert_ids`` to per-slot ids ``[routes_max]``.""" + block_m = int(block_m) + pos = torch.arange(routes_max, device=expert_ids.device, dtype=torch.int64) + block_idx = (pos // block_m).clamp(max=expert_ids.numel() - 1) + return expert_ids[block_idx].to(torch.int32) + + def _expert_per_route(routing: MoERoutingMetadata, routes_max: int) -> torch.Tensor: """Per-route local expert id ``[routes_max]`` from the route-list block metadata.""" if ( @@ -65,10 +77,16 @@ def _expert_per_route(routing: MoERoutingMetadata, routes_max: int) -> torch.Ten or routing.block_size_m <= 0 ): raise ValueError("_expert_per_route requires prepared routing align buffers.") - block_m = int(routing.block_size_m) - pos = torch.arange(routes_max, device=routing.expert_ids.device, dtype=torch.int64) - block_idx = (pos // block_m).clamp(max=routing.expert_ids.numel() - 1) - return routing.expert_ids[block_idx.to(torch.int64)].to(torch.int32) + if ( + routing.slot_expert_ids is not None + and routing.slot_expert_ids.shape[0] == routes_max + ): + return routing.slot_expert_ids + slot_expert_ids = _expand_expert_ids_per_slot( + routing.expert_ids, routing.block_size_m, routes_max + ) + routing.slot_expert_ids = slot_expert_ids + return slot_expert_ids def _env_int(name: str, default: int) -> int: v = os.environ.get(name) @@ -250,6 +268,12 @@ def prepare_moe_align(metadata: MoERoutingMetadata, block_m: int) -> MoERoutingM and metadata.block_size_m == block_m and metadata.token_routes is not None ): + if metadata.slot_expert_ids is None: + metadata.slot_expert_ids = _expand_expert_ids_per_slot( + metadata.expert_ids, + block_m, + metadata.sorted_slot_ids.shape[0], + ) return metadata counts, within = _ensure_route_scan(metadata) @@ -272,6 +296,9 @@ def prepare_moe_align(metadata: MoERoutingMetadata, block_m: int) -> MoERoutingM ) metadata.sorted_slot_ids = sorted_slot_ids # [T * min(topk, E)] int32: block-padded slot -> token metadata.expert_ids = expert_ids # [blocks_max] int32: expert owning each block (-1 past end) + metadata.slot_expert_ids = _expand_expert_ids_per_slot( + expert_ids, block_m, sorted_slot_ids.shape[0] + ) metadata.num_tokens_post_padded = num_tokens_post_padded # [1] int32 device scalar: padded extent metadata.block_start = block_start # [E] int32: first block index of each expert (block units) metadata.block_size_m = block_m # int: BLOCK_SIZE_M the layout is padded to From 5a0eaec4e5b9df899a0ad0d4156db4b2ea46e76f Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Thu, 6 Aug 2026 15:26:23 +0000 Subject: [PATCH 38/43] Return tensor directly from permute-free grouped GEMM forward and clean up stale fusion code Removes the `PermuteFreeForwardResult` wrapper so `permute_free_grouped_gemm_forward` returns the output tensor directly, and drops it from the public `moe` exports. Updates `GroupedLinear` and the unit tests to stop accessing `.out`. Revises docstrings/comments in `GroupedLinear`, `moe_routing`, and the permute-free kernels to describe the FC2 standalone gated-activation pass instead of the old fused-prologue path. Removes dead helpers (`_pick_warps`, `_default_block_n`, `_env_int`), the skipped `apply_route_probs` test, and the unused `_FLYDSL_FWD_LARGE_TIER` tier, simplifying forward block-size selection. --- .../benchmark_perm_free_grouped_gemm.py | 2 +- .../pytorch/test_perm_free_grouped_linear.py | 56 ++-------- .../pytorch/module/grouped_linear.py | 35 +++--- transformer_engine/pytorch/moe/__init__.py | 2 - transformer_engine/pytorch/moe/moe_routing.py | 9 +- .../pytorch/moe/permute_free_grouped_gemm.py | 102 ++++++------------ .../pytorch/moe/pf_fwd_wrapper.py | 47 +------- .../pytorch/moe/pf_helper_kernels.py | 19 ++-- 8 files changed, 72 insertions(+), 200 deletions(-) diff --git a/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py index e954c7575..014299d56 100644 --- a/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py +++ b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py @@ -222,7 +222,7 @@ def triton_fn(): fns["ck"] = ck_fn fns["triton"] = triton_fn fns["permute_free"] = lambda: _permute_free_gemm(hidden, weights, routing) - # GateUP-only: permute-free FC1 with the fused gated SiLU epilogue. + # GateUP-only: permute-free FC1 GEMM + standalone SiLU gated activation. if is_gated: fns["permute_free_act"] = lambda: _permute_free_act_gemm(hidden, weights, routing) return fns diff --git a/tests/pytorch/test_perm_free_grouped_linear.py b/tests/pytorch/test_perm_free_grouped_linear.py index 3e4b691f5..6858997cc 100644 --- a/tests/pytorch/test_perm_free_grouped_linear.py +++ b/tests/pytorch/test_perm_free_grouped_linear.py @@ -362,16 +362,15 @@ def test_fc2_backward_dispatch_recompute_matches_stored(): def test_fc1_fc2_gated_pipeline(monkeypatch): - """End-to-end FC1 raw 2F -> FC2 fused activation: forward outputs and backward grads match a - PyTorch reference through the permute-free GroupedLinear modules.""" + """End-to-end FC1 raw 2F -> FC2 standalone gated activation: forward outputs and backward grads + match a PyTorch reference through the permute-free GroupedLinear modules.""" import dataclasses from transformer_engine.pytorch.moe import prepare_moe_align monkeypatch.setenv("NVTE_PERMUTE_FREE_GROUPED_GEMM", "1") torch.manual_seed(37) - # FFN width is a multiple of the FlyDSL block_k (64) so the fused FC2 gated_a prologue - # runs on FlyDSL (the default backend) rather than falling back to Triton. + # FFN width is a multiple of the FlyDSL block_k (64) for kernel alignment. num_recv_tokens, hidden, ffn = 128, 128, 128 num_experts, max_hits = 8, 3 @@ -562,7 +561,7 @@ def test_route_list_fc2_fwd_standalone_act(): out = permute_free_grouped_gemm_forward( preact, weights, routing, activation="silu", dispatched_probs=probs - ).out + ) assert out.shape == (num_recv_tokens, out_features) g = preact[:, :in_features].float() @@ -700,7 +699,7 @@ def test_permute_free_forward_dispatch(): res_fc1 = permute_free_grouped_gemm_forward(hidden, w1, fc1_routing) direct = permute_free_grouped_gemm_bf16(hidden, w1, fc1_routing) valid_fc1 = fc1_routing.sorted_slot_ids < num_recv_tokens - assert _rel_l2(res_fc1.out[valid_fc1], direct[valid_fc1]) < 1e-3 + assert _rel_l2(res_fc1[valid_fc1], direct[valid_fc1]) < 1e-3 fc2_routing = PermuteFreeMetadata( routing_map=routing_map, num_experts=num_experts, route_space=True, activation="silu" @@ -721,7 +720,7 @@ def test_permute_free_forward_dispatch(): w2 = torch.randn(num_experts, hidden_dim, ffn, device="cuda", dtype=torch.bfloat16) out = permute_free_grouped_gemm_forward( preact, w2, fc2_routing, activation="silu", dispatched_probs=probs - ).out + ) assert out.shape == (num_recv_tokens, hidden_dim) g = preact[:, :ffn].float() @@ -940,49 +939,6 @@ def test_route_list_gated_act_bwd_gelu(): assert _rel_l2(dprob, probs.grad) < 3e-2 -@pytest.mark.skip(reason="apply_route_probs not ported to transformer_engine.pytorch.moe yet") -def test_apply_route_probs_fwd_bwd(): - """Fused per-route prob apply (gather+multiply) vs an autograd advanced-index reference.""" - from transformer_engine.pytorch.moe import prepare_moe_align - # apply_route_probs was not ported; kept for when route_prob helper lands in moe/. - _ = prepare_moe_align - - torch.manual_seed(31) - num_recv_tokens, hidden_dim = 128, 192 - num_experts, max_hits = 8, 3 - - routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=9) - routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) - prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) - - em_max = int(routing.sorted_slot_ids.shape[0]) - num_routes = int(routing_map.sum().item()) - act = torch.randn(em_max, hidden_dim, device="cuda", dtype=torch.bfloat16) - probs = torch.rand(num_recv_tokens, num_experts, device="cuda", dtype=torch.float32) - - a1 = act.clone().requires_grad_(True) - p1 = probs.clone().requires_grad_(True) - out = apply_route_probs(a1, p1, routing) - - # Reference: differentiable advanced index over the block-padded slot order. - tok = routing.sorted_slot_ids.to(torch.int64).clamp(0, num_recv_tokens - 1) - exp = _expert_per_route(routing, em_max).to(torch.int64).clamp_min(0) - valid = routing.sorted_slot_ids < num_recv_tokens - a2 = act.clone().requires_grad_(True) - p2 = probs.clone().requires_grad_(True) - pr = torch.where(valid, p2[tok, exp], torch.zeros_like(p2[tok, exp])) - ref = a2 * pr[:, None] - - assert _rel_l2(out[valid], ref[valid]) < 2e-2 - - g = torch.randn(em_max, hidden_dim, device="cuda", dtype=torch.bfloat16) - g[num_routes:] = 0 - out.backward(g) - ref.backward(g) - assert _rel_l2(a1.grad[:num_routes], a2.grad[:num_routes]) < 2e-2 - assert _rel_l2(p1.grad, p2.grad) < 2e-2 - - # --------------------------------------------------------------------------- # Module-level weight-gradient accumulation: grouped weight (-> main_grad) vs. # separate per-expert weights (-> autograd .grad). diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 2a24d1ca9..299d656e5 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -540,7 +540,7 @@ def forward( perm_free_route_space = ( getattr(routing_metadata, "route_space", False) if routing_metadata is not None else False ) - # FC2 with a fused gated prologue consumes FC1's raw 2F [gate|up] buffer (width 2F). + # FC2 with gated activation consumes FC1's raw 2F [gate|up] buffer (width 2F). expect_in_features = ( 2 * in_features if ( @@ -671,14 +671,13 @@ def forward( # FC1 emits raw 2F [gate|up]; the ``activation`` hint on the metadata is consumed on # FC2, which applies the gated activation in a standalone pass and then runs a plain # GEMM (the fused-prologue path regressed throughput). Route probs ride with FC2 too. - pf_result = permute_free_grouped_gemm_forward( + out = permute_free_grouped_gemm_forward( inputmats[0], weights_fp8, routing_metadata, activation=perm_free_activation if perm_free_route_space else None, dispatched_probs=dispatched_probs if perm_free_route_space else None, ) - out = pf_result.out elif use_grouped_gemm_triton: general_grouped_gemm_func = general_grouped_gemm_triton kwargs = {"m_splits_tensor": m_splits_tensor} @@ -1126,7 +1125,7 @@ def backward( ): grouped_weight.grad_added_to_main_grad = True else: - # FC2 (transpose) or no direct kernel: sink the returned stacked wgrad. + # Fresh wgrad tensor: sink into main_grad / .grad. dW = pf_result.wgrad_stacked if ctx.grouped_fuse_wgrad: main_grad = ctx.grouped_main_grad_func().view(dW.shape) @@ -2018,22 +2017,20 @@ def forward( False`` (FC1) gathers per expert into the worst-case padded ``[T * min(topk, E), out_features]`` route buffer (valid rows are the compact range ``[0, num_routes)``; the tail is inert zero padding); - ``route_space=True`` (FC2) reads route-ordered input and fuses the - scatter back to token order, returning ``[num_recv_tokens, - out_features]``. TE builds/caches the expert-sorted alignment buffers - on the metadata. The gated-activation fusion hint - (``permute_free_metadata.activation`` = ``"silu"`` / ``"gelu"``) rides on - this object: when set on the FC1 direction (``route_space=False``) it - fuses the **gated** activation into the GEMM epilogue (weight output dim - is the gate+up width ``2F``, laid out as ``[gate | up]``; the returned - buffer is the ``F``-wide ``act(gate) * up`` and the separate activation - pass is skipped). Ignored on other paths. + ``route_space=True`` (FC2) reads route-ordered input, applies the + standalone gated activation (when ``activation`` is set), runs a plain + GEMM, and scatter-combines back to token order, returning + ``[num_recv_tokens, out_features]``. TE builds/caches the expert-sorted + alignment buffers on the metadata. The ``activation`` hint + (``"silu"`` / ``"gelu"``) is carried on this object for the FC2 + direction: FC1 (``route_space=False``) emits raw ``2F`` ``[gate | up]``; + FC2 recompute applies ``act(gate) * up`` (and optionally route probs) + before the GEMM. dispatched_probs : torch.Tensor, optional - ``[num_recv_tokens, num_local_experts]`` gating probabilities. When given - with a fused ``activation``, each route's ``prob[token, expert]`` is - multiplied into the activation in-kernel (skipping the separate route-prob - pass); its gradient is returned to the router through autograd. Must be a - leaf/differentiable tensor for training. + ``[num_recv_tokens, num_local_experts]`` gating probabilities. On the FC2 + permute-free path, multiplied into the gated activation during the + standalone recompute pass; its gradient is returned to the router through + autograd. Must be a leaf/differentiable tensor for training. """ debug = self.is_debug_iter() is_grad_enabled = torch.is_grad_enabled() diff --git a/transformer_engine/pytorch/moe/__init__.py b/transformer_engine/pytorch/moe/__init__.py index ef5a96fb0..4b627b2fd 100644 --- a/transformer_engine/pytorch/moe/__init__.py +++ b/transformer_engine/pytorch/moe/__init__.py @@ -7,7 +7,6 @@ from .permute_free_grouped_gemm import ( MoERoutingMetadata, PermuteFreeBackwardResult, - PermuteFreeForwardResult, PermuteFreeMetadata, is_permute_free_grouped_gemm_enabled, permute_free_grouped_gemm_backward, @@ -26,7 +25,6 @@ __all__ = [ "MoERoutingMetadata", "PermuteFreeBackwardResult", - "PermuteFreeForwardResult", "PermuteFreeMetadata", "is_permute_free_grouped_gemm_enabled", "permute_free_grouped_gemm_backward", diff --git a/transformer_engine/pytorch/moe/moe_routing.py b/transformer_engine/pytorch/moe/moe_routing.py index 79217e252..db5a0ebc0 100644 --- a/transformer_engine/pytorch/moe/moe_routing.py +++ b/transformer_engine/pytorch/moe/moe_routing.py @@ -121,12 +121,13 @@ class PermuteFreeMetadata(MoERoutingMetadata): reused for FC1 and FC2 (e.g. via ``dataclasses.replace(meta, route_space=True)``), avoiding a duplicate align build. - Fusion hint (optional): + Activation hint (optional): activation: - Gated activation to fuse into the FC2 GEMM prologue -- ``"silu"`` or ``"gelu"``. - ``None`` leaves the activation to the caller (no fusion). FC1 emits raw ``2F``; - this hint is consumed on the FC2 direction (``route_space=True``). + Gated activation for the FC2 standalone pass -- ``"silu"`` or ``"gelu"``. + ``None`` leaves activation to the caller. FC1 emits raw ``2F``; this hint is + consumed on the FC2 direction (``route_space=True``) to run + :func:`permute_free_gated_act_recompute` before the plain FC2 GEMM. (The per-route gating probabilities are *not* carried here: they need a gradient, so they are passed as a separate autograd tensor argument to the module rather than as metadata.) diff --git a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py index 2e7f532f6..9ae495b36 100644 --- a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py +++ b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py @@ -36,7 +36,6 @@ "permute_free_grouped_gemm_bf16_wgrad", "permute_free_grouped_gemm_forward", "permute_free_grouped_gemm_backward", - "PermuteFreeForwardResult", "PermuteFreeBackwardResult", "prepare_moe_align", "is_permute_free_grouped_gemm_enabled", @@ -47,7 +46,6 @@ # Minimum v3 gather/dgrad block_m (128x256 MFMA floor). _FLYDSL_MIN_BLOCK_M = 128 _FLYDSL_FWD_BLOCK_M = 256 -_FLYDSL_FWD_LARGE_TIER = 128 def _get_flydsl_fwd(): @@ -62,9 +60,20 @@ def _expand_expert_ids_per_slot( block_m: int, routes_max: int, ) -> torch.Tensor: - """Broadcast block-level ``expert_ids`` to per-slot ids ``[routes_max]``.""" + """Broadcast block-level ``expert_ids`` to per-slot ids ``[routes_max]``. + + ``prepare_moe_align`` lays routes out expert-by-expert, padding each expert's count up to + a multiple of ``block_m``. ``expert_ids[b]`` records which expert owns M-block ``b`` (an + expert with more than ``block_m`` routes spans several consecutive blocks, all with the same + id). FlyDSL GEMM reads ``expert_ids`` at block granularity; the standalone gated-activation + kernels index row-by-row and need ``expert[slot]`` to look up ``dispatched_probs[token, e]``. + + This is a lookup, not ``expert = slot // block_m``: ``slot // block_m`` is the block index, + then ``expert_ids[block_idx]`` is the owner assigned during align. + """ block_m = int(block_m) pos = torch.arange(routes_max, device=expert_ids.device, dtype=torch.int64) + # Map each padded slot to its M-block; clamp covers the static over-allocated tail. block_idx = (pos // block_m).clamp(max=expert_ids.numel() - 1) return expert_ids[block_idx].to(torch.int32) @@ -88,13 +97,6 @@ def _expert_per_route(routing: MoERoutingMetadata, routes_max: int) -> torch.Ten routing.slot_expert_ids = slot_expert_ids return slot_expert_ids -def _env_int(name: str, default: int) -> int: - v = os.environ.get(name) - try: - return int(v) if v is not None and v.strip() != "" else default - except ValueError: - return default - def _get_flydsl_wgrad(): """Return the FC1 FlyDSL wgrad launcher.""" @@ -144,20 +146,13 @@ def _fwd_align_block_size_m( ``block_m`` is baked into ``prepare_moe_align`` (it sets the row padding) and is *shared* by a layer's FC1 fwd, FC1 dgrad and FC2 fwd (they reuse ``routing.block_size_m``), so it is chosen - once here. When the FlyDSL backend will run and the workload is on the large-token tier - (``num_tokens >= _FLYDSL_FWD_LARGE_TIER``), prefer :data:`_FLYDSL_FWD_BLOCK_M` (256): it lifts - FC1 fwd, FC1 dgrad, and FC2 onto the faster ``256x32`` tile. Small/medium tiers default to - :data:`_FLYDSL_MIN_BLOCK_M` (128), the v3 gather/dgrad floor. + once here. For ``num_tokens >= _FLYDSL_MIN_BLOCK_M`` (128), offer :data:`_FLYDSL_FWD_BLOCK_M` + (256, MegaMOE tile) alongside the v3 floor; smaller batches stay at 128 only. """ default_block_m = ( - _FLYDSL_FWD_BLOCK_M if num_tokens >= _FLYDSL_FWD_LARGE_TIER else _FLYDSL_MIN_BLOCK_M + _FLYDSL_FWD_BLOCK_M if num_tokens >= _FLYDSL_MIN_BLOCK_M else _FLYDSL_MIN_BLOCK_M ) - # Offer the default and any smaller floor (so the picker can drop to a smaller LDS-valid - # block_m), plus the 256 bump on the large-token tier, where it is a measured win and the - # extra E_local*128 align padding is a negligible fraction of the routed work. candidates = {c for c in (_FLYDSL_MIN_BLOCK_M, default_block_m) if c <= default_block_m} - if num_tokens >= _FLYDSL_FWD_LARGE_TIER: - candidates.add(_FLYDSL_FWD_BLOCK_M) from .pf_fwd_wrapper import _v3_enabled # The v3 gather/dgrad tile requires block_m >= 128 (128x256 MFMA minimum), so on a @@ -1000,19 +995,6 @@ def permute_free_grouped_gemm_bf16_fc2_wgrad( # metadata so callers (e.g. GroupedLinear) don't have to branch on route_space / # activation themselves. # --------------------------------------------------------------------------- -@dataclass -class PermuteFreeForwardResult: - """Output of :func:`permute_free_grouped_gemm_forward`. - - ``preact`` is the raw ``2F`` pre-activation saved for the FC1 gated-activation backward; - it is ``None`` on every other path (FC2, or FC1 without a fused activation / without a - backward). - """ - - out: torch.Tensor - preact: Optional[torch.Tensor] = None - - @dataclass class PermuteFreeBackwardResult: """Gradients from :func:`permute_free_grouped_gemm_backward`. @@ -1021,16 +1003,15 @@ class PermuteFreeBackwardResult: ``wgrad_stacked`` is the weight gradient as a single ``[E, out, in]`` tensor (the natural kernel output): accumulate it directly into a grouped param's ``main_grad``, or split it into per-expert views (``list(wgrad_stacked)``) for the positional autograd return. ``grad_probs`` - is the route-prob gradient for the FC1 fused-prob path (``None`` otherwise). + is the route-prob gradient from the standalone FC2 gated-activation backward (``None`` + otherwise). """ dgrad: Optional[torch.Tensor] = None grad_probs: Optional[torch.Tensor] = None wgrad_stacked: Optional[torch.Tensor] = None - # True when the wgrad was written directly into the caller-provided ``wgrad_out`` (FC1 - # path), so the caller must not re-apply it. False when ``wgrad_stacked`` is a fresh tensor - # the caller still needs to sink (FC2 -- its transpose precludes an in-place accumulate -- - # or the plain positional-return path). + # True when the wgrad was written directly into the caller-provided ``wgrad_out``; + # False when ``wgrad_stacked`` is a fresh tensor the caller still needs to sink. wgrad_applied: bool = False @@ -1041,39 +1022,27 @@ def permute_free_grouped_gemm_forward( *, activation: Optional[str] = None, dispatched_probs: Optional[torch.Tensor] = None, -) -> PermuteFreeForwardResult: +) -> torch.Tensor: """Dispatch the permute-free grouped-GEMM forward from the routing direction + fusion hints. ``routing.route_space`` selects the direction: - ``True`` (FC2): route-ordered ``2F`` pre-activation (FC1 output). The gated activation ``act(gate)*up[*prob]`` is applied in a standalone pass into an ``F``-wide transient, then - a plain gather-GEMM + gather-combine -> ``[num_recv, out]``. (Fusing the activation into - the FC2 GEMM prologue regressed throughput, so FC2 is kept as a plain DMA GEMM mirroring - FC1.) The ``F``-wide transient is freed after FC2; only the ``2F`` pre-activation is - checkpointed for backward (which recomputes the ``F`` activation just-in-time). - - ``False`` (FC1): gather-in-GEMM -> padded ``[T * min(topk, E), 2F]`` raw ``[gate | up]`` - pre-activation (no activation fusion here; the ``activation`` hint is consumed on FC2). + a plain gather-GEMM + gather-combine -> ``[num_recv, out]``. + - ``False`` (FC1): gather-in-GEMM -> padded ``[em_max, out]`` raw ``[gate | up]`` (``2F``). """ if getattr(routing, "route_space", False): fc2_input = hidden_states if activation is not None: - # Standalone gated-activation pass on the raw 2F [gate|up] pre-activation, producing - # the F-wide FC2 operand. This is the same route-wise kernel the backward uses to - # recompute the activation, so the caller can keep checkpointing only the 2F - # pre-activation (this F-wide buffer is transient and freed after the FC2 GEMM). fc2_input = permute_free_gated_act_recompute( hidden_states, routing, activation=activation, dispatched_probs=dispatched_probs, ) - return PermuteFreeForwardResult( - permute_free_grouped_gemm_bf16_fc2(fc2_input, weights, routing) - ) - return PermuteFreeForwardResult( - permute_free_grouped_gemm_bf16(hidden_states, weights, routing) - ) + return permute_free_grouped_gemm_bf16_fc2(fc2_input, weights, routing) + return permute_free_grouped_gemm_bf16(hidden_states, weights, routing) def permute_free_grouped_gemm_backward( @@ -1085,8 +1054,6 @@ def permute_free_grouped_gemm_backward( hidden_states: Optional[torch.Tensor] = None, requires_dgrad: bool = False, requires_wgrad: bool = False, - fc1_activation: Optional[str] = None, - preact: Optional[torch.Tensor] = None, dispatched_probs: Optional[torch.Tensor] = None, fc2_activation: Optional[str] = None, wgrad_out: Optional[torch.Tensor] = None, @@ -1094,20 +1061,14 @@ def permute_free_grouped_gemm_backward( ) -> PermuteFreeBackwardResult: """Dispatch the permute-free grouped-GEMM backward (mirror of the forward dispatch). - ``routing.route_space`` picks the FC2 vs FC1 dgrad/wgrad kernels. On the FC1 - gated-activation path (``fc1_activation`` set) the raw ``2F`` GEMM-output gradient (and the - route-prob gradient) is first reconstructed from the saved ``preact`` and fed to the - *unchanged* dgrad/wgrad. ``hidden_states`` (the forward input) is only needed for wgrad. - - On the FC2 path (``route_space=True``), when ``fc2_activation`` is set together with - ``preact`` the wgrad rebuilds its F-wide input (``act(gate)*up*prob``) in-flight from the FC1 - pre-activation instead of consuming a stored ``hidden_states`` -- the FC1->FC2 recompute - handoff, letting the FC2 forward skip saving its F-wide input. + ``routing.route_space`` picks the FC2 vs FC1 dgrad/wgrad kernels. On the FC2 path, + ``fc2_activation`` runs the standalone gated-activation backward after ``fc2_dgrad``; + ``hidden_states`` is the saved ``2F`` pre-activation for act-bwd / wgrad recompute. ``wgrad_out`` (optional ``[E, out, in]`` accumulator) folds the wgrad straight into the caller's buffer (``+=`` if ``wgrad_accumulate`` else ``=``) instead of returning a fresh - tensor. FC2 applies via transpose after the FlyDSL kernel. When used, ``result.wgrad_applied`` - is ``True`` and ``result.wgrad_stacked`` is that same buffer. + tensor. When used, ``result.wgrad_applied`` is ``True`` and ``result.wgrad_stacked`` is + that same buffer. """ grad_output = grad_output.contiguous() route_space = getattr(routing, "route_space", False) @@ -1117,9 +1078,8 @@ def permute_free_grouped_gemm_backward( wgrad_applied = False if route_space: - # FC2: grad is token-space [num_recv, out]. dgrad gathers back to the compact route - # buffer [T * min(topk, E), F] (GEMM operand), then the gated-activation backward maps - # dL/dF -> dL/d(2F) (+ dprob) when the forward fused the FC2 prologue. + # FC2: grad is token-space [num_recv, out]. dgrad gathers back to the route buffer, + # then the standalone gated-activation backward maps dL/dF -> dL/d(2F) (+ dprob). if requires_dgrad: weights_stacked = _stack_expert_weights(weights) # [E, H, F], zero-copy when grouped dgrad_f = permute_free_grouped_gemm_bf16_fc2_dgrad( diff --git a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py index 3189e46ec..731043707 100644 --- a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py +++ b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py @@ -32,9 +32,6 @@ _LDS_PAD = 8 _LDS_LIMIT = 163840 # gfx950 per-workgroup LDS (160 KB) -_WARP_CHOICES = [1, 2, 4, 8, 16] - - def _env_flag(name: str, default: bool) -> bool: v = os.environ.get(name) if v is None: @@ -140,37 +137,6 @@ def _warp_valid(block_m, block_n, block_k, wm, wn, transpose_b=False): return True -def _pick_warps(block_m: int, block_n: int, block_k: int, transpose_b=False): - """Pick (warps_m, warps_n) balancing per-warp MFMA tile (M_STEPS x N_STEPS) vs occupancy. - - Prefers keeping the per-warp atom counts moderate (fewer accumulators -> more waves) while - landing a 256-512 thread workgroup, which measured fastest across the Qwen MoE shapes. - """ - wmma = _mfma_dim(transpose_b) - best = None - for wm in _WARP_CHOICES: - for wn in _WARP_CHOICES: - if not _warp_valid(block_m, block_n, block_k, wm, wn, transpose_b): - continue - n_threads = wm * wn * _WARP - if n_threads > 512: - continue - m_steps = block_m // (wm * wmma) - n_steps = block_n // (wn * wmma) - # Favor a small, balanced per-warp atom footprint (fewer accumulators -> more - # waves), then a 256-512 thread workgroup. Ties broken toward |m_steps-n_steps| - # small (balanced reuse of A and B fragments). - score = ( - m_steps * n_steps, - abs(m_steps - n_steps), - 0 if 256 <= n_threads <= 512 else 1, - n_threads, - ) - if best is None or score < best[0]: - best = (score, (wm, wn)) - return best[1] if best is not None else None - - def _fwd_buffering(): """(num_buffers, lds_pad) for the production DMA+swizzle fill path. @@ -191,14 +157,6 @@ def _lds_bytes(block_m, block_n, block_k, transpose_b=False): return (a_tile + b_tile) * nbuf * 2 -def _default_block_n(block_m, block_k, transpose_b=False): - """Widest N tile (128 then 64) that fits LDS -- wider N raises arithmetic intensity.""" - for bn in (128, 64): - if _lds_bytes(block_m, bn, block_k, transpose_b) <= _LDS_LIMIT and bn % _mfma_dim(transpose_b) == 0: - return bn - return 64 - - def flydsl_moe_fwd_supported( A: torch.Tensor, B: torch.Tensor, @@ -270,9 +228,8 @@ def flydsl_moe_fwd( # Measured on FC1 no-act (qwen235b, block_m=256), idle machine, alias scopes on: # 256x32 w4x4 ~1628us, 128x64 w4x4 ~1694us, 128x64 w8x2 ~1726us. # -# An exhaustive sweep of the valid space (all block_n in 64..512 x block_k in 32..256 x warp -# grids >=4 waves; benchmarks/microbenchmarks/sweep_fc1_configs.py) found nothing better, so -# the list below is not missing a winner. Two directions are dead ends and are deliberately +# An exhaustive sweep of the valid tile space found nothing better, so the list below is not +# missing a winner. Two directions are dead ends and are deliberately # absent: block_n>=384 costs 4-22x (per-wave accumulator spill plus 120-144KB LDS pinning # occupancy to 1 workgroup), and trading block_m down to reach block_k=128 -- 4x fewer # barriers, which the ATT trace makes look attractive -- costs 59% (2581us at block_m=128 diff --git a/transformer_engine/pytorch/moe/pf_helper_kernels.py b/transformer_engine/pytorch/moe/pf_helper_kernels.py index 6a106c1d1..ad6793d8b 100644 --- a/transformer_engine/pytorch/moe/pf_helper_kernels.py +++ b/transformer_engine/pytorch/moe/pf_helper_kernels.py @@ -613,8 +613,7 @@ def _expert_meta_kernel( tl.store(block_start_ptr + offs_e, block_start, mask=mask_e) tl.store(ntpp_ptr, total_blocks * BLOCK_SIZE_M) - # expert_ids[b] = #{e : cblocks[e] <= b} (== searchsorted(cblocks, b, right=True)), - # then -1 for blocks past the real extent. + # expert_ids[b] = #{e : cblocks[e] <= b} then -1 for blocks past the real extent. offs_b = pid * BLOCK_B + tl.arange(0, BLOCK_B) mask_b = offs_b < blocks_max cblocks_valid = tl.where(mask_e, cblocks, (1 << 30)) @@ -751,11 +750,14 @@ def route_list_align( counts, within = scan # Per-expert placement metadata + per-block expert ids (single launch). - blocks_per_expert = torch.empty((E,), dtype=torch.int32, device=device) - block_start = torch.empty((E,), dtype=torch.int32, device=device) - expert_ids = torch.empty((blocks_max,), dtype=torch.int32, device=device) - num_tokens_post_padded = torch.empty((1,), dtype=torch.int32, device=device) + blocks_per_expert = torch.empty((E,), dtype=torch.int32, device=device) # [E]: ceil(count[e]/block_m) + block_start = torch.empty((E,), dtype=torch.int32, device=device) # [E]: first M-block index of expert e + expert_ids = torch.empty((blocks_max,), dtype=torch.int32, device=device) # [blocks_max]: owner of M-block b (-1 past end) + num_tokens_post_padded = torch.empty((1,), dtype=torch.int32, device=device) # [1] device scalar: real em (padded route count) block_b = 256 + # From per-expert route counts, derive the expert-sorted block layout: how many M-blocks + # each expert needs, where each expert's slots start (block_start), which expert owns each + # M-block (expert_ids), and the real padded route extent (num_tokens_post_padded). _expert_meta_kernel[(triton.cdiv(blocks_max, block_b),)]( counts, blocks_per_expert, @@ -781,8 +783,9 @@ def route_list_align( token_routes = torch.empty((1,), dtype=torch.int32, device=device) # unused stub token_route_count = token_routes - # Scatter each routed cell into its deterministic block-padded slot. The sentinel-init - # buffer is filled once (one launch); the place kernel then overwrites the routed slots. + # One program per token: for each routed (t, e), write sorted_slot_ids[slot] = t where + # slot = block_start[e] * block_m + within[e, t]. Unwritten slots stay at sentinel T. + # Optionally emits token_routes / token_route_count (token -> padded slot indices). sorted_slot_ids = torch.full((em_max,), T, dtype=torch.int32, device=device) _route_list_place_kernel[(T,)]( routing_map, From 998de36ee76c23b3f8392bb4d459d4554a531d41 Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Thu, 6 Aug 2026 16:32:50 +0000 Subject: [PATCH 39/43] Share wgrad align buffers between FC1 and FC2 in MoE routing metadata Introduces a mutable `WgradAlign` holder in `MoERoutingMetadata` so the original metadata and its `dataclasses.replace` copy (FC2 route-space view) build the block-`CONTRACT_M` align buffers only once. Replaces the standalone `wgrad_*` fields with the shared holder and updates the permute-free grouped GEMM wgrad path to read from it. --- transformer_engine/pytorch/moe/moe_routing.py | 36 +++++++++++++++---- .../pytorch/moe/permute_free_grouped_gemm.py | 33 +++++++++-------- 2 files changed, 48 insertions(+), 21 deletions(-) diff --git a/transformer_engine/pytorch/moe/moe_routing.py b/transformer_engine/pytorch/moe/moe_routing.py index db5a0ebc0..19a717167 100644 --- a/transformer_engine/pytorch/moe/moe_routing.py +++ b/transformer_engine/pytorch/moe/moe_routing.py @@ -12,6 +12,24 @@ import torch +@dataclass +class WgradAlign: + """Lazily-built block-``CONTRACT_M`` align buffers for the route-list wgrad kernel. + + Held in a *mutable* container so that a metadata and its ``dataclasses.replace`` copy + (e.g. the FC2 ``route_space=True`` view) share one instance by reference. The wgrad align + is a pure function of ``routing_map`` (identical for FC1 and FC2), but it is built lazily in + the backward -- after the ``replace`` that shares the fwd align has already happened. Sharing + this holder lets whichever backward runs first build the buffers once; the other reuses them, + avoiding a duplicate align build (and a second live copy). + """ + + sorted_slot_ids: Optional[torch.Tensor] = None + block_start: Optional[torch.Tensor] = None + blocks_per_expert: Optional[torch.Tensor] = None + block_size: Optional[int] = None + + @dataclass class MoERoutingMetadata: """Routing tensors for the route-list gather-in-GEMM MoE path. @@ -62,8 +80,10 @@ class MoERoutingMetadata: are unused padding. block_size_m: ``BLOCK_SIZE_M`` used to build the fwd/dgrad align buffers. - wgrad_*: - Separate block-``CONTRACT_M`` align buffers for the route-list wgrad kernel. + wgrad_align: + Shared, lazily-built block-``CONTRACT_M`` align buffers for the route-list wgrad kernel + (see :class:`WgradAlign`). The holder is shared by reference across a metadata and its + ``dataclasses.replace`` copy, so FC1 and FC2 build the (identical) wgrad align only once. route_counts / route_within: Cached block-size-independent scan (per-expert counts and within-expert ranks) shared by the fwd/dgrad and wgrad align builds. @@ -80,10 +100,10 @@ class MoERoutingMetadata: token_routes: Optional[torch.Tensor] = None token_route_count: Optional[torch.Tensor] = None block_size_m: Optional[int] = None - wgrad_sorted_slot_ids: Optional[torch.Tensor] = None - wgrad_block_start: Optional[torch.Tensor] = None - wgrad_blocks_per_expert: Optional[torch.Tensor] = None - wgrad_block_size: Optional[int] = None + # Shared mutable holder so a metadata and its ``dataclasses.replace`` copy (FC2 route-space + # view) reuse one wgrad align build. ``replace`` copies this reference, and ``__post_init__`` + # only creates a fresh holder when the field is genuinely absent, preserving the shared one. + wgrad_align: Optional[WgradAlign] = None # Block-size-independent scan (per-expert counts + within-expert ranks), shared by the # fwd/dgrad and wgrad align builds so it is computed once per routing map. route_counts: Optional[torch.Tensor] = None @@ -94,6 +114,10 @@ def __post_init__(self): # expert), so the caller can pass just ``routing_map``. if self.num_experts is None: self.num_experts = int(self.routing_map.size(1)) + # Fresh holder only for a genuinely new metadata; a ``replace`` copy passes the existing + # (possibly already-populated) holder through so FC1 and FC2 share the same wgrad align. + if self.wgrad_align is None: + self.wgrad_align = WgradAlign() @property def num_recv_tokens(self) -> int: diff --git a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py index 9ae495b36..e3e261635 100644 --- a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py +++ b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py @@ -305,11 +305,14 @@ def prepare_moe_align(metadata: MoERoutingMetadata, block_m: int) -> MoERoutingM def _prepare_wgrad_align( metadata: MoERoutingMetadata, contract_m: int ) -> MoERoutingMetadata: - """Build and cache the block-``contract_m`` align buffers for the wgrad kernel.""" - if ( - metadata.wgrad_sorted_slot_ids is not None - and metadata.wgrad_block_size == contract_m - ): + """Build and cache the block-``contract_m`` align buffers for the wgrad kernel. + + The buffers live in the shared :class:`WgradAlign` holder, so the FC1 metadata and its + ``dataclasses.replace`` FC2 route-space copy reuse a single build (whichever backward runs + first fills the holder; the other hits the cache below). + """ + cache = metadata.wgrad_align + if cache.sorted_slot_ids is not None and cache.block_size == contract_m: return metadata ( @@ -327,10 +330,10 @@ def _prepare_wgrad_align( scan=_ensure_route_scan(metadata), topk=metadata.topk, ) - metadata.wgrad_sorted_slot_ids = sorted_slot_ids - metadata.wgrad_block_start = block_start - metadata.wgrad_blocks_per_expert = blocks_per_expert - metadata.wgrad_block_size = contract_m + cache.sorted_slot_ids = sorted_slot_ids + cache.block_start = block_start + cache.blocks_per_expert = blocks_per_expert + cache.block_size = contract_m return metadata @@ -725,9 +728,9 @@ def permute_free_grouped_gemm_bf16_wgrad( x, grad_output, out, - routing.wgrad_sorted_slot_ids, - routing.wgrad_block_start, - routing.wgrad_blocks_per_expert, + routing.wgrad_align.sorted_slot_ids, + routing.wgrad_align.block_start, + routing.wgrad_align.blocks_per_expert, grad_base, num_recv_tokens=routing.num_recv_tokens, accumulate=bool(accumulate), @@ -744,9 +747,9 @@ def permute_free_grouped_gemm_bf16_wgrad( x, grad_output, dW, - routing.wgrad_sorted_slot_ids, - routing.wgrad_block_start, - routing.wgrad_blocks_per_expert, + routing.wgrad_align.sorted_slot_ids, + routing.wgrad_align.block_start, + routing.wgrad_align.blocks_per_expert, grad_base, num_recv_tokens=routing.num_recv_tokens, accumulate=bool(accumulate), From 7444757adec88ed9dcacbb9dbe1ff3142adfef3a Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Thu, 6 Aug 2026 16:59:33 +0000 Subject: [PATCH 40/43] Fuse FC2 activation recompute into the gated-activation backward pass Adds a `return_fc2_input` option to `permute_free_gated_act_bwd` (and `emit_act` in the Triton kernel) that re-materializes the `F`-wide FC2 input while the kernel already streams the `2F` preactivation. `permute_free_grouped_gemm_backward` now uses this fused path when both dgrad and wgrad are required, letting the wgrad consume the emitted activation instead of running a separate recompute. Renames `permute_free_gated_act_recompute` to `permute_free_gated_act_fwd` everywhere and adds a unit test verifying the fused dgrad+wgrad output matches the split references. --- .../benchmark_perm_free_grouped_gemm.py | 4 +- .../pytorch/test_perm_free_grouped_linear.py | 60 ++++++++++++ transformer_engine/pytorch/moe/__init__.py | 4 +- transformer_engine/pytorch/moe/moe_routing.py | 2 +- .../pytorch/moe/permute_free_grouped_gemm.py | 97 +++++++++++++------ .../pytorch/moe/pf_helper_kernels.py | 44 ++++++++- 6 files changed, 172 insertions(+), 39 deletions(-) diff --git a/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py index 014299d56..59ea1d05d 100644 --- a/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py +++ b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py @@ -187,12 +187,12 @@ def _permute_free_act_gemm( ``weights`` is the gate+up projection ``[E, 2F, K]``; returns the F-wide activated buffer. """ from transformer_engine.pytorch.moe import ( - permute_free_gated_act_recompute, + permute_free_gated_act_fwd, permute_free_grouped_gemm_bf16, ) preact = permute_free_grouped_gemm_bf16(hidden, weights, routing) - return permute_free_gated_act_recompute(preact, routing, activation="silu") + return permute_free_gated_act_fwd(preact, routing, activation="silu") def _build_backend_fns( diff --git a/tests/pytorch/test_perm_free_grouped_linear.py b/tests/pytorch/test_perm_free_grouped_linear.py index 6858997cc..1d42f48ec 100644 --- a/tests/pytorch/test_perm_free_grouped_linear.py +++ b/tests/pytorch/test_perm_free_grouped_linear.py @@ -361,6 +361,66 @@ def test_fc2_backward_dispatch_recompute_matches_stored(): assert _rel_l2(recompute.wgrad_stacked, stored.wgrad_stacked) < 1e-2 +def test_fc2_backward_dispatch_fused_dgrad_wgrad_matches_split(): + """FC2 backward with both grads + a gated activation: the act-bwd re-emits the F-wide + fc2_input for the wgrad (one fused pass over the 2F preact). Both dgrad and wgrad must match + running dgrad-only and wgrad-only (separate recompute) independently.""" + from transformer_engine.pytorch.moe import ( + permute_free_grouped_gemm_backward, + prepare_moe_align, + ) + + torch.manual_seed(43) + num_recv_tokens, in_features, out_features = 128, 96, 128 # F=in, H=out (W2 is [E, H, F]) + num_experts, max_hits = 8, 4 + + routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=12) + routing = PermuteFreeMetadata( + routing_map=routing_map, num_experts=num_experts, route_space=True, activation="silu" + ) + prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) + + em_max = int(routing.sorted_slot_ids.shape[0]) + valid = routing.sorted_slot_ids < num_recv_tokens + + preact = torch.zeros(em_max, 2 * in_features, device="cuda", dtype=torch.bfloat16) + preact[valid] = torch.randn( + int(valid.sum().item()), 2 * in_features, device="cuda", dtype=torch.bfloat16 + ) + grad_output = torch.randn(num_recv_tokens, out_features, device="cuda", dtype=torch.bfloat16) + probs = torch.rand( + num_recv_tokens, num_experts, device="cuda", dtype=torch.float32, requires_grad=True + ) + weights = [ + torch.randn(out_features, in_features, device="cuda", dtype=torch.bfloat16) + for _ in range(num_experts) + ] + + common = dict( + routing=routing, weights=weights, num_gemms=num_experts, + hidden_states=preact, fc2_activation="silu", dispatched_probs=probs, + ) + # Fused: dgrad + wgrad in one call -> the wgrad reuses the act-bwd's fc2_input. + fused = permute_free_grouped_gemm_backward( + grad_output, requires_dgrad=True, requires_wgrad=True, **common + ) + # Split references: dgrad-only (act-bwd) and wgrad-only (own recompute-from-preact). + dref = permute_free_grouped_gemm_backward( + grad_output, requires_dgrad=True, requires_wgrad=False, **common + ) + wref = permute_free_grouped_gemm_backward( + grad_output, requires_dgrad=False, requires_wgrad=True, **common + ) + + assert fused.dgrad.shape == (em_max, 2 * in_features) + assert fused.wgrad_stacked.shape == (num_experts, out_features, in_features) + # dgrad is bit-identical (same act-bwd, EMIT_ACT only adds a store). + assert torch.equal(fused.dgrad, dref.dgrad) + assert torch.equal(fused.grad_probs, dref.grad_probs) + # wgrad from the fused (re-emitted) act matches the standalone recompute to bf16 rounding. + assert _rel_l2(fused.wgrad_stacked, wref.wgrad_stacked) < 1e-2 + + def test_fc1_fc2_gated_pipeline(monkeypatch): """End-to-end FC1 raw 2F -> FC2 standalone gated activation: forward outputs and backward grads match a PyTorch reference through the permute-free GroupedLinear modules.""" diff --git a/transformer_engine/pytorch/moe/__init__.py b/transformer_engine/pytorch/moe/__init__.py index 4b627b2fd..e4313ae45 100644 --- a/transformer_engine/pytorch/moe/__init__.py +++ b/transformer_engine/pytorch/moe/__init__.py @@ -18,7 +18,7 @@ permute_free_grouped_gemm_bf16_wgrad, permute_free_grouped_gemm_forward, permute_free_gated_act_bwd, - permute_free_gated_act_recompute, + permute_free_gated_act_fwd, prepare_moe_align, ) @@ -36,6 +36,6 @@ "permute_free_grouped_gemm_bf16_wgrad", "permute_free_grouped_gemm_forward", "permute_free_gated_act_bwd", - "permute_free_gated_act_recompute", + "permute_free_gated_act_fwd", "prepare_moe_align", ] diff --git a/transformer_engine/pytorch/moe/moe_routing.py b/transformer_engine/pytorch/moe/moe_routing.py index 19a717167..ff51c36e6 100644 --- a/transformer_engine/pytorch/moe/moe_routing.py +++ b/transformer_engine/pytorch/moe/moe_routing.py @@ -151,7 +151,7 @@ class PermuteFreeMetadata(MoERoutingMetadata): Gated activation for the FC2 standalone pass -- ``"silu"`` or ``"gelu"``. ``None`` leaves activation to the caller. FC1 emits raw ``2F``; this hint is consumed on the FC2 direction (``route_space=True``) to run - :func:`permute_free_gated_act_recompute` before the plain FC2 GEMM. + :func:`permute_free_gated_act_fwd` before the plain FC2 GEMM. (The per-route gating probabilities are *not* carried here: they need a gradient, so they are passed as a separate autograd tensor argument to the module rather than as metadata.) diff --git a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py index e3e261635..aff17feda 100644 --- a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py +++ b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py @@ -473,7 +473,8 @@ def permute_free_gated_act_bwd( *, activation: str, dispatched_probs: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + return_fc2_input: bool = False, +): """Activation (+ route-prob) backward for the fused permute-free FC1 epilogue. Reconstructs the raw ``2F`` GEMM-output gradient ``dpre = [d_gate | d_up]`` (and, when the @@ -491,10 +492,16 @@ def permute_free_gated_act_bwd( dispatched_probs: ``[num_recv_tokens, E]`` gating probs (or ``None`` for a silu/gelu-only fusion). + ``return_fc2_input`` additionally re-materialises the ``F``-wide forward activation + ``fc2_input = act(gate)*up[*prob]`` inside this pass (one extra multiply + store from values + already in registers) so the FC2 wgrad can consume it directly instead of re-streaming the + ``2F`` preact through a separate recompute -- eliminating a full ``2F`` HBM read. + Returns ------- - ``(dpre, dprob)`` -- ``dpre`` is ``[T * min(topk, E), 2F]`` bf16; ``dprob`` matches - ``dispatched_probs`` (or ``None`` when no probs were fused). + ``(dpre, dprob)`` by default, or ``(dpre, dprob, fc2_input)`` when ``return_fc2_input`` -- + ``dpre`` is ``[T * min(topk, E), 2F]`` bf16; ``dprob`` matches ``dispatched_probs`` (or + ``None`` when no probs were fused); ``fc2_input`` is ``[em_max, F]`` bf16. """ # Block-padded canonical layout: grad_out (FC2 dgrad) and preact both live in the [em_max] # padded slot order, and the kernel indexes grad_out/preact/dpre by the same row as @@ -506,7 +513,7 @@ def permute_free_gated_act_bwd( expert = _expert_per_route(routing, em_max) # ``num_tokens_post_padded`` is a device scalar bounding the real padded extent, so the tail # programs exit early (sync-free) instead of streaming the padding through HBM. - return fused_gated_act_prob_bwd( + dpre, dprob, fc2_input = fused_gated_act_prob_bwd( grad_output.contiguous(), preact, token, @@ -515,23 +522,29 @@ def permute_free_gated_act_bwd( activation=activation, dispatched_probs=dispatched_probs, num_routes_bound=routing.num_tokens_post_padded, + emit_act=return_fc2_input, ) + if return_fc2_input: + return dpre, dprob, fc2_input + return dpre, dprob -def permute_free_gated_act_recompute( +def permute_free_gated_act_fwd( preact: torch.Tensor, routing: MoERoutingMetadata, *, activation: str, dispatched_probs: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Rebuild the fused FC1 activation ``act(gate) * up * prob`` from the saved pre-activation. + """Apply the fused gated activation ``act(gate) * up * prob`` to the ``2F`` pre-activation. - Lets the backward checkpoint only the ``2F`` pre-activation (``[T * min(topk, E), 2F]``) and - reconstruct the ``F``-wide activation just-in-time for the FC2 wgrad, feeding the - *unchanged* full-speed wgrad kernel a transient buffer (freed right after) instead of - persisting the activation across the fwd/bwd boundary. Uses the same per-route routing - arrays as :func:`permute_free_gated_act_bwd`. + Maps the raw ``[T * min(topk, E), 2F]`` ``[gate | up]`` GEMM output to the ``F``-wide FC2 + input in the block-padded slot layout. Used both by the forward (to form the FC2 input) and, + on the backward, to reconstruct that ``F``-wide activation just-in-time for the FC2 wgrad -- + so the backward can checkpoint only the ``2F`` pre-activation and feed the *unchanged* + full-speed wgrad kernel a transient buffer (freed right after) instead of persisting the + activation across the fwd/bwd boundary. Uses the same per-route routing arrays as + :func:`permute_free_gated_act_bwd`. Returns ------- @@ -540,7 +553,7 @@ def permute_free_gated_act_recompute( forward output and the layout the wgrad route-reads). """ # Block-padded canonical layout: emit one row per padded slot (keyed by sorted_slot_ids) so - # the rebuilt activation lines up with the [em_max] grad the wgrad route-reads. Padding slots + # the activation lines up with the [em_max] grad the wgrad route-reads. Padding slots # (token sentinel >= num_recv_tokens) are masked to zero by the kernel. em_max = int(routing.sorted_slot_ids.shape[0]) token = routing.sorted_slot_ids.to(torch.int32) @@ -766,7 +779,7 @@ def permute_free_grouped_gemm_bf16_fc2( """Route-list FC2 forward with gather-combine to token order (bf16 MoE). ``fc2_input`` is the route-ordered ``F``-wide FC1 activation (or a transient buffer rebuilt - from the saved ``2F`` pre-activation via :func:`permute_free_gated_act_recompute`). For each + from the saved ``2F`` pre-activation via :func:`permute_free_gated_act_fwd`). For each route the kernel reads its row (``index_a_by_route_pos=True``) and computes the compact per-route GEMM; a separate contention-free gather-combine pass sums each token's routes into its output row. @@ -957,7 +970,7 @@ def permute_free_grouped_gemm_bf16_fc2_wgrad( if preact is not None: if activation is None: raise ValueError("permute_free FC2 wgrad recompute requires an activation.") - fc2_input = permute_free_gated_act_recompute( + fc2_input = permute_free_gated_act_fwd( preact, routing, activation=activation, dispatched_probs=dispatched_probs ) @@ -1038,7 +1051,7 @@ def permute_free_grouped_gemm_forward( if getattr(routing, "route_space", False): fc2_input = hidden_states if activation is not None: - fc2_input = permute_free_gated_act_recompute( + fc2_input = permute_free_gated_act_fwd( hidden_states, routing, activation=activation, @@ -1083,33 +1096,55 @@ def permute_free_grouped_gemm_backward( if route_space: # FC2: grad is token-space [num_recv, out]. dgrad gathers back to the route buffer, # then the standalone gated-activation backward maps dL/dF -> dL/d(2F) (+ dprob). + recompute = fc2_activation is not None and hidden_states is not None + # When both grads are wanted with a gated activation, fuse the wgrad's activation + # recompute into the act-bwd (which is already streaming the 2F preact) so the wgrad + # can reuse the F-wide fc2_input instead of re-reading the 2F preact a second time. + fused_fc2_input = None if requires_dgrad: weights_stacked = _stack_expert_weights(weights) # [E, H, F], zero-copy when grouped dgrad_f = permute_free_grouped_gemm_bf16_fc2_dgrad( grad_output, weights_stacked, routing ) - if fc2_activation is not None and hidden_states is not None: - dgrad, grad_probs = permute_free_gated_act_bwd( - dgrad_f, - hidden_states, - routing, - activation=fc2_activation, - dispatched_probs=dispatched_probs, - ) + if recompute: + if requires_wgrad: + dgrad, grad_probs, fused_fc2_input = permute_free_gated_act_bwd( + dgrad_f, + hidden_states, + routing, + activation=fc2_activation, + dispatched_probs=dispatched_probs, + return_fc2_input=True, + ) + else: + dgrad, grad_probs = permute_free_gated_act_bwd( + dgrad_f, + hidden_states, + routing, + activation=fc2_activation, + dispatched_probs=dispatched_probs, + ) if not (dispatched_probs is not None and dispatched_probs.requires_grad): grad_probs = None else: dgrad = dgrad_f if requires_wgrad: weights_shape = (num_gemms, weights[0].size(0), weights[0].size(1)) - recompute = fc2_activation is not None and hidden_states is not None - dW = permute_free_grouped_gemm_bf16_fc2_wgrad( - hidden_states, grad_output, weights_shape, routing, - out=wgrad_out, accumulate=wgrad_accumulate, - preact=hidden_states if recompute else None, - dispatched_probs=dispatched_probs if recompute else None, - activation=fc2_activation if recompute else None, - ) # [E, H, F] + if fused_fc2_input is not None: + # Reuse the F-wide activation emitted by the act-bwd above: a plain stored-act + # wgrad, no second pass over the 2F preact. + dW = permute_free_grouped_gemm_bf16_fc2_wgrad( + fused_fc2_input, grad_output, weights_shape, routing, + out=wgrad_out, accumulate=wgrad_accumulate, + ) # [E, H, F] + else: + dW = permute_free_grouped_gemm_bf16_fc2_wgrad( + hidden_states, grad_output, weights_shape, routing, + out=wgrad_out, accumulate=wgrad_accumulate, + preact=hidden_states if recompute else None, + dispatched_probs=dispatched_probs if recompute else None, + activation=fc2_activation if recompute else None, + ) # [E, H, F] wgrad_stacked = dW wgrad_applied = wgrad_out is not None else: diff --git a/transformer_engine/pytorch/moe/pf_helper_kernels.py b/transformer_engine/pytorch/moe/pf_helper_kernels.py index ad6793d8b..79fd1bd28 100644 --- a/transformer_engine/pytorch/moe/pf_helper_kernels.py +++ b/transformer_engine/pytorch/moe/pf_helper_kernels.py @@ -152,6 +152,7 @@ def _gated_act_prob_bwd_kernel( expert_ptr, # [routes_max] route -> local expert dpre_ptr, # [T * min(topk, E), 2F] out: grad wrt the raw 2F GEMM output grad_probs_ptr, # [num_recv_tokens, E] out (fp32) + act_out_ptr, # [T * min(topk, E), F] out: fused fc2_input = act(g)*u*prob (EMIT_ACT only) nbound_ptr, # [1] int32 device scalar: dynamic upper bound on compact routes num_recv_tokens, F, @@ -165,8 +166,11 @@ def _gated_act_prob_bwd_kernel( stride_dpreh, stride_gpm, stride_gpe, + stride_aom, + stride_aoh, ACTIVATION: tl.constexpr, HAS_PROBS: tl.constexpr, + EMIT_ACT: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_H: tl.constexpr, ): @@ -184,6 +188,12 @@ def _gated_act_prob_bwd_kernel( over-allocated routes carry the ``token == num_recv_tokens`` sentinel and are masked out of the prob load/store; their ``dpre`` rows are inert (ignored downstream). + When ``EMIT_ACT`` is set the kernel *also* stores the ``F``-wide forward activation + ``fc2_input = act(gate)*up[*prob]`` to ``act_out_ptr`` -- the same buffer the FC2 wgrad + would otherwise recompute in a second full pass over the ``2F`` preact. Since ``act(gate)``, + ``up`` and ``prob`` are already live in registers here, this is one extra multiply + store + and removes the redundant ``2F`` HBM read (see :func:`fused_gated_act_prob_fwd`). + The compact route buffers are statically over-allocated to the worst-case ``routes_max = T * topk`` (sync-free shape bound), but the real routes occupy only the dense head ``[0, num_routes)`` -- under expert parallelism this can be ~topk*E_local/E @@ -242,6 +252,17 @@ def _gated_act_prob_bwd_kernel( d_up.to(dpre_ptr.dtype.element_ty), mask=m, ) + if EMIT_ACT: + # Re-emit the F-wide forward activation from values already in registers, so the + # FC2 wgrad can consume it directly instead of re-streaming the 2F preact. + fi = act_g * u + if HAS_PROBS: + fi = fi * prob[:, None] + tl.store( + act_out_ptr + r_offs[:, None] * stride_aom + offs[None, :] * stride_aoh, + fi.to(act_out_ptr.dtype.element_ty), + mask=m, + ) if HAS_PROBS: tl.store( @@ -262,6 +283,7 @@ def fused_gated_act_prob_bwd( dispatched_probs: Optional[torch.Tensor] = None, grad_probs_shape: Optional[torch.Size] = None, num_routes_bound: Optional[torch.Tensor] = None, + emit_act: bool = False, ): """Backward of the fused gated-activation (+ route-prob) FC1 epilogue. @@ -283,12 +305,18 @@ def fused_gated_act_prob_bwd( populated; passing the actual extent (e.g. ``num_tokens_post_padded`` from the routing metadata) lets tail programs exit early instead of streaming the padding through HBM. When ``None`` the full static ``routes_max`` is used (no early exit). + emit_act: + When ``True`` the kernel additionally re-materialises the ``F``-wide forward activation + ``fc2_input = act(gate)*up[*prob]`` (``[T * min(topk, E), F]`` bf16) from the values it + already computes, so the FC2 wgrad can consume it directly instead of re-reading the ``2F`` + preact in a separate recompute pass. Returned as the third element (``None`` otherwise). Returns ------- - (dpre, grad_probs) + (dpre, grad_probs, fc2_input) ``dpre`` is ``[T * min(topk, E), 2F]`` (bf16), grad wrt the raw GEMM output; ``grad_probs`` is - ``[num_recv_tokens, E]`` (matching ``dispatched_probs.dtype``) or ``None``. + ``[num_recv_tokens, E]`` (matching ``dispatched_probs.dtype``) or ``None``; ``fc2_input`` is + ``[T * min(topk, E), F]`` (bf16) when ``emit_act`` else ``None``. """ em_max, F = grad_out.shape if preact.shape[0] != em_max or preact.shape[1] != 2 * F: @@ -299,6 +327,12 @@ def fused_gated_act_prob_bwd( has_probs = dispatched_probs is not None dpre = torch.empty((em_max, 2 * F), dtype=torch.bfloat16, device=grad_out.device) + if emit_act: + act_out = torch.empty((em_max, F), dtype=torch.bfloat16, device=grad_out.device) + stride_aom, stride_aoh = act_out.stride(0), act_out.stride(1) + else: + act_out = None + stride_aom = stride_aoh = 0 if has_probs: grad_probs = torch.zeros(dispatched_probs.shape, dtype=torch.float32, device=grad_out.device) probs_ptr = dispatched_probs @@ -323,6 +357,7 @@ def fused_gated_act_prob_bwd( expert, dpre, grad_probs if has_probs else grad_out, + act_out if emit_act else dpre, nbound, num_recv_tokens, F, @@ -336,12 +371,15 @@ def fused_gated_act_prob_bwd( dpre.stride(1), stride_gpm, stride_gpe, + stride_aom, + stride_aoh, ACTIVATION=act_id, HAS_PROBS=has_probs, + EMIT_ACT=emit_act, ) if has_probs: grad_probs = grad_probs.to(dispatched_probs.dtype) - return dpre, grad_probs + return dpre, grad_probs, act_out @triton.autotune(configs=_gated_act_bwd_autotune_configs(), key=["F", "HAS_PROBS"]) From d730141b98d8c5fe1f4f4a19b15826f3646c6a9e Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Thu, 6 Aug 2026 18:27:49 +0000 Subject: [PATCH 41/43] Remove FlyDSL permute-free forward autotuner and stale parameters Switches the permute-free grouped GEMM forward path to the fixed-tile v3 kernel directly, removing `flydsl_moe_fwd_autotuned`, the `_FWD_CACHE`, and unused tile knobs (`block_n`, `block_k`, `warps_m`, `warps_n`) from the wrapper and kernel signatures. Also drops the unused `c_m` argument from the FlyDSL compile helpers, the `block_start` routing argument, the `perm_free_route_space` variable, and unused imports/constants across the MoE routing, grouped linear, helper kernels, and benchmark files. --- .../benchmark_perm_free_grouped_gemm.py | 2 - .../pytorch/test_perm_free_grouped_linear.py | 2 +- .../permute_free_grouped_gemm/pf_dgrad.py | 2 - .../permute_free_grouped_gemm/pf_fwd.py | 2 - .../pytorch/module/grouped_linear.py | 1 - transformer_engine/pytorch/moe/moe_routing.py | 2 +- .../pytorch/moe/permute_free_grouped_gemm.py | 15 ++-- .../pytorch/moe/pf_fwd_wrapper.py | 89 ++----------------- .../pytorch/moe/pf_helper_kernels.py | 4 +- 9 files changed, 15 insertions(+), 104 deletions(-) diff --git a/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py index 59ea1d05d..ea9c61721 100644 --- a/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py +++ b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py @@ -27,8 +27,6 @@ from torch.utils.cpp_extension import IS_HIP_EXTENSION from benchmark_grouped_gemm import ( - EP_SIZE_LIST, - GROUPED_GEMM_M_SIZE_LIST, generate_deepseekv2_lite_test_cases, generate_deepseekv2_test_cases, generate_deepseekv3_test_cases, diff --git a/tests/pytorch/test_perm_free_grouped_linear.py b/tests/pytorch/test_perm_free_grouped_linear.py index 1d42f48ec..05949a8b0 100644 --- a/tests/pytorch/test_perm_free_grouped_linear.py +++ b/tests/pytorch/test_perm_free_grouped_linear.py @@ -953,7 +953,7 @@ def test_permute_free_backward_fc2_dgrad_probs(): def test_route_list_gated_act_bwd_gelu(): - """Gelu (tanh) variant of the standalone gated-activation backward (ACT_GELU path).""" + """Gelu (tanh) variant of the standalone gated-activation backward (gelu path).""" torch.manual_seed(79) num_recv_tokens, in_features = 128, 128 num_experts, max_hits = 8, 3 diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py index ddf48cd3d..acd91faf7 100644 --- a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py @@ -213,7 +213,6 @@ def grouped_gemm_dgrad_k( NUM_TILE_BLOCKS: fx.Int32, SORTED: fx.Tensor, A_ELEMS: fx.Int32, - c_m: fx.Int32, ): n_blocks = ceildiv(fx.Int32(Kout), BLOCK_N) lds = fx.SharedAllocator().allocate(SharedStorage).peek() @@ -286,7 +285,6 @@ def launch(GRAD_Y, WEIGHT, DX, TILE_TO_GROUP, NUM_TILE_BLOCKS: fx.Int32, SORTED, NUM_TILE_BLOCKS, SORTED, A_ELEMS, - c_m, value_attrs=make_value_attrs(waves_per_eu, agpr_alloc, "512,512"), ).launch(grid=(grid_x, 1, 1), block=(512, 1, 1), stream=stream) diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py index c70590cda..16fb7d903 100644 --- a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py @@ -228,7 +228,6 @@ def grouped_gemm_gather_k( NUM_TILE_BLOCKS: fx.Int32, SORTED: fx.Tensor, A_ELEMS: fx.Int32, - c_m: fx.Int32, c_n: fx.Int32, ): n_blocks = ceildiv(c_n, BLOCK_N) @@ -309,7 +308,6 @@ def launch(A, B, C, TILE_TO_GROUP, NUM_TILE_BLOCKS: fx.Int32, SORTED, A_ELEMS: f NUM_TILE_BLOCKS, SORTED, A_ELEMS, - c_m, c_n, value_attrs=make_value_attrs(waves_per_eu, agpr_alloc, "512,512"), ).launch(grid=(grid_x, 1, 1), block=(512, 1, 1), stream=stream) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 299d656e5..ab2947105 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -667,7 +667,6 @@ def forward( # Perform GEMM if use_perm_free_grouped_gemm: - perm_free_route_space = getattr(routing_metadata, "route_space", False) # FC1 emits raw 2F [gate|up]; the ``activation`` hint on the metadata is consumed on # FC2, which applies the gated activation in a standalone pass and then runs a plain # GEMM (the fused-prologue path regressed throughput). Route probs ride with FC2 too. diff --git a/transformer_engine/pytorch/moe/moe_routing.py b/transformer_engine/pytorch/moe/moe_routing.py index ff51c36e6..1f66c381e 100644 --- a/transformer_engine/pytorch/moe/moe_routing.py +++ b/transformer_engine/pytorch/moe/moe_routing.py @@ -7,7 +7,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Optional, Tuple +from typing import Optional import torch diff --git a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py index aff17feda..e884d1580 100644 --- a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py +++ b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py @@ -15,7 +15,7 @@ import os from dataclasses import dataclass -from typing import List, Optional, Tuple +from typing import Optional, Tuple import torch @@ -49,10 +49,10 @@ def _get_flydsl_fwd(): - """Return ``(flydsl_moe_fwd_autotuned, flydsl_moe_fwd_supported)``.""" - from .pf_fwd_wrapper import flydsl_moe_fwd_autotuned, flydsl_moe_fwd_supported + """Return ``(flydsl_moe_fwd, flydsl_moe_fwd_supported)``.""" + from .pf_fwd_wrapper import flydsl_moe_fwd, flydsl_moe_fwd_supported - return flydsl_moe_fwd_autotuned, flydsl_moe_fwd_supported + return flydsl_moe_fwd, flydsl_moe_fwd_supported def _expand_expert_ids_per_slot( @@ -115,21 +115,20 @@ def _pf_moe_fwd( block_m: int, index_a_by_route_pos: bool = False, ) -> None: - """Run the route-list gather-GEMM (forward or dgrad) via FlyDSL autotuning.""" - autotuned, supported = _get_flydsl_fwd() + """Run the route-list gather-GEMM (forward or dgrad) via the FlyDSL v3 kernels.""" + launch, supported = _get_flydsl_fwd() block_m = int(block_m) if not supported(A, B, block_m=block_m): raise RuntimeError( "FlyDSL grouped GEMM does not support these operands " f"(block_m={block_m}, A={tuple(A.shape)}, B={tuple(B.shape)})." ) - autotuned( + launch( A, B, C, routing.sorted_slot_ids, routing.expert_ids, - routing.block_start, num_recv_tokens=num_recv_tokens, block_m=block_m, index_a_by_route_pos=index_a_by_route_pos, diff --git a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py index 731043707..593431c5d 100644 --- a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py +++ b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py @@ -21,12 +21,10 @@ __all__ = [ "flydsl_moe_fwd", - "flydsl_moe_fwd_autotuned", "flydsl_moe_fwd_supported", "flydsl_moe_fwd_pick_block_m", ] -_WMMA = 16 _FILL_V = 8 _WARP = 64 _LDS_PAD = 8 @@ -162,8 +160,6 @@ def flydsl_moe_fwd_supported( B: torch.Tensor, *, block_m: int, - block_n: int = 64, - block_k: int = 64, ) -> bool: """Whether the in-tree FlyDSL fwd/dgrad kernels can handle these operands.""" if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: @@ -172,8 +168,7 @@ def flydsl_moe_fwd_supported( return False if B.stride(2) != 1 and B.stride(1) != 1: return False - em_max = None # only known at launch for C; skip em_max % block_m precheck here - _ = (block_n, block_k, em_max) + # em_max % block_m is only known at launch (from C), so it is not prechecked here. return True @@ -183,14 +178,9 @@ def flydsl_moe_fwd( C: torch.Tensor, sorted_slot_ids: torch.Tensor, expert_ids: torch.Tensor, - block_start: torch.Tensor, *, num_recv_tokens: int, block_m: int, - block_n: Optional[int] = None, - block_k: int = 64, - warps_m: Optional[int] = None, - warps_n: Optional[int] = None, index_a_by_route_pos: bool = False, ) -> None: """Route-list gather-GEMM forward, writing the compact ``C[em_max, WIDTH_N]`` in place. @@ -199,6 +189,9 @@ def flydsl_moe_fwd( ``index_a_by_route_pos=False``, else read at the compact route row). ``B`` is ``[num_experts, N_OUT, K]`` (contiguous inner ``K``). Gated activation and route-prob apply are **not** fused here; use the standalone helpers in :mod:`pf_helper_kernels`. + + The v3 (MegaMOE-ported) kernels use a fixed tile geometry, so only ``block_m`` is a tunable + knob (chosen once by the align via :func:`flydsl_moe_fwd_pick_block_m`). """ assert A.dtype == B.dtype == C.dtype == torch.bfloat16 assert A.stride(1) == 1, "A must be contiguous along the contraction (K)" @@ -256,10 +249,6 @@ def flydsl_moe_fwd( (256, 64, 2, 4), ] -# Winner cache: {shape/mode key -> (block_n, block_k, warps_m, warps_n)}. -_FWD_CACHE: dict = {} - - def _valid_config(block_m, bn, bk, wm, wn, K, transpose_b=False): if K % bk != 0 or bn % _mfma_dim(transpose_b) != 0: return False @@ -277,7 +266,7 @@ def flydsl_moe_fwd_pick_block_m( """Largest ``block_m`` in ``candidates`` the FlyDSL fwd can actually run for these operands, or ``None`` if the operands are unsupported at every candidate. - "Can run" == bf16 operands contiguous along the contraction AND at least one autotuner tile + "Can run" == bf16 operands contiguous along the contraction AND at least one candidate tile in :data:`_FWD_TUNE_CONFIGS` fits the per-workgroup LDS budget. Callers should include their default among ``candidates`` (the picker only walks high->low over what is passed), so a small-token workload is never padded up beyond what the caller offered. @@ -296,71 +285,3 @@ def flydsl_moe_fwd_pick_block_m( ): return block_m return None - - -def flydsl_moe_fwd_autotuned( - A: torch.Tensor, - B: torch.Tensor, - C: torch.Tensor, - sorted_slot_ids: torch.Tensor, - expert_ids: torch.Tensor, - block_start: torch.Tensor, - *, - num_recv_tokens: int, - block_m: int, - block_k: int = 64, - index_a_by_route_pos: bool = False, - warmup: int = 3, - iters: int = 10, -) -> None: - """Shape-autotuned :func:`flydsl_moe_fwd`. - - On the first call for a given (block_m, GEMM shape) the valid subset of - ``_FWD_TUNE_CONFIGS`` is benchmarked and the fastest ``(block_n, block_k, warps_m, warps_n)`` - is cached. The production DMA+swizzle fill path is always used; only tile geometry is swept. - """ - N_OUT, K = int(B.shape[1]), int(B.shape[2]) - width_n = int(C.shape[1]) - key = (int(block_m), N_OUT, K, width_n, bool(index_a_by_route_pos)) - - def _launch(bn, bk, wm, wn): - flydsl_moe_fwd( - A, B, C, sorted_slot_ids, expert_ids, block_start, - num_recv_tokens=num_recv_tokens, block_m=block_m, block_n=bn, block_k=bk, - warps_m=wm, warps_n=wn, index_a_by_route_pos=index_a_by_route_pos, - ) - - best = _FWD_CACHE.get(key) - if best is None: - candidates = [ - (bn, bk, wm, wn) - for (bn, bk, wm, wn) in _FWD_TUNE_CONFIGS - if _valid_config(block_m, bn, bk, wm, wn, K) - ] - if not candidates: - _launch(None, block_k, None, None) # heuristic fallback - return - best_t = None - for cfg in candidates: - try: - for _ in range(warmup): - _launch(*cfg) - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - torch.cuda.synchronize() - start.record() - for _ in range(iters): - _launch(*cfg) - end.record() - torch.cuda.synchronize() - t = start.elapsed_time(end) / iters - except Exception: # noqa: BLE001 -- skip configs that fail to compile/run - continue - if best_t is None or t < best_t: - best_t, best = t, cfg - if best is None: - _launch(None, block_k, None, None) - return - _FWD_CACHE[key] = best - - _launch(*best) diff --git a/transformer_engine/pytorch/moe/pf_helper_kernels.py b/transformer_engine/pytorch/moe/pf_helper_kernels.py index 79fd1bd28..c43dfc886 100644 --- a/transformer_engine/pytorch/moe/pf_helper_kernels.py +++ b/transformer_engine/pytorch/moe/pf_helper_kernels.py @@ -12,7 +12,7 @@ from __future__ import annotations -from typing import Any, Dict, Optional +from typing import Optional import torch import triton @@ -42,8 +42,6 @@ def _gelu_tanh(x): # Activation selector for the fused epilogue. Keep the set small and explicit; add a helper # above and an entry here to extend. Passed to the kernel as a compile-time int so each # activation specializes to its own instance (no runtime branch in the hot path). -ACT_SILU: tl.constexpr = 0 -ACT_GELU: tl.constexpr = 1 _ACT_IDS = {"silu": 0, "gelu": 1} From ad7831e6c8795f9319262a512ddcbad74939cc25 Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Thu, 6 Aug 2026 21:30:32 +0000 Subject: [PATCH 42/43] Drop remaining v3 naming and stale permute-free wrapper code Removes leftover "v3" references from docstrings and comments across the FlyDSL permute-free kernels and tests. Simplifies forward block_m selection to a single token-threshold helper and deletes the unused picker/autotuner remnants in `pf_fwd_wrapper`. Unifies dgrad dispatch by passing the forward weight tensor directly with a new `dgrad` flag instead of building a transposed view, and renames the wrapper dispatch helper from `_run_v3_fwd` to `_run_gather_gemm`. --- .../pytorch/test_perm_free_grouped_linear.py | 9 +- .../permute_free_grouped_gemm/pf_dgrad.py | 6 +- .../permute_free_grouped_gemm/pf_fwd.py | 23 +- .../pytorch/moe/permute_free_grouped_gemm.py | 66 ++---- .../pytorch/moe/pf_fwd_wrapper.py | 223 +++--------------- 5 files changed, 66 insertions(+), 261 deletions(-) diff --git a/tests/pytorch/test_perm_free_grouped_linear.py b/tests/pytorch/test_perm_free_grouped_linear.py index 05949a8b0..c5e623fc4 100644 --- a/tests/pytorch/test_perm_free_grouped_linear.py +++ b/tests/pytorch/test_perm_free_grouped_linear.py @@ -438,8 +438,7 @@ def test_fc1_fc2_gated_pipeline(monkeypatch): fc1_meta = PermuteFreeMetadata( routing_map=routing_map, num_experts=num_experts, activation="silu" ) - # Prepare at the v3-compatible forward block_m (>= 128, multiple of 128); the token-count - # default (64 at 128 tokens) is below the v3 gather/dgrad tile minimum. FC1 and FC2 share the + # Prepare at the FlyDSL-compatible forward block_m (>= 128, multiple of 128). FC1 and FC2 share the # same block-padded slot layout, so the derived fc2_meta inherits this align. prepare_moe_align(fc1_meta, _FLYDSL_FWD_BLOCK_M) fc2_meta = dataclasses.replace(fc1_meta, route_space=True) @@ -510,7 +509,7 @@ def test_fc1_fc2_gated_pipeline(monkeypatch): def test_route_list_fwd_dgrad_wgrad_consistency(): """fwd + dgrad + wgrad against a single autograd reference in compact route space.""" torch.manual_seed(23) - # The v3 gather/dgrad tiles need both contraction dims (in for fwd, out for dgrad) to be a + # The gather/dgrad tiles need both contraction dims (in for fwd, out for dgrad) to be a # multiple of the FlyDSL block_k (64) and >= 128 (K_ITERS >= 2), so use 128-wide features. num_recv_tokens, in_features, out_features = 96, 128, 128 num_experts, max_hits = 6, 3 @@ -1025,7 +1024,7 @@ def test_grouped_weight_main_grad_matches_ungrouped_grad(monkeypatch): monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") torch.manual_seed(41) - # The v3 fwd/dgrad tiles need both contraction dims (in for fwd, out for dgrad) >= 128 and a + # The fwd/dgrad tiles need both contraction dims (in for fwd, out for dgrad) >= 128 and a # multiple of block_k (64); use 128-wide features. num_recv_tokens, in_features, out_features = 128, 128, 128 num_experts, max_hits = 8, 3 @@ -1102,7 +1101,7 @@ def test_grouped_weight_nonfused_grad_matches_ungrouped(monkeypatch): monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") torch.manual_seed(43) - # 128-wide features to satisfy the v3 fwd/dgrad tile (K>=128, multiple of block_k=64). + # 128-wide features to satisfy the fwd/dgrad tile (K>=128, multiple of block_k=64). num_recv_tokens, in_features, out_features = 128, 128, 128 num_experts, max_hits = 8, 3 m_splits = [num_recv_tokens // num_experts] * num_experts diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py index acd91faf7..4fd51f8d7 100644 --- a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. -"""Permute-free MoE data-gradient (dgrad) grouped-GEMM (v3): MegaMOE's fast bf16 NN GEMM. +"""Permute-free MoE data-gradient (dgrad) grouped-GEMM: MegaMOE's fast bf16 NN GEMM. Backward companion to ``pf_fwd``. dgrad contracts the incoming grad against the weight over the output-feature axis (NN layout):: @@ -219,7 +219,7 @@ def grouped_gemm_dgrad_k( group_res = create_buffer_resource(TILE_TO_GROUP, max_size=True) sorted_res = create_buffer_resource(SORTED, max_size=True) # Real BLOCK_M tile count as a scalar (host-known, capture-safe -- no per-call device - # tensor, which a HIP graph capture forbids). Matches the v2 int contract. + # tensor, which a HIP graph capture forbids). Host-known int scalar. real_tiles = NUM_TILE_BLOCKS real_grid = real_tiles * n_blocks @@ -248,7 +248,7 @@ def _emit(): block_m = first_pid_m + (pid_in_group % group_size_m) block_n = pid_in_group // group_size_m g_idx = buffer_load(group_res, block_m, vec_width=1, dtype=fx.T.i32()) - # PF padding blocks mark expert_ids=-1; skip the full Mega pipeline (v2 parity). + # PF padding blocks mark expert_ids=-1; skip the full Mega pipeline. if g_idx >= fx.Int32(0): gbase = g_idx * fx.Int32(N) * fx.Int32(Kout) sorted_row_base = block_m * fx.Int32(BLOCK_M) diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py index 16fb7d903..b54d218be 100644 --- a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py @@ -1,31 +1,28 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. -"""Permute-free MoE forward *gather* grouped-GEMM (v3): MegaMOE's fast bf16 GEMM + PF gather. +"""Permute-free MoE forward gather grouped-GEMM: MegaMOE's fast bf16 GEMM + route-list gather. This is a port of the MegaMOE grouped bf16 GEMM (32x32x16 MFMA, 8-wave / 512-thread workgroup, deep distance-2 3-buffer DMA ring, XCD swizzle with ``GROUP_M`` front-loading) into Transformer Engine FlyDSL, with a *single* change: the A operand is fetched through a per-row gather index instead of a contiguous pool. -Motivation. The v2 permute-free kernel (16x16x32, gather-in-K-loop) trails MegaMOE's dense -grouped GEMM by ~13%. Swapping v2's MFMA atom to 32x32x16 or deepening ``block_k`` both -regressed, so the lever is not the atom or LDS depth -- it is the gather structure vs. Mega's -pipeline. v3 tests the opposite direction: keep Mega's pipeline *verbatim* and only redirect the -A fetch through ``sorted_slot_ids``, so we pay Mega's throughput while keeping PF's memory model -(no pre-permutation, gather-on-demand). +The gather is folded into MegaMOE's pipeline rather than implemented as a separate +pre-permute: we keep Mega's DMA/LDS/MFMA path verbatim and only redirect the A fetch +through ``sorted_slot_ids``, preserving Mega's throughput while keeping the permute-free +memory model (no pre-permutation, gather-on-demand). -Contract (mirrors MegaMOE ``grouped_gemm_bf16_only`` + aiter PF routing metadata): +Contract (mirrors MegaMOE ``grouped_gemm_bf16_only`` + permute-free routing metadata): * ``A`` [num_recv, K] bf16 received-token activations, UNPERMUTED (gather source) * ``B`` [E, N, K] bf16 per-expert weights, NT (contiguous inner K) * ``C`` [em_max, N] bf16 compact expert-major output (block-padded, in place) * ``sorted_slot_ids`` [em_max] i32 received-token row per padded slot (sentinel = num_recv) * ``expert_ids`` [num_m_blocks] i32 expert id per ``BLOCK_M`` output block (padding tail: - ``-1``; those blocks early-exit like v2) + ``-1``; those blocks early-exit in-kernel) * ``num_tile_blocks`` [1] i32 real (non-padding) ``BLOCK_M`` block count (device) -Only the plain (non-activation) forward GEMM is ported here; the fused gated epilogue lives in -v2 and can be layered on later once the GEMM-throughput parity is confirmed. +Gated activation and route-prob apply live in standalone Triton helpers, not here. """ from __future__ import annotations @@ -235,7 +232,7 @@ def grouped_gemm_gather_k( group_res = create_buffer_resource(TILE_TO_GROUP, max_size=True) sorted_res = create_buffer_resource(SORTED, max_size=True) # Real (non-padding) BLOCK_M tile count as a scalar (host-known, capture-safe -- avoids a - # per-call device tensor that a HIP graph capture forbids). Matches the v2 int contract. + # per-call device tensor that a HIP graph capture forbids). Host-known int scalar. real_tiles = NUM_TILE_BLOCKS # XCD-swizzle over the REAL tile range only (front-loaded); swizzling the full padded # pool scatters real tiles -> ~2x slower. @@ -263,7 +260,7 @@ def _emit(): block_m = first_pid_m + (pid_in_group % group_size_m) block_n = pid_in_group // group_size_m g_idx = buffer_load(group_res, block_m, vec_width=1, dtype=fx.T.i32()) - # PF padding blocks mark expert_ids=-1; skip the full Mega pipeline (v2 parity). + # PF padding blocks mark expert_ids=-1; skip the full Mega pipeline. if g_idx >= fx.Int32(0): gbase = g_idx * fx.Int32(K) * c_n # Worst-case pool (cap*N > 2^31): rebase C per tile in int64, int32 in-resource diff --git a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py index e884d1580..a3ebcaff6 100644 --- a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py +++ b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py @@ -43,7 +43,7 @@ _WGRAD_CONTRACT_M = 32 -# Minimum v3 gather/dgrad block_m (128x256 MFMA floor). +# Minimum gather/dgrad align block_m (128x256 MFMA floor). _FLYDSL_MIN_BLOCK_M = 128 _FLYDSL_FWD_BLOCK_M = 256 @@ -114,8 +114,9 @@ def _pf_moe_fwd( num_recv_tokens: int, block_m: int, index_a_by_route_pos: bool = False, + dgrad: bool = False, ) -> None: - """Run the route-list gather-GEMM (forward or dgrad) via the FlyDSL v3 kernels.""" + """Run the route-list gather-GEMM (forward or dgrad) via the FlyDSL permute-free kernels.""" launch, supported = _get_flydsl_fwd() block_m = int(block_m) if not supported(A, B, block_m=block_m): @@ -132,44 +133,13 @@ def _pf_moe_fwd( num_recv_tokens=num_recv_tokens, block_m=block_m, index_a_by_route_pos=index_a_by_route_pos, + dgrad=dgrad, ) -def _fwd_align_block_size_m( - A: torch.Tensor, - B: torch.Tensor, - *, - num_tokens: int, -) -> int: - """Pick the forward align/kernel ``block_m``, favoring the backend that will run. - - ``block_m`` is baked into ``prepare_moe_align`` (it sets the row padding) and is *shared* by a - layer's FC1 fwd, FC1 dgrad and FC2 fwd (they reuse ``routing.block_size_m``), so it is chosen - once here. For ``num_tokens >= _FLYDSL_MIN_BLOCK_M`` (128), offer :data:`_FLYDSL_FWD_BLOCK_M` - (256, MegaMOE tile) alongside the v3 floor; smaller batches stay at 128 only. - """ - default_block_m = ( - _FLYDSL_FWD_BLOCK_M if num_tokens >= _FLYDSL_MIN_BLOCK_M else _FLYDSL_MIN_BLOCK_M - ) - candidates = {c for c in (_FLYDSL_MIN_BLOCK_M, default_block_m) if c <= default_block_m} - from .pf_fwd_wrapper import _v3_enabled - - # The v3 gather/dgrad tile requires block_m >= 128 (128x256 MFMA minimum), so on a - # small-token tier where the default would be < 128 we must still floor the align - # block_m to 128 rather than emit a sub-tile the kernel cannot run. - if _v3_enabled(): - candidates = {c for c in candidates if c >= _FLYDSL_MIN_BLOCK_M} or {_FLYDSL_MIN_BLOCK_M} - - # Prefer the in-tree FlyDSL block_m picker when available. - try: - from .pf_fwd_wrapper import flydsl_moe_fwd_pick_block_m - except Exception: # pylint: disable=broad-except - flydsl_moe_fwd_pick_block_m = None - if flydsl_moe_fwd_pick_block_m is not None: - picked = flydsl_moe_fwd_pick_block_m(A, B, candidates=tuple(candidates)) - if picked is not None: - return picked - return default_block_m +def _fwd_align_block_size_m(num_tokens: int) -> int: + """Forward align ``block_m`` (128 or 256); shared by FC1 fwd/dgrad and FC2 fwd.""" + return _FLYDSL_FWD_BLOCK_M if num_tokens >= _FLYDSL_MIN_BLOCK_M else _FLYDSL_MIN_BLOCK_M def _ensure_fwd_align( @@ -180,7 +150,7 @@ def _ensure_fwd_align( """Return ``routing.block_size_m``, building fwd align buffers when missing.""" if routing.sorted_slot_ids is not None and routing.block_size_m is not None: return int(routing.block_size_m) - block_m = _fwd_align_block_size_m(A, B, num_tokens=routing.num_recv_tokens) + block_m = _fwd_align_block_size_m(routing.num_recv_tokens) prepare_moe_align(routing, block_m) return block_m @@ -607,12 +577,10 @@ def permute_free_grouped_gemm_bf16_dgrad( f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." ) - # Contract over out_features via the [E, in, out] transposed *view* (stride relabel only). if weights_stacked.stride(-1) != 1: weights_stacked = weights_stacked.contiguous() - weights_t = weights_stacked.transpose(1, 2) - block_size_m = _ensure_fwd_align(routing, grad_output, weights_t) + block_size_m = _ensure_fwd_align(routing, grad_output, weights_stacked) # Two-stage dgrad reduction (contention-free): (1) plain compact store of each per-route # dX[route] = grad[route] @ W1[e] into an [T * min(topk, E), in] bf16 buffer (coalesced, no atomics), @@ -626,12 +594,13 @@ def permute_free_grouped_gemm_bf16_dgrad( ) _pf_moe_fwd( grad_output, - weights_t, + weights_stacked, compact, routing, num_recv_tokens=routing.num_recv_tokens, block_m=block_size_m, index_a_by_route_pos=True, + dgrad=True, ) return route_gather_combine( compact, @@ -893,33 +862,32 @@ def permute_free_grouped_gemm_bf16_fc2_dgrad( f"num_experts mismatch: weights have {num_experts}, routing has {routing.num_experts}." ) - # Contract over out_features via the [E, in, out] transposed view (stride relabel only). if weights_stacked.stride(-1) != 1: weights_stacked = weights_stacked.contiguous() - weights_t = weights_stacked.transpose(1, 2) - block_size_m = _ensure_fwd_align(routing, grad_output, weights_t) + block_size_m = _ensure_fwd_align(routing, grad_output, weights_stacked) # Padded route-order output; the gather-GEMM writes only the compact [0, num_routes) # range and never visits the tail (bounded by num_tokens_post_padded, sentinel-masked). # Left uninitialized (torch.empty) to skip the em_max*N zero-init; consumers read only the # compact range via the routing metadata, so the garbage tail is never observed. em_max = routing.sorted_slot_ids.shape[0] - dgrad = torch.empty( + dx = torch.empty( (em_max, in_features), dtype=torch.bfloat16, device=grad_output.device, ) _pf_moe_fwd( grad_output, - weights_t, - dgrad, + weights_stacked, + dx, routing, num_recv_tokens=grad_output.shape[0], block_m=block_size_m, index_a_by_route_pos=False, + dgrad=True, ) - return dgrad + return dx def permute_free_grouped_gemm_bf16_fc2_wgrad( diff --git a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py index 593431c5d..397369ada 100644 --- a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py +++ b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py @@ -7,99 +7,67 @@ metadata (``sorted_slot_ids`` / ``expert_ids`` / ``block_start``) can be reused verbatim. Writes the block-padded ``[em_max, WIDTH_N]`` slot output in place. -Forward gather-GEMM is FlyDSL-only (v3 MegaMOE-ported plain GEMM). Gated activation and +Forward gather-GEMM is FlyDSL-only (MegaMOE-ported plain GEMM). Gated activation and route-prob apply live in standalone Triton helpers (:mod:`pf_helper_kernels`), not in the GEMM kernel. """ from __future__ import annotations -import os -from typing import Optional - import torch __all__ = [ "flydsl_moe_fwd", "flydsl_moe_fwd_supported", - "flydsl_moe_fwd_pick_block_m", ] -_FILL_V = 8 -_WARP = 64 -_LDS_PAD = 8 -_LDS_LIMIT = 163840 # gfx950 per-workgroup LDS (160 KB) - -def _env_flag(name: str, default: bool) -> bool: - v = os.environ.get(name) - if v is None: - return default - return v.strip().lower() not in ("0", "false", "no", "off", "") - - -def _v3_enabled() -> bool: - """Whether to route plain GEMMs through the in-tree v3 (MegaMOE-ported) kernels.""" - return _env_flag("AITER_MOE_FLYDSL_V3", True) - +# MegaMOE hand-tuned bf16 grouped-GEMM tile geometry (offline-swept in primus-turbo's +# bench_mega_moe: BLOCK_N=256, GROUP_M=4 fwd / GROUP_M=8 FC1 NN dgrad, num_xcd=1). +# ``block_m`` comes from the route-list align (128 or 256); N-tile and grouping are fixed here. +_PF_BLOCK_N = 256 +_PF_GROUP_M = 4 +_PF_DGRAD_FC1_GROUP_M = 8 # bench_mega_moe grouped_gemm_combine L1-dgrad (NN) sweep +_PF_NUM_XCD = 1 -# MegaMOE's hand-tuned bf16 grouped-GEMM tile geometry (offline-swept in primus-turbo's -# bench_mega_moe: BLOCK_M/BLOCK_N=256, GROUP_M=4 fwd / GROUP_M=8 FC1 NN dgrad, num_xcd=1, -# nt_vmcnt=3). TE's FlyDSL fwd align already bumps FC1 / dgrad / FC2 to block_size_m=256; -# v3 pins the same Mega M-tile rather than trusting the wrapper arg. -_V3_BLOCK_M = 256 -_V3_BLOCK_N = 256 -_V3_GROUP_M = 4 -_V3_DGRAD_FC1_GROUP_M = 8 # bench_mega_moe grouped_gemm_combine L1-dgrad (NN) sweep -_V3_NUM_XCD = 1 - -def _run_v3_fwd( +def _run_gather_gemm( A, B, C, sorted_slot_ids, expert_ids, *, num_recv_tokens, block_m, transpose_b, index_a_by_route_pos, ): - """Dispatch a plain fwd/dgrad GEMM to the v3 (MegaMOE-ported) kernels. + """Dispatch fwd/dgrad to the permute-free FlyDSL gather-GEMM kernels. - Covers the three paths: FC1 gather (``index_a_by_route_pos=False``), FC2 route-read - (``index_a_by_route_pos=True``) and dgrad (``transpose_b``). v3 uses MegaMOE's fixed tile - geometry (32x32x16, ``BLOCK_N=256``, ``GROUP_M=4``, ``num_xcd=1``), so the wrapper's - ``block_n``/``block_k``/warp args are ignored here. + Covers FC1 gather (``index_a_by_route_pos=False``), FC2 route-read + (``index_a_by_route_pos=True``), and dgrad (``transpose_b``). Uses MegaMOE's fixed tile + geometry (32x32x16 MFMA, ``BLOCK_N=256``, ``GROUP_M=4``, ``num_xcd=1``). """ block_m = int(block_m) em_max = int(sorted_slot_ids.shape[0]) if em_max % block_m != 0: raise ValueError( - f"v3 expects em_max ({em_max}) divisible by block_m ({block_m}); " + f"permute-free gather-GEMM expects em_max ({em_max}) divisible by block_m ({block_m}); " "check routing align block_size_m." ) num_tile_blocks = em_max // block_m expert_ids_i32 = expert_ids.to(torch.int32) if expert_ids_i32.numel() != num_tile_blocks: raise ValueError( - f"v3 expects expert_ids length {num_tile_blocks} (em_max={em_max}, BLOCK_M={block_m}), " - f"got {expert_ids_i32.numel()}; routing align block_size_m must match block_m={block_m}" + f"permute-free gather-GEMM expects expert_ids length {num_tile_blocks} " + f"(em_max={em_max}, BLOCK_M={block_m}), got {expert_ids_i32.numel()}; " + f"routing align block_size_m must match block_m={block_m}" ) if transpose_b: - # dgrad: v3 is a native NN GEMM contracting the incoming grad against the weight over - # the output-feature axis. The wrapper's B is the transposed-weight *view* [E, out, in] - # (stride relabel); transpose(1,2) recovers the [E, N=out, K=in] forward weight v3 - # dgrad reads NN. FC1 dgrad (index_a_by_route_pos=True) is compact route-read (grad rows - # == dx rows); FC2 dgrad (index_a_by_route_pos=False) gathers the token-space grad + # dgrad: NN GEMM contracting the incoming grad against the forward weight [E, N, K] + # over the output-feature axis. FC1 dgrad (index_a_by_route_pos=True) is compact + # route-read; FC2 dgrad (index_a_by_route_pos=False) gathers token-space grad # [num_recv, N] into the compact route output [em_max, K]. from ..flydsl_kernels.permute_free_grouped_gemm.pf_dgrad import grouped_gemm_dgrad_bf16 - # TE passes a stride-relabelled transpose *view* (``stride(2) != 1``); undo the view to - # recover the forward ``[E, N, K]`` storage without copying (``transpose`` twice == id). - weight = B.transpose(1, 2) if B.stride(2) != 1 else B - # Real M-tile count for the dgrad grid: derive from the padded pool shape (host-known, - # graph-capture safe). Padding tail blocks (expert_ids=-1) early-exit in-kernel, same as - # the forward v3 path. Do not count active blocks on-device ((expert_ids>=0).sum().item()) - # -- that syncs the GPU and breaks HIP/CUDA graph capture. - dgrad_group_m = _V3_DGRAD_FC1_GROUP_M if index_a_by_route_pos else _V3_GROUP_M + dgrad_group_m = _PF_DGRAD_FC1_GROUP_M if index_a_by_route_pos else _PF_GROUP_M grouped_gemm_dgrad_bf16( - A, weight, C, expert_ids_i32, num_tile_blocks, sorted_slot_ids, + A, B, C, expert_ids_i32, num_tile_blocks, sorted_slot_ids, gather=not index_a_by_route_pos, - BLOCK_M=block_m, BLOCK_N=_V3_BLOCK_N, GROUP_M=dgrad_group_m, num_xcd=_V3_NUM_XCD, + BLOCK_M=block_m, BLOCK_N=_PF_BLOCK_N, GROUP_M=dgrad_group_m, num_xcd=_PF_NUM_XCD, ) return @@ -108,53 +76,10 @@ def _run_v3_fwd( grouped_gemm_gather_bf16( A, B, C, expert_ids_i32, num_tile_blocks, sorted_slot_ids, gather=not index_a_by_route_pos, - BLOCK_M=block_m, BLOCK_N=_V3_BLOCK_N, GROUP_M=_V3_GROUP_M, num_xcd=_V3_NUM_XCD, + BLOCK_M=block_m, BLOCK_N=_PF_BLOCK_N, GROUP_M=_PF_GROUP_M, num_xcd=_PF_NUM_XCD, ) -def _mfma_dim(transpose_b: bool) -> int: - """MFMA output-tile edge: forward GEMM uses the 32x32x16 atom, dgrad the 16x16x32 atom. - - Kept in sync with the kernel's ``MOE_FWD_MFMA32`` escape hatch so the autotuner sweeps the - tile-divisibility that the compiled atom actually requires. - """ - if transpose_b: - return 16 - return 32 if _env_flag("MOE_FWD_MFMA32", False) else 16 - - -def _warp_valid(block_m, block_n, block_k, wm, wn, transpose_b=False): - n_threads = wm * wn * _WARP - wmma = _mfma_dim(transpose_b) - if block_m % (wm * wmma) or block_n % (wn * wmma): - return False - if (block_m * block_k) % (n_threads * _FILL_V): - return False - if (block_n * block_k) % (n_threads * _FILL_V): - return False - return True - - -def _fwd_buffering(): - """(num_buffers, lds_pad) for the production DMA+swizzle fill path. - - Both default on (opt out with ``MOE_FWD_DMA=0`` / ``MOE_FWD_SWZ=0``). The DMA path runs - a distance-2, 3-buffer ring; the register fallback keeps 2-buffer ping/pong. - """ - use_dma = _env_flag("MOE_FWD_DMA", True) - swz = _env_flag("MOE_FWD_SWZ", True) - pad = 0 if (use_dma or swz) else _LDS_PAD - return (3 if use_dma else 2), pad - - -def _lds_bytes(block_m, block_n, block_k, transpose_b=False): - nbuf, pad = _fwd_buffering() - a_tile = block_m * (block_k + pad) - # dgrad stages B as [k, n] (row stride = block_n+pad); fwd as [n, k] (block_k+pad). - b_tile = block_k * (block_n + pad) if transpose_b else block_n * (block_k + pad) - return (a_tile + b_tile) * nbuf * 2 - - def flydsl_moe_fwd_supported( A: torch.Tensor, B: torch.Tensor, @@ -166,7 +91,7 @@ def flydsl_moe_fwd_supported( return False if A.stride(1) != 1: return False - if B.stride(2) != 1 and B.stride(1) != 1: + if B.stride(2) != 1: return False # em_max % block_m is only known at launch (from C), so it is not prechecked here. return True @@ -182,106 +107,22 @@ def flydsl_moe_fwd( num_recv_tokens: int, block_m: int, index_a_by_route_pos: bool = False, + dgrad: bool = False, ) -> None: """Route-list gather-GEMM forward, writing the compact ``C[em_max, WIDTH_N]`` in place. ``A`` is ``[*, K]`` (received-token acts, gathered by ``sorted_slot_ids`` when ``index_a_by_route_pos=False``, else read at the compact route row). ``B`` is - ``[num_experts, N_OUT, K]`` (contiguous inner ``K``). Gated activation and route-prob - apply are **not** fused here; use the standalone helpers in :mod:`pf_helper_kernels`. - - The v3 (MegaMOE-ported) kernels use a fixed tile geometry, so only ``block_m`` is a tunable - knob (chosen once by the align via :func:`flydsl_moe_fwd_pick_block_m`). + ``[num_experts, N_OUT, K]`` (contiguous inner ``K``). Set ``dgrad=True`` for the + data-gradient path (same ``B`` layout). Gated activation and route-prob apply are **not** + fused here; use the standalone helpers in :mod:`pf_helper_kernels`. """ assert A.dtype == B.dtype == C.dtype == torch.bfloat16 assert A.stride(1) == 1, "A must be contiguous along the contraction (K)" - # B is [E, N_OUT, K]: contiguous along K (fwd) or along N (dgrad transposed-weight view). - transpose_b = B.stride(2) != 1 - if transpose_b: - assert B.stride(1) == 1, "transposed B must be contiguous along N (dgrad view)" + assert B.stride(2) == 1, "B must be [E, N, K] with contiguous inner K" - # Plain GEMM: in-tree v3 kernels (pf_fwd / pf_dgrad). - _run_v3_fwd( + _run_gather_gemm( A, B, C, sorted_slot_ids, expert_ids, num_recv_tokens=num_recv_tokens, block_m=block_m, - transpose_b=transpose_b, index_a_by_route_pos=index_a_by_route_pos, + transpose_b=dgrad, index_a_by_route_pos=index_a_by_route_pos, ) - - -# Tile/warp configs the autotuner sweeps: (block_n, block_k, warps_m, warps_n). Fill path -# (DMA+swizzle, 3-buffer) is fixed at the env defaults above -- only tile geometry is tuned. -# -# Two shapes matter for the wide-M (block_m=256) MoE cases, where every block_k=128 and -# 256x64 entry below is filtered out by the LDS limit, leaving only 128x64 candidates: -# * a square-ish w4x4 warp grid, which spreads the cooperative A-fill (2 fills) / B-fill -# (1 fill) DMA and the LDS fragment reads more evenly than the tall w8x2 layout; -# * block_n=256 with block_k=32, which fits LDS at 3 buffers and raises arithmetic -# intensity to block_m*block_n/(2*(block_m+block_n)) = 64 MAC/byte vs 42.7 at block_n=128 -# (the ratio is independent of block_k, so widening N is what pays). -# Measured on FC1 no-act (qwen235b, block_m=256), idle machine, alias scopes on: -# 256x32 w4x4 ~1628us, 128x64 w4x4 ~1694us, 128x64 w8x2 ~1726us. -# -# An exhaustive sweep of the valid tile space found nothing better, so the list below is not -# missing a winner. Two directions are dead ends and are deliberately -# absent: block_n>=384 costs 4-22x (per-wave accumulator spill plus 120-144KB LDS pinning -# occupancy to 1 workgroup), and trading block_m down to reach block_k=128 -- 4x fewer -# barriers, which the ATT trace makes look attractive -- costs 59% (2581us at block_m=128 -# bk=128) because arithmetic intensity falls faster than barrier count. Note block_k>=128 does -# not fit LDS at all once NUM_BUF>=3, which the kernel enforces. -# -# Do not add 256x384x32 w1x4 or 512x32 w1x4 / w2x2: they abort the backend outright -# ("Bad machine code: Virtual register defs don't dominate all uses"), which would take the -# autotuner down with them rather than being skipped. Pre-existing, unrelated to alias scopes. -_FWD_TUNE_CONFIGS = [ - (64, 64, 2, 2), - (128, 64, 2, 2), - (128, 64, 4, 2), - (128, 64, 2, 4), - (128, 64, 4, 4), - (128, 64, 8, 2), - (256, 32, 4, 4), - (256, 32, 2, 4), - (256, 32, 4, 2), - (256, 64, 4, 2), - (128, 128, 2, 2), - (128, 128, 4, 2), - (64, 128, 2, 2), - (256, 64, 2, 4), -] - -def _valid_config(block_m, bn, bk, wm, wn, K, transpose_b=False): - if K % bk != 0 or bn % _mfma_dim(transpose_b) != 0: - return False - if not _warp_valid(block_m, bn, bk, wm, wn, transpose_b): - return False - return _lds_bytes(block_m, bn, bk, transpose_b) <= _LDS_LIMIT - - -def flydsl_moe_fwd_pick_block_m( - A: torch.Tensor, - B: torch.Tensor, - *, - candidates=(256, 128), -) -> Optional[int]: - """Largest ``block_m`` in ``candidates`` the FlyDSL fwd can actually run for these operands, - or ``None`` if the operands are unsupported at every candidate. - - "Can run" == bf16 operands contiguous along the contraction AND at least one candidate tile - in :data:`_FWD_TUNE_CONFIGS` fits the per-workgroup LDS budget. Callers should include their - default among ``candidates`` (the picker only walks high->low over what is passed), so a - small-token workload is never padded up beyond what the caller offered. - """ - if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: - return None - if A.stride(1) != 1: # A must be contiguous along the contraction (K) - return None - if B.stride(2) != 1 and B.stride(1) != 1: # B contiguous along K (fwd) or N (dgrad view) - return None - K = int(B.shape[2]) - for block_m in sorted({int(c) for c in candidates}, reverse=True): - if any( - _valid_config(block_m, bn, bk, wm, wn, K) - for (bn, bk, wm, wn) in _FWD_TUNE_CONFIGS - ): - return block_m - return None From 28451fa6c43f00ec7e7214eeb93d7b2bce919e8c Mon Sep 17 00:00:00 2001 From: sudhu2k Date: Thu, 6 Aug 2026 21:37:57 +0000 Subject: [PATCH 43/43] Standardize permute-free MoE layout terminology and drop FlyDSL forward support check Replaces the ambiguous "compact" terminology across the permute-free MoE grouped GEMM code with explicit layout names: **dense route-ordered** ``[num_routes, F]``, **block-padded route-ordered** ``[em_max, F]``, and **token-ordered** ``[num_recv_tokens, F]``. Renames the test helper ``_compact_route_order`` to ``_dense_route_order`` and updates docstrings, comments, and error messages in the FlyDSL kernels, wrappers, ``GroupedLinear``, and routing metadata to use the new vocabulary. Removes the ``flydsl_moe_fwd_supported`` helper and the runtime capability check in ``_pf_moe_fwd``; the wrapper now calls the FlyDSL launcher directly. Also renames the intermediate dgrad buffer from ``compact`` to ``route_buf`` for clarity. --- .../benchmark_perm_free_grouped_gemm.py | 5 +- .../pytorch/test_perm_free_grouped_linear.py | 32 +++--- .../permute_free_grouped_gemm/pf_dgrad.py | 21 ++-- .../permute_free_grouped_gemm/pf_fwd.py | 8 +- .../permute_free_grouped_gemm/pf_wgrad.py | 2 +- .../pytorch/module/grouped_linear.py | 10 +- transformer_engine/pytorch/moe/moe_routing.py | 14 +-- .../pytorch/moe/permute_free_grouped_gemm.py | 105 +++++++++--------- .../pytorch/moe/pf_fwd_wrapper.py | 42 ++----- .../pytorch/moe/pf_helper_kernels.py | 14 +-- 10 files changed, 115 insertions(+), 138 deletions(-) diff --git a/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py index ea9c61721..b26c109c8 100644 --- a/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py +++ b/benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py @@ -267,11 +267,12 @@ def _build_train_phase_fns( # Run the permute-free forward once (outside timing) to finalize the block-padded align on # ``routing`` (the fwd enforces a v3 block-size floor, which fixes em_max) and to learn the # padded slot extent. The dgrad/wgrad now consume the block-padded ``[em_max, out_features]`` - # slot gradient -- i.e. the gradient of this fwd output -- not a compact ``[num_routes]`` one. + # slot gradient -- i.e. the gradient of this fwd output -- not a dense route-ordered + # ``[num_routes]`` one. em_max = int(permute_free_grouped_gemm_bf16(hidden, weights, routing).shape[0]) # Upstream gradients. ``grad_out_pf`` is the block-padded slot grad for the permute-free - # dgrad/wgrad; the traditional path keeps its compact expert-sorted ``[total_m, n]`` grad. + # dgrad/wgrad; the traditional path keeps its expert-sorted ``[total_m, n]`` grad. # Allocated once, outside the timed region. grad_out_pf = torch.randn(em_max, n, dtype=dtype, device=device) grad_out_perm = torch.randn(total_m, n, dtype=dtype, device=device) diff --git a/tests/pytorch/test_perm_free_grouped_linear.py b/tests/pytorch/test_perm_free_grouped_linear.py index c5e623fc4..7bfa0a679 100644 --- a/tests/pytorch/test_perm_free_grouped_linear.py +++ b/tests/pytorch/test_perm_free_grouped_linear.py @@ -62,8 +62,8 @@ def _random_routing_map(num_recv_tokens, num_experts, max_hits, device, seed): return routing_map -def _compact_route_order(routing_map): - """Expert-sorted (token, expert) route lists, matching ``moe_align_route_list``.""" +def _dense_route_order(routing_map): + """Expert-sorted (token, expert) dense route list, matching ``moe_align_route_list``.""" tok, exp = routing_map.nonzero(as_tuple=True) order = torch.argsort(exp, stable=True) return tok[order].to(torch.int64), exp[order].to(torch.int64) @@ -86,7 +86,7 @@ def test_route_list_fwd(): out_full = permute_free_grouped_gemm_bf16(hidden, weights, routing) - route_to_token, route_expert = _compact_route_order(routing_map) + route_to_token, route_expert = _dense_route_order(routing_map) num_routes = route_to_token.numel() # Block-padded canonical layout: the output is [em_max, out] in expert-sorted, block-padded # slot order. Padded slot s holds hidden[sorted_slot_ids[s]] @ W[expert(s)]^T for the valid @@ -491,12 +491,12 @@ def test_fc1_fc2_gated_pipeline(monkeypatch): assert getattr(fc1, f"weight{i}").grad is not None assert getattr(fc2, f"weight{i}").grad is not None assert getattr(fc2, f"weight{i}").grad.abs().sum() > 0 - # Sanity: compact route head got a non-zero FC1 output grad (FC2 act-bwd -> FC1). + # Sanity: dense route head got a non-zero FC1 output grad (FC2 act-bwd -> FC1). assert preact.grad[:num_routes].abs().sum() > 0 # Numerical check on probs.grad (locks the gated-act backward's per-slot (token, expert) # mapping): with loss = out.sum(), dL/dprob[t, e] = routing_map[t, e] * . A wrong expert-per-slot map (e.g. compact route index misread as a + # sum_h w2[e, h, :]>. A wrong expert-per-slot map (e.g. dense route index misread as a # block-padded slot) silently corrupts this even though probs.grad stays non-zero. with torch.no_grad(): act_noprob = torch.nn.functional.silu(gate) * up # [T, E, F] (CPU float) @@ -507,7 +507,7 @@ def test_fc1_fc2_gated_pipeline(monkeypatch): def test_route_list_fwd_dgrad_wgrad_consistency(): - """fwd + dgrad + wgrad against a single autograd reference in compact route space.""" + """fwd + dgrad + wgrad against a single autograd reference in dense route-ordered space.""" torch.manual_seed(23) # The gather/dgrad tiles need both contraction dims (in for fwd, out for dgrad) to be a # multiple of the FlyDSL block_k (64) and >= 128 (K_ITERS >= 2), so use 128-wide features. @@ -516,9 +516,9 @@ def test_route_list_fwd_dgrad_wgrad_consistency(): routing_map = _random_routing_map(num_recv_tokens, num_experts, max_hits, "cuda", seed=5) routing = MoERoutingMetadata(routing_map=routing_map, num_experts=num_experts) - route_to_token, route_expert = _compact_route_order(routing_map) + route_to_token, route_expert = _dense_route_order(routing_map) num_routes = route_to_token.numel() - # Prepare the fwd/dgrad align (block-padded slot layout) up front so we can map the compact + # Prepare the fwd/dgrad align (block-padded slot layout) up front so we can map the dense # per-route grad into padded-slot order for the block-padded dgrad. routing = prepare_moe_align(routing, _FLYDSL_FWD_BLOCK_M) em_max = int(routing.sorted_slot_ids.shape[0]) @@ -531,20 +531,20 @@ def test_route_list_fwd_dgrad_wgrad_consistency(): ) # Forward output is block-padded [em_max]; the valid padded slots are expert-ascending, the - # same order as the compact route list, so out_full[valid] lines up with the compact ref. + # same order as the dense route list, so out_full[valid] lines up with the dense ref. out_full = permute_free_grouped_gemm_bf16(hidden, weights, routing) out = out_full[valid] # One per-route gradient in the block-padded [em_max] slot layout that both the dgrad and the # wgrad now route-read (padded slot = block_start[e]*block_size_m + within-rank). The valid - # slots are expert-ascending, matching the compact autograd reference order. ``grad`` keeps a - # compact [num_routes] copy only to seed the compact reference backward. + # slots are expert-ascending, matching the dense autograd reference order. ``grad`` keeps a + # dense route-ordered [num_routes] copy only to seed the dense reference backward. grad = torch.randn(num_routes, out_features, device="cuda", dtype=torch.bfloat16) grad_bp = torch.zeros(em_max, out_features, device="cuda", dtype=torch.bfloat16) grad_bp[valid] = grad dA = permute_free_grouped_gemm_bf16_dgrad(grad_bp, weights, routing) dW = permute_free_grouped_gemm_bf16_wgrad(hidden, grad_bp, weights.shape, routing) - # Autograd reference in compact route space. + # Autograd reference in dense route-ordered space. ref_hidden = hidden.float().clone().requires_grad_(True) ref_w = weights.float().clone().requires_grad_(True) ref_out = torch.einsum( @@ -829,8 +829,8 @@ def test_slot_expert_ids_cache(): # ``_expert_per_route`` returns the *cached* tensor (no redundant recompute). assert _expert_per_route(routing, em_max) is routing.slot_expert_ids - # Valid slots carry the expert-sorted compact route experts. - _, route_expert = _compact_route_order(routing_map) + # Valid slots carry the expert-sorted dense route experts. + _, route_expert = _dense_route_order(routing_map) valid = routing.sorted_slot_ids < num_recv_tokens assert torch.equal(routing.slot_expert_ids[valid].to(torch.int64), route_expert) @@ -1041,7 +1041,7 @@ def test_grouped_weight_main_grad_matches_ungrouped_grad(monkeypatch): ) # Shared initial weights for both modules (so the kernel dW is identical). W = torch.randn(num_experts, out_features, in_features, device="cuda", dtype=torch.bfloat16) - # Shared upstream gradient over the compact route range. + # Shared upstream gradient over the dense route-ordered range. g = torch.randn(num_routes, out_features, device="cuda", dtype=torch.bfloat16) # --- ungrouped: separate per-expert params, wgrad via autograd .grad --- @@ -1054,7 +1054,7 @@ def test_grouped_weight_main_grad_matches_ungrouped_grad(monkeypatch): routing_a = PermuteFreeMetadata(routing_map=routing_map, num_experts=num_experts) out_a = mod_a(inp, m_splits, permute_free_metadata=routing_a) # Block-padded canonical: the module output is [em_max, out] in padded slot order; the valid - # (expert-ascending) slots line up with the compact upstream grad g. Slicing [:num_routes] + # (expert-ascending) slots line up with the dense upstream grad g. Slicing [:num_routes] # would cut across the padding gaps, so gather the valid slots instead. valid_a = routing_a.sorted_slot_ids < num_recv_tokens (out_a[valid_a].float() * g.float()).sum().backward() diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py index 4fd51f8d7..c46f6f271 100644 --- a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_dgrad.py @@ -10,11 +10,12 @@ There are two flavours, matching the TE permute-free contract (``index_a_by_route_pos``): -* **route-read** (FC1 dgrad, ``gather=False``): ``grad`` is already compact route order - ``[em_max, N]``; row ``s`` is read directly. ``grad`` and ``dXrow`` share the row count. -* **gather** (FC2 dgrad, ``gather=True``): ``grad`` is token-space ``[num_recv, N]`` and each +* **route-read** (FC1 dgrad, ``gather=False``): ``grad`` is **dense route-ordered** + ``[num_routes, N]`` (or block-padded with only the dense head populated); row ``s`` is read + directly. ``dx`` is **block-padded route-ordered** ``[em_max, K]``. +* **gather** (FC2 dgrad, ``gather=True``): ``grad`` is **token-ordered** ``[num_recv, N]`` and each route slot ``s`` gathers ``grad[SORTED[s]]`` (sentinel ``SORTED[s] == num_recv`` -> 0), writing - the compact route-order ``dXrow[em_max, K]``. Mirrors the forward NT gather, one row map + block-padded route-ordered ``dXrow[em_max, K]``. Mirrors the forward NT gather, one row map redirecting the two LDS A half-tiles, only on the NN tile. The weight is bit-identical to the forward ``[E, N, K]`` (forward reads it NT, dgrad NN). @@ -22,7 +23,7 @@ Contract: * ``grad_y`` [rows, N] bf16 incoming grad (rows = em_max route-read / num_recv gather) * ``weight`` [E, N, K] bf16 per-expert weights (shared with forward) - * ``dx`` [em_max, K] bf16 compact expert-major input grad per slot (in place) + * ``dx`` [em_max, K] bf16 block-padded route-ordered input grad per slot (in place) * ``expert_ids`` [num_m_blocks] i32 expert id per BLOCK_M slot block * ``num_tile_blocks`` [1] i32 real (non-padding) BLOCK_M block count (device) * ``sorted_slot_ids`` [em_max] i32 gather index (gather=True); unused for route-read @@ -170,7 +171,7 @@ def _dgrad_routeread_body( GY_flat, GY_tile, WEIGHT, DX_tile, lds, sorted_res, sorted_row_base, block_n, gbase, *, N, Kout, BLOCK_M, BLOCK_N, out_fp16, nt_vmcnt, ): - """FC1 dgrad tile: read the compact route grad directly (plain Mega NN tile).""" + """FC1 dgrad tile: read the dense route-ordered grad directly (plain Mega NN tile).""" gemm_bf16_nn_tile( GY_tile, WEIGHT, DX_tile, fx.Int32(BLOCK_M), fx.Int32(Kout), lds, fx.Int32(0), block_n, K=N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, out_fp16=out_fp16, nt_vmcnt=nt_vmcnt, @@ -194,9 +195,9 @@ def compile_grouped_gemm_dgrad_bf16( ): """Compile (cached) the grouped BF16 NN dgrad launcher for one ``(N, Kout, tile)`` combo. - ``gather=False`` (FC1 dgrad): ``grad`` is compact route order, read per tile (plain Mega NN - tile). ``gather=True`` (FC2 dgrad): ``grad`` is token-space, gathered via ``SORTED`` into the - compact route output (NN gather tile). Both rebase the C tile in i64 to survive worst-case + ``gather=False`` (FC1 dgrad): ``grad`` is dense route-ordered, read per tile (plain Mega NN + tile). ``gather=True`` (FC2 dgrad): ``grad`` is token-ordered, gathered via ``SORTED`` into the + block-padded route-ordered output (NN gather tile). Both rebase the C tile in i64 to survive worst-case pools; the grid front-loads via XCD swizzle over the real tile range. """ SharedStorage = _make_shared_storage(BLOCK_M, BLOCK_N) @@ -294,7 +295,7 @@ def launch(GRAD_Y, WEIGHT, DX, TILE_TO_GROUP, NUM_TILE_BLOCKS: fx.Int32, SORTED, def grouped_gemm_dgrad_bf16( grad_y, # [rows, N] bf16 incoming grad (rows = em_max route-read / num_recv gather) weight, # [E, N, K] bf16 per-expert weights (shared with forward) - dx, # [em_max, K] bf16 compact expert-major input grad per slot (in place) + dx, # [em_max, K] bf16 block-padded route-ordered input grad per slot (in place) expert_ids, # [num_m_blocks] i32 num_tile_blocks, # int real BLOCK_M block count (host scalar; capture-safe) sorted_slot_ids=None, # [em_max] i32 gather index (required when gather=True) diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py index b54d218be..211a8978c 100644 --- a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py @@ -16,7 +16,7 @@ Contract (mirrors MegaMOE ``grouped_gemm_bf16_only`` + permute-free routing metadata): * ``A`` [num_recv, K] bf16 received-token activations, UNPERMUTED (gather source) * ``B`` [E, N, K] bf16 per-expert weights, NT (contiguous inner K) - * ``C`` [em_max, N] bf16 compact expert-major output (block-padded, in place) + * ``C`` [em_max, N] bf16 block-padded route-ordered output (expert-major, in place) * ``sorted_slot_ids`` [em_max] i32 received-token row per padded slot (sentinel = num_recv) * ``expert_ids`` [num_m_blocks] i32 expert id per ``BLOCK_M`` output block (padding tail: ``-1``; those blocks early-exit in-kernel) @@ -315,7 +315,7 @@ def launch(A, B, C, TILE_TO_GROUP, NUM_TILE_BLOCKS: fx.Int32, SORTED, A_ELEMS: f def grouped_gemm_gather_bf16( A, # [num_recv, K] bf16 received-token activations (UNPERMUTED gather source) weight, # [E, N, K] bf16 per-expert B (NT) - output, # [em_max, N] bf16 compact expert-major C (in place) + output, # [em_max, N] bf16 block-padded route-ordered C (in place) expert_ids, # [num_m_blocks] i32 expert per BLOCK_M output block num_tile_blocks, # int real BLOCK_M block count (host scalar; capture-safe) sorted_slot_ids, # [em_max] i32 received-token row per padded slot (sentinel = num_recv) @@ -330,8 +330,8 @@ def grouped_gemm_gather_bf16( gather=True, ): """Host entry: grouped bf16 NT GEMM. With ``gather=True`` (FC1) ``C[pos] = A[SORTED[pos]] - @ B[expert]^T``; with ``gather=False`` (FC2 route-read) ``A`` is the compact ``[em_max, K]`` - pool read at the route row (``sorted_slot_ids`` unused, may be a dummy). + @ B[expert]^T``; with ``gather=False`` (FC2 route-read) ``A`` is **block-padded route-ordered** + ``[em_max, K]`` read at the route slot (``sorted_slot_ids`` unused, may be a dummy). ``output`` is written in place over its full padded ``[em_max, N]`` extent (padding rows carry dead values, ignored by downstream stages keyed on the same routing metadata). ``c_m`` is the diff --git a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py index 6c81fd6d9..1275f060f 100644 --- a/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py +++ b/transformer_engine/pytorch/flydsl_kernels/permute_free_grouped_gemm/pf_wgrad.py @@ -32,7 +32,7 @@ columns out of the shared tile. ``warps_n = warps_k = 1`` reduces to the single-warp v2. Fixed configuration (no runtime toggles): - - bf16 inputs, bf16 output (FC1: compact grad + token-gathered ``x``) + - bf16 inputs, bf16 output (FC1: block-padded route-ordered grad + token-gathered ``x``) - optional ``accumulate``: overwrite (default) or read-modify-write into ``dW`` - DMA + XOR chunk swizzle fill, 3-stage LDS pipeline, LLVM DMA alias scopes """ diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index ab2947105..d3b7df98c 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -650,7 +650,7 @@ def forward( biases = [cast_if_needed(bias, bias_dtype) for bias in biases] if use_bias else biases # Initialize output tensor. The permute-free path allocates its own worst-case padded # [T * min(topk, E), out_features] output inside permute_free_grouped_gemm_bf16 (valid rows are - # the compact route range [0, num_routes); the tail is inert zero padding). + # the dense route range [0, num_routes); the tail is inert zero padding). if not use_perm_free_grouped_gemm: out = torch.empty( [sum(m_splits), weights_fp8[0].size(0)], @@ -855,7 +855,7 @@ def forward( ctx.reduce_and_update_bwd_fp8_tensors = False # [*, in_features] -> [*, out_features], or worst-case padded [T * min(topk, E), out_features] - # (permute-free route-list path; valid rows are the compact range [0, num_routes)). + # (permute-free route-list path; valid rows are the dense route range [0, num_routes)). if use_perm_free_grouped_gemm: return out, new_workspaces return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @@ -1549,7 +1549,7 @@ class GroupedLinear(TransformerEngineBaseModule): ``[num_recv_tokens, num_local_experts]`` + a ``route_space`` direction) instead of permuting activations before this module. The caller must skip ``moe_permute``. FC1 (``route_space=False``) takes ``[num_recv_tokens, in_features]`` and produces the - worst-case padded ``[T * min(topk, E), out_features]`` route buffer (valid rows are the compact + worst-case padded ``[T * min(topk, E), out_features]`` route buffer (valid rows are the dense route range ``[0, num_routes)``; the tail is inert zero padding); FC2 (``route_space=True``) takes the route-ordered ``[T * min(topk, E), in_features]`` and fuses the scatter back to token order, returning ``[num_recv_tokens, out_features]``. Requires @@ -2014,8 +2014,8 @@ def forward( direction). When set with ``NVTE_PERMUTE_FREE_GROUPED_GEMM=1``, run the route-list GEMM on unpermuted bf16 activations. ``route_space= False`` (FC1) gathers per expert into the worst-case padded - ``[T * min(topk, E), out_features]`` route buffer (valid rows are the compact - range ``[0, num_routes)``; the tail is inert zero padding); + ``[T * min(topk, E), out_features]`` route buffer (valid rows are the dense + route range ``[0, num_routes)``; the tail is inert zero padding); ``route_space=True`` (FC2) reads route-ordered input, applies the standalone gated activation (when ``activation`` is set), runs a plain GEMM, and scatter-combines back to token order, returning diff --git a/transformer_engine/pytorch/moe/moe_routing.py b/transformer_engine/pytorch/moe/moe_routing.py index 1f66c381e..4f06127c9 100644 --- a/transformer_engine/pytorch/moe/moe_routing.py +++ b/transformer_engine/pytorch/moe/moe_routing.py @@ -131,15 +131,15 @@ class PermuteFreeMetadata(MoERoutingMetadata): Extends :class:`MoERoutingMetadata` - - ``route_space=False`` (FC1): the input lives in **received-token order** + - ``route_space=False`` (FC1): the input lives in **token-ordered** ``[num_recv_tokens, in]``. The forward *gathers* per expert (``index_a_by_route_pos= - False``) into the compact/padded ``[T * min(topk, E), out]`` route buffer; the dgrad combines - the input gradient back to token rows (contention-free gather-combine). - - ``route_space=True`` (FC2): the input is already in **route order** - ``[T * min(topk, E), in]`` (FC1's output). The forward reads by route position + False``) into the **block-padded route-ordered** ``[em_max, out]`` buffer; the dgrad + combines the input gradient back to token rows (contention-free gather-combine). + - ``route_space=True`` (FC2): the input is already **block-padded route-ordered** + ``[em_max, in]`` (FC1's output). The forward reads by route slot (``index_a_by_route_pos=True``) and combines each token's routes back to - ``[num_recv_tokens, out]`` (contention-free gather-combine); the dgrad gathers - the token-space grad back into the compact route buffer. + **token-ordered** ``[num_recv_tokens, out]`` (contention-free gather-combine); the dgrad + gathers the token-ordered grad back into the block-padded route-ordered buffer. The align buffers are identical for both directions, so a single built metadata can be reused for FC1 and FC2 (e.g. via ``dataclasses.replace(meta, route_space=True)``), diff --git a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py index a3ebcaff6..275d8dd75 100644 --- a/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py +++ b/transformer_engine/pytorch/moe/permute_free_grouped_gemm.py @@ -4,11 +4,16 @@ """Permute-free route-list grouped GEMM for MoE (bf16). -- TE builds one expert-sorted ``sorted_slot_ids`` (received-token row per block-padded slot) - plus the per-expert ``block_start`` (block units), then runs the gather-GEMM. -- FC1 fwd output is worst-case padded ``[T * min(topk, E), out_features]`` in expert order (valid rows - are the compact route range ``[0, num_routes)``, tail is inert zero padding); dgrad returns - ``dA = [num_recv_tokens, in_features]`` (scatter-add of the per-route gradients). +Layout vocabulary: + * **token-ordered** ``[num_recv_tokens, F]`` — one row per received token. + * **dense route-ordered** ``[num_routes, F]`` — one row per ``(token, expert)`` route, no padding. + * **block-padded route-ordered** ``[em_max, F]`` — expert-sorted routes with per-expert + ``block_m`` padding (sync-free over-allocated upper bound). + +TE builds expert-sorted ``sorted_slot_ids`` (token row per route slot) plus per-expert +``block_start``, then runs gather-in-GEMM. FC1 fwd output is block-padded route-ordered +(valid route slots are ``[0, num_routes)``; tail is inert padding). FC1 dgrad returns +token-ordered ``dA = [num_recv_tokens, in_features]`` via gather-combine. """ from __future__ import annotations @@ -49,10 +54,10 @@ def _get_flydsl_fwd(): - """Return ``(flydsl_moe_fwd, flydsl_moe_fwd_supported)``.""" - from .pf_fwd_wrapper import flydsl_moe_fwd, flydsl_moe_fwd_supported + """Return the permute-free FlyDSL gather-GEMM launcher.""" + from .pf_fwd_wrapper import flydsl_moe_fwd - return flydsl_moe_fwd, flydsl_moe_fwd_supported + return flydsl_moe_fwd def _expand_expert_ids_per_slot( @@ -117,13 +122,8 @@ def _pf_moe_fwd( dgrad: bool = False, ) -> None: """Run the route-list gather-GEMM (forward or dgrad) via the FlyDSL permute-free kernels.""" - launch, supported = _get_flydsl_fwd() + launch = _get_flydsl_fwd() block_m = int(block_m) - if not supported(A, B, block_m=block_m): - raise RuntimeError( - "FlyDSL grouped GEMM does not support these operands " - f"(block_m={block_m}, A={tuple(A.shape)}, B={tuple(B.shape)})." - ) launch( A, B, @@ -266,7 +266,7 @@ def prepare_moe_align(metadata: MoERoutingMetadata, block_m: int) -> MoERoutingM metadata.num_tokens_post_padded = num_tokens_post_padded # [1] int32 device scalar: padded extent metadata.block_start = block_start # [E] int32: first block index of each expert (block units) metadata.block_size_m = block_m # int: BLOCK_SIZE_M the layout is padded to - metadata.token_routes = token_routes # [T, min(topk, E)] int32: token -> its compact route ids + metadata.token_routes = token_routes # [T, min(topk, E)] int32: token -> padded route slot indices metadata.token_route_count = token_route_count # [T] int32: number of routes per token return metadata @@ -415,7 +415,7 @@ def permute_free_grouped_gemm_bf16( # Worst-case (sync-free) allocation: size the output to the block-padded upper bound # em_max = sorted_slot_ids.shape[0], which is derived purely from shapes (num_recv_tokens, # num_experts, block_size) and is always >= num_routes. This avoids the device->host sync - # (.item()) that a compact [num_routes, N] allocation would require. + # (.item()) that a dense route-ordered [num_routes, N] allocation would require. em_max = routing.sorted_slot_ids.shape[0] output = torch.empty( (em_max, out_features), @@ -475,7 +475,7 @@ def permute_free_gated_act_bwd( # Block-padded canonical layout: grad_out (FC2 dgrad) and preact both live in the [em_max] # padded slot order, and the kernel indexes grad_out/preact/dpre by the same row as # token/expert -- so we must feed the padded slot arrays (sorted_slot_ids + slot expert), - # not the compact route arrays, or the padded valid slots beyond routes_max never get a dpre + # not the dense route arrays, or the padded valid slots beyond routes_max never get a dpre # row (and each route pairs a correct token with the wrong padded grad row). em_max = int(routing.sorted_slot_ids.shape[0]) token = routing.sorted_slot_ids.to(torch.int32) @@ -545,9 +545,9 @@ def permute_free_grouped_gemm_bf16_dgrad( ) -> torch.Tensor: """Route-list gather-in-GEMM dgrad (FC1 backward wrt input). - ``grad_output`` is the compact per-route gradient ``[num_routes, out_features]``. The - kernel contracts over ``out_features`` to produce per-route ``dX = [num_routes, in]``, - which is then scatter-added back onto the received tokens. + ``grad_output`` is the **dense route-ordered** gradient ``[num_routes, out_features]``. The + kernel contracts over ``out_features`` to produce per-route ``dX`` in a block-padded + route-ordered buffer, then gather-combine sums back to **token-ordered** ``dA``. Returns ------- @@ -560,7 +560,7 @@ def permute_free_grouped_gemm_bf16_dgrad( ) if grad_output.dim() != 2: raise ValueError( - f"grad_output must be compact [num_routes, out_features], got {grad_output.shape}." + f"grad_output must be dense route-ordered [num_routes, out_features], got {grad_output.shape}." ) if not grad_output.is_contiguous(): grad_output = grad_output.contiguous() @@ -582,12 +582,10 @@ def permute_free_grouped_gemm_bf16_dgrad( block_size_m = _ensure_fwd_align(routing, grad_output, weights_stacked) - # Two-stage dgrad reduction (contention-free): (1) plain compact store of each per-route - # dX[route] = grad[route] @ W1[e] into an [T * min(topk, E), in] bf16 buffer (coalesced, no atomics), - # then (2) a token-parallel gather-combine summing each token's route rows. This replaces - # the fused atomic scatter-to-token, whose per-token contention dominated the FC1 dgrad. + # Two-stage dgrad: (1) store per-route dX into block-padded route-ordered [em_max, in], + # then (2) token-parallel gather-combine back to token-ordered dA. Replaces atomic scatter. em_max = routing.sorted_slot_ids.shape[0] - compact = torch.empty( + route_buf = torch.empty( (em_max, in_features), dtype=torch.bfloat16, device=grad_output.device, @@ -595,7 +593,7 @@ def permute_free_grouped_gemm_bf16_dgrad( _pf_moe_fwd( grad_output, weights_stacked, - compact, + route_buf, routing, num_recv_tokens=routing.num_recv_tokens, block_m=block_size_m, @@ -603,7 +601,7 @@ def permute_free_grouped_gemm_bf16_dgrad( dgrad=True, ) return route_gather_combine( - compact, + route_buf, routing.token_routes, routing.token_route_count, routing.num_recv_tokens, @@ -656,7 +654,7 @@ def permute_free_grouped_gemm_bf16_wgrad( raise TypeError("permute_free_grouped_gemm_bf16_wgrad requires bf16 inputs.") if grad_output.dim() != 2: raise ValueError( - f"grad_output must be compact [num_routes, out_features], got {grad_output.shape}." + f"grad_output must be dense route-ordered [num_routes, out_features], got {grad_output.shape}." ) num_experts, out_features, in_features = (int(v) for v in weights_shape) @@ -744,18 +742,18 @@ def permute_free_grouped_gemm_bf16_fc2( weights: torch.Tensor | list[torch.Tensor], routing: MoERoutingMetadata, ) -> torch.Tensor: - """Route-list FC2 forward with gather-combine to token order (bf16 MoE). + """Route-list FC2 forward with gather-combine to token-ordered output (bf16 MoE). - ``fc2_input`` is the route-ordered ``F``-wide FC1 activation (or a transient buffer rebuilt - from the saved ``2F`` pre-activation via :func:`permute_free_gated_act_fwd`). For each - route the kernel reads its row (``index_a_by_route_pos=True``) and computes the compact - per-route GEMM; a separate contention-free gather-combine pass sums each token's routes - into its output row. + ``fc2_input`` is **block-padded route-ordered** ``F``-wide FC1 activation (or a transient + buffer rebuilt from the saved ``2F`` preact via :func:`permute_free_gated_act_fwd`). For each + route slot the kernel reads its row (``index_a_by_route_pos=True``) and writes + block-padded route-ordered per-route GEMM output; gather-combine then sums each token's + routes into **token-ordered** ``[num_recv_tokens, out]``. Parameters ---------- fc2_input: - Route-ordered activations ``[em_max, in_features(F)]``, bf16. + Block-padded route-ordered activations ``[em_max, in_features(F)]``, bf16. weights: Expert weights ``[num_experts, out_features(H), in_features(F)]`` (W2) or list of ``[H, F]``. routing: @@ -764,7 +762,7 @@ def permute_free_grouped_gemm_bf16_fc2( Returns ------- torch.Tensor - ``[num_recv_tokens, out_features]``, bf16 (token order); ready for the cross-rank + ``[num_recv_tokens, out_features]``, bf16 (**token-ordered**); ready for the cross-rank combine. No separate unpermute is needed. """ if fc2_input.dtype != torch.bfloat16: @@ -791,13 +789,10 @@ def permute_free_grouped_gemm_bf16_fc2( block_size_m = _ensure_fwd_align(routing, fc2_input, weights_stacked) - # Two-stage combine (contention-free): (1) plain compact store of each per-route result - # y[route] = fc2_input[route] @ W2[e] into an [T * min(topk, E), out] bf16 buffer (coalesced, no - # atomics), then (2) a token-parallel gather-combine that sums each token's route rows. - # This replaces the fused atomic scatter-to-token, whose per-token atomic contention made - # the FC2 forward ~2x the FC1 forward. Gather has no contention, so the combine is ~free. + # Two-stage combine: (1) block-padded route-ordered per-route GEMM into [em_max, out], + # then (2) gather-combine to token-ordered output. Replaces atomic scatter-to-token. em_max = routing.sorted_slot_ids.shape[0] - compact = torch.empty( + route_buf = torch.empty( (em_max, out_features), dtype=torch.bfloat16, device=fc2_input.device, @@ -805,14 +800,14 @@ def permute_free_grouped_gemm_bf16_fc2( _pf_moe_fwd( fc2_input, weights_stacked, - compact, + route_buf, routing, num_recv_tokens=routing.num_recv_tokens, block_m=block_size_m, index_a_by_route_pos=True, ) return route_gather_combine( - compact, + route_buf, routing.token_routes, routing.token_route_count, routing.num_recv_tokens, @@ -825,18 +820,18 @@ def permute_free_grouped_gemm_bf16_fc2_dgrad( weights: torch.Tensor | list[torch.Tensor], routing: MoERoutingMetadata, ) -> torch.Tensor: - """FC2 dgrad: gather the token-space grad back into the compact route buffer. + """FC2 dgrad: gather **token-ordered** grad back into **block-padded route-ordered** dX. - ``grad_output`` is the token-space gradient ``[num_recv_tokens, out_features]`` (the grad - of FC2's fused-scatter forward output). For each route the kernel gathers its received - token's grad row (``index_a_by_route_pos=False``) and computes ``grad[token] @ W2[e]`` - (contracting over ``out_features``), writing the compact per-route - ``d(fc2_input) = [T * min(topk, E), in_features]``. + ``grad_output`` is **token-ordered** ``[num_recv_tokens, out_features]`` (grad of FC2's + gather-combine output). For each route slot the kernel gathers its token's grad row + (``index_a_by_route_pos=False``) and computes ``grad[token] @ W2[e]^T``, writing + block-padded route-ordered ``d(fc2_input)`` ``[em_max, in_features]``. Returns ------- torch.Tensor - ``[T * min(topk, E), in_features]``, bf16 (route order; valid rows are ``[0, num_routes)``). + ``[em_max, in_features]``, bf16 (**block-padded route-ordered**; valid slots are the + dense route range ``[0, num_routes)``). """ if grad_output.dtype != torch.bfloat16: raise TypeError( @@ -867,10 +862,10 @@ def permute_free_grouped_gemm_bf16_fc2_dgrad( block_size_m = _ensure_fwd_align(routing, grad_output, weights_stacked) - # Padded route-order output; the gather-GEMM writes only the compact [0, num_routes) - # range and never visits the tail (bounded by num_tokens_post_padded, sentinel-masked). + # Block-padded route-ordered output; the gather-GEMM writes only the dense route range + # [0, num_routes) and never visits the tail (bounded by num_tokens_post_padded, sentinel-masked). # Left uninitialized (torch.empty) to skip the em_max*N zero-init; consumers read only the - # compact range via the routing metadata, so the garbage tail is never observed. + # dense route range via the routing metadata, so the garbage tail is never observed. em_max = routing.sorted_slot_ids.shape[0] dx = torch.empty( (em_max, in_features), diff --git a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py index 397369ada..ca26af741 100644 --- a/transformer_engine/pytorch/moe/pf_fwd_wrapper.py +++ b/transformer_engine/pytorch/moe/pf_fwd_wrapper.py @@ -7,9 +7,7 @@ metadata (``sorted_slot_ids`` / ``expert_ids`` / ``block_start``) can be reused verbatim. Writes the block-padded ``[em_max, WIDTH_N]`` slot output in place. -Forward gather-GEMM is FlyDSL-only (MegaMOE-ported plain GEMM). Gated activation and -route-prob apply live in standalone Triton helpers (:mod:`pf_helper_kernels`), not in the -GEMM kernel. +Forward gather-GEMM is FlyDSL-only (MegaMOE-ported plain GEMM). """ from __future__ import annotations @@ -18,7 +16,6 @@ __all__ = [ "flydsl_moe_fwd", - "flydsl_moe_fwd_supported", ] # MegaMOE hand-tuned bf16 grouped-GEMM tile geometry (offline-swept in primus-turbo's @@ -58,9 +55,8 @@ def _run_gather_gemm( if transpose_b: # dgrad: NN GEMM contracting the incoming grad against the forward weight [E, N, K] - # over the output-feature axis. FC1 dgrad (index_a_by_route_pos=True) is compact - # route-read; FC2 dgrad (index_a_by_route_pos=False) gathers token-space grad - # [num_recv, N] into the compact route output [em_max, K]. + # route-read from block-padded route-ordered grad; FC2 dgrad gathers token-ordered + # grad [num_recv, N] into block-padded route-ordered dX [em_max, K]. from ..flydsl_kernels.permute_free_grouped_gemm.pf_dgrad import grouped_gemm_dgrad_bf16 dgrad_group_m = _PF_DGRAD_FC1_GROUP_M if index_a_by_route_pos else _PF_GROUP_M @@ -80,23 +76,6 @@ def _run_gather_gemm( ) -def flydsl_moe_fwd_supported( - A: torch.Tensor, - B: torch.Tensor, - *, - block_m: int, -) -> bool: - """Whether the in-tree FlyDSL fwd/dgrad kernels can handle these operands.""" - if A.dtype != torch.bfloat16 or B.dtype != torch.bfloat16: - return False - if A.stride(1) != 1: - return False - if B.stride(2) != 1: - return False - # em_max % block_m is only known at launch (from C), so it is not prechecked here. - return True - - def flydsl_moe_fwd( A: torch.Tensor, B: torch.Tensor, @@ -109,13 +88,14 @@ def flydsl_moe_fwd( index_a_by_route_pos: bool = False, dgrad: bool = False, ) -> None: - """Route-list gather-GEMM forward, writing the compact ``C[em_max, WIDTH_N]`` in place. - - ``A`` is ``[*, K]`` (received-token acts, gathered by ``sorted_slot_ids`` when - ``index_a_by_route_pos=False``, else read at the compact route row). ``B`` is - ``[num_experts, N_OUT, K]`` (contiguous inner ``K``). Set ``dgrad=True`` for the - data-gradient path (same ``B`` layout). Gated activation and route-prob apply are **not** - fused here; use the standalone helpers in :mod:`pf_helper_kernels`. + """Route-list gather-GEMM, writing block-padded route-ordered ``C[em_max, WIDTH_N]`` in place. + + ``A`` is token-ordered ``[num_recv, K]`` gathered via ``sorted_slot_ids`` when + ``index_a_by_route_pos=False`` (FC1), or block-padded route-ordered read by route slot when + ``index_a_by_route_pos=True`` (FC2). ``B`` is ``[num_experts, N_OUT, K]`` (contiguous inner + ``K``). Set ``dgrad=True`` for the data-gradient path (same ``B`` layout). Gated activation + and route-prob apply are **not** fused here; use the standalone helpers in + :mod:`pf_helper_kernels`. """ assert A.dtype == B.dtype == C.dtype == torch.bfloat16 assert A.stride(1) == 1, "A must be contiguous along the contraction (K)" diff --git a/transformer_engine/pytorch/moe/pf_helper_kernels.py b/transformer_engine/pytorch/moe/pf_helper_kernels.py index c43dfc886..8ce1bc1ae 100644 --- a/transformer_engine/pytorch/moe/pf_helper_kernels.py +++ b/transformer_engine/pytorch/moe/pf_helper_kernels.py @@ -7,7 +7,7 @@ Grouped GEMM for FC1/FC2 forward, backward (dgrad), and wgrad is FlyDSL-only (``pf_fwd_wrapper``, ``pf_wgrad_wrapper``). This module retains Triton kernels for routing metadata construction, gated-activation recompute/bwd, and the token-order -gather-combine pass that follows the compact route-list GEMM outputs. +gather-combine pass that follows the block-padded route-ordered GEMM outputs. """ from __future__ import annotations @@ -151,7 +151,7 @@ def _gated_act_prob_bwd_kernel( dpre_ptr, # [T * min(topk, E), 2F] out: grad wrt the raw 2F GEMM output grad_probs_ptr, # [num_recv_tokens, E] out (fp32) act_out_ptr, # [T * min(topk, E), F] out: fused fc2_input = act(g)*u*prob (EMIT_ACT only) - nbound_ptr, # [1] int32 device scalar: dynamic upper bound on compact routes + nbound_ptr, # [1] int32 device scalar: dynamic upper bound on route slots num_recv_tokens, F, stride_gom, @@ -192,9 +192,9 @@ def _gated_act_prob_bwd_kernel( ``up`` and ``prob`` are already live in registers here, this is one extra multiply + store and removes the redundant ``2F`` HBM read (see :func:`fused_gated_act_prob_fwd`). - The compact route buffers are statically over-allocated to the worst-case + The block-padded route buffers are statically over-allocated to the worst-case ``routes_max = T * topk`` (sync-free shape bound), but the real routes occupy only the - dense head ``[0, num_routes)`` -- under expert parallelism this can be ~topk*E_local/E + dense route-ordered head ``[0, num_routes)`` -- under expert parallelism this can be ~topk*E_local/E times smaller. ``nbound_ptr`` carries the actual (block-padded) route extent as a device scalar so tail programs beyond it exit before touching HBM, instead of grinding through the padding at full memory bandwidth. @@ -298,7 +298,7 @@ def fused_gated_act_prob_bwd( route-prob multiply. When given, its gradient is returned. num_routes_bound: Optional ``[1]`` int32 device scalar giving a (block-padded) upper bound on the number - of *real* compact routes. The route buffers are statically sized to the worst case + of *real* route slots. The route buffers are statically sized to the worst case ``routes_max = T * topk``, but under expert parallelism only a small dense head is populated; passing the actual extent (e.g. ``num_tokens_post_padded`` from the routing metadata) lets tail programs exit early instead of streaming the padding through HBM. @@ -388,7 +388,7 @@ def _gated_act_prob_fwd_kernel( token_ptr, # [routes_max] route -> received-token row expert_ptr, # [routes_max] route -> local expert act_ptr, # [routes_max, F] out: act(gate) * up * prob - nbound_ptr, # [1] int32 device scalar: dynamic upper bound on compact routes + nbound_ptr, # [1] int32 device scalar: dynamic upper bound on route slots num_recv_tokens, F, stride_prem, @@ -410,7 +410,7 @@ def _gated_act_prob_fwd_kernel( checkpoint only the ``2F`` preact (never the ``F``-wide act) and rebuild it just-in-time for the FC2 wgrad. Padded / over-allocated routes carry the ``token == num_recv_tokens`` sentinel and get ``prob == 0`` (their act rows are ignored downstream). ``nbound_ptr`` - bounds tail programs to the real compact route extent (sync-free, EP-friendly early exit). + bounds tail programs to the real block-padded route extent (sync-free, EP-friendly early exit). """ pid = tl.program_id(axis=0) num_routes_bound = tl.load(nbound_ptr)