quantbench-engine 0.1.0__py3-none-any.whl

quantbench/quant/algorithms/smoothquant.py

@@ -0,0 +1,132 @@
+# Copyright 2026 Prajwal7-2002
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models.
+
+Mathematical Foundation:
+    LLM activations exhibit severe outlier channels whose magnitudes can be 100x
+    larger than typical channels, making activation quantization to INT8 prone to
+    catastrophic precision loss under per-tensor or per-token absmax scaling.
+    Conversely, weights exhibit smooth channel distributions and are easy to quantize.
+
+    SmoothQuant introduces a mathematically invariant per-channel migration scale s:
+        Y = X * W = (X * diag(s)^-1) * (diag(s) * W) = X_hat * W_hat
+
+    The scale vector s in R^C is chosen to balance the maximum dynamic ranges:
+        s_j = (max(|X_j|)^alpha) / (max(|W_j|)^(1 - alpha))
+
+    where alpha in [0, 1] controls migration strength:
+        - alpha = 0.5: Symmetrically divides difficulty between activations and weights.
+        - alpha = 1.0: Migrates all difficulty to the weights (activations become easy to quantize).
+        - alpha = 0.0: Preserves weights completely (no difficulty migration).
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+
+def compute_smoothquant_scale(
+    act_max: torch.Tensor,
+    weight_max: torch.Tensor,
+    alpha: float = 0.5,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    """Compute the per-channel SmoothQuant migration scale vector.
+
+    Args:
+        act_max: Tensor of shape (C,) representing the maximum absolute activation
+            magnitude per input channel across calibration tokens.
+        weight_max: Tensor of shape (C,) representing the maximum absolute weight
+            magnitude per input channel: max_i |W_{i, j}|.
+        alpha: Migration strength hyperparameter in [0.0, 1.0].
+        eps: Small epsilon to prevent division by zero or degenerate scales.
+
+    Returns:
+        Tensor of shape (C,) containing per-channel migration scales s_j.
+    """
+    if not (0.0 <= alpha <= 1.0):
+        raise ValueError(f"alpha must be in [0.0, 1.0], got {alpha}")
+
+    act_max = act_max.clamp(min=eps)
+    weight_max = weight_max.clamp(min=eps)
+
+    # s_j = (s_X^alpha) / (s_W^(1 - alpha))
+    numerator = act_max.pow(alpha)
+    denominator = weight_max.pow(1.0 - alpha)
+    scale = (numerator / denominator).clamp(min=eps)
+    return scale
+
+
+def extract_activation_channel_scales(
+    module: nn.Module,
+    calibration_inputs: list[torch.Tensor] | torch.Tensor,
+) -> torch.Tensor:
+    """Profile activation outlier channel scales across calibration data.
+
+    Args:
+        module: Neural network module (e.g. nn.Linear).
+        calibration_inputs: A tensor or list of tensors of shape (B, S, C_in) or (N, C_in).
+
+    Returns:
+        Tensor of shape (C_in,) containing per-channel max absolute activation values.
+    """
+    if isinstance(calibration_inputs, torch.Tensor):
+        calibration_inputs = [calibration_inputs]
+
+    channel_max: Optional[torch.Tensor] = None
+
+    for batch in calibration_inputs:
+        # Reshape to (N_tokens, C_in)
+        flat = batch.detach().view(-1, batch.shape[-1]).abs()
+        batch_max = flat.max(dim=0).values
+        if channel_max is None:
+            channel_max = batch_max
+        else:
+            channel_max = torch.maximum(channel_max, batch_max)
+
+    if channel_max is None:
+        raise ValueError("calibration_inputs cannot be empty.")
+
+    return channel_max
+
+
+def smooth_linear_weights(
+    weight: torch.Tensor,
+    act_scales: torch.Tensor,
+    alpha: float = 0.5,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Apply SmoothQuant per-channel migration to weight tensor.
+
+    Args:
+        weight: Float tensor of shape (out_features, in_features).
+        act_scales: Per-channel activation max tensor of shape (in_features,).
+        alpha: Migration factor in [0.0, 1.0].
+
+    Returns:
+        Tuple of (smoothed_weight, migration_scale):
+            - smoothed_weight: W * diag(s) of shape (out_features, in_features).
+            - migration_scale: s of shape (in_features,).
+    """
+    # W shape: (out_features, in_features)
+    # Channel max across output rows: max_i |W_{i, j}| -> shape (in_features,)
+    weight_max = weight.abs().max(dim=0).values
+    scale = compute_smoothquant_scale(act_scales, weight_max, alpha=alpha)
+
+    # W_hat = W * diag(s) -> broadcasting along dim 1
+    smoothed_weight = weight * scale.unsqueeze(0)
+    return smoothed_weight, scale

quantbench/quant/blockwise.py

@@ -0,0 +1,115 @@
+"""Group-wise (Block-wise) Quantization Primitives.
+
+Partitions tensors into independent contiguous blocks of size G (e.g. 32, 64, 128)
+to isolate channel outliers and preserve precision across wide transformer layers.
+"""
+
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+
+from quantbench.quant.affine import (
+    dequantize_absmax,
+    dequantize_affine,
+    quantize_absmax,
+    quantize_affine,
+)
+
+
+def quantize_blockwise(
+    x: torch.Tensor,
+    group_size: int = 128,
+    num_bits: int = 4,
+    symmetric: bool = True,
+    signed: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Dict[str, Any]]:
+    """Quantize a tensor into independent group-wise blocks.
+
+    Divides flattened tensor into blocks of size `group_size`. If the tensor element
+    count is not an exact multiple of `group_size`, pads with zeros.
+
+    Args:
+        x: Input tensor of arbitrary shape.
+        group_size: Number of consecutive elements per block (32, 64, 128).
+        num_bits: Bit precision per element (e.g. 4 or 8).
+        symmetric: If True, uses symmetric AbsMax per block (zero_point = 0).
+                   If False, uses asymmetric affine (scale and zero_point).
+        signed: Whether integer values are signed (only used when symmetric=False).
+
+    Returns:
+        Tuple of (q_blocks, scales, zero_points, metadata_dict).
+    """
+    orig_shape = tuple(x.shape)
+    flat_x = x.contiguous().view(-1)
+    n_elements = flat_x.numel()
+
+    # Calculate padding if necessary
+    remainder = n_elements % group_size
+    pad_len = (group_size - remainder) % group_size
+    if pad_len > 0:
+        flat_x = torch.nn.functional.pad(flat_x, (0, pad_len), value=0.0)
+
+    # Reshape into blocks: [num_blocks, group_size]
+    blocked_x = flat_x.view(-1, group_size)
+
+    if symmetric:
+        q_blocks, scales = quantize_absmax(
+            blocked_x,
+            num_bits=num_bits,
+            dim=1,
+            symmetric_range=True,
+        )
+        zero_points = None
+    else:
+        q_blocks, scales, zero_points = quantize_affine(
+            blocked_x,
+            num_bits=num_bits,
+            dim=1,
+            signed=signed,
+        )
+
+    meta = {
+        "original_shape": orig_shape,
+        "n_elements": n_elements,
+        "pad_len": pad_len,
+        "group_size": group_size,
+        "num_bits": num_bits,
+        "symmetric": symmetric,
+        "dtype": x.dtype,
+        "device": x.device,
+    }
+
+    return q_blocks, scales, zero_points, meta
+
+
+def dequantize_blockwise(
+    q_blocks: torch.Tensor,
+    scales: torch.Tensor,
+    zero_points: Optional[torch.Tensor],
+    meta: Dict[str, Any],
+) -> torch.Tensor:
+    """Reconstruct an original floating-point tensor from block-wise quantized representations.
+
+    Args:
+        q_blocks: Quantized integer blocks of shape [num_blocks, group_size].
+        scales: Scale factors of shape [num_blocks, 1].
+        zero_points: Optional zero-points of shape [num_blocks, 1].
+        meta: Metadata dictionary returned by quantize_blockwise.
+
+    Returns:
+        Reconstructed tensor matching original shape and device.
+    """
+    if meta["symmetric"]:
+        x_recon_blocked = dequantize_absmax(q_blocks, scales)
+    else:
+        assert zero_points is not None, (
+            "zero_points required for asymmetric blockwise dequantization"
+        )
+        x_recon_blocked = dequantize_affine(q_blocks, scales, zero_points)
+
+    flat_recon = x_recon_blocked.view(-1)
+    # Remove padding if added during quantization
+    if meta["pad_len"] > 0:
+        flat_recon = flat_recon[: meta["n_elements"]]
+
+    return flat_recon.view(meta["original_shape"]).to(meta["dtype"])

quantbench/quant/fp8.py

@@ -0,0 +1,166 @@
+"""OCP FP8 (E4M3FN and E5M2) Emulation and Simulation Primitives.
+
+Implements bit-exact simulation of the two standardized 8-bit floating-point formats:
+1. E4M3FN: 1 sign, 4 exponent, 3 mantissa (bias 7, max ±448.0, no infinity, higher precision).
+2. E5M2: 1 sign, 5 exponent, 2 mantissa (bias 15, max ±57344.0, IEEE inf/nan, wider dynamic range).
+"""
+
+from enum import Enum
+from typing import Optional, Tuple
+
+import torch
+
+
+class FP8Format(str, Enum):
+    """Standardized Open Compute Project (OCP) 8-bit floating point formats."""
+
+    E4M3 = "e4m3fn"
+    E5M2 = "e5m2"
+
+
+FP8_MAX_VALUES = {
+    FP8Format.E4M3: 448.0,
+    FP8Format.E5M2: 57344.0,
+}
+
+
+def _emulate_fp8(
+    x: torch.Tensor,
+    exp_bits: int,
+    mantissa_bits: int,
+    bias: int,
+    max_val: float,
+    has_inf: bool = False,
+) -> torch.Tensor:
+    """Analytical bit-level emulator for floating-point formats with arbitrary bitwidths.
+
+    Simulates the exact IEEE / OCP quantization, subnormal behavior, and saturation.
+    """
+    sign = torch.sign(x)
+    abs_x = torch.abs(x)
+
+    # Saturate / clamp to max representable value
+    abs_x = torch.clamp(abs_x, max=max_val)
+
+    # Handle zeros
+    is_zero = abs_x == 0.0
+
+    # Calculate base-2 exponent
+    # e = floor(log2(x))
+    log2_x = torch.log2(torch.clamp(abs_x, min=1e-30))
+    exp = torch.floor(log2_x)
+
+    # Minimum normal exponent
+    min_exp = 1 - bias
+    max_exp = (1 << exp_bits) - 1 - bias
+
+    # Separate normals from subnormals
+    is_subnormal = exp < min_exp
+    exp = torch.clamp(exp, min=min_exp, max=max_exp)
+
+    # Mantissa quantization step size
+    # For normals: step = 2^(exp - mantissa_bits)
+    # For subnormals: step = 2^(min_exp - mantissa_bits)
+    eff_exp = torch.where(is_subnormal, torch.full_like(exp, min_exp), exp)
+    step = torch.pow(2.0, eff_exp - mantissa_bits)
+
+    # Discretize mantissa via round-to-nearest
+    q_abs = torch.round(abs_x / step) * step
+
+    # Re-apply clamp
+    q_abs = torch.clamp(q_abs, max=max_val)
+
+    # Zero out true zeros
+    q_abs = torch.where(is_zero, torch.zeros_like(q_abs), q_abs)
+
+    return sign * q_abs
+
+
+def quantize_fp8(
+    x: torch.Tensor,
+    format: FP8Format = FP8Format.E4M3,
+    scale: Optional[torch.Tensor] = None,
+    eps: float = 1e-8,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Quantize floating-point tensor to FP8 with dynamic or user-specified scaling.
+
+    Computes:
+        scale = max(|x|) / max_representable_fp8 (if scale is None)
+        scaled_x = x / scale
+        q_fp8 = quantize(scaled_x)
+
+    Args:
+        x: Input tensor.
+        format: FP8Format.E4M3 or FP8Format.E5M2.
+        scale: Optional precomputed scale factor tensor.
+        eps: Small value to avoid division by zero.
+
+    Returns:
+        Tuple of (quantized_fp8_tensor, scale).
+    """
+    max_repr = FP8_MAX_VALUES[format]
+
+    if scale is None:
+        max_val = x.abs().max()
+        scale = max_val / float(max_repr)
+        scale = torch.clamp(scale, min=eps)
+
+    scaled_x = x / scale
+
+    # Check if native torch float8 dtype is supported and available
+    native_dtype = None
+    if hasattr(torch, "float8_e4m3fn") and format == FP8Format.E4M3:
+        native_dtype = torch.float8_e4m3fn
+    elif hasattr(torch, "float8_e5m2") and format == FP8Format.E5M2:
+        native_dtype = torch.float8_e5m2
+
+    if native_dtype is not None:
+        try:
+            # Native PyTorch conversion
+            q = scaled_x.to(native_dtype)
+            return q, scale
+        except Exception:
+            pass
+
+    # Analytical fallback emulation
+    if format == FP8Format.E4M3:
+        q = _emulate_fp8(
+            scaled_x,
+            exp_bits=4,
+            mantissa_bits=3,
+            bias=7,
+            max_val=448.0,
+            has_inf=False,
+        )
+    else:
+        q = _emulate_fp8(
+            scaled_x,
+            exp_bits=5,
+            mantissa_bits=2,
+            bias=15,
+            max_val=57344.0,
+            has_inf=True,
+        )
+
+    return q, scale
+
+
+def dequantize_fp8(
+    q_fp8: torch.Tensor,
+    scale: torch.Tensor,
+    target_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Dequantize an FP8 tensor back to target floating-point precision.
+
+    Computes:
+        x_hat = q_fp8 * scale
+
+    Args:
+        q_fp8: Quantized FP8 tensor.
+        scale: Scale factor tensor.
+        target_dtype: Desired output floating-point precision (e.g. float16, float32).
+
+    Returns:
+        Dequantized tensor.
+    """
+    return q_fp8.to(target_dtype) * scale.to(target_dtype)

quantbench/quant/packing/__init__.py

@@ -0,0 +1,27 @@
+"""Bit-packing runtimes for sub-byte quantization formats.
+
+Provides vectorized bit-shift packing and unpacking for:
+- INT4 (4-bit integers: unsigned [0, 15] or signed [-8, 7]) -> 2 values per uint8 byte.
+- Ternary (BitNet 1.58b: {-1, 0, +1}) -> 4 values per uint8 byte.
+- 2:4 Structured Sparsity Metadata -> two 2-bit indices per 4-bit nibble.
+"""
+
+from quantbench.quant.packing.int4_pack import pack_int4, unpack_int4
+from quantbench.quant.packing.sparse_pack import (
+    decompress_2_4,
+    pack_2_4_indices,
+    prune_and_compress_2_4,
+    unpack_2_4_indices,
+)
+from quantbench.quant.packing.ternary_pack import pack_ternary, unpack_ternary
+
+__all__ = [
+    "pack_int4",
+    "unpack_int4",
+    "pack_ternary",
+    "unpack_ternary",
+    "pack_2_4_indices",
+    "unpack_2_4_indices",
+    "prune_and_compress_2_4",
+    "decompress_2_4",
+]

quantbench/quant/packing/int4_pack.py

@@ -0,0 +1,169 @@
+"""Vectorized INT4 bit-packing runtime for sub-byte quantization.
+
+Packs two 4-bit values (either unsigned [0, 15] or signed [-8, 7]) into a single uint8
+byte container, delivering a 4x memory footprint reduction over FP16 and 2x over INT8.
+
+Bit layout within each packed uint8 byte:
+- Bits [0:3] (Low nibble): Even element index (2k)
+- Bits [4:7] (High nibble): Odd element index (2k + 1)
+
+Mathematical mapping:
+    byte = (v_even & 0x0F) | ((v_odd & 0x0F) << 4)
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def pack_nibbles(low: int, high: int) -> int:
+    """Pack two 4-bit integer scalars into a single 8-bit unsigned byte.
+
+    Args:
+        low: Integer value for low nibble (bits 0..3), range [0, 15] or [-8, 7].
+        high: Integer value for high nibble (bits 4..7), range [0, 15] or [-8, 7].
+
+    Returns:
+        Unsigned 8-bit integer in range [0, 255].
+    """
+    u_low = low & 0x0F
+    u_high = high & 0x0F
+    return (u_low) | (u_high << 4)
+
+
+def unpack_nibbles(byte: int, signed: bool = False) -> tuple[int, int]:
+    """Unpack a single 8-bit unsigned byte into two 4-bit integer scalars.
+
+    Args:
+        byte: Unsigned 8-bit integer in range [0, 255].
+        signed: If True, sign-extends 4-bit values into signed range [-8, 7].
+            If False, returns unsigned values in range [0, 15].
+
+    Returns:
+        Tuple of (low_nibble, high_nibble).
+    """
+    low = byte & 0x0F
+    high = (byte >> 4) & 0x0F
+    if signed:
+        if low >= 8:
+            low -= 16
+        if high >= 8:
+            high -= 16
+    return low, high
+
+
+def pack_int4(
+    tensor: torch.Tensor,
+    signed: bool = False,
+    dim: int = -1,
+) -> tuple[torch.Tensor, int]:
+    """Vectorized bit-shift packing of INT4 values into a uint8 tensor.
+
+    Args:
+        tensor: Integer PyTorch tensor containing 4-bit values.
+            If signed=False: values must be in [0, 15].
+            If signed=True: values must be in [-8, 7].
+        signed: Whether the input values are signed two's complement.
+        dim: Dimension along which to pack adjacent pairs. Defaults to last dimension (-1).
+
+    Returns:
+        Tuple of (packed_uint8_tensor, original_dim_length).
+        The packed tensor will have size ceil(original_length / 2) along `dim`.
+    """
+    if not tensor.dtype.is_floating_point and tensor.dtype != torch.bool:
+        t_int = tensor.to(torch.int32)
+    else:
+        t_int = tensor.round().to(torch.int32)
+
+    # Validate bounds
+    if signed:
+        if (t_int < -8).any() or (t_int > 7).any():
+            raise ValueError("Signed INT4 values must fall strictly within range [-8, 7].")
+    else:
+        if (t_int < 0).any() or (t_int > 15).any():
+            raise ValueError("Unsigned INT4 values must fall strictly within range [0, 15].")
+
+    # Normalize dim
+    dim = dim if dim >= 0 else tensor.ndim + dim
+    orig_len = tensor.shape[dim]
+
+    # Convert to 4-bit unsigned representation
+    u4 = t_int & 0x0F
+
+    # Pad by 1 if odd length along dim
+    if orig_len % 2 != 0:
+        pad_shape = list(u4.shape)
+        pad_shape[dim] = 1
+        pad_tensor = torch.zeros(pad_shape, dtype=u4.dtype, device=u4.device)
+        u4 = torch.cat([u4, pad_tensor], dim=dim)
+
+    # Slice into even and odd indices
+    # We move target dim to last position for easy contiguous slicing
+    if dim != u4.ndim - 1:
+        u4 = u4.transpose(dim, -1)
+
+    low = u4[..., 0::2]
+    high = u4[..., 1::2]
+
+    packed = (low & 0x0F) | ((high & 0x0F) << 4)
+    packed_uint8 = packed.to(torch.uint8)
+
+    if dim != tensor.ndim - 1:
+        packed_uint8 = packed_uint8.transpose(dim, -1)
+
+    return packed_uint8, orig_len
+
+
+def unpack_int4(
+    packed: torch.Tensor,
+    original_len: int | None = None,
+    signed: bool = False,
+    dim: int = -1,
+) -> torch.Tensor:
+    """Vectorized unpacking of a uint8 packed tensor back into INT4 values.
+
+    Args:
+        packed: Tensor of dtype torch.uint8 containing packed pairs.
+        original_len: Optional original dimension length before packing.
+            Used to remove padding if original length was odd.
+        signed: If True, sign-extends values to signed range [-8, 7].
+            If False, returns unsigned values in [0, 15].
+        dim: Dimension along which elements were packed. Defaults to -1.
+
+    Returns:
+        Unpacked tensor with dtype torch.int8 (if signed) or torch.uint8 (if unsigned),
+        restored to the original length along `dim`.
+    """
+    if packed.dtype != torch.uint8:
+        raise TypeError(f"Packed tensor must have dtype torch.uint8, got {packed.dtype}.")
+
+    dim = dim if dim >= 0 else packed.ndim + dim
+    if dim != packed.ndim - 1:
+        packed = packed.transpose(dim, -1)
+
+    # Extract low nibble and high nibble
+    p32 = packed.to(torch.int32)
+    low = p32 & 0x0F
+    high = (p32 >> 4) & 0x0F
+
+    # Interleave low and high along a new sub-dimension
+    # Shape becomes (*batch_dims, packed_len, 2) -> flatten to (*batch_dims, packed_len * 2)
+    interleaved = torch.stack([low, high], dim=-1)
+    flat_shape = list(interleaved.shape[:-2]) + [interleaved.shape[-2] * 2]
+    unpacked = interleaved.view(flat_shape)
+
+    # Truncate padding if original_len is specified
+    if original_len is not None:
+        unpacked = unpacked[..., :original_len]
+
+    if signed:
+        # Values in 8..15 have the sign bit set; subtract 16
+        unpacked = torch.where(unpacked >= 8, unpacked - 16, unpacked)
+        unpacked = unpacked.to(torch.int8)
+    else:
+        unpacked = unpacked.to(torch.uint8)
+
+    if dim != packed.ndim - 1:
+        unpacked = unpacked.transpose(dim, -1)
+
+    return unpacked