diff --git a/bench/gemm/gemm_ops.py b/bench/gemm/gemm_ops.py index 9f93c163..6c27cf8d 100644 --- a/bench/gemm/gemm_ops.py +++ b/bench/gemm/gemm_ops.py @@ -9,16 +9,11 @@ from enum import auto, Enum import torch + +import mslk.gemm # noqa: F401 — ensure mslk namespace exists from mslk.bench.common.utils import BenchOptions, do_bench from mslk.flydsl.common import is_flydsl_available from mslk.gemm.triton.fp8_gemm import matmul_fp8_block, matmul_fp8_row, to_mxfp8 - -if is_flydsl_available(): - from mslk.gemm.flydsl.preshuffle_gemm import ( - flydsl_preshuffle, - flydsl_preshuffle_batched_gemm, - flydsl_preshuffle_gemm, - ) from mslk.gemm.triton.grouped_gemm import grouped_gemm, grouped_gemm_fp8_rowwise from mslk.quantize.shuffle import ( ck_preshuffle, @@ -42,6 +37,20 @@ ) from mslk.utils.device import is_cuda, is_gfx942, is_gfx950, is_rocm +if is_flydsl_available(): + from mslk.gemm.flydsl.preshuffle_gemm import ( # noqa: E402 + flydsl_preshuffle, + flydsl_preshuffle_batched_gemm, + flydsl_preshuffle_gemm, + ) + +_STUB_SCHEMAS = [ + ("f8i4bf16_rowwise", "(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor w_zp) -> Tensor"), +] +for _name, _schema in _STUB_SCHEMAS: + if not hasattr(torch.ops, "mslk") or not hasattr(torch.ops.mslk, _name): + torch.library.define(f"mslk::{_name}", _schema) + try: from tinygemm.utils import group_quantize_tensor @@ -1956,6 +1965,23 @@ def weight_bytes_per_element(self) -> float: return 0.5 +@register_gemm_op +class TritonFP8Int4Rowwise(CutlassFP8Int4Rowwise): + """ROCm Triton FP8xINT4 rowwise GEMM.""" + + def compute(self, xq, wq, x_scale, w_scale, w_zp): + from mslk.gemm.triton.f8i4bf16_rowwise_gemm import matmul_f8i4bf16_rowwise + + return matmul_f8i4bf16_rowwise(xq, wq, x_scale, w_scale, w_zp) + + def quantize_and_compute(self, xq, wq, x_scale, w_scale, w_zp): + return self.compute(xq, wq, x_scale, w_scale, w_zp) + + @property + def supported_accelerators(self) -> set[Accelerator]: + return {Accelerator.AMD_GFX942, Accelerator.AMD_GFX950} + + @register_gemm_op class TinyGemmBF16Int4Groupwise(GemmOpBase): """ diff --git a/mslk/gemm/__init__.py b/mslk/gemm/__init__.py index 95c5eb9a..4c738d07 100644 --- a/mslk/gemm/__init__.py +++ b/mslk/gemm/__init__.py @@ -39,6 +39,7 @@ # module, which overrides the default (non-existent) CUDA impl so that # torch.ops.mslk.* dispatches to the Triton kernel on AMD. from .triton import ( # noqa: F401 + f8i4bf16_rowwise_gemm as _f8i4bf16_rowwise_gemm, fp8_groupwise_gemm, fp8_groupwise_grouped_gemm, grouped_gemm as _grouped_gemm, diff --git a/mslk/gemm/triton/f8i4bf16_rowwise_gemm.py b/mslk/gemm/triton/f8i4bf16_rowwise_gemm.py new file mode 100644 index 00000000..0fd066a8 --- /dev/null +++ b/mslk/gemm/triton/f8i4bf16_rowwise_gemm.py @@ -0,0 +1,159 @@ +# 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 x INT4 weight-only GEMM for ROCm/AMD GPUs, producing BF16 output. + +Thin wrapper around the unified _bf16i4_rowwise_kernel (int4_gemm.py). +The kernel is shared with the BF16xINT4 path; the only differences are: + 1. Activations are FP8 — byte-split and reinterpreted as FP8 TensorWrappers + here, then upcast to bfloat16 inside the kernel via .to(tl.bfloat16). + 2. HAS_X_SCALE=True — the kernel multiplies the accumulator by a per-row + activation dequant scale before storing to workspace. + +Op signature (mslk::f8i4bf16_rowwise): + XQ : [M, K] FP8 activations (float8_e4m3fnuz on AMD) + WQ : [N, K//2] int8 packed INT4 (lo nibble = even K, hi = odd K) + x_scale : [M] float32 per-row activation dequant scale + w_scale : [num_groups, N] float32/bf16/fp16 per-group weight scale + w_zp : [num_groups, N] same dtype as w_scale, per-group zero point + output : [M, N] bfloat16 +""" + +import torch +import triton # @manual +from mslk.gemm.triton.int4_gemm import ( + _bf16i4_rowwise_kernel, + _bf16i4_splitk_reduce, + _MAX_SPLIT_K, + _TL_CAT_HAS_DIM, +) +from mslk.utils.triton.fp8_utils import get_fp8_constants, reinterpret_fp8_type + + +def matmul_f8i4bf16_rowwise( + XQ: torch.Tensor, + WQ: torch.Tensor, + x_scale: torch.Tensor, + w_scale: torch.Tensor, + w_zp: torch.Tensor, +) -> torch.Tensor: + """ + FP8 activation x INT4 weight GEMM with per-row x_scale and per-group w_scale/w_zp. + + Args: + XQ : [M, K] FP8 activations (float8_e4m3fnuz on AMD) + WQ : [N, K//2] int8 packed INT4 (lo nibble=even K, hi=odd K) + x_scale : [M] float32 per-row activation dequant scale + w_scale : [num_groups, N] float32/bf16/fp16 per-group weight scale + w_zp : [num_groups, N] same dtype, per-group weight zero point + + Returns: + Y : [M, N] bfloat16 + """ + assert XQ.ndim == 2, f"XQ must be 2D [M, K], got shape {XQ.shape}" + M, K = XQ.shape + N = WQ.shape[0] + K2 = K // 2 + num_groups = w_scale.shape[0] + group_size = K // num_groups + + assert WQ.shape == (N, K2), f"WQ must be [N, K//2], got {WQ.shape}" + assert x_scale.shape == (M,), f"x_scale must be [M]={M}, got {x_scale.shape}" + assert w_scale.shape == (num_groups, N), ( + f"w_scale must be [num_groups, N]={num_groups, N}, got {w_scale.shape}" + ) + assert w_zp.shape == (num_groups, N), ( + f"w_zp must be [num_groups, N]={num_groups, N}, got {w_zp.shape}" + ) + assert group_size % 64 == 0, ( + f"group_size={group_size} must be divisible by 64 (2 * BLOCK_K_min=32)" + ) + + _, tl_fp8_dtype, _, _ = get_fp8_constants() + XQ_int8 = XQ.view(torch.int8) + x_even_t = XQ_int8[:, 0::2].contiguous() + x_odd_t = XQ_int8[:, 1::2].contiguous() + x_even = reinterpret_fp8_type(x_even_t, tl_fp8_dtype) + x_odd = reinterpret_fp8_type(x_odd_t, tl_fp8_dtype) + + if w_scale.dtype != torch.float32: + w_scale = w_scale.to(torch.float32) + if w_zp.dtype != torch.float32: + w_zp = w_zp.to(torch.float32) + if x_scale.dtype != torch.float32: + x_scale = x_scale.to(torch.float32) + + def grid(meta): + return ( + triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]), + meta["SPLIT_K"], + ) + + workspace = torch.empty((_MAX_SPLIT_K, M, N), dtype=torch.float32, device=XQ.device) + + _bf16i4_rowwise_kernel[grid]( + x_even, + x_odd, + WQ, + workspace, + w_scale, + w_zp, + x_scale, + M, + N, + K2, + group_size, + x_even.stride(0), + x_even.stride(1), + WQ.stride(0), + WQ.stride(1), + M * N, + N, + 1, + w_scale.stride(0), + w_scale.stride(1), + FUSE_DOT=_TL_CAT_HAS_DIM, + HAS_X_SCALE=True, + ) + + split_k = _bf16i4_rowwise_kernel.best_config.kwargs["SPLIT_K"] + + if split_k == 1: + return workspace[0].to(torch.bfloat16) + + Y_bf16 = torch.empty((M, N), dtype=torch.bfloat16, device=XQ.device) + reduce_grid = (triton.cdiv(M, 32), triton.cdiv(N, 32)) + _bf16i4_splitk_reduce[reduce_grid]( + workspace, + Y_bf16, + M, + N, + SPLIT_K=split_k, + BLOCK_M=32, + BLOCK_N=32, + ) + return Y_bf16 + + +# --------------------------------------------------------------------------- +# Register as ROCm implementation of mslk::f8i4bf16_rowwise +# --------------------------------------------------------------------------- + +if torch.version.hip is not None and hasattr(torch.ops, "mslk"): + if hasattr(torch.ops.mslk, "f8i4bf16_rowwise"): + + @torch.library.impl("mslk::f8i4bf16_rowwise", "CUDA") + def _f8i4bf16_rowwise_rocm( + XQ: torch.Tensor, + WQ: torch.Tensor, + x_scale: torch.Tensor, + w_scale: torch.Tensor, + w_zp: torch.Tensor, + ) -> torch.Tensor: + return matmul_f8i4bf16_rowwise(XQ, WQ, x_scale, w_scale, w_zp) diff --git a/mslk/gemm/triton/int4_gemm.py b/mslk/gemm/triton/int4_gemm.py index 992c7fa0..cb83384d 100644 --- a/mslk/gemm/triton/int4_gemm.py +++ b/mslk/gemm/triton/int4_gemm.py @@ -93,12 +93,14 @@ def _prune_configs( M = named_args["M"] N = named_args["N"] K2 = named_args["K2"] + has_x_scale = named_args.get("HAS_X_SCALE", False) pruned = [] for c in configs: bm = c.kwargs["BLOCK_M"] bn = c.kwargs["BLOCK_N"] bk = c.kwargs["BLOCK_K"] sk = c.kwargs.get("SPLIT_K", 1) + gsm = c.kwargs.get("GROUP_SIZE_M", 4) if group_size % (2 * bk) != 0: continue if K2 % (bk * sk) != 0: @@ -107,6 +109,15 @@ def _prune_configs( continue if bm > max(M, 32) or bn > max(N, 32): continue + if has_x_scale: + if bn < 128 and N >= 128: + continue + if bk > 64: + continue + if sk == 2: + continue + if gsm != 4: + continue pruned.append(c) return pruned @@ -132,12 +143,13 @@ def _prune_configs( ) @triton.jit def _bf16i4_rowwise_kernel( - X_even_ptr, # [M, K//2] bfloat16 — even K columns of activations - X_odd_ptr, # [M, K//2] bfloat16 — odd K columns of activations + X_even_ptr, # [M, K//2] bfloat16 or FP8 — even K columns of activations + X_odd_ptr, # [M, K//2] bfloat16 or FP8 — odd K columns of activations W_ptr, # [N, K//2] int8 packed (lo nibble = even K, hi = odd K) - Y_ptr, # [M, N] bfloat16 + Y_ptr, # [SPLIT_K, M, N] float32 workspace scale_ptr, # [num_groups, N] zero_ptr, # [num_groups, N] + x_scale_ptr, # [M] float32 per-row activation scale (used when HAS_X_SCALE) M, N, K2, # K // 2 @@ -160,6 +172,7 @@ def _bf16i4_rowwise_kernel( EVEN_MN: tl.constexpr, GRID_MN: tl.constexpr, FUSE_DOT: tl.constexpr, + HAS_X_SCALE: tl.constexpr = False, ) -> None: """ Computes Y[M, N] = X[M, K] @ dequant(W)[K, N]. @@ -232,10 +245,10 @@ def _bf16i4_rowwise_kernel( if EVEN_MN and EVEN_K: x_even = tl.load( X_even_ptr + offs_m[:, None] * stride_xm + offs_k2[None, :] * stride_xk, - ).to(tl.bfloat16) + ) x_odd = tl.load( X_odd_ptr + offs_m[:, None] * stride_xm + offs_k2[None, :] * stride_xk, - ).to(tl.bfloat16) + ) w_q = tl.load( W_ptr + offs_n[:, None] * stride_wn + offs_k2[None, :] * stride_wk, ).to(tl.int32) @@ -245,12 +258,12 @@ def _bf16i4_rowwise_kernel( X_even_ptr + offs_m[:, None] * stride_xm + offs_k2[None, :] * stride_xk, mask=k_mask, other=0.0, - ).to(tl.bfloat16) + ) x_odd = tl.load( X_odd_ptr + offs_m[:, None] * stride_xm + offs_k2[None, :] * stride_xk, mask=k_mask, other=0.0, - ).to(tl.bfloat16) + ) w_q = tl.load( W_ptr + offs_n[:, None] * stride_wn + offs_k2[None, :] * stride_wk, mask=k_mask, @@ -263,12 +276,12 @@ def _bf16i4_rowwise_kernel( X_even_ptr + offs_m[:, None] * stride_xm + offs_k2[None, :] * stride_xk, mask=xk_mask, other=0.0, - ).to(tl.bfloat16) + ) x_odd = tl.load( X_odd_ptr + offs_m[:, None] * stride_xm + offs_k2[None, :] * stride_xk, mask=xk_mask, other=0.0, - ).to(tl.bfloat16) + ) w_q = tl.load( W_ptr + offs_n[:, None] * stride_wn + offs_k2[None, :] * stride_wk, mask=wk_mask, @@ -281,17 +294,19 @@ def _bf16i4_rowwise_kernel( X_even_ptr + offs_m[:, None] * stride_xm + offs_k2[None, :] * stride_xk, mask=xk_mask, other=0.0, - ).to(tl.bfloat16) + ) x_odd = tl.load( X_odd_ptr + offs_m[:, None] * stride_xm + offs_k2[None, :] * stride_xk, mask=xk_mask, other=0.0, - ).to(tl.bfloat16) + ) w_q = tl.load( W_ptr + offs_n[:, None] * stride_wn + offs_k2[None, :] * stride_wk, mask=wk_mask, other=0, ).to(tl.int32) + x_even = x_even.to(tl.bfloat16) + x_odd = x_odd.to(tl.bfloat16) group_idx = (k2_slice_start * 2) // group_size if EVEN_MN: s = tl.load( @@ -332,12 +347,12 @@ def _bf16i4_rowwise_kernel( X_even_ptr + offs_m[:, None] * stride_xm + next_offs_k2[None, :] * stride_xk, - ).to(tl.bfloat16) + ) next_x_odd = tl.load( X_odd_ptr + offs_m[:, None] * stride_xm + next_offs_k2[None, :] * stride_xk, - ).to(tl.bfloat16) + ) next_w_q = tl.load( W_ptr + offs_n[:, None] * stride_wn + next_offs_k2[None, :] * stride_wk, ).to(tl.int32) @@ -349,14 +364,14 @@ def _bf16i4_rowwise_kernel( + next_offs_k2[None, :] * stride_xk, mask=next_k_mask, other=0.0, - ).to(tl.bfloat16) + ) next_x_odd = tl.load( X_odd_ptr + offs_m[:, None] * stride_xm + next_offs_k2[None, :] * stride_xk, mask=next_k_mask, other=0.0, - ).to(tl.bfloat16) + ) next_w_q = tl.load( W_ptr + offs_n[:, None] * stride_wn + next_offs_k2[None, :] * stride_wk, mask=next_k_mask, @@ -371,14 +386,14 @@ def _bf16i4_rowwise_kernel( + next_offs_k2[None, :] * stride_xk, mask=next_xk_mask, other=0.0, - ).to(tl.bfloat16) + ) next_x_odd = tl.load( X_odd_ptr + offs_m[:, None] * stride_xm + next_offs_k2[None, :] * stride_xk, mask=next_xk_mask, other=0.0, - ).to(tl.bfloat16) + ) next_w_q = tl.load( W_ptr + offs_n[:, None] * stride_wn + next_offs_k2[None, :] * stride_wk, mask=next_wk_mask, @@ -393,14 +408,14 @@ def _bf16i4_rowwise_kernel( + next_offs_k2[None, :] * stride_xk, mask=next_xk_mask, other=0.0, - ).to(tl.bfloat16) + ) next_x_odd = tl.load( X_odd_ptr + offs_m[:, None] * stride_xm + next_offs_k2[None, :] * stride_xk, mask=next_xk_mask, other=0.0, - ).to(tl.bfloat16) + ) next_w_q = tl.load( W_ptr + offs_n[:, None] * stride_wn + next_offs_k2[None, :] * stride_wk, mask=next_wk_mask, @@ -424,7 +439,20 @@ def _bf16i4_rowwise_kernel( mask=offs_n_raw < N, ).to(tl.float32) - if FUSE_DOT: + if HAS_X_SCALE: + acc = tl.dot( + x_even, + tl.trans(w_lo_dq.to(tl.bfloat16)), + acc, + out_dtype=tl.float32, + ) + acc = tl.dot( + x_odd, + tl.trans(w_hi_dq.to(tl.bfloat16)), + acc, + out_dtype=tl.float32, + ) + elif FUSE_DOT: # `dim` exists only on newer triton (guarded by FUSE_DOT/_TL_CAT_HAS_DIM); # the pinned stable stub lacks it, so Pyre flags the keyword. x_fused = tl.cat(x_even, x_odd, dim=1) # pyre-ignore[28] @@ -448,6 +476,9 @@ def _bf16i4_rowwise_kernel( out_dtype=tl.float32, ) + next_x_even = next_x_even.to(tl.bfloat16) + next_x_odd = next_x_odd.to(tl.bfloat16) + # Rotate buffers x_even = next_x_even x_odd = next_x_odd @@ -455,6 +486,16 @@ def _bf16i4_rowwise_kernel( s = next_s z = next_z + # ---- per-row activation scale (FP8 path) ---- + if HAS_X_SCALE: + if EVEN_MN: + xs = tl.load(x_scale_ptr + offs_m).to(tl.float32) + else: + xs = tl.load(x_scale_ptr + offs_m, mask=offs_m < M, other=1.0).to( + tl.float32 + ) + acc = acc * xs[:, None] + # ---- store output ---- # Write partial sum for this K-slice into workspace[pid_z, m, n]. # When SPLIT_K=1, pid_z=0 and this is a direct store to [M, N]. @@ -565,6 +606,7 @@ def grid(meta): # workspace[pid_z]. Allocated with empty (no zeroing) since the reduction kernel # only reads slices [0:split_k] — unwritten slices beyond split_k are never touched. workspace = torch.empty((_MAX_SPLIT_K, M, N), dtype=torch.float32, device=X.device) + _dummy_x_scale = torch.empty(1, dtype=torch.float32, device=X.device) _bf16i4_rowwise_kernel[grid]( x_even, @@ -573,6 +615,7 @@ def grid(meta): workspace, w_scale_group, w_zero_group, + _dummy_x_scale, M, N, K2, @@ -587,6 +630,7 @@ def grid(meta): w_scale_group.stride(0), w_scale_group.stride(1), FUSE_DOT=_TL_CAT_HAS_DIM, + HAS_X_SCALE=False, ) split_k = _bf16i4_rowwise_kernel.best_config.kwargs["SPLIT_K"]