diff --git a/.github/workflows/mslk_ci_rocm.yml b/.github/workflows/mslk_ci_rocm.yml index 806716e1..546b3fb3 100644 --- a/.github/workflows/mslk_ci_rocm.yml +++ b/.github/workflows/mslk_ci_rocm.yml @@ -37,6 +37,8 @@ on: - 'mslk/attention/flydsl/**' - 'test/attention/flydsl/**' - 'test/flydsl/**' + # FlyDSL GEMM ops + - 'mslk/gemm/flydsl/**' # GEMM tests - 'test/gemm/gemm_test.py' # AMD/ROCm Triton GEMM kernels diff --git a/csrc/gemm/gemm_ops.cpp b/csrc/gemm/gemm_ops.cpp index 66d3b633..5d9361b2 100644 --- a/csrc/gemm/gemm_ops.cpp +++ b/csrc/gemm/gemm_ops.cpp @@ -59,11 +59,17 @@ TORCH_LIBRARY_FRAGMENT(mslk, m) { // Triton implementation registered by fp8_groupwise_gemm.py. m.def( "f8f8bf16_groupwise(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale) -> Tensor"); - // FP8 groupwise grouped GEMM: shared schema; CUDA uses CUTLASS, ROCm uses - // the Triton implementation registered by fp8_groupwise_grouped_gemm.py. + // FP8 groupwise grouped GEMM: shared schema; CUDA uses CUTLASS, ROCm uses the + // FlyDSL implementation registered by + // mslk/gemm/flydsl/fp8_groupwise_grouped_gemm.py. m.def( "f8f8bf16_groupwise_grouped(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor M_sizes) -> Tensor"); #ifdef USE_ROCM + // Sibling of f8f8bf16_groupwise_grouped taking weights already swizzled into + // the MFMA B layout; schema only on ROCm, implemented by the same FlyDSL + // module via torch.library.impl at Python import time. + m.def( + "f8f8bf16_groupwise_grouped_preshuffle(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor M_sizes) -> Tensor"); m.def( "f8f8f16_rowwise(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor? bias=None, bool use_fast_accum=True) -> Tensor"); m.def( diff --git a/mslk/flydsl/kernels/__init__.py b/mslk/flydsl/kernels/__init__.py new file mode 100644 index 00000000..581f84e4 --- /dev/null +++ b/mslk/flydsl/kernels/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict diff --git a/mslk/flydsl/kernels/common/__init__.py b/mslk/flydsl/kernels/common/__init__.py new file mode 100644 index 00000000..581f84e4 --- /dev/null +++ b/mslk/flydsl/kernels/common/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict diff --git a/mslk/flydsl/kernels/common/kernels_common.py b/mslk/flydsl/kernels/common/kernels_common.py new file mode 100644 index 00000000..91f40095 --- /dev/null +++ b/mslk/flydsl/kernels/common/kernels_common.py @@ -0,0 +1,169 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Common helpers shared by kernel modules. + +Keep helper naming consistent with other kernel helpers (e.g. `mfma_preshuffle_pipeline.py`), +but this module is intentionally small and MLIR-dialect facing. +""" + +from contextlib import contextmanager + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import arith as _std_arith +from flydsl._mlir.dialects import builtin +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import gpu as _gpu +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import scf as _scf +from flydsl.expr import arith as _expr_arith +from flydsl.expr import buffer_ops, const_expr +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch, is_rdna_arch + + +def get_llvm_ptr(ptr, offset, dtype_bytes, ptr_type=None): + """Build a global (address-space 1) ``!llvm.ptr`` at ``ptr + offset*dtype_bytes``. + + Shared home for the LLVM-ptr arithmetic used by atomic/global accesses + (previously duplicated in hgemm_splitk.py and rmsnorm_kernel.py). + """ + if ptr_type is None: + ptr_type = ir.Type.parse("!llvm.ptr<1>") + base_ptr = _fly.extract_aligned_pointer_as_index(ptr_type, ptr) + base_ptr = _llvm.PtrToIntOp(T.i64, base_ptr).result + byte_offset = _expr_arith.index_cast(T.i64, fx.Index(offset) * fx.Index(dtype_bytes)) + llvm_ptr = _llvm.AddOp(base_ptr, byte_offset, _llvm.IntegerOverflowFlags(0)).result + llvm_ptr = _llvm.IntToPtrOp(ptr_type, llvm_ptr).result + return llvm_ptr._value if const_expr(hasattr(llvm_ptr, "_value")) else llvm_ptr + + +def atomic_add( + dst, + offset, + value, + *, + dtype_bytes=4, + syncscope="agent", + ordering=None, + alignment=None, + ptr_type=None, +): + """Atomically add ``value`` into ``dst[offset]`` in global memory. + + Wraps the ``get_llvm_ptr`` + ``llvm.atomicrmw`` pair that kernels used to + inline (rmsnorm backward ``dweight`` accumulation, hgemm split-K epilogue and + semaphore). Selects ``fadd`` for a floating-point operand and integer ``add`` + otherwise, from the operand's IR type, so a single call covers both cases. + Returns the atomicrmw result (the value previously stored at ``dst[offset]``). + + ``dtype_bytes`` sizes the byte offset and, unless ``alignment`` is given, is + reused as the access alignment. + """ + ptr = get_llvm_ptr(dst, offset, dtype_bytes, ptr_type=ptr_type) + val = value.ir_value() if const_expr(hasattr(value, "ir_value")) else value + elem_ty = val.type.element_type if isinstance(val.type, ir.VectorType) else val.type + bin_op = _llvm.AtomicBinOp.fadd if isinstance(elem_ty, ir.FloatType) else _llvm.AtomicBinOp.add + if ordering is None: + ordering = _llvm.AtomicOrdering.monotonic + if alignment is None: + alignment = dtype_bytes + return _llvm.AtomicRMWOp( + bin_op, + ptr, + val, + ordering, + syncscope=syncscope, + alignment=alignment, + ).result + + +@contextmanager +def _if_then(if_op, scf=None): + """Context manager for SCF IfOp then-region across old/new Python APIs. + + Ensures the then block always ends with a YieldOp. + The optional *scf* parameter is accepted for backward compatibility + but ignored — the module-level import is used. + """ + with ir.InsertionPoint(if_op.then_block): + try: + yield if_op.then_block + finally: + blk = if_op.then_block + if (not blk.operations) or not isinstance(blk.operations[-1], _scf.YieldOp): + _scf.YieldOp([]) + + +@contextmanager +def _if_else(if_op, scf=None): + """Context manager for SCF IfOp else-region across old/new Python APIs. + + Ensures the else block always ends with a YieldOp. The optional *scf* + parameter is accepted for backward compatibility but ignored. + """ + if getattr(if_op, "else_block", None) is None: + raise RuntimeError("IfOp has no else block") + with ir.InsertionPoint(if_op.else_block): + try: + yield if_op.else_block + finally: + blk = if_op.else_block + if (not blk.operations) or not isinstance(blk.operations[-1], _scf.YieldOp): + _scf.YieldOp([]) + + +_VALID_A_DTYPES = frozenset(("fp8", "fp16", "int8", "fp4")) +_VALID_B_DTYPES = frozenset(("fp8", "fp16", "int8", "int4", "fp4")) + + +def validate_moe_dtypes(a_dtype: str, b_dtype: str) -> None: + """Validate a_dtype/b_dtype strings for mixed MoE kernels.""" + if a_dtype not in _VALID_A_DTYPES: + raise ValueError(f"a_dtype must be one of {tuple(sorted(_VALID_A_DTYPES))}, got {a_dtype!r}") + if b_dtype not in _VALID_B_DTYPES: + raise ValueError(f"b_dtype must be one of {tuple(sorted(_VALID_B_DTYPES))}, got {b_dtype!r}") + + +def dtype_to_elem_type(dtype_str: str): + """Map a dtype string to its FlyDSL numeric type. + + Supported: 'f32', 'f16', 'bf16', 'fp8' (OCP e4m3fn, not the fnuz variant). + """ + if dtype_str == "f32": + return fx.Float32 + if dtype_str == "f16": + return fx.Float16 + if dtype_str == "bf16": + return fx.BFloat16 + if dtype_str == "fp8": + return fx.Float8E4M3FN + raise ValueError(f"unsupported dtype: {dtype_str!r} (expected 'f32', 'f16', 'bf16', or 'fp8')") + + +def get_warp_size(arch=None): + """Return the wavefront/warp size for the given GPU architecture. + + CDNA (gfx9xx) uses wave64, RDNA (gfx10xx/gfx11xx/gfx12xx) uses wave32. + """ + if arch is None: + arch = get_rocm_arch() + return 32 if is_rdna_arch(arch) else 64 + + +def _create_llvm_ptr(value, address_space: int = 1): + value = buffer_ops._unwrap_value(value) + if isinstance(value.type, ir.IndexType): + i64_type = T.i64 + value = buffer_ops._unwrap_value(_std_arith.IndexCastOp(i64_type, value).result) + ptr_type = ir.Type.parse(f"!llvm.ptr<{address_space}>") + return _llvm.IntToPtrOp(ptr_type, value).result + + +def stream_ptr_to_async_token(stream_ptr_value, loc=None, ip=None): + stream_llvm_ptr = _create_llvm_ptr(stream_ptr_value) + + async_token_type = _gpu.AsyncTokenType.get() + cast_op = builtin.UnrealizedConversionCastOp([async_token_type], [stream_llvm_ptr], loc=loc, ip=ip) + return cast_op.results[0] diff --git a/mslk/flydsl/kernels/gemm/__init__.py b/mslk/flydsl/kernels/gemm/__init__.py new file mode 100644 index 00000000..581f84e4 --- /dev/null +++ b/mslk/flydsl/kernels/gemm/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict diff --git a/mslk/flydsl/kernels/gemm/fp8_gemm_utils.py b/mslk/flydsl/kernels/gemm/fp8_gemm_utils.py new file mode 100644 index 00000000..c3eadf3d --- /dev/null +++ b/mslk/flydsl/kernels/gemm/fp8_gemm_utils.py @@ -0,0 +1,255 @@ +# 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 + + +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 ceildiv(a: int, b: int) -> int: + return (a + b - 1) // b + + +def divmod(a: int, b: int) -> tuple[int, int]: + return (a // b, a % b) + + +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)]) diff --git a/mslk/flydsl/kernels/gemm/grouped_gemm_blockscale_common.py b/mslk/flydsl/kernels/gemm/grouped_gemm_blockscale_common.py new file mode 100644 index 00000000..82d93d1b --- /dev/null +++ b/mslk/flydsl/kernels/gemm/grouped_gemm_blockscale_common.py @@ -0,0 +1,1167 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Shared building blocks for the grouped FP8 blockscale GEMM kernels. + +Used by the grouped_gemm_blockscale contiguous and masked kernels. Holds the +parts of the two kernels that are byte-identical (parameter validation, +compile-time scalar constants, helper closures) so they live in one place. + +scale_b is indexed as [num_groups, scale_k, scale_n] (per-group, per-K-block, +per-N-block); scale_a is [scale_k, M] (transposed, per-token per-K-block). +""" + +from collections import namedtuple + +import flydsl.expr as fx +from flydsl._mlir.dialects import math as math_dialect +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector +from flydsl.expr.arith import ArithValue +from flydsl.expr.typing import T, Vector +from mslk.flydsl.kernels.mma.mfma_preshuffle_pipeline import ( + crd2idx, + lds_store_16b_xor16, + load_b_pack_k32, + make_preshuffle_b_layout, + swizzle_xor16, + tile_chunk_coord_i32, +) + +CompileConstants = namedtuple( + "CompileConstants", + [ + "total_threads", + "elem_bytes", + "num_k_tiles", + "scale_k", + "scale_n", + "sb_per_tile", + "k_unroll", + "kpack_bytes", + "tile_k_bytes", + "tile_k_dwords", + "bytes_a_per_tile", + "bytes_per_thread_a", + "a_load_bytes", + "chunk_i32_a", + "num_a_loads", + "chunk_i32_b", + "num_b_loads", + ], +) + + +def validate_params(*, n, k, tile_n, tile_k, scale_block_k, scale_block_n, out_dtype): + """Validate the divisibility constraints and out_dtype choice shared by + both grouped GEMM blockscale kernels.""" + if k % tile_k != 0: + raise ValueError(f"k ({k}) must be divisible by tile_k ({tile_k})") + if n % tile_n != 0: + raise ValueError(f"n ({n}) must be divisible by tile_n ({tile_n})") + if tile_k % scale_block_k != 0: + raise ValueError(f"tile_k ({tile_k}) must be divisible by scale_block_k ({scale_block_k})") + if tile_n % scale_block_n != 0: + raise ValueError(f"tile_n ({tile_n}) must be divisible by scale_block_n ({scale_block_n})") + if out_dtype not in ("bf16", "f16"): + raise ValueError(f"out_dtype must be 'bf16' or 'f16', got {out_dtype!r}") + + +# Conservative default for architectures missing from FlyDSL's capacity table. +_LDS_CAPACITY_FALLBACK_BYTES = 64 * 1024 + + +def lds_capacity_bytes(arch=None): + """LDS bytes available to one workgroup on ``arch``. + + This is arch-dependent and the difference is large: CDNA3 (gfx942/MI300) has + 64 KiB while CDNA4 (gfx950/MI350) has 160 KiB. Sourced from FlyDSL's + SMEM_CAPACITY_MAP so the limit stays in sync with the compiler that enforces + it; unknown architectures fall back to the conservative 64 KiB. + """ + try: + from flydsl.utils.smem_allocator import SMEM_CAPACITY_MAP + except Exception: + return _LDS_CAPACITY_FALLBACK_BYTES + return SMEM_CAPACITY_MAP.get(str(arch), _LDS_CAPACITY_FALLBACK_BYTES) + + +def _check_lds_budget(*, variant, total, detail, tile_m, tile_n, tile_k, arch): + """Raise ValueError if ``total`` LDS bytes overflow the device capacity. + + Must run BEFORE the kernel is traced: the compiler backend reports an LDS + overflow as a hard error that aborts the process, which an autotuner cannot + catch and skip. + """ + capacity = lds_capacity_bytes(arch) + if total > capacity: + raise ValueError( + f"{variant} LDS budget {total} bytes exceeds {capacity} on " + f"{arch or 'unknown arch'} ({detail}) for tile_m={tile_m} " + f"tile_n={tile_n} tile_k={tile_k}. Reduce tile_m/tile_n/tile_k." + ) + + +def validate_lds_budget_preshuffle(*, tile_m, tile_n, tile_k, elem_bytes=1, arch=None): + """Check the preshuffle-B kernel's LDS budget fits in one workgroup's LDS. + + B is loaded HBM->registers here, so only the ping-pong A tiles occupy LDS + during the K-loop; the CShuffle epilogue output aliases that same arena. + Budget = max(A ping-pong, epilogue out). + """ + lds_a_bytes = 2 * tile_m * tile_k * elem_bytes # ping-pong A + lds_out_bytes = tile_m * tile_n * 2 # bf16/f16 epilogue output, aliases base + _check_lds_budget( + variant="preshuffle-B", + total=max(lds_a_bytes, lds_out_bytes), + detail=f"A ping-pong {lds_a_bytes}, epilogue {lds_out_bytes}", + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + arch=arch, + ) + + +def validate_lds_budget_plain( + *, tile_m, tile_n, tile_k, elem_bytes=1, b_pingpong=False, arch=None +): + """Check the plain-B kernel's LDS budget fits in one workgroup's LDS. + + The plain-B kernel stages BOTH A (ping-pong) and B through LDS during the + K-loop, so they coexist; the epilogue output aliases that same arena (it + runs after the final barrier). Budget = max(K-loop staging, epilogue out). + Raises ValueError if a tile config would overflow LDS. + """ + lds_a_bytes = 2 * tile_m * tile_k * elem_bytes # ping-pong A + b_buffers = 2 if b_pingpong else 1 + lds_b_bytes = b_buffers * tile_n * tile_k * elem_bytes + lds_out_bytes = tile_m * tile_n * 2 # bf16/f16 epilogue output, aliases base + kloop_bytes = lds_a_bytes + lds_b_bytes + _check_lds_budget( + variant="plain-B", + total=max(kloop_bytes, lds_out_bytes), + detail=( + f"A ping-pong {lds_a_bytes} + B {lds_b_bytes} = {kloop_bytes}, " + f"epilogue {lds_out_bytes}, b_pingpong={b_pingpong}" + ), + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + arch=arch, + ) + + +def out_mlir_for(out_dtype): + """Return a zero-arg callable that yields the MLIR element type for the + chosen output dtype. Matches the original local `out_mlir` lambda exactly + so MLIR emission is unchanged.""" + return lambda: T.bf16 if out_dtype == "bf16" else T.f16 + + +def compute_compile_constants(*, n, k, tile_m, tile_n, tile_k, scale_block_k, scale_block_n): + """Compute the compile-time scalar constants shared by both kernels. + + Returns a `CompileConstants` namedtuple. Pure-Python — no MLIR ops emitted. + """ + total_threads = 256 + elem_bytes = 1 # FP8 + num_k_tiles = k // tile_k + scale_k = k // scale_block_k + scale_n = n // scale_block_n + sb_per_tile = tile_k // scale_block_k # scale blocks per K-tile + k_unroll = tile_k // 64 # K64-byte micro-steps (for K32 MFMA pairs) + kpack_bytes = 16 # 16-byte packs for FP8 + + tile_k_bytes = tile_k * elem_bytes + tile_k_dwords = tile_k_bytes // 4 + bytes_a_per_tile = tile_m * tile_k * elem_bytes + bytes_per_thread_a = bytes_a_per_tile // total_threads + a_load_bytes = 16 # 16-byte loads (dwordx4) + chunk_i32_a = a_load_bytes // 4 # 4 dwords per load + num_a_loads = bytes_per_thread_a // a_load_bytes + + # Plain-B staging (non-preshuffle kernel): B tile is [tile_n, tile_k], + # loaded HBM->LDS just like A but with N in place of M. + bytes_b_per_tile = tile_n * tile_k * elem_bytes + bytes_per_thread_b = bytes_b_per_tile // total_threads + chunk_i32_b = a_load_bytes // 4 # same 16-byte dwordx4 load + num_b_loads = bytes_per_thread_b // a_load_bytes + + return CompileConstants( + total_threads=total_threads, + elem_bytes=elem_bytes, + num_k_tiles=num_k_tiles, + scale_k=scale_k, + scale_n=scale_n, + sb_per_tile=sb_per_tile, + k_unroll=k_unroll, + kpack_bytes=kpack_bytes, + tile_k_bytes=tile_k_bytes, + tile_k_dwords=tile_k_dwords, + bytes_a_per_tile=bytes_a_per_tile, + bytes_per_thread_a=bytes_per_thread_a, + a_load_bytes=a_load_bytes, + chunk_i32_a=chunk_i32_a, + num_a_loads=num_a_loads, + chunk_i32_b=chunk_i32_b, + num_b_loads=num_b_loads, + ) + + +def setup_lds_allocation(*, allocator, tile_m, tile_k, tile_n, elem_bytes): + """Reserve LDS for ping-pong A tiles and the CShuffle epilogue output. + + The ping-pong A buffers and the FP16/BF16 epilogue output share the same + LDS arena (alias), so we reserve the max of the two. Returns + `(lds_alloc_offset, lds_tile_elems)` where `lds_tile_elems` is the + A-element stride between the ping and pong halves. + """ + lds_a_bytes = tile_m * tile_k * elem_bytes + lds_pingpong_bytes = 2 * lds_a_bytes + lds_out_bytes = tile_m * tile_n * 2 # bf16/f16 = 2 bytes per element + lds_total_bytes = max(lds_pingpong_bytes, lds_out_bytes) + lds_alloc_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_alloc_offset + lds_total_bytes + lds_tile_elems = tile_m * tile_k # element offset between ping and pong + return lds_alloc_offset, lds_tile_elems + + +def setup_lds_allocation_plain(*, allocator, tile_m, tile_n, tile_k, elem_bytes, b_pingpong=False): + """Reserve LDS for the plain-B kernel: ping-pong A + (single/ping-pong) B + + aliased CShuffle epilogue output. + + Unlike the preshuffle kernel (which loads B straight to registers and needs + no B LDS), plain B is staged HBM->LDS->registers alongside A, so A and B + LDS coexist during the K-loop. The epilogue output aliases the whole arena + (offset 0) since it runs after the final K-loop barrier. + + Returns `(lds_alloc_offset, lds_tile_elems, lds_b_offset_elems)` where: + - `lds_alloc_offset` is the byte base of the arena (A ping half at 0), + - `lds_tile_elems` is the A ping<->pong element stride (= tile_m*tile_k), + - `lds_b_offset_elems` is the element offset (from arena base) to the B + buffer, i.e. just past the A ping-pong region. + """ + lds_a_elems = tile_m * tile_k + lds_a_pingpong_elems = 2 * lds_a_elems + b_buffers = 2 if b_pingpong else 1 + lds_b_elems = b_buffers * tile_n * tile_k + kloop_elems = lds_a_pingpong_elems + lds_b_elems # FP8: 1 byte/elem + lds_out_bytes = tile_m * tile_n * 2 + lds_total_bytes = max(kloop_elems * elem_bytes, lds_out_bytes) + lds_alloc_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_alloc_offset + lds_total_bytes + lds_tile_elems = lds_a_elems + lds_b_offset_elems = lds_a_pingpong_elems + return lds_alloc_offset, lds_tile_elems, lds_b_offset_elems + + +def make_a_tile_loaders( + *, + a_rsrc, + lds_a, + layout_lds, + bx_m, + tx, + tile_m, + tile_k, + tile_k_bytes, + tile_k_dwords, + chunk_i32_a, + num_a_loads, + total_threads, + elem_bytes, + k_in, + m_in=None, + group_idx=None, +): + """Build the prefetch + LDS-store closures for the A tile. + + Returns `(prefetch_a_tile, store_a_tile_to_lds, a_row_local, + a_col_local_i32, k_blocks16)`. When `m_in` and `group_idx` are both + None (contig path) no group offset is emitted; when both are provided + (masked path), `group_idx * m_in * (k_in/4)` is added as the leading + term inside `prefetch_a_tile`, exactly matching the original masked + code so the resulting MLIR (and ISA) is byte-identical. `k_blocks16` + is returned for reuse by the downstream LDS-load helper. + """ + layout_a_tile_div4 = fx.make_layout((tile_m, tile_k_dwords), stride=(tile_k_dwords, 1)) + c_chunk_a = fx.Index(chunk_i32_a) + tx_i32_base = tx * c_chunk_a + _k_div4_factor = k_in // fx.Index(4) + if m_in is not None and group_idx is not None: + a_tile_offset_div4 = group_idx * m_in * _k_div4_factor # 3D A Offset + else: + a_tile_offset_div4 = None + k_blocks16 = arith.index(tile_k_bytes // 16) + c4_bytes = fx.Index(4) + + a_row_local = [] + a_col_local_i32 = [] + for i in range_constexpr(num_a_loads): + row_local, col_local_i32 = tile_chunk_coord_i32( + arith, + tx_i32_base=tx_i32_base, + i=i, + total_threads=total_threads, + layout_tile_div4=layout_a_tile_div4, + chunk_i32=chunk_i32_a, + ) + a_row_local.append(row_local) + a_col_local_i32.append(col_local_i32) + + def prefetch_a_tile(k_tile_idx_py): + """Load A tile from global memory into VGPRs.""" + base_k_div4 = fx.Index(k_tile_idx_py * tile_k_dwords) + parts = [] + for i in range_constexpr(num_a_loads): + row_global = bx_m + a_row_local[i] + if a_tile_offset_div4 is None: + idx_i32 = row_global * _k_div4_factor + base_k_div4 + a_col_local_i32[i] + else: + idx_i32 = a_tile_offset_div4 + row_global * _k_div4_factor + base_k_div4 + a_col_local_i32[i] + a_vec = buffer_ops.buffer_load(a_rsrc, idx_i32, vec_width=4, dtype=T.i32) + parts.append(Vector(a_vec).bitcast(fx.Int32)) + return parts + + def store_a_tile_to_lds(a_parts, lds_base): + """Write prefetched A tile from VGPRs into LDS with XOR16 swizzle.""" + for i in range_constexpr(num_a_loads): + lds_store_16b_xor16( + arith, + vector, + lds_memref=lds_a, + vec16_ty=T.f8x16, + layout_lds=layout_lds, + row_local=a_row_local[i], + col_local_i32=a_col_local_i32[i], + tx_c4=c4_bytes, + k_blocks16=k_blocks16, + lds_base=lds_base, + vec_part_i32x4=a_parts[i], + elem_bytes=elem_bytes, + ) + + return prefetch_a_tile, store_a_tile_to_lds, a_row_local, a_col_local_i32, k_blocks16 + + +def make_b_tile_loaders( + *, + b_rsrc, + lds_b, + layout_lds_b, + by_n, + group_idx, + tx, + tile_n, + tile_k, + tile_k_bytes, + tile_k_dwords, + chunk_i32_b, + num_b_loads, + total_threads, + elem_bytes, + n_in, + k_in, +): + """Build the prefetch + LDS-store closures for a PLAIN (non-preshuffled) + B tile [tile_n, tile_k]. + + Mirror of `make_a_tile_loaders` with N in place of M. B is `[G, N, K]` + row-major, so the per-tile global base adds the group offset + `group_idx * n_in * (k_in/4)` (always present — B is always grouped) plus + the N-tile base `by_n` (the block's N-block start). Coalesced 16-byte + (dwordx4) loads via `tile_chunk_coord_i32`; LDS store uses the same XOR16 + swizzle as A. Returns `(prefetch_b_tile, store_b_tile_to_lds, b_row_local, + b_col_local_i32, k_blocks16_b)`. + """ + layout_b_tile_div4 = fx.make_layout((tile_n, tile_k_dwords), stride=(tile_k_dwords, 1)) + c_chunk_b = fx.Index(chunk_i32_b) + tx_i32_base = tx * c_chunk_b + _k_div4_factor = k_in // fx.Index(4) + # B is [G, N, K]: leading offset selects this tile's group and N-block base. + b_tile_offset_div4 = group_idx * n_in * _k_div4_factor + k_blocks16_b = arith.index(tile_k_bytes // 16) + c4_bytes = fx.Index(4) + + b_row_local = [] + b_col_local_i32 = [] + for i in range_constexpr(num_b_loads): + row_local, col_local_i32 = tile_chunk_coord_i32( + arith, + tx_i32_base=tx_i32_base, + i=i, + total_threads=total_threads, + layout_tile_div4=layout_b_tile_div4, + chunk_i32=chunk_i32_b, + ) + b_row_local.append(row_local) + b_col_local_i32.append(col_local_i32) + + def prefetch_b_tile(k_tile_idx_py): + """Load plain B tile from global memory into VGPRs (coalesced dwordx4).""" + base_k_div4 = fx.Index(k_tile_idx_py * tile_k_dwords) + parts = [] + for i in range_constexpr(num_b_loads): + row_global = by_n + b_row_local[i] # global N row + idx_i32 = b_tile_offset_div4 + row_global * _k_div4_factor + base_k_div4 + b_col_local_i32[i] + b_vec = buffer_ops.buffer_load(b_rsrc, idx_i32, vec_width=4, dtype=T.i32) + parts.append(Vector(b_vec).bitcast(fx.Int32)) + return parts + + def store_b_tile_to_lds(b_parts, lds_base): + """Write prefetched B tile from VGPRs into LDS with XOR16 swizzle.""" + for i in range_constexpr(num_b_loads): + lds_store_16b_xor16( + arith, + vector, + lds_memref=lds_b, + vec16_ty=T.f8x16, + layout_lds=layout_lds_b, + row_local=b_row_local[i], + col_local_i32=b_col_local_i32[i], + tx_c4=c4_bytes, + k_blocks16=k_blocks16_b, + lds_base=lds_base, + vec_part_i32x4=b_parts[i], + elem_bytes=elem_bytes, + ) + + return prefetch_b_tile, store_b_tile_to_lds, b_row_local, b_col_local_i32, k_blocks16_b + + +def make_lds_loader(*, lds_a, layout_lds, k_blocks16): + """Build the LDS-side A K64 pack loader. + + Returns `lds_load_packs_k64(curr_row_a_lds, col_base_bytes, lds_base)` + which loads 16B from LDS with the XOR16 swizzle and returns the two + i64 halves. + """ + + def lds_load_packs_k64(curr_row_a_lds, col_base_bytes, lds_base): + col_base_swz_bytes = swizzle_xor16(curr_row_a_lds, col_base_bytes, k_blocks16) + idx_a16 = crd2idx((curr_row_a_lds, col_base_swz_bytes), layout_lds) + lds_base + loaded_a16 = Vector.load(T.vec(16, T.f8), lds_a, [idx_a16]) + a_i64x2 = loaded_a16.bitcast(fx.Int64) + return a_i64x2[0], a_i64x2[1] + + return lds_load_packs_k64 + + +def make_lds_b_loader(*, lds_b, layout_lds_b, k_blocks16_b): + """Build the LDS-side plain-B K64 pack loader (mirror of `make_lds_loader`). + + Returns `lds_load_b_packs_k64(row_n_lds, col_base_bytes, lds_base_b)` which + loads 16B from the B LDS buffer with the same XOR16 swizzle A uses and + returns the two i64 halves — the exact fragment form the MFMA consumes for + the B operand. + """ + + def lds_load_b_packs_k64(row_n_lds, col_base_bytes, lds_base_b): + col_base_swz_bytes = swizzle_xor16(row_n_lds, col_base_bytes, k_blocks16_b) + idx_b16 = crd2idx((row_n_lds, col_base_swz_bytes), layout_lds_b) + lds_base_b + loaded_b16 = Vector.load(T.vec(16, T.f8), lds_b, [idx_b16]) + b_i64x2 = loaded_b16.bitcast(fx.Int64) + return b_i64x2[0], b_i64x2[1] + + return lds_load_b_packs_k64 + + +def make_plain_b_tile( + *, lds_load_b_packs_k64, lane_mod_16, n_tile_base, col_offset_base_bytes, k_unroll, num_acc_n +): + """Build the plain-B tile assembler that reads B from LDS (mirror of the + preshuffle `make_b_loader`, but sourcing from LDS instead of HBM). + + Returns `load_b_tile_from_lds(lds_base_b)` producing the SAME structure the + preshuffle path did — a list of length `k_unroll` where each entry is + `(packs0[ni], packs1[ni])` (two i64 halves per K64 micro-step, per N-acc) — + so `make_compute_tile` consumes it unchanged. + + N-row addressing must match what the MFMA B operand expects, i.e. the same + N-column `make_n_block_coords` uses (common.py): + col = by_n + n_tile_base + ni*16 + lane_mod_16 + where `n_tile_base = wave_mod_4 * n_per_wave` is this WAVE's N sub-range. + Since the B LDS buffer holds the block's [tile_n, tile_k] tile (row 0 = the + block's `by_n`), the LDS N-row for accumulator `ni` is the tile-LOCAL row + `n_tile_base + ni*16 + lane_mod_16` (by_n is the tile base, already 0 in the + LDS-local frame). Missing `n_tile_base` gives the wrong N per wave. + + K addressing mirrors A exactly: per-pack column base is + `col_offset_base_bytes + ku*64` (`col_offset_base_bytes = lane_div_16*16`). + """ + + def load_b_tile_from_lds(lds_base_b): + b_tile = [] + for ku in range_constexpr(k_unroll): + col_base_bytes = col_offset_base_bytes + fx.Index(ku * 64) + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + row_n_lds = n_tile_base + (ni * 16) + lane_mod_16 + b0, b1 = lds_load_b_packs_k64(row_n_lds, col_base_bytes, lds_base_b) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + return load_b_tile_from_lds + + +def make_b_loader( + *, + arg_b, + b_rsrc, + layout_b, + n_blk_list, + n_intra_list, + lane_div_16, + kpack_bytes, + elem_bytes, + k_unroll, + num_acc_n, +): + """Build the B-tile loader closure. + + Returns `load_b_tile(base_k)` which loads all B packs for one K-tile, + returning a list of length `k_unroll` where each entry is + `(packs_half0[ni], packs_half1[ni])` for one K64 micro-step. + """ + + def load_b_pack(base_k, ki_step, ni): + return load_b_pack_k32( + buffer_ops, + arith, + vector, + arg_b=arg_b, + b_rsrc=b_rsrc, + layout_b=layout_b, + base_k=base_k, + ki_step=ki_step, + n_blk=n_blk_list[ni], + n_intra=n_intra_list[ni], + lane_div_16=lane_div_16, + elem_type=T.f8, + kpack_bytes=kpack_bytes, + elem_bytes=elem_bytes, + ) + + def load_b_tile(base_k): + b_tile = [] + for ku in range_constexpr(k_unroll): + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + ki0 = (ku * 2) + 0 + ki1 = (ku * 2) + 1 + b0 = load_b_pack(base_k, ki0, ni) + b1 = load_b_pack(base_k, ki1, ni) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + return load_b_tile + + +def pack_i64x4_to_i32x8(x0, x1, x2, x3): + """Pack four i64 values into a single i32x8 vector via i64x4 bitcast. + + Used to assemble the K=128 MFMA A/B operands on gfx950. + """ + v4 = Vector.from_elements([x0, x1, x2, x3], fx.Int64) + return v4.bitcast(fx.Int32) + + +def make_hot_loop_scheduler( + *, + _use_hw_scale, + sb_per_tile, + m_repeat, + num_acc_n, + k_unroll, + num_a_loads, + ku_per_sb, +): + """Build the per-tile sched_group_barrier scheduler closure. + + Emits the dsrd / mfma / vmem_rd / dswr group barriers in the order + matching the MoE stage-2 pattern. Returns a zero-arg closure to be + invoked once per K-tile body inside the ping-pong loop. + """ + + def hot_loop_scheduler(): + mfma_group = num_acc_n + if _use_hw_scale: + total_mfma = sb_per_tile * m_repeat * mfma_group + else: + total_mfma = k_unroll * m_repeat * mfma_group * 2 + rocdl.sched_group_barrier(rocdl.mask_dsrd, ku_per_sb * m_repeat, 0) + rocdl.sched_group_barrier(rocdl.mask_mfma, total_mfma, 1) + rocdl.sched_group_barrier(rocdl.mask_vmem_rd, num_a_loads, 2) + rocdl.sched_group_barrier(rocdl.mask_dswr, num_a_loads, 3) + rocdl.sched_barrier(0) + + return hot_loop_scheduler + + +def make_prefetch_scales( + *, + _use_hw_scale, + sa_rsrc, + sb_rsrc, + group_idx, + scale_n, + scale_k, + c_scale_k, + n_block_for_scale, + bx_m, + lane_mod_16, + m_in, + sb_per_tile, + m_repeat, + num_acc_n, + sa_group_off=None, +): + """Build the cross-tile E8M0 scale prefetch closure (gfx950 HW path). + + Returns `prefetch_scales(k_tile_idx_py)` that returns + `(sa_pf, sb_pf)` — outer index = sb (sb_per_tile), inner = + m_repeat / num_acc_n. Returns None on the gfx942 SW path (where + scales are loaded inside compute_tile instead). + + `sa_group_off` is None for the contig path (no addition emitted) + and `group_idx * c_scale_k * m_in` for the masked path. Using a + Python `is None` guard keeps the contig MLIR identical to the + pre-extraction code. + """ + + def prefetch_scales(k_tile_idx_py): + if not _use_hw_scale: + return None + sa_pf = [] + sb_pf = [] + # scale_b layout is [num_groups, scale_k, scale_n]. + sb_group_offset = group_idx * fx.Index(scale_k * scale_n) + for sb in range_constexpr(sb_per_tile): + kb = fx.Index(k_tile_idx_py * sb_per_tile + sb) + if sa_group_off is None: + sa_base_pf = kb * m_in + else: + sa_base_pf = sa_group_off + kb * m_in + + sa_sb = [] + for mi in range_constexpr(m_repeat): + sa_row = bx_m + (mi * 16) + lane_mod_16 + sa_idx = sa_base_pf + sa_row + sa_i8 = buffer_ops.buffer_load(sa_rsrc, sa_idx, vec_width=1, dtype=T.i8) + sa_e8m0 = ArithValue(sa_i8).extui(T.i32) + sa_sb.append(sa_e8m0) + sa_pf.append(sa_sb) + + sb_sb = [] + for ni in range_constexpr(num_acc_n): + sb_idx = sb_group_offset + kb * fx.Index(scale_n) + n_block_for_scale[ni] + sb_i8 = buffer_ops.buffer_load(sb_rsrc, sb_idx, vec_width=1, dtype=T.i8) + sb_i32 = ArithValue(sb_i8).extui(T.i32) + sb_e8m0 = rocdl.readfirstlane(T.i32, sb_i32) + sb_sb.append(sb_e8m0) + sb_pf.append(sb_sb) + return (sa_pf, sb_pf) + + return prefetch_scales + + +def make_compute_tile( + *, + _use_hw_scale, + _is_gfx950=False, + lds_load_packs_k64, + sa_rsrc, + sb_rsrc, + group_idx, + scale_n, + scale_k, + c_scale_k, + n_block_for_scale, + bx_m, + lane_mod_16, + lane_div_16, + m_in, + sb_per_tile, + m_repeat, + num_acc_n, + ku_per_sb, + col_offset_base_bytes, + mfma_res_ty, + acc_init, + sa_group_off=None, + group_m_start=None, + group_m_size=None, +): + """Build the per-K-tile compute closure. + + Returns `compute_tile(accs_in, k_tile_idx_py, lds_base, b_tile_in, + scales_pf, *, a0_prefetch=None)` which advances the accumulators by + one K-tile of MFMA work. `scales_pf` is the prefetched scales for + the gfx950 HW path; None for the gfx942 SW path (which loads scales + locally inside the closure). + + `sa_group_off` is None for the contig path (no addition emitted) + and `group_idx * c_scale_k * m_in` for the masked path; only the + gfx942 SW path uses it. + """ + + def compute_tile(accs_in, k_tile_idx_py, lds_base, b_tile_in, scales_pf, *, a0_prefetch=None): + current_accs = list(accs_in) + + for sb in range_constexpr(sb_per_tile): + kb = fx.Index(k_tile_idx_py * sb_per_tile + sb) + + s_a_vecs = [] + s_b_vals = [] + if not _use_hw_scale: + if group_m_start is not None: + # quantize_fp8_group(m_sizes=...) stores scale_a as per-group + # blocks: group g starts at M_start*scale_k and holds element + # (local_m, k_g) at local_m + k_g*M_g. This is NOT a global + # [scale_k, TotalM] transpose -- the two only coincide when + # there is one group or one K-block. + # sa_idx adds row_global (= M_start + local_m) below, so fold + # the -M_start into the base: M_start*(scale_k-1) + k_g*M_g. + sa_base = group_m_start * fx.Index(scale_k - 1) + kb * group_m_size + elif sa_group_off is None: + sa_base = kb * m_in + else: + sa_base = sa_group_off + kb * m_in + row_off_base = lane_div_16 * fx.Index(4) + for mi in range_constexpr(m_repeat): + s_a_row = [] + for ii in range_constexpr(4): + row_in_tile = (mi * 16) + row_off_base + fx.Index(ii) + row_global = bx_m + row_in_tile + sa_idx = sa_base + row_global + s_a_val = buffer_ops.buffer_load(sa_rsrc, sa_idx, vec_width=1, dtype=T.f32) + s_a_row.append(s_a_val) + s_a_vec4 = Vector.from_elements(s_a_row, fx.Float32) + s_a_vecs.append(s_a_vec4) + + # scale_b layout is [num_groups, scale_k, scale_n]: + # element (g, kb, n_blk) at g*scale_k*scale_n + kb*scale_n + n_blk. + sb_group_offset = group_idx * fx.Index(scale_k * scale_n) + for ni in range_constexpr(num_acc_n): + sb_idx = sb_group_offset + kb * fx.Index(scale_n) + n_block_for_scale[ni] + s_b_val = buffer_ops.buffer_load(sb_rsrc, sb_idx, vec_width=1, dtype=T.f32) + s_b_val = rocdl.readfirstlane(T.f32, s_b_val) + s_b_vals.append(s_b_val) + + if _use_hw_scale: + sa_pf, sb_pf = scales_pf + sa_e8m0_list = sa_pf[sb] + sb_e8m0_list = sb_pf[sb] + + ku0 = sb * ku_per_sb + ku1 = ku0 + 1 + b0_packs0, b0_packs1 = b_tile_in[ku0] + b1_packs0, b1_packs1 = b_tile_in[ku1] + col_base0 = col_offset_base_bytes + fx.Index(ku0 * 64) + col_base1 = col_offset_base_bytes + fx.Index(ku1 * 64) + + for mi in range_constexpr(m_repeat): + curr_row_a_lds = lane_mod_16 + (mi * 16) + if a0_prefetch is not None and sb == 0 and mi == 0: + a0, a1 = a0_prefetch + else: + a0, a1 = lds_load_packs_k64(curr_row_a_lds, col_base0, lds_base) + a2, a3 = lds_load_packs_k64(curr_row_a_lds, col_base1, lds_base) + a128 = pack_i64x4_to_i32x8(a0, a1, a2, a3) + + for ni in range_constexpr(num_acc_n): + b128 = pack_i64x4_to_i32x8( + b0_packs0[ni], + b0_packs1[ni], + b1_packs0[ni], + b1_packs1[ni], + ) + acc_idx = mi * num_acc_n + ni + current_accs[acc_idx] = rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [a128, b128, current_accs[acc_idx], 0, 0, 0, sa_e8m0_list[mi], 0, sb_e8m0_list[ni]], + ) + elif _is_gfx950: + # gfx950: use the wide 16x16x128 MFMA with a neutral E8M0 scale + # (0x7F7F7F7F = no-op hardware scaling), accumulate a whole + # scale-block into block_accs, then apply the FP32 scales in + # software once per scale-block. This avoids both the 4x-narrower + # 16x16x32 MFMA and the per-K-step VALU scale cost of the path + # below. + combined_scales = [] + for mi in range_constexpr(m_repeat): + mi_combined = [] + for ni in range_constexpr(num_acc_n): + s_b_bc = Vector.filled((4,), fx.Float32(s_b_vals[ni]), fx.Float32) + mi_combined.append(ArithValue(s_a_vecs[mi]) * ArithValue(s_b_bc)) + combined_scales.append(mi_combined) + + block_accs = [acc_init] * (num_acc_n * m_repeat) + ku0 = sb * ku_per_sb + ku1 = ku0 + 1 + b0_packs0, b0_packs1 = b_tile_in[ku0] + b1_packs0, b1_packs1 = b_tile_in[ku1] + col_base0 = col_offset_base_bytes + fx.Index(ku0 * 64) + col_base1 = col_offset_base_bytes + fx.Index(ku1 * 64) + + for mi in range_constexpr(m_repeat): + curr_row_a_lds = lane_mod_16 + (mi * 16) + if a0_prefetch is not None and sb == 0 and mi == 0: + a0, a1 = a0_prefetch + else: + a0, a1 = lds_load_packs_k64(curr_row_a_lds, col_base0, lds_base) + a2, a3 = lds_load_packs_k64(curr_row_a_lds, col_base1, lds_base) + a128 = pack_i64x4_to_i32x8(a0, a1, a2, a3) + + for ni in range_constexpr(num_acc_n): + b128 = pack_i64x4_to_i32x8( + b0_packs0[ni], b0_packs1[ni], b1_packs0[ni], b1_packs1[ni] + ) + acc_idx = mi * num_acc_n + ni + block_accs[acc_idx] = rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [a128, b128, block_accs[acc_idx], 0, 0, 0, 0x7F7F7F7F, 0, 0x7F7F7F7F], + ) + + for mi in range_constexpr(m_repeat): + for ni in range_constexpr(num_acc_n): + acc_idx = mi * num_acc_n + ni + current_accs[acc_idx] = math_dialect.fma( + block_accs[acc_idx], + combined_scales[mi][ni], + current_accs[acc_idx], + ) + else: + for ku_local in range_constexpr(ku_per_sb): + ku = sb * ku_per_sb + ku_local + k_offset_bytes = ku * 64 + b_packs0, b_packs1 = b_tile_in[ku] + + for mi in range_constexpr(m_repeat): + if a0_prefetch is not None and sb == 0 and ku_local == 0 and mi == 0: + a0, a1 = a0_prefetch + else: + row_a_lds = lane_mod_16 + (mi * 16) + col_a_base_bytes = lane_div_16 * fx.Index(16) + fx.Index(k_offset_bytes) + a0, a1 = lds_load_packs_k64(row_a_lds, col_a_base_bytes, lds_base) + + for ni in range_constexpr(num_acc_n): + acc_idx = mi * num_acc_n + ni + + mfma_fn = rocdl.mfma_f32_16x16x32_fp8_fp8 + mfma_mid = mfma_fn(T.f32x4, [a0, b_packs0[ni], acc_init, 0, 0, 0]) + mfma_result = mfma_fn(T.f32x4, [a1, b_packs1[ni], mfma_mid, 0, 0, 0]) + + s_a_v4 = s_a_vecs[mi] + s_b_bc = Vector.filled((4,), fx.Float32(s_b_vals[ni]), fx.Float32) + scaled = ArithValue(mfma_result) * ArithValue(s_a_v4) + current_accs[acc_idx] = math_dialect.fma(scaled, s_b_bc, current_accs[acc_idx]) + + return current_accs + + return compute_tile + + +def make_kloop_plain( + *, + num_k_tiles, + tile_k, + prefetch_a_tile, + store_a_tile_to_lds, + prefetch_b_tile, + store_b_tile_to_lds, + load_b_tile_from_lds, + prefetch_scales, + compute_tile, + lds_base_pong, + lds_base_b, +): + """Single-LDS-buffer K-loop for the plain-B kernel. + + Plain B is staged HBM->LDS->registers each K-tile, alongside A. A and B each + get one LDS buffer rather than the ping-pong pair `make_pingpong_kloop` uses, + which keeps the LDS budget within reach at wide tile_n but costs two barriers + per K-tile: one after the stores, and one at the end of the iteration before + the buffers are overwritten. Global-load latency is still overlapped, by + issuing the next tile's HBM loads into registers before computing the current + one. `lds_base_pong` is reused as the single A buffer base. + """ + + def run_kloop(accs): + if num_k_tiles == 0: + return accs + + a_regs = prefetch_a_tile(0) + b_regs = prefetch_b_tile(0) + for kt in range_constexpr(num_k_tiles): + # Publish this tile's A/B (already in VGPRs) to LDS. + store_a_tile_to_lds(a_regs, lds_base_pong) + store_b_tile_to_lds(b_regs, lds_base_b) + scales_pf = prefetch_scales(kt) + gpu.barrier() + + # Issue next tile's HBM loads NOW so they overlap the compute below. + if kt + 1 < num_k_tiles: + a_regs = prefetch_a_tile(kt + 1) + b_regs = prefetch_b_tile(kt + 1) + + # Read B fragment from LDS, compute the tile. + b_tile = load_b_tile_from_lds(lds_base_b) + accs = compute_tile(accs, kt, lds_base_pong, b_tile, scales_pf) + # Barrier before next iter overwrites the shared A/B buffers. + gpu.barrier() + return accs + + return run_kloop + + +def make_pingpong_kloop( + *, + num_k_tiles, + tile_k, + prefetch_a_tile, + store_a_tile_to_lds, + load_b_tile, + prefetch_scales, + compute_tile, + hot_loop_scheduler, + lds_load_packs_k64, + lds_base_pong, + lds_base_ping, + row_a_lds_base, + col_offset_base_bytes, +): + """Build the ping-pong K-loop driver. + + Returns `run_kloop(accs)` which advances `accs` through all + K-tiles using the prologue + alternating ping/pong stages. + Loop body is byte-identical between contig and masked, so this + factory has no offset parameters. + """ + + def run_kloop(accs): + # Prologue: prefetch first A tile into VGPRs, store to LDS, load B + scales + a_regs0 = prefetch_a_tile(0) + store_a_tile_to_lds(a_regs0, lds_base_pong) + b_tile_pong = load_b_tile(fx.Index(0)) + scales_pong_pf = prefetch_scales(0) + gpu.barrier() + + # Prefetch first A pack from pong (hides LDS latency behind upcoming VMEM) + a0_prefetch_pong = lds_load_packs_k64(row_a_lds_base, col_offset_base_bytes, lds_base_pong) + + for k_pair in range_constexpr(0, num_k_tiles, 2): + # Prefetch the next scales before the B-tile VMEM so the scale-load + # latency hides behind it; then the A+B registers. + if k_pair + 1 < num_k_tiles: + scales_ping_pf = prefetch_scales(k_pair + 1) + a_regs_ping = prefetch_a_tile(k_pair + 1) + b_tile_ping = load_b_tile(fx.Index((k_pair + 1) * tile_k)) + + # Compute current tile from pong LDS + accs = compute_tile(accs, k_pair, lds_base_pong, b_tile_pong, scales_pong_pf, a0_prefetch=a0_prefetch_pong) + a0_prefetch_pong = None + + # Store next A to LDS (ds_write after compute, overlaps with trailing MFMAs) + if k_pair + 1 < num_k_tiles: + store_a_tile_to_lds(a_regs_ping, lds_base_ping) + hot_loop_scheduler() + gpu.barrier() + + if k_pair + 1 < num_k_tiles: + # Prefetch first A pack from ping + a0_prefetch_ping = lds_load_packs_k64(row_a_lds_base, col_offset_base_bytes, lds_base_ping) + + # Prefetch next scales + A+B + if k_pair + 2 < num_k_tiles: + scales_pong_pf = prefetch_scales(k_pair + 2) + a_regs_pong = prefetch_a_tile(k_pair + 2) + b_tile_pong = load_b_tile(fx.Index((k_pair + 2) * tile_k)) + + # Compute current tile from ping LDS + accs = compute_tile( + accs, k_pair + 1, lds_base_ping, b_tile_ping, scales_ping_pf, a0_prefetch=a0_prefetch_ping + ) + a0_prefetch_ping = None + + # Store next A to LDS + if k_pair + 2 < num_k_tiles: + store_a_tile_to_lds(a_regs_pong, lds_base_pong) + hot_loop_scheduler() + gpu.barrier() + + # Prefetch first A pack from pong for next iteration + if k_pair + 2 < num_k_tiles: + a0_prefetch_pong = lds_load_packs_k64(row_a_lds_base, col_offset_base_bytes, lds_base_pong) + + return accs + + return run_kloop + + +def make_epilogue_writers( + *, + accs, + d_rsrc, + out_mlir, + e_vec, + c_n, + d_group_off=None, +): + """Build the CShuffle-epilogue writer closures. + + Returns `(write_row_to_lds, store_pair)` to be passed to + `mfma_epilog`. `d_group_off` is None for the contig path (no + addition emitted) and `group_idx * m_in * n_in` for the masked + path. Using a Python `is None` guard keeps the contig MLIR + identical to the pre-extraction code. + """ + + def write_row_to_lds( + *, + mi, + ii, + row_in_tile, + row, + row_base_lds, + col_base_local, + num_acc_n, + lds_out, + ): + for ni in range_constexpr(num_acc_n): + col_local = col_base_local + (ni * 16) + acc_idx = mi * num_acc_n + ni + acc = accs[acc_idx] + val = Vector(acc)[ii] + v_out = arith.trunc_f(out_mlir(), val) + lds_idx = row_base_lds + col_local + v1 = Vector.from_elements([v_out]) + v1.store(lds_out, [lds_idx], alignment=2) + + def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): + if d_group_off is None: + idx_out = row * c_n + col_g0 + else: + idx_out = d_group_off + row * c_n + col_g0 + byte_off = idx_out * 2 + if e_vec == 4: + frag_i32x2 = Vector(frag).bitcast(fx.Int32) + buffer_ops.buffer_store(frag_i32x2, d_rsrc, byte_off, offset_is_bytes=True) + else: + frag_i32x1 = Vector(frag).bitcast(fx.Int32) + frag_i32 = frag_i32x1[0] + buffer_ops.buffer_store(frag_i32, d_rsrc, byte_off, offset_is_bytes=True) + + return write_row_to_lds, store_pair + + +MfmaTilingConstants = namedtuple( + "MfmaTilingConstants", + ["m_repeat", "num_waves", "n_per_wave", "num_acc_n", "num_accs"], +) + + +def compute_mfma_tiling(*, tile_m, tile_n): + """Pure-Python derivation of the MFMA tiling constants. + + Returns an `MfmaTilingConstants` namedtuple with `m_repeat`, + `num_waves`, `n_per_wave`, `num_acc_n`, `num_accs`. Emits no MLIR. + """ + m_repeat = tile_m // 16 # 8 for tile_m=128 + num_waves = 4 + n_per_wave = tile_n // num_waves # 32 for tile_n=128 + num_acc_n = n_per_wave // 16 # 2 for n_per_wave=32 + num_accs = m_repeat * num_acc_n + return MfmaTilingConstants( + m_repeat=m_repeat, + num_waves=num_waves, + n_per_wave=n_per_wave, + num_acc_n=num_acc_n, + num_accs=num_accs, + ) + + +def init_accumulators(num_accs): + """Emit the FP32 zero-vector accumulator constant and replicate it + for all `num_accs` MFMA result slots. Returns `(acc_init, accs)`.""" + acc_init = arith.constant_vector(0.0, T.f32x4) + accs = [acc_init] * num_accs + return acc_init, accs + + +NBlockCoords = namedtuple( + "NBlockCoords", + ["n_tile_base", "n_block_for_scale", "layout_b", "n_blk_list", "n_intra_list", "c_scale_k"], +) + + +def make_n_block_coords( + *, + wave_id, + by_n, + group_idx, + num_groups_in, + n_in, + k_in, + lane_mod_16, + kpack_bytes, + elem_bytes, + scale_block_n, + scale_k, + n_per_wave, + num_acc_n, +): + """Compute per-wave N-tile base, scale_b N-block indices, the + preshuffle B layout, and the per-MFMA (n_blk, n_intra) coordinate + lists for all groups concatenated along N. + + Byte-identical between contig and masked. Returns an `NBlockCoords` + namedtuple matching the original local variable names so the caller + can keep referring to them unchanged. + """ + wave_mod_4 = wave_id % fx.Index(4) + n_tile_base = wave_mod_4 * fx.Index(n_per_wave) + + c_scale_block_n = fx.Index(scale_block_n) + c_scale_k = fx.Index(scale_k) + n_block_for_scale = [] + for ni in range_constexpr(num_acc_n): + col_base = by_n + n_tile_base + (ni * 16) + n_blk = col_base // c_scale_block_n + n_block_for_scale.append(n_blk) + + c_n_total = num_groups_in * n_in + b_layout = make_preshuffle_b_layout( + arith, + c_n=c_n_total, + c_k=k_in, + kpack_bytes=kpack_bytes, + elem_bytes=elem_bytes, + ) + layout_b = b_layout.layout_b + + c_n0 = c_n_total // fx.Index(16) + c_n0_i32 = arith.index_cast(T.i32, c_n0) + layout_n_blk_intra = fx.make_layout((c_n0_i32, 16), stride=(16, 1)) + n_blk_list = [] + n_intra_list = [] + group_n_off = group_idx * n_in + for ni in range_constexpr(num_acc_n): + col_global = group_n_off + by_n + n_tile_base + (ni * 16) + lane_mod_16 + coord_ni = fx.idx2crd(fx.Int32(col_global), layout_n_blk_intra) + n_blk_list.append(fx.get(coord_ni, 0)) + n_intra_list.append(fx.get(coord_ni, 1)) + + return NBlockCoords( + n_tile_base=n_tile_base, + n_block_for_scale=n_block_for_scale, + layout_b=layout_b, + n_blk_list=n_blk_list, + n_intra_list=n_intra_list, + c_scale_k=c_scale_k, + ) diff --git a/mslk/flydsl/kernels/gemm/grouped_gemm_blockscale_contiguous.py b/mslk/flydsl/kernels/gemm/grouped_gemm_blockscale_contiguous.py new file mode 100644 index 00000000..34410d3e --- /dev/null +++ b/mslk/flydsl/kernels/gemm/grouped_gemm_blockscale_contiguous.py @@ -0,0 +1,617 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Contiguous Grouped FP8 GEMM kernel with block scaling. + +Groups are concatenated along M with arbitrary (not tile-aligned) per-group row +counts, and the output is compact [M_total, N]. Each output M-tile belongs to +exactly one group, so a tile never spans a group boundary. The kernel resolves +the owning group for its M-tile from m_sizes, and rows at or beyond that group's +end are the partial-tile tail and are masked out of the store. + +Scales are FP32 (software scaling) on all architectures. + +Tensors: + - A: [M_total, K] FP8 - concatenated rows from all groups + - scale_a: FP32 per-token, per-128K scales, laid out as per-group blocks (as + written by quantize_fp8_group with m_sizes): group g's block begins at + m_start * scale_k, and within it element (local_m, k_block) sits at + local_m + k_block * M_g. This is not a global [scale_k, M_total] transpose. + - B: [num_groups, N, K] FP8 - one weight matrix per group, preshuffled + - scale_b: [num_groups, scale_k, scale_n] FP32 - per-block scales + - m_sizes: [num_groups] INT64 - rows per group (sum to M_total) + - D: [M_total, N] BF16 - output + +Block scaling granularity: + - A: (1, 128) - per-token, per-128-K-elements + - B: (128, 128) - per-128-N, per-128-K block +""" + +import functools + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl, vector +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr +from mslk.flydsl.kernels.gemm.grouped_gemm_blockscale_common import ( + compute_compile_constants, + compute_mfma_tiling, + init_accumulators, + make_a_tile_loaders, + make_b_loader, + make_b_tile_loaders, + make_compute_tile, + make_epilogue_writers, + make_hot_loop_scheduler, + make_kloop_plain, + make_lds_b_loader, + make_lds_loader, + make_n_block_coords, + make_pingpong_kloop, + make_plain_b_tile, + make_prefetch_scales, + out_mlir_for, + setup_lds_allocation, + setup_lds_allocation_plain, + validate_lds_budget_plain, + validate_lds_budget_preshuffle, + validate_params, +) +from mslk.flydsl.kernels.mma.mfma_epilogues import mfma_epilog + + +@functools.lru_cache(maxsize=128) +def compile_grouped_gemm_blockscale_contiguous( + *, + n: int, + k: int, + num_groups: int, + tile_m: int = 128, + tile_n: int = 128, + tile_k: int = 128, + scale_block_k: int = 128, + scale_block_n: int = 128, + out_dtype: str = "bf16", + waves_per_eu: int | None = None, + b_preshuffled: bool = True, +): + """Compile grouped FP8 GEMM kernel and return the JIT launcher. + + Args: + n: N dimension (output columns per group) + k: K dimension (reduction dimension) + num_groups: Number of groups (experts) + tile_m: M tile size (default 128) + tile_n: N tile size (default 128) + tile_k: K tile size (default 128) + scale_block_k: K-dimension scale block size (default 128) + scale_block_n: N-dimension scale block size (default 128) + out_dtype: Output data type ("bf16" or "f16") + b_preshuffled: When True (default) B is expected pre-swizzled into the + MFMA layout and loaded HBM->registers (no B LDS). When False, B is + plain row-major [num_groups, N, K] and is staged HBM->LDS->registers + like A. The two paths share the entire kernel body (tile-map group + dispatch, FP32 block scaling, wide-MFMA, CShuffle epilogue); only the + B load stage and its LDS allocation differ. + + Returns: + JIT launcher function. + """ + gpu_arch = get_hip_arch() + # This FP8 kernel always uses the FP32 software-scaling path; the shared + # helpers' hardware E8M0 microscaling path is not used here. + _use_hw_scale = False + # On gfx950 the SW path still uses the wide 16x16x128 MFMA with a neutral + # E8M0 scale (no-op HW scaling) and applies FP32 scales in software; gfx942 + # lacks that instruction and falls back to the narrow 16x16x32 path. + _is_gfx950 = str(gpu_arch).startswith("gfx95") + + _sym = "smem_grouped_gemm" if b_preshuffled else "smem_grouped_gemm_plain" + allocator = SmemAllocator(None, arch=gpu_arch, global_sym_name=_sym) + + validate_params( + n=n, + k=k, + tile_n=tile_n, + tile_k=tile_k, + scale_block_k=scale_block_k, + scale_block_n=scale_block_n, + out_dtype=out_dtype, + ) + # Check the LDS budget before tracing: the compiler treats an overflow as a + # hard error that kills the process, which an autotuner cannot skip. Capacity + # is arch-dependent (64 KiB gfx942, 160 KiB gfx950). + if b_preshuffled: + # Preshuffled B goes HBM->registers; only A ping-pong / epilogue use LDS. + validate_lds_budget_preshuffle( + tile_m=tile_m, tile_n=tile_n, tile_k=tile_k, arch=gpu_arch + ) + else: + # Plain B needs its own LDS buffer alongside A. + validate_lds_budget_plain( + tile_m=tile_m, tile_n=tile_n, tile_k=tile_k, b_pingpong=False, arch=gpu_arch + ) + out_mlir = out_mlir_for(out_dtype) + + _c = compute_compile_constants( + n=n, + k=k, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + scale_block_k=scale_block_k, + scale_block_n=scale_block_n, + ) + total_threads = _c.total_threads + elem_bytes = _c.elem_bytes + num_k_tiles = _c.num_k_tiles + scale_k = _c.scale_k + scale_n = _c.scale_n + sb_per_tile = _c.sb_per_tile + k_unroll = _c.k_unroll + kpack_bytes = _c.kpack_bytes + tile_k_bytes = _c.tile_k_bytes + tile_k_dwords = _c.tile_k_dwords + chunk_i32_a = _c.chunk_i32_a + num_a_loads = _c.num_a_loads + chunk_i32_b = _c.chunk_i32_b + num_b_loads = _c.num_b_loads + + if b_preshuffled: + lds_alloc_offset, lds_tile_elems = setup_lds_allocation( + allocator=allocator, + tile_m=tile_m, + tile_k=tile_k, + tile_n=tile_n, + elem_bytes=elem_bytes, + ) + lds_b_offset_elems = None + else: + lds_alloc_offset, lds_tile_elems, lds_b_offset_elems = setup_lds_allocation_plain( + allocator=allocator, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + elem_bytes=elem_bytes, + b_pingpong=False, + ) + + # Module name for caching + _variant = "contiguous_pingpong" if b_preshuffled else "plain" + module_name = ( + f"grouped_gemm_blockscale_{_variant}_{out_dtype}" + f"_n{n}_k{k}_g{num_groups}" + f"_t{tile_m}x{tile_n}x{tile_k}" + ).replace("-", "_") + + @flyc.kernel(name=module_name) + def grouped_gemm_blockscale_contiguous_kernel( + arg_d: fx.Tensor, + arg_a: fx.Tensor, + arg_b: fx.Tensor, + arg_scale_a: fx.Tensor, + arg_scale_b: fx.Tensor, + arg_m_sizes: fx.Tensor, + i32_m: fx.Int32, + i32_n: fx.Int32, + i32_k: fx.Int32, + i32_num_groups: fx.Int32, + ): + # Convert runtime parameters to index type + m_in = fx.Index(i32_m) + n_in = fx.Index(i32_n) + k_in = fx.Index(i32_k) + num_groups_in = fx.Index(i32_num_groups) + + # Thread and block IDs + tx = gpu.thread_id("x") + by = gpu.block_id("x") # N-block index + bx = gpu.block_id("y") # M-tile index (into the per-tile dispatch map) + + # N-block position; bx_m (global row base) is loaded from the tile map below. + by_n = by * fx.Index(tile_n) + + # Wave/lane decomposition (256 threads = 4 waves x 64 lanes) + layout_wave_lane = fx.make_layout((4, 64), stride=(64, 1)) + coord_wave_lane = fx.idx2crd(fx.Int32(tx), layout_wave_lane) + wave_id = fx.get(coord_wave_lane, 0) + lane_id = fx.get(coord_wave_lane, 1) + + # Lane decomposition for MFMA (lane_id -> lane_div_16, lane_mod_16) + layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane_id), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # LDS setup: ping-pong A buffers (preshuffle) or A ping-pong + single B + # buffer (plain). B LDS is only needed for the plain path. + base_ptr = allocator.get_base() + lds_a = SmemPtr(base_ptr, lds_alloc_offset, T.f8, shape=(2 * tile_m * tile_k,)).get() + lds_stride = tile_k + layout_lds = fx.make_layout((tile_m, tile_k), stride=(lds_stride, 1)) + lds_base_pong = fx.Index(0) + lds_base_ping = fx.Index(lds_tile_elems) + + if const_expr(not b_preshuffled): + # Plain-B LDS buffer, placed just past the A ping-pong region. + lds_b = SmemPtr( + base_ptr, lds_alloc_offset, T.f8, shape=((lds_b_offset_elems + tile_n * tile_k),) + ).get() + layout_lds_b = fx.make_layout((tile_n, tile_k), stride=(tile_k, 1)) + lds_base_b = fx.Index(lds_b_offset_elems) + + # CShuffle epilogue LDS (aliased from same base, out-dtype element type) + lds_out = SmemPtr(base_ptr, lds_alloc_offset, out_mlir(), shape=(tile_m * tile_n,)).get() + + # Buffer resources + a_nbytes = m_in * k_in + a_rsrc = buffer_ops.create_buffer_resource(arg_a, max_size=False, num_records_bytes=a_nbytes) + + b_nbytes = num_groups_in * n_in * k_in + b_rsrc = buffer_ops.create_buffer_resource(arg_b, max_size=False, num_records_bytes=b_nbytes) + + d_nbytes = m_in * n_in * fx.Index(2) # bf16/f16 = 2 bytes + d_rsrc = buffer_ops.create_buffer_resource(arg_d, max_size=False, num_records_bytes=d_nbytes) + + # Scale buffers — gfx950 HW E8M0 path consumes int8 (one byte/scale, + # pre-packed on host); gfx942 SW path consumes f32. + scale_byte_size = 1 if _use_hw_scale else 4 + + # scale_a: [scale_k, M] - transposed layout + sa_nbytes = fx.Index(scale_k) * m_in * fx.Index(scale_byte_size) + sa_rsrc = buffer_ops.create_buffer_resource(arg_scale_a, max_size=False, num_records_bytes=sa_nbytes) + + # scale_b: [num_groups, scale_n, scale_k] + sb_nbytes = num_groups_in * fx.Index(scale_n * scale_k * scale_byte_size) + sb_rsrc = buffer_ops.create_buffer_resource(arg_scale_b, max_size=False, num_records_bytes=sb_nbytes) + + # Resolve which group owns this flat M-tile id (bx) from m_sizes. Doing it + # here rather than from a host-built dispatch map keeps the launch free of + # helper kernels, which matters under CUDA-graph capture where each one is + # replayed per call. num_groups is a compile-time constant, so the loop + # unrolls to a few scalar ops. acc_m/acc_t are the running m_start and + # tile_start prefixes; tiles beyond the real tile count (the grid extent is + # an upper bound) match no group and stay marked -1. + ms_rsrc = buffer_ops.create_buffer_resource( + arg_m_sizes, max_size=False, num_records_bytes=num_groups_in * fx.Index(8) + ) + + def _i32(v): # raw i32 constant (arith.* requires unwrapped MLIR values) + return arith.constant(int(v), type=T.i32) + + bx_i32 = arith.index_cast(T.i32, bx) + tile_m_c = _i32(tile_m) + tile_m_bump = _i32(tile_m - 1) + acc_m = _i32(0) # cumulative rows before group g (m_starts[g]) + acc_t = _i32(0) # cumulative tiles before group g (tile_starts[g]) + group_id_i32 = _i32(-1) + row_start_i32 = _i32(0) + row_limit_i32 = _i32(0) + group_m_start_i32 = _i32(0) # first global row of the owning group + group_m_size_i32 = _i32(0) # row count of the owning group + for _g in range_constexpr(num_groups): + # m_sizes is int64; read the low dword of element _g (index _g*2 in + # dwords). Row counts fit in int32, so the high dword is always zero + # and no host-side narrowing kernel is needed. + m_g = buffer_ops.buffer_load(ms_rsrc, _g * 2, vec_width=1, dtype=T.i32) + tiles_g = arith.divui(arith.addi(m_g, tile_m_bump), tile_m_c) + acc_t_next = arith.addi(acc_t, tiles_g) + in_grp = arith.andi( + arith.cmpi(arith.CmpIPredicate.sge, bx_i32, acc_t), + arith.cmpi(arith.CmpIPredicate.slt, bx_i32, acc_t_next), + ) + rs = arith.addi(acc_m, arith.muli(arith.subi(bx_i32, acc_t), tile_m_c)) + rl = arith.addi(acc_m, m_g) + group_id_i32 = arith.select(in_grp, _i32(_g), group_id_i32) + row_start_i32 = arith.select(in_grp, rs, row_start_i32) + row_limit_i32 = arith.select(in_grp, rl, row_limit_i32) + group_m_start_i32 = arith.select(in_grp, acc_m, group_m_start_i32) + group_m_size_i32 = arith.select(in_grp, m_g, group_m_size_i32) + acc_m = arith.addi(acc_m, m_g) + acc_t = acc_t_next + + is_valid = arith.cmpi(arith.CmpIPredicate.sge, group_id_i32, _i32(0)) + + # Early exit for surplus/no-op tiles. + if is_valid: + group_idx = fx.Index(group_id_i32) + + # Global row base of this tile and the exclusive row end of its group + # (the group end masks the partial-tile tail in the epilogue store). + bx_m = fx.Index(row_start_i32) + + _t = compute_mfma_tiling(tile_m=tile_m, tile_n=tile_n) + m_repeat = _t.m_repeat + n_per_wave = _t.n_per_wave + num_acc_n = _t.num_acc_n + + acc_init, accs = init_accumulators(_t.num_accs) + + _nb = make_n_block_coords( + wave_id=wave_id, + by_n=by_n, + group_idx=group_idx, + num_groups_in=num_groups_in, + n_in=n_in, + k_in=k_in, + lane_mod_16=lane_mod_16, + kpack_bytes=kpack_bytes, + elem_bytes=elem_bytes, + scale_block_n=scale_block_n, + scale_k=scale_k, + n_per_wave=n_per_wave, + num_acc_n=num_acc_n, + ) + n_tile_base = _nb.n_tile_base + n_block_for_scale = _nb.n_block_for_scale + layout_b = _nb.layout_b + n_blk_list = _nb.n_blk_list + n_intra_list = _nb.n_intra_list + c_scale_k = _nb.c_scale_k + + prefetch_a_tile, store_a_tile_to_lds, a_row_local, a_col_local_i32, k_blocks16 = make_a_tile_loaders( + a_rsrc=a_rsrc, + lds_a=lds_a, + layout_lds=layout_lds, + bx_m=bx_m, + tx=tx, + tile_m=tile_m, + tile_k=tile_k, + tile_k_bytes=tile_k_bytes, + tile_k_dwords=tile_k_dwords, + chunk_i32_a=chunk_i32_a, + num_a_loads=num_a_loads, + total_threads=total_threads, + elem_bytes=elem_bytes, + k_in=k_in, + ) + + lds_load_packs_k64 = make_lds_loader( + lds_a=lds_a, + layout_lds=layout_lds, + k_blocks16=k_blocks16, + ) + + # Base coordinates for A0 prefetch (mi=0, ku=0) + row_a_lds_base = lane_mod_16 # mi=0 + col_offset_base_bytes = lane_div_16 * fx.Index(16) # ku=0 + + # ---- B load path: preshuffled (HBM->registers) vs plain (HBM->LDS->registers) ---- + if const_expr(b_preshuffled): + load_b_tile = make_b_loader( + arg_b=arg_b, + b_rsrc=b_rsrc, + layout_b=layout_b, + n_blk_list=n_blk_list, + n_intra_list=n_intra_list, + lane_div_16=lane_div_16, + kpack_bytes=kpack_bytes, + elem_bytes=elem_bytes, + k_unroll=k_unroll, + num_acc_n=num_acc_n, + ) + else: + prefetch_b_tile, store_b_tile_to_lds, _b_row_local, _b_col_local_i32, k_blocks16_b = ( + make_b_tile_loaders( + b_rsrc=b_rsrc, + lds_b=lds_b, + layout_lds_b=layout_lds_b, + by_n=by_n, + group_idx=group_idx, + tx=tx, + tile_n=tile_n, + tile_k=tile_k, + tile_k_bytes=tile_k_bytes, + tile_k_dwords=tile_k_dwords, + chunk_i32_b=chunk_i32_b, + num_b_loads=num_b_loads, + total_threads=total_threads, + elem_bytes=elem_bytes, + n_in=n_in, + k_in=k_in, + ) + ) + lds_load_b_packs_k64 = make_lds_b_loader( + lds_b=lds_b, + layout_lds_b=layout_lds_b, + k_blocks16_b=k_blocks16_b, + ) + load_b_tile_from_lds = make_plain_b_tile( + lds_load_b_packs_k64=lds_load_b_packs_k64, + lane_mod_16=lane_mod_16, + n_tile_base=n_tile_base, + col_offset_base_bytes=col_offset_base_bytes, + k_unroll=k_unroll, + num_acc_n=num_acc_n, + ) + + mfma_res_ty = T.f32x4 + + ku_per_sb = scale_block_k // 64 + rocdl.sched_barrier(0) + + if const_expr(b_preshuffled): + hot_loop_scheduler = make_hot_loop_scheduler( + _use_hw_scale=_use_hw_scale, + sb_per_tile=sb_per_tile, + m_repeat=m_repeat, + num_acc_n=num_acc_n, + k_unroll=k_unroll, + num_a_loads=num_a_loads, + ku_per_sb=ku_per_sb, + ) + + prefetch_scales = make_prefetch_scales( + _use_hw_scale=_use_hw_scale, + sa_rsrc=sa_rsrc, + sb_rsrc=sb_rsrc, + group_idx=group_idx, + scale_n=scale_n, + scale_k=scale_k, + c_scale_k=c_scale_k, + n_block_for_scale=n_block_for_scale, + bx_m=bx_m, + lane_mod_16=lane_mod_16, + m_in=m_in, + sb_per_tile=sb_per_tile, + m_repeat=m_repeat, + num_acc_n=num_acc_n, + ) + + compute_tile = make_compute_tile( + _use_hw_scale=_use_hw_scale, + _is_gfx950=_is_gfx950, + lds_load_packs_k64=lds_load_packs_k64, + sa_rsrc=sa_rsrc, + sb_rsrc=sb_rsrc, + group_idx=group_idx, + scale_n=scale_n, + scale_k=scale_k, + c_scale_k=c_scale_k, + n_block_for_scale=n_block_for_scale, + bx_m=bx_m, + lane_mod_16=lane_mod_16, + lane_div_16=lane_div_16, + m_in=m_in, + sb_per_tile=sb_per_tile, + m_repeat=m_repeat, + num_acc_n=num_acc_n, + ku_per_sb=ku_per_sb, + col_offset_base_bytes=col_offset_base_bytes, + mfma_res_ty=mfma_res_ty, + acc_init=acc_init, + group_m_start=fx.Index(group_m_start_i32), + group_m_size=fx.Index(group_m_size_i32), + ) + + if const_expr(b_preshuffled): + run_kloop = make_pingpong_kloop( + num_k_tiles=num_k_tiles, + tile_k=tile_k, + prefetch_a_tile=prefetch_a_tile, + store_a_tile_to_lds=store_a_tile_to_lds, + load_b_tile=load_b_tile, + prefetch_scales=prefetch_scales, + compute_tile=compute_tile, + hot_loop_scheduler=hot_loop_scheduler, + lds_load_packs_k64=lds_load_packs_k64, + lds_base_pong=lds_base_pong, + lds_base_ping=lds_base_ping, + row_a_lds_base=row_a_lds_base, + col_offset_base_bytes=col_offset_base_bytes, + ) + else: + run_kloop = make_kloop_plain( + num_k_tiles=num_k_tiles, + tile_k=tile_k, + prefetch_a_tile=prefetch_a_tile, + store_a_tile_to_lds=store_a_tile_to_lds, + prefetch_b_tile=prefetch_b_tile, + store_b_tile_to_lds=store_b_tile_to_lds, + load_b_tile_from_lds=load_b_tile_from_lds, + prefetch_scales=prefetch_scales, + compute_tile=compute_tile, + lds_base_pong=lds_base_pong, + lds_base_b=lds_base_b, + ) + accs = run_kloop(accs) + + # ===== Epilogue: CShuffle vectorized stores ===== + c_n = n_in + e_vec = 4 if (tile_n % (32 * 4)) == 0 else 2 + + write_row_to_lds, store_pair = make_epilogue_writers( + accs=accs, + d_rsrc=d_rsrc, + out_mlir=out_mlir, + e_vec=e_vec, + c_n=c_n, + ) + + # Mask the partial-tile tail: skip stores for global rows at or beyond + # the owning group's end. Returning (ctx, pred) lets the epilogue skip + # the whole N-store loop for out-of-group rows. + def precompute_row(*, row_local, row): + row_i32 = arith.index_cast(T.i32, row) + row_valid = arith.cmpi(arith.CmpIPredicate.ult, row_i32, row_limit_i32) + return (None, row_valid) + + mfma_epilog( + use_cshuffle=True, + arith=arith, + vector=vector, + gpu=gpu, + scf=scf, + range_constexpr=range_constexpr, + tile_m=tile_m, + tile_n=tile_n, + e_vec=e_vec, + m_repeat=m_repeat, + num_acc_n=num_acc_n, + tx=tx, + lane_div_16=lane_div_16, + lane_mod_16=lane_mod_16, + bx_m=bx_m, + by_n=by_n, + n_tile_base=n_tile_base, + lds_out=lds_out, + frag_elem_type=out_mlir(), + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) + + # ===== JIT Launcher ===== + @flyc.jit + def launch_grouped_gemm_blockscale_contiguous( + arg_d: fx.Tensor, + arg_a: fx.Tensor, + arg_b: fx.Tensor, + arg_scale_a: fx.Tensor, + arg_scale_b: fx.Tensor, + arg_m_sizes: fx.Tensor, + i32_m: fx.Int32, + i32_n: fx.Int32, + i32_k: fx.Int32, + i32_num_groups: fx.Int32, + i32_num_m_tiles: fx.Int32, + stream: fx.Stream, + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + # Grid dimensions. The M axis enumerates output M-tiles; its extent is a + # host-known upper bound on the tile count, and tiles past the real count + # match no group and exit early. + n_in = fx.Index(i32_n) + gx = n_in // fx.Index(tile_n) # N-blocks + gy = fx.Index(i32_num_m_tiles) # M-tiles + + launcher = grouped_gemm_blockscale_contiguous_kernel( + arg_d, + arg_a, + arg_b, + arg_scale_a, + arg_scale_b, + arg_m_sizes, + i32_m, + i32_n, + i32_k, + i32_num_groups, + ) + if waves_per_eu is not None: + _wpe = int(waves_per_eu) + if _wpe >= 1: + for op in ctx.gpu_module_body.operations: + if hasattr(op, "attributes") and op.OPERATION_NAME == "gpu.func": + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe) + launcher.launch(grid=(gx, gy, 1), block=(total_threads, 1, 1), stream=stream) + + return launch_grouped_gemm_blockscale_contiguous diff --git a/mslk/flydsl/kernels/mma/__init__.py b/mslk/flydsl/kernels/mma/__init__.py new file mode 100644 index 00000000..581f84e4 --- /dev/null +++ b/mslk/flydsl/kernels/mma/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict diff --git a/mslk/flydsl/kernels/mma/mfma_epilogues.py b/mslk/flydsl/kernels/mma/mfma_epilogues.py new file mode 100644 index 00000000..b71d2313 --- /dev/null +++ b/mslk/flydsl/kernels/mma/mfma_epilogues.py @@ -0,0 +1,449 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Reusable epilogue helpers for MFMA 16x16-based kernels. + +This module provides: + +- `mfma_epilog(...)` + A single entrypoint that dispatches to either the default row-epilogue or the + LDS CShuffle epilogue based on input parameters. + +- `default_epilog(...)` (implementation helper) + A lightweight row-iterator for the common MFMA accumulator-to-output mapping + (mi in [0,m_repeat), ii in [0,4), row = bx_m + mi*16 + lane_div_16*4 + ii). + The caller supplies `body_row(...)` that performs the per-row epilogue work + (e.g. loads scales once, loops over ni, stores). + +- `c_shuffle_epilog(...)` (implementation helper) + A LDS CShuffle epilogue skeleton: + 1) call `write_row_to_lds(...)` for each MFMA output row to populate `lds_out` + in row-major [tile_m, tile_n] order + 2) barrier + 3) remap threads into (MLane, NLane) = (8,32) and read half2 from LDS, + then call `store_pair(...)` to emit the final global store/atomic. + + When ``lds_out_split`` is provided, the epilogue runs in split-LDS mode: + waves are partitioned into two groups (group A uses ``lds_out``, group B + uses ``lds_out_split``), each handling half of the N dimension. + +These helpers are intentionally *dialect-agnostic*: callers pass the dialect +modules (`arith`, `vector`, `gpu`) and the `range_constexpr` iterator. +""" + +from __future__ import annotations + +from typing import Callable + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects.arith import CmpIPredicate +from flydsl.expr.typing import T +from mslk.flydsl.kernels.common.kernels_common import _if_then + + +def default_epilog( + *, + arith, + range_constexpr, + m_repeat: int, + lane_div_16, + bx_m, + body_row: Callable, +): + """Iterate the standard MFMA 16x16 row mapping and call `body_row(...)`. + + The mapping matches the common MFMA fragment layout used across kernels in this repo. + + Args: + arith: flydsl arith ext module. + range_constexpr: compile-time unrolled range helper. + m_repeat: tile_m // 16 (python int). + lane_div_16: index Value (0..3). + bx_m: base row (index Value). For MoE, this is the base sorted-row for the tile. + body_row: callback invoked as: + body_row(mi=, ii=, row_in_tile=, row=) + """ + bx_m_v = bx_m + lane_div_16_mul4 = lane_div_16 * 4 + ii_idx_list = [fx.Index(ii) for ii in range(4)] + + for mi in range_constexpr(m_repeat): + mi_base = arith.constant(mi * 16, index=True) + for ii in range_constexpr(4): + row_off = lane_div_16_mul4 + ii_idx_list[ii] + row_in_tile = mi_base + row_off + row = bx_m_v + row_in_tile + body_row(mi=mi, ii=ii, row_in_tile=row_in_tile, row=row) + + +def c_shuffle_epilog( + *, + arith, + vector, + gpu, + scf=None, + range_constexpr, + # Tile params + tile_m: int, + tile_n: int, + e_vec: int = 2, + cshuffle_nlane: int = 32, + block_size: int = 256, + m_repeat: int, + num_acc_n: int, + # Thread mapping inputs + tx, + lane_div_16, + lane_mod_16, + bx_m, + by_n, + n_tile_base, + # LDS buffer (f16 view, row-major [tile_m, tile_n] flattened) + lds_out, + # Element type for LDS loads (defaults to f16). Pass bf16 to support bf16 epilogues. + frag_elem_type: ir.Type | None = None, + # Callbacks + write_row_to_lds: Callable, + precompute_row: Callable | None = None, + store_pair: Callable, + # When LDS overflows, split lds_out across two buffers by wave-group. + # Pass the second buffer here; first buffer is `lds_out`. + lds_out_split=None, + # Row offset in lds_out for 8-wave mode (MLIR index value). + # Shifts both write and read LDS indices by lds_row_offset * tile_n elements. + lds_row_offset=None, +): + """LDS CShuffle epilogue skeleton. + + Call pattern: + - `write_row_to_lds(...)` is called once per MFMA row produced by this thread. + It is responsible for writing all ni columns for that row into `lds_out`. + - `store_pair(...)` is called for each (row_local, col_pair0) half2 after shuffle. + + `store_pair` can implement either global stores or atomics. + """ + if int(block_size) <= 0 or (int(block_size) % int(cshuffle_nlane)) != 0: + raise ValueError(f"block_size ({block_size}) must be divisible by cshuffle_nlane ({cshuffle_nlane})") + cshuffle_mlane = int(block_size) // int(cshuffle_nlane) + if (int(tile_m) % cshuffle_mlane) != 0: + raise ValueError(f"tile_m must be divisible by CShuffleMLane ({cshuffle_mlane}), got tile_m={tile_m}") + if int(e_vec) <= 0: + raise ValueError(f"e_vec must be positive, got {e_vec}") + if (int(tile_n) % (int(cshuffle_nlane) * int(e_vec))) != 0: + raise ValueError( + f"tile_n must be divisible by (CShuffleNLane*EVec) = {cshuffle_nlane*e_vec}, got tile_n={tile_n}" + ) + + # ===================== Split-LDS mode (early return) ===================== + # When lds_out_split is provided, waves are divided into two groups: + # Group A (waves 0..N/2-1) uses lds_out, columns [0, tile_n/2) + # Group B (waves N/2..N-1) uses lds_out_split, columns [tile_n/2, tile_n) + # Each group writes/reads independently; same barriers synchronise all waves. + if lds_out_split is not None: + if scf is None: + raise ValueError("scf module is required for split-LDS cshuffle") + + _half_n = int(tile_n) // 2 + _half_threads = int(block_size) // 2 + EVec = int(e_vec) + + CShuffleNLane_s = min(int(cshuffle_nlane), _half_n // EVec) + if _half_threads % CShuffleNLane_s != 0: + raise ValueError(f"half_threads={_half_threads} not divisible by CShuffleNLane_split={CShuffleNLane_s}") + CShuffleMLane_s = _half_threads // CShuffleNLane_s + if int(tile_m) % CShuffleMLane_s != 0: + raise ValueError(f"tile_m={tile_m} not divisible by CShuffleMLane_split={CShuffleMLane_s}") + m_reps_s = int(tile_m) // CShuffleMLane_s + n_reps_s = _half_n // (CShuffleNLane_s * EVec) + + _half_n_idx = arith.constant(_half_n, index=True) + _half_thr_idx = arith.constant(_half_threads, index=True) + _zero_idx = arith.constant(0, index=True) + + _is_group_b = arith.cmpi(CmpIPredicate.uge, tx, _half_thr_idx) + + # -- write phase (all waves, each to its group's LDS buffer) -- + n_tile_base_v = n_tile_base + col_base_local_a = n_tile_base_v + lane_mod_16 + col_base_local_b = col_base_local_a - _half_n_idx + + def _write_row_split(mi: int, ii: int, row_in_tile, row): + row_base_lds = row_in_tile * _half_n_idx + _if_g = scf.IfOp(_is_group_b) + with ir.InsertionPoint(_if_g.then_block): + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local_b, + num_acc_n=num_acc_n, + lds_out=lds_out_split, + ) + scf.YieldOp([]) + with ir.InsertionPoint(_if_g.else_block): + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local_a, + num_acc_n=num_acc_n, + lds_out=lds_out, + ) + scf.YieldOp([]) + + gpu.barrier() + default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=_write_row_split, + ) + gpu.barrier() + + # -- read phase (each group reads from its own LDS buffer) -- + tx_local = tx - arith.select(_is_group_b, _half_thr_idx, _zero_idx) + c_nlane_s = arith.constant(CShuffleNLane_s, index=True) + m_lane_s = tx_local / c_nlane_s + n_lane_s = tx_local % c_nlane_s + c_evec = arith.constant(EVec, index=True) + + if frag_elem_type is None: + frag_elem_type = T.f16 + vec_frag = T.vec(EVec, frag_elem_type) + bx_m_v = bx_m + by_n_v = by_n + + _precomputed_rows_s = [] + for mr in range_constexpr(m_reps_s): + row_base_m = arith.constant(mr * CShuffleMLane_s, index=True) + row_local = row_base_m + m_lane_s + row = bx_m_v + row_local + row_ctx_raw = precompute_row(row_local=row_local, row=row) if precompute_row is not None else None + row_ctx = row_ctx_raw + row_pred = None + if scf is not None and row_ctx_raw is not None and isinstance(row_ctx_raw, tuple) and len(row_ctx_raw) == 2: + row_ctx, row_pred = row_ctx_raw + _precomputed_rows_s.append((row_local, row, row_ctx, row_pred)) + + for mr in range_constexpr(m_reps_s): + row_local, row, row_ctx, row_pred = _precomputed_rows_s[mr] + + def _do_store_row_split(): + row_base_lds = row_local * _half_n_idx + for nr in range_constexpr(n_reps_s): + col_base_nr = arith.constant(nr * (CShuffleNLane_s * EVec), index=True) + col_pair0_local = col_base_nr + (n_lane_s * c_evec) + lds_idx = row_base_lds + col_pair0_local + + _if_ld = scf.IfOp(_is_group_b, [vec_frag]) + with ir.InsertionPoint(_if_ld.then_block): + fb = vector.load_op(vec_frag, lds_out_split, [lds_idx]) + scf.YieldOp([fb]) + with ir.InsertionPoint(_if_ld.else_block): + fa = vector.load_op(vec_frag, lds_out, [lds_idx]) + scf.YieldOp([fa]) + frag = _if_ld.results[0] + + col_pair0 = col_pair0_local + arith.select(_is_group_b, _half_n_idx, _zero_idx) + store_pair( + row_local=row_local, + row=row, + row_ctx=row_ctx, + col_pair0=col_pair0, + col_g0=by_n_v + col_pair0, + frag=frag, + ) + + if row_pred is not None: + _if_row = scf.IfOp(row_pred) + with _if_then(_if_row, scf): + _do_store_row_split() + else: + _do_store_row_split() + + return # split path complete + + # ===================== Standard (non-split) path below ===================== + + # ---------------- Step 1: write C tile to LDS (row-major, fp16) ---------------- + tile_n_idx = arith.constant(int(tile_n), index=True) + n_tile_base_v = n_tile_base + col_base_local = n_tile_base_v + lane_mod_16 # index within [0,tile_n) + + _lds_row_base_offset = lds_row_offset * tile_n_idx if lds_row_offset is not None else None + + def _write_row(mi: int, ii: int, row_in_tile, row): + row_base_lds = row_in_tile * tile_n_idx + if _lds_row_base_offset is not None: + row_base_lds = row_base_lds + _lds_row_base_offset + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local, + num_acc_n=num_acc_n, + lds_out=lds_out, + ) + + # Ensure all LDS reads finished before the lds write. + gpu.barrier() + default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=_write_row, + ) + + # Ensure all LDS writes are visible before the shuffle-read. + gpu.barrier() + + # ---------------- Step 2: shuffle mapping + half2 store/atomic ---------------- + CShuffleNLane = int(cshuffle_nlane) + CShuffleMLane = int(cshuffle_mlane) + EVec = int(e_vec) + + m_reps_shuffle = int(tile_m) // CShuffleMLane + n_reps_shuffle = int(tile_n) // (CShuffleNLane * EVec) + + c_nlane = fx.Index(CShuffleNLane) + m_lane = tx // c_nlane + n_lane = tx % c_nlane + c_evec = fx.Index(EVec) + + if frag_elem_type is None: + frag_elem_type = T.f16 + vec_frag = T.vec(EVec, frag_elem_type) + bx_m_v = bx_m + by_n_v = by_n + + # Batch-precompute all row contexts (sorted_idx loads) before the store loop. + # This issues all buffer_load instructions upfront so the compiler can pipeline + # them instead of serializing each load with s_waitcnt vmcnt(0). + _precomputed_rows = [] + for mr in range_constexpr(m_reps_shuffle): + row_base_m = arith.constant(mr * CShuffleMLane, index=True) + row_local = row_base_m + m_lane + row = bx_m_v + row_local + + row_ctx_raw = precompute_row(row_local=row_local, row=row) if precompute_row is not None else None + + # Optional row-level predicate: if `precompute_row` returns `(ctx, pred_i1)` and `scf` + # is provided, we can skip the entire N-loop for invalid rows (cheaper than per-store checks). + row_ctx = row_ctx_raw + row_pred = None + if scf is not None and row_ctx_raw is not None and isinstance(row_ctx_raw, tuple) and len(row_ctx_raw) == 2: + row_ctx, row_pred = row_ctx_raw + + _precomputed_rows.append((row_local, row, row_ctx, row_pred)) + + # Now perform LDS reads and stores using the pre-fetched row contexts. + for mr in range_constexpr(m_reps_shuffle): + row_local, row, row_ctx, row_pred = _precomputed_rows[mr] + + def _do_store_row(): + row_base_lds = row_local * tile_n_idx + if _lds_row_base_offset is not None: + row_base_lds = row_base_lds + _lds_row_base_offset + for nr in range_constexpr(n_reps_shuffle): + col_base_nr = arith.constant(nr * (CShuffleNLane * EVec), index=True) + col_pair0 = col_base_nr + (n_lane * c_evec) # even col within tile + + lds_idx_pair = row_base_lds + col_pair0 + frag = vector.load_op(vec_frag, lds_out, [lds_idx_pair]) + + store_pair( + row_local=row_local, + row=row, + row_ctx=row_ctx, + col_pair0=col_pair0, + col_g0=by_n_v + col_pair0, + frag=frag, + ) + + if row_pred is not None: + _if_row = scf.IfOp(row_pred) + with _if_then(_if_row, scf): + _do_store_row() + else: + _do_store_row() + + +def mfma_epilog( + *, + use_cshuffle: bool, + # Common (always required) + arith, + range_constexpr, + m_repeat: int, + lane_div_16, + bx_m, + # Default epilog (required when use_cshuffle=False) + body_row: Callable | None = None, + # CShuffle epilog (required when use_cshuffle=True) + vector=None, + gpu=None, + scf=None, + tile_m: int | None = None, + tile_n: int | None = None, + e_vec: int = 2, + cshuffle_nlane: int = 32, + block_size: int = 256, + num_acc_n: int | None = None, + tx=None, + lane_mod_16=None, + by_n=None, + n_tile_base=None, + lds_out=None, + write_row_to_lds: Callable | None = None, + precompute_row: Callable | None = None, + store_pair: Callable | None = None, + frag_elem_type: ir.Type | None = None, +): + if not use_cshuffle: + if body_row is None: + raise ValueError("mfma_epilog(use_cshuffle=False) requires `body_row`.") + return default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=body_row, + ) + + return c_shuffle_epilog( + arith=arith, + vector=vector, + gpu=gpu, + scf=scf, + range_constexpr=range_constexpr, + tile_m=int(tile_m), + tile_n=int(tile_n), + e_vec=int(e_vec), + cshuffle_nlane=int(cshuffle_nlane), + block_size=int(block_size), + m_repeat=m_repeat, + num_acc_n=int(num_acc_n), + tx=tx, + lane_div_16=lane_div_16, + lane_mod_16=lane_mod_16, + bx_m=bx_m, + by_n=by_n, + n_tile_base=n_tile_base, + lds_out=lds_out, + frag_elem_type=frag_elem_type, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) diff --git a/mslk/flydsl/kernels/mma/mfma_preshuffle_pipeline.py b/mslk/flydsl/kernels/mma/mfma_preshuffle_pipeline.py new file mode 100644 index 00000000..6c9ef6ec --- /dev/null +++ b/mslk/flydsl/kernels/mma/mfma_preshuffle_pipeline.py @@ -0,0 +1,893 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Shared MFMA preshuffle helpers for preshuffle GEMM kernels. + +Key primitives: +- B preshuffle layout builder (supports byte-packed element types, incl. packed int4) +- B pack load for MFMA K32 micro-steps (8B output pack; optional int4->int8 unpack) +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl.expr import arith as _arith +from flydsl.expr.typing import T + + +def crd2idx(crd, layout): + """crd2idx returning an index-typed ir.Value (unwraps fly.int_tuple).""" + scalar = fx.get_scalar(fx.crd2idx(crd, layout)).ir_value() + if isinstance(scalar.type, ir.IndexType): + return scalar + return _arith.IndexCastOp(T.index, scalar).result + + +def swizzle_xor16(row, col, k_blocks16): + """XOR-with-row swizzle on the K dimension at 16B granularity. + + Computes: col XOR ((row & (k_blocks16 - 1)) * 16) + + k_blocks16 is always a power of 2 (tile_k_bytes / 16), so use + bitwise AND instead of remui to save ~10 VALU cycles on CDNA. + """ + from flydsl.expr import arith as _swz_arith + + mask = k_blocks16 - _swz_arith.index(1) + rem = _swz_arith.andi(row, mask) + return col ^ (rem * 16) + + +def lds_row_major_idx(row, col, row_stride, base=None): + """Linearize a 2D LDS coordinate with explicit index arithmetic.""" + idx = row * row_stride + col + return idx if base is None else idx + base + + +def split_row_major_2d(index, minor_extent): + """Split a linear row-major index into (major, minor).""" + return index // minor_extent, index % minor_extent + + +def _buffer_load_vec( + buffer_ops, + vector, + rsrc, + idx, + *, + elem_type, + vec_elems, + elem_bytes, + offset_in_bytes, + cache_modifier=0, +): + """Load vec_elems elements via buffer_load dwordx[1,2,4] + bitcast.""" + from flydsl.expr import arith as _ld_arith + + elem_size = int(elem_bytes) + load_bytes = int(vec_elems) * elem_size + vec_width = load_bytes // 4 + + if offset_in_bytes: + idx_i32 = _ld_arith.shrui(idx, _ld_arith.index(2)) + elif elem_bytes == 2: + idx_i32 = _ld_arith.shrui(idx, _ld_arith.index(1)) + else: + idx_i32 = idx + + i32_val = buffer_ops.buffer_load( + rsrc, + idx_i32, + vec_width=vec_width, + dtype=T.i32, + cache_modifier=cache_modifier, + ) + if vec_width == 1: + i32_vec = vector.from_elements(T.vec(1, T.i32), [i32_val]) + else: + i32_vec = i32_val + return vector.bitcast(T.vec(int(vec_elems), elem_type), i32_vec) + + +@dataclass(frozen=True) +class PreshuffleScaleLayout: + """Container returned by `make_preshuffle_scale_layout`. + + The scale layout is ``(c_mn1, c_k1, 4, 16) : (stride_n0, stride_k0, stride_klane, 1)``. + Callers compute flat index directly with plain arith:: + + idx = mni * stride_n0 + ku * stride_k0 + k_lane * stride_klane + n_lane + """ + + layout_scale: object + stride_n0: object + stride_k0: object + stride_klane: object + + +def make_preshuffle_scale_layout( + arith, + *, + c_mn: ir.Value, + c_k: ir.Value, + mn_pack: int = 2, + k_pack: int = 2, + elem_bytes: int = 4, + scale_block_size: int = 32, +) -> PreshuffleScaleLayout: + """Build scale layout matching aiter/CK preshuffle for FP4/FP8 microscale. + + Layout shape: ``(c_mn1, c_k1, 4, 16)`` where + ``c_mn1 = c_mn / 16 / mn_pack`` and ``c_k1 = (c_k / scale_block_size) / 4 / k_pack``. + """ + c16 = fx.Index(16) + c4 = fx.Index(4) + c_k_scale = c_k // fx.Index(scale_block_size) + + c_mn1 = (c_mn // c16) // fx.Index(mn_pack) + c_k1 = (c_k_scale // c4) // fx.Index(k_pack) + if elem_bytes != mn_pack * k_pack: + raise ValueError(f"elem_bytes of scale must be {mn_pack} * {k_pack}, got {elem_bytes!r}") + + stride_klane = c16 + stride_k0 = c4 * stride_klane + stride_n0 = c_k1 * stride_k0 + + c_mn1_i32 = arith.index_cast(T.i32, c_mn1) + c_k1_i32 = arith.index_cast(T.i32, c_k1) + stride_n0_i32 = arith.index_cast(T.i32, stride_n0) + stride_k0_i32 = arith.index_cast(T.i32, stride_k0) + stride_klane_i32 = arith.index_cast(T.i32, stride_klane) + + layout_scale = fx.make_layout( + (c_mn1_i32, c_k1_i32, 4, 16), + stride=(stride_n0_i32, stride_k0_i32, stride_klane_i32, 1), + ) + + return PreshuffleScaleLayout( + layout_scale=layout_scale, + stride_n0=stride_n0, + stride_k0=stride_k0, + stride_klane=stride_klane, + ) + + +@dataclass(frozen=True) +class PreshuffleBLayout: + """Container returned by `make_preshuffle_b_layout`.""" + + layout_b: object + kpack_bytes: int + + +def make_preshuffle_b_layout( + arith, + *, + c_n: ir.Value, + c_k: ir.Value, + kpack_bytes: int = 16, + elem_bytes: int = 1, + k_major: bool = False, +) -> PreshuffleBLayout: + """Build B layout matching aiter/CK preshuffle for A8 MFMA kernels. + + When *k_major* is True the block-level order is K-major (``k_blk`` outermost), + matching the ``(0,3,1,4,2,5)`` shuffle permutation. The default N-major + order (``k_major=False``) matches the legacy ``(0,1,3,4,2,5)`` permutation. + """ + if kpack_bytes not in (8, 16): + raise ValueError(f"kpack_bytes must be 8 or 16, got {kpack_bytes!r}") + + c16 = fx.Index(16) + c_kpack = fx.Index(kpack_bytes) + + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + c_k_bytes = c_k * arith.constant(int(elem_bytes), index=True) + n0 = c_n // c16 + + c_kpack_elems = c_kpack if elem_bytes == 1 else (c_kpack // arith.constant(int(elem_bytes), index=True)) + + stride_nlane = c_kpack_elems + + if k_major: + c32 = fx.Index(32) + c2 = fx.Index(2) + c_k0 = c_k_bytes // c32 + klane_dim = 2 + stride_klane = c16 * stride_nlane + stride_n0 = c2 * stride_klane + stride_k0 = n0 * stride_n0 + else: + c64 = fx.Index(64) + c4 = fx.Index(4) + c_k0 = c_k_bytes // c64 + klane_dim = 4 + stride_klane = c16 * stride_nlane + stride_k0 = c4 * stride_klane + stride_n0 = c_k0 * stride_k0 + + kpack_elems_static = kpack_bytes if elem_bytes == 1 else kpack_bytes // elem_bytes + n0_i32 = arith.index_cast(T.i32, n0) + c_k0_i32 = arith.index_cast(T.i32, c_k0) + stride_n0_i32 = arith.index_cast(T.i32, stride_n0) + stride_k0_i32 = arith.index_cast(T.i32, stride_k0) + stride_klane_i32 = arith.index_cast(T.i32, stride_klane) + stride_nlane_i32 = arith.index_cast(T.i32, stride_nlane) + + stride_b = (stride_n0_i32, stride_k0_i32, stride_klane_i32, stride_nlane_i32, 1) + layout_b = fx.make_layout((n0_i32, c_k0_i32, klane_dim, 16, kpack_elems_static), stride_b) + return PreshuffleBLayout(layout_b=layout_b, kpack_bytes=kpack_bytes) + + +def _unpack_int4_to_int8_pair(packed32): + """Split packed int4 dword into two int8 dwords (even/odd nibbles). + + 7-op bit manipulation shared by all int4 unpack paths (W4A8, W4A16, W4A_FP8). + """ + c_08 = fx.Int32(0x08080808) + c_0f = fx.Int32(0x0F0F0F0F) + c_1e = fx.Int32(0x1E) + c_4 = fx.Int32(4) + s0 = (packed32 & c_08) * c_1e + even = (packed32 & c_0f) | s0 + t = packed32 >> c_4 + s1 = (t & c_08) * c_1e + odd = (t & c_0f) | s1 + return even, odd + + +def _pack_i32_pair_to_i64(lo, hi, vector): + """Pack two i32 values into one i64 via vector bitcast.""" + v2 = vector.from_elements(T.vec(2, T.i32), [lo, hi]) + v64 = vector.bitcast(T.vec(1, T.i64), v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def _i8x4_in_i32_to_bf16x4_i64(val_i32, arith, vector, scale_val=None): + """Convert one i32 (4 signed int8 bytes) to 4 bf16 packed as i64. + + Uses shift-based f32->bf16 truncation (lshr 16) instead of arith.truncf + which on gfx942 expands to ~5 VALU per element. The shift is exact for + unscaled int8 values and introduces <0.5 ULP error for scaled values. + """ + v1 = vector.from_elements(T.vec(1, T.i32), [val_i32]) + i8x4 = vector.bitcast(T.i8x4, v1) + + f32_vals = [] + for i in range(4): + val_i8 = vector.extract(i8x4, static_position=[i], dynamic_position=[]) + v = arith.sitofp(T.f32, val_i8) + if scale_val is not None: + v = v * scale_val + f32_vals.append(v) + + c16 = fx.Int32(16) + c_ffff0000 = fx.Int32(0xFFFF0000) + bits = [arith.bitcast(T.i32, f) for f in f32_vals] + i32_lo = (bits[0] >> c16) | (bits[1] & c_ffff0000) + i32_hi = (bits[2] >> c16) | (bits[3] & c_ffff0000) + return _pack_i32_pair_to_i64(i32_lo, i32_hi, vector) + + +def load_b_raw_w4a16( + buffer_ops, + arith, + vector, + *, + arg_b, + b_rsrc, + layout_b, + base_k: ir.Value, + ku: int, + n_blk: ir.Value, + n_intra: ir.Value, + lane_div_16: ir.Value, + elem_type: ir.Type, + kpack_bytes: int = 8, +): + """Phase 1 of W4A16 B load: issue buffer_load_dword, return raw packed i32. + + Same address calculation as the int4 unpack path in load_b_pack_k32 + but using ku-based indexing for 2-phase latency hiding. + """ + if kpack_bytes != 8: + raise ValueError(f"W4A16 requires kpack_bytes=8, got {kpack_bytes!r}") + + c64 = fx.Index(64) + half_bytes = kpack_bytes // 2 + c2_idx = fx.Index(2) + c4_idx = fx.Index(4) + + k0_base = base_k // c64 + + k1_layout_offset = ku * 2 + lane_div_32 = lane_div_16 // c2_idx + total_k1 = fx.Index(k1_layout_offset) + lane_div_32 + k0 = k0_base + (total_k1 // c4_idx) + k1_local = total_k1 % c4_idx + lane_odd = lane_div_16 % c2_idx + k2_base = lane_odd * fx.Index(half_bytes) + + coord_pack = (n_blk, k0, k1_local, n_intra, fx.Index(0)) + idx_pack = crd2idx(tuple(fx.Int32(c) for c in coord_pack), layout_b) + idx_bytes = idx_pack + k2_base + + b4 = _buffer_load_vec( + buffer_ops, + vector, + b_rsrc, + idx_bytes, + elem_type=elem_type, + vec_elems=4, + elem_bytes=1, + offset_in_bytes=True, + ) + packed32 = vector.extract( + vector.bitcast(T.vec(1, T.i32), b4), + static_position=[0], + dynamic_position=[], + ) + return packed32 + + +def _int4_to_bf16x4_i64_gfx950(packed32, nibble_offsets, arith, vector, scale_val=None, defer_scale16=False): + """Convert 4 int4 nibbles to 4 bf16 packed as i64 using gfx950 instructions. + + Uses v_cvt_off_f32_i4_sdwa with byte_sel to avoid per-nibble shifts. + Even nibbles (0,2,4,6) → SDWA BYTE_0/1/2/3 on original src. + Odd nibbles (1,3,5,7) → SDWA BYTE_0/1/2/3 on (src >> 4). + Only 1 shift total instead of 7. + + When defer_scale16=True, the ×16 correction factor for v_cvt_off_f32_i4 is + omitted and must be applied later (e.g. in the epilogue). This saves VALU + in the hot loop and uses v_cvt_pk_bf16_f32 for proper f32→bf16 conversion. + """ + from flydsl._mlir.dialects._arith_ops_gen import MulFOp as _MulFOp + from flydsl.expr import rocdl + + _uw = _arith._to_raw + _av = _arith.ArithValue + + src_even = packed32 + src_odd = packed32 >> fx.Int32(4) + + f32_vals = [] + for nib in nibble_offsets: + byte_idx = nib // 2 + src = src_odd if (nib % 2) else src_even + v = rocdl.cvt_off_f32_i4(src, byte_sel=byte_idx) + f32_vals.append(v) + + if defer_scale16: + # Skip ×16; multiply by scale_val only if groupwise. + if scale_val is not None: + raw_scale = _uw(scale_val) + f32_vals = [_MulFOp(v, raw_scale).result for v in f32_vals] + # Use v_cvt_pk_bf16_f32 for proper f32→bf16 (no bit-shift trick needed). + i32_lo = rocdl.cvt_pk_bf16_f32(f32_vals[0], f32_vals[1]) + i32_hi = rocdl.cvt_pk_bf16_f32(f32_vals[2], f32_vals[3]) + else: + c16 = fx.Float32(16.0) + effective_scale = scale_val * c16 if scale_val is not None else c16 + raw_scale = _uw(effective_scale) + f32_vals = [_MulFOp(v, raw_scale).result for v in f32_vals] + # Truncate f32→bf16 via bit-shift (exact for scaled int values). + c16_shift = fx.Int32(16) + c_ffff0000 = fx.Int32(0xFFFF0000) + bf16_vals = [arith.bitcast(T.i32, _av(v)) for v in f32_vals] + i32_lo = (bf16_vals[0] >> c16_shift) | (bf16_vals[1] & c_ffff0000) + i32_hi = (bf16_vals[2] >> c16_shift) | (bf16_vals[3] & c_ffff0000) + + return _pack_i32_pair_to_i64(i32_lo, i32_hi, vector) + + +def unpack_b_w4a16(packed32, arith, vector, scale_val=None, use_gfx950_cvt=False, defer_scale16=False): + """Phase 2 of W4A16 B load: unpack int4->int8 + convert int8->bf16. + + Takes raw packed32 from load_b_raw_w4a16 and produces (b0, b1) -- + two i64 values each containing 4 bf16 for one MFMA. + + When use_gfx950_cvt=True, uses v_cvt_off_f32_i4 + v_cvt_pk_bf16_f32 + for ~2x fewer VALU instructions. + + When defer_scale16=True (requires use_gfx950_cvt=True), the ×16 + correction for v_cvt_off_f32_i4 is omitted; caller must apply it + in the epilogue. + """ + if use_gfx950_cvt: + b0 = _int4_to_bf16x4_i64_gfx950(packed32, [0, 2, 4, 6], arith, vector, scale_val, defer_scale16=defer_scale16) + b1 = _int4_to_bf16x4_i64_gfx950(packed32, [1, 3, 5, 7], arith, vector, scale_val, defer_scale16=defer_scale16) + return (b0, b1) + even, odd = _unpack_int4_to_int8_pair(packed32) + b0 = _i8x4_in_i32_to_bf16x4_i64(even, arith, vector, scale_val=scale_val) + b1 = _i8x4_in_i32_to_bf16x4_i64(odd, arith, vector, scale_val=scale_val) + return (b0, b1) + + +def load_b_pack_k32( + buffer_ops, + arith, + vector, + *, + arg_b, + b_rsrc, + layout_b, + base_k: ir.Value, + ki_step: int, + n_blk: ir.Value, + n_intra: ir.Value, + lane_div_16: ir.Value, + elem_type: ir.Type, + kpack_bytes: int = 16, + elem_bytes: int = 1, + unpack_int4: bool = False, +) -> ir.Value: + """Load one B pack for one MFMA(x32) micro-step. + + Returns an i64 Value containing 8 bytes consumed by MFMA. + """ + if kpack_bytes not in (8, 16): + raise ValueError(f"kpack_bytes must be 8 or 16, got {kpack_bytes!r}") + if unpack_int4 and kpack_bytes != 8: + raise ValueError("unpack_int4 requires kpack_bytes=8 (packed int4 layout)") + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + + c64 = fx.Index(64) + base_k_bytes = base_k * arith.constant(int(elem_bytes), index=True) + k0_base = base_k_bytes // c64 + k0 = k0_base + arith.constant(ki_step // 2, index=True) + k1 = lane_div_16 + half_bytes = kpack_bytes // 2 + k2_base = arith.constant((ki_step % 2) * half_bytes, index=True) + + coord_pack = (n_blk, k0, k1, n_intra, fx.Index(0)) + idx_pack = crd2idx(tuple(fx.Int32(c) for c in coord_pack), layout_b) + + if unpack_int4: + idx_bytes = idx_pack + k2_base + b4 = _buffer_load_vec( + buffer_ops, + vector, + b_rsrc, + idx_bytes, + elem_type=elem_type, + vec_elems=4, + elem_bytes=1, + offset_in_bytes=True, + ) + packed32 = vector.extract( + vector.bitcast(T.vec(1, T.i32), b4), + static_position=[0], + dynamic_position=[], + ) + even, odd = _unpack_int4_to_int8_pair(packed32) + return _pack_i32_pair_to_i64(even, odd, vector) + + vec_elems = kpack_bytes // int(elem_bytes) + b16 = _buffer_load_vec( + buffer_ops, + vector, + b_rsrc, + idx_pack, + elem_type=elem_type, + vec_elems=vec_elems, + elem_bytes=elem_bytes, + offset_in_bytes=(elem_bytes == 1), + ) + + b_i32x4 = vector.bitcast(T.i32x4, b16) + + base = (ki_step % 2) * 2 + d0 = vector.extract(b_i32x4, static_position=[base], dynamic_position=[]) + d1 = vector.extract(b_i32x4, static_position=[base + 1], dynamic_position=[]) + return _pack_i32_pair_to_i64(d0, d1, vector) + + +def tile_chunk_coord_i32( + arith, + *, + tx_i32_base: ir.Value, + i: int, + total_threads: int, + layout_tile_div4, + chunk_i32: int = 4, +): + """Map (thread, chunk_id) -> (row_local, col_local_i32) for X/A loads.""" + if chunk_i32 not in (1, 2, 4): + raise ValueError(f"chunk_i32 must be one of (1,2,4), got {chunk_i32!r}") + chunk_off_i32 = arith.constant(i * total_threads * chunk_i32, index=True) + tile_idx_i32 = tx_i32_base + chunk_off_i32 + coord_local = fx.idx2crd(fx.Int32(tile_idx_i32), layout_tile_div4) + row_local = fx.get(coord_local, 0) + col_local_i32 = fx.get(coord_local, 1) + return row_local, col_local_i32 + + +def buffer_copy_gmem16_dwordx4( + buffer_ops, + vector, + *, + elem_type, + idx_i32: ir.Value, + rsrc, + vec_elems: int = 16, + elem_bytes: int = 1, +): + """Copy 16 bytes from global memory into regs via buffer-load dwordx4 lowering.""" + if int(vec_elems) <= 0: + raise ValueError(f"vec_elems must be > 0, got {vec_elems!r}") + return _buffer_load_vec( + buffer_ops, + vector, + rsrc, + idx_i32, + elem_type=elem_type, + vec_elems=vec_elems, + elem_bytes=elem_bytes, + offset_in_bytes=False, + ) + + +def _lds_store_xor16( + vector, + *, + lds_memref, + vec_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part: ir.Value, + elem_bytes: int, +): + """Store one chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + col_swz_bytes = swizzle_xor16(row_local, col_local_i32 * tx_c4, k_blocks16) + col_swz = col_swz_bytes if elem_bytes == 1 else col_swz_bytes // 2 + idx0 = crd2idx((fx.Int32(row_local), fx.Int32(col_swz)), layout_lds) + lds_base + vector.store(vector.bitcast(vec_ty, vec_part), lds_memref, [idx0]) + + +def lds_store_16b_xor16( + arith, + vector, + *, + lds_memref, + vec16_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part_i32x4: ir.Value, + elem_bytes: int = 1, +): + """Store one 16B chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + _lds_store_xor16( + vector, + lds_memref=lds_memref, + vec_ty=vec16_ty, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=tx_c4, + k_blocks16=k_blocks16, + lds_base=lds_base, + vec_part=vec_part_i32x4, + elem_bytes=elem_bytes, + ) + + +def lds_store_8b_xor16( + arith, + vector, + *, + lds_memref, + vec8_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part_i32x2: ir.Value, + elem_bytes: int = 1, +): + """Store one 8B chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + _lds_store_xor16( + vector, + lds_memref=lds_memref, + vec_ty=vec8_ty, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=tx_c4, + k_blocks16=k_blocks16, + lds_base=lds_base, + vec_part=vec_part_i32x2, + elem_bytes=elem_bytes, + ) + + +def lds_store_4b_xor16( + arith, + vector, + *, + lds_memref, + vec4_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part_i32x1: ir.Value, + elem_bytes: int = 1, +): + """Store one 4B chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + _lds_store_xor16( + vector, + lds_memref=lds_memref, + vec_ty=vec4_ty, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=tx_c4, + k_blocks16=k_blocks16, + lds_base=lds_base, + vec_part=vec_part_i32x1, + elem_bytes=elem_bytes, + ) + + +def lds_load_pack_k32( + arith, + vector, + *, + lds_memref, + layout_lds, + k_blocks16: ir.Value, + curr_row_a_lds: ir.Value, + col_base: ir.Value, + half: int, + lds_base: ir.Value, + ck_lds128: bool, + vec16_ty, + vec8_ty, + vec2_i64_ty, + vec1_i64_ty, +): + """Load one i64 A-pack for an MFMA K32 micro-step from LDS.""" + col_base_swz = swizzle_xor16(curr_row_a_lds, col_base, k_blocks16) + if ck_lds128: + coord_a16 = (curr_row_a_lds, col_base_swz) + idx_a16 = crd2idx(tuple(fx.Int32(c) for c in coord_a16), layout_lds) + lds_base + loaded_a16 = vector.load_op(vec16_ty, lds_memref, [idx_a16]) + a_vec128 = vector.bitcast(vec2_i64_ty, loaded_a16) + return vector.extract(a_vec128, static_position=[half], dynamic_position=[]) + else: + col_swizzled = col_base_swz + (half * 8) + coord_a = (curr_row_a_lds, col_swizzled) + idx_a = crd2idx(tuple(fx.Int32(c) for c in coord_a), layout_lds) + lds_base + loaded_a8 = vector.load_op(vec8_ty, lds_memref, [idx_a]) + a_vec64 = vector.bitcast(vec1_i64_ty, loaded_a8) + return vector.extract(a_vec64, static_position=[0], dynamic_position=[]) + + +@flyc.jit +def xcd_remap_bx_by( + bx, + by, + c_m, + *, + tile_m: int, + tile_n: int, + N: int, + xcd_swizzle: int, + num_xcds: int = 8, +): + if xcd_swizzle <= 0: + return bx, by + + # Keep the whole remap in i32 (grid dims fit): gpu.block_id yields index, so + # cast the block ids so every derived value (and both ternary branches) is i32. + bx = fx.Int32(bx) + by = fx.Int32(by) + + gx = N // tile_n + gy = (c_m + tile_m - 1) // tile_m + + linear_id = bx * gx + by + num_wgs = gx * gy + + q = num_wgs // num_xcds + r = num_wgs % num_xcds + xcd = linear_id % num_xcds + in_xcd = linear_id // num_xcds + xcd_lt_r = xcd < r + clip = xcd if xcd_lt_r else r + wgid = xcd * q + clip + in_xcd + + num_wgid_in_group = xcd_swizzle * gx + group_id = wgid // num_wgid_in_group + first_pid_m = group_id * xcd_swizzle + remaining_m = gy - first_pid_m + cmp_m = remaining_m < xcd_swizzle + group_size_m = remaining_m if cmp_m else fx.Int32(xcd_swizzle) + + wgid_in_group = wgid % num_wgid_in_group + new_bx = first_pid_m + (wgid_in_group % group_size_m) + new_by = wgid_in_group // group_size_m + return new_bx, new_by + + +__all__ = [ + "PreshuffleBLayout", + "PreshuffleScaleLayout", + "buffer_copy_gmem16_dwordx4", + "lds_load_pack_k32", + "lds_row_major_idx", + "lds_store_4b_xor16", + "lds_store_8b_xor16", + "lds_store_16b_xor16", + "make_preshuffle_b_layout", + "make_preshuffle_scale_layout", + "load_b_pack_k32", + "split_row_major_2d", + "swizzle_xor16", + "tile_chunk_coord_i32", + "unpack_b_w4a16", + "xcd_remap_bx_by", +] + + +# --------------------------------------------------------------------------- +# Groupwise scale load helper (shared by W4A16 and W4A8 groupwise paths) +# --------------------------------------------------------------------------- + + +def _load_groupwise_scale( + buffer_ops, + arith, + *, + scale_rsrc, + expert_offset, + n_blk, + n_intra, + k_pos, + num_groups: int, + group_size: int, + n_per_expert: int, + scale_dtype=None, +): + """Load one per-group scale value from the scale buffer. + + Computes the linear index into the scale tensor from expert offset, + N position, and group index derived from ``k_pos``. + + For bf16 scales the tensor uses ``(E, G//2, N, 2)`` layout — two + adjacent groups for the same N position are packed into one dword. + We load the raw i32 dword (no extraction) so it can be carried as + loop state without register copies. Use :func:`extract_bf16_scale` + in the compute phase to obtain the f32 value. + """ + c16 = fx.Index(16) + n_global = n_blk * c16 + n_intra + c_group_size = fx.Index(group_size) + c_npe = fx.Index(n_per_expert) + group_idx = k_pos // c_group_size + if scale_dtype is None: + scale_dtype = T.f32 + + if scale_dtype == T.bf16: + # (E, G//2, N, 2) layout: same flat formula but with G//2 pairs of groups. + pair_idx = group_idx >> fx.Index(1) # group_idx // 2 + num_pairs = num_groups // 2 + c_npm1 = fx.Index(num_pairs - 1) + dword_base = expert_offset * c_npm1 + n_global + dword_elem = dword_base + pair_idx * c_npe + dword_idx = arith.index_cast(T.i32, dword_elem) + scale_val = buffer_ops.buffer_load(scale_rsrc, dword_idx, vec_width=1, dtype=T.i32) + else: + # (E, G, N) layout with f32 dtype + c_gm1 = fx.Index(num_groups - 1) + base_scale = expert_offset * c_gm1 + n_global + elem_idx = base_scale + group_idx * c_npe + scale_idx_i32 = arith.index_cast(T.i32, elem_idx) + scale_val = buffer_ops.buffer_load(scale_rsrc, scale_idx_i32, vec_width=1, dtype=T.f32) + return scale_val + + +def extract_bf16_scale(arith, scale_raw_i32, ku: int): + """Extract f32 scale from raw i32 dword loaded by bf16 groupwise path. + + In the ``(E, G//2, N, 2)`` layout two adjacent groups share one dword. + ``ku`` determines which half: even ku → low bf16, odd ku → high bf16. + """ + if ku % 2 == 0: + # Low bf16: shift left by 16 to place in upper 16 bits → f32 + return arith.bitcast(T.f32, scale_raw_i32 << fx.Int32(16)) + else: + # High bf16: mask upper 16 bits → f32 + return arith.bitcast(T.f32, scale_raw_i32 & fx.Int32(0xFFFF0000)) + + +# --------------------------------------------------------------------------- +# W4A16 groupwise load / unpack helpers +# --------------------------------------------------------------------------- + + +def load_b_raw_w4a16_groupwise( + buffer_ops, + arith, + vector, + *, + arg_b, + b_rsrc, + layout_b, + base_k, + ku: int, + n_blk, + n_intra, + lane_div_16, + elem_type, + scale_rsrc, + expert_offset, + num_groups: int, + group_size: int, + n_per_expert: int, + kpack_bytes: int = 8, + scale_dtype=None, +): + """Phase 1 of W4A16 groupwise B load: buffer_loads for weight + scale. + + Reuses :func:`load_b_raw_w4a16` for the weight load, then issues an + additional ``buffer_load_dword`` for the per-group scale. + + Returns ``(packed32, scale_val)``. + """ + packed32 = load_b_raw_w4a16( + buffer_ops, + arith, + vector, + arg_b=arg_b, + b_rsrc=b_rsrc, + layout_b=layout_b, + base_k=base_k, + ku=ku, + n_blk=n_blk, + n_intra=n_intra, + lane_div_16=lane_div_16, + elem_type=elem_type, + kpack_bytes=kpack_bytes, + ) + k_pos = base_k + fx.Index(ku * 32) + scale_val = _load_groupwise_scale( + buffer_ops, + arith, + scale_rsrc=scale_rsrc, + expert_offset=expert_offset, + n_blk=n_blk, + n_intra=n_intra, + k_pos=k_pos, + num_groups=num_groups, + group_size=group_size, + n_per_expert=n_per_expert, + scale_dtype=scale_dtype, + ) + return (packed32, scale_val) + + +def unpack_b_w4a16_groupwise(packed32, scale_val, arith, vector, use_gfx950_cvt=False): + """Phase 2 of W4A16 groupwise: unpack + scale + convert to bf16.""" + return unpack_b_w4a16(packed32, arith, vector, scale_val=scale_val, use_gfx950_cvt=use_gfx950_cvt) diff --git a/mslk/gemm/__init__.py b/mslk/gemm/__init__.py index 6ddb2329..dc25d03c 100644 --- a/mslk/gemm/__init__.py +++ b/mslk/gemm/__init__.py @@ -40,3 +40,10 @@ mx8mx4_gemm, mx8mx8_gemm, ) + + from mslk.flydsl.common import is_flydsl_available + + if is_flydsl_available(): + # Registers mslk::f8f8bf16_groupwise_grouped and its _preshuffle + # sibling (FlyDSL). + from .flydsl import fp8_groupwise_grouped_gemm as _flydsl_groupwise_grouped diff --git a/mslk/gemm/flydsl/fp8_groupwise_grouped_gemm.py b/mslk/gemm/flydsl/fp8_groupwise_grouped_gemm.py new file mode 100644 index 00000000..b41ddebc --- /dev/null +++ b/mslk/gemm/flydsl/fp8_groupwise_grouped_gemm.py @@ -0,0 +1,302 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +"""FP8 groupwise-scaled grouped GEMM via FlyDSL. + +Registers two ops, both backed by the same kernel: + +* ``mslk::f8f8bf16_groupwise_grouped`` -- the ROCm implementation of the plain + op, taking row-major ``[G, N, K]`` weights. +* ``mslk::f8f8bf16_groupwise_grouped_preshuffle`` -- a sibling that consumes + weights already in the MFMA B-preshuffle layout (see + ``mslk.quantize.shuffle.preshuffle_b_mfma``). Callers shuffle once at load + time; the op does no shuffling. + +Tensor contract: + XQ : [TotalM, K] FP8 -- all groups concatenated along M + WQ : [G, N, K] FP8 -- per-group weights, MFMA-preshuffled + for the preshuffle op + x_scale : FP32 -- per-token per-128K scales, in the + per-group block layout produced by + quantize_fp8_group(m_sizes=...) + w_scale : [G, K//128, N//128] FP32 -- per-group per-block scales + M_sizes : [G] int64 -- rows per group (sum to TotalM) + Output : [TotalM, N] BF16 +""" + +import os + +import torch + +from mslk.flydsl.common import is_flydsl_available +from mslk.flydsl.jit import run_compiled +from mslk.utils.device import supports_float8_fnuz + +_OP_NAME = "mslk::f8f8bf16_groupwise_grouped_preshuffle" + +# Only the scale-block granularity is fixed; tile_m/tile_n/tile_k are chosen per +# call -- either by FlyDSL autotune (MSLK_AUTOTUNE_ENABLE set) or a fixed default. +_SCALE_BLOCK = 128 + +# Default tile when autotuning is disabled. Valid for any supported shape +# (tile_n=tile_k=128 divide every supported N/K, incl. small N=128). This is the +# CI / no-benchmark path -- matches the CUTLASS heuristic fallback tile. +_DEFAULT_TILE = (128, 128, 128) + +# Candidate tile space swept by autotune. tile_n must be a multiple of +# scale_block_n=128 so a tile never straddles a weight scale block, and tile_k is +# pinned to the same granularity. Configs that do not divide a given shape are +# pruned before benchmarking, and ones that overflow LDS are rejected at compile. +_AUTOTUNE_TILES = ( + (64, 128, 128), + (128, 128, 128), + (256, 128, 128), + (64, 256, 128), + (128, 256, 128), + (256, 256, 128), +) + + +def _next_pow2(x: int) -> int: + """Smallest power of two >= x (x>=1). Buckets TotalM for the autotune key so + nearby token counts share one tuned config -- matching the CUDA-graph capture + buckets a server pre-captures, and bounding the pre-warm set.""" + if x <= 1: + return 1 + return 1 << (int(x) - 1).bit_length() + + +def _launch_kernel( + XQ, WQ, x_scale, w_scale, m_sizes, output, *, tile_m, tile_n, tile_k, b_preshuffled +): + """Compile (cached) and launch the grouped GEMM for one tile config. Shared + by the autotune target and the fixed-config path. Writes into `output`.""" + from mslk.flydsl.kernels.gemm.grouped_gemm_blockscale_contiguous import ( + compile_grouped_gemm_blockscale_contiguous, + ) + + TotalM, K = XQ.shape + G, N, _ = WQ.shape + # Grid M-extent: host-known upper bound (each group wastes at most one partial + # tile). The kernel resolves group ownership from M_sizes and self-skips + # surplus tiles. + num_m_tiles = TotalM // tile_m + G + launcher = compile_grouped_gemm_blockscale_contiguous( + n=N, + k=K, + num_groups=G, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + scale_block_k=_SCALE_BLOCK, + scale_block_n=_SCALE_BLOCK, + out_dtype="bf16", + b_preshuffled=b_preshuffled, + ) + # Operands keep their natural shape: argument marshalling packs each memref + # extent as int32, which a flattened view overflows at 2**31 elements. The + # kernel addresses them as flat byte buffers regardless. FP8 is viewed as + # int8 for the handoff. + run_compiled( + launcher, + output, + XQ.contiguous().view(torch.int8), + WQ.contiguous().view(torch.int8), + x_scale.contiguous(), + w_scale.contiguous(), + m_sizes, + TotalM, + N, + K, + G, + num_m_tiles, + torch.cuda.current_stream(), + ) + return output + +def _f8f8bf16_groupwise_grouped_preshuffle_meta( + XQ: torch.Tensor, + WQ: torch.Tensor, + x_scale: torch.Tensor, + w_scale: torch.Tensor, + M_sizes: torch.Tensor, +) -> torch.Tensor: + TotalM = XQ.shape[0] + N = WQ.shape[1] + return XQ.new_empty((TotalM, N), dtype=torch.bfloat16) + + +def _autotune_target( + XQ, WQ, x_scale, w_scale, m_sizes, output, m_bucket, n, k, b_preshuffled, + *, tile_m, tile_n, tile_k, +): + """FlyDSL @autotune benchmarks this per candidate tile. Keyed on + (m_bucket, n, k, b_preshuffled): m_bucket=nextPow2(TotalM) buckets token + counts; n/k separate the problem shapes (gate/up vs down-proj want different + tiles); b_preshuffled distinguishes the two kernels (different B-load path, + can't share a tuned config). Key args are otherwise passed straight through. + tile_* arrive as Config kwargs.""" + return _launch_kernel( + XQ, WQ, x_scale, w_scale, m_sizes, output, + tile_m=tile_m, tile_n=tile_n, tile_k=tile_k, b_preshuffled=b_preshuffled, + ) + + +def _prune_tiles(configs, named_args, **kwargs): + """Drop tile configs invalid for this shape (tile_n must divide N, tile_k + must divide K) before benchmarking. + + FlyDSL's Autotuner calls this as ``(configs, sig_args)``; ``**kwargs`` keeps + it compatible with the Triton-style ``(configs, named_args, **meta)`` form. + """ + WQ = named_args.get("WQ") + XQ = named_args.get("XQ") + if WQ is None or XQ is None: + return configs + N = WQ.shape[1] + K = XQ.shape[1] + kept = [ + c for c in configs + if N % c.kwargs["tile_n"] == 0 and K % c.kwargs["tile_k"] == 0 + ] + return kept or configs + + +# Single autotuner for both B-layout variants, built lazily (flydsl.autotune only +# imports when FlyDSL is present). b_preshuffled is a KEY arg (not a Config kwarg) +# so the two kernels get separate tuned entries in one shared disk cache. +_AUTOTUNER = None + + +def _get_autotuner(): + global _AUTOTUNER + if _AUTOTUNER is None: + from flydsl.autotune import Config, autotune + + configs = [ + Config(tile_m=tm, tile_n=tn, tile_k=tk) + for (tm, tn, tk) in _AUTOTUNE_TILES + ] + _AUTOTUNER = autotune( + configs=configs, + key=["m_bucket", "n", "k", "b_preshuffled"], + prune_configs_by=_prune_tiles, + )(_autotune_target) + return _AUTOTUNER + + +def _dispatch_grouped_gemm( + XQ: torch.Tensor, + WQ: torch.Tensor, + x_scale: torch.Tensor, + w_scale: torch.Tensor, + M_sizes: torch.Tensor, + *, + b_preshuffled: bool, +) -> torch.Tensor: + """Shared dispatch for both grouped ops. WQ is already in the layout the + variant expects (MFMA-preshuffled if b_preshuffled else plain [G,N,K]). + + Tile selection follows the CUTLASS precedent: when MSLK_AUTOTUNE_ENABLE is + set, FlyDSL autotune benchmarks the candidate tiles on a cache-miss and + persists the winner (keyed on nextPow2(TotalM) and b_preshuffled); otherwise + a fixed default tile is used with no benchmarking (the CI / graph-capture-safe + path). + """ + assert XQ.ndim == 2, f"XQ must be [TotalM, K], got {XQ.shape}" + assert WQ.ndim == 3, f"WQ must be [G, N, K], got {WQ.shape}" + assert M_sizes.ndim == 1, f"M_sizes must be [G], got {M_sizes.shape}" + TotalM, K = XQ.shape + G, N, Kw = WQ.shape + assert Kw == K, f"K mismatch: XQ K={K}, WQ K={Kw}" + assert M_sizes.shape[0] == G, f"M_sizes length {M_sizes.shape[0]} must equal G={G}" + # The MFMA instructions read the operands in the arch's native FP8 format, and + # the kernel passes them through as raw bytes, so an fnuz/OCP mismatch would + # be applied with the wrong exponent bias rather than rejected. + expected_fp8 = ( + torch.float8_e4m3fnuz if supports_float8_fnuz() else torch.float8_e4m3fn + ) + assert XQ.dtype == expected_fp8, f"XQ must be {expected_fp8}, got {XQ.dtype}" + assert WQ.dtype == expected_fp8, f"WQ must be {expected_fp8}, got {WQ.dtype}" + assert M_sizes.dtype == torch.int64, ( + f"M_sizes must be int64, got {M_sizes.dtype}" + ) + + output = torch.empty((TotalM, N), dtype=torch.bfloat16, device=XQ.device) + if TotalM == 0 or N == 0 or K == 0 or G == 0: + return output + + if os.environ.get("MSLK_AUTOTUNE_ENABLE"): + # FlyDSL's Autotuner discards the tuned function's return value, so read + # the result from the output buffer the kernel wrote into. + _get_autotuner()( + XQ, WQ, x_scale, w_scale, M_sizes, output, + _next_pow2(TotalM), N, K, b_preshuffled, + ) + return output + + tile_m, tile_n, tile_k = _DEFAULT_TILE + assert N % tile_n == 0, f"N={N} must be a multiple of tile_n={tile_n}" + assert K % tile_k == 0, f"K={K} must be a multiple of tile_k={tile_k}" + return _launch_kernel( + XQ, WQ, x_scale, w_scale, M_sizes, output, + tile_m=tile_m, tile_n=tile_n, tile_k=tile_k, b_preshuffled=b_preshuffled, + ) + + +def matmul_f8f8bf16_groupwise_grouped_preshuffle( + XQ: torch.Tensor, + WQ: torch.Tensor, + x_scale: torch.Tensor, + w_scale: torch.Tensor, + M_sizes: torch.Tensor, +) -> torch.Tensor: + """Preshuffled-B grouped groupwise FP8 GEMM (WQ already MFMA-preshuffled).""" + return _dispatch_grouped_gemm( + XQ, WQ, x_scale, w_scale, M_sizes, b_preshuffled=True + ) + + +def matmul_f8f8bf16_groupwise_grouped( + XQ: torch.Tensor, + WQ: torch.Tensor, + x_scale: torch.Tensor, + w_scale: torch.Tensor, + M_sizes: torch.Tensor, +) -> torch.Tensor: + """Plain (non-preshuffled) grouped groupwise FP8 GEMM via FlyDSL. + + Same contract as the preshuffle sibling, but WQ is plain row-major + ``[G, N, K]``. Uses the shared kernel with ``b_preshuffled=False``, which + stages B through LDS instead of loading it straight to registers. + """ + return _dispatch_grouped_gemm( + XQ, WQ, x_scale, w_scale, M_sizes, b_preshuffled=False + ) + + +if is_flydsl_available() and torch.version.hip is not None and hasattr(torch.ops, "mslk"): + # FlyDSL supplies the ROCm implementation of both ops; their schemas are + # declared in csrc/gemm/gemm_ops.cpp. Skip an op whose schema is missing, as + # in a python-only build, and tolerate a repeat import rebinding it. + def _register(op_name, cuda_fn, meta_fn=None) -> None: + if not hasattr(torch.ops.mslk, op_name.split("::")[1]): + return + try: + torch.library.impl(op_name, "CUDA")(cuda_fn) + if meta_fn is not None: + torch.library.impl(op_name, "Meta")(meta_fn) + except RuntimeError: + pass + + _register( + _OP_NAME, + matmul_f8f8bf16_groupwise_grouped_preshuffle, + _f8f8bf16_groupwise_grouped_preshuffle_meta, + ) + _register("mslk::f8f8bf16_groupwise_grouped", matmul_f8f8bf16_groupwise_grouped) diff --git a/mslk/gemm/triton/fp8_groupwise_grouped_gemm.py b/mslk/gemm/triton/fp8_groupwise_grouped_gemm.py index 2e699434..127f2672 100644 --- a/mslk/gemm/triton/fp8_groupwise_grouped_gemm.py +++ b/mslk/gemm/triton/fp8_groupwise_grouped_gemm.py @@ -287,8 +287,18 @@ def matmul_f8f8bf16_groupwise_grouped( # The C++ schema is declared in gemm_ops.cpp (shared between CUDA and ROCm). # On ROCm the CUDA C++ implementation is absent; this module registers the # Triton kernel above as the dispatch target via torch.library.impl. - -if torch.version.hip is not None and hasattr(torch.ops, "mslk"): +# +# FALLBACK ONLY: FlyDSL owns this op where it is available (see +# mslk/gemm/flydsl/fp8_groupwise_grouped_gemm.py), so Triton registers it only +# when FlyDSL is absent. Only one CUDA implementation can win, hence the explicit +# gate rather than a dependence on import order. +from mslk.flydsl.common import is_flydsl_available as _is_flydsl_available + +if ( + torch.version.hip is not None + and hasattr(torch.ops, "mslk") + and not _is_flydsl_available() +): if hasattr(torch.ops.mslk, "f8f8bf16_groupwise_grouped"): try: diff --git a/mslk/quantize/shuffle.py b/mslk/quantize/shuffle.py index 423bb87e..2d004092 100644 --- a/mslk/quantize/shuffle.py +++ b/mslk/quantize/shuffle.py @@ -205,3 +205,42 @@ def ck_preshuffle(src: torch.Tensor, NXdl: int = 16) -> torch.Tensor: # Reshape to original input shape. dst = dst.reshape(N, K) return dst + + +def preshuffle_b_mfma(src: torch.Tensor) -> torch.Tensor: + """Swizzle FP8 weights into the layout MFMA A8 GEMM kernels load B from. + + Rearranges the trailing ``(N, K)`` of ``src`` into + ``(N0, K0, KLane, NLane, KPack)`` order, which places each wave's 16x64 B + fragment contiguously so it can be read without a transpose. A leading group + dimension is preserved, so both a single weight ``[N, K]`` and per-group + weights ``[G, N, K]`` are accepted. + + The tile dimensions are fixed at ``NLane=16``, ``KLane=4``, ``KPack=16`` + because the consuming kernels build their B layout with those extents + hardcoded (see ``make_preshuffle_b_layout``); any other choice produces a + layout they cannot read. + + Args: + src: FP8 weights, ``[N, K]`` or ``[G, N, K]``. + + Returns: + The swizzled tensor, same shape and dtype as ``src``. + """ + NLane, KLane, KPack = 16, 4, 16 + N, K = src.shape[-2], src.shape[-1] + if N % NLane or K % (KLane * KPack): + raise ValueError( + f"preshuffle_b_mfma needs N % {NLane} == 0 and K % {KLane * KPack} == 0, " + f"got N={N} K={K}" + ) + K0 = K // (KLane * KPack) + lead = src.shape[:-2] + src = src.reshape(*lead, N // NLane, NLane, K0, KLane, KPack) + ndim = src.ndim + # Permute only the trailing 5 dims: (N0, NLane, K0, KLane, KPack) -> + # (N0, K0, KLane, NLane, KPack), keeping any leading group dim in place. + lead_axes = tuple(range(ndim - 5)) + n0, nlane, k0, klane, kpack = range(ndim - 5, ndim) + dst = src.permute(*lead_axes, n0, k0, klane, nlane, kpack).contiguous() + return dst.reshape(*lead, N, K) diff --git a/test/gemm/gemm_test.py b/test/gemm/gemm_test.py index 3c4643e1..6318c633 100644 --- a/test/gemm/gemm_test.py +++ b/test/gemm/gemm_test.py @@ -888,19 +888,32 @@ def test_f8f8bf16_groupwise(self, M: int, N: int, K: int) -> None: @parameterized.expand( [ - # (m_values, N, K) - ([128, 64], 128, 256), # small, 2 groups - ([512, 256, 128], 256, 512), # medium, 3 groups - ([1, 128, 256], 128, 256), # decode + prefill mix - ([2048, 1024], 512, 512), # large, 2 groups + (*case, preshuffle) + for preshuffle in [False, True] + for case in [ + # (m_values, N, K) + ([128, 64], 128, 256), # small, 2 groups + ([512, 256, 128], 256, 512), # medium, 3 groups + ([1, 128, 256], 128, 256), # decode + prefill mix + ([2048, 1024], 512, 512), # large, 2 groups + ] ] ) - def test_f8f8bf16_groupwise_grouped(self, m_values: list, N: int, K: int) -> None: + def test_f8f8bf16_groupwise_grouped( + self, m_values: list, N: int, K: int, preshuffle: bool + ) -> None: from mslk.quantize.triton.fp8_quantize import ( quantize_fp8_block, quantize_fp8_group, ) + # FlyDSL backs both variants on ROCm, and the preshuffle op exists only + # there, so skip when it is unavailable. + from mslk.flydsl.common import is_flydsl_available + + if (preshuffle or torch.version.hip is not None) and not is_flydsl_available(): + self.skipTest("FlyDSL not available") + G = len(m_values) device = self.device m_sizes = torch.tensor(m_values, dtype=torch.int64, device=device) @@ -925,9 +938,18 @@ def test_f8f8bf16_groupwise_grouped(self, m_values: list, N: int, K: int) -> Non # Quantize activations: xq [TotalM, K], x_scale [K//128, TotalM]. xq, x_scale = quantize_fp8_group(x, m_sizes=m_sizes) - out = torch.ops.mslk.f8f8bf16_groupwise_grouped( - xq, wq, x_scale, w_scale, m_sizes - ) + if preshuffle: + from mslk.quantize.shuffle import preshuffle_b_mfma + + # The preshuffle op consumes weights already in the MFMA B layout; + # callers shuffle once at prep time. + out = torch.ops.mslk.f8f8bf16_groupwise_grouped_preshuffle( + xq, preshuffle_b_mfma(wq), x_scale, w_scale, m_sizes + ) + else: + out = torch.ops.mslk.f8f8bf16_groupwise_grouped( + xq, wq, x_scale, w_scale, m_sizes + ) # BF16 reference: compute per group and concatenate. ref_parts = [] @@ -939,7 +961,7 @@ def test_f8f8bf16_groupwise_grouped(self, m_values: list, N: int, K: int) -> Non self.assertFalse(out.isnan().any().item(), "Output contains NaN") self.assertFalse(out.isinf().any().item(), "Output contains Inf") - torch.testing.assert_close(out, ref, atol=8.0e-2, rtol=8.0e-2) + torch.testing.assert_close(out, ref, atol=1.0e-2, rtol=4.0e-2) @skipUnlessCuda()