[cuda backend] int5 quantization support

backends/aoti/common_shims.cpp

@@ -175,6 +175,10 @@ int32_t aoti_torch_dtype_bfloat16() {
   return 15; // PyTorch's bfloat16 dtype code
 }
 
+int32_t aoti_torch_dtype_float16() {
+  return 5; // PyTorch's float16 (Half) dtype code
+}
+
 int32_t aoti_torch_dtype_uint8() {
   return 0; // PyTorch's uint8 dtype code
 }

backends/aoti/common_shims.h

@@ -64,6 +64,7 @@ AOTI_SHIM_EXPORT int32_t aoti_torch_device_type_cpu();
 AOTI_SHIM_EXPORT int32_t aoti_torch_layout_strided();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float32();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_bfloat16();
+AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float16();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_bool();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int8();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_uint8();

backends/aoti/common_shims_slim.cpp

@@ -120,6 +120,10 @@ int32_t aoti_torch_dtype_bfloat16() {
   return 15; // ScalarType::BFloat16
 }
 
+int32_t aoti_torch_dtype_float16() {
+  return 5; // ScalarType::Half
+}
+
 int32_t aoti_torch_dtype_int64() {
   return 4; // ScalarType::Long
 }

backends/aoti/common_shims_slim.h

@@ -72,6 +72,7 @@ aoti_torch_get_device_index(Tensor* tensor, int32_t* ret_device_index);
 
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float32();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_bfloat16();
+AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float16();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int64();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int32();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int16();

backends/cuda/CMakeLists.txt

@@ -114,6 +114,7 @@ if(CMAKE_CUDA_COMPILER)
     _aoti_cuda_shim_sources
     runtime/shims/int4mm.cu
     runtime/shims/int4_plain_mm.cu
+    runtime/shims/int5_plain_mm.cu
     runtime/shims/int6_plain_mm.cu
     runtime/shims/int8_plain_mm.cu
     runtime/shims/sort.cu

backends/cuda/coalesced_int4_tensor.py

@@ -14,20 +14,43 @@
 ``Int4Tensor`` weights keep falling back to torchao's default (mslk/tinygemm)
 path.
 
+Metadata encoding: the SCALE is a per-group **uint8 code** with a
+**per-256-super-block fp16 step** (``scale = code * step``); the ZERO now uses
+the SAME per-256-super-block fp16 step (per-group uint8 code + per-256 fp16
+step). group_size is 32, so a 256-weight super-block spans 8 groups and there
+are ``K/256`` scale steps and ``K/256`` zero steps per row. Dequant is
+unchanged: ``w = (q - zero) * scale`` with ``zero = min/scale``.
+
+This mirrors GGUF Q4_K's per-super-block fp16 ``d`` for both the scale and the
+zero: the finer per-256 step (vs the previous per-row step) is what lifts
+whole-weight dequant SNR to ~45.89 dB (vs 45.15 dB for the per-row-zero-step
+encoding, +0.74 dB). The scale/zero codes stay single coalesced uint8s (one
+byte/group) so the decode kernel reads exactly one scale byte and one zero byte
+per group — no bit-plane reconstruct — and both per-256 steps are loaded once
+per super-block. Both steps MUST be fp16 (bf16 for the per-256 step costs ~0.05
+dB on the scale; the per-256 zero step is +0.74 dB at fp16 vs +0.52 at bf16).
+
 Layout difference from torchao ``Int4Tensor``:
     qdata      : packed int4 weight (N, K/2), nibble-packed (same as Int4Tensor)
-    scale      : (N, n_groups) — the *coalesced* layout, transposed from
-                 torchao's documented (n_groups, N)
-    zero_point : (N, n_groups) — coalesced, transposed from (n_groups, N)
+    scale      : (N, n_groups) uint8 — per-group scale *codes*, coalesced
+                 (transposed from torchao's (n_groups, N))
+    scale_step : (N, K/256) fp16 — per-256-super-block scale step; the real
+                 per-group scale is ``scale_code * scale_step[:, g // 8]``.
+    zero_point : (N, n_groups) uint8 — per-group zero codes
+    zero_point_step : (N, K/256) fp16 — per-256-super-block zero step; the real
+                 per-group zero is ``zero_code * zero_point_step[:, g // 8]``.
+
+Bits-per-weight: 4.0 (qdata) + 8/32 (scale codes) + 16/256 (fp16 scale step) +
+8/32 (uint8 zero codes) + 16/256 (fp16 zero step) = 4.625 bpw.
 
 The coalesced [N, n_groups] layout is exactly what the W4A8 dp4a matvec kernel
 (``executorch_cuda::int4_plain_mm`` / ``int4_plain_mm.cuh``) reads row-for-row
 with qdata, so the exported decode graph carries no per-step transpose. The
-transpose is owned by :meth:`from_int4_tensor` so it is baked into the
-serialized weight constant once at pack time.
+transpose (and the uint8 re-encoding) is owned by :meth:`from_int4_tensor` so it
+is baked into the serialized weight constant once at pack time.
 """
 
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 import torch
 from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor
@@ -37,17 +60,72 @@
     "CudaCoalescedInt4Tensor",
 ]
 
+_CODE_MAX = 255  # uint8 code range [0, 255] (both scale and zero)
+_SUPER_BLOCK = 256  # weights per super-block (GGUF Q4_K QK_K); scale step is per this
+
+
+def _encode_uint8_per_super(
+    x: torch.Tensor,
+    group_size: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Encode a (n_groups, N) non-negative tensor to per-super-block uint8 codes.
+
+    Used for both the scale and the zero. A super-block is ``_SUPER_BLOCK``
+    (256) weights = ``groups_per_super = 256 // group_size`` groups (8 for
+    group_size=32). Returns ``(codes, step)`` where ``codes`` is
+    ``(N, n_groups)`` uint8 (transposed to the coalesced layout) and ``step`` is
+    ``(N, n_super)`` fp16 with ``n_super = n_groups // groups_per_super = K //
+    256``, such that ``code * step[:, g // groups_per_super] ≈ x.t()``. The step
+    is the per-256-super-block max / 255 rounded to fp16. Rounding uses the
+    fp16-rounded step (what the kernel reads) so encode and decode agree.
+    """
+    xt = x.t().contiguous().float()  # (N, n_groups)
+    N, n_groups = int(xt.shape[0]), int(xt.shape[1])
+    groups_per_super = _SUPER_BLOCK // int(group_size)
+    if groups_per_super < 1:
+        raise ValueError(
+            f"group_size={group_size} must be <= {_SUPER_BLOCK} for the per-256 "
+            "scale step"
+        )
+    if n_groups % groups_per_super != 0:
+        raise ValueError(
+            f"n_groups={n_groups} must be a multiple of {groups_per_super} "
+            f"(K must be a multiple of {_SUPER_BLOCK}) for group_size={group_size}"
+        )
+    n_super = n_groups // groups_per_super
+    xb = xt.reshape(N, n_super, groups_per_super)  # (N, n_super, gps)
+    block_max = xb.amax(dim=2, keepdim=True).clamp_min(1e-30)  # (N, n_super, 1)
+    step = (block_max / _CODE_MAX).to(torch.float16)  # (N, n_super, 1) fp16
+    step_f = step.float().clamp_min(1e-30)
+    codes = torch.round(xb / step_f).clamp_(0, _CODE_MAX).to(torch.uint8)
+    codes = codes.reshape(N, n_groups).contiguous()
+    return codes, step.squeeze(2).contiguous()
+
+
+def _unpack_nibble_qdata(qdata: torch.Tensor, N: int, K: int) -> torch.Tensor:
+    """Unpack nibble-packed int4 qdata ``(N, K/2)`` -> ``(N, K)`` uint8 [0, 15]."""
+    qu = qdata.to(torch.uint8)
+    even = qu & 0xF
+    odd = (qu >> 4) & 0xF
+    return torch.stack([even, odd], dim=-1).reshape(N, K)
+
 
 class CudaCoalescedInt4Tensor(TorchAOBaseTensor):
-    """INT4 weight with scale/zero_point in the coalesced [N, n_groups] layout.
+    """INT4 weight, uint8 scale + per-256 fp16 step, uint8 zero + per-256 fp16 step.
 
     ExecuTorch-internal; see the module docstring. Mirrors torchao
     ``Int4Tensor``'s data/attribute layout (so the common tensor utilities and
-    serialization work) but owns the [n_groups, N] -> [N, n_groups] transpose
-    of scale/zero_point via :meth:`from_int4_tensor`.
+    serialization work) but owns the [n_groups, N] -> [N, n_groups] transpose and
+    the uint8 re-encoding via :meth:`from_int4_tensor`.
     """
 
-    tensor_data_names = ["qdata", "scale", "zero_point"]
+    tensor_data_names = [
+        "qdata",
+        "scale",
+        "scale_step",
+        "zero_point",
+        "zero_point_step",
+    ]
     tensor_attribute_names = ["block_size", "shape"]
     optional_tensor_data_names = ["act_pre_scale"]
     optional_tensor_attribute_names = ["activation_dtype"]
@@ -56,23 +134,29 @@ def __new__(
         cls,
         qdata: torch.Tensor,
         scale: torch.Tensor,
+        scale_step: torch.Tensor,
         zero_point: torch.Tensor,
+        zero_point_step: torch.Tensor,
         block_size: List[int],
         shape: torch.Size,
         act_pre_scale: Optional[torch.Tensor] = None,
         activation_dtype: Optional[torch.dtype] = None,
     ):
         kwargs = {}
         kwargs["device"] = qdata.device
-        kwargs["dtype"] = scale.dtype
+        kwargs["dtype"] = (
+            activation_dtype if activation_dtype is not None else torch.bfloat16
+        )
         kwargs["requires_grad"] = False
         return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
 
     def __init__(
         self,
         qdata: torch.Tensor,
         scale: torch.Tensor,
+        scale_step: torch.Tensor,
         zero_point: torch.Tensor,
+        zero_point_step: torch.Tensor,
         block_size: List[int],
         shape: torch.Size,
         act_pre_scale: Optional[torch.Tensor] = None,
@@ -81,7 +165,9 @@ def __init__(
         super().__init__()
         self.qdata = qdata
         self.scale = scale
+        self.scale_step = scale_step
         self.zero_point = zero_point
+        self.zero_point_step = zero_point_step
         self.block_size = block_size
         self.activation_dtype = (
             activation_dtype if activation_dtype is not None else torch.bfloat16
@@ -98,21 +184,57 @@ def _quantization_type(self):
     def from_int4_tensor(cls, t: Int4Tensor) -> "CudaCoalescedInt4Tensor":
         """Build a coalesced tensor from a torchao ``Int4Tensor``.
 
-        Owns the transpose: torchao stores scale/zero_point as (n_groups, N);
-        the CUDA decode kernel reads (N, n_groups). The ``.t().contiguous()``
-        here is baked into the serialized weight constant so the exported
-        decode graph has no per-step transpose/clone.
+        Owns the transpose AND the uint8 re-encoding: torchao stores
+        scale/zero_point as (n_groups, N) bf16. The CUDA decode kernel reads the
+        (N, n_groups) uint8 scale/zero *codes* plus per-256-super-block
+        (N, K/256) fp16 ``scale_step`` / ``zero_point_step`` (scale = scale_code *
+        scale_step[:, g//8], zero = zero_code * zero_point_step[:, g//8]). The
+        transpose + encode here is baked into the serialized weight constant so
+        the exported decode graph has no per-step transpose/clone.
         """
+        scale_codes, scale_step = _encode_uint8_per_super(t.scale, t.block_size[-1])
+        zero_codes, zero_point_step = _encode_uint8_per_super(
+            t.zero_point, t.block_size[-1]
+        )
         return cls(
             t.qdata,
-            t.scale.t().contiguous(),
-            t.zero_point.t().contiguous(),
+            scale_codes,
+            scale_step,
+            zero_codes,
+            zero_point_step,
             t.block_size,
             t.shape,
             t.act_pre_scale,
             t.activation_dtype,
         )
 
+    def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor:
+        """Dequantize to a dense tensor: ``w = (q - zero) * scale``.
+
+        Reconstructs the per-group scale from the uint8 codes and the per-256
+        fp16 step, and the per-group zero from the uint8 codes and the per-256
+        fp16 step. Used as the numerical reference and for the tied lm_head /
+        token embedding.
+        """
+        dtype = output_dtype if output_dtype is not None else torch.bfloat16
+        N, K = int(self.shape[0]), int(self.shape[1])
+        gs = self.block_size[-1]
+        n_groups = K // gs
+        n_super = int(self.scale_step.shape[1])
+        groups_per_super = n_groups // n_super
+
+        q = _unpack_nibble_qdata(self.qdata, N, K).to(torch.float32)
+        scale_code = self.scale.to(torch.float32)  # (N, n_groups)
+        scale_step = self.scale_step.float().repeat_interleave(groups_per_super, dim=1)
+        scale = (scale_code * scale_step).repeat_interleave(gs, dim=1)  # (N, K)
+
+        zero_code = self.zero_point.to(torch.float32)  # (N, n_groups)
+        zero_point_step = self.zero_point_step.float().repeat_interleave(
+            groups_per_super, dim=1
+        )
+        zero = (zero_code * zero_point_step).repeat_interleave(gs, dim=1)  # (N, K)
+        return ((q - zero) * scale).to(dtype)
+
 
 # Allow a model with CudaCoalescedInt4Tensor weights to be loaded with
 # `weights_only=True` (mirrors torchao Int4Tensor).

backends/cuda/cuda_backend.py

@@ -31,7 +31,6 @@
 from torch._inductor.decomposition import conv1d_to_conv2d
 from torch.nn.attention import SDPBackend
 
-
 # ---------------------------------------------------------------------------
 # AOTI compile-time CPU clones for mutated buffers
 # ---------------------------------------------------------------------------
@@ -415,8 +414,7 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
                 mode = spec.value.decode("utf-8").upper()
                 if mode not in ["ON", "OFF"]:
                     raise ValueError(
-                        f"Invalid triton_kernel_mode: {mode}. "
-                        f"Expected 'ON' or 'OFF'."
+                        f"Invalid triton_kernel_mode: {mode}. Expected 'ON' or 'OFF'."
                     )
                 triton_kernel_mode = mode
         passes = [MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass()]
@@ -469,12 +467,19 @@ def get_aoti_compile_options(
                 torch.ops.executorch_cuda.int4_plain_mm.default: [
                     "AOTITorchError aoti_torch_cuda_int4_plain_mm("
                     "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
+                    "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
+                    "int64_t, AtenTensorHandle*)"
+                ],
+                torch.ops.executorch_cuda.int5_plain_mm.default: [
+                    "AOTITorchError aoti_torch_cuda_int5_plain_mm("
+                    "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
+                    "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
                     "AtenTensorHandle, int64_t, AtenTensorHandle*)"
                 ],
                 torch.ops.executorch_cuda.int6_plain_mm.default: [
                     "AOTITorchError aoti_torch_cuda_int6_plain_mm("
                     "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, "
-                    "AtenTensorHandle, int64_t, AtenTensorHandle*)"
+                    "AtenTensorHandle, AtenTensorHandle, int64_t, AtenTensorHandle*)"
                 ],
                 torch.ops.executorch_cuda.int8_plain_mm.default: [
                     "AOTITorchError aoti_torch_cuda_int8_plain_mm("