mslk-cuda-nightly 2026.1.19__cp310-cp310-manylinux_2_28_x86_64.whl

mslk/bench/gemm/gemm_ops.py

@@ -0,0 +1,3342 @@
+# 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.
+
+# Keep a registry of all quantize operators.
+import abc
+import functools
+from enum import auto, Enum
+
+import numpy as np
+import torch
+from mslk.bench.common.utils import BenchOptions, do_bench
+from mslk.gemm.triton.fp8_gemm import matmul_fp8_block, matmul_fp8_row, to_mxfp8
+from mslk.gemm.triton.grouped_gemm import grouped_gemm, grouped_gemm_fp8_rowwise
+from mslk.quantize.shuffle import ck_preshuffle, quantize_int4_preshuffle
+from mslk.quantize.triton.fp4_quantize import (
+    _to_blocked,
+    calculate_group_max,
+    mega_fp4_pack,
+    mega_fp4_quantize_kernel,
+    mega_fp4_unpack,
+    triton_quantize_mx4_unpack,
+    triton_quantize_nvfp4,
+    triton_scale_nvfp4_quant_rms,
+    triton_scale_nvfp4_quant_silu,
+)
+from mslk.quantize.triton.fp8_quantize import (
+    quantize_fp8_block,
+    quantize_fp8_group,
+    quantize_fp8_row,
+    scale_fp8_row,
+    triton_quantize_fp8_row,
+)
+from mslk.utils.triton.fp8_utils import get_fp8_constants
+
+try:
+    from gen_ai.llm_inference.fb.llm.kernel.rms_norm import rms_norm
+    from gen_ai.llm_inference.fb.llm.kernel.silu_mul import silu_mul
+except ImportError:
+    # Above is used for some experiments, but the quantize is not relying on them. Okay to just skip.
+    pass
+
+try:
+    from tinygemm.utils import group_quantize_tensor
+
+    if torch.cuda.is_available() and torch.version.cuda:
+        torch.ops.load_library("//tinygemm:tinygemm")
+    TINYGEMM_ENABLED = True
+except ImportError:
+    TINYGEMM_ENABLED = False
+
+# Marlin currently only is supported only internally at Meta.
+try:
+    from marlin.quantize import marlin_quantize
+
+    torch.ops.load_library("//ai_codesign/gen_ai/marlin:marlin_ops")
+    MARLIN_ENABLED = True
+except ImportError:
+    MARLIN_ENABLED = False
+
+try:
+    from deep_gemm import (
+        gemm_fp8_fp8_bf16_nt,
+        get_col_major_tma_aligned_tensor,
+        m_grouped_gemm_fp8_fp8_bf16_nt_contiguous,
+        m_grouped_gemm_fp8_fp8_bf16_nt_masked,
+    )
+
+    DEEPGEMM_ENABLED = True
+except ImportError:
+    DEEPGEMM_ENABLED = False
+
+
+# Machete is also only supported internally at Meta for now.
+try:
+    from machete.machete import machete_gemm
+    from machete.quantize import machete_quantize_and_pack
+
+    MACHETE_ENABLED = True
+except ImportError:
+    MACHETE_ENABLED = False
+
+
+class Accelerator(Enum):
+    NVIDIA_SM90 = auto()
+    NVIDIA_SM100 = auto()
+    NVIDIA_SM103 = auto()
+    AMD_MI300X = auto()
+
+
+class GemmType(Enum):
+    REGULAR = auto()
+    GROUPED = auto()
+
+
+class ComputeDtype(Enum):
+    FP32 = auto()
+    TF32 = auto()
+    BF16 = auto()
+    FP8 = auto()
+    FP4 = auto()
+
+
+@functools.cache
+def get_current_accelerator() -> Accelerator | None:
+    if not torch.cuda.is_available():
+        raise Exception("Cannot run gemm_bench without accelerator.")
+
+    if torch.version.hip is not None:
+        device_name = torch.cuda.get_device_name()
+        if "MI300X" in device_name.upper():
+            return Accelerator.AMD_MI300X
+    elif torch.version.cuda is not None:
+        major, minor = torch.cuda.get_device_capability()
+        if major == 9 and minor == 0:
+            return Accelerator.NVIDIA_SM90
+        elif major == 10 and minor == 0:
+            return Accelerator.NVIDIA_SM100
+        elif major == 10 and minor == 3:
+            return Accelerator.NVIDIA_SM103
+
+    raise Exception("Cannot detect hardware that is supported by gemm_bench.")
+
+
+gemm_op_registry = []
+
+
+class GemmOpBase(metaclass=abc.ABCMeta):
+    """Helper abstract class to define expected methods of quantize ops."""
+
+    @abc.abstractmethod
+    def quantize(self, *args):
+        """Function which quantizes inputs."""
+        pass
+
+    @abc.abstractmethod
+    def compute(self, *args):
+        """Function which performs main compute operation."""
+        pass
+
+    @abc.abstractmethod
+    def quantize_and_compute(self, *args):
+        """Function which quantizes inputs and performs main compute operation."""
+        pass
+
+    def preprocess(self, *args):
+        """Preprocess inputs before benchmarking. These outputs will be passed to quantize."""
+        return args
+
+    def benchmark(
+        self,
+        *args,
+        opts: BenchOptions,
+        bench_quantize: bool,
+    ) -> float:
+        """Benchmark runtime of this operator."""
+        t = do_bench(
+            lambda *a: self.quantize_and_compute(*a)
+            if bench_quantize
+            else self.compute(*a),
+            args,
+            opts,
+        )
+        return t
+
+    @abc.abstractproperty
+    def name(self) -> str:
+        """Name of the operator."""
+        pass
+
+    @abc.abstractproperty
+    def supported_accelerators(self) -> set[Accelerator]:
+        pass
+
+    @property
+    def supported(self) -> bool:
+        """Whether this op will run on the current device."""
+        accelerator = get_current_accelerator()
+        return accelerator in self.supported_accelerators
+
+    @abc.abstractproperty
+    def supported_gemm_types(self) -> set[GemmType]:
+        pass
+
+    @abc.abstractproperty
+    def compute_dtype(self) -> ComputeDtype:
+        """The dtype used by tensor cores for the compute."""
+        pass
+
+
+def register_gemm_op(op):
+    """Decorator function for assembling all quantize ops."""
+    gemm_op_registry.append(op())
+    return op
+
+
+def get_gemm_ops() -> list[GemmOpBase]:
+    """Get all registered quantize ops."""
+    return gemm_op_registry
+
+
+@register_gemm_op
+class FP32Baseline(GemmOpBase):
+    """
+    FP32 matmul baseline.
+    """
+
+    def quantize(self, x, w):
+        if isinstance(x, list):
+            x = [i.float() for i in x]
+            w = [torch.transpose(i, -2, -1).float() for i in w]
+        else:
+            x = x.float()
+            w = torch.transpose(w, -2, -1).float()
+        return x, w
+
+    def compute(self, x, w):
+        # Handle both grouped and standard gemm.
+        if isinstance(x, list):
+            output = []
+            for i in range(len(x)):
+                output.append(torch.matmul(x[i], w[i]))
+            return output
+        return torch.matmul(x, w)
+
+    def quantize_and_compute(self, x, w):
+        return self.compute(*self.quantize(x, w))
+
+    @property
+    def name(self) -> str:
+        return "fp32_baseline"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return set(Accelerator)
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR, GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP32
+
+
+@register_gemm_op
+class TF32Baseline(GemmOpBase):
+    """
+    TF32 matmul baseline.
+    """
+
+    def quantize(self, x, w):
+        if isinstance(x, list):
+            x = [i.float() for i in x]
+            w = [torch.transpose(i, -2, -1).float() for i in w]
+        else:
+            x = x.float()
+            w = torch.transpose(w, -2, -1).float()
+        return x, w
+
+    def compute(self, x, w):
+        # Handle both grouped and standard gemm.
+        original_precision = torch.get_float32_matmul_precision()
+        torch.set_float32_matmul_precision("high")
+        if isinstance(x, list):
+            output = []
+            for i in range(len(x)):
+                output.append(torch.matmul(x[i], w[i]))
+            return output
+        out = torch.matmul(x, w)
+        torch.set_float32_matmul_precision(original_precision)
+        return out
+
+    def quantize_and_compute(self, x, w):
+        return self.compute(*self.quantize(x, w))
+
+    @property
+    def name(self) -> str:
+        return "tf32_baseline"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return set(Accelerator)
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR, GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.TF32
+
+
+@register_gemm_op
+class BF16Baseline(GemmOpBase):
+    """
+    Baseline BF16 matmul.
+    """
+
+    def quantize(self, x, w):
+        if isinstance(x, list):
+            x = [i.bfloat16() for i in x]
+            w = [torch.transpose(i, -2, -1).bfloat16() for i in w]
+        else:
+            x = x.bfloat16()
+            w = torch.transpose(w, -2, -1).bfloat16()
+        return x, w
+
+    def compute(self, x, w):
+        # Handle both grouped and standard gemm.
+        if isinstance(x, list):
+            output = []
+            for i in range(len(x)):
+                output.append(torch.matmul(x[i], w[i]))
+            return output
+        return torch.matmul(x, w)
+
+    def quantize_and_compute(self, x, w):
+        return self.compute(*self.quantize(x, w))
+
+    @property
+    def name(self) -> str:
+        return "bf16_baseline"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return set(Accelerator)
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR, GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class ScaledMMBaseline(GemmOpBase):
+    """
+    Reference FP8 matmul implemented in native torch with cublas or hipblas.
+    """
+
+    def __init__(self):
+        self.fp8_dtype, _, _, _ = get_fp8_constants()
+        self.E4M3_MAX_POS: float = torch.finfo(self.fp8_dtype).max
+        self.E5M2_MAX_POS: float = torch.finfo(torch.float8_e5m2).max
+        self.FP16_MAX_POS: float = torch.finfo(torch.float16).max
+        self.EPS: float = 1e-12
+        self.fast_accum = True
+
+    def _amax_to_scale(
+        self, amax: torch.Tensor, float8_dtype: torch.dtype, orig_dtype: torch.dtype
+    ) -> torch.Tensor:
+        # To make scale dtype to be fp32 for accuracy
+        amax = amax.float()
+        if float8_dtype == self.fp8_dtype:
+            # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
+            res = self.E4M3_MAX_POS / torch.clamp(amax, min=self.EPS)
+        else:  # e5m2
+            # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
+            res = self.E5M2_MAX_POS / torch.clamp(amax, min=self.EPS)
+
+        # pyre-fixme[7]: Expected `Tensor` but got `Union[float, Tensor]`.
+        return res
+
+    def _to_fp8_saturated(
+        self, x: torch.Tensor, float8_dtype: torch.dtype
+    ) -> torch.Tensor:
+        if float8_dtype == torch.float8_e4m3fn:
+            x = x.clamp(min=-1 * self.E4M3_MAX_POS, max=self.E4M3_MAX_POS)
+        else:
+            x = x.clamp(min=-1 * self.E5M2_MAX_POS, max=self.E5M2_MAX_POS)
+        return x.to(float8_dtype)
+
+    def _quantize_tensor(self, x):
+        x_amax = torch.max(torch.abs(x))
+        scale = self._amax_to_scale(x_amax, self.fp8_dtype, x.dtype)
+        scaled_x = self._to_fp8_saturated(x * scale, self.fp8_dtype)
+        x_inverse_scale = scale.reciprocal()
+        return scaled_x, x_inverse_scale
+
+    def quantize(self, x, w):
+        xq, x_scale = self._quantize_tensor(x)
+        wq, w_scale = self._quantize_tensor(w.t())
+        return xq, wq, x_scale, w_scale
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        output = torch._scaled_mm(
+            xq,
+            wq,
+            bias=None,
+            out_dtype=torch.bfloat16,
+            scale_a=x_scale,
+            scale_b=w_scale,
+            scale_result=None,
+            use_fast_accum=self.fast_accum,
+        )
+        return output
+
+    def quantize_and_compute(self, x, w):
+        return self.compute(*self.quantize(x, w))
+
+    @property
+    def name(self) -> str:
+        return "scaled_mm"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return set(Accelerator)
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class BF16X9Baseline(GemmOpBase):
+    """
+    FP32 matmul implemented with BF16X9 emulation.
+    """
+
+    def quantize(self, x, w):
+        if isinstance(x, list):
+            x = [i.float() for i in x]
+            w = [i.float() for i in w]
+        else:
+            x = x.float()
+            w = w.float()
+        return x, w
+
+    def compute(self, x, w):
+        # Handle both grouped and standard gemm.
+        if isinstance(x, list):
+            output = []
+            for i in range(len(x)):
+                output.append(torch.ops.mslk.bf16x9_gemm(x[i], w[i]))
+            return output
+        return torch.ops.mslk.bf16x9_gemm(x, w)
+
+    def quantize_and_compute(self, x, w):
+        return self.compute(*self.quantize(x, w))
+
+    @property
+    def name(self) -> str:
+        return "bf16x9_gemm"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.NVIDIA_SM100,
+            Accelerator.NVIDIA_SM103,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR, GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class ScaledMMRowwise(GemmOpBase):
+    def __init__(self):
+        self.fast_accum = True
+        self.torch_compile = False
+
+    def quantize(self, x, w):
+        xq, x_scale = quantize_fp8_row(x)
+        wq, w_scale = quantize_fp8_row(w)
+        return xq, wq.t(), x_scale.unsqueeze(1), w_scale.unsqueeze(0)
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        if self.torch_compile:
+            f = torch.compile(
+                torch._scaled_mm,
+                options={
+                    "max_autotune": True,
+                    "max_autotune_gemm_backends": "TRITON,CK,CUTLASS,ATEN",
+                },
+            )
+        else:
+            f = torch._scaled_mm
+
+        return f(
+            xq,
+            wq,
+            bias=None,
+            out_dtype=torch.bfloat16,
+            scale_a=x_scale,
+            scale_b=w_scale,
+            scale_result=None,
+            use_fast_accum=self.fast_accum,
+        )
+
+    def quantize_and_compute(self, x, w):
+        return self.compute(*self.quantize(x, w))
+
+    @property
+    def name(self) -> str:
+        return "scaled_mm_rowwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return set(Accelerator)
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class ScaledMMMXFP8(GemmOpBase):
+    def __init__(self):
+        self.torch_compile = False
+
+    def quantize(self, x, w):
+        x_scale, xq = to_mxfp8(x)
+        x_scale = _to_blocked(x_scale)
+        w_scale, wq = to_mxfp8(w)
+        w_scale = _to_blocked(w_scale)
+        return xq, wq.t(), x_scale, w_scale
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        if self.torch_compile:
+            f = torch.compile(
+                torch._scaled_mm,
+                options={
+                    "max_autotune": True,
+                    "max_autotune_gemm_backends": "TRITON,CK,CUTLASS,ATEN",
+                },
+            )
+        else:
+            f = torch._scaled_mm
+
+        return f(
+            xq,
+            wq,
+            bias=None,
+            out_dtype=torch.bfloat16,
+            scale_a=x_scale,
+            scale_b=w_scale,
+        )
+
+    def quantize_and_compute(self, x, w):
+        return self.compute(*self.quantize(x, w))
+
+    @property
+    def name(self) -> str:
+        return "scaled_mm_mxfp8"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM100, Accelerator.NVIDIA_SM103}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class ScaledMMNVFP4(GemmOpBase):
+    def __init__(self):
+        self.torch_compile = False
+
+    def quantize(self, x, w):
+        x_global_scale = torch.tensor([1.0], device=x.device, dtype=torch.float32)
+        w_global_scale = torch.tensor([1.0], device=w.device, dtype=torch.float32)
+
+        xq, x_scale = triton_quantize_nvfp4(x, x_global_scale)
+        wq, w_scale = triton_quantize_nvfp4(w, w_global_scale)
+
+        return (
+            xq.view(torch.float4_e2m1fn_x2),
+            wq.view(torch.float4_e2m1fn_x2),
+            x_scale.view(torch.float8_e4m3fn),
+            w_scale.view(torch.float8_e4m3fn),
+        )
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        if self.torch_compile:
+            f = torch.compile(
+                torch._scaled_mm,
+                options={
+                    "max_autotune": True,
+                    "max_autotune_gemm_backends": "TRITON,CK,CUTLASS,ATEN",
+                },
+            )
+        else:
+            f = torch._scaled_mm
+
+        return f(
+            xq,
+            wq.t(),
+            bias=None,
+            out_dtype=torch.bfloat16,
+            scale_a=x_scale,
+            scale_b=w_scale,
+        )
+
+    def quantize_and_compute(self, x, w):
+        return self.compute(*self.quantize(x, w))
+
+    @property
+    def name(self) -> str:
+        return "scaled_mm_nvfp4"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM100, Accelerator.NVIDIA_SM103}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP4
+
+
+@register_gemm_op
+class FP8TensorwiseGemm(GemmOpBase):
+    """
+    FP8 matmul with tensorwise scaling.
+    """
+
+    def quantize(self, x, w):
+        # Quantize both input tensors.
+        xq, x_scale = torch.ops.mslk.quantize_fp8_per_tensor(x)
+        wq, w_scale = torch.ops.mslk.quantize_fp8_per_tensor(w)
+        return xq, wq, x_scale, w_scale
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        return torch.ops.mslk.f8f8bf16(xq, wq, x_scale * w_scale)
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale = self.quantize(x, w)
+        return self.compute(xq, wq, x_scale, w_scale)
+
+    @property
+    def name(self) -> str:
+        return "cutlass_tensorwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM90}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class FP8CublasRowwiseGemm(GemmOpBase):
+    """
+    FP8 cublas matmul with rowwise scaling.
+    """
+
+    def quantize(self, x, w):
+        # Quantize both input tensors.
+        xq, x_scale = torch.ops.mslk.quantize_fp8_per_row(x)
+        wq, w_scale = torch.ops.mslk.quantize_fp8_per_row(w)
+        return xq, wq, x_scale, w_scale
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        out = torch.ops.mslk.f8f8bf16_cublas(xq, wq)
+        scaled_out = scale_fp8_row(out, x_scale, w_scale)
+        return scaled_out
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale = self.quantize(x, w)
+        return self.compute(xq, wq, x_scale, w_scale)
+
+    @property
+    def name(self) -> str:
+        return "cublas_rowwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.NVIDIA_SM100,
+            Accelerator.NVIDIA_SM103,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class FP8CublasTensorwiseGemm(GemmOpBase):
+    """
+    FP8 cublas matmul with tensorwise scaling.
+    """
+
+    def quantize(self, x, w):
+        # Quantize both input tensors.
+        xq, x_scale = torch.ops.mslk.quantize_fp8_per_tensor(x)
+        wq, w_scale = torch.ops.mslk.quantize_fp8_per_tensor(w)
+        return xq, wq, x_scale, w_scale
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        return torch.ops.mslk.f8f8bf16_cublas(xq, wq, x_scale * w_scale)
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale = self.quantize(x, w)
+        return self.compute(xq, wq, x_scale * w_scale)
+
+    @property
+    def name(self) -> str:
+        return "cublas_tensorwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.NVIDIA_SM100,
+            Accelerator.NVIDIA_SM103,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class FP8RowwiseGemm(GemmOpBase):
+    """
+    FP8 matmul with rowwise scaling.
+    """
+
+    def __init__(self):
+        self.fast_accum = True
+        self.gemm_op = torch.ops.mslk.f8f8bf16_rowwise
+
+    def preprocess(self, x, w):
+        # Prequantize weights.
+        if isinstance(w, (list, tuple)):
+            wq, w_scale = zip(*[quantize_fp8_row(i) for i in w])
+        else:
+            wq, w_scale = quantize_fp8_row(w)
+            if wq.dim() == 3:
+                w_scale = w_scale.view(wq.size(0), -1)
+        return x, wq, w_scale
+
+    def quantize(self, x, wq, w_scale):
+        # Quantize both input tensors.
+        # Handle both grouped and standard gemm.
+        if isinstance(x, (list, tuple)):
+            xq, x_scale = zip(*[quantize_fp8_row(i) for i in x])
+        else:
+            xq, x_scale = quantize_fp8_row(x)
+            # Set proper batch dimension shapes.
+            if xq.dim() == 3:
+                x_scale = x_scale.view(xq.size(0), -1)
+        return xq, wq, x_scale, w_scale
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        # Handle group gemm if inputs are grouped.
+        if isinstance(xq, (list, tuple)):
+            output = []
+            for i in range(len(xq)):
+                output.append(
+                    self.gemm_op(
+                        xq[i],
+                        wq[i],
+                        x_scale[i],
+                        w_scale[i],
+                        use_fast_accum=self.fast_accum,
+                    )
+                )
+            return output
+        # Unroll batched gemm if needed.
+        elif xq.dim() == 3 and wq.dim() == 3:
+            B, M, _ = xq.shape
+            _, N, _ = wq.shape
+            y = torch.empty((B, M, N), device=xq.device, dtype=torch.bfloat16)
+            for i in range(B):
+                y[i] = self.gemm_op(
+                    xq[i], wq[i], x_scale[i], w_scale[i], use_fast_accum=self.fast_accum
+                )
+            return y
+        # Otherwise return normal gemm result.
+        return self.gemm_op(xq, wq, x_scale, w_scale, use_fast_accum=self.fast_accum)
+
+    def quantize_and_compute(self, x, wq, w_scale):
+        xq, wq, x_scale, w_scale = self.quantize(x, wq, w_scale)
+        return self.compute(xq, wq, x_scale, w_scale)
+
+    @property
+    def name(self) -> str:
+        if torch.version.cuda:
+            return "cutlass_rowwise"
+        else:
+            return "ck_rowwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.NVIDIA_SM100,
+            Accelerator.NVIDIA_SM103,
+            Accelerator.AMD_MI300X,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class FP8RowwisePreshuffleGemm(FP8RowwiseGemm):
+    """
+    FP8 matmul with rowwise scaling and preshuffling of input B.
+    """
+
+    def __init__(self):
+        self.fast_accum = True
+        if self.supported:
+            self.gemm_op = torch.ops.mslk.f8f8bf16_rowwise_preshuffle
+
+    def preprocess(self, x, w):
+        x, wq, w_scale = super().preprocess(x, w)
+        return x, ck_preshuffle(wq, 16), w_scale
+
+    @property
+    def name(self) -> str:
+        return "ck_rowwise_preshuffle"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.AMD_MI300X}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class FP8RowwiseGroupedGemm(GemmOpBase):
+    """
+    FP8 grouped matmul with rowwise scaling.
+    """
+
+    def preprocess(self, x, w):
+        # Apply sparsity to inputs if appropriate.
+        # First check if N and K are fixed.
+        m_values = [i.shape[0] for i in x]
+        n_values = [i.shape[0] for i in w]
+        k_values = [i.shape[1] for i in w]
+        # If so, do specialized version of initialization.
+        if len(np.unique(n_values)) == 1 and len(np.unique(k_values)) == 1:
+            m_values = [i.shape[0] for i in x]
+            # Inputs for fixed nk mode must be contiguous, however in the benchmark
+            # script they typically are not. Do a little special processing to make them
+            # work. In practice this wont be needed.
+            # Start by padding along m dimension with zeros.
+            max_m = max(m_values)
+            x = [
+                torch.nn.functional.pad(i, (0, 0, 0, max_m - i.shape[0]), value=0)
+                for i in x
+            ]
+            # Stack inputs into groups.
+            x = torch.stack(x).contiguous()
+            w = torch.stack(w).contiguous()
+
+            # Preapply weight quantization.
+            wq, w_scale = quantize_fp8_row(w)
+            # Return processed tensors.
+            return (
+                x,
+                wq,
+                w_scale,
+                torch.tensor(m_values).to(dtype=torch.int64, device=x[0].device),
+            )
+        # Otherwise run without sparsity.
+        wq, w_scale = zip(*[quantize_fp8_row(i) for i in w])
+        return x, wq, w_scale, None
+
+    def quantize(self, x, wq, w_scale, m_values=None):
+        # Handle case where inputs are explicitly grouped and non-sparse.
+        if isinstance(x, (tuple, list)):
+            xq, x_scale = zip(*[triton_quantize_fp8_row(i) for i in x])
+            return xq, wq, x_scale, w_scale, m_values
+        # Otherwise inputs are unified tensors and sparse.
+        else:
+            B = x.shape[0]
+            xq, x_scale = triton_quantize_fp8_row(x, zero_start_index_M=m_values)
+            x_scale = x_scale.view(B, -1)
+            return xq, wq, x_scale, w_scale, m_values
+
+    def compute(self, xq, wq, x_scale, w_scale, m_values):
+        if m_values is None:
+            return torch.ops.mslk.f8f8bf16_rowwise_grouped(
+                xq,
+                wq,
+                x_scale,
+                w_scale,
+            )
+        else:
+            # Break tensor into groups, simulates what is done e2e.
+            return torch.ops.mslk.f8f8bf16_rowwise_grouped_dynamic(
+                xq,
+                wq,
+                x_scale,
+                w_scale,
+                zero_start_index_M=m_values,
+            )
+
+    def quantize_and_compute(self, x, wq, w_scale, m_values=None):
+        xq, wq, x_scale, w_scale, m_values = self.quantize(x, wq, w_scale, m_values)
+        return self.compute(xq, wq, x_scale, w_scale, m_values)
+
+    @property
+    def name(self) -> str:
+        if torch.version.cuda:
+            return "cutlass_rowwise_grouped"
+        else:
+            return "ck_rowwise_grouped"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.NVIDIA_SM100,
+            Accelerator.NVIDIA_SM103,
+            Accelerator.AMD_MI300X,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class BF16TritonStackedGroupedGemm(GemmOpBase):
+    """
+    BF16 grouped matmul with stacked inputs implemented with triton.
+    """
+
+    def preprocess(self, x, w):
+        m_values = [i.shape[0] for i in x]
+        # Convert m_values into offsets into grouped tensor.
+        m_sizes = torch.tensor(m_values).to(dtype=torch.int32, device=x[0].device)
+        w = torch.concat(w, dim=0).contiguous()
+        # Also view input as flattened.
+        x = torch.concat(x, dim=0).contiguous()
+        # Return processed tensors.
+        return x, w, m_sizes
+
+    def quantize(self, x, w, m_sizes):
+        return x, w, m_sizes
+
+    def compute(self, x, w, m_sizes):
+        return grouped_gemm(x, w, m_sizes, _use_warp_specialization=True)
+
+    def quantize_and_compute(self, x, w, m_sizes):
+        x, w, m_sizes = self.quantize(x, w, m_sizes)
+        return self.compute(x, w, m_sizes)
+
+    @property
+    def name(self) -> str:
+        return "triton_bf16_grouped_stacked"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return set(Accelerator)
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class BF16TritonStackedGroupedGemmFuseScatterAdd(BF16TritonStackedGroupedGemm):
+    """
+    BF16 grouped matmul with stacked inputs implemented with triton. Fused with ScatterAdd.
+    """
+
+    def preprocess(self, x, w):
+        x, w, m_sizes = super().preprocess(x, w)
+        M = x.shape[0]
+        N = w.shape[0] // m_sizes.shape[0]
+        output = torch.zeros(M, N, dtype=torch.bfloat16, device=x.device)
+        indices = torch.randperm(M, dtype=torch.int32, device=x.device)
+        return x, w, m_sizes, output, indices
+
+    def quantize(self, x, w, m_sizes, *args):
+        return *super().quantize(x, w, m_sizes), *args
+
+    def compute(self, x, w, m_sizes, output, indices):
+        return grouped_gemm(
+            x,
+            w,
+            m_sizes,
+            _use_warp_specialization=True,
+            _output_tensor=output,
+            _scatter_add_indices=indices,
+        )
+
+    def quantize_and_compute(self, x, w, m_sizes, *args):
+        x, w, m_sizes, *ret = self.quantize(x, w, m_sizes, *args)
+        return self.compute(x, w, m_sizes, *ret)
+
+    @property
+    def name(self) -> str:
+        return "triton_bf16_grouped_stacked_fuse_scatter_add"
+
+
+@register_gemm_op
+class FP8TritonStackedGroupedGemm(GemmOpBase):
+    """
+    FP8 grouped matmul with rowwise scaling and stacked inputs implemented with triton.
+    """
+
+    def preprocess(self, x, w):
+        m_values = [i.shape[0] for i in x]
+        # Convert m_values into offsets into grouped tensor.
+        m_sizes = torch.tensor(m_values).to(dtype=torch.int32, device=x[0].device)
+        # Quantize weights.
+        wq, w_scale = zip(*[quantize_fp8_row(i) for i in w])
+        # Group weights as single tensor.
+        wq = torch.concat(wq, dim=0).contiguous()
+        w_scale = torch.concat(w_scale, dim=0).contiguous()
+        # Also view input as flattened.
+        x = torch.concat(x, dim=0).contiguous()
+        # Return processed tensors.
+        return x, wq, w_scale, m_sizes
+
+    def quantize(self, x, wq, w_scale, m_sizes):
+        B = x.shape[0]
+        xq, x_scale = triton_quantize_fp8_row(x)
+        x_scale = x_scale.view(B, -1)
+        return xq, wq, x_scale, w_scale, m_sizes
+
+    def compute(self, xq, wq, x_scale, w_scale, m_sizes):
+        return grouped_gemm_fp8_rowwise(
+            xq, wq, m_sizes, x_scale, w_scale, _use_warp_specialization=True
+        )
+
+    def quantize_and_compute(self, x, wq, w_scale, m_sizes):
+        xq, wq, x_scale, w_scale, m_sizes = self.quantize(x, wq, w_scale, m_sizes)
+        return self.compute(xq, wq, x_scale, w_scale, m_sizes)
+
+    @property
+    def name(self) -> str:
+        return "triton_grouped_stacked"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return set(Accelerator)
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class FP8TritonStackedGroupedGemmFuseScatterAdd(FP8TritonStackedGroupedGemm):
+    """
+    FP8 grouped matmul with stacked inputs implemented with triton. Fused with ScatterAdd.
+    """
+
+    def preprocess(self, x, w):
+        x, wq, w_scale, m_sizes = super().preprocess(x, w)
+        M = x.shape[0]
+        N = wq.shape[0] // m_sizes.shape[0]
+        output = torch.zeros(M, N, dtype=torch.bfloat16, device=x.device)
+        indices = torch.randperm(M, dtype=torch.int32, device=x.device)
+        return x, wq, w_scale, m_sizes, output, indices
+
+    def quantize(self, x, wq, w_scale, m_sizes, *args):
+        return *super().quantize(x, wq, w_scale, m_sizes), *args
+
+    def compute(self, xq, wq, x_scale, w_scale, m_sizes, output, indices):
+        return grouped_gemm_fp8_rowwise(
+            xq,
+            wq,
+            m_sizes,
+            x_scale,
+            w_scale,
+            _use_warp_specialization=True,
+            _output_tensor=output,
+            _scatter_add_indices=indices,
+        )
+
+    def quantize_and_compute(self, x, wq, w_scale, m_sizes, *args):
+        xq, wq, x_scale, w_scale, m_sizes, *ret = self.quantize(
+            x, wq, w_scale, m_sizes, *args
+        )
+        return self.compute(xq, wq, x_scale, w_scale, m_sizes, *ret)
+
+    @property
+    def name(self) -> str:
+        return "triton_grouped_stacked_fuse_scatter_add"
+
+
+@register_gemm_op
+class DeepGemmStacked(GemmOpBase):
+    """
+    FP8 grouped matmul with blockwise scaling implemented with DeepGemm.
+    """
+
+    def preprocess(self, x, w):
+        m_values = [i.shape[0] for i in x]
+        # Convert m_values into offsets into grouped tensor.
+        indices = torch.arange(len(m_values))
+        m_indices = indices.repeat_interleave(torch.tensor(m_values)).to(
+            device=x[0].device, dtype=torch.int
+        )
+        # Quantize weights.
+        wq, w_scale = zip(*[quantize_fp8_block(i, block_k=128, block_m=128) for i in w])
+        # Group weights as single tensor.
+        wq = torch.stack(wq, dim=0).contiguous()
+        w_scale = torch.stack(w_scale, dim=0).contiguous()
+        # Also view input as flattened.
+        x = torch.concat(x, dim=0).contiguous()
+        # Return processed tensors.
+        return x, wq, w_scale, m_indices
+
+    def quantize(self, x, wq, w_scale, m_indices):
+        xq, x_scale = quantize_fp8_block(x, block_m=1, block_k=128)
+        # Pretranspose scales to deepgemm format.
+        x_scale = get_col_major_tma_aligned_tensor(x_scale)
+        return xq, wq, x_scale, w_scale, m_indices
+
+    def compute(self, xq, wq, x_scale, w_scale, m_indices):
+        # Preallocate output.
+        out = torch.empty(
+            [xq.shape[0], wq.shape[1]], device=xq.device, dtype=torch.bfloat16
+        )
+        m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
+            (xq, x_scale), (wq, w_scale), out, m_indices
+        )
+        return out
+
+    def quantize_and_compute(self, x, wq, w_scale, m_indices):
+        xq, wq, x_scale, w_scale, m_indices = self.quantize(x, wq, w_scale, m_indices)
+        return self.compute(xq, wq, x_scale, w_scale, m_indices)
+
+    @property
+    def name(self) -> str:
+        return "deepgemm_stacked"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        if DEEPGEMM_ENABLED:
+            return {Accelerator.NVIDIA_SM90}
+        return set()
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class DeepGemmMaskedStacked(DeepGemmStacked):
+    def preprocess(self, x, w):
+        # Quantize weights.
+        wq, w_scale = zip(*[quantize_fp8_block(i, block_k=128, block_m=128) for i in w])
+        # Group weights as single tensor.
+        wq = torch.stack(wq, dim=0).contiguous()
+        w_scale = torch.stack(w_scale, dim=0).contiguous()
+
+        # Also view input as flattened.
+        m_values = [i.shape[0] for i in x]
+        expected_m = max(m_values)
+        padded_m_max = ((max(m_values) + 127) // 128) * 128
+        masked_m = torch.tensor(m_values).to(dtype=torch.int32, device=x[0].device)
+
+        num_groups = len(m_values)
+        k = x[0].shape[1]
+        x_padded = torch.zeros(
+            [num_groups, padded_m_max, k], device=x[0].device, dtype=x[0].dtype
+        )
+        for g in range(num_groups):
+            x_padded[g, : m_values[g], :] = x[g]
+
+        # Return processed tensors.
+        return x_padded, wq, w_scale, masked_m, expected_m, m_values
+
+    def quantize(self, x, wq, w_scale, masked_m, expected_m, m_values):
+        g, m_max, k = x.shape
+        xq, x_scale = quantize_fp8_block(x.view(-1, k), block_m=1, block_k=128)
+        # Pretranspose scales to deepgemm format.
+        x_scale = get_col_major_tma_aligned_tensor(x_scale)
+        return (
+            xq.view(g, m_max, -1),
+            wq,
+            x_scale.view(g, m_max, -1),
+            w_scale,
+            masked_m,
+            expected_m,
+            m_values,
+        )
+
+    def compute(self, xq, wq, x_scale, w_scale, masked_m, expected_m, m_values):
+        # Preallocate output.
+        out = torch.empty(
+            [xq.shape[0], xq.shape[1], wq.shape[1]],
+            device=xq.device,
+            dtype=torch.bfloat16,
+        )
+        m_grouped_gemm_fp8_fp8_bf16_nt_masked(
+            (xq, x_scale), (wq, w_scale), out, masked_m, expected_m
+        )
+        num_groups = xq.shape[0]
+        out_list = [out[g, : m_values[g], :] for g in range(num_groups)]
+        return out_list
+
+    def quantize_and_compute(self, x, wq, w_scale, masked_m, expected_m, m_values):
+        xq, wq, x_scale, w_scale, masked_m, expected_m = self.quantize(
+            x, wq, w_scale, masked_m, expected_m, m_values
+        )
+        return self.compute(xq, wq, x_scale, w_scale, masked_m, expected_m, m_values)
+
+    @property
+    def name(self) -> str:
+        return "deepgemm_masked_stacked"
+
+
+@register_gemm_op
+class DeepGemmBlockwise(GemmOpBase):
+    """
+    FP8 matmul with blockwise scaling implemented with DeepGemm.
+    """
+
+    def preprocess(self, x, w):
+        # Quantize weights.
+        wq, w_scale = quantize_fp8_block(w, block_m=128, block_k=128)
+        # allocate output.
+        out = torch.empty(
+            x.shape[0], wq.shape[0], device=x.device, dtype=torch.bfloat16
+        )
+        # Return processed tensors.
+        return x, wq, w_scale, out
+
+    def quantize(self, x, wq, w_scale, out):
+        xq, x_scale = quantize_fp8_group(x, group_size=128)
+        return xq, wq, x_scale, w_scale, out
+
+    def compute(self, xq, wq, x_scale, w_scale, out):
+        gemm_fp8_fp8_bf16_nt((xq, x_scale), (wq, w_scale), out)
+        return out
+
+    def quantize_and_compute(self, x, wq, w_scale, out):
+        xq, wq, x_scale, w_scale, out = self.quantize(x, wq, w_scale, out)
+        return self.compute(xq, wq, x_scale, w_scale, out)
+
+    @property
+    def name(self) -> str:
+        return "deepgemm_blockwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        if DEEPGEMM_ENABLED:
+            return {Accelerator.NVIDIA_SM90}
+        return set()
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class DeepGemmRowwise(GemmOpBase):
+    """
+    FP8 matmul with rowwise scaling implemented with DeepGemm.
+    """
+
+    def preprocess(self, x, w):
+        # Quantize weights.
+        wq, w_scale = quantize_fp8_row(w)
+        # allocate output.
+        out = torch.empty(
+            x.shape[0], wq.shape[0], device=x.device, dtype=torch.bfloat16
+        )
+        # Return processed tensors.
+        return x, wq, w_scale, out
+
+    def quantize(self, x, wq, w_scale, out):
+        xq, x_scale = quantize_fp8_row(x)
+        # Pretranspose scales to deepgemm format.
+        x_scale = get_col_major_tma_aligned_tensor(x_scale, rowwise_scaling=True)
+        return xq, wq, x_scale, w_scale, out
+
+    def compute(self, xq, wq, x_scale, w_scale, out):
+        gemm_fp8_fp8_bf16_nt((xq, x_scale), (wq, w_scale), out)
+        return out
+
+    def quantize_and_compute(self, x, wq, w_scale, out):
+        xq, wq, x_scale, w_scale, out = self.quantize(x, wq, w_scale, out)
+        return self.compute(xq, wq, x_scale, w_scale, out)
+
+    @property
+    def name(self) -> str:
+        return "deepgemm_rowwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        if DEEPGEMM_ENABLED:
+            return {Accelerator.NVIDIA_SM90}
+        return set()
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class FP8StackedGroupedGemm(GemmOpBase):
+    """
+    FP8 grouped matmul with rowwise scaling and stacked inputs.
+    """
+
+    def preprocess(self, x, w):
+        m_values = [i.shape[0] for i in x]
+        m_sizes = torch.tensor(m_values).to(dtype=torch.int64, device=x[0].device)
+        # Quantize weights.
+        wq, w_scale = zip(*[quantize_fp8_row(i) for i in w])
+        # Group weights as single tensor.
+        wq = torch.stack(wq, dim=0).contiguous()
+        w_scale = torch.stack(w_scale, dim=0).contiguous()
+        # Also view input as flattened.
+        x = torch.concat(x, dim=0).contiguous()
+        # Return processed tensors.
+        return x, wq, w_scale, m_sizes
+
+    def quantize(self, x, wq, w_scale, m_sizes):
+        B = x.shape[0]
+        xq, x_scale = triton_quantize_fp8_row(x)
+        x_scale = x_scale.view(B, -1)
+        return xq, wq, x_scale, w_scale, m_sizes
+
+    def compute(self, xq, wq, x_scale, w_scale, m_sizes):
+        return torch.ops.mslk.f8f8bf16_rowwise_grouped_stacked(
+            xq, wq, x_scale, w_scale, m_sizes
+        )
+
+    def quantize_and_compute(self, x, wq, w_scale, m_sizes):
+        xq, wq, x_scale, w_scale, m_sizes = self.quantize(x, wq, w_scale, m_sizes)
+        return self.compute(xq, wq, x_scale, w_scale, m_sizes)
+
+    @property
+    def name(self) -> str:
+        if torch.version.cuda:
+            return "cutlass_grouped_stacked"
+        else:
+            return "ck_grouped_stacked"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.NVIDIA_SM100,
+            Accelerator.NVIDIA_SM103,
+            Accelerator.AMD_MI300X,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class FP8StackedGroupedGemmTorch(FP8StackedGroupedGemm):
+    def quantize(self, x, wq, w_scale, m_sizes):
+        xq, wq, x_scale, w_scale, m_sizes = super().quantize(x, wq, w_scale, m_sizes)
+        offsets = torch.cumsum(m_sizes, dim=0, dtype=torch.int32)
+        out = torch.empty(
+            (xq.shape[0], wq.shape[1]), dtype=torch.bfloat16, device=xq.device
+        )
+        x_scale = x_scale.view(x_scale.shape[0])
+        return xq, wq, x_scale, w_scale, offsets, out
+
+    def compute(self, xq, wq, x_scale, w_scale, offsets, out):
+        return torch.ops.mslk.f8f8bf16_rowwise_grouped_mm(
+            xq, wq, x_scale, w_scale, offsets, out
+        )
+
+    def quantize_and_compute(self, x, wq, w_scale, m_sizes):
+        xq, wq, x_scale, w_scale, offsets, out = self.quantize(x, wq, w_scale, m_sizes)
+        return self.compute(xq, wq, x_scale, w_scale, offsets, out)
+
+    @property
+    def name(self) -> str:
+        return "ck_grouped_stacked_torch_2d3d"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.AMD_MI300X}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class ScaledGroupedMMRowwise(FP8StackedGroupedGemmTorch):
+    def __init__(self):
+        self.fast_accum = True
+        self.torch_compile = False
+
+    def compute(self, xq, wq, x_scale, w_scale, offsets, _):
+        if self.torch_compile:
+            f = torch.compile(
+                torch._scaled_grouped_mm,
+                options={
+                    "max_autotune": True,
+                    "max_autotune_gemm_backends": "TRITON,CK,CUTLASS,ATEN",
+                },
+            )
+        else:
+            f = torch._scaled_grouped_mm
+
+        return f(
+            xq,
+            wq.transpose(-2, -1),
+            offs=offsets,
+            out_dtype=torch.bfloat16,
+            scale_a=x_scale,
+            scale_b=w_scale,
+            scale_result=None,
+            use_fast_accum=self.fast_accum,
+        )
+
+    @property
+    def name(self) -> str:
+        return "scaled_grouped_mm_rowwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.AMD_MI300X,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class FP8StackedGroupwiseGroupedGemm(GemmOpBase):
+    """
+    FP8 grouped matmul with groupwise scaling and stacked inputs.
+    """
+
+    def preprocess(self, x, w):
+        m_values = [i.shape[0] for i in x]
+        m_sizes = torch.tensor(m_values).to(dtype=torch.int64, device=x[0].device)
+        # Quantize weights.
+        wq, w_scale = zip(
+            *[quantize_fp8_block(i, block_m=128, block_k=128, k_major=False) for i in w]
+        )
+        # Group weights as single tensor.
+        wq = torch.stack(wq, dim=0).contiguous()
+        w_scale = torch.stack(w_scale, dim=0).contiguous()
+        # Also view input as flattened.
+        x = torch.concat(x, dim=0).contiguous()
+        # Return processed tensors.
+        return x, wq, w_scale, m_sizes
+
+    def quantize(self, x, wq, w_scale, m_sizes):
+        xq, x_scale = quantize_fp8_group(x, m_sizes=m_sizes)
+        return xq, wq, x_scale, w_scale, m_sizes
+
+    def compute(self, xq, wq, x_scale, w_scale, m_sizes):
+        return torch.ops.mslk.f8f8bf16_groupwise_grouped(
+            xq, wq, x_scale, w_scale, m_sizes
+        )
+
+    def quantize_and_compute(self, x, wq, w_scale, m_sizes):
+        xq, wq, x_scale, w_scale, m_sizes = self.quantize(x, wq, w_scale, m_sizes)
+        return self.compute(xq, wq, x_scale, w_scale, m_sizes)
+
+    @property
+    def name(self) -> str:
+        return "cutlass_groupwise_grouped"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM90}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class BF16GroupedGemm(GemmOpBase):
+    """
+    BF16 grouped matmul implemented with CK or Cutlass.
+    """
+
+    def preprocess(self, x, w):
+        # Apply sparsity to inputs if appropriate.
+        # First check if N and K are fixed.
+        m_values = [i.shape[0] for i in x]
+        n_values = [i.shape[0] for i in w]
+        k_values = [i.shape[1] for i in w]
+        # If so, do specialized version of initialization.
+        if len(np.unique(n_values)) == 1 and len(np.unique(k_values)) == 1:
+            m_values = [i.shape[0] for i in x]
+            # Inputs for fixed nk mode must be contiguous, however in the benchmark
+            # script they typically are not. Do a little special processing to make them
+            # work. In practice this wont be needed.
+            # Start by padding along m dimension with zeros.
+            max_m = max(m_values)
+            x = [
+                torch.nn.functional.pad(i, (0, 0, 0, max_m - i.shape[0]), value=0)
+                for i in x
+            ]
+            # Stack inputs into groups.
+            x = torch.stack(x).contiguous()
+            w = torch.stack(w).contiguous()
+            return (
+                x,
+                w,
+                torch.tensor(m_values).to(dtype=torch.int64, device=x[0].device),
+            )
+        return x, w, None
+
+    def quantize(self, x, w, m_values=None):
+        # No action required.
+        return x, w, m_values
+
+    def compute(self, x, w, m_values):
+        if m_values is None:
+            return torch.ops.mslk.bf16bf16bf16_grouped(x, w)
+        else:
+            return torch.ops.mslk.bf16bf16bf16_grouped_dynamic(x, w, m_values)
+
+    def quantize_and_compute(self, x, w, m_values):
+        return self.compute(x, w, m_values)
+
+    @property
+    def name(self) -> str:
+        if torch.version.cuda:
+            return "cutlass_bf16_grouped"
+        else:
+            return "ck_bf16_grouped"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.NVIDIA_SM100,
+            Accelerator.NVIDIA_SM103,
+            Accelerator.AMD_MI300X,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class FP8RowwiseBatchedGemm(GemmOpBase):
+    """
+    FP8 batched matmul with rowwise scaling.
+    """
+
+    def quantize(self, x, w):
+        assert isinstance(x, list) and isinstance(w, list)
+        x = torch.stack(x, dim=0)
+        w = torch.stack(w, dim=0)
+        # Quantize both input tensors.
+        xq, x_scale = quantize_fp8_row(x)
+        wq, w_scale = quantize_fp8_row(w)
+        return xq, wq, x_scale, w_scale
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        return torch.ops.mslk.f8f8bf16_rowwise_batched(xq, wq, x_scale, w_scale)
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale = self.quantize(x, w)
+        return self.compute(xq, wq, x_scale, w_scale)
+
+    @property
+    def name(self) -> str:
+        if torch.version.cuda:
+            return "cutlass_rowwise_batched"
+        else:
+            return "ck_rowwise_batched"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.NVIDIA_SM100,
+            Accelerator.NVIDIA_SM103,
+            Accelerator.AMD_MI300X,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+# This kernel is broken and causes GPU to lock up, needs some investigation
+# @register_gemm_op
+class TritonFP8RowwiseGemm(GemmOpBase):
+    """
+    FP8 matmul with rowwise scaling implemented with Triton.
+    """
+
+    def __init__(self):
+        self.fast_accum = True
+
+    def quantize(self, x, w):
+        # Quantize both input tensors.
+        xq, x_scale = quantize_fp8_row(x)
+        wq, w_scale = quantize_fp8_row(w)
+        bias = torch.randn(w.shape[0], device=x.device, dtype=torch.float32)
+        return xq, wq, x_scale, w_scale, bias
+
+    def compute(self, xq, wq, x_scale, w_scale, bias):
+        return matmul_fp8_row(
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            bias=bias,
+            fp8_fast_accum=self.fast_accum,
+            use_warp_specialization=True,
+        )
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale = self.quantize(x, w)
+        return self.compute(xq, wq, x_scale, w_scale)
+
+    @property
+    def name(self) -> str:
+        return "triton_rowwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return set(Accelerator)
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class FP8TritonBlockwiseGemm(GemmOpBase):
+    """
+    FP8 matmul with block scaling.
+    """
+
+    def quantize(self, x, w):
+        # Quantize both input tensors.
+        xq, x_scale = quantize_fp8_block(x, 128, 128)
+        wq, w_scale = quantize_fp8_block(w, 128, 128)
+        return xq, wq, x_scale, w_scale
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        return matmul_fp8_block(xq, wq, x_scale, w_scale, 128, 128, 128)
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale = self.quantize(x, w)
+        return self.compute(xq, wq, x_scale, w_scale)
+
+    @property
+    def name(self) -> str:
+        return "triton_blockwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return set(Accelerator)
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class FP8CutlassBlockwiseGemm(GemmOpBase):
+    """
+    FP8 matmul with block scaling.
+    """
+
+    def quantize(self, x, w):
+        # Quantize both input tensors.
+        xq, x_scale = quantize_fp8_block(x, 128, 128)
+        wq, w_scale = quantize_fp8_block(w, 128, 128)
+        return xq, wq, x_scale, w_scale
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        return torch.ops.mslk.f8f8bf16_blockwise(
+            xq, wq, x_scale, w_scale, 128, 128, 128
+        )
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale = self.quantize(x, w)
+        return self.compute(xq, wq, x_scale, w_scale)
+
+    @property
+    def name(self) -> str:
+        if torch.version.cuda:
+            return "cutlass_blockwise"
+        else:
+            return "ck_blockwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.AMD_MI300X,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class FP8CutlassGroupwiseGemm(GemmOpBase):
+    """
+    FP8 matmul with group / block scaling.
+    """
+
+    def preprocess(self, x, w):
+        # Quantize weights.
+        # Scale is expected to be in [K, N] layout (N Major).
+        wq, w_scale = quantize_fp8_block(w, block_m=128, block_k=128, k_major=False)
+        # Return processed tensors.
+        return x, wq, w_scale
+
+    def quantize(self, x, wq, w_scale):
+        # Scale is expected to be in [K, M] layout (M Major).
+        xq, x_scale = quantize_fp8_group(x, k_major=False)
+        # Pretranspose scales to deepgemm format.
+        return xq, wq, x_scale, w_scale
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        return torch.ops.mslk.f8f8bf16_groupwise(xq, wq, x_scale, w_scale)
+
+    def quantize_and_compute(self, x, wq, w_scale):
+        xq, wq, x_scale, w_scale = self.quantize(x, wq, w_scale)
+        return self.compute(xq, wq, x_scale, w_scale)
+
+    @property
+    def name(self) -> str:
+        return "cutlass_groupwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM90}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class F8I4RowwiseGemm(GemmOpBase):
+    """
+    Mixed Precision FP8 Activations with Int4 Weights.
+    """
+
+    def _int4_row_quantize(
+        self,
+        x: torch.Tensor,
+        group_size: int = 128,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        n_bit = 4  # Number of target bits.
+        to_quant = x.reshape(-1, group_size).to(torch.float)
+
+        max_val = to_quant.amax(dim=1, keepdim=True)
+        min_val = to_quant.amin(dim=1, keepdim=True)
+        max_int = 2**n_bit - 1
+        min_int = 0
+        scales = (max_val - min_val).clamp(min=1e-6) / max_int
+
+        zeros = min_val + scales * (2 ** (n_bit - 1))
+
+        out = to_quant.sub(min_val).div(scales).round().clamp_(min_int, max_int)
+
+        # Recenter output and move to int8.
+        out = (out - 2 ** (n_bit - 1)).to(dtype=torch.int8).reshape(x.shape)
+
+        # Cutlass expects column major layout for scale and zero point,
+        # so we transpose here and make them contiguous.
+        scales = scales.view(x.shape[0], -1).t().contiguous()
+        zeros = zeros.view(x.shape[0], -1).t().contiguous()
+
+        return out, scales, zeros
+
+    def _pack_int4(self, x: torch.Tensor) -> torch.Tensor:
+        # Given int8 x, pack adjacent int4 values into a single int8.
+        low_x = x[:, ::2]
+        high_x = x[:, 1::2]
+
+        # High bits need to left shift, this also masks off extra bits.
+        high_x = torch.bitwise_left_shift(high_x, 4)
+        # Low bits need to have sign bits removed.
+        low_x = torch.bitwise_and(low_x, 0xF)
+
+        # Recombine into a single value with bitwise or.
+        return torch.bitwise_or(low_x, high_x).contiguous()
+
+    def quantize(self, x, w):
+        # Quantize both input tensors.
+        xq, x_scale = quantize_fp8_row(x)
+        wq, w_scale, w_zp = self._int4_row_quantize(w)
+        # Pack int4 values together.
+        wq = self._pack_int4(wq)
+        return xq, wq, x_scale, w_scale, w_zp
+
+    def compute(self, xq, wq, x_scale, w_scale, w_zp):
+        return torch.ops.mslk.f8i4bf16_rowwise(xq, wq, x_scale, w_scale, w_zp)
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale, w_zp = self.quantize(x, w)
+        return self.compute(xq, wq, x_scale, w_scale, w_zp)
+
+    @property
+    def name(self) -> str:
+        return "cutlass_f8i4_rowwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM90}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class F8I4ShuffledGemm(GemmOpBase):
+    def preprocess(self, x, w):
+        # Prequantize and pack weights.
+        wq, (group_scale, row_scale) = quantize_int4_preshuffle(w)
+        return x, wq, row_scale, group_scale
+
+    def quantize(self, x, wq, row_scale, group_scale):
+        # Quantize both input tensors.
+        xq, x_scale = quantize_fp8_row(x)
+        return xq, wq, x_scale, row_scale, group_scale
+
+    def compute(self, xq, wq, x_scale, row_scale, group_scale):
+        # Handle batched cases by looping over each batch.
+        if xq.dim() == 3:
+            B, M, _ = xq.shape
+            _, N, _ = wq.shape
+            y = torch.empty((B, M, N), device=xq.device, dtype=torch.bfloat16)
+            for i in range(B):
+                y[i] = torch.ops.mslk.f8i4bf16_shuffled(
+                    xq[i], wq[i], x_scale[i], row_scale[i], group_scale[i]
+                )
+            return y
+        # Otherwise run gemm normally.
+        return torch.ops.mslk.f8i4bf16_shuffled(xq, wq, x_scale, row_scale, group_scale)
+
+    def quantize_and_compute(self, x, wq, row_scale, group_scale):
+        xq, wq, x_scale, row_scale, group_scale = self.quantize(
+            x, wq, row_scale, group_scale
+        )
+        return self.compute(xq, wq, x_scale, row_scale, group_scale)
+
+    @property
+    def name(self) -> str:
+        return "cutlass_f8i4_preshuffle"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM90}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class BF16I4ShuffledGemm(GemmOpBase):
+    def preprocess(self, x, w):
+        # Prequantize and pack weights.
+        wq, (group_scale, group_zero) = quantize_int4_preshuffle(w, dtype="bf16")
+        return x, wq, group_scale, group_zero
+
+    def quantize(self, x, wq, group_scale, group_zero):
+        # No extra action required.
+        return x, wq, group_scale, group_zero
+
+    def compute(self, x, wq, group_scale, group_zero):
+        # Handle batched cases by looping over each batch.
+        if x.dim() == 3:
+            B, M, _ = x.shape
+            _, N, _ = wq.shape
+            y = torch.empty((B, M, N), device=x.device, dtype=torch.bfloat16)
+            for i in range(B):
+                y[i] = torch.ops.mslk.bf16i4bf16_shuffled(
+                    x[i], wq[i], group_scale[i], group_zero[i]
+                )
+            return y
+        # Otherwise run gemm normally.
+        return torch.ops.mslk.bf16i4bf16_shuffled(x, wq, group_scale, group_zero)
+
+    def quantize_and_compute(self, x, wq, group_scale, group_zero):
+        x, wq, group_scale, group_zero = self.quantize(x, wq, group_scale, group_zero)
+        return self.compute(x, wq, group_scale, group_zero)
+
+    @property
+    def name(self) -> str:
+        return "cutlass_bf16i4_preshuffle"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM90}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class BF16I4ShuffledBatchedGemm(GemmOpBase):
+    """
+    BF16 x INT4 mixed dtype batched gemm with preshuffling.
+    """
+
+    def preprocess(self, x, w):
+        assert isinstance(x, list) and isinstance(w, list)
+        x = torch.stack(x, dim=0)
+        w = torch.stack(w, dim=0)
+        # Prequantize and pack weights.
+        wq, (group_scale, group_zero) = quantize_int4_preshuffle(w, dtype="bf16")
+        return x, wq, group_scale, group_zero
+
+    def quantize(self, x, wq, group_scale, group_zero):
+        # No extra action required.
+        return x, wq, group_scale, group_zero
+
+    def compute(self, x, wq, group_scale, group_zero):
+        return torch.ops.mslk.bf16i4bf16_shuffled_batched(
+            x, wq, group_scale, group_zero
+        )
+
+    def quantize_and_compute(self, x, wq, group_scale, group_zero):
+        x, wq, group_scale, group_zero = self.quantize(x, wq, group_scale, group_zero)
+        return self.compute(x, wq, group_scale, group_zero)
+
+    @property
+    def name(self) -> str:
+        return "cutlass_bf16i4_preshuffle_batched"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM90}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class F8I4ShuffledGroupedGemm(GemmOpBase):
+    """
+    FP8 x Int4 mixed dtype grouped gemm with preshuffling.
+    """
+
+    def preprocess(self, x, w):
+        assert isinstance(x, list) and isinstance(w, list), (
+            "Only supported for grouped inputs."
+        )
+        m_values = [i.shape[0] for i in x]
+        # Convert m_values into offsets into grouped tensor.
+        m_sizes = torch.tensor(m_values).to(dtype=torch.int32, device=x[0].device)
+        # Quantize weights.
+        wq, scales = zip(*[quantize_int4_preshuffle(i) for i in w])
+        group_scale, row_scale = zip(*scales)
+        # Group weights as single tensor.
+        wq = torch.stack(wq, dim=0).contiguous()
+        row_scale = torch.stack(row_scale, dim=0).contiguous()
+        group_scale = torch.stack(group_scale, dim=0).contiguous()
+        # Also view input as flattened.
+        x = torch.concat(x, dim=0).contiguous()
+        # Return processed tensors.
+        return x, wq, row_scale, group_scale, m_sizes
+
+    def quantize(self, x, wq, row_scale, group_scale, m_sizes):
+        B = x.shape[0]
+        xq, x_scale = triton_quantize_fp8_row(x)
+        x_scale = x_scale.view(B, -1)
+        return xq, wq, x_scale, row_scale, group_scale, m_sizes
+
+    def compute(self, xq, wq, x_scale, row_scale, group_scale, m_sizes):
+        out = torch.ops.mslk.f8i4bf16_shuffled_grouped(
+            xq, wq, x_scale, row_scale, group_scale, m_sizes
+        )
+        return out
+
+    def quantize_and_compute(self, x, wq, row_scale, group_scale, m_sizes):
+        xq, wq, x_scale, row_scale, group_scale, m_sizes = self.quantize(
+            x, wq, row_scale, group_scale, m_sizes
+        )
+        return self.compute(xq, wq, x_scale, row_scale, group_scale, m_sizes)
+
+    @property
+    def name(self) -> str:
+        if torch.version.cuda:
+            return "cutlass_f8i4_grouped_preshuffle"
+        else:
+            return "ck_f8i4_grouped_preshuffle"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM90}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class BF16I4ShuffledGroupedGemm(GemmOpBase):
+    """
+    BF16 x Int4 mixed dtype grouped gemm with preshuffling.
+    """
+
+    def preprocess(self, x, w):
+        assert isinstance(x, list) and isinstance(w, list), (
+            "Only supported for grouped inputs."
+        )
+        m_values = [i.shape[0] for i in x]
+        # Convert m_values into offsets into grouped tensor.
+        m_sizes = torch.tensor(m_values).to(dtype=torch.int32, device=x[0].device)
+        # Quantize weights.
+        wq, scales = zip(
+            *[quantize_int4_preshuffle(i, dtype="bf16", use_zp=False) for i in w]
+        )
+        # Group weights as single tensor.
+        group_scale, group_zero = zip(*scales)
+        wq = torch.stack(wq, dim=0).contiguous()
+        group_scale = torch.stack(group_scale, dim=0).contiguous()
+        group_zero = torch.stack(group_zero, dim=0).contiguous()
+        # Also view input as flattened.
+        x = torch.concat(x, dim=0).contiguous()
+        # Return processed tensors.
+        return x, wq, group_scale, group_zero, m_sizes
+
+    def quantize(self, x, wq, group_scale, group_zero, m_sizes):
+        return x, wq, group_scale, group_zero, m_sizes
+
+    def compute(self, x, wq, group_scale, group_zero, m_sizes):
+        # TODO Zero points arent currently supported in grouped gemm.
+        # We leave them as inputs for future compatibility but they are ignored.
+        return torch.ops.mslk.bf16i4bf16_shuffled_grouped(
+            x, wq, group_scale, group_zero, m_sizes
+        )
+
+    def quantize_and_compute(self, x, wq, group_scale, group_zero, m_sizes):
+        x, wq, group_scale, group_zero, m_sizes = self.quantize(
+            x, wq, group_scale, group_zero, m_sizes
+        )
+        return self.compute(x, wq, group_scale, group_zero, m_sizes)
+
+    @property
+    def name(self) -> str:
+        return "cutlass_bf16i4_grouped_preshuffle"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM90}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class BF16GroupedGrad(GemmOpBase):
+    """
+    BF16 grouped matmul with dgrad inputs in pretraining backed by cutlass
+    """
+
+    def preprocess(self, x, w):
+        m_values = [i.shape[0] for i in x]
+        # Convert m_values into offsets into grouped tensor.
+        m_sizes = torch.tensor(m_values).to(dtype=torch.int64, device=x[0].device)
+        # Group weights as single tensor.
+        w = torch.stack(w, dim=0).contiguous()
+        # Prepare online dgrad during pretraining backward.
+        w_perm = w.permute(0, 2, 1).contiguous()
+        # w.contiguous() is very expensive so handling it inside the gmm kernel for free
+        w = w_perm.permute(0, 2, 1)
+
+        # Also view input as flattened.
+        x = torch.concat(x, dim=0).contiguous()
+        # Return processed tensors.
+        return x, w, m_sizes
+
+    def quantize(self, x, w, m_sizes):
+        return x, w, m_sizes
+
+    def compute(self, x, w, m_sizes):
+        return torch.ops.mslk.bf16bf16bf16_grouped_grad(x, w, m_sizes)
+
+    def quantize_and_compute(self, x, w, m_sizes):
+        x, w, m_sizes = self.quantize(x, w, m_sizes)
+        return self.compute(x, w, m_sizes)
+
+    @property
+    def name(self) -> str:
+        return "bf16_grouped_grad"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.NVIDIA_SM100,
+            Accelerator.NVIDIA_SM103,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class BF16GroupedWGrad(GemmOpBase):
+    """
+    BF16 grouped matmul with wgrad inputs in pretraining backed by cutlass
+    """
+
+    def preprocess(self, x, w):
+        # Get K values for each group
+        k_values = [xi.shape[1] for xi in x]  # K dimension for each group
+
+        # Convert k_values into sizes tensor
+        k_sizes = torch.tensor(k_values).to(dtype=torch.int64, device=x[0].device)
+
+        x = torch.concat(x, dim=1).contiguous()  # shape: (M, G*K)
+        w = torch.concat(w, dim=1).contiguous()  # shape: (N, G*K)
+
+        # Transpose the follows to simulate wgrad shapes
+        x = x.t().contiguous()  # shape: (G*K, M)
+        w = w.t().contiguous()  # shape: (G*K, N)
+
+        # Return processed tensors
+        return x, w, k_sizes
+
+    def quantize(self, x, w, k_sizes):
+        return x, w, k_sizes
+
+    def compute(self, x, w, k_sizes):
+        return torch.ops.mslk.bf16bf16bf16_grouped_wgrad(x, w, k_sizes)
+
+    def quantize_and_compute(self, x, w, k_sizes):
+        x, w, k_sizes = self.quantize(x, w, k_sizes)
+        return self.compute(x, w, k_sizes)
+
+    @property
+    def name(self) -> str:
+        return "bf16_grouped_wgrad"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.NVIDIA_SM100,
+            Accelerator.NVIDIA_SM103,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class BF16GroupedStacked(GemmOpBase):
+    """
+    BF16 grouped matmul with stacked inputs backed by cutlass or ck.
+    """
+
+    def preprocess(self, x, w):
+        m_values = [i.shape[0] for i in x]
+        # Convert m_values into offsets into grouped tensor.
+        m_sizes = torch.tensor(m_values).to(dtype=torch.int64, device=x[0].device)
+        # Group weights as single tensor.
+        w = torch.stack(w, dim=0).contiguous()
+        # Also view input as flattened.
+        x = torch.concat(x, dim=0).contiguous()
+        # Return processed tensors.
+        return x, w, m_sizes
+
+    def quantize(self, x, w, m_sizes):
+        return x, w, m_sizes
+
+    def compute(self, x, w, m_sizes):
+        return torch.ops.mslk.bf16bf16bf16_grouped_stacked(x, w, m_sizes)
+
+    def quantize_and_compute(self, x, w, m_sizes):
+        x, w, m_sizes = self.quantize(x, w, m_sizes)
+        return self.compute(x, w, m_sizes)
+
+    @property
+    def name(self) -> str:
+        return "bf16_grouped_stacked"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {
+            Accelerator.NVIDIA_SM90,
+            Accelerator.NVIDIA_SM100,
+            Accelerator.NVIDIA_SM103,
+            Accelerator.AMD_MI300X,
+        }
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class BF16I4RowwiseGemm(F8I4RowwiseGemm):
+    """
+    Mixed Precision BF16 Activations with Int4 Weights.
+    """
+
+    def quantize(self, x, w):
+        # Quantize both input tensors.
+        wq, w_scale, w_zp = self._int4_row_quantize(w)
+        # Pack int4 values together.
+        wq = self._pack_int4(wq)
+        return (
+            x.to(torch.bfloat16),
+            wq,
+            w_scale,
+            w_zp,
+        )
+
+    def compute(self, x, wq, w_scale, w_zp):
+        return torch.ops.mslk.bf16i4bf16_rowwise(x, wq, w_scale, w_zp)
+
+    def quantize_and_compute(self, x, w):
+        x, wq, w_scale, w_zp = self.quantize(x, w)
+        return self.compute(x, wq, w_scale, w_zp)
+
+    @property
+    def name(self) -> str:
+        return "cutlass_bf16i4_rowwise"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM90}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class TinyGemmBF16I4(GemmOpBase):
+    """
+    Mixed Precision BF16 Activations with Int4 Weights using tinygemm.
+    """
+
+    def quantize(self, x, w):
+        # Quantize and pack weights to int4 using tinygemm utils.
+        w_int32, w_scales_and_zeros = group_quantize_tensor(
+            w, n_bit=4, q_group_size=128
+        )
+        wq = torch.ops.tinygemm.convert_matrix_to_m16n8k16_Aint4_layout(w_int32, 4)
+        return x, wq, w_scales_and_zeros
+
+    def compute(self, x, wq, scale):
+        return torch.ops.tinygemm.tinygemm_y_f16RM_x_f16RM_w_int4TC(
+            wq, x, 128, scale, False
+        )
+
+    def quantize_and_compute(self, x, w):
+        x, wq, scale = self.quantize(x, w)
+        return self.compute(x, wq, scale)
+
+    @property
+    def name(self) -> str:
+        return "tinygemm_bf16i4"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        if TINYGEMM_ENABLED:
+            return {Accelerator.NVIDIA_SM90}
+        return set()
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class MarlinBF16I4(GemmOpBase):
+    """
+    Mixed Precision BF16 Activations with Int4 Weights using Marlin.
+    """
+
+    def quantize(self, x, w):
+        # Marlin quantize expects weights in [K, N] layout.
+        _, wq, scale = marlin_quantize(w.t().contiguous(), 128)
+        return x, wq, scale
+
+    def compute(self, x, wq, scale):
+        return torch.ops.marlin.marlin_gemm(x, wq, scale)
+
+    def quantize_and_compute(self, x, w):
+        x, wq, scale = self.quantize(x, w)
+        return self.compute(x, wq, scale)
+
+    @property
+    def name(self) -> str:
+        return "marlin_bf16i4"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        if MARLIN_ENABLED:
+            return {
+                Accelerator.NVIDIA_SM90,
+                Accelerator.NVIDIA_SM100,
+                Accelerator.NVIDIA_SM103,
+            }
+        return set()
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class MacheteBF16I4(GemmOpBase):
+    """
+    Mixed Precision BF16 Activations with Int4 Weights using Machete.
+    """
+
+    def quantize(self, x, w):
+        # Marlin quantize expects weights in [K, N] layout.
+        _, wq, scale, _ = machete_quantize_and_pack(
+            w.t().contiguous(), bits=4, groupsize=128
+        )
+        return x, wq, scale
+
+    def compute(self, x, wq, scale):
+        return machete_gemm(x, wq, bits=4, groupsize=128, scales=scale)
+
+    def quantize_and_compute(self, x, w):
+        x, wq, scale = self.quantize(x, w)
+        return self.compute(x, wq, scale)
+
+    @property
+    def name(self) -> str:
+        return "machete_bf16i4"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        if MACHETE_ENABLED:
+            return {Accelerator.NVIDIA_SM90}
+        return set()
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class NVFP4Gemm(GemmOpBase):
+    """
+    NVFP4 matmul with block-wise scaling.
+    """
+
+    def quantize(self, x, w):
+        x_global_scale = (448.0 * 6.0) / torch.amax(torch.abs(x.flatten()), dim=-1).to(
+            torch.float32
+        )
+        w_global_scale = (448.0 * 6.0) / torch.amax(torch.abs(w.flatten()), dim=-1).to(
+            torch.float32
+        )
+        global_scale = 1 / (x_global_scale * w_global_scale)
+
+        xq, x_scale = triton_quantize_nvfp4(x, x_global_scale)
+        wq, w_scale = triton_quantize_nvfp4(w, w_global_scale)
+
+        return xq, wq, x_scale, w_scale, global_scale
+
+    def compute(self, xq, wq, x_scale, w_scale, global_scale):
+        return torch.ops.mslk.f4f4bf16(
+            xq, wq, x_scale, w_scale, global_scale=global_scale
+        )
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale, global_scale = self.quantize(x, w)
+        return self.compute(xq, wq, x_scale, w_scale, global_scale=global_scale)
+
+    @property
+    def name(self) -> str:
+        return "cutlass_nv_f4f4bf16"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM100, Accelerator.NVIDIA_SM103}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP4
+
+
+@register_gemm_op
+class NVFP4Quantize(GemmOpBase):
+    """
+    NVFP4 quantization with block-wise scaling.
+    """
+
+    def quantize_rms(self, x, w):
+        M, N = w.shape
+        group_size = 16
+        w = torch.randn(group_size, dtype=torch.bfloat16, device=w.device)
+        x_global_scale = torch.tensor([448.0 * 6.0]).to(
+            device=x.device, dtype=torch.float32
+        ) / torch.amax(torch.abs(x.flatten()), dim=-1).to(torch.float32)
+        xq_ref, x_scale_ref = triton_scale_nvfp4_quant_rms(
+            x,
+            w.repeat(M * N // group_size),
+            x_global_scale,
+            group_size=group_size,
+            EPS=1e-5,
+        )
+
+        intermediate = rms_norm(x.reshape(-1, 16), w, eps=1e-5)
+        intermediate = intermediate.to(torch.bfloat16).reshape(M, N)
+        xq, x_scale = triton_quantize_nvfp4(
+            intermediate,
+            x_global_scale,
+            group_size=group_size,
+        )
+
+    def quantize_silu(self, x, w):
+        M, N = x.shape
+        group_size = 16
+        x_global_scale = torch.tensor([448.0 * 6.0]).to(
+            device=x.device, dtype=torch.float32
+        ) / torch.amax(torch.abs(x.flatten()), dim=-1).to(torch.float32)
+        xq_ref, x_scale_ref = triton_scale_nvfp4_quant_silu(
+            x,
+            w,
+            x_global_scale,
+            group_size=group_size,
+        )
+
+        intermediate = silu_mul(x.reshape(-1, 16), w.reshape(-1, 16))
+        intermediate = intermediate.to(torch.bfloat16).reshape(M, N)
+        xq, x_scale = triton_quantize_nvfp4(
+            intermediate,
+            x_global_scale,
+            group_size=group_size,
+        )
+
+    def quantize(self, x, w):
+        x_global_scale = (448.0 * 6.0) / torch.amax(torch.abs(x.flatten()), dim=-1).to(
+            torch.float32
+        )
+        w_global_scale = (448.0 * 6.0) / torch.amax(torch.abs(w.flatten()), dim=-1).to(
+            torch.float32
+        )
+        global_scale = 1 / (x_global_scale * w_global_scale)
+
+        xq, x_scale = triton_quantize_nvfp4(x, x_global_scale)
+        wq, w_scale = triton_quantize_nvfp4(w, w_global_scale)
+        return xq, wq, x_scale, w_scale, global_scale
+
+    def compute(self, xq, wq, x_scale, w_scale, global_scale):
+        return torch.ops.mslk.f4f4bf16(
+            xq, wq, x_scale, w_scale, global_scale=global_scale
+        )
+
+    def quantize_and_compute(self, x, w):
+        return self.quantize(x, w)
+
+    @property
+    def name(self) -> str:
+        return "nvfp4_quantize"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM100, Accelerator.NVIDIA_SM103}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP4
+
+
+@register_gemm_op
+class MXFP4Gemm(GemmOpBase):
+    """
+    MXFP4 matmul with block-wise scaling.
+    """
+
+    def quantize(self, x, w):
+        xq, x_scale = triton_quantize_mx4_unpack(x)
+        wq, w_scale = triton_quantize_mx4_unpack(w)
+        return xq, wq, x_scale, w_scale
+
+    def compute(self, xq, wq, x_scale, w_scale):
+        return torch.ops.mslk.f4f4bf16(xq, wq, x_scale, w_scale)
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale = self.quantize(x, w)
+        return self.compute(xq, wq, x_scale, w_scale)
+
+    @property
+    def name(self) -> str:
+        return "cutlass_f4f4bf16"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM100, Accelerator.NVIDIA_SM103}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.REGULAR}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP4
+
+
+@register_gemm_op
+class MXFP4StackedGroupedGemm(GemmOpBase):
+    """
+    MXFP4 grouped matmul with blockwise scaling and stacked inputs.
+    """
+
+    def preprocess(self, x, w):
+        m_values = [i.shape[0] for i in x]
+        m_sizes = torch.tensor(m_values).to(dtype=torch.int64, device=x[0].device)
+        wq, w_scale = zip(*[triton_quantize_mx4_unpack(i) for i in w])
+        wq = torch.stack(wq, dim=0).contiguous()
+        w_scale = torch.stack(w_scale, dim=0).contiguous()
+        return x, wq, w_scale, m_sizes
+
+    def quantize(self, x, wq, w_scale, m_sizes):
+        starting_row_after_padding_list = [0]
+        xq_list = []
+        x_scale_list = []
+        for i in range(m_sizes.shape[0]):
+            scale_slice = x[i]
+            if m_sizes[i].item() != 0:
+                xq, x_scale = triton_quantize_mx4_unpack(scale_slice)
+                xq_list.append(xq)
+                x_scale_list.append(x_scale)
+                starting_row_after_padding_list.append(
+                    starting_row_after_padding_list[i]
+                    + x_scale.numel() // (x[0].shape[1] // 32)
+                )
+            else:
+                starting_row_after_padding_list.append(
+                    starting_row_after_padding_list[i]
+                )
+        xq = torch.cat(xq_list, dim=0).contiguous()
+        x_scale = torch.cat(x_scale_list, dim=0).contiguous()
+        x_scale = x_scale.reshape(-1, x[0].shape[-1] // 32)
+        xq = xq.view(-1, xq.shape[-1])
+        return (
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            m_sizes,
+            torch.tensor(starting_row_after_padding_list, device=xq.device),
+        )
+
+    def compute(self, xq, wq, x_scale, w_scale, m_sizes, starting_row_after_padding):
+        return torch.ops.mslk.f4f4bf16_grouped_stacked(
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            m_sizes,
+            starting_row_after_padding=starting_row_after_padding,
+        )
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale, m_sizes, starting_row_after_padding = self.quantize(
+            x, w
+        )
+        return self.compute(
+            xq, wq, x_scale, w_scale, m_sizes, starting_row_after_padding
+        )
+
+    @property
+    def name(self) -> str:
+        return "cutlass_f4f4bf16_grouped_stacked"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM100, Accelerator.NVIDIA_SM103}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP4
+
+
+@register_gemm_op
+class NVFP4StackedGroupedGemm(GemmOpBase):
+    """
+    NVFP4 grouped matmul with blockwise scaling and stacked inputs.
+    """
+
+    def preprocess(self, x, w):
+        m_values = [i.shape[0] for i in x]
+        m_sizes = torch.tensor(m_values).to(dtype=torch.int64, device=x[0].device)
+        x = torch.concat(x, dim=0).contiguous()
+
+        def get_global_scale(x, w, m_sizes):
+            G = len(w)
+            w_global_scale = []
+            global_scale = []
+
+            cumulative_sum = torch.zeros(
+                m_sizes.shape[0] + 1, dtype=torch.int64, device=m_sizes.device
+            )
+            cumulative_sum[1:] = torch.cumsum(m_sizes, dim=0)
+
+            x_global_scale, tensor_idx = calculate_group_max(x, m_sizes=m_sizes)
+
+            for i in range(G):
+                w_global_scale_ = (448.0 * 6.0) / torch.amax(
+                    torch.abs(w[i].flatten()), dim=-1
+                ).to(torch.float32)
+
+                global_scale_ = 1 / (x_global_scale[i] * w_global_scale_)
+
+                w_global_scale.append(w_global_scale_)
+                global_scale.append(global_scale_)
+
+            return x_global_scale, w_global_scale, global_scale, tensor_idx
+
+        # Compute global scale for each group
+        G = m_sizes.numel()
+        x_global_scale, w_global_scale, global_scale, tensor_idx = get_global_scale(
+            x, w, m_sizes
+        )
+        global_scale = torch.stack(global_scale, dim=0).contiguous()
+
+        wq, w_scale = zip(
+            *[triton_quantize_nvfp4(w[i], w_global_scale[i]) for i in range(G)]
+        )
+        wq = torch.stack(wq, dim=0).contiguous()
+        w_scale = torch.stack(w_scale, dim=0).contiguous()
+
+        return x, wq, w_scale, x_global_scale, global_scale, m_sizes, tensor_idx
+
+    def quantize(
+        self, x, wq, w_scale, x_global_scale, global_scale, m_sizes, tensor_idx
+    ):
+        # alternative methods, may be useful in some scenarios
+        """
+        starting_row_after_padding, belong_indices, row_within_tensor = (
+            nvfp4_fused_padding_cumsum_and_segmented_arange(m_sizes, x.shape[0])
+            # fused_single_block_cumsum_and_segmented_arange(m_sizes, x.shape[0])
+        )
+
+        xq, x_scale = triton_nvfp4_quant_stacked(
+            x,
+            x_global_scale[0],
+            belong_indices,
+            starting_row_after_padding,
+            row_within_tensor,
+        )
+        """
+
+        # we can optionally set optional_tensor_idx to None to run the alternative method
+        xq, x_scale, starting_row_after_padding = mega_fp4_quantize_kernel(
+            m_sizes, x, x_global_scale, optional_tensor_idx=tensor_idx
+        )
+
+        x_scale = x_scale.reshape(-1, x.shape[1] // 16)
+        return (
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            m_sizes,
+            global_scale,
+            starting_row_after_padding,
+        )
+
+    def compute(
+        self,
+        xq,
+        wq,
+        x_scale,
+        w_scale,
+        m_sizes,
+        global_scale,
+        starting_row_after_padding,
+    ):
+        gemm_result = torch.ops.mslk.f4f4bf16_grouped_stacked(
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            m_sizes,
+            global_scale,
+            starting_row_after_padding,
+            use_mx=False,
+        )
+        return gemm_result
+
+    def quantize_and_compute(
+        self, x, wq, w_scale, x_global_scale, global_scale, m_sizes, tensor_idx
+    ):
+        (
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            m_sizes,
+            global_scale,
+            starting_row_after_padding,
+        ) = self.quantize(
+            x, wq, w_scale, x_global_scale, global_scale, m_sizes, tensor_idx
+        )
+        return self.compute(
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            m_sizes,
+            global_scale,
+            starting_row_after_padding,
+        )
+
+    @property
+    def name(self) -> str:
+        return "cutlass_nv_f4f4bf16_grouped_stacked"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM100, Accelerator.NVIDIA_SM103}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP4
+
+
+# Broken with cuda graph
+# @register_gemm_op
+class NVFP4StackedGroupedGemmPackUnpack(GemmOpBase):
+    """
+    NVFP4 grouped matmul with blockwise scaling and stacked inputs.
+    """
+
+    def preprocess(self, x, w):
+        m_values = [i.shape[0] for i in x]
+        m_sizes = torch.tensor(m_values).to(dtype=torch.int64, device=x[0].device)
+        x = torch.concat(x, dim=0).contiguous()
+
+        def get_global_scale(x, w):
+            G = len(w)
+            x_global_scale = []
+            w_global_scale = []
+            global_scale = []
+
+            x_global_scale_ = (448.0 * 6.0) / torch.amax(
+                torch.abs(x.flatten()), dim=-1
+            ).to(torch.float32)
+
+            for i in range(G):
+                w_global_scale_ = (448.0 * 6.0) / torch.amax(
+                    torch.abs(w[i].flatten()), dim=-1
+                ).to(torch.float32)
+
+                global_scale_ = 1 / (x_global_scale_ * w_global_scale_)
+
+                x_global_scale.append(x_global_scale_)
+                w_global_scale.append(w_global_scale_)
+                global_scale.append(global_scale_)
+
+            return x_global_scale, w_global_scale, global_scale
+
+        # Compute global scale for each group
+        G = m_sizes.numel()
+        x_global_scale, w_global_scale, global_scale = get_global_scale(x, w)
+
+        global_scale = torch.stack(global_scale, dim=0).contiguous()
+
+        wq, w_scale = zip(
+            *[triton_quantize_nvfp4(w[i], w_global_scale[i]) for i in range(G)]
+        )
+        wq = torch.stack(wq, dim=0).contiguous()
+        w_scale = torch.stack(w_scale, dim=0).contiguous()
+        x_global_scale = torch.tensor(x_global_scale, device=m_sizes.device)
+        return (
+            x,
+            wq,
+            w_scale,
+            x_global_scale,
+            global_scale,
+            m_sizes,
+        )
+
+    def quantize(self, x, wq, w_scale, x_global_scale, global_scale, m_sizes):
+        # alternative packing methods that only uses the overall global scale rather than per tensor
+        """
+        packed = mega_fp4_pack(x, x_global_scale[0])
+        """
+        packed = mega_fp4_pack(
+            x,
+            x_global_scale,
+            per_tensor=True,
+            m_sizes=m_sizes,
+        )
+        xq, x_scale, starting_row_after_padding = mega_fp4_unpack(m_sizes, packed)
+        xq_other, x_scale_other, starting_row_after_padding_other = (
+            mega_fp4_quantize_kernel(
+                m_sizes,
+                x,
+                x_global_scale,
+            )
+        )
+
+        x_scale = x_scale.reshape(-1, x.shape[1] // 16)
+        x_scale_other = x_scale_other.reshape(-1, x.shape[1] // 16)
+        return (
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            m_sizes,
+            global_scale,
+            starting_row_after_padding,
+            xq_other,
+            x_scale_other,
+            starting_row_after_padding_other,
+        )
+
+    def compute(
+        self,
+        xq,
+        wq,
+        x_scale,
+        w_scale,
+        m_sizes,
+        global_scale,
+        starting_row_after_padding,
+        xq_other,
+        x_scale_other,
+        starting_row_after_padding_other,
+    ):
+        ref_solution = torch.ops.mslk.f4f4bf16_grouped_stacked(
+            xq_other,
+            wq,
+            x_scale_other,
+            w_scale,
+            m_sizes,
+            global_scale,
+            starting_row_after_padding_other,
+            use_mx=False,
+        )
+        gemm_result = torch.ops.mslk.f4f4bf16_grouped_stacked(
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            m_sizes,
+            global_scale,
+            starting_row_after_padding,
+            use_mx=False,
+        )
+        assert torch.allclose(ref_solution, gemm_result)
+
+        return gemm_result
+
+    def quantize_and_compute(
+        self, x, wq, w_scale, x_global_scale, global_scale, m_sizes
+    ):
+        (
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            m_sizes,
+            global_scale,
+            starting_row_after_padding,
+            xq_other,
+            x_scale_other,
+            starting_row_after_padding_other,
+        ) = self.quantize(x, wq, w_scale, x_global_scale, global_scale, m_sizes)
+        return self.compute(
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            m_sizes,
+            global_scale,
+            starting_row_after_padding,
+            xq_other,
+            x_scale_other,
+            starting_row_after_padding_other,
+        )
+
+    @property
+    def name(self) -> str:
+        return "cutlass_nv_f4f4bf16_grouped_stacked_pack_unpack"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM100, Accelerator.NVIDIA_SM103}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP4
+
+
+@register_gemm_op
+class BF16GroupedGemm2d3d(GemmOpBase):
+    """
+    Torch BF16 grouped GEMM with 2D inputs and 3D weights.
+    """
+
+    def preprocess(self, x, w):
+        assert isinstance(x, list)
+        assert isinstance(w, list)
+        offs = torch.tensor(
+            [i.shape[0] for i in x], dtype=torch.int32, device=x[0].device
+        )
+        offs = torch.cumsum(offs, dim=0).to(torch.int32)
+        x = torch.cat(x, dim=0).contiguous()  # (G * M, K)
+        w = torch.stack(w, dim=0).contiguous()  # (G, N, K)
+        return x, w, offs
+
+    def quantize(self, x, w, offs):
+        return x, w, offs
+
+    def compute(self, x, w, offs):
+        return torch._grouped_mm(
+            x,
+            w.transpose(-2, -1),
+            offs=offs,
+        )
+
+    def quantize_and_compute(self, x, w, offs):
+        x, w, offs = self.quantize(x, w)
+        return self.compute(x, w, offs)
+
+    @property
+    def name(self) -> str:
+        return "bf16_baseline_grouped_2d_3d"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return set(Accelerator)
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.BF16
+
+
+@register_gemm_op
+class MXFP8GroupedGemm2d3d(GemmOpBase):
+    """
+    MXFP8 grouped GEMM with 2D inputs and 3D weights.
+    """
+
+    def preprocess(self, x, w):
+        assert isinstance(x, list)
+        assert isinstance(w, list)
+        x = torch.cat(x, dim=0).contiguous()  # (G * M, K)
+        w = torch.stack(w, dim=0).contiguous()  # (G, N, K)
+        return x, w
+
+    def quantize(self, x, w):
+        block_size = 32
+        G, N, K = w.shape
+        total_M = x.shape[0]
+        group_size = total_M // G
+        input_group_end_offsets = torch.arange(
+            group_size, total_M + 1, group_size, dtype=torch.int32, device=x.device
+        )
+
+        # For each constituent 2d subtensor in the 3d weights, quantize and convert scale to blocked format separately,
+        # as they each used for independent gemm in the grouped gemm.
+        wq_list = []
+        w_scale_list = []
+        for i in range(G):
+            w_scale, wq = to_mxfp8(w[i])
+            w_scale = _to_blocked(w_scale)
+            wq_list.append(wq)
+            w_scale_list.append(w_scale)
+        wq = torch.stack(wq_list, dim=0).contiguous()
+        w_scale = torch.stack(w_scale_list, dim=0).contiguous()
+
+        # For each group along `total_M` in the 2D tensor, quantize and convert scale to blocked format separately,
+        # as they each used for independent gemm in the grouped gemm.
+        xq_list = []
+        x_scale_list = []
+        for i in range(G):
+            prev_group_end = 0 if i == 0 else input_group_end_offsets[i - 1]
+            curr_group_end = input_group_end_offsets[i]
+            group_size = curr_group_end - prev_group_end
+            if group_size > 0:
+                x_slice = x[prev_group_end:curr_group_end, :]
+                x_scale, xq = to_mxfp8(x_slice)
+                x_scale = _to_blocked(x_scale)
+                xq_list.append(xq)
+                x_scale_list.append(x_scale)
+        xq = torch.cat(xq_list, dim=0).contiguous()
+        x_scale = torch.cat(x_scale_list, dim=0).contiguous()
+        x_scale = x_scale.reshape(-1, K // block_size)
+        xq = xq.view(-1, xq.shape[-1])
+        return xq, wq, x_scale, w_scale, input_group_end_offsets
+
+    def compute(self, xq, wq, x_scale, w_scale, input_group_end_offsets):
+        return torch.ops.mslk.mx8mx8bf16_grouped_mm(
+            xq,
+            wq.transpose(-2, -1),
+            x_scale,
+            w_scale,
+            input_group_end_offsets,
+        )
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale, input_group_end_offsets = self.quantize(x, w)
+        return self.compute(
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            input_group_end_offsets,
+        )
+
+    @property
+    def name(self) -> str:
+        return "cutlass_mx8mx8bf16_grouped_mm_2d_3d"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM100, Accelerator.NVIDIA_SM103}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8
+
+
+@register_gemm_op
+class MXFP8GroupedGemm2d2d(GemmOpBase):
+    """
+    MXFP8 grouped GEMM with 2D inputs and 3D weights.
+    """
+
+    def preprocess(self, x, w):
+        assert isinstance(x, list)
+        assert isinstance(w, list)
+        G = len(x)
+        x = torch.cat(x, dim=1).contiguous()  # (M, total_K)
+        w = torch.cat(w, dim=1).contiguous()  # (N, total_K)
+        return x, w, G
+
+    def quantize(self, x, w, G):
+        # Simulate 2d-2d grouped gemm in backward pass `grad_weight = grad_output_t @ input`,
+        # where we use "K" as the contracting dim which has "G" groups.
+        M, total_K = x.shape
+        N, _ = w.shape
+        group_size = total_K // G
+        input_group_end_offsets = torch.arange(
+            group_size, total_K + 1, group_size, dtype=torch.int32, device=x.device
+        )
+
+        # Convert scales to blocked format.
+        x_list = []
+        w_list = []
+        x_blocked_scale_list = []
+        w_blocked_scale_list = []
+
+        def round_up(x: int, y: int) -> int:
+            return ((x + y - 1) // y) * y
+
+        for group_idx in range(G):
+            # to_mxfp8 per group
+            prev_group_end_offset = (
+                0 if group_idx == 0 else input_group_end_offsets[group_idx - 1]
+            )
+            curr_group_end_offset = input_group_end_offsets[group_idx]
+            group_size = curr_group_end_offset - prev_group_end_offset
+            if group_size > 0:
+                x_slice = x[
+                    :, prev_group_end_offset:curr_group_end_offset
+                ].contiguous()  # (M, K_group)
+                w_slice = w[
+                    :, prev_group_end_offset:curr_group_end_offset
+                ].contiguous()  # (N, K_group)
+                x_scale_slice, xq_slice = to_mxfp8(
+                    x_slice
+                )  # scale shape -> (M, K_group // 32)
+                w_scale_slice, wq_slice = to_mxfp8(
+                    w_slice
+                )  # scale shape -> (N, K_group // 32)
+                x_list.append(xq_slice)
+                w_list.append(wq_slice)
+
+                # Convert scales to blocked format.
+                x_scale_slice_blocked = _to_blocked(
+                    x_scale_slice
+                )  # (round_up(M, 128), round_up(K_group//32, 4))
+                w_scale_slice_blocked = _to_blocked(
+                    w_scale_slice
+                )  # (round_up(N, 128), round_up(K_group//32, 4))
+                x_blocked_scale_list.append(x_scale_slice_blocked)
+                w_blocked_scale_list.append(w_scale_slice_blocked)
+
+        # Assemble the full XQ and WQ
+        xq = torch.cat(x_list, dim=1).contiguous()
+        wq = torch.cat(w_list, dim=1).contiguous()
+
+        # Combine all XQ groups blocked scales into one tensor.
+        x_blocked_scales = torch.cat(x_blocked_scale_list, dim=0)
+        M_rounded = round_up(M, 128)
+        x_blocked_scales = x_blocked_scales.reshape(M_rounded, -1)
+
+        # Combine all WQ groups blocked scales into one tensor.
+        w_blocked_scales = torch.cat(w_blocked_scale_list, dim=0)
+        N_rounded = round_up(N, 128)
+        w_blocked_scales = w_blocked_scales.reshape(N_rounded, -1)
+        return xq, wq, x_blocked_scales, w_blocked_scales, input_group_end_offsets
+
+    def compute(self, xq, wq, x_scale, w_scale, input_group_end_offsets):
+        return torch.ops.mslk.mx8mx8bf16_grouped_mm(
+            xq,
+            wq.transpose(-2, -1),
+            x_scale,
+            w_scale,
+            input_group_end_offsets,
+        )
+
+    def quantize_and_compute(self, x, w):
+        xq, wq, x_scale, w_scale, input_group_end_offsets = self.quantize(x, w)
+        return self.compute(
+            xq,
+            wq,
+            x_scale,
+            w_scale,
+            input_group_end_offsets,
+        )
+
+    @property
+    def name(self) -> str:
+        return "cutlass_mx8mx8bf16_grouped_mm_2d_2d"
+
+    @property
+    def supported_accelerators(self) -> set[Accelerator]:
+        return {Accelerator.NVIDIA_SM100, Accelerator.NVIDIA_SM103}
+
+    @property
+    def supported_gemm_types(self) -> set[GemmType]:
+        return {GemmType.GROUPED}
+
+    @property
+    def compute_dtype(self) -> ComputeDtype:
+        return ComputeDtype.FP8