Add files using upload-large-folder tool

lib/python3.12/site-packages/torchao/dtypes/__init__.py

@@ -0,0 +1,64 @@
+from . import affine_quantized_tensor_ops
+from .affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    to_affine_quantized_floatx,
+    to_affine_quantized_floatx_static,
+    # experimental, will be merged into floatx in the future
+    to_affine_quantized_fpx,
+    to_affine_quantized_intx,
+    to_affine_quantized_intx_static,
+)
+from .floatx import (
+    CutlassSemiSparseLayout,
+    Float8Layout,
+)
+from .nf4tensor import NF4Tensor, to_nf4
+from .uintx import (
+    BlockSparseLayout,
+    CutlassInt4PackedLayout,
+    Int4CPULayout,
+    Int4XPULayout,
+    MarlinQQQLayout,
+    MarlinQQQTensor,
+    MarlinSparseLayout,
+    PackedLinearInt8DynamicActivationIntxWeightLayout,
+    QDQLayout,
+    SemiSparseLayout,
+    TensorCoreTiledLayout,
+    UintxLayout,
+    to_marlinqqq_quantized_intx,
+)
+from .utils import (
+    Layout,
+    PlainLayout,
+)
+
+__all__ = [
+    "NF4Tensor",
+    "to_nf4",
+    "AffineQuantizedTensor",
+    "to_affine_quantized_intx",
+    "to_affine_quantized_intx_static",
+    "to_affine_quantized_fpx",
+    "to_affine_quantized_floatx",
+    "to_affine_quantized_floatx_static",
+    "to_marlinqqq_quantized_intx",
+    "Layout",
+    "PlainLayout",
+    "SemiSparseLayout",
+    "TensorCoreTiledLayout",
+    "Float8Layout",
+    "MarlinSparseLayout",
+    "affine_quantized_tensor_ops",
+    "BlockSparseLayout",
+    "UintxLayout",
+    "MarlinQQQTensor",
+    "MarlinQQQLayout",
+    "Int4CPULayout",
+    "CutlassInt4PackedLayout",
+    "CutlassSemiSparseLayout",
+    "QDQLayout",
+    "PackedLinearInt8DynamicActivationIntxWeightLayout",
+    "to_affine_quantized_packed_linear_int8_dynamic_activation_intx_weight",
+    "Int4XPULayout",
+]

lib/python3.12/site-packages/torchao/dtypes/_nf4tensor_api.py

@@ -0,0 +1,39 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+
+from torchao.core.config import AOBaseConfig
+from torchao.dtypes.nf4tensor import NF4Tensor
+from torchao.quantization.transform_module import (
+    register_quantize_module_handler,
+)
+
+
+class NF4WeightOnlyConfig(AOBaseConfig):
+    """
+    Note: the file location of this workflow is temporary.
+    TODO(future PR): integrate this properly into torchao's directory structure
+    """
+
+    block_size: int = 64
+    scaler_block_size: int = 256
+
+
+# for bc
+nf4_weight_only = NF4WeightOnlyConfig
+
+
+@register_quantize_module_handler(NF4WeightOnlyConfig)
+def _nf4_weight_only_transform(
+    module: torch.nn.Module,
+    config: NF4WeightOnlyConfig,
+) -> torch.nn.Module:
+    block_size = config.block_size
+    scaler_block_size = config.scaler_block_size
+
+    new_weight = NF4Tensor.from_tensor(module.weight, block_size, scaler_block_size)
+    module.weight = torch.nn.Parameter(new_weight, requires_grad=False)
+    return module

lib/python3.12/site-packages/torchao/dtypes/affine_quantized_tensor.py

@@ -0,0 +1,507 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+import logging
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+
+from torchao.dtypes.utils import (
+    AQTTensorImpl,
+    Layout,
+    PlainLayout,
+)
+from torchao.quantization.quant_primitives import (
+    FP8_TYPES,
+    MappingType,
+    ZeroPointDomain,
+    choose_qparams_affine,
+    choose_qparams_affine_floatx,
+    choose_qparams_and_quantize_affine_hqq,
+    dequantize_affine,
+    dequantize_affine_floatx,
+    quantize_affine,
+    quantize_affine_floatx,
+)
+from torchao.utils import (
+    TORCH_VERSION_AT_LEAST_2_5,
+    TorchAOBaseTensor,
+)
+
+logger = logging.getLogger(__name__)
+aten = torch.ops.aten
+
+__all__ = [
+    "AffineQuantizedTensor",
+    "register_layout",
+    "to_affine_quantized_intx",
+    "to_affine_quantized_floatx",
+    "to_affine_quantized_intx_static",
+    "to_affine_quantized_floatx_static",
+    "to_affine_quantized_fpx",
+]
+
+
+##############################
+# Tensor Subclass Definition #
+##############################
+class AffineQuantizedTensor(TorchAOBaseTensor):
+    """Affine quantized tensor subclass. Affine quantization means we quantize the floating point tensor with an affine transformation:
+    quantized_tensor = float_tensor / scale + zero_point
+
+    To see what happens during choose_qparams, quantization and dequantization for affine quantization,
+    please checkout https://github.com/pytorch/ao/blob/main/torchao/quantization/quant_primitives.py
+    and check the three quant primitive ops: choose_qparams_affine, quantize_affine qand dequantize_affine
+
+    The shape and dtype of the tensor subclass represent how the tensor subclass looks externally,
+    regardless of the internal representation's type or orientation.
+
+    fields:
+        - tensor_impl (AQTTensorImpl): tensor that serves as a general tensor impl storage for the quantized data,
+            e.g. storing plain tensors (int_data, scale, zero_point) or packed formats depending on device and operator/kernel
+        - block_size (Tuple[int, ...]): granularity of quantization, this means the size of the tensor elements that's sharing the same qparam
+            e.g. when size is the same as the input tensor dimension, we are using per tensor quantization
+        - shape (torch.Size): the shape for the original high precision Tensor
+        - quant_min (Optional[int]): minimum quantized value for the Tensor, if not specified, it will be derived from dtype of `int_data`
+        - quant_max (Optional[int]): maximum quantized value for the Tensor, if not specified, it will be derived from dtype of `int_data`
+        - zero_point_domain (ZeroPointDomain): the domain that zero_point is in, should be either integer or float
+            if zero_point is in integer domain, zero point is added to the quantized integer value during quantization
+            if zero_point is in floating point domain, zero point is subtracted from the floating point (unquantized) value during quantization
+            default is ZeroPointDomain.INT
+        - dtype: dtype for original high precision tensor, e.g. torch.float32
+    """
+
+    @staticmethod
+    def __new__(
+        cls,
+        tensor_impl: AQTTensorImpl,
+        block_size: Tuple[int, ...],
+        shape: torch.Size,
+        quant_min: Optional[Union[int, float]] = None,
+        quant_max: Optional[Union[int, float]] = None,
+        zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT,
+        dtype=None,
+        strides=None,
+    ):
+        if zero_point_domain is None:
+            raise ValueError("please use ZeroPointDomain.NONE instead of None")
+        kwargs = {}
+        kwargs["device"] = tensor_impl.device
+        kwargs["layout"] = (
+            kwargs.get("layout") if kwargs.get("layout", False) else tensor_impl.layout
+        )
+        kwargs["dtype"] = dtype
+        if strides is not None:
+            kwargs["strides"] = strides
+        kwargs["requires_grad"] = False
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        tensor_impl: AQTTensorImpl,
+        block_size: Tuple[int, ...],
+        shape: torch.Size,
+        quant_min: Optional[Union[int, float]] = None,
+        quant_max: Optional[Union[int, float]] = None,
+        zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT,
+        dtype=None,
+        strides=None,
+    ):
+        self.tensor_impl = tensor_impl
+        self.block_size = block_size
+        self.quant_min = quant_min
+        self.quant_max = quant_max
+        self.zero_point_domain = zero_point_domain
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(tensor_impl={self.tensor_impl}, block_size={self.block_size}, "
+            f"shape={self.shape}, device={self.device}, dtype={self.dtype}, requires_grad={self.requires_grad})"
+        )
+
+    def _quantization_type(self):
+        return f"shape={self.shape}, block_size={self.block_size}, device={self.device}, _layout={self._layout}, tensor_impl_dtype={self.tensor_impl.dtype}, quant_min={self.quant_min}, quant_max={self.quant_max}"
+
+    def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor:
+        if output_dtype is None:
+            output_dtype = self.dtype
+
+        from torchao.dtypes.floatx import FloatxTensorCoreLayout
+
+        if isinstance(self._layout, FloatxTensorCoreLayout):
+            int_data, scale = self.tensor_impl.get_plain()
+            return dequantize_affine_floatx(
+                int_data,
+                scale,
+                self._layout.ebits,
+                self._layout.mbits,
+                output_dtype=output_dtype,
+            )
+        else:
+            data, scale, zero_point = self.tensor_impl.get_plain()
+            dq = dequantize_affine(
+                data,
+                self.block_size,
+                scale,
+                zero_point,
+                data.dtype,
+                self.quant_min,
+                self.quant_max,
+                self.zero_point_domain,
+                output_dtype=output_dtype,
+            )
+            from torchao.dtypes.uintx import TensorCoreTiledLayout
+
+            if isinstance(self._layout, TensorCoreTiledLayout):
+                # need to return to original shape if tensor was padded
+                # in preprocessing
+                # TODO: we could add an API for this if there are more use cases
+                # (e.g. dequant_post_process) in TensorImpl or Layout
+                for dim, dim_size in enumerate(self.shape):
+                    dq = dq.narrow(dim, 0, dim_size)
+            return dq
+
+    def __tensor_flatten__(self):
+        # This is used in rumtime to unwrap AffineQuantizedTensor activations.
+        # AffineQuantizedTensor has __torch_function__ override:
+        # Each getattr will go through it, which is up to 10x slower than default attribute access.
+        with torch._C.DisableTorchFunctionSubclass():
+            return ["tensor_impl"], [
+                self.block_size,
+                self.shape,
+                self.quant_min,
+                self.quant_max,
+                self.zero_point_domain,
+                self.dtype,
+            ]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        tensor_impl = tensor_data_dict["tensor_impl"]
+        block_size, shape, quant_min, quant_max, zero_point_domain, dtype = (
+            tensor_attributes
+        )
+        return cls(
+            tensor_impl,
+            block_size,
+            shape if outer_size is None else outer_size,
+            quant_min,
+            quant_max,
+            zero_point_domain,
+            dtype=dtype,
+            strides=outer_stride,
+        )
+
+    @classmethod
+    def from_hp_to_intx(
+        cls,
+        input_float: torch.Tensor,
+        mapping_type: MappingType,
+        block_size: Tuple[int, ...],
+        target_dtype: torch.dtype,
+        quant_min: Optional[int] = None,
+        quant_max: Optional[int] = None,
+        eps: Optional[float] = None,
+        scale_dtype: Optional[torch.dtype] = None,
+        zero_point_dtype: Optional[torch.dtype] = None,
+        preserve_zero: bool = True,
+        zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT,
+        _layout: Layout = PlainLayout(),
+        use_hqq: bool = False,
+    ):
+        """Convert a high precision tensor to an integer affine quantized tensor."""
+        original_shape = input_float.shape
+        input_float = _layout.pre_process(input_float)
+
+        if use_hqq:
+            assert (
+                zero_point_domain == ZeroPointDomain.FLOAT
+                and mapping_type == MappingType.ASYMMETRIC
+                and quant_min == 0
+            ), "Invalid input parameters for HQQ quantization."
+            nbits = int(math.log2(quant_max + 1))
+            axis = 1 if (block_size[0] == 1) else 0
+            group_size = max(block_size)
+            compute_dtype = (
+                zero_point_dtype
+                if (zero_point_dtype is not None)
+                else input_float.dtype
+            )
+            device = input_float.device
+            from torchao.dtypes import Int4CPULayout
+            from torchao.dtypes.uintx import TensorCoreTiledLayout
+
+            data, scale, zero_point, _ = choose_qparams_and_quantize_affine_hqq(
+                input_float,
+                nbits=nbits,
+                group_size=group_size,
+                axis=axis,
+                compute_dtype=compute_dtype,
+                device=device,
+                verbose=False,
+                raw_output=not isinstance(
+                    _layout, (TensorCoreTiledLayout, PlainLayout, Int4CPULayout)
+                ),
+                # raw_output=False is basically the 'convert to TensorCoreTiledLayout zero_point version' option (add scale*midpoint)
+                # note in choose_qparams_affine, preserve_zero = False does this same thing while also controlling whether
+                # zero is preserved.
+                # TODO uncouple preserve_zero and conversion of zero_point to TensorCoreTiledLayout version
+                # TODO move the conversion of zero_point out of quant_primitives and into TensorCoreTiledLayout.from_plain
+                # TODO change PlainLayout to use raw_output.
+            )
+            data = data.to(target_dtype)
+        else:
+            scale, zero_point = choose_qparams_affine(
+                input_float,
+                mapping_type,
+                block_size,
+                target_dtype,
+                quant_min,
+                quant_max,
+                eps,
+                scale_dtype,
+                zero_point_dtype,
+                preserve_zero,
+                zero_point_domain,
+            )
+            # choose_qparams_affine is a custom op that does support returning optional Tensors. We thus set the zero_point to None if its domain is None
+            if zero_point_domain == ZeroPointDomain.NONE:
+                zero_point = None
+            data = quantize_affine(
+                input_float,
+                block_size,
+                scale,
+                zero_point,
+                target_dtype,
+                quant_min,
+                quant_max,
+                zero_point_domain,
+            )
+            # Note: output will be uint8 tensor for sub byte tensors for now
+
+        data, scale, zero_point = _layout.post_process(
+            data, scale, zero_point, block_size
+        )
+        tensor_impl_ctr = get_tensor_impl_constructor(type(_layout))
+        tensor_impl = tensor_impl_ctr(data, scale, zero_point, _layout)
+        return cls(
+            tensor_impl,
+            block_size,
+            original_shape,
+            quant_min,
+            quant_max,
+            zero_point_domain,
+            dtype=input_float.dtype,
+        )
+
+    @classmethod
+    def from_hp_to_intx_static(
+        cls,
+        input_float: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        block_size: Tuple[int, ...],
+        target_dtype: torch.dtype,
+        quant_min: Optional[int] = None,
+        quant_max: Optional[int] = None,
+        zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT,
+        _layout: Layout = PlainLayout(),
+    ):
+        """Create an integer AffineQuantizedTensor from a high precision tensor using static parameters."""
+        if zero_point_domain is None:
+            raise ValueError("please use ZeroPointDomain.NONE instead of None")
+        elif zero_point_domain is ZeroPointDomain.NONE and zero_point is not None:
+            raise ValueError("zero_point should be None when zero_point_domain is NONE")
+        if target_dtype not in FP8_TYPES:
+            assert zero_point is not None, (
+                "zero_point must be specified for non-fp8 types"
+            )
+        original_shape = input_float.shape
+        input_float, scale, zero_point = _layout.pre_process_static(
+            input_float, scale, zero_point, block_size
+        )
+
+        int_data = quantize_affine(
+            input_float,
+            block_size,
+            scale,
+            zero_point,
+            target_dtype,
+            quant_min,
+            quant_max,
+            zero_point_domain,
+        )
+
+        int_data, scale, zero_point = _layout.post_process(
+            int_data,
+            scale,
+            zero_point,
+            block_size,
+        )
+
+        tensor_impl_ctr = get_tensor_impl_constructor(type(_layout))
+        tensor_impl = tensor_impl_ctr(int_data, scale, zero_point, _layout)
+        return cls(
+            tensor_impl,
+            block_size,
+            original_shape,
+            quant_min,
+            quant_max,
+            zero_point_domain,
+            dtype=input_float.dtype,
+        )
+
+    @classmethod
+    def from_hp_to_floatx(
+        cls,
+        input_float: torch.Tensor,
+        block_size: Tuple[int, ...],
+        target_dtype: torch.dtype,
+        _layout: Layout,
+        scale_dtype: Optional[torch.dtype] = None,
+    ):
+        """Convert a high precision tensor to a float8 quantized tensor."""
+        if target_dtype in FP8_TYPES:
+            return cls.from_hp_to_intx(
+                input_float=input_float,
+                mapping_type=MappingType.SYMMETRIC,
+                block_size=block_size,
+                target_dtype=target_dtype,
+                quant_min=math.ceil(torch.finfo(target_dtype).min),
+                quant_max=math.ceil(torch.finfo(target_dtype).max),
+                eps=torch.finfo(torch.float32).eps,
+                scale_dtype=scale_dtype,
+                zero_point_dtype=None,
+                preserve_zero=True,
+                zero_point_domain=ZeroPointDomain.NONE,
+                _layout=_layout,
+                use_hqq=False,
+            )
+        else:
+            raise NotImplementedError(
+                f"Unsupported dtype {target_dtype} for from_hp_to_floatx"
+            )
+
+    @classmethod
+    def from_hp_to_floatx_static(
+        cls,
+        input_float: torch.Tensor,
+        scale: torch.Tensor,
+        block_size: Tuple[int, ...],
+        target_dtype: torch.dtype,
+        _layout: Layout,
+    ):
+        """Create a float8 AffineQuantizedTensor from a high precision tensor using static parameters."""
+        if target_dtype in FP8_TYPES:
+            return cls.from_hp_to_intx_static(
+                input_float=input_float,
+                scale=scale,
+                zero_point=None,
+                block_size=block_size,
+                target_dtype=target_dtype,
+                quant_min=math.ceil(torch.finfo(target_dtype).min),
+                quant_max=math.ceil(torch.finfo(target_dtype).max),
+                zero_point_domain=ZeroPointDomain.NONE,
+                _layout=_layout,
+            )
+        else:
+            raise NotImplementedError(
+                f"Unsupported dtype {target_dtype} for from_hp_to_floatx_static"
+            )
+
+    @classmethod
+    def from_hp_to_fpx(
+        cls,
+        input_float: torch.Tensor,
+        _layout: Layout,
+    ):
+        """Create a floatx AffineQuantizedTensor from a high precision tensor. Floatx is represented as ebits and mbits, and supports the representation of float1-float7."""
+        from torchao.dtypes.floatx import FloatxTensorCoreLayout
+
+        assert isinstance(_layout, FloatxTensorCoreLayout), (
+            f"Only FloatxTensorCoreLayout is supported for floatx, got {_layout}"
+        )
+        original_shape = input_float.shape
+        input_float = _layout.pre_process(input_float)
+        # per axis quantization, where axis = 1
+        block_size = list(input_float.shape)
+        block_size[1] = 1
+
+        ebits, mbits = _layout.ebits, _layout.mbits
+        # Note: these ops are hardcoded to have per axis quantization (axis=1) right now
+        scale = choose_qparams_affine_floatx(input_float, ebits, mbits)
+        floatx_unpacked = quantize_affine_floatx(input_float, scale, ebits, mbits)
+        floatx_packed, scale, _ = _layout.post_process(
+            floatx_unpacked, scale, None, block_size
+        )
+
+        tensor_impl_ctr = get_tensor_impl_constructor(type(_layout))
+        tensor_impl = tensor_impl_ctr(floatx_packed, scale, None, _layout)
+        return cls(tensor_impl, block_size, original_shape, dtype=input_float.dtype)
+
+    @property
+    def _layout(self) -> Layout:
+        return self.tensor_impl._layout
+
+    def to(self, *args, **kwargs):
+        kwargs = self._get_to_kwargs(*args, **kwargs)
+        device = kwargs.pop("device")
+        return self.__class__(
+            self.tensor_impl.to(device),
+            self.block_size,
+            self.shape,
+            self.quant_min,
+            self.quant_max,
+            self.zero_point_domain,
+            **kwargs,
+        )
+
+    def _apply_fn_to_data(self, fn):
+        return self.__class__(
+            fn(self.tensor_impl),
+            self.block_size,
+            self.shape,
+            self.quant_min,
+            self.quant_max,
+            self.zero_point_domain,
+            dtype=self.dtype,
+            strides=self.stride(),
+        )
+
+    # following are the comments for __torch_function__/__torch_dispatch__, -> this is defined in affine_quantized_tensor_ops.py
+    # a bit later
+    # Note: we only added cpu path here for 8da4w, this is for executorch, in the future
+    # 1. we'll add cpu/cuda version (int4mm etc.)
+    # 2. we'll need to hide the 8da4w executorch version under things like layouts (we also have multiple impl for cpu kernel as Michael mentioned), so it will be something like
+    #   cpu device + et laytout --> gives current 8da4w executorch representation
+    #   cpu device + avx layout --> gives optimized kernel for 8da4w in avx cpu etc.
+    #   cuda device + some layout --> gives cuda kernel
+
+    # two scenarios where we currently fall back to vanilla mm:
+    # 1 - when tensor is on CUDA: we'll add this later, we'll also enable dispatching to optimized
+    #     kernels in CPU as well, see the note above
+    # 2 - we're given non-floats - quantizing long to int8 is crazy
+
+
+######################################################
+# Layout and TensorImpl Subclass Registration #
+######################################################
+register_layout = AffineQuantizedTensor.register_layout
+get_tensor_impl_constructor = AffineQuantizedTensor.get_tensor_impl_constructor
+
+
+to_affine_quantized_intx = AffineQuantizedTensor.from_hp_to_intx
+to_affine_quantized_intx_static = AffineQuantizedTensor.from_hp_to_intx_static
+to_affine_quantized_floatx = AffineQuantizedTensor.from_hp_to_floatx
+to_affine_quantized_floatx_static = AffineQuantizedTensor.from_hp_to_floatx_static
+# experimental will be merged in to floatx
+to_affine_quantized_fpx = AffineQuantizedTensor.from_hp_to_fpx
+
+if TORCH_VERSION_AT_LEAST_2_5:
+    # Allow a model with AffineQuantizedTensor weights to be loaded with `weights_only=True`
+    torch.serialization.add_safe_globals([AffineQuantizedTensor])

lib/python3.12/site-packages/torchao/dtypes/affine_quantized_tensor_ops.py

@@ -0,0 +1,580 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+import logging
+
+import torch
+from torch.utils._python_dispatch import return_and_correct_aliasing
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+)
+from torchao.dtypes.floatx.cutlass_semi_sparse_layout import (
+    _linear_fp8_act_fp8_weight_sparse_cutlass_check,
+    _linear_fp8_act_fp8_weight_sparse_cutlass_impl,
+)
+from torchao.dtypes.floatx.float8_layout import (
+    _linear_fp8_act_fp8_weight_check,
+    _linear_fp8_act_fp8_weight_impl,
+    _linear_fp_act_fp8_weight_check,
+    _linear_fp_act_fp8_weight_impl,
+)
+from torchao.dtypes.floatx.floatx_tensor_core_layout import (
+    _linear_f16_bf16_act_floatx_weight_check,
+    _linear_f16_bf16_act_floatx_weight_impl,
+)
+from torchao.dtypes.uintx.block_sparse_layout import (
+    _linear_int8_act_int8_weight_block_sparse_check,
+    _linear_int8_act_int8_weight_block_sparse_impl,
+)
+from torchao.dtypes.uintx.cutlass_int4_packed_layout import (
+    _linear_int4_act_int4_weight_cutlass_check,
+    _linear_int4_act_int4_weight_cutlass_impl,
+    _linear_int8_act_int4_weight_cutlass_check,
+    _linear_int8_act_int4_weight_cutlass_impl,
+)
+from torchao.dtypes.uintx.gemlite_layout import (
+    _linear_fp_act_int4_weight_gemlite_check,
+    _linear_fp_act_int4_weight_gemlite_impl,
+)
+from torchao.dtypes.uintx.int4_cpu_layout import (
+    _linear_fp_act_uint4_weight_cpu_check,
+    _linear_fp_act_uint4_weight_cpu_impl,
+)
+from torchao.dtypes.uintx.int4_xpu_layout import (
+    _linear_bf16_act_uint4_weight_float_zero_check,
+    _linear_bf16_act_uint4_weight_float_zero_impl,
+    _linear_bf16_act_uint4_weight_int8_zero_check,
+    _linear_bf16_act_uint4_weight_int8_zero_impl,
+)
+from torchao.dtypes.uintx.marlin_qqq_tensor import (
+    _linear_int8_act_int4_weight_marlin_qqq_check,
+    _linear_int8_act_int4_weight_marlin_qqq_impl,
+)
+from torchao.dtypes.uintx.marlin_sparse_layout import (
+    _linear_fp_act_int4_weight_sparse_marlin_check,
+    _linear_fp_act_int4_weight_sparse_marlin_impl,
+)
+from torchao.dtypes.uintx.packed_linear_int8_dynamic_activation_intx_weight_layout import (
+    _linear_check as _linear_int8_act_intx_weight_packed_check,
+)
+from torchao.dtypes.uintx.packed_linear_int8_dynamic_activation_intx_weight_layout import (
+    _linear_impl as _linear_int8_act_intx_weight_packed_impl,
+)
+from torchao.dtypes.uintx.plain_layout import (
+    PlainAQTTensorImpl,
+    _linear_fp_act_int8_weight_check,
+    _linear_fp_act_int8_weight_impl,
+    _linear_int8_act_int8_weight_check,
+    _linear_int8_act_int8_weight_impl,
+)
+from torchao.dtypes.uintx.q_dq_layout import (
+    _embedding_check as _embedding_q_dq_check,
+)
+from torchao.dtypes.uintx.q_dq_layout import (
+    _embedding_impl as _embedding_q_dq_impl,
+)
+from torchao.dtypes.uintx.q_dq_layout import (
+    _linear_check as _linear_q_dq_check,
+)
+from torchao.dtypes.uintx.q_dq_layout import (
+    _linear_impl as _linear_q_dq_impl,
+)
+from torchao.dtypes.uintx.semi_sparse_layout import (
+    _linear_int8_act_int8_weight_semi_structured_sparse_check,
+    _linear_int8_act_int8_weight_semi_structured_sparse_impl,
+)
+from torchao.dtypes.uintx.tensor_core_tiled_layout import (
+    _linear_bf16_act_uint4_weight_check,
+    _linear_bf16_act_uint4_weight_impl,
+)
+from torchao.quantization.quant_primitives import dequantize_affine
+from torchao.utils import (
+    fill_defaults,
+)
+
+logger = logging.getLogger(__name__)
+
+
+aten = torch.ops.aten
+
+
+_AQT_QLINEAR_DISPATCH_TABLE = {}
+
+
+def register_aqt_quantized_linear_dispatch(dispatch_condition, impl):
+    """Register a dispatch for quantized linear op with dispatch_condition function and impl function
+    both takes three arguments:
+      input_tensor: dimension is (M1, M2, ..., in_features)
+      weight_tensor: dimension is (out_features, in_features)
+      bias: dimension is (out_features,)
+      so that these can be shared by F.linear, aten.mm, aten.addmm dispatches
+
+    Args:
+        `dispatch_condition` (Callable[[torch.Tensor, torch.Tensor, torch.Tensor], bool]: the dispatch
+            condition for a specialized quantized linear implementation, e.g. bfloat16 activation + uint4 weight
+        `impl` (Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]: the specialized
+            quantized linear implementation
+    """
+    _AQT_QLINEAR_DISPATCH_TABLE[dispatch_condition] = impl
+
+
+def deregister_aqt_quantized_linear_dispatch(dispatch_condition):
+    if dispatch_condition in _AQT_QLINEAR_DISPATCH_TABLE:
+        del _AQT_QLINEAR_DISPATCH_TABLE[dispatch_condition]
+    else:
+        logger.warn(
+            f"Attempting to remove non-existant dispatch condition {dispatch_condition}"
+        )
+
+
+def _same_metadata(self: AffineQuantizedTensor, src: AffineQuantizedTensor):
+    return (
+        isinstance(self, AffineQuantizedTensor)
+        and isinstance(src, AffineQuantizedTensor)
+        and all(
+            [
+                getattr(self, attr) == getattr(src, attr)
+                for attr in [
+                    "block_size",
+                    "shape",
+                    "quant_min",
+                    "quant_max",
+                    "zero_point_domain",
+                    "dtype",
+                ]
+            ]
+        )
+        and isinstance(self.tensor_impl, type(src.tensor_impl))
+    )
+
+
+class QuantizedLinearNotImplementedError(NotImplementedError):
+    """Thin wrapper around NotImplementedError to make it easier to catch this error in the dispatch table"""
+
+    pass
+
+
+# input_tensor: dimension is (M1, M2, ..., in_features)
+# weight_tensor: dimension is (out_features, in_features)
+# bias: dimension is (out_features,)
+@staticmethod
+def _quantized_linear_op(input_tensor, weight_tensor, bias):
+    for dispatch_condition, impl in _AQT_QLINEAR_DISPATCH_TABLE.items():
+        if dispatch_condition(input_tensor, weight_tensor, bias):
+            return impl(input_tensor, weight_tensor, bias)
+    raise QuantizedLinearNotImplementedError(
+        "No specialized dispatch found for quantized linear op"
+    )
+
+
+# Attach the _quantized_linear_op to the AffineQuantizedTensor class
+AffineQuantizedTensor._quantized_linear_op = _quantized_linear_op
+
+
+# _register_aqt_quantized_linear_dispatches function has a list of (dispatch_condition, implementation) functions, defined in their dtype layout classes, that takes the following args:
+# input_tensor: dimension is (M1, M2, ..., in_features)
+# weight_tensor: dimension is (out_features, in_features)
+# bias: dimension is (out_features,)
+# so that these can be shared by F.linear, aten.mm, aten.addmm dispatches
+def _register_aqt_quantized_linear_dispatches():
+    for dispatch_condition, impl in [
+        (_linear_int8_act_int8_weight_check, _linear_int8_act_int8_weight_impl),
+        (
+            _linear_int8_act_int8_weight_semi_structured_sparse_check,
+            _linear_int8_act_int8_weight_semi_structured_sparse_impl,
+        ),
+        (
+            _linear_int8_act_int8_weight_block_sparse_check,
+            _linear_int8_act_int8_weight_block_sparse_impl,
+        ),
+        (_linear_fp8_act_fp8_weight_check, _linear_fp8_act_fp8_weight_impl),
+        (_linear_fp_act_fp8_weight_check, _linear_fp_act_fp8_weight_impl),
+        (_linear_bf16_act_uint4_weight_check, _linear_bf16_act_uint4_weight_impl),
+        (_linear_fp_act_int8_weight_check, _linear_fp_act_int8_weight_impl),
+        (
+            _linear_f16_bf16_act_floatx_weight_check,
+            _linear_f16_bf16_act_floatx_weight_impl,
+        ),
+        (
+            _linear_fp_act_int4_weight_sparse_marlin_check,
+            _linear_fp_act_int4_weight_sparse_marlin_impl,
+        ),
+        (
+            _linear_int8_act_int4_weight_marlin_qqq_check,
+            _linear_int8_act_int4_weight_marlin_qqq_impl,
+        ),
+        (
+            _linear_fp_act_int4_weight_gemlite_check,
+            _linear_fp_act_int4_weight_gemlite_impl,
+        ),
+        (
+            _linear_int8_act_int4_weight_cutlass_check,
+            _linear_int8_act_int4_weight_cutlass_impl,
+        ),
+        (
+            _linear_int4_act_int4_weight_cutlass_check,
+            _linear_int4_act_int4_weight_cutlass_impl,
+        ),
+        (
+            _linear_fp8_act_fp8_weight_sparse_cutlass_check,
+            _linear_fp8_act_fp8_weight_sparse_cutlass_impl,
+        ),
+        (
+            _linear_fp_act_uint4_weight_cpu_check,
+            _linear_fp_act_uint4_weight_cpu_impl,
+        ),
+        (
+            _linear_int8_act_intx_weight_packed_check,
+            _linear_int8_act_intx_weight_packed_impl,
+        ),
+        (
+            _linear_q_dq_check,
+            _linear_q_dq_impl,
+        ),
+        (
+            _linear_bf16_act_uint4_weight_int8_zero_check,
+            _linear_bf16_act_uint4_weight_int8_zero_impl,
+        ),
+        (
+            _linear_bf16_act_uint4_weight_float_zero_check,
+            _linear_bf16_act_uint4_weight_float_zero_impl,
+        ),
+    ]:
+        register_aqt_quantized_linear_dispatch(dispatch_condition, impl)
+
+
+_register_aqt_quantized_linear_dispatches()
+
+implements = AffineQuantizedTensor.implements
+
+
+@implements([torch.nn.functional.linear, aten.linear.default])
+def _(func, types, args, kwargs):
+    input_tensor, weight_tensor, bias = (
+        args[0],
+        args[1],
+        args[2] if len(args) > 2 else None,
+    )
+    if not input_tensor.is_floating_point():
+        raise NotImplementedError(
+            f"{func} is not implemented for non floating point input"
+        )
+
+    # using try/except here so that we can have a general fallback when input_tensor/weight_tensor
+    # is not picked up by any of the dispatch paths in `_quantized_linear_op`, this allows us to
+    # make the branches easier to understand in `_quantized_linear_op`
+    try:
+        return weight_tensor._quantized_linear_op(input_tensor, weight_tensor, bias)
+    except QuantizedLinearNotImplementedError as e:
+        # fallback path is only called when user did not specify a specfic quantized linear implementation with `_layout.quantized_linear_impl`
+        if (
+            isinstance(weight_tensor, AffineQuantizedTensor)
+            and hasattr(weight_tensor._layout, "quantized_linear_impl")
+            and weight_tensor._layout.quantized_linear_impl is not None
+        ):
+            raise e
+
+        if isinstance(input_tensor, AffineQuantizedTensor):
+            input_tensor = input_tensor.dequantize()
+        if isinstance(weight_tensor, AffineQuantizedTensor):
+            weight_tensor = weight_tensor.dequantize()
+        return torch.nn.functional.linear(input_tensor, weight_tensor, bias)
+
+
+@implements(torch.nn.functional.embedding)
+def _(func, types, args, kwargs):
+    if _embedding_q_dq_check(args, kwargs):
+        return _embedding_q_dq_impl(args, kwargs)
+
+    # new_arg1 = args[1].dequantize()
+    # return torch.nn.embedding(args[0], new_arg1, *args[2:], **kwargs)
+    assert isinstance(args[1].tensor_impl, PlainAQTTensorImpl), (
+        f"embedding only works with PlainAQTTensorImpl but got {type(args[1].tensor_impl)}"
+    )
+    assert (
+        kwargs["padding_idx"] is None
+        and kwargs["max_norm"] is None
+        and not kwargs["scale_grad_by_freq"]
+        and not kwargs["sparse"]
+        and kwargs["norm_type"] == 2.0
+    )
+    idx = args[0]
+    int_data, scale, zero_point = args[1].tensor_impl.get_plain()
+
+    sliced_data, sliced_scale, sliced_zero_point = (
+        int_data[idx],
+        scale[idx],
+        zero_point[idx],
+    )
+    # Block size is expecting 2 dimensions [1, group size] but
+    # batchsize or other dims gets added to sliced_data, sliced_scale and sliced_zero_point so
+    # we need to increase block size to correct dim
+    new_blocks = idx.dim() - 1
+    return dequantize_affine(
+        sliced_data,
+        new_blocks * [1] + list(args[1].block_size),
+        sliced_scale,
+        sliced_zero_point,
+        sliced_data.dtype,
+        args[1].quant_min,
+        args[1].quant_max,
+        args[1].zero_point_domain,
+        output_dtype=sliced_scale.dtype,
+    )
+
+
+@implements(aten.addmm.default)
+def _(func, types, args, kwargs):
+    input_tensor, weight_tensor, bias = (
+        args[1],
+        args[2],
+        args[0],
+    )
+    if not input_tensor.is_floating_point():
+        raise NotImplementedError(
+            f"{func} is not implemented for non floating point input"
+        )
+
+    assert input_tensor.shape[-1] == weight_tensor.shape[0], (
+        f"need mat1 shape: {input_tensor.shape} final dim"
+        f"to match mat2 shape: {weight_tensor.shape} first dim"
+    )
+
+    # using try/except here so that we can have a general fallback when input_tensor/weight_tensor
+    # is not picked up by any of the dispatch paths in `_quantized_linear_op`, this allows us to
+    # make the branches easier to understand in `_quantized_linear_op`
+    try:
+        transposed_weight_tensor = weight_tensor.t()
+        return weight_tensor._quantized_linear_op(
+            input_tensor, transposed_weight_tensor, bias
+        )
+    except QuantizedLinearNotImplementedError as e:
+        # fallback path is only called when user did not specify a specfic quantized linear implementation with `_layout.quantized_linear_impl`
+        if (
+            isinstance(weight_tensor, AffineQuantizedTensor)
+            and hasattr(weight_tensor._layout, "quantized_linear_impl")
+            and weight_tensor._layout.quantized_linear_impl is not None
+        ):
+            raise e
+
+        if isinstance(input_tensor, AffineQuantizedTensor):
+            input_tensor = input_tensor.dequantize()
+        if isinstance(weight_tensor, AffineQuantizedTensor):
+            weight_tensor = weight_tensor.dequantize()
+        return func(bias, input_tensor, weight_tensor)
+
+
+@implements(aten.mm.default)
+def _(func, types, args, kwargs):
+    input_tensor, weight_tensor, bias = (args[0], args[1], None)
+    if not input_tensor.is_floating_point():
+        raise NotImplementedError(
+            f"{func} is not implemented for non floating point input"
+        )
+
+    assert input_tensor.shape[-1] == weight_tensor.shape[0], (
+        f"need mat1 shape: {input_tensor.shape} final dim"
+        f"to match mat2 shape: {weight_tensor.shape} first dim"
+    )
+
+    try:
+        transposed_weight_tensor = weight_tensor.t()
+        return weight_tensor._quantized_linear_op(
+            input_tensor, transposed_weight_tensor, bias
+        )
+    except QuantizedLinearNotImplementedError as e:
+        # fallback path is only called when user did not specify a specfic quantized linear implementation with `_layout.quantized_linear_impl`
+        if (
+            isinstance(weight_tensor, AffineQuantizedTensor)
+            and hasattr(weight_tensor._layout, "quantized_linear_impl")
+            and weight_tensor._layout.quantized_linear_impl is not None
+        ):
+            raise e
+
+        if isinstance(input_tensor, AffineQuantizedTensor):
+            input_tensor = input_tensor.dequantize()
+        if isinstance(weight_tensor, AffineQuantizedTensor):
+            weight_tensor = weight_tensor.dequantize()
+        return func(input_tensor, weight_tensor)
+
+
+@implements([aten.detach.default, aten.alias.default])
+def _(func, types, args, kwargs):
+    return return_and_correct_aliasing(
+        func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+    )
+
+
+@implements(aten.clone.default)
+def _(func, types, args, kwargs):
+    return return_and_correct_aliasing(
+        func, args, kwargs, args[0]._apply_fn_to_data(torch.clone)
+    )
+
+
+@implements(aten.copy_.default)
+def _(func, types, args, kwargs):
+    return return_and_correct_aliasing(
+        func, args, kwargs, args[1]._apply_fn_to_data(torch.clone)
+    )
+
+
+@implements(aten._to_copy.default)
+def _(func, types, args, kwargs):
+    return return_and_correct_aliasing(
+        func,
+        args,
+        kwargs,
+        args[0].to(*args[1:], **kwargs)._apply_fn_to_data(torch.clone),
+    )
+
+
+@implements(aten.copy_.default)
+def _(func, types, args, kwargs):
+    self = args[0]
+    src = args[1]
+    if _same_metadata(self, src):
+        self_tensors = self.__tensor_flatten__()[0]
+        for tensor_name in self_tensors:
+            getattr(self, tensor_name).copy_(getattr(src, tensor_name))
+        return
+    raise ValueError(
+        f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}"
+    )
+
+
+@implements(aten.t.default)
+def _(func, types, args, kwargs):
+    block_size = args[0].block_size
+    assert len(block_size) == 2
+    transposed_block_size = (block_size[1], block_size[0])
+    tensor = args[0]
+    shape = tensor.shape[::-1]
+    new = tensor.__class__(
+        tensor.tensor_impl.t(),
+        transposed_block_size,
+        shape,
+        tensor.quant_min,
+        tensor.quant_max,
+        tensor.zero_point_domain,
+        dtype=tensor.dtype,
+        strides=tensor.stride(),
+    )
+    return return_and_correct_aliasing(func, args, kwargs, new)
+
+
+@implements(aten.slice.Tensor)
+def _(func, types, args, kwargs):
+    self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1])
+    assert step == 1
+    assert dim == 0 or dim == 1, f"Only dim==0 or 1 are supported, got: {dim}"
+    if end >= self.shape[dim]:
+        end = self.shape[dim]
+    shape = list(self.shape)
+    shape[dim] = end - start
+    block_size = self.block_size
+    assert len(block_size) in [
+        2,
+        3,
+    ], f"Slice only works for 2 and 3d block_size right now, got: {block_size}"
+    # with slice, some shape dimension might be smaller than block_size dimension, so
+    # we need to make sure there is no overflow
+    if len(block_size) == 2:
+        block_size = (min(shape[0], block_size[0]), min(shape[1], block_size[1]))
+
+    new = self.__class__(
+        aten.slice.Tensor(self.tensor_impl, dim, start, end, step),
+        block_size,
+        shape,
+        self.quant_min,
+        self.quant_max,
+        self.zero_point_domain,
+        dtype=self.dtype,
+        strides=self.stride() if len(block_size) == 2 else None,
+    )
+    return return_and_correct_aliasing(func, args, kwargs, new)
+
+
+@implements(aten.index.Tensor)
+def _(func, types, args, kwargs):
+    self, indices = args
+    assert len(indices) == 1, (
+        f"op {func} currently only implemented for single dimensional indexing but got indices: {indices}"
+    )
+    new_tensor_impl = aten.index.Tensor(self.tensor_impl, indices)
+    shape = tuple([indices[0].numel(), *self.shape[1:]])
+
+    block_size = self.block_size
+    new = self.__class__(
+        new_tensor_impl,
+        block_size,
+        shape,
+        self.quant_min,
+        self.quant_max,
+        self.zero_point_domain,
+        dtype=self.dtype,
+    )
+    return return_and_correct_aliasing(func, args, kwargs, new)
+
+
+@implements(aten.select.int)
+def _(func, types, args, kwargs):
+    self, dim, index = fill_defaults(args, 3, [0, 0])
+    assert dim == 0, f"op {func} currently only implemented for dim=0 but got dim={dim}"
+    assert self.dim() == 3, (
+        f"op {func} currently only implemented for 3 dimensional tensors but got shape={self.shape}"
+    )
+
+    new_tensor_impl = aten.select.int(self.tensor_impl, dim, index)
+
+    shape = self.shape[1:]
+    block_size = self.block_size[1:]
+    new = self.__class__(
+        new_tensor_impl,
+        block_size,
+        shape,
+        self.quant_min,
+        self.quant_max,
+        self.zero_point_domain,
+        dtype=self.dtype,
+    )
+    return return_and_correct_aliasing(func, args, kwargs, new)
+
+
+# this is needed for DTensor.from_local() and for flattening tensor
+@implements(aten.view.default)
+def _(func, types, args, kwargs):
+    self, shape = args
+
+    if tuple(self.shape) == tuple(shape):
+        return self.__class__(
+            self.tensor_impl,
+            self.block_size,
+            self.shape,
+            self.quant_min,
+            self.quant_max,
+            self.zero_point_domain,
+            dtype=self.dtype,
+            strides=self.stride(),
+        )
+
+    if len(shape) == 1 and shape[0] == -1:
+        assert len(self.block_size) == 2 and self.block_size[0] == 1
+        block_size = (self.block_size[1],)
+        return self.__class__(
+            self.tensor_impl,
+            block_size,
+            (self.numel(),),
+            self.quant_min,
+            self.quant_max,
+            self.zero_point_domain,
+            dtype=self.dtype,
+            strides=self.stride(),
+        )
+
+    raise ValueError(
+        f"{self.__class__.__name__} only supports .view() with same shape or shape=[-1]"
+    )

lib/python3.12/site-packages/torchao/dtypes/floatx/__init__.py

@@ -0,0 +1,17 @@
+from .cutlass_semi_sparse_layout import (
+    CutlassSemiSparseLayout,
+)
+from .float8_layout import Float8Layout
+from .floatx_tensor_core_layout import (
+    FloatxTensorCoreLayout,
+    from_scaled_tc_floatx,
+    to_scaled_tc_floatx,
+)
+
+__all__ = [
+    "FloatxTensorCoreLayout",
+    "to_scaled_tc_floatx",
+    "from_scaled_tc_floatx",
+    "Float8Layout",
+    "CutlassSemiSparseLayout",
+]

lib/python3.12/site-packages/torchao/dtypes/floatx/cutlass_semi_sparse_layout.py

@@ -0,0 +1,212 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch.utils._python_dispatch import (
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+from torchao.dtypes.utils import AQTTensorImpl, Layout
+from torchao.ops import (
+    rowwise_scaled_linear_sparse_cutlass_f8f8,
+    to_sparse_semi_structured_cutlass_sm9x_f8,
+)
+
+aten = torch.ops.aten
+
+
+def _same_metadata(
+    self: "CutlassSemiSparseTensorImpl", src: "CutlassSemiSparseTensorImpl"
+) -> bool:
+    return (
+        isinstance(self, CutlassSemiSparseTensorImpl)
+        and isinstance(src, CutlassSemiSparseTensorImpl)
+        and self.shape == src.shape
+        and self.sparse.shape == src.sparse.shape
+        and self.meta.shape == src.meta.shape
+        and self.scale.shape == src.scale.shape
+        and type(self._layout) == type(src._layout)
+    )
+
+
+@dataclass(frozen=True)
+class CutlassSemiSparseLayout(Layout):
+    """Layout class for float8 2:4 sparsity layout for affine quantized tensor, for cutlass kernel."""
+
+    def pre_process(self, dense: torch.Tensor) -> torch.Tensor:
+        # prune to 2:4 if not already
+        from torchao.sparsity.utils import mask_creator
+
+        return dense * mask_creator(dense).bool()
+
+
+@register_layout(CutlassSemiSparseLayout)
+class CutlassSemiSparseTensorImpl(AQTTensorImpl):
+    @staticmethod
+    def __new__(
+        cls,
+        sparse: torch.Tensor,
+        meta: torch.Tensor,
+        scale: torch.Tensor,
+        _layout: Layout,
+    ):
+        kwargs = {}
+        kwargs["device"] = sparse.device
+        kwargs["layout"] = (
+            kwargs.get("layout") if kwargs.get("layout", False) else sparse.layout
+        )
+        kwargs["dtype"] = sparse.dtype
+        kwargs["requires_grad"] = False
+        shape = (sparse.shape[0], 2 * sparse.shape[-1])
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        sparse: torch.Tensor,
+        meta: torch.Tensor,
+        scale: torch.Tensor,
+        _layout: Layout,
+    ):
+        self.sparse = sparse
+        self.meta = meta
+        self.scale = scale
+        self._layout = _layout
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+        elif func is aten.copy_.default:
+            self = args[0]
+            src = args[1]
+            if _same_metadata(self, src):
+                self_tensors = self.__tensor_flatten__()[0]
+                for tensor_name in self_tensors:
+                    getattr(self, tensor_name).copy_(getattr(src, tensor_name))
+                return
+            raise ValueError(
+                f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}"
+            )
+
+        raise NotImplementedError(
+            f"CutlassSemiSparseTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    def __tensor_flatten__(self):
+        return ["sparse", "meta", "scale"], [self._layout]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        sparse = tensor_data_dict["sparse"]
+        meta = tensor_data_dict["meta"]
+        scale = tensor_data_dict["scale"]
+        (_layout,) = tensor_attributes
+        return cls(sparse, meta, scale, _layout)
+
+    def get_plain(self):
+        # No support in CUTLASS to convert back to dense from sparse
+        # semi-structured format, so multiplying with identity matrix,
+        # and using identity scale factors, for the conversion.
+        cols = self.shape[1]
+        input = torch.eye(cols, dtype=self.sparse.dtype, device=self.sparse.device)
+        input_scale = torch.ones(
+            (cols,), dtype=self.scale.dtype, device=self.sparse.device
+        )
+        sparse_scale = torch.ones_like(self.scale)
+        out_dtype = torch.bfloat16
+        dense = (
+            rowwise_scaled_linear_sparse_cutlass_f8f8(
+                input,
+                input_scale,
+                self.sparse,
+                self.meta,
+                sparse_scale,
+                out_dtype=out_dtype,
+            )
+            .to(self.dtype)
+            .t()
+            .contiguous()
+        )
+
+        return dense, self.scale, None
+
+    @classmethod
+    def from_plain(
+        cls,
+        dense: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        assert zero_point is None or torch.all(zero_point == 0)
+
+        sparse, meta = to_sparse_semi_structured_cutlass_sm9x_f8(dense)
+
+        return cls(
+            sparse,
+            meta,
+            scale,
+            _layout,
+        )
+
+    def get_layout(self) -> Layout:
+        return self._layout
+
+    def _apply_fn_to_data(self, fn):
+        self.sparse = fn(self.sparse)
+        self.meta = fn(self.meta)
+        self.scale = fn(self.scale)
+        return self
+
+
+def _linear_fp8_act_fp8_weight_sparse_cutlass_check(input_tensor, weight_tensor, bias):
+    from torchao.dtypes.floatx import Float8Layout
+
+    return (
+        isinstance(input_tensor, AffineQuantizedTensor)
+        and isinstance(input_tensor._layout, Float8Layout)
+        and input_tensor.dtype in (torch.float16, torch.bfloat16)
+        and len(input_tensor.shape) >= 2
+        and input_tensor.tensor_impl.scale.dtype == torch.float32
+        and len(input_tensor.tensor_impl.scale.shape) == len(input_tensor.shape) - 1
+        and isinstance(weight_tensor, AffineQuantizedTensor)
+        and isinstance(weight_tensor._layout, CutlassSemiSparseLayout)
+        and weight_tensor.dtype == input_tensor.dtype
+        and len(weight_tensor.shape) == 2
+        and weight_tensor.tensor_impl.scale.dtype == torch.float32
+        and len(weight_tensor.tensor_impl.scale.shape) == 1
+        and (bias is None or bias.dtype == input_tensor.dtype)
+        and (bias is None or len(bias.shape) == 1)
+    )
+
+
+def _linear_fp8_act_fp8_weight_sparse_cutlass_impl(input_tensor, weight_tensor, bias):
+    from torchao.ops import rowwise_scaled_linear_sparse_cutlass_f8f8
+
+    input = input_tensor.tensor_impl.float8_data
+    input_scale = input_tensor.tensor_impl.scale
+    weight = weight_tensor.tensor_impl.sparse
+    weight_meta = weight_tensor.tensor_impl.meta
+    weight_scale = weight_tensor.tensor_impl.scale
+    out_dtype = input_tensor.dtype
+
+    out = rowwise_scaled_linear_sparse_cutlass_f8f8(
+        input, input_scale, weight, weight_meta, weight_scale, bias, out_dtype
+    )
+
+    return out

lib/python3.12/site-packages/torchao/dtypes/floatx/float8_layout.py

@@ -0,0 +1,386 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+from torch.utils._python_dispatch import (
+    is_traceable_wrapper_subclass,
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+from torchao.dtypes.utils import AQTTensorImpl, Layout, get_out_shape
+from torchao.float8.inference import (
+    Float8MMConfig,
+    _is_rowwise_scaled,
+    addmm_float8_unwrapped_inference,
+    preprocess_data,
+)
+from torchao.utils import _is_float8_type, fill_defaults
+
+aten = torch.ops.aten
+
+
+def _same_metadata(self: "Float8AQTTensorImpl", src: "Float8AQTTensorImpl") -> bool:
+    # Special handling for transposed attribute
+    transposed_match = (self.transposed == src.transposed) or (
+        self.transposed is False and src.transposed is None
+    )
+
+    return (
+        isinstance(self, Float8AQTTensorImpl)
+        and isinstance(src, Float8AQTTensorImpl)
+        and self.shape == src.shape
+        and self.float8_data.shape == src.float8_data.shape
+        and self.scale.shape == src.scale.shape
+        and transposed_match
+        and type(self._layout) == type(src._layout)
+    )
+
+
+@dataclass(frozen=True)
+class Float8Layout(Layout):
+    """Represents the layout configuration for Float8 affine quantized tensors.
+
+    Attributes:
+        mm_config (Optional[Float8MMConfig]): Configuration for matrix multiplication operations involving Float8 tensors. If None, default settings are used.
+    """
+
+    mm_config: Optional[Float8MMConfig] = None
+
+
+_fallback_warning_shown = False
+
+
+@register_layout(Float8Layout)
+class Float8AQTTensorImpl(AQTTensorImpl):
+    """
+    TensorImpl for float8 layout affine quantized tensor
+
+    Note: technically we should not create a new layout for float8 we should merge this into
+    plain layout
+    """
+
+    float8_data: torch.Tensor
+    scale: torch.Tensor
+    transposed: bool
+
+    def __new__(
+        cls,
+        float8_data: torch.Tensor,
+        scale: torch.Tensor,
+        transposed: bool,
+        _layout: Layout,
+    ):
+        kwargs = {}
+        kwargs["device"] = float8_data.device
+        kwargs["layout"] = (
+            kwargs.get("layout") if kwargs.get("layout", False) else float8_data.layout
+        )
+        kwargs["dtype"] = float8_data.dtype
+        kwargs["requires_grad"] = False
+        shape = float8_data.shape
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        float8_data: torch.Tensor,
+        scale: torch.Tensor,
+        transposed: bool,
+        _layout: Layout,
+    ):
+        self.float8_data = float8_data
+        self.scale = scale
+        self.transposed = transposed
+        self._layout = _layout
+
+    def _apply_fn_to_data(self, fn):
+        """Applys a fn to all tensor components stored on this class"""
+        global _fallback_warning_shown
+
+        try:
+            return self.__class__(
+                fn(self.float8_data),
+                fn(self.scale),
+                self.transposed,
+                self._layout,
+            )
+        except RuntimeError as e:
+            if '"index_cuda" not implemented for ' in str(e):
+                if not _fallback_warning_shown:
+                    import warnings
+
+                    warnings.warn(
+                        f"When trying to index Float8AQTTensorImpl, got known error {e}, will use slower fallback but "
+                        + "note: You can torch.compile the model to avoid this problem.",
+                        UserWarning,
+                    )
+                    _fallback_warning_shown = True
+
+                return self.__class__(  # do indexing in bfloat16 then convert back
+                    fn(self.float8_data.to(torch.bfloat16)).to(self.float8_data.dtype),
+                    fn(self.scale),
+                    self.transposed,
+                    self._layout,
+                )
+            else:
+                raise e
+
+    def to(self, *args, **kwargs):
+        kwargs = self._get_to_kwargs(*args, **kwargs)
+        return self.__class__(
+            self.float8_data.to(kwargs["device"]),
+            self.scale.to(kwargs["device"]),
+            self.transposed,
+            self._layout,
+        )
+
+    def __tensor_flatten__(self):
+        return ["float8_data", "scale"], [self.transposed, self._layout]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        float8_data, scale = tensor_data_dict["float8_data"], tensor_data_dict["scale"]
+        (
+            transposed,
+            _layout,
+        ) = tensor_attributes
+        return cls(float8_data, scale, transposed, _layout)
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+        elif func is aten.clone.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.clone)
+            )
+        elif func is aten.t.default:
+            """we don't need to repack the weight and just rely on external
+            shape being changed and record the status of transpose/no-transpose
+            """
+            args[0].transposed = not args[0].transposed
+            return return_and_correct_aliasing(func, args, kwargs, args[0])
+        elif func is aten.copy_.default:
+            self = args[0]
+            src = args[1]
+            if _same_metadata(self, src):
+                self_tensors = self.__tensor_flatten__()[0]
+                for tensor_name in self_tensors:
+                    getattr(self, tensor_name).copy_(getattr(src, tensor_name))
+                return
+            raise ValueError(
+                f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}"
+            )
+        elif func in [aten.select.int, aten.index.Tensor]:
+            return return_and_correct_aliasing(
+                func,
+                args,
+                kwargs,
+                args[0]._apply_fn_to_data(lambda x: func(x, *args[1:], **kwargs)),
+            )
+        elif func is aten.slice.Tensor:
+            self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1])
+            if dim == 0:
+                # TODO: scale replecation should be dependent on block size
+                if self.scale.ndim == 1:
+                    return return_and_correct_aliasing(
+                        func,
+                        args,
+                        kwargs,
+                        args[0]._apply_fn_to_data(
+                            lambda x: aten.slice.Tensor(x, dim, start, end, step)
+                        ),
+                    )
+                elif self.scale.ndim == 0:
+                    return return_and_correct_aliasing(
+                        func,
+                        args,
+                        kwargs,
+                        Float8AQTTensorImpl(
+                            aten.slice.Tensor(self.float8_data, dim, start, end, step),
+                            self.scale,
+                            None,
+                            self._layout,
+                        ),
+                    )
+                else:
+                    raise NotImplementedError(
+                        f"Float8AQTTensorImpl dispatch: attempting to run {func}, with scale ndim={dim}, that is not supported"
+                    )
+            elif dim == 1:
+                return return_and_correct_aliasing(
+                    func,
+                    args,
+                    kwargs,
+                    Float8AQTTensorImpl(
+                        aten.slice.Tensor(
+                            self.float8_data, dim, start, end, step
+                        ).contiguous(),
+                        self.scale,
+                        None,
+                        self._layout,
+                    ),
+                )
+            else:
+                raise NotImplementedError(
+                    f"Float8AQTTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported"
+                )
+        else:
+            raise NotImplementedError(
+                f"Float8AQTTensorImpl dispatch: attempting to run {func}, this is not supported"
+            )
+
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
+    def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        return self.float8_data, self.scale, None
+
+    def get_layout(self) -> Layout:
+        return self._layout
+
+    @classmethod
+    def from_plain(
+        cls,
+        data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        """Main entrypoint for constructing Float8TensorImpl"""
+        assert _is_float8_type(data.dtype), (
+            f"Float8 TensorImpl must be constructed from float8 dtype but got {data.dtype}"
+        )
+        assert isinstance(_layout, Float8Layout), (
+            f"Float8 TensorImpl must be constructed from Float8Layout but got {_layout}"
+        )
+        return cls(data, scale, False, _layout)
+
+    def __repr__(self):
+        float8_data, scale, _ = self.get_plain()
+        _layout = self.get_layout()
+        return (
+            f"{self.__class__.__name__}(\n"
+            f"float8_data={float8_data},\n"
+            f"scale={scale},\n"
+            f"transposed={self.transposed}, "
+            f"_layout={_layout})"
+        )
+
+
+##########################
+# Float8 Dispatch Kernels
+##########################
+
+
+def _linear_fp8_act_fp8_weight_check(
+    input_tensor: Union[torch.Tensor, "AffineQuantizedTensor"],
+    weight_tensor: Union[torch.Tensor, "AffineQuantizedTensor"],
+    bias: Optional[torch.Tensor],
+) -> bool:
+    def check_aqt(aqt: Union[torch.Tensor, AffineQuantizedTensor]) -> bool:
+        return (
+            isinstance(aqt, AffineQuantizedTensor)
+            and isinstance(aqt._layout, Float8Layout)
+            and aqt.tensor_impl.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]
+            and (aqt.shape == aqt.block_size or _is_rowwise_scaled(aqt))
+        )
+
+    return check_aqt(input_tensor) and check_aqt(weight_tensor)
+
+
+def preprocess_scale(input_scale: torch.Tensor, input_shape: Tuple[int]):
+    """Ensures input tensor is correctly formated for _scaled_mm"""
+    input_scale = input_scale.unsqueeze(-1)
+
+    if input_scale.dim() > 2:
+        input_scale = input_scale.reshape(-1, input_scale.shape[-1])
+
+    return input_scale
+
+
+def _linear_fp8_act_fp8_weight_impl(
+    input_tensor: "AffineQuantizedTensor",
+    weight_tensor: "AffineQuantizedTensor",
+    bias: Optional[torch.Tensor],
+):
+    """Implements matmul between FP8 input and FP8 weight with compute using _scaled_mm"""
+    scaled_mm_config = weight_tensor._layout.mm_config
+    assert scaled_mm_config is not None
+    out_shape = get_out_shape(input_tensor.shape, weight_tensor.shape)
+
+    # Weight tensor preprocessing
+    w_tensor_impl = weight_tensor.tensor_impl
+    assert not w_tensor_impl.transposed, "Weight tensor must be contiguous"
+    w_data = w_tensor_impl.float8_data
+    w_scale = w_tensor_impl.scale
+
+    # Input tensor preprocessing
+    inpt_data = input_tensor.tensor_impl.float8_data
+    input_scale = input_tensor.tensor_impl.scale
+    # Handle case where input tensor is more than 2D
+    inpt_data = inpt_data.reshape(-1, inpt_data.shape[-1])
+
+    # Handle rowwise case
+    if _is_rowwise_scaled(weight_tensor):
+        assert _is_rowwise_scaled(input_tensor), (
+            "Input tensor must be rowwise block size"
+        )
+        w_scale = w_scale.unsqueeze(-1).T
+        input_scale = preprocess_scale(input_scale, input_tensor.shape)
+
+    # Preprocess data
+    inpt_data, w_data = preprocess_data(inpt_data, w_data.T, scaled_mm_config)
+
+    # Perform the computation
+    return addmm_float8_unwrapped_inference(
+        inpt_data,
+        input_scale,
+        w_data,
+        w_scale,
+        output_dtype=input_tensor.dtype,
+        bias=bias,
+        use_fast_accum=scaled_mm_config.use_fast_accum,
+    ).reshape(out_shape)
+
+
+def _linear_fp_act_fp8_weight_check(
+    input_tensor: Union[torch.Tensor, "AffineQuantizedTensor"],
+    weight_tensor: Union[torch.Tensor, "AffineQuantizedTensor"],
+    bias: Optional[torch.Tensor],
+) -> bool:
+    return (
+        # input is native float tensor
+        not is_traceable_wrapper_subclass(input_tensor)
+        and input_tensor.is_floating_point()
+        and
+        # weight is float8 quantized affine quantized tensor
+        isinstance(weight_tensor, AffineQuantizedTensor)
+        and isinstance(weight_tensor._layout, Float8Layout)
+        and weight_tensor.tensor_impl.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]
+        and (
+            weight_tensor.shape == weight_tensor.block_size
+            or _is_rowwise_scaled(weight_tensor)
+        )
+    )
+
+
+def _linear_fp_act_fp8_weight_impl(
+    input_tensor: torch.Tensor,
+    weight_tensor: "AffineQuantizedTensor",
+    bias: Optional[torch.Tensor],
+):
+    return torch.nn.functional.linear(input_tensor, weight_tensor.dequantize(), bias)

lib/python3.12/site-packages/torchao/dtypes/floatx/floatx_tensor_core_layout.py

@@ -0,0 +1,666 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from functools import reduce
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+from torch.utils._python_dispatch import (
+    is_traceable_wrapper_subclass,
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+from torchao.dtypes.utils import (
+    AQTTensorImpl,
+    Layout,
+)
+from torchao.prototype.custom_fp_utils import (
+    _f32_to_floatx_unpacked,
+    _floatx_unpacked_to_f32,
+    _n_ones,
+)
+
+aten = torch.ops.aten
+_ONES_TABLE = [_n_ones(i) for i in range(8)]
+
+
+def _pack(x: Tensor, n_bits: int) -> Tensor:
+    return reduce(
+        torch.bitwise_or,
+        [
+            x[..., i :: (8 // n_bits)] << (8 - (i + 1) * n_bits)
+            for i in range(8 // n_bits)
+        ],
+    )
+
+
+def _unpack(x: Tensor, n_bits: int) -> Tensor:
+    return torch.stack(
+        [
+            (x >> (8 - (i + 1) * n_bits)) & ((1 << n_bits) - 1)
+            for i in range(8 // n_bits)
+        ],
+        dim=-1,
+    ).flatten(-2)
+
+
+# https://github.com/usyd-fsalab/fp6_llm/blob/5df6737cca32f604e957e3f63f03ccc2e4d1df0d/fp6_llm/csrc/utils/weight_prepacking.h#L87-L116
+def _bit_interleave(x: Tensor, n_bits: int, undo: bool = False) -> Tensor:
+    # the original code unpacks/packs the values from/to uint32 while we unpack/pack the values from/to uint8
+    # thus, we need to reverse byte order within a uint32 word.
+    x = x.reshape(-1, 4).flip(1)
+
+    x = _unpack(x, n_bits)
+    x = x.view(-1, 4 * (8 // n_bits))
+
+    if not undo:
+        bit_order = {
+            1: [
+                1,
+                5,
+                9,
+                13,
+                17,
+                21,
+                25,
+                29,
+                3,
+                7,
+                11,
+                15,
+                19,
+                23,
+                27,
+                31,
+                0,
+                4,
+                8,
+                12,
+                16,
+                20,
+                24,
+                28,
+                2,
+                6,
+                10,
+                14,
+                18,
+                22,
+                26,
+                30,
+            ],
+            2: [1, 5, 9, 13, 3, 7, 11, 15, 0, 4, 8, 12, 2, 6, 10, 14],
+            4: [1, 5, 3, 7, 0, 4, 2, 6],
+        }[n_bits]
+
+    else:
+        # this is inverse of the above, obtained by running
+        # [v.index(i) for i in range(len(v))]
+        bit_order = {
+            1: [
+                16,
+                0,
+                24,
+                8,
+                17,
+                1,
+                25,
+                9,
+                18,
+                2,
+                26,
+                10,
+                19,
+                3,
+                27,
+                11,
+                20,
+                4,
+                28,
+                12,
+                21,
+                5,
+                29,
+                13,
+                22,
+                6,
+                30,
+                14,
+                23,
+                7,
+                31,
+                15,
+            ],
+            2: [8, 0, 12, 4, 9, 1, 13, 5, 10, 2, 14, 6, 11, 3, 15, 7],
+            4: [4, 0, 6, 2, 5, 1, 7, 3],
+        }[n_bits]
+
+    x = x[:, bit_order]
+    x = _pack(x, n_bits)
+
+    # reverse byte order within a uint32 word again.
+    x = x.reshape(-1, 4).flip(1)
+    return x.flatten()
+
+
+# this is a literal adaptation of FP6-LLM ahead-of-time bit-level pre-packing
+# https://github.com/usyd-fsalab/fp6_llm/blob/5df6737cca32f604e957e3f63f03ccc2e4d1df0d/fp6_llm/csrc/utils/weight_prepacking.h
+def _pack_tc_floatx(tensor: Tensor, nbits: int) -> Tensor:
+    assert tensor.ndim == 2, tensor.dtype == torch.uint8
+    M, N = tensor.shape
+    assert (M % 64 == 0) and (N % 64 == 0)
+
+    # Pass 1 from original code
+    tensor = tensor.view(M // 64, 4, 2, 8, N // 16, 2, 8)
+    tensor = tensor.permute(0, 4, 1, 5, 2, 3, 6)
+    tensor = tensor.reshape(-1, 32, 2)
+    tensor = tensor.permute(1, 0, 2)
+    tensor = tensor.flatten()
+
+    used_bits = 0
+    fragments = []
+
+    for y in [1, 2, 4]:
+        if nbits & y:
+            mask = (1 << y) - 1
+            tensor_ybit = (tensor >> (nbits - used_bits - y)) & mask
+            tensor_ybit = _pack(tensor_ybit, y)
+
+            tensor_ybit = (
+                tensor_ybit.view(32, -1, 4).permute(1, 0, 2).flip(2)
+            )  # Pass 2 from original code
+            tensor_ybit = _bit_interleave(
+                tensor_ybit.flatten(), y
+            )  # Pass 3 from original code
+            fragments.append(tensor_ybit)
+            used_bits += y
+
+    return torch.cat(fragments, dim=0).view(M, -1)
+
+
+# more optimized version of _pack_tc_floatx() for FP6 by merging ops
+def _pack_tc_fp6(tensor: Tensor) -> Tensor:
+    assert tensor.ndim == 2, tensor.dtype == torch.uint8
+    M, N = tensor.shape
+    assert (M % 64 == 0) and (N % 64 == 0)
+
+    tensor = tensor.view(M // 64, 2, 2, 2, 8, N // 16, 2, 8)
+    tensor = tensor.flip(3)
+
+    tensor_2bit = (tensor >> 4) & 0b11
+    tensor_2bit = tensor_2bit.permute(0, 5, 1, 4, 7, 3, 2, 6)
+    tensor_2bit = _pack(tensor_2bit.flatten(), 2)
+
+    tensor_4bit = tensor & 0b1111
+    tensor_4bit = tensor_4bit.permute(0, 5, 1, 2, 4, 7, 3, 6)
+    tensor_4bit = _pack(tensor_4bit.flatten(), 4)
+
+    return torch.cat([tensor_2bit, tensor_4bit], dim=0).view(M, -1)
+
+
+# currently only optimize for TC-FP6 packing
+def pack_tc_floatx(tensor: Tensor, nbits: int) -> Tensor:
+    if nbits == 6:
+        return _pack_tc_fp6(tensor)
+    return _pack_tc_floatx(tensor, nbits)
+
+
+def to_scaled_tc_floatx(
+    tensor: Tensor, ebits: int, mbits: int
+) -> Tuple[Tensor, Tensor]:
+    # _n_ones() is not compatible with torch.compile() due to << operator
+    # https://github.com/pytorch/pytorch/issues/119152
+    # exp_bias = _n_ones(ebits - 1)
+    # max_normal = 2 ** (_n_ones(ebits) - exp_bias) * (_n_ones(mbits + 1) / (2 ** mbits))
+
+    # workaround: global lookup table
+    exp_bias = _ONES_TABLE[ebits - 1]
+    max_normal = 2 ** (_ONES_TABLE[ebits] - exp_bias) * (
+        _ONES_TABLE[mbits + 1] / (2**mbits)
+    )
+
+    dtype = tensor.dtype
+    tensor = tensor.float()
+    scale = tensor.abs().amax(1).clamp(min=1e-12) / max_normal
+    tensor_floatx = _f32_to_floatx_unpacked(tensor / scale.view(-1, 1), ebits, mbits)
+    tensor_tc_floatx = pack_tc_floatx(tensor_floatx, 1 + ebits + mbits)
+    return tensor_tc_floatx, scale.to(dtype)
+
+
+# inverse of _pack_tc_floatx()
+def _unpack_tc_floatx(tensor: Tensor, nbits: int) -> Tensor:
+    assert tensor.ndim == 2 and tensor.dtype == torch.uint8
+    M = tensor.shape[0]
+    size = tensor.numel()
+    tensor = tensor.flatten()
+    offset = 0
+    used_bits = 0
+
+    tensor_floatx = None
+
+    for y in [1, 2, 4]:
+        if nbits & y:
+            size_ybit = size // nbits * y
+            tensor_ybit = tensor[offset : offset + size_ybit]
+            offset += size_ybit
+
+            tensor_ybit = _bit_interleave(tensor_ybit, y, undo=True)  # undo Pass 3
+            tensor_ybit = (
+                tensor_ybit.view(-1, 32, 4).flip(2).permute(1, 0, 2)
+            )  # undo Pass 2
+
+            tensor_ybit = _unpack(tensor_ybit.flatten(), y)
+            tensor_ybit = tensor_ybit << (nbits - used_bits - y)
+            used_bits += y
+
+            if tensor_floatx is None:
+                tensor_floatx = tensor_ybit
+            else:
+                tensor_floatx |= tensor_ybit
+
+    # undo Pass 1
+    tensor_floatx = tensor_floatx.view(32, -1, 2).permute(1, 0, 2)
+    tensor_floatx = tensor_floatx.reshape(M // 64, -1, 4, 2, 2, 8, 8)
+    tensor_floatx = tensor_floatx.permute(0, 2, 4, 5, 1, 3, 6)
+    tensor_floatx = tensor_floatx.reshape(M, -1)
+    return tensor_floatx
+
+
+# more optimized version of _unpack_tc_floatx() for FP6 by merging ops
+# inverse of _unpack_tc_fp6()
+def _unpack_tc_fp6(tensor: Tensor) -> Tensor:
+    assert tensor.ndim == 2 and tensor.dtype == torch.uint8
+    M = tensor.shape[0]
+    N = tensor.shape[1] // 3 * 4
+    assert (M % 64 == 0) and (N % 64 == 0)
+    size_2bit = M * N // 4
+    size_4bit = M * N // 2
+    tensor = tensor.view(-1)
+    assert tensor.numel() == size_2bit + size_4bit
+
+    tensor_2bit, tensor_4bit = tensor.split([size_2bit, size_4bit])
+
+    tensor_2bit = _unpack(tensor_2bit, 2)
+    tensor_2bit = tensor_2bit.view(M // 64, N // 16, 2, 8, 8, 2, 2, 2)
+    tensor_2bit = tensor_2bit.permute(0, 2, 6, 5, 3, 1, 7, 4)
+
+    tensor_4bit = _unpack(tensor_4bit, 4)
+    tensor_4bit = tensor_4bit.view(M // 64, N // 16, 2, 2, 8, 8, 2, 2)
+    tensor_4bit = tensor_4bit.permute(0, 2, 3, 6, 4, 1, 7, 5)
+
+    tensor_fp6 = (tensor_2bit << 4) | tensor_4bit
+    tensor_fp6 = tensor_fp6.flip(3).reshape(M, N)
+    return tensor_fp6
+
+
+def unpack_tc_floatx(tensor: Tensor, nbits: int) -> Tensor:
+    if nbits == 6:
+        return _unpack_tc_fp6(tensor)
+    return _unpack_tc_floatx(tensor, nbits)
+
+
+def from_scaled_tc_floatx(tensor: Tensor, ebits: int, mbits: int, scale=None) -> Tensor:
+    floatx_unpacked = unpack_tc_floatx(tensor, 1 + ebits + mbits)
+    tensor = _floatx_unpacked_to_f32(floatx_unpacked, ebits, mbits)
+    if scale is not None:
+        tensor = tensor * scale.float().view(-1, 1)
+    return tensor
+
+
+# https://github.com/microsoft/DeepSpeed/blob/3a3a6db3332e339cc9fd94efd4982f6d60635a3d/deepspeed/inference/v2/kernels/core_ops/cuda_linear/cuda_linear.py
+_SPLIT_K_MAP = [
+    {  # tokens: [1, 64]
+        3072: 18,
+        4096: 13,
+        5120: 10,
+        6144: 9,
+        8192: 6,
+        10240: 5,
+        14336: 7,
+        28672: 7,
+        57344: 7,
+    },
+    {  # tokens: [65:128]
+        3072: 9,
+        4096: 6,
+        5120: 5,
+        6144: 9,
+        8192: 3,
+        10240: 5,
+        14336: 7,
+        28672: 7,
+        57344: 6,
+    },
+    {  # tokens: [129:192]
+        3072: 6,
+        4096: 4,
+        5120: 7,
+        6144: 3,
+        8192: 2,
+        10240: 5,
+        14336: 5,
+        28672: 5,
+        57344: 4,
+    },
+    {  # tokens: [193:256]
+        3072: 9,
+        4096: 3,
+        5120: 5,
+        6144: 2,
+        8192: 5,
+        10240: 4,
+        14336: 8,
+        28672: 6,
+        57344: 4,
+    },
+    {  # tokens: [257:320]
+        3072: 7,
+        4096: 5,
+        5120: 2,
+        6144: 5,
+        8192: 4,
+        10240: 1,
+        14336: 3,
+        28672: 3,
+        57344: 4,
+    },
+    {  # tokens: [321:384]
+        3072: 3,
+        4096: 2,
+        5120: 5,
+        6144: 3,
+        8192: 1,
+        10240: 8,
+        14336: 3,
+        28672: 4,
+        57344: 3,
+    },
+    {  # tokens: [385:448]
+        3072: 5,
+        4096: 7,
+        5120: 3,
+        6144: 5,
+        8192: 7,
+        10240: 3,
+        14336: 1,
+        28672: 1,
+        57344: 3,
+    },
+    {  # tokens: [449:512]
+        3072: 2,
+        4096: 5,
+        5120: 4,
+        6144: 1,
+        8192: 5,
+        10240: 2,
+        14336: 6,
+        28672: 4,
+        57344: 1,
+    },
+    {  # tokens: [513:576]
+        3072: 2,
+        4096: 3,
+        5120: 1,
+        6144: 1,
+        8192: 3,
+        10240: 3,
+        14336: 3,
+        28672: 1,
+        57344: 1,
+    },
+    {  # tokens: [577:640]
+        3072: 5,
+        4096: 4,
+        5120: 1,
+        6144: 4,
+        8192: 2,
+        10240: 1,
+        14336: 1,
+        28672: 1,
+        57344: 1,
+    },
+    {  # tokens: [641:704]
+        3072: 3,
+        4096: 1,
+        5120: 2,
+        6144: 2,
+        8192: 1,
+        10240: 2,
+        14336: 1,
+        28672: 1,
+        57344: 1,
+    },
+    {  # tokens: [705:768]
+        3072: 3,
+        4096: 1,
+        5120: 3,
+        6144: 2,
+        8192: 1,
+        10240: 1,
+        14336: 1,
+        28672: 1,
+        57344: 1,
+    },
+]
+
+
+# quantization api integrations
+@dataclass(frozen=True)
+class FloatxTensorCoreLayout(Layout):
+    """FloatxTensorCoreLayout is a data class that defines the layout for a tensor with a specific number of exponent bits (ebits) and mantissa bits (mbits).
+    This layout is used in the context of quantization and packing of tensors optimized for TensorCore operations.
+    """
+
+    ebits: int
+    mbits: int
+
+
+@register_layout(FloatxTensorCoreLayout)
+class FloatxTensorCoreAQTTensorImpl(AQTTensorImpl):
+    """FloatxTensorCoreAQTTensorImpl represents a Tensor with dtype floatx(ebits=a, mbits=b),
+    it has a internal tensor field of "packed_floatx_data", which is packed from the
+    uint8 unpacked data (the output of `quantize_affine_floatx` operator)
+
+    The packing is optimized for TensorCore, from the fp6-llm paper: https://arxiv.org/abs/2401.14112
+    github repo: https://github.com/usyd-fsalab/fp6_llm, now renamed to quant-llm
+
+    At a high level packing is done by grouping bits into 1 bit fragments (shards), 2 bit fragments and
+    4 bit fragments each fragments are packed separately and concatenated together.
+    For example for 6 bit dtype, we can extract the first 4 bits for all elements and pack them together
+    in a fragment, and extract the last 2 bits for all elements and pack them into fragment, in the end
+    we concatenate the fragments together.
+
+    If original Tensor shape is (M, N), and the data is in nbit, the shape of the packed data will be
+    (M, N // 8 * nbit)
+
+    FloatxTensorCoreAQTTensorImpl.from_plain takes an unpacked uint8 floatx Tensor of shape (M, N), with format of
+    (zero padding bits + sign bit + exponent bits + mantissa bits), e.g. 00SEEEMM for fp6_e3_m2
+    it will then pack the weight and instantiate the FloatxTensorCoreAQTTensorImpl tensor
+    FloatxTensorCoreAQTTensorImpl.__init__() takes a packed floatx Tensor of shape (M, N // 8 * nbit)
+    """
+
+    def __new__(
+        cls,
+        packed_floatx_data: torch.Tensor,
+        scale: torch.Tensor,
+        _layout: Layout,
+    ):
+        assert packed_floatx_data.ndim == 2
+        assert packed_floatx_data.dtype == torch.uint8
+        shape = (
+            packed_floatx_data.shape[0],
+            packed_floatx_data.shape[1] // (1 + _layout.ebits + _layout.mbits) * 8,
+        )
+        kwargs = {}
+        kwargs["device"] = packed_floatx_data.device
+        kwargs["layout"] = (
+            kwargs.get("layout")
+            if kwargs.get("layout", False)
+            else packed_floatx_data.layout
+        )
+        kwargs["dtype"] = packed_floatx_data.dtype
+        kwargs["requires_grad"] = False
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        packed_floatx_data: torch.Tensor,
+        scale: torch.Tensor,
+        _layout: Layout,
+    ):
+        self.packed_floatx_data = packed_floatx_data
+        self.scale = scale
+        self._layout = _layout
+
+    def __tensor_flatten__(self):
+        return ["packed_floatx_data", "scale"], [self._layout]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        packed_floatx_data, scale = (
+            tensor_data_dict["packed_floatx_data"],
+            tensor_data_dict["scale"],
+        )
+        (_layout,) = tensor_attributes
+        return cls(packed_floatx_data, scale, _layout)
+
+    def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        unpacked_floatx_data = unpack_tc_floatx(
+            self.packed_floatx_data, 1 + self._layout.ebits + self._layout.mbits
+        )
+        return unpacked_floatx_data, self.scale
+
+    @classmethod
+    def from_plain(
+        cls,
+        unpacked_floatx_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        """
+        Format for `unpacked_floatx_data` will be:
+        zero padding bits | sign bit | exponent bits | mantissa bits
+
+        For example for fp6_e3_m2, the format will be: `00SEEEMM`, where S is sign bit, E is exponent
+        bit, M is mantissa bit
+        """
+        assert isinstance(_layout, FloatxTensorCoreLayout)
+        packed_floatx_data = pack_tc_floatx(
+            unpacked_floatx_data, 1 + _layout.ebits + _layout.mbits
+        )
+        return cls(packed_floatx_data, scale, _layout)
+
+    def __repr__(self):
+        unpacked_floatx_data, scale = self.get_plain()
+        _layout = self.get_layout()
+        return f"{self.__class__.__name__}(unpacked_floatx_data={unpacked_floatx_data}, scale={scale}, _layout={_layout})"
+
+    def _apply_fn_to_data(self, fn):
+        return self.__class__(
+            fn(self.packed_floatx_data),
+            fn(self.scale),
+            self._layout,
+        )
+
+    def to(self, *args, **kwargs):
+        kwargs = self._get_to_kwargs(*args, **kwargs)
+        device = kwargs.pop("device")
+        return self.__class__(
+            self.packed_floatx_data.to(device),
+            self.scale.to(device),
+            self._layout,
+        )
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+        elif func is aten.clone.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.clone)
+            )
+        elif func is aten._to_copy.default:
+            return return_and_correct_aliasing(
+                func,
+                args,
+                kwargs,
+                args[0]._apply_fn_to_data(
+                    lambda x: x.to(device=kwargs.pop("device", None))
+                ),
+            )
+
+        raise NotImplementedError(
+            f"FloatxTensorCoreAQTTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
+    def get_layout(self) -> Layout:
+        return self._layout
+
+
+def _linear_f16_bf16_act_floatx_weight_check(input_tensor, weight_tensor, bias):
+    from torchao.dtypes.floatx import FloatxTensorCoreLayout
+
+    return (
+        # input is native float32 tensor
+        not is_traceable_wrapper_subclass(input_tensor)
+        and input_tensor.is_floating_point()
+        and input_tensor.dtype in (torch.float16, torch.bfloat16)
+        and
+        # weight is floatx Tensor
+        isinstance(weight_tensor, AffineQuantizedTensor)
+        and isinstance(weight_tensor._layout, FloatxTensorCoreLayout)
+        and (
+            # weight is using fp6 quantization
+            (weight_tensor._layout.ebits == 3 and weight_tensor._layout.mbits == 2)
+            or (weight_tensor._layout.ebits == 2 and weight_tensor._layout.mbits == 3)
+            or
+            # weight is using fp5 quantization
+            (weight_tensor._layout.ebits == 2 and weight_tensor._layout.mbits == 2)
+            or (weight_tensor._layout.ebits == 3 and weight_tensor._layout.mbits == 1)
+        )
+    )
+
+
+def _linear_f16_bf16_act_floatx_weight_impl(input_tensor, weight_tensor, bias):
+    from torchao.ops import quant_llm_linear
+
+    act = input_tensor
+    weight = weight_tensor
+
+    out_dim, in_dim = weight.shape
+    act_reshaped = act.view(-1, in_dim)
+
+    # https://github.com/microsoft/DeepSpeed/blob/3a3a6db3332e339cc9fd94efd4982f6d60635a3d/deepspeed/inference/v2/kernels/core_ops/cuda_linear/cuda_linear.py
+    bsize = act_reshaped.shape[0]
+    splitK = _SPLIT_K_MAP[(bsize - 1) // 64].get(out_dim, 1) if bsize <= 768 else 1
+
+    out = quant_llm_linear(
+        weight._layout.ebits,
+        weight._layout.mbits,
+        act_reshaped,
+        weight.tensor_impl.packed_floatx_data,
+        weight.tensor_impl.scale,
+        splitK=splitK,
+    )
+
+    if bias is not None:
+        out += bias
+
+    return out.view(*act.shape[:-1], out_dim).to(act.dtype)

lib/python3.12/site-packages/torchao/dtypes/nf4tensor.py

@@ -0,0 +1,1086 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+import functools
+import math
+import sys
+from dataclasses import dataclass, replace
+from enum import Enum, auto
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch._prims_common import make_contiguous_strides_for
+from torch.distributed.device_mesh import DeviceMesh
+
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
+
+aten = torch.ops.aten
+
+c10d_functional = torch.ops.c10d_functional
+
+
+NF4_OPS_TABLE: Dict[Any, Any] = {}
+
+
+_INNER_TENSOR_NAMES_FOR_SHARDING = [
+    "quantized_scalers",
+    "quantization_factor",
+    "quantized_data",
+]
+
+# Note: Quantize in Chunks
+# During quantization to NF4, one of the steps to convert from the original float number
+# to the index of the nearest value in the NF4 format. This can cause a large memory spike
+# Due to intermediates of the quantization process. Instead we process the original
+# tensor in chunks. This is a tradeoff between memory and speed. This number seems to
+# strike a good balance between memory and speed
+CHUNK_SIZE = 1024**2
+
+
+def same_metadata(a: "NF4Tensor", b: "NF4Tensor"):
+    both_nf4 = isinstance(a, NF4Tensor) and isinstance(b, NF4Tensor)
+    return (
+        both_nf4
+        and a.block_size == b.block_size
+        and a.scaler_block_size == b.scaler_block_size
+        and a.n_blocks == b.n_blocks
+    )
+
+
+def implements(aten_ops):
+    """Use this decorator to implement a function for an aten op in __torch_dispatch__"""
+
+    def decorator(func):
+        for op in aten_ops:
+            NF4_OPS_TABLE[op] = func
+        return func
+
+    return decorator
+
+
+def construct_nf4_args(nf4tensor: "NF4Tensor", kwargs: Optional[Dict[str, Any]] = None):
+    if kwargs is None:
+        kwargs = {}
+    tensor_meta = SubclassTensorArgs(
+        kwargs.get("size", nf4tensor.size()),
+        kwargs.get("stride", nf4tensor.stride()),
+        kwargs.get("storage_offset", nf4tensor.storage_offset()),
+        kwargs.get("dtype", nf4tensor.dtype),
+        kwargs.get("device", nf4tensor.device),
+        kwargs.get("requires_grad", nf4tensor.requires_grad),
+    )
+    return (
+        tensor_meta,
+        kwargs.get("block_size", nf4tensor.block_size),
+        kwargs.get("n_blocks", nf4tensor.n_blocks),
+        kwargs.get("scaler_block_size", nf4tensor.scaler_block_size),
+        kwargs.get("quantized_scalers", nf4tensor.quantized_scalers),
+        kwargs.get("quantization_factor", nf4tensor.quantization_factor),
+        kwargs.get("scaler_mean", nf4tensor.scaler_mean),
+        kwargs.get("quantized_data", nf4tensor.quantized_data),
+        kwargs.get("nf4", nf4tensor.nf4),
+    )
+
+
+# __torch_dispatch__ utils: apply aten op to inner tensors
+def apply_to_inner_tensors(nf4tensor: "NF4Tensor", aten_op, args, kwargs):
+    attr_to_tensor = {}
+    for attr in _INNER_TENSOR_NAMES_FOR_SHARDING:
+        attr_to_tensor[attr] = aten_op(getattr(nf4tensor, attr), *args, **kwargs)
+    return attr_to_tensor
+
+
+# __torch_function__ utils: call tensor ops from inner tensors
+def call_from_inner_tensors(nf4tensor: "NF4Tensor", method_name: str, args, kwargs):
+    attr_to_tensor = {}
+    for attr in _INNER_TENSOR_NAMES_FOR_SHARDING:
+        inner_tensor = getattr(nf4tensor, attr)
+        func = getattr(inner_tensor, method_name)
+        attr_to_tensor[attr] = func(*args, **kwargs)
+    return attr_to_tensor
+
+
+class CompareOp(Enum):
+    EQ = auto()
+    LT = auto()
+
+
+def expect_num_of_args(op: CompareOp, num: int, msg: str):
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(aten_op, args, kwargs=None):
+            if op == CompareOp.LT and not (len(args) < num):
+                raise NotImplementedError(msg)
+            return func(aten_op, args, kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def expect_arg_value_at_k(k: int, op: CompareOp, value: Any, msg: str):
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(aten_op, args, kwargs=None):
+            if op == CompareOp.EQ and not (args[k] == value):
+                raise NotImplementedError(msg + str(args[k]))
+            return func(aten_op, args, kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def expect_args_len_at_k(k: int, op: CompareOp, value: Any, msg: str):
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(aten_op, args, kwargs=None):
+            if op == CompareOp.LT and not (len(args[k]) < value):
+                raise NotImplementedError(msg + str(len(args[k])))
+            elif op == CompareOp.EQ and not (len(args[k]) == value):
+                raise NotImplementedError(msg + str(len(args[k])))
+            return func(aten_op, args, kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+@implements([torch.ops.aten.detach])
+def noop_detach(func, *args, **kwargs):
+    return args[0][0]
+
+
+@implements([torch.ops.aten.clone.default])
+def clone(func, *args, **kwargs):
+    return to_nf4(args[0][0].get_original_weight())
+
+
+@implements(
+    [
+        aten.detach.default,
+    ]
+)
+def nf4_detach(aten_op, args, kwargs=None):
+    nf4tensor = args[0]
+    updated_attrs = apply_to_inner_tensors(nf4tensor, aten_op, args[1:], kwargs)
+    return NF4Tensor(*construct_nf4_args(nf4tensor, updated_attrs))
+
+
+@implements(
+    [
+        aten.empty_like.default,
+    ]
+)
+def nf4_empty_like(aten_op, args, kwargs=None):
+    nf4tensor = args[0]
+    updated_attrs = apply_to_inner_tensors(nf4tensor, aten_op, args[1:], kwargs)
+    if kwargs is not None and len(kwargs):
+        for key, value in kwargs.items():
+            updated_attrs[key] = value
+    return NF4Tensor(*construct_nf4_args(nf4tensor, updated_attrs))
+
+
+@implements(
+    [
+        aten.split.Tensor,
+    ]
+)
+def nf4_split(aten_op, args, kwargs=None):
+    if len(args) == 3 and args[2] != 0:
+        raise NotImplementedError(f"aten.split(NF4Tensor, dim={args[2]})")
+    nf4tensor = args[0]
+    num_chunks = nf4tensor.size(0) // args[1]
+
+    attr_to_chunks = {}
+    for attr in _INNER_TENSOR_NAMES_FOR_SHARDING:
+        inner_tensor = getattr(nf4tensor, attr)
+        assert inner_tensor.numel() % num_chunks == 0, (
+            f"{attr}.numel() not divisible by {num_chunks}"
+        )
+        chunks = aten_op(inner_tensor, inner_tensor.numel() // num_chunks, **kwargs)
+        attr_to_chunks[attr] = chunks
+
+    orig_dim = nf4tensor.dim()
+    if orig_dim == 1:
+        chunked_size = (nf4tensor.size(0) // num_chunks,)
+    elif orig_dim == 2:
+        chunked_size = (nf4tensor.size(0) // num_chunks, nf4tensor.size(1))
+    else:
+        chunked_size = ()
+        raise NotImplementedError(
+            f"aten.split(NF4Tensor) wherer NF4Tensor.dim() = {orig_dim}"
+        )
+
+    nf4_chunks = []
+    for idx in range(num_chunks):
+        updated_attrs = {"size": chunked_size}
+        for attr, chunks in attr_to_chunks.items():
+            updated_attrs[attr] = chunks[idx]
+        nf4_chunks.append(NF4Tensor(*construct_nf4_args(nf4tensor, updated_attrs)))
+    return nf4_chunks
+
+
+@implements(
+    [
+        aten.new_zeros.default,
+    ]
+)
+@expect_args_len_at_k(1, CompareOp.LT, 3, "aten.view(NF4Tensor) with len(size)=")
+def nf4_new_zeros(aten_op, args, kwargs=None):
+    nf4tensor = args[0]
+    new_size = tuple(args[1])
+
+    if nf4tensor.numel() % math.prod(new_size) != 0:
+        raise NotImplementedError(f"aten.new_zeros(NF4Tensor) with new size {new_size}")
+    ratio = nf4tensor.numel() // math.prod(new_size)
+
+    updated_attrs = {}
+    for attr in _INNER_TENSOR_NAMES_FOR_SHARDING:
+        inner_tensor = getattr(nf4tensor, attr)
+        assert inner_tensor.size(0) % ratio == 0, (
+            f"{attr}.numel() must be divisible by {ratio}"
+        )
+        inner_tensor = aten_op(inner_tensor, [inner_tensor.size(0) // ratio], **kwargs)
+        updated_attrs[attr] = inner_tensor
+    updated_attrs["size"] = new_size
+
+    return NF4Tensor(*construct_nf4_args(nf4tensor, updated_attrs))
+
+
+@implements(
+    [
+        aten.slice.Tensor,
+    ]
+)
+@expect_num_of_args(CompareOp.LT, 5, "aten.slice(NF4Tensor) with customized step")
+@expect_arg_value_at_k(1, CompareOp.EQ, 0, "aten.slice(NF4Tensor) with dim=")
+@expect_arg_value_at_k(2, CompareOp.EQ, 0, "aten.slice(NF4Tensor) with start=")
+def nf4_slice(aten_op, args, kwargs=None):
+    nf4tensor = args[0]
+    # for tensor 512 x 512, tensor[:, :512] dispatch to
+    # aten.slice(dim = 0, end=sys.maxsize)
+    if args[3] not in [nf4tensor.size(0), sys.maxsize]:
+        raise NotImplementedError(f"aten.slice(NF4Tensor) with end={args[3]}")
+    return NF4Tensor(*construct_nf4_args(nf4tensor))
+
+
+@implements(
+    [
+        aten.view.default,
+    ]
+)
+@expect_args_len_at_k(1, CompareOp.EQ, 1, "aten.view(NF4Tensor) with len(size)=")
+def nf4_view(aten_op, args, kwargs=None):
+    nf4tensor = args[0]
+    size = args[1]
+    if size[0] != -1:
+        raise NotImplementedError(f"aten.view(NF4Tensor) with size={size}")
+    updated_attrs = apply_to_inner_tensors(nf4tensor, aten_op, args[1:], kwargs)
+    updated_attrs.update(
+        {
+            "size": [nf4tensor.numel()],
+            "stride": (1,),
+        }
+    )
+    return NF4Tensor(*construct_nf4_args(nf4tensor, updated_attrs))
+
+
+@implements(
+    [
+        aten.as_strided.default,
+    ]
+)
+@expect_args_len_at_k(
+    1, CompareOp.LT, 3, "aten.as_strided(NF4Tensor) only support dim <= 2 but got dim="
+)
+def nf4_as_strided(aten_op, args, kwargs=None):
+    nf4tensor = args[0]
+    size = args[1]
+    stride = tuple(args[2])
+    storage_offset = args[3]
+    if math.prod(size) != nf4tensor.numel():
+        raise NotImplementedError(
+            f"aten.as_strided(NF4Tensor) different numel={nf4tensor.numel()} and size={size}"
+        )
+    if stride != make_contiguous_strides_for(size):
+        raise NotImplementedError(
+            f"aten.as_strided(NF4Tensor) only support continuous stride={make_contiguous_strides_for(size)} but got stride={stride}"
+        )
+    if nf4tensor.storage_offset() != storage_offset:
+        raise NotImplementedError(
+            f"aten.as_strided(NF4Tensor) only support original storage offset {nf4tensor.storage_offset()} but got {storage_offset}"
+        )
+    kwargs = {
+        "size": torch.Size(size),
+        "stride": stride,
+        "storage_offset": storage_offset,
+    }
+    return NF4Tensor(*construct_nf4_args(nf4tensor, kwargs))
+
+
+@implements([torch.ops.aten._to_copy.default])
+def _to_copy(func, *args, **kwargs):
+    if not args[0][0].is_contiguous():
+        assert args[0][0].t().is_contiguous()
+        return func(args[0][0].t()).t()
+    out = args[0][0].get_original_weight().to(args[1]["dtype"])
+    if "device" in args[1]:
+        out = out.to(args[1]["device"])
+    return out
+
+
+@implements([torch.ops.aten.to.dtype])
+def to_dtype(func, *args, **kwargs):
+    if not args[0][0].is_contiguous():
+        assert args[0][0].t().is_contiguous()
+        return torch.ops.aten.to.dtype(args[0][0].t(), args[0][1]).t()
+    return args[0][0].get_original_weight().to(args[0][1])
+
+
+@implements([torch.ops.aten.t.default])
+def t_default(func, *args, **kwargs):
+    a = args[0][0]
+    tensor_meta = SubclassTensorArgs(
+        a.size(),
+        (a.stride(1), a.stride(0)),
+        a.storage_offset(),
+        a.dtype,
+        a.device,
+        a.requires_grad,
+    )
+    b = NF4Tensor(
+        tensor_meta,
+        a.block_size,
+        a.n_blocks,
+        a.scaler_block_size,
+        a.quantized_scalers,
+        a.quantization_factor,
+        a.scaler_mean,
+        a.quantized_data,
+        a.nf4,
+    )
+    return b
+
+
+@implements([torch.ops.aten.mm.default])
+def mm_default(func, *args, **kwargs):
+    return linear_nf4(args[0][0], args[0][1])
+
+
+@implements(
+    [
+        aten.copy_.default,
+    ]
+)
+def copy_(func, *args, **kwargs):
+    original: NF4Tensor = args[0][0]
+    copy_in: torch.Tensor = args[0][1]
+
+    # Base Case
+
+    if same_metadata(original, copy_in):
+        original_tensors = original.__tensor_flatten__()[0]
+        for tensor_name in original_tensors:
+            getattr(original, tensor_name).copy_(getattr(copy_in, tensor_name))
+        return
+
+    # Convert Non NF4Tensor into NF4 for copy in
+    if not isinstance(copy_in, NF4Tensor):
+        copy_in_nf4 = NF4Tensor.from_tensor(
+            copy_in.to(original.device), original.block_size, original.scaler_block_size
+        )
+        return original.copy_(copy_in_nf4)
+
+    # Other Tensor is not a NF4Tensor
+    full_precision = copy_in.get_original_weight()
+    same_meta_nf4 = NF4Tensor.from_tensor(
+        full_precision, original.block_size, original.scaler_block_size
+    )
+    return original.copy_(same_meta_nf4)
+
+
+@implements(
+    [
+        aten.is_pinned.default,
+    ]
+)
+def nf4_is_pinned(aten_op, args, kwargs=None):
+    nf4tensor = args[0]
+    for attr in _INNER_TENSOR_NAMES_FOR_SHARDING:
+        inner_tensor = getattr(nf4tensor, attr)
+        if not aten_op(inner_tensor, *(args[1:]), **kwargs):
+            return False
+    return True
+
+
+@implements(
+    [
+        aten._pin_memory.default,
+    ]
+)
+def nf4_pin_memory(aten_op, args, kwargs=None):
+    nf4tensor = args[0]
+    updated_attrs = apply_to_inner_tensors(nf4tensor, aten_op, args[1:], kwargs)
+    return NF4Tensor(*construct_nf4_args(nf4tensor, updated_attrs))
+
+
+@implements(
+    [
+        aten.cat.default,
+    ]
+)
+def nf4_cat(aten_op: torch._ops.OpOverload, args, kwargs=None):
+    tensors_to_cat = args[0]
+    assert all(isinstance(t, torch.Tensor) for t in tensors_to_cat)
+    remaining_args = args[1:]
+
+    ts = []
+    for t in tensors_to_cat:
+        assert isinstance(t, torch.Tensor)
+
+        if isinstance(t, NF4Tensor):
+            ts.append(t.get_original_weight())
+        else:
+            ts.append(t)
+
+    dtype = ts[0].dtype
+    assert all(t.dtype == dtype for t in ts)
+
+    if kwargs is None:
+        kwargs = {}
+
+    tensors = aten_op(ts, *remaining_args, **kwargs)
+    return tensors
+
+
+@dataclass(frozen=True)
+class SubclassTensorArgs:
+    original_shape: torch.Size
+
+    original_strides: Tuple
+    storage_offset: int
+    dtype: torch.dtype
+    device: torch.device
+    requires_grad: bool
+
+
+def get_block_absmax(input_tensor: torch.Tensor, block_size: int) -> torch.Tensor:
+    """Iterate through a flattened tensor getting the absmax scalers for each block
+
+    Args:
+        input_tensor: Input tensor to get scalers for
+        block_size: Block size for the scanning window
+    Returns:
+        torch.Tensor: Tensor of scalers for each block
+    """
+    assert input_tensor.dim() == 1, "Input tensor must be flattened"
+    assert (input_tensor.numel() % block_size) == 0, (
+        f"Input tensor must be divisible by block size, got {input_tensor.numel()} and {block_size}"
+    )
+
+    n_blocks = input_tensor.numel() // block_size
+    blocks = input_tensor.view(n_blocks, block_size)
+    block_scalers = blocks.abs().max(dim=1).values
+    return block_scalers
+
+
+class NF4Tensor(torch.Tensor):
+    """NF4Tensor class for converting a weight to the QLoRA NF4 format"""
+
+    @torch._dynamo.disable
+    def __new__(
+        cls,
+        # Args related for base tensor construction
+        tensor_meta: SubclassTensorArgs,
+        # Args stored on the instance
+        block_size: int,
+        n_blocks: int,
+        scaler_block_size: int,
+        quantized_scalers: torch.Tensor,
+        quantization_factor: torch.Tensor,
+        scaler_mean: torch.Tensor,
+        quantized_data: torch.Tensor,
+        nf4: torch.Tensor,
+    ):
+        """Create a new NF4Tensor object
+        Args:
+            tensor_meta: Metadata for the tensor
+            block_size: Size of the quantization block
+            n_blocks: Number of blocks to cover the full tensor
+            scaler_block_size: Block size for the scalar quantization
+            quantized_scalers: Quantized scalers data' represented a uint8 tensor
+            quantization_factor: Quantization factor, single scalar represented as torch.Tensor
+            scaler_mean: Mean of the scalers
+            quantized_data: Quantized data represented as uint8 tensor
+            nf4: NF4 tensor LUT for the quantization and dequantization
+
+        """
+
+        nf4tensor = torch.Tensor._make_wrapper_subclass(
+            cls,
+            tensor_meta.original_shape,
+            tensor_meta.original_strides,
+            tensor_meta.storage_offset,
+            # Picked some floating dtype, but we need dtype extensibility
+            dtype=tensor_meta.dtype,
+            device=tensor_meta.device,
+            requires_grad=tensor_meta.requires_grad,
+        )
+        return nf4tensor
+
+    @torch._dynamo.disable
+    def __init__(
+        self,
+        tensor_meta: SubclassTensorArgs,
+        block_size: int,
+        n_blocks: int,
+        scaler_block_size: int,
+        quantized_scalers: torch.Tensor,
+        quantization_factor: torch.Tensor,
+        scaler_mean: torch.Tensor,
+        quantized_data: torch.Tensor,
+        nf4: torch.Tensor,
+    ):
+        """Initialize the NF4Tensor class"""
+        self.block_size = block_size
+        self.n_blocks = n_blocks
+        self.scaler_block_size = scaler_block_size
+        self.quantized_scalers = quantized_scalers
+        self.quantization_factor = quantization_factor
+        self.scaler_mean = scaler_mean
+        self.quantized_data = quantized_data
+        self.nf4 = nf4
+
+    @classmethod
+    @torch.no_grad()
+    def from_tensor(
+        cls,
+        input_tensor: torch.Tensor,
+        block_size: int,
+        scaler_block_size: int,
+    ):
+        assert input_tensor.dim() <= 2, (
+            f"expect input tensor dim <= 2 but got dim = {input_tensor.dim()}"
+        )
+        assert input_tensor.numel() % block_size == 0, (
+            f"Input tensor must be divisible by block size, got {input_tensor.numel()} and {block_size}"
+        )
+        assert input_tensor.is_contiguous, "Input tensor must be contiguous!"
+        # I think I want do this
+        # assert not input_tensor.requires_grad, "Input tensor must not require grad"
+        device = input_tensor.device
+        # Cache the tensor on the class def
+        nf4 = torch.tensor(
+            [
+                -1.0000,
+                -0.6962,
+                -0.5251,
+                -0.3949,
+                -0.2844,
+                -0.1848,
+                -0.0911,
+                0.0000,
+                0.0796,
+                0.1609,
+                0.2461,
+                0.3379,
+                0.4407,
+                0.5626,
+                0.7230,
+                1.0000,
+            ],
+            device=device,
+            dtype=input_tensor.dtype,
+        )
+        n_blocks = input_tensor.numel() // block_size
+        # Double quantization
+        (
+            quantized_scalers,
+            quantization_factor,
+            scaler_mean,
+        ) = cls.double_quantize_scalers(
+            input_tensor.flatten(), block_size, scaler_block_size
+        )
+        quantized_data = cls.convert_to_norm_float_weight(
+            input_tensor, n_blocks, block_size, nf4
+        )
+        tensor_meta = SubclassTensorArgs(
+            input_tensor.size(),
+            input_tensor.stride(),
+            input_tensor.storage_offset(),
+            input_tensor.dtype,
+            input_tensor.device,
+            input_tensor.requires_grad,
+        )
+        return cls(
+            tensor_meta,
+            block_size,
+            n_blocks,
+            scaler_block_size,
+            quantized_scalers,
+            quantization_factor,
+            scaler_mean,
+            quantized_data,
+            nf4=nf4,
+        )
+
+    @staticmethod
+    def double_quantize_scalers(
+        input_tensor: torch.Tensor,
+        block_size: int,
+        scaler_block_size: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Used to achieve the double quantization of the scalers
+        We take the input tensor first calculate the absmax quantization factors for each block.
+        We then find the mean of our positive absmax scalers. We subtract this mean from the scalers
+        And then we calculate the absmax quantization factors for each block again. We then quantize the scalers to int8.
+
+        Args:
+            input_tensor: Input tensor to convert to QLoRA format, typically a weight tensor
+
+        Returns:
+            torch.Tensor: Tensor of per_block quantization factors stored in int8 format
+                size: (n_blocks)
+            torch.Tensor: Tensor of per_scaler_block quantization factors stored in int16 format
+                size: (n_scaler_blocks)
+        """
+        assert input_tensor.dim() == 1, "Input tensor must be flattened"
+        assert (input_tensor.numel() % scaler_block_size) == 0, (
+            f"Input tensor must be divisible by block size, got {input_tensor.numel()} and {scaler_block_size}"
+        )
+
+        # First round of quantization
+        # Produces: A tensor of size (n_blocks) of input_tensor.dtype
+        scalers_1 = get_block_absmax(input_tensor, block_size)
+        scalers_1_mean = scalers_1.mean()
+        scalers_1 = scalers_1 - scalers_1_mean
+        # Second round of quantization
+        assert scalers_1.numel() % scaler_block_size == 0, (
+            f"Number of scalers must be divisible by scaler block size, got {scalers_1.numel()} scaler_block_size {scaler_block_size} "
+        )
+        n_scaler_blocks = scalers_1.numel() // scaler_block_size
+        scaler_blocks = scalers_1.view(n_scaler_blocks, scaler_block_size)
+
+        scaler_absmax = get_block_absmax(scalers_1, scaler_block_size)
+        scaler_absmax = scaler_absmax.unsqueeze(-1).expand(
+            n_scaler_blocks, scaler_block_size
+        )
+
+        quantization_factor = 256 / (2 * scaler_absmax)
+        # Length equal to weight numel // block_size
+        quantized_scaler_blocks = scaler_blocks * quantization_factor
+        quantized_scaler_blocks = quantized_scaler_blocks.round()
+        quantized_scaler_blocks = quantized_scaler_blocks.clamp(-128, 127)
+
+        # This is needed to make sure that quantization_factor remains a repeated view of n_scaler_blocks
+        # For some reason the 127/scaler_absmax realizes n_scaler entries when only n_scaler_blocks are needed
+        # The following will grab the first entry for the n_scaler_blocks which is the same across the scaler_block_size
+
+        quantization_factor = quantization_factor[:, 0]
+
+        return (
+            quantized_scaler_blocks.flatten().to(torch.int8),
+            quantization_factor.view(n_scaler_blocks).contiguous(),
+            scalers_1_mean,
+        )
+
+    def dequantize_scalers(
+        self,
+        input_tensor: torch.Tensor,
+        quantization_factor: torch.Tensor,
+        scaler_block_size: int,
+    ) -> torch.Tensor:
+        """Used to unpack the double quantized scalers
+
+        Args:
+            input_tensor: Input tensor to convert to QLoRA format this is the quantized scalers in int8 format
+            quantization_factor: Tensor of per_scaler_block quantization factors stored in inpt_weight.dtype
+            scaler_block_size: Scaler block size to use for double quantization.
+
+        """
+        assert input_tensor.dim() == 1, "Input tensor must be flattened"
+        assert (input_tensor.numel() % scaler_block_size) == 0, (
+            f"Input tensor must be divisible by block size, got {input_tensor.numel()} and {scaler_block_size}"
+        )
+        n_scaler_blocks = input_tensor.numel() // scaler_block_size
+        input_tensor = input_tensor.view(n_scaler_blocks, scaler_block_size)
+        dequantized = (input_tensor / quantization_factor.unsqueeze(-1)).flatten().to(
+            self.dtype
+        ) + self.scaler_mean
+        return dequantized
+
+    @staticmethod
+    def convert_to_norm_float_weight(
+        input_tensor: torch.Tensor, n_blocks: int, block_size: int, nf4: torch.Tensor
+    ) -> torch.Tensor:
+        """Convert a tensor to the normalized float weight format"""
+        flattened_tensor = input_tensor.flatten()
+        #  Since we are using uint8 we will encode 2 entries per byte
+        numel = input_tensor.numel()
+        assert numel % 2 == 0, (
+            "Number of elements must be even just to not have to think about the end"
+        )
+        # Reshape the flattened tensor into blocks of size self.block_size
+        blocks = flattened_tensor.view(n_blocks, block_size)
+
+        # Scale the blocks
+        scalers = get_block_absmax(input_tensor.flatten(), block_size)
+        scales = scalers.unsqueeze(-1).expand(n_blocks, block_size)
+        scaled_blocks = blocks / scales
+
+        # Returns a flattened tensor with each element quantized to nf4 index
+        # See Note: Quantize in Chunks
+        quantized_blocks = torch.empty(
+            numel, dtype=torch.uint8, device=input_tensor.device
+        )
+        flattened = scaled_blocks.flatten()
+        for chunk_num in range(math.ceil(numel / CHUNK_SIZE)):
+            start = chunk_num * CHUNK_SIZE
+            end = min(start + CHUNK_SIZE, numel)
+            quantized_blocks[start:end] = NF4Tensor.quantize_tensor_nearest(
+                flattened[start:end], nf4
+            ).to(torch.uint8)
+
+        # Combine the quantized elements into uint8 values
+        # This lays out two consecutive elements in the same byte
+        # [a, b, c, d] -> [ab, cd]
+        # The size of combined blocks will be half the size of the original tensor
+        combined_blocks = quantized_blocks[::2] << 4 | quantized_blocks[1::2]
+
+        return combined_blocks.to(torch.uint8)
+
+    def get_original_weight(self) -> torch.Tensor:
+        """Get the original weight from the normalized float weight format"""
+        # Since we are using uint8 we will decode 2 entries per byte
+        # Shift elements down 4 and select out the bottom 4 bits
+        first_elements = (self.quantized_data >> 4).to(torch.long)
+        second_elements = (self.quantized_data & 0b1111).to(torch.long)
+
+        # Dequantize every element
+        dequantized_first = self.dequantize(first_elements, self.nf4)
+        dequantized_second = self.dequantize(second_elements, self.nf4)
+
+        # Build up matrix of scalers repeated for each element in the block
+        # Since first and second elements make up a full block
+        # we expand out to half the size of the full block
+        scalers = self.dequantize_scalers(
+            self.quantized_scalers, self.quantization_factor, self.scaler_block_size
+        )
+        repeated = scalers.unsqueeze(-1).expand(scalers.size(0), self.block_size // 2)
+
+        scaled_first = dequantized_first * repeated.flatten()
+        scaled_second = dequantized_second * repeated.flatten()
+
+        # Flip them to be vertical and them stack them together horizontally
+        # Upon flattening this will interleave the elements
+        scaled_first = scaled_first.unsqueeze(-1).transpose(0, 1)
+        scaled_second = scaled_second.unsqueeze(-1).transpose(0, 1)
+        return torch.stack([scaled_first, scaled_second], dim=-1).reshape(self.shape)
+
+    @staticmethod
+    def quantize_tensor_nearest(value: torch.Tensor, nf4: torch.Tensor) -> torch.Tensor:
+        """Quantize a float16 tensor to nf4 format to nearest and not rounded up"""
+        value = value.unsqueeze(-1)  # (numel, 1)
+        # Compare the value tensor with the nf4 tensor element-wise
+        diff = (value - nf4).abs()
+        closest_nf4 = diff.min(dim=-1).indices
+        return closest_nf4
+
+    @staticmethod
+    def dequantize(value: torch.Tensor, nf4: torch.Tensor) -> torch.Tensor:
+        """Dequantize a nf4 value to bfloat16 format"""
+        # return nf4.index_select(0, value)
+        return nf4[value]
+
+    def __repr__(self) -> str:
+        return f"Quantized Data: {self.quantized_data}\nScalers: {self.quantized_scalers}\n"
+
+    def __str__(self) -> str:
+        return f"NF4Tensor({self.shape}, {self.block_size})"
+
+    def __tensor_flatten__(self):
+        tensor_meta = SubclassTensorArgs(
+            self.shape,
+            self.stride(),
+            self.storage_offset(),
+            self.dtype,
+            self.device,
+            self.requires_grad,
+        )
+        ctx = {
+            "block_size": self.block_size,
+            "n_blocks": self.n_blocks,
+            "scaler_block_size": self.scaler_block_size,
+            "tensor_meta": tensor_meta,
+        }
+        return [
+            "quantized_data",
+            "scaler_mean",
+            "quantization_factor",
+            "quantized_scalers",
+            "nf4",
+        ], ctx
+
+    @staticmethod
+    def __tensor_unflatten__(inner_tensors: Dict, metadata, outer_size, outer_stride):
+        assert len(inner_tensors) == 5, "Expected 5 inner tensors"
+        return NF4Tensor(
+            metadata["tensor_meta"],
+            metadata["block_size"],
+            metadata["n_blocks"],
+            metadata["scaler_block_size"],
+            inner_tensors["quantized_scalers"],
+            inner_tensors["quantization_factor"],
+            inner_tensors["scaler_mean"],
+            inner_tensors["quantized_data"],
+            inner_tensors["nf4"],
+        )
+
+    @classmethod
+    @torch._dynamo.disable
+    def __torch_dispatch__(cls, func, types, args, kwargs=None):
+        """TODO we are not supporting torch dispatch at the moment
+        instead we have created a Autograd.Function to handle the linear
+        """
+        # All ops in the  NF4_OPS_TABLE expect NF4 Tensors as inputs
+        # And don't support mixed tensor subclasses. This will trigger the handler for
+        # the next type in the dispatch list
+
+        def allowed_subclasses(type):
+            return (
+                issubclass(cls, type)
+                or issubclass(torch._subclasses.fake_tensor.FakeTensor, type)
+                or issubclass(
+                    torch._subclasses.functional_tensor.FunctionalTensor, type
+                )
+            )
+
+        if not all(allowed_subclasses(t) for t in types):
+            return NotImplemented("Up to the next one to handle")
+
+        if func in NF4_OPS_TABLE:
+            return NF4_OPS_TABLE[func](func, args, kwargs)
+        raise NotImplementedError(
+            f"NF4Tensor dispatch: attempting to run {func}, this is not supported"
+        )
+
+    # Do not force the Float8Tensor type on the returned tensor
+
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        if kwargs is None:
+            kwargs = {}
+
+        try:
+            if func in NF4_TORCH_FUNCTIONS:
+                return NF4_TORCH_FUNCTIONS[func](*args, **kwargs)
+        except NotImplementedError:
+            pass
+
+        with torch._C.DisableTorchFunctionSubclass():
+            return func(*args, **kwargs)
+
+    def fsdp_pre_all_gather(
+        self, mesh: DeviceMesh
+    ) -> Tuple[Tuple[torch.Tensor, ...], Any]:
+        return (
+            self.quantized_scalers,
+            self.quantization_factor,
+            self.quantized_data,
+        ), (
+            SubclassTensorArgs(
+                self.size(),
+                self.stride(),
+                self.storage_offset(),
+                self.dtype,
+                self.device,
+                self.requires_grad,
+            ),
+            self.block_size,
+            self.n_blocks,
+            self.scaler_block_size,
+            self.scaler_mean,
+            self.nf4,
+            mesh.get_group().size(),
+        )
+
+    def fsdp_post_all_gather(
+        self,
+        all_gather_outputs: Tuple[torch.Tensor, ...],
+        metadata: Any,
+        param_dtype: torch.dtype,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> Union[Tuple[torch.Tensor, Tuple[torch.Tensor, ...]], None]:
+        (quantized_scalers, quantization_factor, quantized_data) = all_gather_outputs
+        (
+            tensor_meta,
+            block_size,
+            n_blocks,
+            scaler_block_size,
+            scaler_mean,
+            nf4,
+            pg_size,
+        ) = metadata
+        if len(tensor_meta.original_shape) != 2:
+            raise NotImplementedError(
+                f"only support 2D shape but got dim={len(tensor_meta.original_shape)}"
+            )
+
+        new_shape = torch.Size(
+            (tensor_meta.original_shape[0] * pg_size, tensor_meta.original_shape[1])
+        )
+        new_tensor_meta = replace(tensor_meta, original_shape=new_shape)
+        if out is not None:
+            # TODO: add param dtype for mixed precision
+            assert isinstance(out, NF4Tensor), f"{type(out)}"
+            assert (
+                quantized_scalers.untyped_storage().data_ptr()
+                == out.quantized_scalers.untyped_storage().data_ptr()
+                and quantization_factor.untyped_storage().data_ptr()
+                == out.quantization_factor.untyped_storage().data_ptr()
+                and quantized_data.untyped_storage().data_ptr()
+                == out.quantized_data.untyped_storage().data_ptr()
+            ), "Expects out's data to be the all-gather output"
+            return
+
+        return nf4_constructor(
+            new_tensor_meta,
+            block_size,
+            n_blocks,
+            scaler_block_size,
+            quantized_scalers,
+            quantization_factor,
+            scaler_mean,
+            quantized_data,
+            nf4,
+        ), (quantized_scalers, quantization_factor, quantized_data)
+
+
+class LinearNF4(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight: NF4Tensor):
+        """Save the quantized nf4 weight for backward pass"""
+        ctx.save_for_backward(weight)
+        return F.linear(input, weight.to(input.dtype))
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """The nf4 weight will never require grad so we can just return the grad_output @ weight.to(grad_output.dtype)"""
+        weight: NF4Tensor = ctx.saved_tensors[0]
+        return grad_output @ weight.to(grad_output.dtype), None
+
+
+def linear_nf4(input: torch.Tensor, weight: NF4Tensor) -> torch.Tensor:
+    """Apply a linear operation with the NF4Tensor weight
+
+    Args:
+        input: Input tensor
+        weight: NF4Tensor weight
+    """
+    return LinearNF4.apply(input, weight)
+
+
+def to_nf4(tensor, block_size: int = 64, scaler_block_size: int = 256):
+    """Convert a given tensor to normalized float 4-bit tensor."""
+    return NF4Tensor.from_tensor(tensor, block_size, scaler_block_size)
+
+
+NF4_TORCH_FUNCTIONS = {}
+
+
+def implements_torch_function(torch_function):
+    def decorator(func):
+        functools.update_wrapper(func, torch_function)
+        NF4_TORCH_FUNCTIONS[torch_function] = func
+        return func
+
+    return decorator
+
+
+@implements_torch_function(torch.Tensor.to)
+def function_to_dtype(*args, **kwargs):
+    tensor = args[0]
+    device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(
+        *args[1:], **kwargs
+    )
+
+    # dtype is specified -> dequantize
+    if dtype is not None:
+        return tensor.get_original_weight().to(
+            device, dtype, non_blocking, memory_format=convert_to_format
+        )
+
+    # dtype is not specified -> keep NF4
+    updated_attrs = dict(device=device)
+    tensor_attrs, _ = tensor.__tensor_flatten__()
+    for attr in tensor_attrs:
+        inner_tensor = getattr(tensor, attr)
+        updated_attrs[attr] = inner_tensor.to(
+            device, dtype, non_blocking, memory_format=convert_to_format
+        )
+    return NF4Tensor(*construct_nf4_args(tensor, updated_attrs))
+
+
+@implements_torch_function(torch.Tensor.cpu)
+def function_cpu(*args, **kwargs):
+    # Tensor.cpu(self, memory_format)
+    return args[0].to("cpu", *args[1:], **kwargs)
+
+
+@implements_torch_function(torch.Tensor.cuda)
+def function_cuda(*args, **kwargs):
+    # Tensor.cuda(self, device, non_blocking, memory_format)
+    tensor = args[0]
+    updated_attrs = dict()
+    tensor_attrs, _ = tensor.__tensor_flatten__()
+    for attr in tensor_attrs:
+        inner_tensor = getattr(tensor, attr)
+        updated_attrs[attr] = inner_tensor.cuda(*args[1:], **kwargs)
+    updated_attrs["device"] = updated_attrs[tensor_attrs[0]].device
+    return NF4Tensor(*construct_nf4_args(tensor, updated_attrs))
+
+
+@implements_torch_function(F.linear)
+def _(*args, **kwargs):
+    input = args[0]
+    weight = args[1]
+    bias = args[2] if len(args) > 2 else None
+    out = LinearNF4.apply(input, weight)
+    if bias is not None:
+        out = out + bias
+    return out
+
+
+@torch._dynamo.allow_in_graph
+def nf4_constructor(
+    tensor_meta: SubclassTensorArgs,
+    block_size: int,
+    n_blocks: int,
+    scaler_block_size: int,
+    quantized_scalers: torch.Tensor,
+    quantization_factor: torch.Tensor,
+    scaler_mean: torch.Tensor,
+    quantized_data: torch.Tensor,
+    nf4: torch.Tensor,
+):
+    return NF4Tensor(
+        tensor_meta,
+        block_size,
+        n_blocks,
+        scaler_block_size,
+        quantized_scalers,
+        quantization_factor,
+        scaler_mean,
+        quantized_data,
+        nf4,
+    )
+
+
+if TORCH_VERSION_AT_LEAST_2_5:
+    torch.serialization.add_safe_globals([NF4Tensor])
+    torch.serialization.add_safe_globals([NF4Tensor])

lib/python3.12/site-packages/torchao/dtypes/uintx/__init__.py

@@ -0,0 +1,51 @@
+from .block_sparse_layout import (
+    BlockSparseLayout,
+)
+from .cutlass_int4_packed_layout import (
+    CutlassInt4PackedLayout,
+)
+from .int4_cpu_layout import (
+    Int4CPULayout,
+)
+from .int4_xpu_layout import (
+    Int4XPULayout,
+)
+from .marlin_qqq_tensor import (
+    MarlinQQQLayout,
+    MarlinQQQTensor,
+    to_marlinqqq_quantized_intx,
+)
+from .marlin_sparse_layout import (
+    MarlinSparseLayout,
+)
+from .packed_linear_int8_dynamic_activation_intx_weight_layout import (
+    PackedLinearInt8DynamicActivationIntxWeightLayout,
+)
+from .q_dq_layout import (
+    QDQLayout,
+)
+from .semi_sparse_layout import (
+    SemiSparseLayout,
+)
+from .tensor_core_tiled_layout import (
+    TensorCoreTiledLayout,
+)
+from .uintx_layout import (
+    UintxLayout,
+)
+
+__all__ = [
+    "UintxLayout",
+    "BlockSparseLayout",
+    "MarlinSparseLayout",
+    "SemiSparseLayout",
+    "TensorCoreTiledLayout",
+    "Int4CPULayout",
+    "MarlinQQQLayout",
+    "MarlinQQQTensor",
+    "to_marlinqqq_quantized_intx",
+    "CutlassInt4PackedLayout",
+    "PackedLinearInt8DynamicActivationIntxWeightLayout",
+    "QDQLayout",
+    "Int4XPULayout",
+]

lib/python3.12/site-packages/torchao/dtypes/uintx/bitpacking.py

@@ -0,0 +1,265 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+from functools import reduce
+from typing import List, Optional
+
+import torch
+
+# for selecting the shards from 8 bits
+maskbits = {
+    1: (0x01,),
+    2: (0x03,),
+    3: (0x03, 0x04),
+    4: (0x0F,),
+    5: (0x0F, 0x10),
+    6: (0x0F, 0x30),
+    7: (0x0F, 0x30, 0x40),
+}
+
+unpack_mask = {
+    1: (0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80),
+    2: (0x03, 0x0C, 0x30, 0xC0),
+    4: (0x0F, 0xF0),
+}
+
+# size of each shard
+numbits = {
+    1: (1,),
+    2: (2,),
+    3: (2, 1),
+    4: (4,),
+    5: (4, 1),
+    6: (4, 2),
+    7: (4, 2, 1),
+}
+
+# shift amount for each shard
+shifts = {
+    1: (0,),
+    2: (0,),
+    3: (0, 2),
+    4: (0,),
+    5: (0, 4),
+    6: (0, 4),
+    7: (0, 4, 6),
+}
+
+
+# for shifting groups left but right if shift is negative
+def abs_lsh(data, shift):
+    if shift == 0:
+        return data
+    elif shift < 0:
+        return data >> -shift
+    else:
+        return data << shift
+
+
+# inverse of abs_lsh for unpacking
+def abs_rsh(data, shift):
+    if shift == 0:
+        return data
+    elif shift < 0:
+        return data << -shift
+    else:
+        return data >> shift
+
+
+def pack_cpu(
+    data: torch.Tensor, elem_size: int, dim: Optional[int] = -1
+) -> List[torch.Tensor]:
+    """
+    Inputs:
+    data: a tensor of sub byte elements in uint8
+    elem_size: the size in bits of the elements to pack
+    dim: the dimension to pack along
+    Returns: a list of packed shards
+
+    ==================================================================================================
+    given an array such as [0x30,0x29,0x17,0x5,0x20,0x16,0x9,0x22] which are 8 uint6 elements
+    first seperate into two shards: the upper 2 bits and the lower 4 bits by using a mask (0x30 and 0x0f respectively)
+    2 bit shard:
+    mask: 0x30
+    [0x30,       0x20,       0x10,       0x00,        0x00,       0x10,       0x00,        0x20    ]
+    [0b00110000, 0b00100000, 0b00010000, 0b00000000, 0b00100000, 0b00010000, 0b00000000, 0b00100000]
+
+    Group elements into subsets that will be shifted to the same position within the 8bit container
+    group1 >> 4,  group2 >> 2, group3 >> 0, group4 << 2
+
+    [0b00000011, 0b00000010, 0b00000100, 0b00000000, 0b00100000, 0b00010000, 0b00000000, 0b10000000]
+    |------ group 1 ------| |------ group 2 ------| |------ group 3 ------| |------ group 4 ------|
+
+    Finally bitwise-or the groups together
+    [0b00000011, 0b00000010,
+     0b00000100, 0b00000000,
+     0b00100000, 0b00010000,
+     0b00000000, 0b01000000]
+
+    [0b00100111, 0b10010010]
+    ==================================================================================================
+    Similarly for 4 bit shards:
+    mask: 0x0f
+    [0x00,       0x09,       0x07,       0x05,       0x00,       0x16,       0x9,        0x02]
+    [0b00000000, 0b00001001, 0b00000111, 0b00000101, 0b00000000, 0b00000110, 0b00001001, 0b00000010]
+
+    group1 << 0, group2 << 4
+    [0b00000000, 0b00001001, 0b00000111, 0b00000101, 0b00000000, 0b01100000, 0b10010000, 0b00100000]
+    |------------------ group 1 ------------------| |------------------ group 2 ------------------|
+
+    bitwise-or:
+    [0b00000000, 0b00001001, 0b00000111, 0b00000101,
+     0b00000000, 0b01100000, 0b10010000, 0b00100000]
+
+    [0b00000000, 0b01101001, 0b10010111, 0b00100101]
+    ==================================================================================================
+    After pack, data went from 8 elements to 6: [[0, 105, 151, 37], [39, 146]]
+    In general this means pack reduces input tensor size from n * 8 to n * elem_size
+    """
+    torch._assert(
+        data.shape[dim] % 8 == 0,
+        f"pack dimension size ({data.shape[dim]}) is not divisble by scale",
+    )
+    torch._assert(data.dtype == torch.uint8, "data must be uint8")
+    output_shape = list(data.shape)
+
+    output = []
+    for i in range(len(numbits[elem_size])):
+        output_shape[dim] = data.shape[dim] * numbits[elem_size][i] // 8
+        shard = torch.zeros(output_shape, dtype=torch.uint8, device=data.device)
+        bit_size = numbits[elem_size][i]
+        rel_pos = shifts[elem_size][i]
+        bits = data & maskbits[elem_size][i]
+        scale = 8 // bit_size
+        slice_len = bits.shape[dim] // scale
+        for j in range(scale):
+            bit_slice = bits.narrow(dim, slice_len * j, slice_len)
+            shard = torch.bitwise_or(shard, abs_lsh(bit_slice, j * bit_size - rel_pos))
+        output.append(shard)
+    return output
+
+
+def unpack_cpu(
+    data: List[torch.Tensor], elem_size: int, dim: Optional[int] = -1
+) -> torch.Tensor:
+    """
+    Unpacks small dtype elements from a larger dtype.
+
+    Inputs:
+    data: - a list of packed shards
+    elem_size: the size in bits of the elements to unpack
+    dim: the dimension to unpack along
+
+    Returns: torch.Tensor - a tensor of the unpacked elements.
+    """
+    # define the output tensor
+    output_shape = list(data[0].shape)
+    output_shape[dim] = data[0].shape[dim] * 8 // numbits[elem_size][0]
+    output = torch.zeros(output_shape, dtype=torch.uint8, device=data[0].device)
+
+    for i in range(len(numbits[elem_size])):
+        # define variables for the current shard
+        bit_size = numbits[elem_size][i]
+        rel_pos = shifts[elem_size][i]
+        scale = 8 // bit_size
+        group_size = bit_size * output_shape[dim] // 8
+        # mask and shift every group of bits to the correct position
+        for j in range(scale):
+            output_narrow = output.narrow(dim, j * group_size, group_size)
+            group = data[i] & unpack_mask[bit_size][j]
+            output_narrow.copy_(
+                torch.bitwise_or(output_narrow, abs_rsh(group, j * bit_size - rel_pos))
+            )
+    return output
+
+
+# these are faster on the GPU
+
+
+def _pack(data, elem_size, scale, dim):
+    """
+    Inner for loop from above pack function
+    """
+    packed_shape = list(data.shape)
+    packed_shape[dim] = packed_shape[dim] // scale
+
+    packed = torch.zeros(packed_shape, dtype=data.dtype, device=data.device)
+
+    for i in range(scale):
+        narrow_slice = data.narrow(
+            dim, data.shape[dim] * i // scale, data.shape[dim] // scale
+        )
+        packed |= narrow_slice << (elem_size * i)
+
+    return packed
+
+
+def _unpack(data, element_size, scale, dim):
+    """
+    Inner for loop from above unpack function
+    """
+    unpacked_shape = list(data.shape)
+    unpacked_shape[dim] *= scale
+
+    nbits = (1 << element_size) - 1  # mask for the last element_size bits
+
+    unpacked_data = torch.zeros(unpacked_shape, dtype=data.dtype, device=data.device)
+
+    for i in range(scale):
+        shift_amt = element_size * i
+        unpacked_data.narrow(
+            dim,
+            unpacked_data.shape[dim] * i // scale,
+            unpacked_data.shape[dim] // scale,
+        ).copy_((data >> shift_amt) & nbits)
+
+    return unpacked_data
+
+
+def pack(
+    data: torch.Tensor, elem_size: int, dim: Optional[int] = -1
+) -> List[torch.Tensor]:
+    """
+    a less branching but more compute version so better for gpu
+    """
+    torch._assert(
+        data.shape[dim] % 8 == 0,
+        f"pack dimension size ({data.shape[dim]}) is not divisble by scale",
+    )
+    torch._assert(data.dtype == torch.uint8, "data must be uint8")
+    container_size = 8
+    shards = [
+        (data & maskbits[elem_size][i]) >> shifts[elem_size][i]
+        for i in range(len(maskbits[elem_size]))
+    ]
+    return tuple(
+        [
+            _pack(
+                shards[i],
+                numbits[elem_size][i],
+                container_size // numbits[elem_size][i],
+                dim,
+            )
+            for i in range(len(maskbits[elem_size]))
+        ]
+    )
+
+
+def unpack(
+    data: List[torch.Tensor], elem_size: int, dim: Optional[int] = 0
+) -> torch.Tensor:
+    """
+    a less branching but more compute version so better for gpu
+    """
+    container_size = 8
+    # unpack each 4,2,1 bit shard and unshift them back to the correct position
+    data = [
+        _unpack(
+            data[i], numbits[elem_size][i], container_size // numbits[elem_size][i], dim
+        )
+        << shifts[elem_size][i]
+        for i in range(len(data))
+    ]
+    return reduce(torch.bitwise_or, data)

lib/python3.12/site-packages/torchao/dtypes/uintx/block_sparse_layout.py

@@ -0,0 +1,233 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+import logging
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch.utils._python_dispatch import (
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+from torchao.dtypes.uintx.plain_layout import (
+    PlainAQTTensorImpl,
+    _aqt_is_int8_reduced_range,
+)
+from torchao.dtypes.utils import (
+    Layout,
+    PlainLayout,
+)
+
+logger = logging.getLogger(__name__)
+
+aten = torch.ops.aten
+
+
+@dataclass(frozen=True)
+class BlockSparseLayout(Layout):
+    """BlockSparseLayout is a data class that represents the layout of a block sparse matrix.
+
+    Attributes:
+        blocksize (int): The size of the blocks in the sparse matrix. Default is 64.
+    """
+
+    blocksize: int = 64
+
+
+@register_layout(BlockSparseLayout)
+class BlockSparseAQTTensorImpl(PlainAQTTensorImpl):
+    bsr_crow_indices: Optional[torch.Tensor]
+    bsr_col_indices: Optional[torch.Tensor]
+    bsr_values: Optional[torch.Tensor]
+    scale: Optional[torch.Tensor]
+    zero_point: Optional[torch.Tensor]
+
+    __slots__ = [
+        "bsr_crow_indices",
+        "bsr_col_indices",
+        "bsr_values",
+        "scale",
+        "zero_point",
+    ]
+
+    @staticmethod
+    def __new__(  # noqa: PYI034
+        cls,
+        shape: torch.Size,
+        bsr_crow_indices: Optional[torch.Tensor],
+        bsr_col_indices: Optional[torch.Tensor],
+        bsr_values: Optional[torch.Tensor],
+        scale: Optional[torch.Tensor],
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+        requires_grad: bool = False,
+    ):
+        if bsr_values is None:
+            raise ValueError("bsr values must be provided!")
+        else:
+            previous_tensor = bsr_values
+
+        kwargs = {
+            "device": previous_tensor.device,
+            "dtype": previous_tensor.dtype,
+            "layout": previous_tensor.layout,
+            "requires_grad": requires_grad,
+        }
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(  # noqa: PYI034
+        self,
+        shape: torch.Size,
+        bsr_crow_indices: Optional[torch.Tensor],
+        bsr_col_indices: Optional[torch.Tensor],
+        bsr_values: Optional[torch.Tensor],
+        scale: Optional[torch.Tensor],
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+        requires_grad: bool = False,
+    ):
+        self.bsr_crow_indices = bsr_crow_indices
+        self.bsr_col_indices = bsr_col_indices
+        self.bsr_values = bsr_values
+        self.scale = scale
+        self.zero_point = zero_point
+        self._layout = _layout
+
+    def __tensor_flatten__(self):
+        inner_tensors = list(
+            filter(lambda x: getattr(self, x) is not None, self.__slots__)
+        )
+        tensor_meta = (self.shape, self._layout, self.requires_grad)
+        return inner_tensors, tensor_meta
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls,
+        inner_tensors,
+        tensor_meta: Tuple[torch.Size, bool],
+        outer_size,
+        outer_stride,
+    ) -> torch.Tensor:
+        shape, _layout, requires_grad = tensor_meta
+        return cls(
+            shape=shape,
+            bsr_crow_indices=inner_tensors.get("bsr_crow_indices", None),
+            bsr_col_indices=inner_tensors.get("bsr_col_indices", None),
+            bsr_values=inner_tensors.get("bsr_values", None),
+            scale=inner_tensors.get("scale", None),
+            zero_point=inner_tensors.get("zero_point", None),
+            _layout=_layout,
+            requires_grad=requires_grad,
+        )
+
+    @classmethod
+    def from_plain(cls, int_data, scale, zero_point, _layout):
+        bsr_tensor = int_data.to_sparse_bsr(_layout.blocksize)
+        return cls(
+            shape=int_data.shape,
+            bsr_crow_indices=bsr_tensor.crow_indices(),
+            bsr_col_indices=bsr_tensor.col_indices(),
+            bsr_values=bsr_tensor.values(),
+            scale=scale,
+            zero_point=zero_point,
+            _layout=_layout,
+            requires_grad=False,
+        )
+
+    def get_plain(self):
+        int_data_expanded = torch.ops.blocksparse.bsr_to_dense(
+            self.crow_indices(),
+            self.col_indices(),
+            self.values(),
+            self.shape[0],
+            self.shape[1],
+        )
+        return int_data_expanded, self.scale, self.zero_point
+
+    def _apply_fn_to_data(self, func):
+        return self.__class__(
+            shape=self.shape,
+            bsr_crow_indices=func(self.bsr_crow_indices),
+            bsr_col_indices=func(self.bsr_col_indices),
+            bsr_values=func(self.bsr_values),
+            scale=self.scale,
+            zero_point=self.zero_point,
+            _layout=self._layout,
+            requires_grad=self.requires_grad,
+        )
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+        if func is aten.clone.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.clone)
+            )
+
+        # Need the following for bsr specific functions
+        if func is aten.crow_indices.default:
+            return args[0].bsr_crow_indices.detach()
+
+        if func is aten.col_indices.default:
+            return args[0].bsr_col_indices.detach()
+
+        if func is aten.values.default:
+            return args[0].bsr_values.detach()
+
+        if func is aten._nnz.default:
+            return args[0].bsr_values.shape[0]
+
+        raise NotImplementedError(
+            f"BlockSparseAQTTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+
+def _linear_int8_act_int8_weight_block_sparse_check(input_tensor, weight_tensor, bias):
+    return (
+        isinstance(input_tensor, AffineQuantizedTensor)
+        and _aqt_is_int8_reduced_range(input_tensor)
+        and isinstance(weight_tensor, AffineQuantizedTensor)
+        and weight_tensor.is_cuda
+        and input_tensor.dtype == weight_tensor.dtype
+        and isinstance(input_tensor._layout, PlainLayout)
+        and isinstance(weight_tensor._layout, BlockSparseLayout)
+    )
+
+
+def _linear_int8_act_int8_weight_block_sparse_impl(input_tensor, weight_tensor, bias):
+    x_vals_int8 = input_tensor.tensor_impl.int_data
+    x_scales = input_tensor.tensor_impl.scale
+    w_vals = weight_tensor.tensor_impl
+    w_scales = weight_tensor.tensor_impl.scale
+    tmp = x_vals_int8.reshape(-1, x_vals_int8.shape[-1])
+    tmp_t = tmp.t()
+
+    y = torch.ops.blocksparse.int_addmm(
+        w_vals.crow_indices(),
+        w_vals.col_indices(),
+        w_vals.values(),
+        tmp_t,
+        w_scales,
+        x_scales.reshape(-1),
+    )
+    y_shape = (*x_vals_int8.shape[:-1], w_scales.shape[-1])
+    y = y.reshape(*y_shape)
+
+    # can downcast only at the very end
+    output_dtype = input_tensor.dtype
+    y = y.to(output_dtype)
+    if bias is not None:
+        y += bias
+    return y

lib/python3.12/site-packages/torchao/dtypes/uintx/cutlass_int4_packed_layout.py

@@ -0,0 +1,224 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch.utils._python_dispatch import (
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+from torchao.dtypes.uintx.plain_layout import (
+    _aqt_is_int8,
+)
+from torchao.dtypes.utils import AQTTensorImpl, Layout, PlainLayout
+
+aten = torch.ops.aten
+
+
+def _aqt_is_int4(aqt):
+    """Check if an AffineQuantizedTensor is int4 quantized Tensor"""
+    # TODO: use torch.int4
+    return (
+        aqt.tensor_impl.dtype == torch.int8
+        and aqt.quant_min == -8
+        and aqt.quant_max == 7
+    )
+
+
+def _same_metadata(self: "Int4PackedTensorImpl", src: "Int4PackedTensorImpl") -> bool:
+    return (
+        isinstance(self, Int4PackedTensorImpl)
+        and isinstance(src, Int4PackedTensorImpl)
+        and self.shape == src.shape
+        and self.int_data.shape == src.int_data.shape
+        and self.scale.shape == src.scale.shape
+        and type(self._layout) == type(src._layout)
+    )
+
+
+@dataclass(frozen=True)
+class CutlassInt4PackedLayout(Layout):
+    """Layout class for int4 packed layout for affine quantized tensor, for cutlass kernel."""
+
+    pass
+
+
+@register_layout(CutlassInt4PackedLayout)
+class Int4PackedTensorImpl(AQTTensorImpl):
+    """
+    TensorImpl storage class for int4 packed layout for affine quantized tensor.
+    """
+
+    @staticmethod
+    def __new__(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        _layout: Layout,
+    ):
+        kwargs = {}
+        kwargs["device"] = int_data.device
+        kwargs["layout"] = (
+            kwargs.get("layout") if kwargs.get("layout", False) else int_data.layout
+        )
+        kwargs["dtype"] = int_data.dtype
+        kwargs["requires_grad"] = False
+        shape = int_data.shape
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        _layout: Layout,
+    ):
+        self.int_data = int_data
+        self.scale = scale
+        self._layout = _layout
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+
+        elif func is aten.copy_.default:
+            self = args[0]
+            src = args[1]
+            if _same_metadata(self, src):
+                self_tensors = self.__tensor_flatten__()[0]
+                for tensor_name in self_tensors:
+                    getattr(self, tensor_name).copy_(getattr(src, tensor_name))
+                return
+            raise ValueError(
+                f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}"
+            )
+
+        raise NotImplementedError(
+            f"Int4PackedTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    def __tensor_flatten__(self):
+        return ["int_data", "scale"], [self._layout]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        int_data = tensor_data_dict["int_data"]
+        scale = tensor_data_dict["scale"]
+        (_layout,) = tensor_attributes
+        return cls(int_data, scale, _layout)
+
+    def get_plain(self):
+        int_data = torch.stack(
+            ((self.int_data << 4) >> 4, self.int_data >> 4), dim=-1
+        ).view(self.int_data.shape[:-1] + (2 * self.int_data.shape[-1],))
+        return int_data, self.scale, None
+
+    @classmethod
+    def from_plain(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        assert zero_point is None or torch.all(zero_point == 0)
+        int_data_s4 = ((int_data[..., 1::2] & 0xF) << 4) | (int_data[..., 0::2] & 0xF)
+        return cls(
+            int_data_s4,
+            scale,
+            _layout,
+        )
+
+    def get_layout(self) -> Layout:
+        return self._layout
+
+    def _apply_fn_to_data(self, fn):
+        self.int_data = fn(self.int_data)
+        self.scale = fn(self.scale)
+        return self
+
+
+def _linear_int8_act_int4_weight_cutlass_check(input_tensor, weight_tensor, bias):
+    return (
+        isinstance(input_tensor, AffineQuantizedTensor)
+        and isinstance(input_tensor._layout, PlainLayout)
+        and _aqt_is_int8(input_tensor)
+        and input_tensor.dtype in (torch.float16, torch.bfloat16)
+        and len(input_tensor.shape) >= 2
+        and input_tensor.tensor_impl.scale.dtype == torch.float32
+        and len(input_tensor.tensor_impl.scale.shape) == len(input_tensor.shape) - 1
+        and isinstance(weight_tensor, AffineQuantizedTensor)
+        and isinstance(weight_tensor._layout, CutlassInt4PackedLayout)
+        and _aqt_is_int4(weight_tensor)
+        and weight_tensor.dtype == input_tensor.dtype
+        and len(weight_tensor.shape) == 2
+        and weight_tensor.tensor_impl.scale.dtype == torch.float32
+        and len(weight_tensor.tensor_impl.scale.shape) == 1
+        and (bias is None or bias.dtype == input_tensor.dtype)
+        and (bias is None or len(bias.shape) == 1)
+    )
+
+
+def _linear_int8_act_int4_weight_cutlass_impl(input_tensor, weight_tensor, bias):
+    from torchao.ops import rowwise_scaled_linear_cutlass_s8s4
+
+    weight = weight_tensor.tensor_impl.int_data
+    weight_scale = weight_tensor.tensor_impl.scale
+    input = input_tensor.tensor_impl.int_data
+    input_scale = input_tensor.tensor_impl.scale
+    out_dtype = input_tensor.dtype
+
+    out = rowwise_scaled_linear_cutlass_s8s4(
+        input, input_scale, weight, weight_scale, bias, out_dtype
+    )
+
+    return out
+
+
+def _linear_int4_act_int4_weight_cutlass_check(input_tensor, weight_tensor, bias):
+    return (
+        isinstance(input_tensor, AffineQuantizedTensor)
+        and isinstance(input_tensor._layout, CutlassInt4PackedLayout)
+        and _aqt_is_int4(input_tensor)
+        and input_tensor.dtype in (torch.float16, torch.bfloat16)
+        and len(input_tensor.shape) >= 2
+        and input_tensor.tensor_impl.scale.dtype == torch.float32
+        and len(input_tensor.tensor_impl.scale.shape) == len(input_tensor.shape) - 1
+        and isinstance(weight_tensor, AffineQuantizedTensor)
+        and isinstance(weight_tensor._layout, CutlassInt4PackedLayout)
+        and _aqt_is_int4(weight_tensor)
+        and weight_tensor.dtype == input_tensor.dtype
+        and len(weight_tensor.shape) == 2
+        and weight_tensor.tensor_impl.scale.dtype == torch.float32
+        and len(weight_tensor.tensor_impl.scale.shape) == 1
+    )
+
+
+def _linear_int4_act_int4_weight_cutlass_impl(input_tensor, weight_tensor, bias):
+    from torchao.ops import rowwise_scaled_linear_cutlass_s4s4
+
+    weight = weight_tensor.tensor_impl.int_data
+    weight_scale = weight_tensor.tensor_impl.scale
+    input = input_tensor.tensor_impl.int_data
+    input_scale = input_tensor.tensor_impl.scale
+    out_dtype = input_tensor.dtype
+
+    out = rowwise_scaled_linear_cutlass_s4s4(
+        input, input_scale, weight, weight_scale, bias, out_dtype
+    )
+
+    return out

lib/python3.12/site-packages/torchao/dtypes/uintx/gemlite_layout.py

@@ -0,0 +1,403 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple
+
+import torch
+from torch.utils._python_dispatch import (
+    is_traceable_wrapper_subclass,
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+from torchao.dtypes.uintx.tensor_core_tiled_layout import TensorCoreTiledAQTTensorImpl
+from torchao.dtypes.utils import Layout
+from torchao.utils import fill_defaults
+
+try:
+    import gemlite
+    from gemlite.core import GemLiteLinearTriton
+except:
+    gemlite = None
+
+
+aten = torch.ops.aten
+
+
+def _same_metadata(
+    self: "GemliteAQTTensorImpl",
+    src: "GemliteAQTTensorImpl",
+) -> bool:
+    kwargs_match = len(self.gemlite_kwargs) == len(src.gemlite_kwargs)
+    for k, v in self.gemlite_kwargs.items():
+        if k != "scale_activations":
+            kwargs_match = kwargs_match and (v == src.gemlite_kwargs[k])
+
+    return (
+        isinstance(self, GemliteAQTTensorImpl)
+        and isinstance(src, GemliteAQTTensorImpl)
+        and self.shape == src.shape
+        and self.packed_weight.shape == src.packed_weight.shape
+        and self.scale.shape == src.scale.shape
+        and self.zero_point.shape == src.zero_point.shape
+        and kwargs_match
+        and type(self._layout) == type(src._layout)
+    )
+
+
+def scale_activations_no_scaling(x):
+    return x, None
+
+
+def scale_activations_int8(x):
+    x_shape = x.shape
+    out_x = x.view(-1, x.shape[-1])
+    scaled_x = torch.abs(out_x).amax(axis=1, keepdim=True) / 127
+    out_x = torch.round(out_x / scaled_x).to(dtype=torch.int8)
+    return out_x.view(x_shape), scaled_x
+
+
+def get_gemlite_quant_kwargs(bit_width, group_size, dtype):
+    from torchao.quantization.quant_primitives import MappingType, ZeroPointDomain
+
+    kwargs = {}
+    if bit_width != 8:
+        kwargs["mapping_type"] = MappingType.ASYMMETRIC
+        kwargs["block_size"] = (1, group_size)
+        kwargs["target_dtype"] = torch.uint8
+        kwargs["eps"] = 1e-6
+        kwargs["quant_min"] = 0
+        kwargs["quant_max"] = (2**bit_width) - 1
+        kwargs["eps"] = 1e-6
+        kwargs["zero_point_dtype"] = dtype
+        kwargs["zero_point_domain"] = ZeroPointDomain.FLOAT
+    elif bit_width == 8:
+        kwargs["mapping_type"] = MappingType.SYMMETRIC
+        kwargs["block_size"] = (1, group_size)
+        kwargs["target_dtype"] = torch.int8
+        kwargs["quant_min"] = -128
+        kwargs["quant_max"] = 127
+        kwargs["eps"] = 1e-5
+        kwargs["zero_point_dtype"] = None
+        kwargs["zero_point_domain"] = ZeroPointDomain.NONE
+    return kwargs
+
+
+def get_gemlite_aqt_kwargs(
+    weight,
+    group_size=64,
+    bit_width=4,
+    packing_bitwidth=32,
+    contiguous=None,
+    use_hqq=True,
+):
+    if gemlite is None:
+        raise ImportError(
+            "Unable to import 'gemlite'. Please ensure it is installed correctly. You can install it with: pip install gemlite"
+        )
+
+    assert bit_width in [
+        4,
+        8,
+    ], f"gemlite only works with bit_width 4,8 but got {bit_width}"
+    assert packing_bitwidth in [
+        8,
+        16,
+        32,
+        None,
+    ], f"gemlite needs packing_bitwidth in [8, 16, 32] but got {packing_bitwidth}"
+    assert weight.dtype in [torch.float16, torch.bfloat16], (
+        f"gemlite only works with dtype torch.float16 or torch.bfloat16 but got {weight.dtype}"
+    )
+    assert group_size in [32, 64, 128, 256, 512, 1024, None]
+    assert group_size is None or bit_width != 8, (
+        "gemlite only works with group_size=None for bit_width=8"
+    )
+
+    out_features, in_features = weight.shape
+    group_size = in_features if group_size is None else group_size
+
+    aqt_kwargs = get_gemlite_quant_kwargs(bit_width, group_size, weight.dtype)
+    aqt_kwargs["_layout"] = GemlitePackedLayout(
+        group_size=group_size,
+        bit_width=bit_width,
+        packing_bitwidth=packing_bitwidth,
+        contiguous=contiguous,
+    )
+    aqt_kwargs["use_hqq"] = use_hqq
+    return aqt_kwargs
+
+
+@dataclass(frozen=True)
+class GemlitePackedLayout(Layout):
+    group_size: Optional[int] = 64
+    bit_width: int = 4
+    packing_bitwidth: int = None
+    contiguous: bool = None
+
+
+@register_layout(GemlitePackedLayout)
+class GemliteAQTTensorImpl(TensorCoreTiledAQTTensorImpl):
+    def __new__(
+        cls,
+        packed_weight: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        gemlite_kwargs: Dict,
+        _layout: Layout,
+    ):
+        kwargs = {}
+        kwargs["device"] = packed_weight.device
+        kwargs["layout"] = (
+            kwargs.get("layout")
+            if kwargs.get("layout", False)
+            else packed_weight.layout
+        )
+        kwargs["dtype"] = packed_weight.dtype
+        kwargs["requires_grad"] = False
+        shape = packed_weight.shape
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        packed_weight: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        gemlite_kwargs: Dict,
+        _layout: Layout,
+    ):
+        self.packed_weight = packed_weight
+        self.scale = scale
+        self.zero_point = zero_point
+        self.gemlite_kwargs = gemlite_kwargs
+        self._layout = _layout
+
+    def __tensor_flatten__(self):
+        return ["packed_weight", "scale", "zero_point"], [
+            self._layout,
+            self.gemlite_kwargs,
+        ]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        packed_weight, scale, zero_point = (
+            tensor_data_dict["packed_weight"],
+            tensor_data_dict["scale"],
+            tensor_data_dict["zero_point"],
+        )
+        _layout, gemlite_kwargs = tensor_attributes
+        return cls(packed_weight, scale, zero_point, gemlite_kwargs, _layout)
+
+    @classmethod
+    def from_plain(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        assert isinstance(_layout, GemlitePackedLayout), (
+            f"GemliteAQTTensorImpl only works with GemliteLinearTriton but got {_layout}"
+        )
+        device = int_data.device
+        if device.type != "cuda":
+            int_data = (
+                int_data.cuda()
+            )  # We need int_data on cuda device because of Triton packing
+
+        group_size, bit_width = _layout.group_size, _layout.bit_width
+        out_features, in_features = int_data.shape
+
+        gemlite_linear = gemlite.helper.A16Wn(device=int_data.device).from_weights(
+            int_data, scale, zero_point, bit_width, group_size, bias=None
+        )
+
+        gemlite_kwargs = {
+            "out_features": out_features,
+            "scaled_activations": gemlite_linear.scaled_activations,
+            "meta_args": gemlite_linear.get_meta_args(),
+        }
+
+        packed_weight, scale, zero_point = gemlite_linear.get_tensor_args()
+        packed_weight = packed_weight.to(device)
+
+        return cls(packed_weight, scale, zero_point, gemlite_kwargs, _layout)
+
+    def to(self, *args, **kwargs):
+        kwargs = self._get_to_kwargs(*args, **kwargs)
+        device = kwargs["device"]
+        return self.__class__(
+            self.packed_weight.to(device),
+            self.scale.to(device),
+            self.zero_point.to(device),
+            self.gemlite_kwargs,
+            self._layout,
+        )
+
+    def _apply_fn_to_data(self, fn):
+        return self.__class__(
+            fn(self.packed_weight),
+            fn(self.scale),
+            fn(self.zero_point),
+            self.gemlite_kwargs,
+            self._layout,
+        )
+
+    def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        device = self.packed_weight.device
+        elements_per_sample = self._layout.packing_bitwidth // self._layout.bit_width
+        in_features = (
+            self.packed_weight.numel() * elements_per_sample
+        ) // self.gemlite_kwargs["out_features"]
+        int_data = (
+            gemlite.bitpack.unpack_over_rows(
+                self.packed_weight.cuda(),
+                W_nbits=self._layout.bit_width,
+                num_output_rows=in_features,
+                dtype=torch.uint8,
+            )
+            .t()
+            .contiguous()
+        ).to(device)
+        scale = self.scale.t().contiguous()
+        zero_point = self.zero_point.t().contiguous()
+
+        return int_data, scale, zero_point
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        # we don't handle transpose operations and just ignore them. In practice the only
+        # reason a transpsoe should occur is because the functional linear
+        # op can decompose into e.g. transpose + addmm so since we want
+        # to use the gemlite matmul kernel, which expects teh weight to be passed in as is,
+        # we ignore the transpose
+        if func is aten.detach.default or func is aten.t.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+
+        if func is aten.clone.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.clone)
+            )
+
+        if func is aten.slice.Tensor:
+            self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1])
+            assert step == 1, "Only step == 1 is supported in slicing right now"
+
+            if dim in [0, 1]:
+                int_data, scale, zero_point = self.get_plain()
+                data_len = int_data.shape[dim]
+                scale_len = scale.shape[dim]
+                ratio = data_len / scale_len
+                start_scale = int(start / ratio)
+                end_scale = int(end / ratio)
+
+                int_data = aten.slice.Tensor(int_data, dim, start, end, step)
+                scale = aten.slice.Tensor(scale, dim, start_scale, end_scale, step)
+                if zero_point is not None and zero_point.numel() > 0:
+                    zero_point = aten.slice.Tensor(
+                        zero_point, dim, start_scale, end_scale, step
+                    )
+                else:
+                    zero_point = None
+                # this is to handle padding
+                int_data, scale, zero_point = self._layout.post_process(
+                    int_data, scale, zero_point, self.block_size
+                )
+
+                sliced = self.from_plain(
+                    int_data, scale, zero_point, self._layout
+                )  # Will be transposed again
+
+                return return_and_correct_aliasing(func, args, kwargs, sliced)
+
+            else:
+                raise NotImplementedError(
+                    f"GemliteAQTTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported"
+                )
+
+        elif func is aten.copy_.default:
+            self = args[0]
+            src = args[1]
+            if _same_metadata(self, src):
+                self_tensors = self.__tensor_flatten__()[0]
+                for tensor_name in self_tensors:
+                    getattr(self, tensor_name).copy_(getattr(src, tensor_name))
+                return
+            raise ValueError(
+                f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}"
+            )
+
+        raise NotImplementedError(
+            f"GemliteAQTTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
+    def get_layout(self) -> Layout:
+        return self._layout
+
+    @property
+    def block_size(self):
+        return (1, self._layout.group_size)
+
+
+# logic taken from gemlite's core.py
+def _matmul_type_fn(batch_size: int, bit_width: int) -> str:
+    if batch_size > 64:
+        return "GEMM"
+    elif batch_size > 1:
+        return "GEMM_SPLITK"
+    else:
+        return gemlite.core.get_default_gemv(bit_width)
+
+
+def _linear_fp_act_int4_weight_gemlite_impl(input_tensor, weight_tensor, bias=None):
+    if hasattr(weight_tensor, "tensor_impl"):
+        weight_impl = weight_tensor.tensor_impl
+    else:
+        weight_impl = weight_tensor
+
+    batch_size = input_tensor.view(-1, input_tensor.shape[-1]).shape[0]
+    matmul_type = _matmul_type_fn(batch_size, weight_impl._layout.bit_width)
+
+    if weight_impl.gemlite_kwargs["scaled_activations"]:
+        scale_activations = scale_activations_int8
+    else:
+        scale_activations = scale_activations_no_scaling
+
+    return GemLiteLinearTriton.forward_functional(
+        x=input_tensor,
+        bias=bias,
+        matmul_type=matmul_type,
+        out_features=weight_impl.gemlite_kwargs["out_features"],
+        scale_activations=scale_activations,
+        meta_args=weight_impl.gemlite_kwargs["meta_args"],
+        tensor_args=(
+            weight_impl.packed_weight,
+            weight_impl.scale,
+            weight_impl.zero_point,
+        ),
+    )
+
+
+def _linear_fp_act_int4_weight_gemlite_check(input_tensor, weight_tensor, bias):
+    return (
+        # input is native fp16 tensor
+        not is_traceable_wrapper_subclass(input_tensor)
+        # and input_tensor.dtype in [torch.float16, torch.bfloat16]
+        # weight is gemlite layout
+        and isinstance(weight_tensor, AffineQuantizedTensor)
+        and isinstance(weight_tensor._layout, GemlitePackedLayout)
+    )

lib/python3.12/site-packages/torchao/dtypes/uintx/int4_cpu_layout.py

@@ -0,0 +1,354 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch.utils._python_dispatch import (
+    is_traceable_wrapper_subclass,
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+from torchao.dtypes.utils import AQTTensorImpl, Layout, is_device
+from torchao.quantization.quant_primitives import ZeroPointDomain
+from torchao.utils import (
+    TORCH_VERSION_AT_LEAST_2_5,
+    TORCH_VERSION_AT_LEAST_2_6,
+    fill_defaults,
+)
+
+aten = torch.ops.aten
+
+
+@dataclass(frozen=True)
+class Int4CPULayout(Layout):
+    """Layout class for int4 CPU layout for affine quantized tensor, used by tinygemm kernels `_weight_int4pack_mm_for_cpu`.
+    Only for PyTorch version at least 2.6
+    """
+
+    pass
+
+
+@register_layout(Int4CPULayout)
+class Int4CPUAQTTensorImpl(AQTTensorImpl):
+    """TensorImpl for int4 CPU layout for affine quantized tensor, this is for int4 only,
+    used by tinygemm kernels `_weight_int4pack_mm_for_cpu`
+    It stores the original tensor of dimension [n][k] (int32 dtype) as packed weight of 2-d tensor of
+    dimension: [n][k / 2] (uint8 dtype)
+    (unpacked Tensor shape is n * k)
+    Note: we also pack scale and zero point together here for tinygemm kernel
+    Note: technically Int4 CPU layout should be the layout for the underlying packed weight
+    (int Tensor) but since the scale and zero_point are also packed into the same tensor here which is not used
+    in plain layout, we just created a layout for AQT right now, this could be improved if we split out
+    int4 aqt into a separate tensor subclass
+    fields:
+      packed_weight (torch.Tensor): the 2-d packed tensor in a Int4 CPU layout
+      scale_and_zero (torch.Tensor): the combined scale Tensor used to map between floating point tensor to quantized tensor and zero_point Tensor
+    """
+
+    def __new__(
+        cls,
+        packed_weight: torch.Tensor,
+        scale_and_zero: torch.Tensor,
+        transposed: bool,
+        _layout: Layout,
+    ):
+        kwargs = {}
+        kwargs["device"] = packed_weight.device
+        kwargs["layout"] = (
+            kwargs.get("layout")
+            if kwargs.get("layout", False)
+            else packed_weight.layout
+        )
+        kwargs["dtype"] = packed_weight.dtype
+        kwargs["requires_grad"] = False
+        shape = packed_weight.shape
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        packed_weight: torch.Tensor,
+        scale_and_zero: torch.Tensor,
+        transposed: bool,
+        _layout: Layout,
+    ):
+        self.packed_weight = packed_weight
+        self.scale_and_zero = scale_and_zero
+        self.transposed = False
+        self._layout = _layout
+
+    def __tensor_flatten__(self):
+        return ["packed_weight", "scale_and_zero"], [self.transposed, self._layout]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        packed_weight, scale_and_zero = (
+            tensor_data_dict["packed_weight"],
+            tensor_data_dict["scale_and_zero"],
+        )
+        (
+            transposed,
+            _layout,
+        ) = tensor_attributes
+        return cls(packed_weight, scale_and_zero, transposed, _layout)
+
+    @classmethod
+    def from_plain(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        assert isinstance(_layout, Int4CPULayout)
+
+        if TORCH_VERSION_AT_LEAST_2_6:
+            assert int_data.dtype == torch.int32, (
+                "torch.ops.aten._convert_weight_to_int4pack_for_cpu expects `int32` dtype"
+            )
+            packed_weight = torch.ops.aten._convert_weight_to_int4pack_for_cpu(
+                int_data,
+                1,  # TODO:remove
+            )
+        elif TORCH_VERSION_AT_LEAST_2_5:
+            int_data = (int_data[::, ::2] << 4 | int_data[::, 1::2]).to(torch.uint8)
+            assert int_data.dtype == torch.uint8, (
+                "torch.ops.aten._convert_weight_to_int4pack in torch 2.5 expects `uint8` dtype"
+            )
+            packed_weight = torch.ops.aten._convert_weight_to_int4pack(
+                int_data, _layout.inner_k_tiles
+            )
+        else:
+            assert int_data.dtype == torch.int32, (
+                "torch.ops.aten._convert_weight_to_int4pack in torch 2.4 expects `int32` dtype"
+            )
+            packed_weight = torch.ops.aten._convert_weight_to_int4pack(
+                int_data, _layout.inner_k_tiles
+            )
+
+        scale = scale.reshape(int_data.shape[0], -1)
+        zero_point = zero_point.reshape(int_data.shape[0], -1)
+        from torchao.quantization.utils import pack_tinygemm_scales_and_zeros
+
+        scale_and_zero = pack_tinygemm_scales_and_zeros(scale, zero_point, scale.dtype)
+        return cls(packed_weight, scale_and_zero, False, _layout)
+
+    def to(self, *args, **kwargs):
+        kwargs = self._get_to_kwargs(*args, **kwargs)
+        device = kwargs["device"]
+        if not is_device(torch.device(self.device).type, device):
+            raise ValueError(
+                f"Int4CPUAQTTensorImpl does not support conversion from {self.device} to {device}"
+            )
+        return self.__class__(
+            self.packed_weight.to(device),
+            self.scale_and_zero.to(device),
+            self.transposed,
+            self._layout,
+        )
+
+    def _apply_fn_to_data(self, fn):
+        return self.__class__(
+            fn(self.packed_weight),
+            fn(self.scale_and_zero),
+            self.transposed,
+            self._layout,
+        )
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+
+        if func is aten.clone.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.clone)
+            )
+
+        if func is aten.t.default:
+            """we don't need to repack the weight and just rely on external
+            shape being changed and record the status of transpose/no-transpose
+            """
+            transposed = Int4CPUAQTTensorImpl(
+                args[0].packed_weight,
+                args[0].scale_and_zero,
+                not args[0].transposed,
+                args[0]._layout,
+            )
+            return return_and_correct_aliasing(func, args, kwargs, transposed)
+
+        if func is aten.slice.Tensor:
+            self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1])
+            if dim in [0, 1]:
+                assert step == 1, "Only step == 1 is supported in slicing right now"
+                int_data, scale, zero_point = self.get_plain()
+                data_len = int_data.shape[dim]
+                scale_len = scale.shape[dim]
+                ratio = data_len / scale_len
+                start_scale = int(start / ratio)
+                end_scale = int(end / ratio)
+
+                int_data = aten.slice.Tensor(int_data, dim, start, end, step)
+                scale = aten.slice.Tensor(scale, dim, start_scale, end_scale, step)
+                zero_point = aten.slice.Tensor(
+                    zero_point, dim, start_scale, end_scale, step
+                )
+                # this is to handle padding
+                int_data, scale, zero_point = self._layout.post_process(
+                    int_data, scale, zero_point, self.block_size
+                )
+                sliced = self.from_plain(int_data, scale, zero_point, self._layout)
+                return return_and_correct_aliasing(func, args, kwargs, sliced)
+            else:
+                raise NotImplementedError(
+                    f"Int4CPUAQTTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported"
+                )
+
+        raise NotImplementedError(
+            f"Int4CPUAQTTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
+    @property
+    def block_size(self):
+        from torchao.quantization.utils import unpack_tinygemm_scales_and_zeros
+
+        scale, zero = unpack_tinygemm_scales_and_zeros(self.scale_and_zero)
+        cur_shape = self.shape
+        assert len(cur_shape) == 4
+        inner_k_tiles = cur_shape[-1] * 2
+        original_shape = (cur_shape[0] * 8, cur_shape[1] * (inner_k_tiles * 16))
+        groupsize = int(original_shape[1] / scale.shape[-2])
+        return (1, groupsize)
+
+    def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        from torchao.quantization.quant_primitives import (
+            ZeroPointDomain,
+            quantize_affine,
+        )
+        from torchao.quantization.utils import unpack_tinygemm_scales_and_zeros
+
+        scale, zero = unpack_tinygemm_scales_and_zeros(self.scale_and_zero)
+
+        cur_shape = self.shape
+        assert len(cur_shape) == 2
+        original_shape = (cur_shape[0], cur_shape[1] * 2)
+        eye_shape = original_shape[1]
+        groupsize = int(original_shape[1] / scale.shape[-2])
+        block_size = (1, groupsize)
+        device = self.device
+        original_dtype = self.scale_and_zero.dtype
+        target_dtype = torch.int32
+        quant_min = 0
+        quant_max = 15
+        zero_point_domain = ZeroPointDomain.FLOAT
+        assert len(block_size) == 2 and block_size[0] == 1
+        dequantized = torch.ops.aten._weight_int4pack_mm_for_cpu(
+            torch.eye(eye_shape, device=device, dtype=original_dtype),
+            self.packed_weight,
+            groupsize,
+            self.scale_and_zero,
+        )
+        dequantized = dequantized.t().contiguous()
+        # TODO: move this to `unpack_tinygemm_scales_and_zeros`?
+        scale = scale.reshape(scale.shape[:-1]).contiguous()
+        zero = zero.reshape(zero.shape[:-1]).contiguous()
+        int_data = quantize_affine(
+            dequantized,
+            block_size,
+            scale,
+            zero,
+            target_dtype,
+            quant_min,
+            quant_max,
+            zero_point_domain,
+        )
+        return int_data, scale, zero
+
+    def get_layout(self) -> Layout:
+        return self._layout
+
+
+def _aqt_is_uint4(aqt):
+    """Check if an AffineQuantizedTensor is uint4 quantized Tensor"""
+    return (
+        aqt.tensor_impl.dtype == torch.uint8
+        and aqt.quant_min == 0
+        and aqt.quant_max == 15
+    )
+
+
+def _is_float(dtype):
+    return dtype in (torch.float, torch.half, torch.bfloat16)
+
+
+def _linear_fp_act_uint4_weight_cpu_check(input_tensor, weight_tensor, bias):
+    return (
+        TORCH_VERSION_AT_LEAST_2_6
+        and is_device(input_tensor.device.type, "cpu")
+        and is_device(weight_tensor.device.type, "cpu")
+        and (bias is None or is_device(bias.device.type, "cpu"))
+        and not is_traceable_wrapper_subclass(input_tensor)
+        and _is_float(input_tensor.dtype)
+        and isinstance(weight_tensor, AffineQuantizedTensor)
+        and _aqt_is_uint4(weight_tensor)
+        and _is_float(weight_tensor.dtype)
+        and len(weight_tensor.shape) == 2
+        and weight_tensor.zero_point_domain == ZeroPointDomain.FLOAT
+        and isinstance(weight_tensor._layout, Int4CPULayout)
+    )
+
+
+def _linear_fp_act_uint4_weight_cpu_impl(input_tensor, weight_tensor, bias):
+    assert TORCH_VERSION_AT_LEAST_2_6, (
+        f"Requires PyTorch version at least 2.6, but got: {torch.__version__}"
+    )
+    assert is_device(input_tensor.device.type, "cpu"), (
+        f"For CPU device only but got: {input_tensor.device}"
+    )
+    assert weight_tensor.block_size[0] == 1, (
+        f"Requires groupwise quantization, got block_size: {weight_tensor.block_size}"
+    )
+    assert input_tensor.shape[-1] == weight_tensor.shape[1], (
+        f"need input_tensor shape: {input_tensor.shape} final"
+        f"dim to match weight_tensor shape: {weight_tensor.shape} second dim "
+    )
+
+    act_mat = input_tensor
+    packed_weight = weight_tensor.tensor_impl.packed_weight
+    scale_and_zero = weight_tensor.tensor_impl.scale_and_zero
+
+    orig_act_size = act_mat.size()
+    orig_dtype = act_mat.dtype
+
+    # reshape to 2D
+    act_mat = act_mat.reshape(-1, act_mat.shape[-1])
+
+    # groupwise int4 quantization
+    groupsize = weight_tensor.block_size[1]
+    y = torch.ops.aten._weight_int4pack_mm_for_cpu(
+        act_mat.contiguous(), packed_weight, groupsize, scale_and_zero
+    )
+
+    # remove out_feature padding
+    orig_out_features = weight_tensor.shape[-2]
+    y = y[:, :orig_out_features]
+    y = y.reshape(*orig_act_size[:-1], orig_out_features)
+
+    if bias is not None:
+        y += bias
+    return y.to(orig_dtype)

lib/python3.12/site-packages/torchao/dtypes/uintx/int4_xpu_layout.py

@@ -0,0 +1,445 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch.utils._python_dispatch import (
+    is_traceable_wrapper_subclass,
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+from torchao.dtypes.utils import AQTTensorImpl, Layout, is_device
+from torchao.quantization.quant_primitives import ZeroPointDomain
+from torchao.utils import (
+    TORCH_VERSION_AT_LEAST_2_8,
+    fill_defaults,
+)
+
+aten = torch.ops.aten
+
+
+def _aqt_is_xpu_layout_uint4(aqt):
+    """Check if an AffineQuantizedTensor is uint4 quantized Tensor"""
+    # TODO: use torch.uint4
+    return (
+        aqt.tensor_impl.dtype == torch.int32
+        and aqt.quant_min == 0
+        and aqt.quant_max == 15
+    )
+
+
+def _linear_bf16_act_uint4_weight_float_zero_check(input_tensor, weight_tensor, bias):
+    return (
+        # input is native bfloat16 tensor
+        not is_traceable_wrapper_subclass(input_tensor)
+        and input_tensor.dtype == torch.bfloat16
+        and
+        # weight is uint4, group quantized tensor_core_tiled tensor impl affine quantized tensor
+        isinstance(weight_tensor, AffineQuantizedTensor)
+        and _aqt_is_xpu_layout_uint4(weight_tensor)
+        and weight_tensor.dtype == torch.bfloat16
+        and len(weight_tensor.shape) == 2
+        and weight_tensor.zero_point_domain == ZeroPointDomain.FLOAT
+        and weight_tensor.tensor_impl.scale_and_zero is not None
+        and weight_tensor.tensor_impl.scale_and_zero.dtype == torch.bfloat16
+        and isinstance(weight_tensor._layout, Int4XPULayout)
+    )
+
+
+def _linear_bf16_act_uint4_weight_float_zero_impl(input_tensor, weight_tensor, bias):
+    assert weight_tensor.block_size[0] == 1, (
+        f"Requires groupwise quantization, got block_size: {weight_tensor.block_size}"
+    )
+    assert input_tensor.shape[-1] == weight_tensor.shape[1], (
+        f"need input_tensor shape: {input_tensor.shape} final"
+        f"dim to match weight_tensor shape: {weight_tensor.shape} second dim "
+    )
+
+    # TODO: check groupsize quantization
+    # avoid circular dep, TODO: move this to a common util.py
+    act_mat = input_tensor
+    if act_mat.is_contiguous() == False:
+        act_mat = act_mat.contiguous()
+    # weight is packed from padded (out_features, in_features) weight tensor
+    # (same dimension requirement as F.linear weight)
+    packed_weight = weight_tensor.tensor_impl.packed_weight
+    scales_and_zeros = weight_tensor.tensor_impl.scale_and_zero
+
+    orig_act_size = act_mat.size()
+    orig_dtype = act_mat.dtype
+
+    act_mat = act_mat.reshape(-1, act_mat.shape[-1]).to(torch.bfloat16)
+
+    # groupwise int4 quantization
+    groupsize = weight_tensor.block_size[1]
+    y = torch.ops.aten._weight_int4pack_mm(
+        act_mat, packed_weight, groupsize, scales_and_zeros
+    )
+
+    # remove out_feature padding
+    orig_out_features = weight_tensor.shape[-2]
+    y = y[:, :orig_out_features]
+    y = y.reshape(*orig_act_size[:-1], orig_out_features)
+
+    if bias is not None:
+        y += bias
+    return y.to(orig_dtype)
+
+
+def _linear_bf16_act_uint4_weight_int8_zero_check(input_tensor, weight_tensor, bias):
+    return (
+        # input is native bfloat16 tensor
+        not is_traceable_wrapper_subclass(input_tensor)
+        and input_tensor.dtype == torch.bfloat16
+        and
+        # weight is uint4, group quantized tensor_core_tiled tensor impl affine quantized tensor
+        isinstance(weight_tensor, AffineQuantizedTensor)
+        and _aqt_is_xpu_layout_uint4(weight_tensor)
+        and weight_tensor.dtype == torch.bfloat16
+        and len(weight_tensor.shape) == 2
+        and weight_tensor.zero_point_domain == ZeroPointDomain.INT
+        and weight_tensor.tensor_impl.scale_and_zero is None
+        and weight_tensor.tensor_impl.scale.dtype == torch.bfloat16
+        and weight_tensor.tensor_impl.zero.dtype == torch.int8
+        and isinstance(weight_tensor._layout, Int4XPULayout)
+    )
+
+
+def _linear_bf16_act_uint4_weight_int8_zero_impl(input_tensor, weight_tensor, bias):
+    assert weight_tensor.block_size[0] == 1, (
+        f"Requires groupwise quantization, got block_size: {weight_tensor.block_size}"
+    )
+    assert input_tensor.shape[-1] == weight_tensor.shape[1], (
+        f"need input_tensor shape: {input_tensor.shape} final"
+        f"dim to match weight_tensor shape: {weight_tensor.shape} second dim "
+    )
+
+    # TODO: check groupsize quantization
+    # avoid circular dep, TODO: move this to a common util.py
+    act_mat = input_tensor
+    # weight is packed from padded (out_features, in_features) weight tensor
+    # (same dimension requirement as F.linear weight)
+    packed_weight = weight_tensor.tensor_impl.packed_weight
+    scale = weight_tensor.tensor_impl.scale
+    zero = weight_tensor.tensor_impl.zero
+
+    orig_act_size = act_mat.size()
+    orig_dtype = act_mat.dtype
+
+    act_mat = act_mat.reshape(-1, act_mat.shape[-1]).to(torch.bfloat16)
+
+    # groupwise int4 quantization
+    groupsize = weight_tensor.block_size[1]
+
+    y = torch.ops.aten._weight_int4pack_mm_with_scales_and_zeros(
+        act_mat, packed_weight, groupsize, scale, zero
+    )
+
+    # remove out_feature padding
+    orig_out_features = weight_tensor.shape[-2]
+    y = y[:, :orig_out_features]
+    y = y.reshape(*orig_act_size[:-1], orig_out_features)
+
+    if bias is not None:
+        y += bias
+    return y.to(orig_dtype)
+
+
+@dataclass(frozen=True)
+class Int4XPULayout(Layout):
+    """Only for PyTorch version at least 2.7"""
+
+    pass
+
+
+@register_layout(Int4XPULayout)
+class Int4XPUAQTTensorImpl(AQTTensorImpl):
+    """
+    TensorImpl for int4 XPU layout for affine quantized tensor, this is for int4 only,
+    used by tinygemm kernels `_weight_int4pack_mm_xpu` and `_weight_int4pack_mm_with_zeros_and_scales` (TBD)
+    It stores the original tensor of dimension [n][k] (int32 dtype) as packed weight of 2-d tensor of
+    dimension: [n][k / 8] (int32 dtype)
+    (unpacked Tensor shape is n * k)
+    Note: we also pack scale and zero point together here for tinygemm kernel
+    Note: technically Int4 XPU layout should be the layout for the underlying packed weight
+    (int Tensor) but since the scale and zero_point are also packed into the same tensor here which is not used
+    in plain layout, we just created a layout for AQT right now, this could be improved if we split out
+    int4 aqt into a separate tensor subclass
+    fields:
+      packed_weight (torch.Tensor): the 2-d packed tensor in a Int4 XPU layout
+      [Optional] scale_and_zero (torch.Tensor): the combined scale Tensor used to map between floating point tensor to quantized tensor and zero_point Tensor
+      [Optional] scale (torch.Tensor): scale tensors, should be the same dtype of packed weight
+      [Optional] zeros (torch.Tensor): can be of the same dtype of packed weight or different dtype
+    """
+
+    def __new__(
+        cls,
+        packed_weight: torch.Tensor,
+        scale_and_zero: torch.Tensor,
+        transposed: bool,
+        _layout: Layout,
+        scale: torch.Tensor = None,
+        zero: torch.Tensor = None,
+    ):
+        kwargs = {}
+        kwargs["device"] = packed_weight.device
+        kwargs["layout"] = (
+            kwargs.get("layout")
+            if kwargs.get("layout", False)
+            else packed_weight.layout
+        )
+        kwargs["dtype"] = packed_weight.dtype
+        kwargs["requires_grad"] = False
+        shape = packed_weight.shape
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        packed_weight: torch.Tensor,
+        scale_and_zero: torch.Tensor,
+        transposed: bool,
+        _layout: Layout,
+        scale: torch.Tensor = None,
+        zero: torch.Tensor = None,
+    ):
+        self.packed_weight = packed_weight
+        self.scale_and_zero = scale_and_zero
+        self.transposed = False
+        self._layout = _layout
+        self.scale = scale
+        self.zero = zero
+
+    def __tensor_flatten__(self):
+        if self.scale_and_zero is not None:
+            return ["packed_weight", "scale_and_zero"], [self.transposed, self._layout]
+        else:
+            return ["packed_weight", "scale", "zero"], [self.transposed, self._layout]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        packed_weight = tensor_data_dict["packed_weight"]
+        scale_and_zero = (
+            tensor_data_dict.get("scale_and_zero")
+            if "scale_and_zero" in tensor_data_dict
+            else None
+        )
+        scale = tensor_data_dict.get("scale") if "scale" in tensor_data_dict else None
+        zero = tensor_data_dict.get("zero") if "zero" in tensor_data_dict else None
+        (
+            transposed,
+            _layout,
+        ) = tensor_attributes
+        return cls(packed_weight, scale_and_zero, transposed, _layout, scale, zero)
+
+    @classmethod
+    def from_plain(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        assert isinstance(_layout, Int4XPULayout)
+
+        from torchao.quantization.utils import convert_weight_to_int4pack_xpu
+
+        if TORCH_VERSION_AT_LEAST_2_8:
+            assert int_data.dtype == torch.int32, (
+                "torch.ops.aten._convert_weight_to_int4pack_for_cpu expects `int32` dtype"
+            )
+            packed_weight = convert_weight_to_int4pack_xpu(
+                int_data, zero_point.dtype != scale.dtype
+            )
+        else:
+            assert False, "INT4 not supported on XPU until 2.8"
+
+        scale = scale.reshape(int_data.shape[0], -1)
+        zero_point = zero_point.reshape(int_data.shape[0], -1)
+        if zero_point.dtype == scale.dtype:
+            from torchao.quantization.utils import pack_tinygemm_scales_and_zeros
+
+            scale_and_zero = pack_tinygemm_scales_and_zeros(scale, zero_point)
+            return cls(packed_weight, scale_and_zero, False, _layout, None, None)
+        else:
+            return cls(
+                packed_weight,
+                None,
+                False,
+                _layout,
+                scale.transpose(0, 1).contiguous(),
+                zero_point.transpose(0, 1).contiguous().to(torch.int8),
+            )
+
+    def to(self, *args, **kwargs):
+        kwargs = self._get_to_kwargs(*args, **kwargs)
+        device = kwargs["device"]
+        if not is_device(torch.device(self.device).type, device):
+            raise ValueError(
+                f"Int4XPUAQTTensorImpl does not support conversion from {self.device} to {device}"
+            )
+        return self.__class__(
+            self.packed_weight.to(device),
+            self.scale_and_zero.to(device) if self.scale_and_zero is not None else None,
+            self.transposed,
+            self._layout,
+            self.scale.to(device) if self.scale is not None else None,
+            self.zero.to(device) if self.zero is not None else None,
+        )
+
+    def _apply_fn_to_data(self, fn):
+        return self.__class__(
+            fn(self.packed_weight),
+            fn(self.scale_and_zero) if self.scale_and_zero is not None else None,
+            self.transposed,
+            self._layout,
+            fn(self.scale) if self.scale is not None else None,
+            fn(self.zero) if self.zero is not None else None,
+        )
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+
+        if func is aten.clone.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.clone)
+            )
+
+        if func is aten.t.default:
+            """we don't need to repack the weight and just rely on external
+            shape being changed and record the status of transpose/no-transpose
+            """
+            transposed = Int4XPUAQTTensorImpl(
+                args[0].packed_weight,
+                args[0].scale_and_zero,
+                not args[0].transposed,
+                args[0]._layout,
+                args[0].scale,
+                args[0].zero,
+            )
+            return return_and_correct_aliasing(func, args, kwargs, transposed)
+
+        if func is torch.ops.aten.copy_.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.clone)
+            )
+
+        if func is aten.slice.Tensor:
+            self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1])
+            if dim == 0:
+                int_data, scale, zero_point = self.get_plain()
+                int_data = aten.slice.Tensor(int_data, dim, start, end, step)
+                # this is to handle padding
+                int_data = self._layout.post_process(int_data)
+                sliced = self.from_plain(int_data, scale, zero_point, self._layout)
+                return return_and_correct_aliasing(func, args, kwargs, sliced)
+            elif dim == 1:
+                int_data, scale, zero_point = self.get_plain()
+                assert step == 1, "Only step == 1 is supported in slicing right now"
+                data_len = int_data.shape[dim]
+                scale_len = scale.shape[dim]
+                ratio = data_len / scale_len
+                start_scale = int(start / ratio)
+                end_scale = int(end / ratio)
+
+                int_data = aten.slice.Tensor(int_data, dim, start, end, step)
+                # this is to handle padding
+                int_data = self._layout.post_process(int_data)
+                scale = aten.slice.Tensor(scale, dim, start_scale, end_scale, step)
+                zero_point = aten.slice.Tensor(
+                    zero_point, dim, start_scale, end_scale, step
+                )
+                sliced = self.from_plain(int_data, scale, zero_point, self._layout)
+                return sliced
+            else:
+                raise NotImplementedError(
+                    f"Int4XPUAQTTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported"
+                )
+
+        raise NotImplementedError(
+            f"Int4XPUAQTTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
+    def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        from torchao.quantization.quant_primitives import (
+            ZeroPointDomain,
+            quantize_affine,
+        )
+        from torchao.quantization.utils import unpack_tinygemm_scales_and_zeros
+
+        if self.scale_and_zero is not None:
+            scale, zero = unpack_tinygemm_scales_and_zeros(self.scale_and_zero)
+        else:
+            scale = self.scale.transpose(0, 1).contiguous()
+            zero = self.zero.transpose(0, 1).contiguous()
+
+        cur_shape = self.shape
+        assert len(cur_shape) == 2
+        original_shape = (cur_shape[0], cur_shape[1] * 8)
+        eye_shape = original_shape[1]
+        groupsize = int(original_shape[1] / scale.shape[1])
+        block_size = (1, groupsize)
+        device = self.device
+        original_dtype = torch.bfloat16
+        target_dtype = torch.int32
+        quant_min = 0
+        quant_max = 15
+        assert len(block_size) == 2 and block_size[0] == 1
+        if self.scale_and_zero is None:
+            zero_point_domain = ZeroPointDomain.INT
+            dequantized = torch.ops.aten._weight_int4pack_mm_with_scales_and_zeros(
+                torch.eye(eye_shape, device=device, dtype=original_dtype),
+                self.packed_weight,
+                groupsize,
+                self.scale,
+                self.zero,
+            )
+            dequantized = dequantized.t().contiguous()
+            int_data = quantize_affine(
+                dequantized,
+                block_size,
+                scale,
+                zero,
+                target_dtype,
+                quant_min,
+                quant_max,
+                zero_point_domain,
+            )
+        else:
+            zero_point_domain = ZeroPointDomain.FLOAT
+            dequantized = torch.ops.aten._weight_int4pack_mm(
+                torch.eye(eye_shape, device=device, dtype=original_dtype),
+                self.packed_weight,
+                groupsize,
+                self.scale_and_zero,
+            )
+            dequantized = dequantized.t().contiguous()
+            # TODO: move this to `unpack_tinygemm_scales_and_zeros`?
+            scale = scale.reshape(scale.shape[:-1]).contiguous()
+            zero = zero.reshape(zero.shape[:-1]).contiguous()
+            int_data = quantize_affine(
+                dequantized,
+                block_size,
+                scale,
+                zero,
+                target_dtype,
+                quant_min,
+                quant_max,
+                zero_point_domain,
+            )
+        return int_data, scale, zero
+
+    def get_layout(self) -> Layout:
+        return self._layout

lib/python3.12/site-packages/torchao/dtypes/uintx/marlin_qqq_tensor.py

@@ -0,0 +1,351 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+import logging
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch.utils._python_dispatch import (
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    get_tensor_impl_constructor,
+    register_layout,
+)
+from torchao.dtypes.uintx.plain_layout import (
+    _aqt_is_int8_reduced_range,
+)
+from torchao.dtypes.utils import AQTTensorImpl, Layout
+from torchao.quantization.quant_primitives import (
+    ZeroPointDomain,
+    choose_qparams_and_quantize_affine_qqq,
+    dequantize_affine_qqq,
+)
+
+logger = logging.getLogger(__name__)
+
+aten = torch.ops.aten
+
+
+class MarlinQQQTensor(AffineQuantizedTensor):
+    """MarlinQQQ quantized tensor subclass which inherits AffineQuantizedTensor class.
+
+    To see what happens during choose_qparams_and_quantize_affine_qqq, quantization and dequantization for marlin qqq quantization,
+    please checkout https://github.com/pytorch/ao/blob/main/torchao/quantization/quant_primitives.py
+    and check the two quant primitive ops: choose_qparams_and_quantize_affine_qqq and dequantize_affine_qqq
+    """
+
+    def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor:
+        if output_dtype is None:
+            output_dtype = self.dtype
+
+        int_data, s_group, s_channel = self.tensor_impl.get_plain()
+        nbits = int(math.log2(self.quant_max - self.quant_min + 1))
+        group_size = max(self.block_size)
+        return dequantize_affine_qqq(
+            int_data, s_group, s_channel, nbits, group_size, output_dtype
+        )
+
+    @classmethod
+    def from_hp_to_intx(
+        cls,
+        input_float: torch.Tensor,
+        block_size: Tuple[int, ...],
+        quant_min: Optional[int] = None,
+        quant_max: Optional[int] = None,
+        zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT,
+        _layout: Optional[Layout] = None,
+    ):
+        """Converts a floating point tensor to a Marlin QQQ quantized tensor."""
+        if zero_point_domain is None:
+            raise ValueError("Please use ZeroPointDomain.NONE instead of None")
+        original_shape = input_float.shape
+        input_float = _layout.pre_process(input_float)
+        nbits = int(math.log2(quant_max - quant_min + 1))
+        group_size = max(block_size)
+        data, s_group, s_channel, _ = choose_qparams_and_quantize_affine_qqq(
+            input_float, nbits, group_size
+        )
+        tensor_impl_ctr = get_tensor_impl_constructor(type(_layout))
+        tensor_impl = tensor_impl_ctr(data, s_group, s_channel, _layout)
+        return cls(
+            tensor_impl,
+            block_size,
+            original_shape,
+            quant_min,
+            quant_max,
+            zero_point_domain,
+            dtype=input_float.dtype,
+        )
+
+
+@dataclass(frozen=True)
+class MarlinQQQLayout(Layout):
+    """MarlinQQQLayout is a layout class for Marlin QQQ quantization."""
+
+    pass
+
+
+@register_layout(MarlinQQQLayout)
+class MarlinQQQAQTTensorImpl(AQTTensorImpl):
+    """
+    TensorImpl storage class for sparse_qqq layout for affine quantized tensor.
+
+    Can only be used with 4 bits quantization for now.
+
+    Original marlin documentation and information:
+    https://github.com/IST-DASLab/marlin/tree/master
+
+    Marlin qqq information:
+    https://github.com/HandH1998/QQQ/tree/main
+    https://arxiv.org/pdf/2406.09904
+
+    fields:
+        original_shape (torch.Size): the original shape of the tensor. used to unpack the tensor to the original shape
+        group_size (int): the group size used to pack the tensor
+        num_bits (int): the number of bits used to quantize the tensor
+    """
+
+    @staticmethod
+    def __new__(
+        cls,
+        int_data: torch.Tensor,
+        s_group: torch.Tensor,
+        s_channel: torch.Tensor,
+        _layout: Layout,
+        original_shape: torch.Size,
+        group_size: int,
+        num_bits: int,
+    ):
+        kwargs = {}
+        kwargs["device"] = int_data.device
+        kwargs["layout"] = (
+            kwargs.get("layout") if kwargs.get("layout", False) else int_data.layout
+        )
+        kwargs["dtype"] = int_data.dtype
+        kwargs["requires_grad"] = False
+        shape = int_data.shape
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        int_data: torch.Tensor,
+        s_group: torch.Tensor,
+        s_channel: torch.Tensor,
+        _layout: Layout,
+        original_shape: torch.Size,
+        group_size: int,
+        num_bits: int,
+    ):
+        self.int_data = int_data
+        self.s_group = s_group
+        self.s_channel = s_channel
+        self._layout = _layout
+        self.original_shape = original_shape
+        self.group_size = group_size
+        self.num_bits = num_bits
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+
+        raise NotImplementedError(
+            f"MarlinQQQAQTTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    def __tensor_flatten__(self):
+        return ["int_data", "s_group", "s_channel"], [
+            self._layout,
+            self.original_shape,
+            self.group_size,
+            self.num_bits,
+        ]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        int_data = tensor_data_dict["int_data"]
+        s_group = tensor_data_dict["s_group"]
+        s_channel = tensor_data_dict["s_channel"]
+        _layout, original_shape, group_size, num_bits = tensor_attributes
+        return cls(
+            int_data, s_group, s_channel, _layout, original_shape, group_size, num_bits
+        )
+
+    def get_plain(self):
+        from torchao.quantization.marlin_qqq import (
+            unpack_from_marlin_qqq,
+        )
+
+        int_data_expanded, s_group_expanded, s_channel_expanded = (
+            unpack_from_marlin_qqq(
+                self.int_data,
+                self.s_group,
+                self.s_channel,
+                self.original_shape,
+                self.num_bits,
+                self.group_size,
+            )
+        )
+        int_data_expanded_t = int_data_expanded.t()
+        s_group_expanded_t = s_group_expanded.t()
+        s_channel_expanded_t = s_channel_expanded.t()
+        return int_data_expanded_t, s_group_expanded_t, s_channel_expanded_t
+
+    @classmethod
+    def from_plain(
+        cls,
+        int_data: torch.Tensor,
+        s_group: torch.Tensor,
+        s_channel: torch.Tensor,
+        _layout: Layout,
+    ):
+        from torchao.quantization.marlin_qqq import (
+            const,
+            pack_to_marlin_qqq,
+        )
+
+        assert isinstance(_layout, MarlinQQQLayout)
+
+        # Linear layers are (in_features, out_features) but the int_data that is reaching this point
+        # is (out_features, in_features). We need to transpose it to match the expected shape in the marlin code.
+        q_w = int_data.t()
+        s_group_t = s_group.t()
+        s_channel_t = s_channel.t()
+
+        if not torch.cuda.get_device_capability()[0] >= 8:
+            raise ValueError(
+                f"Can not use Marlin QQQ int4*int8 kernel with a device of compute capability {torch.cuda.get_device_capability()}, the minimum compute capability is 8.0 for Marlin kernel."
+            )
+
+        if q_w.dtype != torch.int32:
+            raise ValueError("Only `torch.int32` weights are supported.")
+
+        in_features, out_features = q_w.shape
+        # (thread_k, thread_n)
+        thread_config = [(64, 256), (128, 128), (128, 64), (64, 128)]
+        if not any(
+            [
+                in_features % thread_k == 0 and out_features % thread_n == 0
+                for thread_k, thread_n in thread_config
+            ]
+        ):
+            raise ValueError(
+                "Not supported `in_features`: {} and `out_features`: {}.".format(
+                    in_features, out_features
+                )
+            )
+
+        num_bits = 4 if torch.max(q_w) - torch.min(q_w) < 16 else -1
+        if num_bits not in [4]:
+            raise ValueError(f"Only {[4]} bits are supported, got {num_bits}.")
+
+        if s_group.numel() == 0:
+            group_size = -1
+        else:
+            group_size = in_features // s_group_t.shape[0]
+        assert group_size <= in_features, (
+            "Group size must be less than or equal to in_features."
+        )
+
+        if group_size not in const.SUPPORTED_GROUP_SIZES:
+            raise ValueError(
+                f"Only {const.SUPPORTED_GROUP_SIZES} group sizes are supported, got {group_size}."
+            )
+
+        # Compress quantized weight to marlin format
+        marlin_qqq_q_w, marlin_qqq_s_group, marlin_qqq_s_channel = pack_to_marlin_qqq(
+            q_w, s_group_t, s_channel_t, num_bits, group_size
+        )
+
+        return cls(
+            marlin_qqq_q_w,
+            marlin_qqq_s_group,
+            marlin_qqq_s_channel,
+            _layout,
+            q_w.shape,
+            group_size,
+            num_bits,
+        )
+
+    def get_layout(self) -> Layout:
+        return self._layout
+
+    def _apply_fn_to_data(self, fn):
+        self.int_data = fn(self.int_data)
+        self.s_group = fn(self.s_group)
+        self.s_channel = fn(self.s_channel)
+        return self
+
+
+def _linear_int8_act_int4_weight_marlin_qqq_check(input_tensor, weight_tensor, bias):
+    return (
+        isinstance(input_tensor, AffineQuantizedTensor)
+        and _aqt_is_int8_reduced_range(input_tensor)
+        and input_tensor.dtype == torch.float16
+        and input_tensor.tensor_impl.scale.dtype == torch.float32
+        and len(input_tensor.tensor_impl.scale.shape) == len(input_tensor.shape) - 1
+        and isinstance(weight_tensor, AffineQuantizedTensor)
+        and weight_tensor.tensor_impl.dtype == torch.int32
+        and len(weight_tensor.shape) == 2
+        and isinstance(weight_tensor._layout, MarlinQQQLayout)
+    )
+
+
+def _linear_int8_act_int4_weight_marlin_qqq_impl(input_tensor, weight_tensor, bias):
+    from torchao.ops import marlin_qqq_gemm
+    from torchao.quantization.marlin_qqq import marlin_qqq_workspace
+
+    assert isinstance(input_tensor, AffineQuantizedTensor)
+    assert isinstance(weight_tensor, AffineQuantizedTensor)
+
+    input = input_tensor.tensor_impl.int_data
+    input_scale = input_tensor.tensor_impl.scale
+
+    w_int4 = weight_tensor.tensor_impl.int_data
+    s_group = weight_tensor.tensor_impl.s_group
+    s_channel = weight_tensor.tensor_impl.s_channel
+    original_shape = weight_tensor.tensor_impl.original_shape
+
+    # Folds batch dimension into the first dimension
+    input_2d = input.view(-1, input.shape[-1])
+    input_scale = input_scale.view(1, -1)
+
+    size_m = input_2d.shape[0]
+    size_n = s_channel.shape[1]
+    size_k = input_2d.shape[1]
+    workspace_qqq = marlin_qqq_workspace(original_shape[1])
+
+    out = marlin_qqq_gemm(
+        input_2d,
+        w_int4,
+        input_scale,
+        s_channel,
+        s_group,
+        workspace_qqq,
+        size_m,
+        size_n,
+        size_k,
+    )
+
+    # Unfold the batch dimension
+    out = out.reshape(input.shape[:-1] + (s_channel.shape[1],))
+
+    if bias is not None:
+        out += bias.to(out.dtype)
+    return out
+
+
+to_marlinqqq_quantized_intx = MarlinQQQTensor.from_hp_to_intx

lib/python3.12/site-packages/torchao/dtypes/uintx/marlin_sparse_layout.py

@@ -0,0 +1,310 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+
+import torch
+from torch.utils._python_dispatch import (
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+from torchao.dtypes.uintx.tensor_core_tiled_layout import _aqt_is_tensor_core_tile_uint4
+from torchao.dtypes.utils import AQTTensorImpl, Layout
+from torchao.quantization.quant_primitives import (
+    ZeroPointDomain,
+)
+
+aten = torch.ops.aten
+
+
+def _linear_fp_act_int4_weight_sparse_marlin_check(input_tensor, weight_tensor, bias):
+    return (
+        isinstance(weight_tensor, AffineQuantizedTensor)
+        and _aqt_is_tensor_core_tile_uint4(weight_tensor)
+        and input_tensor.dtype == torch.float16
+        and len(weight_tensor.shape) == 2
+        and weight_tensor.zero_point_domain == ZeroPointDomain.INT
+        and isinstance(weight_tensor._layout, MarlinSparseLayout)
+    )
+
+
+def _linear_fp_act_int4_weight_sparse_marlin_impl(input_tensor, weight_tensor, bias):
+    from torchao.ops import marlin_24_gemm
+    from torchao.sparsity.marlin import marlin_24_workspace
+
+    assert isinstance(weight_tensor, AffineQuantizedTensor)
+
+    sparse_w_int4 = weight_tensor.tensor_impl.int_data
+    scale = weight_tensor.tensor_impl.scale
+    meta = weight_tensor.tensor_impl.meta
+    original_shape = weight_tensor.tensor_impl.original_shape
+    num_bits = weight_tensor.tensor_impl.num_bits
+
+    # Folds batch dimension into the first dimension
+    input_2d = input_tensor.view(-1, input_tensor.shape[-1])
+
+    size_m = input_2d.shape[0]
+    size_n = scale.shape[1]
+    size_k = input_2d.shape[1]
+    workspace_24 = marlin_24_workspace(original_shape[1])
+
+    out = marlin_24_gemm(
+        input_2d,
+        sparse_w_int4,
+        meta,
+        scale,
+        workspace_24,
+        num_bits,
+        size_m,
+        size_n,
+        size_k,
+    )
+
+    # Unfold the batch dimension
+    out = out.reshape(input_tensor.shape[:-1] + (scale.shape[1],))
+
+    if bias is not None:
+        out += bias.to(out.dtype)
+    return out
+
+
+@dataclass(frozen=True)
+class MarlinSparseLayout(Layout):
+    """MarlinSparseLayout is a layout class for handling sparse tensor formats
+    specifically designed for the Marlin sparse kernel. This layout is used
+    to optimize the storage and computation of affine quantized tensors with
+    2:4 sparsity patterns.
+
+    The layout ensures that the tensor data is pre-processed and stored in a
+    format that is compatible with the Marlin sparse kernel operations. It
+    provides methods for preprocessing input tensors and managing the layout
+    of quantized tensors.
+    """
+
+    def pre_process(self, input: torch.Tensor) -> torch.Tensor:
+        """Preprocess the input tensor to be in the correct format for the Marlin sparse kernel.
+            - 1º: the input tensor is transposed since the linear layer keeps the weights in a transposed format
+            - 2º: tensor is injected with 2:4 sparsity
+            - 3º: transposes it again because the quantization process will compute the scales for dim=-1
+
+        Args:
+            input (torch.Tensor): the input tensor to preprocess
+
+        Returns:
+            torch.Tensor: the preprocessed tensor
+        """
+        from torchao.sparsity.marlin import inject_24  # avoid circular import
+
+        input_t = input.t()
+        w_24, _ = inject_24(input_t, *input_t.shape)
+        return w_24.t()
+
+
+@register_layout(MarlinSparseLayout)
+class MarlinSparseAQTTensorImpl(AQTTensorImpl):
+    """
+    TensorImpl for sparse_marlin_24 layout for affine quantized tensor.
+
+    Can be used with 4 bits and 8 bits quantization.
+
+    Original marlin documentation and information:
+    https://github.com/IST-DASLab/marlin/tree/master
+
+    Sparse marlin documentation and information:
+    https://github.com/IST-DASLab/Sparse-Marlin?tab=readme-ov-file
+
+    fields:
+        original_shape (torch.Size): the original shape of the tensor. used to unpack the tensor to the original shape
+        group_size (int): the group size used to pack the tensor
+        num_bits (int): the number of bits used to quantize the tensor
+    """
+
+    @staticmethod
+    def __new__(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        meta: torch.Tensor,
+        _layout: Layout,
+        original_shape: torch.Size,
+        group_size: int,
+        num_bits: int,
+    ):
+        kwargs = {}
+        kwargs["device"] = int_data.device
+        kwargs["layout"] = (
+            kwargs.get("layout") if kwargs.get("layout", False) else int_data.layout
+        )
+        kwargs["dtype"] = int_data.dtype
+        kwargs["requires_grad"] = False
+        shape = int_data.shape
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        meta: torch.Tensor,
+        _layout: Layout,
+        original_shape: torch.Size,
+        group_size: int,
+        num_bits: int,
+    ):
+        self.int_data = int_data
+        self.scale = scale
+        self.zero_point = zero_point
+        self.meta = meta
+        self._layout = _layout
+        self.original_shape = original_shape
+        self.group_size = group_size
+        self.num_bits = num_bits
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+
+        raise NotImplementedError(
+            f"MarlinSparseAQTTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    def __tensor_flatten__(self):
+        return ["int_data", "scale", "zero_point", "meta"], [
+            self._layout,
+            self.original_shape,
+            self.group_size,
+            self.num_bits,
+        ]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        int_data = tensor_data_dict["int_data"]
+        scale = tensor_data_dict["scale"]
+        zero_point = tensor_data_dict["zero_point"]
+        meta = tensor_data_dict["meta"]
+        _layout, original_shape, group_size, num_bits = tensor_attributes
+        return cls(
+            int_data,
+            scale,
+            zero_point,
+            meta,
+            _layout,
+            original_shape,
+            group_size,
+            num_bits,
+        )
+
+    def get_plain(self):
+        from torchao.sparsity.marlin import (
+            unpack_from_marlin_24,
+        )
+
+        int_data_expanded, scales_expanded = unpack_from_marlin_24(
+            self.int_data,
+            self.scale,
+            self.meta,
+            self.original_shape,
+            self.group_size,
+            self.num_bits,
+        )
+        int_data_expanded_t = int_data_expanded.t()
+        scales_expanded_t = scales_expanded.t()
+        return int_data_expanded_t, scales_expanded_t, self.zero_point
+
+    @classmethod
+    def from_plain(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        _layout: Layout,
+    ):
+        from torchao.sparsity.marlin import (
+            const,
+            pack_to_marlin_24,
+        )
+
+        assert isinstance(_layout, MarlinSparseLayout)
+
+        # Linear layers are (in_features, out_features) but the int_data that is reaching this point
+        # is (out_features, in_features). We need to transpose it to match the expected shape in the marlin code.
+        q_w_24 = int_data.t()
+        # addressing the case when scale has dimension 1, happens when
+        # weight_shape[-1] == group_size == 128
+        if scale.ndim == 1:
+            scale = scale.reshape(scale.shape[0], -1)
+
+        scale_t = scale.t()
+
+        if not torch.cuda.get_device_capability()[0] >= 8:
+            raise ValueError(
+                f"Can not use Sparse Marlin 2:4 int4*fp16 kernel with a device of compute capability {torch.cuda.get_device_capability()}, the minimum compute capability is 8.0 for Marlin kernel."
+            )
+
+        if q_w_24.dtype != torch.int32:
+            raise ValueError("Only `torch.int32` weights are supported.")
+
+        in_features, out_features = q_w_24.shape
+        if in_features % 128 != 0 or out_features != 256 == 0:
+            raise ValueError(
+                "`in_features` must be divisible by 64 and `out_features` by 256."
+            )
+
+        # NOTE: The current marlin 2:4 kernel supports both 4 and 8 bits quantization but fp8
+        # will require a bit more work to get our current quantization flow to work with it.
+        # Check the link for a reference: https://github.com/neuralmagic/nm-vllm/tree/main
+        num_bits = 4 if torch.max(q_w_24) < 16 else -1
+        if num_bits not in [4]:
+            raise ValueError(f"Only {[4]} bits are supported, got {num_bits}.")
+
+        group_size = in_features // scale_t.shape[0]
+        if group_size == 0:
+            group_size = in_features
+        assert group_size <= in_features, (
+            "Group size must be less than or equal to in_features."
+        )
+
+        if group_size not in const.SUPPORTED_GROUP_SIZES:
+            raise ValueError(
+                f"Only {const.SUPPORTED_GROUP_SIZES} group sizes are supported, got {group_size}."
+            )
+
+        # Compress quantized weight to marlin 2:4 format
+        marlin_24_q_w_comp, marlin_24_s, meta = pack_to_marlin_24(
+            q_w_24, scale_t, num_bits, group_size
+        )
+
+        return cls(
+            marlin_24_q_w_comp,
+            marlin_24_s,
+            zero_point,
+            meta,
+            _layout,
+            q_w_24.shape,
+            group_size,
+            num_bits,
+        )
+
+    def get_layout(self) -> Layout:
+        return self._layout
+
+    def _apply_fn_to_data(self, fn):
+        self.int_data = fn(self.int_data)
+        self.scale = fn(self.scale)
+        self.zero_point = fn(self.zero_point)
+        self.meta = fn(self.meta)
+        return self

lib/python3.12/site-packages/torchao/dtypes/uintx/packed_linear_int8_dynamic_activation_intx_weight_layout.py

@@ -0,0 +1,467 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from enum import Enum, auto
+from typing import Optional, Tuple, Union
+
+import torch
+from torch.utils._python_dispatch import return_and_correct_aliasing
+
+from torchao.dtypes.affine_quantized_tensor import register_layout
+from torchao.dtypes.utils import AQTTensorImpl, Layout
+from torchao.experimental.op_lib_utils import _check_torchao_ops_loaded
+from torchao.quantization.quant_primitives import (
+    _DTYPE_TO_BIT_WIDTH,
+    _DTYPE_TO_QVALUE_BOUNDS,
+    ZeroPointDomain,
+)
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_6
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+import sys
+
+handler = logging.StreamHandler(sys.stdout)
+formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+handler.setFormatter(formatter)
+logger.addHandler(handler)
+
+
+class Target(Enum):
+    """Enum that indicates the backend target"""
+
+    # AUTO target will automatically select a packing format
+    # based on the available hardware.
+    AUTO = auto()
+    UNIVERSAL = auto()
+    KLEIDIAI = auto()
+
+    # ATEN target will use the ATen operator
+    ATEN = auto()
+
+
+_TARGET_AND_STR = [
+    (Target.AUTO, "auto"),
+    (Target.ATEN, "aten"),
+    (Target.UNIVERSAL, "universal"),
+    (Target.KLEIDIAI, "kleidiai"),
+]
+
+
+def target_to_str(target: Target) -> str:
+    target_to_str = {t: s for t, s in _TARGET_AND_STR}
+    return target_to_str[target]
+
+
+def target_from_str(target: str) -> Target:
+    str_to_target = {s: t for t, s in _TARGET_AND_STR}
+    if target.lower() in str_to_target:
+        return str_to_target[target.lower()]
+    raise ValueError(f"Invalid target: {target}")
+
+
+class PackedLinearInt8DynamicActivationIntxWeightLayout(Layout):
+    bit_width: Optional[int]
+    group_size: Optional[int]
+    has_weight_zeros: Optional[bool]
+    has_bias: Optional[bool]
+    target: Optional[Target]
+
+    def __init__(
+        self,
+        target: Union[str, Target] = "auto",
+    ):
+        if isinstance(target, str):
+            target = target_from_str(target)
+        self.target = target
+
+        self.bit_width: Optional[int] = None
+        self.group_size: Optional[int] = None
+        self.has_weight_zeros: Optional[bool] = None
+        # has_bias is whether the packed weights
+        # have bias packed with them, not whether the
+        # linear operator has bias
+        self.has_bias: Optional[bool] = None
+
+    def extra_repr(self):
+        return f"group_size={self.group_size}, bit_width={self.bit_width}, has_weight_zeros={self.has_weight_zeros}, has_bias={self.has_bias}, target={self.target}"
+
+    def has_params_set(self) -> bool:
+        return (
+            (self.bit_width is not None)
+            and (self.group_size is not None)
+            and (self.has_weight_zeros is not None)
+            and (self.has_bias is not None)
+            and (self.target is not None)
+        )
+
+    def set_params(
+        self, bit_width: int, group_size: int, has_weight_zeros: bool, has_bias: bool
+    ):
+        assert bit_width >= 1 and bit_width <= 8, "bit_width must be 1 to 8"
+        assert group_size >= 1, f"group_size must be positive, got {group_size}"
+
+        self.bit_width = bit_width
+        self.group_size = group_size
+        self.has_weight_zeros = has_weight_zeros
+        self.has_bias = has_bias
+        assert self.has_params_set()
+
+
+@register_layout(PackedLinearInt8DynamicActivationIntxWeightLayout)
+class PackedLinearInt8DynamicActivationIntxWeightAQTTensorImpl(AQTTensorImpl):
+    def __new__(
+        cls,
+        packed_weight: torch.Tensor,
+        _layout: Layout,
+    ):
+        kwargs = {}
+        kwargs["device"] = packed_weight.device
+        kwargs["dtype"] = packed_weight.dtype
+        assert not packed_weight.requires_grad
+        kwargs["requires_grad"] = False
+        shape = packed_weight.shape
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        packed_weight: torch.Tensor,
+        _layout: Layout,
+    ):
+        assert isinstance(_layout, PackedLinearInt8DynamicActivationIntxWeightLayout)
+        self.packed_weight = packed_weight
+        self._layout = _layout
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(packed_weight={str(self.packed_weight)}, layout={self.get_layout()})"
+
+    def get_layout(self) -> Layout:
+        return self._layout
+
+    def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        raise NotImplementedError(
+            "get_plain is not implemented for PackedLinearInt8DynamicActivationIntxWeightAQTTensorImpl"
+        )
+
+    @classmethod
+    def from_plain(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        layout: Layout,
+        bias: Optional[torch.Tensor] = None,
+        *,
+        validate_inputs: bool = True,
+    ):
+        assert isinstance(layout, PackedLinearInt8DynamicActivationIntxWeightLayout)
+        assert layout.target in [t for t, _ in _TARGET_AND_STR], (
+            f"Unexpected target: {layout.target}"
+        )
+        assert layout.has_params_set(), (
+            "PackedLinearInt8DynamicActivationIntxWeightLayout params must be set before calling from_plain"
+        )
+
+        if layout.target != Target.ATEN:
+            _check_torchao_ops_loaded()
+        else:
+            assert TORCH_VERSION_AT_LEAST_2_6, (
+                "aten target is requires torch version > 2.6.0"
+            )
+            assert torch.backends.kleidiai.is_available(), (
+                "ATEN target requires torch.backends.kleidiai.is_available()"
+            )
+            layout.bit_width == 4, "ATEN target only supports torch.int4"
+            assert not layout.has_weight_zeros, "ATEN target does not support zeros"
+
+        data_dtype = getattr(torch, f"int{layout.bit_width}")
+        qmin, qmax = _DTYPE_TO_QVALUE_BOUNDS[data_dtype]
+
+        int_types = [torch.int8, torch.int16, torch.int32, torch.int64]
+
+        # Check int_data
+        assert int_data.device == torch.device("cpu")
+        assert int_data.dtype in int_types
+        n, k = int_data.shape
+        assert k % layout.group_size == 0, "k must be divisible by group_size"
+        if validate_inputs:
+            assert int_data.min().item() >= qmin
+            assert int_data.max().item() <= qmax
+        int_data = int_data.to(torch.int8)
+
+        # Check scale
+        assert scale.device == torch.device("cpu")
+        if scale.dtype != torch.float32:
+            logging.info(f"scale has dtype {scale.dtype}, converting to torch.float32")
+            scale = scale.to(torch.float32)
+        n_, _ = scale.shape
+        assert n_ == n
+        assert scale.numel() * layout.group_size == int_data.numel(), (
+            "must have 1 scale per group"
+        )
+        if validate_inputs:
+            assert scale.min().item() > 0
+            # Some targets round scales to bfloat16, give warning if scales are at higher precision
+            scale_is_rounded_to_bf16 = torch.allclose(
+                scale, scale.to(torch.bfloat16).to(torch.float32)
+            )
+            if not scale_is_rounded_to_bf16:
+                if layout.target == Target.ATEN and (layout.group_size < k):
+                    logging.warning(
+                        "When using Target.ATEN with group_size < k, scales will be rounded to bfloat16"
+                    )
+                if layout.target in [Target.AUTO, Target.KLEIDIAI]:
+                    logging.warning(
+                        "When using [Target.AUTO, Target.KLEIDIAI], scales will be rounded to bfloat16"
+                    )
+
+        # Check zero_point
+        if zero_point is None:
+            assert not layout.has_weight_zeros, (
+                "zero_point must be provided if has_weight_zeros=True"
+            )
+        else:
+            assert zero_point.device == torch.device("cpu")
+            assert zero_point.shape == scale.shape
+            assert zero_point.dtype in int_types
+            assert zero_point.numel() * layout.group_size == int_data.numel(), (
+                "must have 1 zero_point per group"
+            )
+            if validate_inputs:
+                zero_point_min = zero_point.min().item()
+                zero_point_max = zero_point.max().item()
+                assert zero_point.min().item() >= qmin
+                assert zero_point.max().item() <= qmax
+                has_weight_zeros = True
+                if zero_point_min == 0 and zero_point_max == 0:
+                    has_weight_zeros = False
+                assert has_weight_zeros == layout.has_weight_zeros, (
+                    "zero_point being all zeros must be consistent with layout.has_weight_zeros"
+                )
+            zero_point = zero_point.to(torch.int8)
+
+        # Check bias
+        has_bias = bias is not None
+        assert has_bias == layout.has_bias, (
+            "bias being None must be consistent with layout.has_bias"
+        )
+        if has_bias:
+            assert bias.device == torch.device("cpu")
+            if bias.dtype != torch.float32:
+                logging.info(
+                    f"bias has dtype {bias.dtype}, converting to torch.float32"
+                )
+                bias = bias.to(torch.float32)
+            assert bias.shape == (n,)
+
+        # Construct packed_weight
+        if layout.target == Target.ATEN:
+            int_data = int_data.add(8)
+            int_data = (int_data[::, 1::2] << 4 | int_data[::, ::2]).to(torch.uint8)
+
+            # If group_size < k, convert scales to bfloat16
+            # to call optimized kernel
+            if layout.group_size < k:
+                scale = scale.to(torch.bfloat16)
+            packed_weight = torch.ops.aten._dyn_quant_pack_4bit_weight(
+                int_data, scale, bias, layout.group_size, k, n
+            )
+            return cls(packed_weight, layout)
+
+        args = [
+            int_data,
+            scale.reshape(-1),
+            zero_point.reshape(-1) if layout.has_weight_zeros else None,
+            layout.group_size,
+            bias,
+            target_to_str(layout.target) if layout.target != Target.AUTO else None,
+        ]
+        packed_weight = getattr(
+            torch.ops.torchao,
+            f"_pack_8bit_act_{layout.bit_width}bit_weight",
+        )(*args)
+
+        return cls(packed_weight, layout)
+
+    def _apply_fn_to_data(self, fn):
+        self.packed_weight = fn(self.packed_weight)
+        return self
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is torch.ops.aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+        if func is torch.ops.aten.clone.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.clone)
+            )
+
+        raise NotImplementedError(
+            f"PackedLinearInt8DynamicActivationIntxWeightAQTTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    def __tensor_flatten__(self):
+        return ["packed_weight"], [self.get_layout()]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        packed_weight = tensor_data_dict["packed_weight"]
+        (layout,) = tensor_attributes
+        return cls(packed_weight, layout)
+
+
+def _linear_check(input_tensor, weight_tensor, bias):
+    layout = weight_tensor.tensor_impl.get_layout()
+    return isinstance(layout, PackedLinearInt8DynamicActivationIntxWeightLayout)
+
+
+def _linear_impl(input_tensor, weight_tensor, bias):
+    def _impl_2d_non_aten(input_tensor, weight_tensor):
+        assert input_tensor.dim() == 2
+        assert weight_tensor.dim() == 2
+
+        m, k = input_tensor.shape
+        n, k_ = weight_tensor.shape
+        assert k_ == k
+        group_size = weight_tensor.tensor_impl.get_layout().group_size
+
+        args = (
+            input_tensor,
+            weight_tensor.tensor_impl.packed_weight,
+            group_size,
+            n,
+            k,
+        )
+
+        assert len(weight_tensor.block_size) == 2
+        assert weight_tensor.block_size[0] == 1
+        assert group_size == weight_tensor.block_size[1]
+        bit_width = weight_tensor.tensor_impl.get_layout().bit_width
+
+        return getattr(torch.ops.torchao, f"_linear_8bit_act_{bit_width}bit_weight")(
+            *args
+        )
+
+    def _impl_2d_aten(input_tensor, weight_tensor):
+        assert input_tensor.dim() == 2
+        assert weight_tensor.dim() == 2
+
+        m, k = input_tensor.shape
+        n, k_ = weight_tensor.shape
+
+        assert k_ == k
+        group_size = weight_tensor.tensor_impl.get_layout().group_size
+        packed_weight = weight_tensor.tensor_impl.packed_weight
+        return torch.ops.aten._dyn_quant_matmul_4bit(
+            input_tensor, packed_weight, group_size, k, n
+        )
+
+    if input_tensor.numel() == 0:
+        return input_tensor
+
+    target = weight_tensor.tensor_impl.get_layout().target
+
+    if weight_tensor.tensor_impl.get_layout().has_bias:
+        assert bias is None, (
+            "bias should be None because it is already packed with the weights (has_bias=True)"
+        )
+
+    if target == Target.ATEN:
+        assert TORCH_VERSION_AT_LEAST_2_6 == 1, "Target.ATEN requires torch >= 2.6.0"
+        _impl_2d = _impl_2d_aten
+    else:
+        _impl_2d = _impl_2d_non_aten
+
+    if input_tensor.dim() == 2:
+        res = _impl_2d(input_tensor, weight_tensor)
+    else:
+        assert input_tensor.dim() >= 3
+        lead_shape = input_tensor.shape[0:-2]
+        m, k = input_tensor.shape[-2], input_tensor.shape[-1]
+        n, k_ = weight_tensor.shape
+        assert k_ == k
+
+        res = _impl_2d(input_tensor.reshape(-1, k), weight_tensor)
+        res = res.reshape(*lead_shape, m, n)
+
+    if bias is not None:
+        res = res + bias
+    return res
+
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+)
+from torchao.dtypes.utils import AQTTensorImpl, Layout
+
+
+def make_packed_linear_int8_dynamic_activation_intx_weight_tensor(
+    int_data: torch.Tensor,
+    scale: torch.Tensor,
+    zero_point: Optional[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    data_dtype: torch.dtype,
+    target: Union[str, Target] = "auto",
+    *,
+    validate_inputs: bool = True,
+) -> AffineQuantizedTensor:
+    """
+    Constructs an AffineQuantizedTensor with PackedLinearInt8DynamicActivationIntxWeightLayout
+    from plain data.
+    """
+    # TORCH_VERSION_AT_LEAST_2_6 is needed for torch.intx with x < 8
+    assert TORCH_VERSION_AT_LEAST_2_6, (
+        "Using PackedLinearInt8DynamicActivationIntxWeightLayout requires torch version > 2.6.0"
+    )
+
+    layout = PackedLinearInt8DynamicActivationIntxWeightLayout(target=target)
+
+    bit_width = _DTYPE_TO_BIT_WIDTH[data_dtype]
+    qmin, qmax = _DTYPE_TO_QVALUE_BOUNDS[data_dtype]
+
+    n, k = int_data.shape
+    n_, groups_per_k = scale.shape
+    assert k % groups_per_k == 0
+    group_size = k // groups_per_k
+
+    has_weight_zeros = True
+    if zero_point is None:
+        has_weight_zeros = False
+    else:
+        zero_point_min = zero_point.min().item()
+        zero_point_max = zero_point.max().item()
+        if zero_point_min == 0 and zero_point_max == 0:
+            has_weight_zeros = False
+
+    has_bias = bias is not None
+
+    layout.set_params(bit_width, group_size, has_weight_zeros, has_bias)
+    assert layout.has_params_set()
+    tensor_impl = PackedLinearInt8DynamicActivationIntxWeightAQTTensorImpl.from_plain(
+        int_data,
+        scale,
+        zero_point,
+        layout,
+        bias,
+        validate_inputs=validate_inputs,
+    )
+
+    return AffineQuantizedTensor(
+        tensor_impl,
+        block_size=(1, group_size),
+        shape=int_data.shape,
+        quant_min=qmin,
+        quant_max=qmax,
+        zero_point_domain=ZeroPointDomain.INT,
+    )

lib/python3.12/site-packages/torchao/dtypes/uintx/plain_layout.py

@@ -0,0 +1,315 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+from typing import Optional, Tuple
+
+import torch
+from torch.utils._python_dispatch import (
+    is_traceable_wrapper_subclass,
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+from torchao.dtypes.utils import AQTTensorImpl, Layout, PlainLayout
+from torchao.kernel import (
+    int_scaled_matmul,
+)
+from torchao.quantization.quant_primitives import (
+    ZeroPointDomain,
+)
+from torchao.utils import fill_defaults
+
+aten = torch.ops.aten
+
+
+def _same_metadata(self: "PlainAQTTensorImpl", src: "PlainAQTTensorImpl") -> bool:
+    return (
+        isinstance(self, PlainAQTTensorImpl)
+        and isinstance(src, PlainAQTTensorImpl)
+        and self.shape == src.shape
+        and self.int_data.shape == src.int_data.shape
+        and self.scale.shape == src.scale.shape
+        and (self.zero_point is None and src.zero_point is None)
+        or (
+            self.zero_point is not None
+            and src.zero_point is not None
+            and self.zero_point.shape == src.zero_point.shape
+        )
+        and type(self._layout) == type(src._layout)
+    )
+
+
+@register_layout(PlainLayout)
+class PlainAQTTensorImpl(AQTTensorImpl):
+    """
+    TensorImpl for plain layout for affine quantized tensor, it stores int_data, scale, zero_point
+    tensors directly as plain tensors.
+
+    fields:
+      int_data (torch.Tensor): the quantized integer data Tensor
+      scale (torch.Tensor): the scale Tensor used to map between floating point tensor to quantized tensor
+      zero_point (torch.Tensor): the zero_point Tensor used to map between floating point tensor to quantized tensor
+    """
+
+    def __new__(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        kwargs = {}
+        kwargs["device"] = int_data.device
+        kwargs["layout"] = (
+            kwargs.get("layout") if kwargs.get("layout", False) else int_data.layout
+        )
+        kwargs["dtype"] = int_data.dtype
+        kwargs["requires_grad"] = False
+        shape = int_data.shape
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        self.int_data = int_data
+        self.scale = scale
+        self.zero_point = zero_point
+        self._layout = _layout
+
+    def __tensor_flatten__(self):
+        if self.zero_point is None:
+            return ["int_data", "scale"], [self._layout]
+        return ["int_data", "scale", "zero_point"], [self._layout]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        int_data, scale, zero_point = (
+            tensor_data_dict["int_data"],
+            tensor_data_dict["scale"],
+            tensor_data_dict.get("zero_point", None),
+        )
+        (_layout,) = tensor_attributes
+        return cls(int_data, scale, zero_point, _layout)
+
+    def to(self, *args, **kwargs):
+        kwargs = self._get_to_kwargs(*args, **kwargs)
+        return self.__class__(
+            self.int_data.to(kwargs["device"]),
+            self.scale.to(kwargs["device"]),
+            self.zero_point.to(kwargs["device"])
+            if self.zero_point is not None
+            else None,
+            self._layout,
+        )
+
+    def _apply_fn_to_data(self, fn):
+        return self.__class__(
+            fn(self.int_data),
+            fn(self.scale),
+            fn(self.zero_point) if self.zero_point is not None else None,
+            self._layout,
+        )
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+
+        elif func is aten.clone.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.clone)
+            )
+
+        elif func is aten.copy_.default:
+            self = args[0]
+            src = args[1]
+            if _same_metadata(self, src):
+                self_tensors = self.__tensor_flatten__()[0]
+                for tensor_name in self_tensors:
+                    getattr(self, tensor_name).copy_(getattr(src, tensor_name))
+                return
+            raise ValueError(
+                f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}"
+            )
+
+        elif func is aten.t.default:
+            tensor = args[0]
+            new = tensor.__class__(
+                tensor.int_data.t(), tensor.scale, tensor.zero_point, tensor._layout
+            )
+            return return_and_correct_aliasing(func, args, kwargs, new)
+
+        elif func in [aten.select.int, aten.index.Tensor]:
+            return return_and_correct_aliasing(
+                func,
+                args,
+                kwargs,
+                args[0]._apply_fn_to_data(lambda x: func(x, *args[1:], **kwargs)),
+            )
+
+        elif func is aten.slice.Tensor:
+            self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1])
+            if dim == 0:
+                return return_and_correct_aliasing(
+                    func,
+                    args,
+                    kwargs,
+                    args[0]._apply_fn_to_data(
+                        lambda x: aten.slice.Tensor(x, dim, start, end, step)
+                    ),
+                )
+            elif dim == 1:
+                assert len(self.scale.shape) == 1, (
+                    f"slice dim==1 only works when len(scale.shape) == 1 currently, got: {self.scale.shape}"
+                )
+                return PlainAQTTensorImpl(
+                    aten.slice.Tensor(self.int_data, dim, start, end, step),
+                    self.scale.view(-1),
+                    self.zero_point.view(-1) if self.zero_point is not None else None,
+                    self._layout,
+                )
+            else:
+                raise NotImplementedError(
+                    f"PlainAQTTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported"
+                )
+
+        raise NotImplementedError(
+            f"PlainAQTTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
+    def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        return self.int_data, self.scale, self.zero_point
+
+    def get_layout(self) -> Layout:
+        return self._layout
+
+    @classmethod
+    def from_plain(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        assert isinstance(_layout, PlainLayout)
+        return cls(int_data, scale, zero_point, _layout)
+
+
+def _aqt_is_int8(aqt):
+    """Check if an AffineQuantizedTensor is int8 quantized Tensor"""
+    return (
+        aqt.tensor_impl.dtype == torch.int8
+        and (aqt.quant_min is None or aqt.quant_min == -128)
+        and (aqt.quant_max is None or aqt.quant_max == 127)
+    )
+
+
+def _aqt_is_int8_reduced_range(aqt):
+    return (
+        aqt.tensor_impl.dtype == torch.int8
+        and aqt.quant_min == -127
+        and (aqt.quant_max is None or aqt.quant_max == 127)
+    )
+
+
+def _linear_fp_act_int8_weight_check(input_tensor, weight_tensor, bias):
+    return (
+        # input is native float tensor
+        not is_traceable_wrapper_subclass(input_tensor)
+        and input_tensor.is_floating_point()
+        and
+        # weight is int8 per channel quantized affine quantized tensor
+        isinstance(weight_tensor, AffineQuantizedTensor)
+        and _aqt_is_int8(weight_tensor)
+        and len(weight_tensor.shape) == 2
+        and len(weight_tensor.block_size) == 2
+        and weight_tensor.block_size[0] == 1
+        and weight_tensor.block_size[1] == weight_tensor.shape[1]
+        and weight_tensor.zero_point_domain == ZeroPointDomain.INT
+        and isinstance(weight_tensor._layout, PlainLayout)
+    )
+
+
+def _linear_fp_act_int8_weight_impl(input_tensor, weight_tensor, bias):
+    # TODO: enable cpu and mps efficient path
+    # is_cpu and is_mps only, some issue with is_contiguous() currently
+    # return torch.ops.aten._weight_int8pack_mm(input_tensor.contiguous(), w_vals_int8_t, weight_tensor.tensor_impl.scale)
+
+    # per channel int8 weight only quantizated mm
+    w_vals_int8_t = weight_tensor.tensor_impl.int_data.t()
+    scale = weight_tensor.tensor_impl.scale
+    m = torch.mm(
+        input_tensor.reshape(-1, input_tensor.shape[-1]),
+        w_vals_int8_t.to(input_tensor.dtype),
+    )
+    y = m * scale.to(m.dtype)
+    y = y.reshape(*input_tensor.shape[:-1], y.shape[-1])
+    if bias is not None:
+        y += bias.to(m.dtype)
+    return y
+
+
+def _linear_int8_act_int8_weight_check(input_tensor, weight_tensor, bias):
+    return (
+        isinstance(input_tensor, AffineQuantizedTensor)
+        and _aqt_is_int8_reduced_range(input_tensor)
+        and isinstance(weight_tensor, AffineQuantizedTensor)
+        and _aqt_is_int8(weight_tensor)
+        and input_tensor.dtype == weight_tensor.dtype
+        and isinstance(input_tensor._layout, PlainLayout)
+        and isinstance(weight_tensor._layout, PlainLayout)
+    )
+
+
+def _linear_int8_act_int8_weight_impl(input_tensor, weight_tensor, bias):
+    #
+    # 1. do the matrix form of dot(X_i, W_j)
+    #
+    #
+    # 2. rescale the output
+    #
+    # in cases with large matrices, y_dot_int32 can grow sufficiently
+    # large that y_dot_int32 * a float16 scale is greater than the maximum
+    # value of a float 16, (which results in a value of inf even if multiplying
+    # by the other scale would bring it within the expected range)
+
+    x_vals_int8 = input_tensor.tensor_impl.int_data
+    x_scales = input_tensor.tensor_impl.scale
+    w_vals_int8_t = weight_tensor.tensor_impl.int_data.contiguous().t()
+    w_scales = weight_tensor.tensor_impl.scale
+    tmp = x_vals_int8.reshape(-1, x_vals_int8.shape[-1])
+    x_scales_dtype = x_scales.dtype
+    # Cast fp16 scale to float to avoid overflow in int_scaled_matmul
+    intermediate_dtype = torch.float if x_scales_dtype == torch.half else x_scales_dtype
+    y_dot_scaled = int_scaled_matmul(
+        tmp, w_vals_int8_t, x_scales.reshape(-1, 1).to(intermediate_dtype)
+    )
+    y_dot_scaled = y_dot_scaled.to(x_scales_dtype)
+
+    y = (y_dot_scaled * w_scales).reshape(
+        *x_vals_int8.shape[:-1], y_dot_scaled.shape[-1]
+    )
+
+    # can downcast only at the very end
+    output_dtype = input_tensor.dtype
+    y = y.to(output_dtype)
+    if bias is not None:
+        y += bias
+    return y

lib/python3.12/site-packages/torchao/dtypes/uintx/q_dq_layout.py

@@ -0,0 +1,242 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+
+import torch
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+import sys
+
+handler = logging.StreamHandler(sys.stdout)
+formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+handler.setFormatter(formatter)
+logger.addHandler(handler)
+
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+from torch.utils._python_dispatch import (
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.utils import AQTTensorImpl, Layout
+from torchao.utils import fill_defaults
+
+aten = torch.ops.aten
+
+
+@dataclass(frozen=True)
+class QDQLayout(Layout):
+    pass
+
+
+def _same_metadata(self: "QDQTensorImpl", src: "QDQTensorImpl") -> bool:
+    return (
+        isinstance(self, QDQTensorImpl)
+        and isinstance(src, QDQTensorImpl)
+        and self.shape == src.shape
+        and self.int_data.shape == src.int_data.shape
+        and self.scale.shape == src.scale.shape
+        and (self.zero_point is None and src.zero_point is None)
+        or (
+            self.zero_point is not None
+            and src.zero_point is not None
+            and self.zero_point.shape == src.zero_point.shape
+        )
+        and type(self._layout) == type(src._layout)
+    )
+
+
+@register_layout(QDQLayout)
+class QDQTensorImpl(AQTTensorImpl):
+    """
+    TensorImpl for QDQLayout layout for affine quantized tensor, it stores int_data, scale, zero_point
+    tensors directly as plain tensors.
+
+    fields:
+      int_data (torch.Tensor): the quantized integer data Tensor
+      scale (torch.Tensor): the scale Tensor used to map between floating point tensor to quantized tensor
+      zero_point (torch.Tensor): the zero_point Tensor used to map between floating point tensor to quantized tensor
+    """
+
+    def __new__(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        kwargs = {}
+        kwargs["device"] = int_data.device
+        kwargs["layout"] = (
+            kwargs.get("layout") if kwargs.get("layout", False) else int_data.layout
+        )
+        kwargs["dtype"] = int_data.dtype
+        kwargs["requires_grad"] = False
+        shape = int_data.shape
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        self.int_data = int_data
+        self.scale = scale
+        self.zero_point = zero_point
+        self._layout = _layout
+
+    def __tensor_flatten__(self):
+        if self.zero_point is None:
+            return ["int_data", "scale"], [self._layout]
+        return ["int_data", "scale", "zero_point"], [self._layout]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        int_data, scale, zero_point = (
+            tensor_data_dict["int_data"],
+            tensor_data_dict["scale"],
+            tensor_data_dict.get("zero_point", None),
+        )
+        (_layout,) = tensor_attributes
+        return cls(int_data, scale, zero_point, _layout)
+
+    def to(self, *args, **kwargs):
+        kwargs = self._get_to_kwargs(*args, **kwargs)
+        return self.__class__(
+            self.int_data.to(kwargs["device"]),
+            self.scale.to(kwargs["device"]),
+            self.zero_point.to(kwargs["device"])
+            if self.zero_point is not None
+            else None,
+            self._layout,
+        )
+
+    def _apply_fn_to_data(self, fn):
+        return self.__class__(
+            fn(self.int_data),
+            fn(self.scale),
+            fn(self.zero_point) if self.zero_point is not None else None,
+            self._layout,
+        )
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+
+        elif func is aten.clone.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.clone)
+            )
+
+        elif func is aten.copy_.default:
+            self = args[0]
+            src = args[1]
+            if _same_metadata(self, src):
+                self_tensors = self.__tensor_flatten__()[0]
+                for tensor_name in self_tensors:
+                    getattr(self, tensor_name).copy_(getattr(src, tensor_name))
+                return
+            raise ValueError(
+                f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}"
+            )
+
+        elif func is aten.t.default:
+            tensor = args[0]
+            new = tensor.__class__(
+                tensor.int_data.t(), tensor.scale, tensor.zero_point, tensor._layout
+            )
+            return return_and_correct_aliasing(func, args, kwargs, new)
+
+        elif func is aten.slice.Tensor:
+            self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1])
+            if dim in [0, 1]:
+                int_data, scale, zero_point = self.get_plain()
+                data_len = int_data.shape[dim]
+                scale_len = scale.shape[dim]
+                ratio = data_len / scale_len
+                start_scale = int(start / ratio)
+                end_scale = int(end / ratio)
+
+                int_data = aten.slice.Tensor(int_data, dim, start, end, step)
+                scale = aten.slice.Tensor(scale, dim, start_scale, end_scale, step)
+                if zero_point is not None:
+                    zero_point = aten.slice.Tensor(
+                        zero_point, dim, start_scale, end_scale, step
+                    )
+                sliced = self.from_plain(int_data, scale, zero_point, self._layout)
+                return return_and_correct_aliasing(func, args, kwargs, sliced)
+            else:
+                raise NotImplementedError(
+                    f"QDQTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported"
+                )
+
+        raise NotImplementedError(
+            f"QDQTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
+    def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        return self.int_data, self.scale, self.zero_point
+
+    def get_layout(self) -> Layout:
+        return self._layout
+
+    @classmethod
+    def from_plain(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        assert isinstance(_layout, QDQLayout)
+        return cls(int_data, scale, zero_point, _layout)
+
+
+def _linear_check(input_tensor, weight_tensor, bias):
+    layout = weight_tensor.tensor_impl.get_layout()
+    return isinstance(layout, QDQLayout)
+
+
+def _linear_impl(input_tensor, weight_tensor, bias):
+    if isinstance(input_tensor, AffineQuantizedTensor):
+        input_tensor = input_tensor.dequantize()
+    if isinstance(weight_tensor, AffineQuantizedTensor):
+        weight_tensor = weight_tensor.dequantize()
+    return torch.nn.functional.linear(input_tensor, weight_tensor, bias)
+
+
+def _embedding_check(args, kwargs):
+    _, weight_tensor = args
+    layout = weight_tensor.tensor_impl.get_layout()
+    return isinstance(layout, QDQLayout)
+
+
+def _embedding_impl(args, kwargs):
+    input_tensor, weight_tensor = args
+    if isinstance(weight_tensor, AffineQuantizedTensor):
+        weight_tensor = weight_tensor.dequantize()
+    return torch.nn.functional.embedding(input_tensor, weight_tensor, **kwargs)

lib/python3.12/site-packages/torchao/dtypes/uintx/semi_sparse_layout.py

@@ -0,0 +1,128 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch.utils._python_dispatch import (
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+from torchao.dtypes.uintx.plain_layout import (
+    PlainAQTTensorImpl,
+    _aqt_is_int8_reduced_range,
+)
+from torchao.dtypes.utils import Layout, PlainLayout
+
+aten = torch.ops.aten
+
+
+def _linear_int8_act_int8_weight_semi_structured_sparse_check(
+    input_tensor, weight_tensor, bias
+):
+    return (
+        isinstance(input_tensor, AffineQuantizedTensor)
+        and _aqt_is_int8_reduced_range(input_tensor)
+        and isinstance(weight_tensor, AffineQuantizedTensor)
+        and weight_tensor.is_cuda
+        and input_tensor.dtype == weight_tensor.dtype
+        and isinstance(input_tensor._layout, PlainLayout)
+        and isinstance(weight_tensor._layout, SemiSparseLayout)
+    )
+
+
+def _linear_int8_act_int8_weight_semi_structured_sparse_impl(
+    input_tensor, weight_tensor, bias
+):
+    x_vals_int8 = input_tensor.tensor_impl.int_data
+    x_scales = input_tensor.tensor_impl.scale
+    w_vals_int8 = weight_tensor.tensor_impl.int_data
+    w_scales = weight_tensor.tensor_impl.scale
+    tmp = x_vals_int8.reshape(-1, x_vals_int8.shape[-1])
+    # must pad
+    row, col = tmp.shape
+    from torch.sparse import SparseSemiStructuredTensorCUSPARSELT
+
+    tmp_padded = SparseSemiStructuredTensorCUSPARSELT._pad_dense_input(tmp)
+    # we fuse one of the scalar matrix multiplications (w_scales) into the sparse mm
+    y_dot_bf16_w_scales_fused = torch._cslt_sparse_mm(
+        w_vals_int8,
+        tmp_padded.t(),
+        alpha=w_scales.to(torch.float32),
+        out_dtype=torch.bfloat16,
+    ).t()[:row, :]
+    y = (y_dot_bf16_w_scales_fused * x_scales.reshape(-1, 1)).reshape(
+        *x_vals_int8.shape[:-1], y_dot_bf16_w_scales_fused.shape[-1]
+    )
+    output_dtype = input_tensor.dtype
+    # TODO: waiting for jesse's test/fix
+    y = y.to(output_dtype).contiguous()
+    if bias is not None:
+        y += bias
+    return y
+
+
+@dataclass(frozen=True)
+class SemiSparseLayout(Layout):
+    """SemiSparseLayout is a layout class for handling semi-structured sparse
+    matrices in affine quantized tensors. This layout is specifically designed
+    to work with the 2:4 sparsity pattern, where two out of every four elements
+    are pruned to zero. This class provides methods for preprocessing input
+    tensors to conform to this sparsity pattern.
+    """
+
+    def pre_process(self, input: torch.Tensor) -> torch.Tensor:
+        # prune to 2:4 if not already
+        temp = input.detach()
+        pruning_inds = temp.abs().view(-1, 4).argsort(dim=1)[:, :2]
+        temp.view(-1, 4).scatter_(1, pruning_inds, value=0)
+        return temp
+
+
+@register_layout(SemiSparseLayout)
+class SemiSparseAQTTensorImpl(PlainAQTTensorImpl):
+    """
+    TensorImpl for semi_sparse_cusparselt layout for affine quantized tensor
+    """
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+
+        raise NotImplementedError(
+            f"SparseAQTTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    def get_plain(self):
+        # Currently we don't have cuSPARSELt expansion routines, so we matmul by
+        # the identity matrix to get the original dense matrix. This is slow though.
+        cols = self.int_data.numel() * 16 // (10 * self.scale.shape[0])
+        int_data_expanded = torch._cslt_sparse_mm(
+            self.int_data,
+            torch.eye(cols, dtype=self.int_data.dtype, device=self.int_data.device).t(),
+        )
+        return int_data_expanded, self.scale, self.zero_point
+
+    @classmethod
+    def from_plain(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        assert isinstance(_layout, SemiSparseLayout)
+        int_data_compressed = torch._cslt_compress(int_data)
+        return cls(int_data_compressed, scale, zero_point, _layout)

lib/python3.12/site-packages/torchao/dtypes/uintx/tensor_core_tiled_layout.py

@@ -0,0 +1,527 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch.utils._python_dispatch import (
+    is_traceable_wrapper_subclass,
+    return_and_correct_aliasing,
+)
+
+from torchao.dtypes.affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    register_layout,
+)
+from torchao.dtypes.utils import AQTTensorImpl, Layout, is_device
+from torchao.quantization.quant_primitives import ZeroPointDomain, _get_reduction_params
+from torchao.utils import (
+    TORCH_VERSION_AT_LEAST_2_5,
+    fill_defaults,
+    find_multiple,
+)
+
+aten = torch.ops.aten
+
+
+def _aqt_is_tensor_core_tile_uint4(aqt):
+    """Check if an AffineQuantizedTensor is uint4 quantized Tensor"""
+    # TODO: use torch.uint4
+    return (
+        aqt.tensor_impl.dtype == torch.int32
+        and aqt.quant_min == 0
+        and aqt.quant_max == 15
+    )
+
+
+def _same_metadata(
+    self: "TensorCoreTiledAQTTensorImpl", src: "TensorCoreTiledAQTTensorImpl"
+) -> bool:
+    return (
+        isinstance(self, TensorCoreTiledAQTTensorImpl)
+        and isinstance(src, TensorCoreTiledAQTTensorImpl)
+        and self.shape == src.shape
+        and self.packed_weight.shape == src.packed_weight.shape
+        and self.scale_and_zero.shape == src.scale_and_zero.shape
+        and self.transposed == src.transposed
+        and type(self._layout) == type(src._layout)
+    )
+
+
+def _linear_bf16_act_uint4_weight_check(input_tensor, weight_tensor, bias):
+    return (
+        # input is native bfloat16 tensor
+        not is_traceable_wrapper_subclass(input_tensor)
+        and input_tensor.dtype == torch.bfloat16
+        and
+        # weight is uint4, group quantized tensor_core_tiled tensor impl affine quantized tensor
+        isinstance(weight_tensor, AffineQuantizedTensor)
+        and _aqt_is_tensor_core_tile_uint4(weight_tensor)
+        and weight_tensor.dtype == torch.bfloat16
+        and len(weight_tensor.shape) == 2
+        and weight_tensor.zero_point_domain == ZeroPointDomain.FLOAT
+        and isinstance(weight_tensor._layout, TensorCoreTiledLayout)
+    )
+
+
+def _linear_bf16_act_uint4_weight_impl(input_tensor, weight_tensor, bias):
+    assert weight_tensor.block_size[0] == 1, (
+        f"Requires groupwise quantization, got block_size: {weight_tensor.block_size}"
+    )
+    assert input_tensor.shape[-1] == weight_tensor.shape[1], (
+        f"need input_tensor shape: {input_tensor.shape} final"
+        f"dim to match weight_tensor shape: {weight_tensor.shape} second dim "
+    )
+
+    # TODO: check groupsize quantization
+    # avoid circular dep, TODO: move this to a common util.py
+    act_mat = input_tensor
+    # weight is packed from padded (out_features, in_features) weight tensor
+    # (same dimension requirement as F.linear weight)
+    packed_weight = weight_tensor.tensor_impl.packed_weight
+    scale_and_zero = weight_tensor.tensor_impl.scale_and_zero
+
+    orig_act_size = act_mat.size()
+    orig_dtype = act_mat.dtype
+
+    # reshape and pad activation
+    act_mat = act_mat.reshape(-1, act_mat.shape[-1]).to(torch.bfloat16)
+    pad_size = find_multiple(act_mat.shape[-1], 1024)
+    act_mat = torch.nn.functional.pad(act_mat, (0, pad_size - act_mat.shape[-1]))
+
+    # groupwise int4 quantization
+    groupsize = weight_tensor.block_size[-1]
+    if act_mat.numel() == 0:  # handling for empty input
+        y = act_mat
+    else:
+        y = torch.ops.aten._weight_int4pack_mm(
+            act_mat.contiguous(), packed_weight, groupsize, scale_and_zero
+        )
+    # remove out_feature padding
+    orig_out_features = weight_tensor.shape[-2]
+    y = y[:, :orig_out_features]
+    y = y.reshape(*orig_act_size[:-1], orig_out_features)
+
+    if bias is not None:
+        y += bias
+    return y.to(orig_dtype)
+
+
+@dataclass(frozen=True)
+class TensorCoreTiledLayout(Layout):
+    """TensorCoreTiledLayout is a layout class for handling tensor core tiled layouts in affine quantized tensors. It provides methods for pre-processing and post-processing tensors to fit the required layout for efficient computation on tensor cores.
+
+    Attributes:
+        inner_k_tiles (int): An internal argument for the packing function of tensor core tiled layout that can affect the performance of the matmul kernel. Defaults to 8.
+    """
+
+    inner_k_tiles: int = 8
+
+    def pre_process(self, input: torch.Tensor) -> torch.Tensor:
+        orig_out_features, orig_in_features = input.shape[-2:]
+        in_features = find_multiple(orig_in_features, 1024)
+        out_features = find_multiple(orig_out_features, 8)
+        input = torch.nn.functional.pad(
+            input,
+            (0, in_features - orig_in_features, 0, out_features - orig_out_features),
+        )
+        return input
+
+    def pre_process_static(
+        self,
+        input: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        block_size: Tuple[int, ...],
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        input = self.pre_process(input)
+        orig_qparam_shape = scale.shape
+        new_qparam_shape, reduction_dims = _get_reduction_params(
+            block_size, input.size()
+        )
+        for dim in reduction_dims:
+            new_qparam_shape.pop(dim)
+        change_in_qparam_shape = [
+            new_dim_size - orig_dim_size
+            for new_dim_size, orig_dim_size in zip(new_qparam_shape, orig_qparam_shape)
+        ]
+        padding_changes = []
+        for dim_change in change_in_qparam_shape:
+            padding_changes = [0, dim_change] + padding_changes
+        scale = torch.nn.functional.pad(scale, padding_changes)
+        zero_point = torch.nn.functional.pad(zero_point, padding_changes)
+        return input, scale, zero_point
+
+    def post_process(
+        self,
+        input: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        block_size: Tuple[int, ...],
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        orig_out_features, orig_in_features = input.shape[-2:]
+        in_features = find_multiple(orig_in_features, 1024)
+        out_features = find_multiple(orig_out_features, 8)
+        input = torch.nn.functional.pad(
+            input,
+            (0, in_features - orig_in_features, 0, out_features - orig_out_features),
+        )
+        assert len(block_size) == 2 or len(block_size) == 3, (
+            f"TensorCoreTiledLayout only supports len(block_size) == 2 or 3, got: {block_size}"
+        )
+        scale_pad_dim_0 = (out_features - orig_out_features) // block_size[-2]
+        scale_pad_dim_1 = (in_features - orig_in_features) // block_size[-1]
+        scale = torch.nn.functional.pad(scale, (0, scale_pad_dim_1, 0, scale_pad_dim_0))
+        zero_point = torch.nn.functional.pad(
+            zero_point, (0, scale_pad_dim_1, 0, scale_pad_dim_0)
+        )
+        return input, scale, zero_point
+
+    def extra_repr(self):
+        return f"inner_k_tiles={self.inner_k_tiles}"
+
+
+@register_layout(TensorCoreTiledLayout)
+class TensorCoreTiledAQTTensorImpl(AQTTensorImpl):
+    """TensorImpl for tensor_core_tiled layout for affine quantized tensor, this is for int4 only,
+    used by tinygemm kernels `_weight_int4pack_mm`
+
+    It stores the original tensor of dimension [n][k] (int32 dtype) as packed weight of 4-d tensor of
+    dimension: [n / 8][k / (inner_k_tiles * 16)][32][inner_k_tiles / 2]
+    (unpacked Tensor shape is n * k)
+    where inner_k_tiles is an internal argument for packing function of tensor core tiled layout
+    that can affect the performance of the matmul kernel (defaults to 8)
+
+    Note: we also pack scale and zero point together here for tinygemm kernel
+
+    Note: technically tensor core tiled layout should be the layout for the underlying packed weight
+    (int Tensor) but since the scale and zero_point are also packed into the same tensor here which is not used
+    in plain layout, we just created a layout for AQT right now, this could be improved if we split out
+    int4 aqt into a separate tensor subclass
+
+    fields:
+      packed_weight (torch.Tensor): the 4-d packed tensor in a tensor_core_tiled layout
+      scale_and_zero (torch.Tensor): the combined scale Tensor used to map between floating point tensor to quantized tensor and zero_point Tensor
+    """
+
+    def __new__(
+        cls,
+        packed_weight: torch.Tensor,
+        scale_and_zero: torch.Tensor,
+        transposed: bool,
+        _layout: Layout,
+    ):
+        kwargs = {}
+        kwargs["device"] = packed_weight.device
+        kwargs["layout"] = (
+            kwargs.get("layout")
+            if kwargs.get("layout", False)
+            else packed_weight.layout
+        )
+        kwargs["dtype"] = packed_weight.dtype
+        kwargs["requires_grad"] = False
+        shape = packed_weight.shape
+        return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
+
+    def __init__(
+        self,
+        packed_weight: torch.Tensor,
+        scale_and_zero: torch.Tensor,
+        transposed: bool,
+        _layout: Layout,
+    ):
+        self.packed_weight = packed_weight
+        self.scale_and_zero = scale_and_zero
+        self.transposed = False
+        self._layout = _layout
+
+    def __tensor_flatten__(self):
+        return ["packed_weight", "scale_and_zero"], [self.transposed, self._layout]
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        packed_weight, scale_and_zero = (
+            tensor_data_dict["packed_weight"],
+            tensor_data_dict["scale_and_zero"],
+        )
+        (
+            transposed,
+            _layout,
+        ) = tensor_attributes
+        return cls(packed_weight, scale_and_zero, transposed, _layout)
+
+    @classmethod
+    def from_plain(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: Optional[torch.Tensor],
+        _layout: Layout,
+    ):
+        assert isinstance(_layout, TensorCoreTiledLayout)
+        assert int_data.dtype == torch.int32, (
+            "torch.ops.aten._convert_weight_to_int4pack in torch 2.4 expects `int32` dtype"
+        )
+
+        def quant_2d(int_data_2d):
+            if TORCH_VERSION_AT_LEAST_2_5:
+                int_data_2d = (int_data_2d[::, ::2] << 4 | int_data_2d[::, 1::2]).to(
+                    torch.uint8
+                )
+            else:
+                assert int_data_2d.dtype == torch.int32, (
+                    "torch.ops.aten._convert_weight_to_int4pack in torch 2.4 expects `int32` dtype"
+                )
+            return torch.ops.aten._convert_weight_to_int4pack(
+                int_data_2d.contiguous(), _layout.inner_k_tiles
+            )
+
+        if int_data.dim() == 3:  # for moe quant
+            num_experts = int_data.shape[0]
+            packed_weight_list = []
+            for expert in range(num_experts):
+                packed_weight_list.append(quant_2d(int_data[expert]).unsqueeze(0))
+            packed_weight = torch.cat(packed_weight_list, dim=0)
+            scale = scale.reshape(int_data.shape[0], int_data.shape[-2], -1)
+            zero_point = (
+                zero_point.reshape(int_data.shape[0], int_data.shape[-2], -1)
+                if zero_point is not None
+                else None
+            )
+        else:
+            assert int_data.dim() == 2
+            packed_weight = quant_2d(int_data)
+            scale = scale.reshape(int_data.shape[0], -1)
+            zero_point = (
+                zero_point.reshape(int_data.shape[0], -1)
+                if zero_point is not None
+                else None
+            )
+        from torchao.quantization.utils import pack_tinygemm_scales_and_zeros
+
+        scale_and_zero = pack_tinygemm_scales_and_zeros(scale, zero_point, scale.dtype)
+        return cls(packed_weight, scale_and_zero, False, _layout)
+
+    def to(self, *args, **kwargs):
+        kwargs = self._get_to_kwargs(*args, **kwargs)
+        device = kwargs["device"]
+        # tensor core tiled layout supports both cpu and cuda but does not support the conversion
+        # between these two devices, in the future we should not use the same layout for
+        # cpu and cuda device: https://github.com/pytorch/ao/issues/1117
+        if not is_device(torch.device(self.device).type, device):
+            raise ValueError(
+                f"TensorCoreTiledAQTTensorImpl does not support conversion from {self.device} to {device}"
+            )
+        return self.__class__(
+            self.packed_weight.to(device),
+            self.scale_and_zero.to(device),
+            self.transposed,
+            self._layout,
+        )
+
+    def _apply_fn_to_data(self, fn):
+        # self.packed_weight = fn(self.packed_weight)
+        # self.scale_and_zero = fn(self.scale_and_zero)
+        # return self
+        return self.__class__(
+            fn(self.packed_weight),
+            fn(self.scale_and_zero),
+            self.transposed,
+            self._layout,
+        )
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):
+        kwargs = {} if kwargs is None else kwargs
+
+        if func is aten.detach.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+            )
+
+        if func is aten.clone.default:
+            return return_and_correct_aliasing(
+                func, args, kwargs, args[0]._apply_fn_to_data(torch.clone)
+            )
+
+        if func is aten.copy_.default:
+            self = args[0]
+            src = args[1]
+            if _same_metadata(self, src):
+                self_tensors = self.__tensor_flatten__()[0]
+                for tensor_name in self_tensors:
+                    getattr(self, tensor_name).copy_(getattr(src, tensor_name))
+                return
+            raise ValueError(
+                f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}"
+            )
+
+        if func in [aten.select.int, aten.index.Tensor]:
+            assert not (func is aten.select.int and args[1] != 0), (
+                "aten.select.int currently only has support for dim=0"
+            )
+            return return_and_correct_aliasing(
+                func,
+                args,
+                kwargs,
+                args[0]._apply_fn_to_data(lambda x: func(x, *args[1:], **kwargs)),
+            )
+
+        if func is aten.t.default:
+            """we don't need to repack the weight and just rely on external
+            shape being changed and record the status of transpose/no-transpose
+            """
+            transposed = TensorCoreTiledAQTTensorImpl(
+                args[0].packed_weight,
+                args[0].scale_and_zero,
+                not args[0].transposed,
+                args[0]._layout,
+            )
+            return return_and_correct_aliasing(func, args, kwargs, transposed)
+
+        if func is aten.slice.Tensor:
+            self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1])
+            cur_shape = self.shape
+            assert len(cur_shape) == 4
+            inner_k_tiles = cur_shape[-1] * 2
+            original_shape = (cur_shape[0] * 8, cur_shape[1] * (inner_k_tiles * 16))
+
+            n_by_8, k_by_inner_tiles, _, _ = self.packed_weight.shape
+            sz_dim1, sz_dim0, _ = self.scale_and_zero.shape
+
+            data_len = original_shape[dim]
+            assert dim in [0, 1], (
+                f"TensorCoreTiledAQTTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported"
+            )
+
+            if dim == 0:
+                pw_len = n_by_8
+                sz_len = sz_dim0
+            else:
+                pw_len = k_by_inner_tiles
+                sz_len = sz_dim1
+
+            if pw_len == 0 or sz_len == 0:
+                return return_and_correct_aliasing(
+                    func,
+                    args,
+                    kwargs,
+                    TensorCoreTiledAQTTensorImpl(
+                        self.packed_weight,
+                        self.scale_and_zero,
+                        self.transposed,
+                        self._layout,
+                    ),
+                )
+
+            pw_ratio = data_len / pw_len
+            start_pw = int(start / pw_ratio)
+            end_pw = int(end / pw_ratio)
+
+            sz_ratio = data_len / sz_len
+            start_sz = int(start / sz_ratio)
+            end_sz = int(end / sz_ratio)
+
+            packed_weight = aten.slice(self.packed_weight, dim, start_pw, end_pw, step)
+            scale_and_zero = aten.slice(
+                self.scale_and_zero, 1 - dim, start_sz, end_sz, step
+            )
+            return return_and_correct_aliasing(
+                func,
+                args,
+                kwargs,
+                TensorCoreTiledAQTTensorImpl(
+                    packed_weight, scale_and_zero, self.transposed, self._layout
+                ),
+            )
+
+        raise NotImplementedError(
+            f"TensorCoreTiledAQTTensorImpl dispatch: attempting to run {func}, this is not supported"
+        )
+
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
+    @property
+    def block_size(self):
+        from torchao.quantization.utils import unpack_tinygemm_scales_and_zeros
+
+        scale, zero = unpack_tinygemm_scales_and_zeros(self.scale_and_zero)
+        cur_shape = self.shape
+        if len(cur_shape) == 5:
+            ones = [1, 1]
+            cur_shape = cur_shape[1:]
+        else:
+            assert len(cur_shape) == 4
+            ones = [1]
+        inner_k_tiles = cur_shape[-1] * 2
+        original_shape = (cur_shape[0] * 8, cur_shape[1] * (inner_k_tiles * 16))
+        groupsize = int(original_shape[1] / scale.shape[-2])
+        return tuple([*ones, groupsize])
+
+    def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        from torchao.quantization.quant_primitives import (
+            ZeroPointDomain,
+            quantize_affine,
+        )
+        from torchao.quantization.utils import unpack_tinygemm_scales_and_zeros
+
+        def dequant_4d(self):
+            cur_shape = self.shape
+            scale, zero = unpack_tinygemm_scales_and_zeros(self.scale_and_zero)
+            assert len(cur_shape) == 4
+            inner_k_tiles = cur_shape[-1] * 2
+            original_shape = (cur_shape[0] * 8, cur_shape[1] * (inner_k_tiles * 16))
+            eye_shape = original_shape[1]
+            groupsize = int(original_shape[1] / scale.shape[-2])
+            block_size = (1, groupsize)
+            original_dtype = torch.bfloat16
+            assert len(block_size) == 2 and block_size[0] == 1
+            dequantized = torch.ops.aten._weight_int4pack_mm(
+                torch.eye(eye_shape, device=self.device, dtype=original_dtype),
+                self.packed_weight,
+                groupsize,
+                self.scale_and_zero,
+            )
+            dequantized = dequantized.t().contiguous()
+            return dequantized
+
+        cur_shape = self.shape
+
+        if len(cur_shape) == 4:
+            dequantized = dequant_4d(self)
+        else:
+            assert len(cur_shape) == 5
+            num_experts = cur_shape[0]
+            dequantized_list = []
+            for expert in range(num_experts):
+                dequantized_list.append(dequant_4d(self[expert]).unsqueeze(0))
+            dequantized = torch.cat(dequantized_list, dim=0)
+
+        scale, zero = unpack_tinygemm_scales_and_zeros(self.scale_and_zero)
+        # TODO: move this to `unpack_tinygemm_scales_and_zeros`?
+        scale = scale.reshape(scale.shape[:-1]).contiguous()
+        zero = zero.reshape(zero.shape[:-1]).contiguous()
+
+        target_dtype = torch.int32
+        quant_min = 0
+        quant_max = 15
+        zero_point_domain = ZeroPointDomain.FLOAT
+        int_data = quantize_affine(
+            dequantized,
+            self.block_size,
+            scale,
+            zero,
+            target_dtype,
+            quant_min,
+            quant_max,
+            zero_point_domain,
+        )
+        return int_data, scale, zero
+
+    def get_layout(self) -> Layout:
+        return self._layout

lib/python3.12/site-packages/torchao/dtypes/uintx/uint4_layout.py

@@ -0,0 +1,321 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+import torch._prims_common as utils
+import torch.utils._pytree as pytree
+from torch.library import Library, impl
+
+from torchao.utils import fill_defaults
+
+
+def down_size(size):
+    assert size[-1] % 2 == 0, f"{size} last dim not divisible by two"
+    return (*size[:-1], size[-1] // 2)
+
+
+def up_size(size):
+    return (*size[:-1], size[-1] * 2)
+
+
+# from
+# https://github.com/drisspg/transformer_nuggets/blob/9ad3a7fc552a954eb702ade0e276b8d8e09c3db6/transformer_nuggets/quant/qlora.py#L233
+
+
+def unpack_uint4(uint8_data) -> torch.Tensor:
+    """Get the original weight from the normalized float weight format"""
+    # since we are using uint8 we will decode 2 entries per byte
+    # Shift elements down 4 and select out the bottom 4 bits
+    shape = uint8_data.shape
+    first_elements = (uint8_data >> 4).to(torch.uint8)
+    second_elements = (uint8_data & 0b1111).to(torch.uint8)
+    return torch.stack([first_elements, second_elements], dim=-1).view(up_size(shape))
+
+
+def pack_uint4(uint8_data) -> torch.Tensor:
+    # converting to uint8 for operations
+    shape = uint8_data.shape
+    assert shape[-1] % 2 == 0
+    uint8_data = uint8_data.contiguous().view(-1)
+    return (uint8_data[::2] << 4 | uint8_data[1::2]).view(down_size(shape))
+
+
+qtensor_lib = Library("qtensors", "DEF")
+qtensor_lib.define(
+    "quantize_per_tensor_uint4(Tensor input, float scale, int zero_point) -> Tensor"
+)
+
+
+@impl(qtensor_lib, "quantize_per_tensor_uint4", "CompositeExplicitAutograd")
+def quantize_per_tensor_uint4(
+    input: torch.Tensor,
+    scale: float,
+    zero_point: int,
+) -> torch.Tensor:
+    inv_scale = 1.0 / scale
+    return pack_uint4(
+        torch.clamp(torch.round(input * inv_scale) + zero_point, 0, 15).to(torch.uint8)
+    )
+
+
+qtensor_lib.define(
+    "dequantize_per_tensor_uint4(Tensor input, float scale, int zero_point) -> Tensor"
+)
+
+
+@impl(qtensor_lib, "dequantize_per_tensor_uint4", "CompositeExplicitAutograd")
+def dequantize_per_tensor_uint4(
+    input: torch.Tensor,
+    scale: float,
+    zero_point: int,
+) -> torch.Tensor:
+    input = unpack_uint4(input)
+    return (input.view(torch.uint8).to(torch.float32) - zero_point) * scale
+
+
+class UInt4Tensor(torch.Tensor):
+    @staticmethod
+    def __new__(cls, elem, **kwargs):
+        assert elem.dtype is torch.uint8
+        assert not kwargs.get("requires_grad", False)
+        kwargs["requires_grad"] = False
+
+        return torch.Tensor._make_wrapper_subclass(
+            cls, up_size(elem.shape), dtype=torch.uint4, **kwargs
+        )
+
+    def __init__(self, elem, **kwargs):
+        self.elem = elem
+
+    @classmethod
+    def from_unpacked(cls, unpacked):
+        return UInt4Tensor(pack_uint4(unpacked))
+
+    def tolist(self):
+        return self.to(torch.uint8).tolist()
+
+    def __tensor_flatten__(self):
+        return ["elem"], None
+
+    @staticmethod
+    def __tensor_unflatten__(flattened, meta, outer_size, outer_stride):
+        assert meta is None
+        elem = flattened["elem"]
+        return UInt4Tensor(elem)
+
+    def __hash__(self):
+        return hash(self.elem)
+
+    def __eq__(self, other):
+        return torch.equal(self.elem, other.elem)
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs=None):
+        if func is torch.ops.aten.view.default:
+            self, size = args
+            size = utils.infer_size(size, self.numel())
+            assert not kwargs
+            # WARNING: views not preserved
+            return UInt4Tensor(self.elem.reshape(down_size(size)))
+        elif func is torch.ops.aten.view.dtype:
+            self, dtype = args
+            if dtype == torch.uint8:
+                return unpack_uint4(self.elem).view(torch.uint8)
+            return NotImplementedError(f"view {args}")
+        elif func is torch.ops.aten.to.dtype:
+            self, dtype = args
+            if dtype == torch.uint8:
+                return unpack_uint4(self.elem).view(torch.uint8)
+            return NotImplementedError(f"to {args}")
+        elif func is torch.ops.aten.eq.Tensor:
+            args = pytree.tree_map_only(
+                UInt4Tensor, lambda x: x.elem.view(torch.uint8), args
+            )
+            kwargs = pytree.tree_map_only(
+                UInt4Tensor, lambda x: x.elem.view(torch.uint8), kwargs
+            )
+            return torch.ops.aten.eq.Tensor(*args, **kwargs)
+        elif func is torch.ops.aten._to_copy.default:
+            (self,) = args
+            if kwargs == {"dtype": torch.uint8}:
+                return unpack_uint4(self.elem).view(self.shape)  # no wrap
+            else:
+                raise NotImplementedError(f"_to_copy {kwargs}")
+        elif func is torch.ops.aten.unbind.int:
+            # This is tricky.  Given torch.tensor([0, 1, 2, 3]) we want to
+            # create four tensors containing one element each.  But we can't
+            # do this with uint4 because such a tensor's size is not divisible
+            # by bytes.  What I am going to do instead is promote to uint8
+            # when this happens
+            self, dim = fill_defaults(args, 2, [0])
+            if dim != self.dim() - 1:
+                raise NotImplementedError(f"unbind dim={dim}")
+            else:
+                # We're unbinding the last dimension, need to promote
+                return torch.ops.aten._to_copy.default(self, dtype=torch.uint8).unbind(
+                    dim
+                )
+        elif func is torch.ops.aten.select.int:
+            self, dim, index = args
+            if dim != self.dim() - 1:
+                return UInt4Tensor(torch.ops.aten.select.int(self.elem, dim, index))
+            else:
+                raise NotImplementedError(f"select dim={dim}")
+        elif func is torch.ops.aten.slice.Tensor:
+            self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1])
+            if dim == self.dim() - 1:
+                # hard case
+                if step != 1:
+                    raise NotImplementedError(f"slice step={step}")
+                assert start % 2 == 0, start
+                assert end >= self.shape[dim] or end % 2 == 0, end
+                return UInt4Tensor(
+                    torch.ops.aten.slice.Tensor(self.elem, dim, start // 2, end // 2, 1)
+                )
+            else:
+                # easy case
+                return UInt4Tensor(
+                    torch.ops.aten.slice.Tensor(self.elem, dim, start, end, step)
+                )
+        elif func is torch.ops.aten.t.default:
+            # assert False, "transpose is not properly implemented currently"
+            (self,) = args
+            unpacked = unpack_uint4(self.elem)
+            transposed = torch.ops.aten.t.default(unpacked)
+            transposed_and_packed = pack_uint4(transposed)
+            return UInt4Tensor(transposed_and_packed)
+        elif func is torch.ops.aten.transpose_copy.int:
+            self, dim0, dim1 = args
+            unpacked = unpack_uint4(self.elem).view(self.shape)
+            transposed = torch.ops.aten.transpose_copy.int(unpacked, dim0, dim1)
+            transposed_and_packed = pack_uint4(transposed)
+            return UInt4Tensor(transposed_and_packed)
+        elif func is torch.ops.aten.as_strided.default:
+            # size, stride, storage_offset are referring to tensor elements, not physical bytes
+            self, size, stride, storage_offset = args
+            size = down_size(size)
+
+            new_stride = []
+            for s in stride:
+                if s != 1:
+                    # since two int4 equals to 1 uint8
+                    new_stride.append(s // 2)
+                else:
+                    new_stride.append(s)
+            stride = new_stride
+
+            storage_offset //= 2
+            return UInt4Tensor(
+                torch.ops.aten.as_strided.default(
+                    self.elem, size, stride, storage_offset
+                )
+            )
+
+        raise NotImplementedError(f"{func}")
+
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
+
+def _dynamically_quantize_per_channel_int4(x, quant_min, quant_max, target_dtype):
+    # assumes symmetric quantization
+    # assumes axis == 0
+    # assumes dense memory format
+    # TODO(future): relax ^ as needed
+
+    # default setup for affine quantization of activations
+    eps = torch.finfo(torch.float32).eps
+
+    # get min and max
+    min_val, max_val = torch.aminmax(x, dim=1)
+
+    # calculate scale and zero point based on min and max
+    # reference: https://fburl.com/code/srbiybme
+    min_val_neg = torch.min(min_val, torch.zeros_like(min_val))
+    max_val_pos = torch.max(max_val, torch.zeros_like(max_val))
+    device = min_val_neg.device
+
+    # reference: https://fburl.com/code/4wll53rk
+    max_val_pos = torch.max(-min_val_neg, max_val_pos)
+    scale = max_val_pos / (float(quant_max - quant_min) / 2)
+    # ensure scale is the same dtype as the original tensor
+    scale = torch.clamp(scale, min=eps).to(x.dtype)
+    zero_point = torch.zeros(min_val_neg.size(), dtype=torch.int64, device=device)
+
+    # quantize based on qmin/qmax/scale/zp
+    # reference: torch/ao/quantization/fx/_decomposed.py?lines=63
+    x_div = x.transpose(0, 1) / scale
+    x_round = torch.round(x_div)
+    x_zp = x_round + zero_point
+    x_zp = x_zp.transpose(0, 1)
+    quant = torch.clamp(x_zp, quant_min, quant_max)
+
+    if target_dtype == torch.uint4:
+        # TODO: simplify (maybe implement to)
+        quant = PerChannelSymmetricWeightUInt4Tensor.from_unpacked(
+            quant.to(torch.uint8), scale
+        )
+    else:
+        quant = quant.to(target_dtype)
+
+    return quant, scale, zero_point
+
+
+class PerChannelSymmetricWeightUInt4Tensor(UInt4Tensor):
+    @staticmethod
+    def __new__(cls, elem, scales, **kwargs):
+        return super().__new__(cls, elem, **kwargs)
+
+    def __init__(self, elem, scales, **kwargs):
+        super().__init__(elem, **kwargs)
+
+        self.scales = scales
+
+    def __tensor_flatten__(self):
+        return ["elem", "scales"], None
+
+    @staticmethod
+    def __tensor_unflatten__(flattened, meta, outer_size, outer_stride):
+        assert meta is None
+        elem = flattened["elem"]
+        scales = flattened["scales"]
+        return PerChannelSymmetricWeightUInt4Tensor(elem, scales)
+
+    @classmethod
+
+    #  inconsistently.
+
+    def from_unpacked(cls, unpacked, scales):
+        return cls(pack_uint4(unpacked), scales)
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs=None):
+        if func is torch.ops.aten.addmm.default:
+            bias, x, weight = args
+            x_view = x.view(-1, x.shape[-1])
+            y = torch.mm(x_view, weight.to(torch.uint8).to(x.dtype)) * weight.scales
+            y = y.reshape(*x.shape[:-1], -1)
+            if bias is not None:
+                y += bias
+            return y
+        elif func is torch.ops.aten.t.default:
+            # TODO: add proper support for transpose
+            (self,) = args
+            unpacked = unpack_uint4(self.elem)
+            transposed = torch.ops.aten.t.default(unpacked)
+            return PerChannelSymmetricWeightUInt4Tensor.from_unpacked(
+                transposed, self.scales
+            )
+        elif func is torch.ops.aten.detach.default:
+            (self,) = args
+            return self
+        return super().__torch_dispatch__(func, types, args, kwargs)
+
+    @classmethod
+    def from_float(cls, w_fp32):
+        w_int4, scales, _zp = _dynamically_quantize_per_channel_int4(
+            w_fp32, 0, 15, torch.uint4
+        )
+        w_int4 = w_int4.to(device=w_fp32.device)
+        return w_int4