From 7d2a3487fe2a334e9c6a9c376351dc3458ac9d7c Mon Sep 17 00:00:00 2001 From: Benson Ma Date: Mon, 24 Aug 2026 13:33:28 -0700 Subject: [PATCH] =?UTF-8?q?Add=20FP8=C3=97INT4=20rowwise=20GEMM=20for=20RO?= =?UTF-8?q?Cm=20(f8i4bf16=5Frowwise)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: NOTE: No linked task. Please associate a task with this diff. Adds a ROCm Triton implementation of `mslk::f8i4bf16_rowwise` for FP8 activations and packed INT4 weights while reusing the existing BF16×INT4 kernel. - Defines the operator schema on both CUDA and ROCm so the HIP Python implementation can register reliably. - Specializes the shared kernel with `HAS_X_SCALE`, two BF16 dot operations for FP8-backed inputs, and fused per-row activation scaling. - Adds FP8-specific autotune candidates and pruning while preserving the prior BF16 tuning space. - Validates tensor shapes, dtypes, devices, grouping invariants, and contiguous scale layouts. - Avoids allocating unused split-K workspace for large-M shapes. - Registers the ROCm benchmark implementation and adds a ROCm correctness/dispatch regression test. Upstream PR: https://github.com/meta-pytorch/MSLK/pull/495 Reviewed By: jwfromm Differential Revision: D116965755 --- bench/gemm/gemm_ops.py | 17 ++ csrc/gemm/gemm_ops.cpp | 4 +- mslk/gemm/__init__.py | 1 + mslk/gemm/triton/f8i4bf16_rowwise_gemm.py | 180 ++++++++++++++++++++++ mslk/gemm/triton/int4_gemm.py | 113 ++++++++++---- test/gemm/gemm_test.py | 31 ++++ 6 files changed, 311 insertions(+), 35 deletions(-) create mode 100644 mslk/gemm/triton/f8i4bf16_rowwise_gemm.py diff --git a/bench/gemm/gemm_ops.py b/bench/gemm/gemm_ops.py index 9f93c163..669d8873 100644 --- a/bench/gemm/gemm_ops.py +++ b/bench/gemm/gemm_ops.py @@ -1956,6 +1956,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/csrc/gemm/gemm_ops.cpp b/csrc/gemm/gemm_ops.cpp index a7d7ed6b..bd236067 100644 --- a/csrc/gemm/gemm_ops.cpp +++ b/csrc/gemm/gemm_ops.cpp @@ -78,6 +78,8 @@ TORCH_LIBRARY_FRAGMENT(mslk, m) { "bf16i4bf16_shuffled_grouped(Tensor X, Tensor WQ, Tensor w_scale_group, Tensor w_zero_group, Tensor M_sizes) -> Tensor"); m.def( "bf16i4bf16_shuffled_batched(Tensor X, Tensor WQ, Tensor w_scale, Tensor w_zp) -> Tensor"); + m.def( + "f8i4bf16_rowwise(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor w_zp) -> Tensor"); m.def("preshuffle_i4(Tensor WQ, Tensor w_scale) -> (Tensor, Tensor)"); #ifdef USE_ROCM // Sibling of f8f8bf16_groupwise_grouped taking weights already swizzled into @@ -113,8 +115,6 @@ TORCH_LIBRARY_FRAGMENT(mslk, m) { m.def( "f4f4bf16_ultra_grouped_mm(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor offsets, Tensor x_global_scale, Tensor w_global_scale, Tensor(a!)? output=None) -> Tensor"); m.def("bf16x9_gemm(Tensor A, Tensor B, Tensor(a!)? output=None) -> Tensor"); - m.def( - "f8i4bf16_rowwise(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor w_zp) -> Tensor"); m.def( "f8i4bf16_shuffled(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor w_scale_group) -> Tensor"); m.def( 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..4cf49b61 --- /dev/null +++ b/mslk/gemm/triton/f8i4bf16_rowwise_gemm.py @@ -0,0 +1,180 @@ +# 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}" + assert WQ.ndim == 2, f"WQ must be 2D [N, K//2], got shape {WQ.shape}" + assert x_scale.ndim == 1, f"x_scale must be 1D [M], got shape {x_scale.shape}" + assert w_scale.ndim == 2, ( + f"w_scale must be 2D [num_groups, N], got shape {w_scale.shape}" + ) + assert w_zp.ndim == 2, f"w_zp must be 2D [num_groups, N], got shape {w_zp.shape}" + M, K = XQ.shape + N = WQ.shape[0] + K2 = K // 2 + num_groups = w_scale.shape[0] + assert K % 2 == 0, f"K={K} must be even for packed INT4 weights" + + 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 WQ.dtype == torch.int8, f"WQ must be int8, got {WQ.dtype}" + assert XQ.is_contiguous(), "XQ must be contiguous" + assert WQ.is_contiguous(), "WQ must be contiguous" + assert all(tensor.device == XQ.device for tensor in (WQ, x_scale, w_scale, w_zp)), ( + "all inputs must be on the same device" + ) + + pt_fp8_dtype, tl_fp8_dtype, _, _ = get_fp8_constants() + assert XQ.dtype == pt_fp8_dtype, f"XQ must be {pt_fp8_dtype}, got {XQ.dtype}" + + if M == 0 or N == 0 or K == 0: + return torch.zeros((M, N), dtype=torch.bfloat16, device=XQ.device) + + assert num_groups > 0, "w_scale must contain at least one quantization group" + assert K % num_groups == 0, f"K={K} must be divisible by num_groups={num_groups}" + group_size = K // num_groups + assert group_size % 64 == 0, ( + f"group_size={group_size} must be divisible by 64 (2 * BLOCK_K_min=32)" + ) + + 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) + + w_scale = w_scale.to(torch.float32).contiguous() + w_zp = w_zp.to(torch.float32).contiguous() + x_scale = x_scale.to(torch.float32).contiguous() + + def grid(meta): + return ( + triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]), + meta["SPLIT_K"], + ) + + workspace_splits = 1 if M >= 512 else _MAX_SPLIT_K + workspace = torch.empty( + (workspace_splits, 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, # pyre-ignore[6] + BLOCK_N=32, # pyre-ignore[6] + ) + 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..1130121f 100644 --- a/mslk/gemm/triton/int4_gemm.py +++ b/mslk/gemm/triton/int4_gemm.py @@ -69,20 +69,21 @@ def _get_configs() -> List[Config]: for bn in [32, 64, 128, 256]: for bk in [32, 64, 128]: for sk in [1, 2, 4, 8]: - for nw in _num_warps_for_tile(bm, bn): - configs.append( - Config( - { - "BLOCK_M": bm, - "BLOCK_N": bn, - "BLOCK_K": bk, - "GROUP_SIZE_M": 8, - "SPLIT_K": sk, - }, - num_warps=nw, - num_stages=2, + for gsm in [4, 8]: + for nw in _num_warps_for_tile(bm, bn): + configs.append( + Config( + { + "BLOCK_M": bm, + "BLOCK_N": bn, + "BLOCK_K": bk, + "GROUP_SIZE_M": gsm, + "SPLIT_K": sk, + }, + num_warps=nw, + num_stages=2, + ) ) - ) return configs @@ -93,12 +94,14 @@ def _prune_configs( M = named_args["M"] N = named_args["N"] K2 = named_args["K2"] + has_x_scale = kwargs.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 +110,17 @@ 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 + elif gsm != 8: + continue pruned.append(c) return pruned @@ -132,12 +146,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 +175,7 @@ def _bf16i4_rowwise_kernel( EVEN_MN: tl.constexpr, GRID_MN: tl.constexpr, FUSE_DOT: tl.constexpr, + HAS_X_SCALE: tl.constexpr, ) -> None: """ Computes Y[M, N] = X[M, K] @ dequant(W)[K, N]. @@ -232,10 +248,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 +261,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 +279,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 +297,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 +350,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 +367,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 +389,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 +411,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 +442,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 +479,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 +489,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 +609,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 +618,7 @@ def grid(meta): workspace, w_scale_group, w_zero_group, + _dummy_x_scale, M, N, K2, @@ -587,6 +633,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"] diff --git a/test/gemm/gemm_test.py b/test/gemm/gemm_test.py index 8db911d0..ce44540b 100644 --- a/test/gemm/gemm_test.py +++ b/test/gemm/gemm_test.py @@ -1365,6 +1365,37 @@ def test_torch_op_dispatch( torch.testing.assert_close(y_op, y_direct, atol=0.0, rtol=0.0) +@skipUnlessRocm() +class FP8Int4TritonROCmTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + from mslk.gemm.triton.f8i4bf16_rowwise_gemm import matmul_f8i4bf16_rowwise + + cls.matmul_rowwise = staticmethod(matmul_f8i4bf16_rowwise) + + def test_rowwise_accuracy_and_dispatch(self) -> None: + M, N, K, group_size = 1, 256, 1024, 128 + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") * 0.1 + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 0.01 + xq, x_scale = quantize_fp8_row(x) + wq_unpacked, w_scale, w_zp = int4_row_quantize(w, group_size) + wq = pack_int4(wq_unpacked).contiguous().to(device="cuda") + w_scale = w_scale.contiguous().to(device="cuda") + w_zp = w_zp.contiguous().to(device="cuda") + + y_direct = self.matmul_rowwise(xq, wq, x_scale, w_scale, w_zp) + y_op = torch.ops.mslk.f8i4bf16_rowwise(xq, wq, x_scale, w_scale, w_zp) + x_dequant = xq.float() * x_scale[:, None] + w_dequant = ( + wq_unpacked.reshape(N, -1, group_size).float() * w_scale.T[..., None] + + w_zp.T[..., None] + ).reshape(N, K) + y_ref = (x_dequant @ w_dequant.T).to(torch.bfloat16) + + torch.testing.assert_close(y_direct, y_ref, atol=8.0e-2, rtol=8.0e-2) + torch.testing.assert_close(y_op, y_direct, atol=0.0, rtol=0.0) + + @skipUnlessRocm() class BF16Int4TritonROCmGroupedTests(unittest.TestCase): """