quantbench-engine 0.1.0__py3-none-any.whl

quantbench/quant/affine.py

@@ -0,0 +1,188 @@
+"""Uniform Affine and Symmetric AbsMax Quantization Primitives.
+
+Implements the fundamental mathematical mappings between continuous floating-point
+representations and discrete integer quantization grids.
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+
+
+def quantize_affine(
+    x: torch.Tensor,
+    num_bits: int = 8,
+    dim: Optional[Union[int, Tuple[int, ...]]] = None,
+    signed: bool = False,
+    eps: float = 1e-8,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Quantize a floating-point tensor using uniform affine (asymmetric) quantization.
+
+    Maps [x_min, x_max] to integer range [q_min, q_max] using:
+        scale = (x_max - x_min) / (2^num_bits - 1)
+        zero_point = round(-x_min / scale) + q_min
+        q = clamp(round(x / scale) + zero_point, q_min, q_max)
+
+    Args:
+        x: Input floating-point tensor.
+        num_bits: Target integer precision (typically 4 or 8).
+        dim: Dimension(s) along which to reduce for per-channel scaling.
+             If None, computes a single scalar per-tensor.
+        signed: Whether the integer grid is signed (e.g. [-128, 127]) or unsigned ([0, 255]).
+        eps: Small epsilon to prevent division by zero for constant tensors.
+
+    Returns:
+        Tuple of (quantized_tensor, scale, zero_point).
+    """
+    if signed:
+        q_min = -(1 << (num_bits - 1))
+        q_max = (1 << (num_bits - 1)) - 1
+    else:
+        q_min = 0
+        q_max = (1 << num_bits) - 1
+
+    if dim is None:
+        x_min = x.min()
+        x_max = x.max()
+    else:
+        x_min = x.amin(dim=dim, keepdim=True)
+        x_max = x.amax(dim=dim, keepdim=True)
+
+    # Standard zero-preserving affine quantization
+    # (ensures 0.0 maps exactly to an integer zero_point)
+    zeros = torch.zeros_like(x_min)
+    x_min_adj = torch.minimum(x_min, zeros)
+    x_max_adj = torch.maximum(x_max, zeros)
+
+    # Calculate scale factor
+    scale = (x_max_adj - x_min_adj) / float(q_max - q_min)
+    scale = torch.clamp(scale, min=eps)
+
+    # Calculate zero-point
+    zero_point = torch.round(-x_min_adj / scale) + q_min
+    zero_point = torch.clamp(zero_point, min=q_min, max=q_max)
+
+    # Quantize and clamp
+    q = torch.round(x / scale) + zero_point
+    q = torch.clamp(q, min=q_min, max=q_max)
+
+    # Cast to integer representation
+    dtype = torch.int8 if signed else torch.uint8
+    if num_bits <= 8:
+        q = q.to(dtype)
+    else:
+        q = q.to(torch.int32)
+
+    return q, scale, zero_point
+
+
+def dequantize_affine(
+    q: torch.Tensor,
+    scale: torch.Tensor,
+    zero_point: torch.Tensor,
+) -> torch.Tensor:
+    """Dequantize an affine-quantized integer tensor back to floating-point.
+
+    Computes:
+        x_hat = (q - zero_point) * scale
+
+    Args:
+        q: Integer quantized tensor.
+        scale: Floating-point scale tensor (broadcastable to q).
+        zero_point: Integer or floating-point zero-point tensor.
+
+    Returns:
+        Reconstructed floating-point tensor.
+    """
+    return (q.to(scale.dtype) - zero_point.to(scale.dtype)) * scale
+
+
+def quantize_absmax(
+    x: torch.Tensor,
+    num_bits: int = 8,
+    dim: Optional[Union[int, Tuple[int, ...]]] = None,
+    symmetric_range: bool = True,
+    eps: float = 1e-8,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Quantize a floating-point tensor using symmetric AbsMax quantization (zero_point = 0).
+
+    Maps [-max(|x|), max(|x|)] to signed integer range [-q_max, q_max] using:
+        scale = max(|x|) / q_max
+        q = clamp(round(x / scale), q_min, q_max)
+
+    Args:
+        x: Input floating-point tensor.
+        num_bits: Target bit precision (typically 4 or 8).
+        dim: Dimension(s) along which to reduce for per-channel scaling.
+             If None, computes per-tensor scale.
+        symmetric_range: If True, restricts range to [-(2^(b-1)-1), 2^(b-1)-1] to preserve
+                         exact symmetry around zero. If False, allows [-2^(b-1), 2^(b-1)-1].
+        eps: Epsilon to avoid division by zero.
+
+    Returns:
+        Tuple of (quantized_tensor, scale).
+    """
+    q_max = (1 << (num_bits - 1)) - 1
+    q_min = -q_max if symmetric_range else -(1 << (num_bits - 1))
+
+    if dim is None:
+        max_val = x.abs().max()
+    else:
+        max_val = x.abs().amax(dim=dim, keepdim=True)
+
+    scale = max_val / float(q_max)
+    scale = torch.clamp(scale, min=eps)
+
+    q = torch.round(x / scale)
+    q = torch.clamp(q, min=q_min, max=q_max)
+
+    if num_bits <= 8:
+        q = q.to(torch.int8)
+    else:
+        q = q.to(torch.int32)
+
+    return q, scale
+
+
+def dequantize_absmax(
+    q: torch.Tensor,
+    scale: torch.Tensor,
+) -> torch.Tensor:
+    """Dequantize an AbsMax-quantized tensor back to floating-point.
+
+    Computes:
+        x_hat = q * scale
+
+    Args:
+        q: Integer quantized tensor.
+        scale: Floating-point scale tensor (broadcastable to q).
+
+    Returns:
+        Reconstructed floating-point tensor.
+    """
+    return q.to(scale.dtype) * scale
+
+
+def check_quantization_invariants(
+    x: torch.Tensor,
+    x_hat: torch.Tensor,
+    scale: torch.Tensor,
+    tolerance_factor: float = 0.51,
+) -> bool:
+    """Verify the fundamental quantization invariant: error <= scale / 2.
+
+    For uniform round-to-nearest quantization, the maximum discretization error
+    for any element within the clamp range is bounded by scale / 2.
+
+    Args:
+        x: Original floating-point tensor.
+        x_hat: Dequantized floating-point tensor.
+        scale: Scale tensor used for quantization.
+        tolerance_factor: Discretization error threshold multiplier (default 0.51).
+
+    Returns:
+        True if all elements within dynamic range satisfy the invariant.
+    """
+    err = (x - x_hat).abs()
+    max_allowed_err = scale * tolerance_factor
+    return bool(torch.all(err <= max_allowed_err + 1e-6))

quantbench/quant/algorithms/__init__.py

@@ -0,0 +1,41 @@
+# 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.
+
+"""Advanced calibration algorithms for post-training quantization."""
+
+from quantbench.quant.algorithms.awq import (
+    awq_quantize_weights,
+    compute_activation_channel_means,
+    search_best_awq_scale,
+)
+from quantbench.quant.algorithms.gptq import (
+    compute_activation_hessian,
+    gptq_quantize_weights,
+)
+from quantbench.quant.algorithms.smoothquant import (
+    compute_smoothquant_scale,
+    extract_activation_channel_scales,
+    smooth_linear_weights,
+)
+
+__all__ = [
+    "compute_smoothquant_scale",
+    "extract_activation_channel_scales",
+    "smooth_linear_weights",
+    "compute_activation_channel_means",
+    "search_best_awq_scale",
+    "awq_quantize_weights",
+    "compute_activation_hessian",
+    "gptq_quantize_weights",
+]

quantbench/quant/algorithms/awq.py

@@ -0,0 +1,184 @@
+# 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.
+
+"""AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration.
+
+Mathematical Foundation:
+    Not all weights in a neural network are equally critical to perplexity.
+    AWQ discovers that protecting the top ~1% of salient weights avoids perplexity
+    degradation even at aggressive INT4/INT3 regimes. Crucially, salient weights
+    directly correspond to channels that process activations with high magnitudes:
+        s_X = (1 / N) * sum(|X|)   (mean absolute activation per input channel)
+
+    Instead of maintaining mixed-precision floating-point formats (which incur runtime
+    overhead and degrade hardware SIMD execution efficiency), AWQ applies a per-channel
+    protection scale s:
+        W' = W * diag(s)
+        X' = diag(s)^-1 * X
+
+    When W' is quantized to sub-byte integers (e.g. INT4 blockwise), the relative quantization
+    error on salient channels is inversely proportional to s_j:
+        Delta W_j / W_j ~ 1 / s_j
+
+    AWQ searches for the optimal per-channel scale s_j = (s_{X, j})^alpha over a grid of
+    alpha in [0.0, 1.0] by directly minimizing the Frobenius output reconstruction error:
+        min_alpha || W * X - [Quantize(W * diag(s)) * diag(s)^-1] * X ||_F^2
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Optional, Tuple
+
+import torch
+
+from quantbench.quant.blockwise import dequantize_blockwise, quantize_blockwise
+
+
+def compute_activation_channel_means(
+    calibration_inputs: list[torch.Tensor] | torch.Tensor,
+) -> torch.Tensor:
+    """Calculate mean absolute activation magnitude per input channel.
+
+    Args:
+        calibration_inputs: Tensor of shape (..., in_features) or list of tensors.
+
+    Returns:
+        Tensor of shape (in_features,) containing per-channel mean absolute activation.
+    """
+    if isinstance(calibration_inputs, torch.Tensor):
+        calibration_inputs = [calibration_inputs]
+
+    total_sum: Optional[torch.Tensor] = None
+    total_tokens = 0
+
+    for batch in calibration_inputs:
+        flat = batch.detach().view(-1, batch.shape[-1]).abs()
+        batch_sum = flat.sum(dim=0)
+        n_tokens = flat.shape[0]
+
+        if total_sum is None:
+            total_sum = batch_sum
+        else:
+            total_sum = total_sum + batch_sum
+        total_tokens += n_tokens
+
+    if total_sum is None or total_tokens == 0:
+        raise ValueError("calibration_inputs cannot be empty.")
+
+    return total_sum / total_tokens
+
+
+def search_best_awq_scale(
+    weight: torch.Tensor,
+    act_means: torch.Tensor,
+    calibration_x: torch.Tensor,
+    quant_fn: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+    n_grid: int = 20,
+    max_alpha: float = 1.0,
+) -> Tuple[torch.Tensor, float]:
+    """Perform grid search for optimal AWQ protection scale alpha.
+
+    Args:
+        weight: Float tensor of shape (out_features, in_features).
+        act_means: Channel-wise activation mean of shape (in_features,).
+        calibration_x: Sample calibration activations of shape (N_samples, in_features).
+        quant_fn: Quantize-and-dequantize callable. Defaults to INT4 blockwise (group_size=128).
+        n_grid: Number of search points in alpha grid [0, max_alpha].
+        max_alpha: Maximum alpha value for search.
+
+    Returns:
+        Tuple of (optimal_scale_vector, best_alpha):
+            - optimal_scale_vector: Per-channel scale s in R^{in_features}.
+            - best_alpha: Optimal alpha hyperparameter.
+    """
+    if quant_fn is None:
+        def default_quant(w: torch.Tensor) -> torch.Tensor:
+            q_b, s_b, zp_b, meta = quantize_blockwise(
+                w, group_size=min(128, w.shape[1]), num_bits=4
+            )
+            return dequantize_blockwise(q_b, s_b, zp_b, meta)
+        quant_fn = default_quant
+
+    # Reference unquantized linear output
+    # calibration_x: (N, in_features), weight: (out_features, in_features)
+    # y_ref: (N, out_features)
+    y_ref = calibration_x @ weight.t()
+
+    best_error = float("inf")
+    best_scale = torch.ones(weight.shape[1], device=weight.device, dtype=weight.dtype)
+    best_alpha = 0.0
+
+    act_means = act_means.clamp(min=1e-5)
+
+    alphas = torch.linspace(0.0, max_alpha, steps=n_grid)
+
+    for alpha_val in alphas:
+        alpha = float(alpha_val.item())
+        # Candidate scale s = act_means^alpha
+        cand_scale = act_means.pow(alpha)
+        # Normalize scale by geometric mean of min/max to maintain weight scale stability
+        geom_mean = (cand_scale.max() * cand_scale.min()).sqrt().clamp(min=1e-5)
+        cand_scale = (cand_scale / geom_mean).clamp(min=1e-4)
+
+        # Scale weights: W' = W * diag(s)
+        # weight: (out, in), cand_scale: (in,) -> broadcast along rows
+        w_scaled = weight * cand_scale.unsqueeze(0)
+
+        # Quantize and dequantize
+        w_dequant = quant_fn(w_scaled)
+
+        # Inverse scale: W_cand = w_dequant * diag(s)^-1
+        w_cand = w_dequant / cand_scale.unsqueeze(0)
+
+        # Output error || y_ref - x @ w_cand.t() ||_F^2
+        y_cand = calibration_x @ w_cand.t()
+        err = torch.norm(y_ref - y_cand, p="fro").item()
+
+        if err < best_error:
+            best_error = err
+            best_scale = cand_scale
+            best_alpha = alpha
+
+    return best_scale, best_alpha
+
+
+def awq_quantize_weights(
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    quant_fn: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Apply optimal AWQ scale and quantize weight matrix.
+
+    Args:
+        weight: Original weight of shape (out_features, in_features).
+        scale: Optimal per-channel protection scale of shape (in_features,).
+        quant_fn: Quantization-dequantization function.
+
+    Returns:
+        Tuple of (dequantized_weight, scale).
+    """
+    if quant_fn is None:
+        def default_quant(w: torch.Tensor) -> torch.Tensor:
+            q_b, s_b, zp_b, meta = quantize_blockwise(
+                w, group_size=min(128, w.shape[1]), num_bits=4
+            )
+            return dequantize_blockwise(q_b, s_b, zp_b, meta)
+        quant_fn = default_quant
+
+    w_scaled = weight * scale.unsqueeze(0)
+    w_dequant = quant_fn(w_scaled)
+    # Invert scale so layer can be run directly with unscaled input:
+    # y = x @ (w_dequant / s).t()
+    w_effective = w_dequant / scale.unsqueeze(0)
+    return w_effective, scale

quantbench/quant/algorithms/gptq.py

@@ -0,0 +1,181 @@
+# 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.
+
+"""GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers.
+
+Mathematical Foundation:
+    Post-training quantization aims to minimize the layer-wise Frobenius reconstruction error:
+        argmin_{W_hat} || W * X - W_hat * X ||_F^2
+
+    Using a second-order Taylor expansion around the continuous weights W, the loss
+    increase induced by quantizing weights is governed by the input activation Hessian:
+        H = 2 * X * X^T in R^{d_in x d_in}
+
+    Optimal Brain Surgeon (OBS) establishes that rounding column j to W_hat_{:, j} produces
+    an error delta w_j = W_{:, j} - W_hat_{:, j}. The optimal linear compensation applied to
+    the remaining unquantized columns F = {j+1, ..., d_in - 1} is given by:
+        Delta W_{:, F} = - (W_{:, j} - W_hat_{:, j}) / [H^-1]_{j, j} * [H^-1]_{j, F}
+
+    To prevent numerical divergence from zero/near-zero eigenvalues, the empirical Hessian
+    is regularized with diagonal damping before inversion:
+        H_damp = H + lambda * mean(diag(H)) * I    (typically lambda = 0.01)
+
+    GPTQ implements a vectorized block-wise update (block size B, e.g. 32 or 128)
+    converting column-by-column corrections into dense matrix multiplications (GEMM)
+    while avoiding O(d^3) naive Cholesky recomputations.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from quantbench.quant.affine import dequantize_absmax, quantize_absmax
+from quantbench.quant.blockwise import dequantize_blockwise, quantize_blockwise
+
+
+def compute_activation_hessian(
+    calibration_inputs: list[torch.Tensor] | torch.Tensor,
+    in_features: int,
+    device: Optional[torch.device] = None,
+) -> torch.Tensor:
+    """Compute empirical input activation Hessian H = 2 * X * X^T across calibration batches.
+
+    Args:
+        calibration_inputs: Tensor of shape (..., in_features) or list of tensors.
+        in_features: Number of input features (columns of weight matrix).
+        device: Torch device for computation.
+
+    Returns:
+        Tensor of shape (in_features, in_features) representing the accumulated Hessian.
+    """
+    if isinstance(calibration_inputs, torch.Tensor):
+        calibration_inputs = [calibration_inputs]
+
+    dev = device or calibration_inputs[0].device
+    H = torch.zeros((in_features, in_features), device=dev, dtype=torch.float32)
+    n_samples = 0
+
+    for batch in calibration_inputs:
+        # Reshape to (N_tokens, in_features)
+        x = batch.to(dev).float().view(-1, in_features)
+        # H += 2 * x^T * x
+        H += 2.0 * (x.t() @ x)
+        n_samples += x.shape[0]
+
+    if n_samples == 0:
+        raise ValueError("calibration_inputs contains 0 tokens.")
+
+    # Normalize by total sample count
+    return H / n_samples
+
+
+def gptq_quantize_weights(
+    weight: torch.Tensor,
+    H: torch.Tensor,
+    bits: int = 4,
+    group_size: int = 128,
+    block_size: int = 32,
+    damping_factor: float = 0.01,
+) -> torch.Tensor:
+    """Quantize weight matrix using GPTQ second-order Hessian error compensation.
+
+    Args:
+        weight: Float tensor of shape (out_features, in_features).
+        H: Input Hessian matrix of shape (in_features, in_features).
+        bits: Target integer bit-width (e.g. 4 or 8).
+        group_size: Quantization block/group size for scales (-1 for per-channel).
+        block_size: GPTQ column block update size (e.g. 32).
+        damping_factor: Diagonal damping regularization factor lambda.
+
+    Returns:
+        Quantized and dequantized weight tensor in continuous precision with
+        second-order error compensation applied.
+    """
+    out_features, in_features = weight.shape
+    device = weight.device
+    dtype = weight.dtype
+
+    # Work in float32 for numerical stability
+    W = weight.clone().float()
+    H = H.clone().float()
+
+    # Damping regularization: H_reg = H + lambda * mean(diag(H)) * I
+    diag_mean = torch.mean(torch.diag(H)).item()
+    if diag_mean <= 0.0:
+        diag_mean = 1.0
+    damp = damping_factor * diag_mean
+    H += damp * torch.eye(in_features, device=device)
+
+    # Invert Hessian via Cholesky decomposition: H = L * L^T -> H^-1 = L^-T * L^-1
+    try:
+        L = torch.linalg.cholesky(H)
+        L_inv = torch.linalg.inv(L)
+        H_inv = L_inv.t() @ L_inv
+    except RuntimeError:
+        # Fallback to pseudo-inverse if Cholesky fails due to singular sub-blocks
+        H_inv = torch.linalg.pinv(H)
+
+    # Quantization helper for column vectors / sub-blocks
+    def quantize_dequantize_cols(col_w: torch.Tensor) -> torch.Tensor:
+        # col_w: (out_features, num_cols)
+        if group_size > 0:
+            q_b, s_b, zp_b, meta = quantize_blockwise(
+                col_w, group_size=group_size, num_bits=bits
+            )
+            return dequantize_blockwise(q_b, s_b, zp_b, meta)
+        else:
+            q, scale = quantize_absmax(col_w, num_bits=bits, dim=0)
+            return dequantize_absmax(q, scale)
+
+    Q = torch.zeros_like(W)
+
+    # Iterate through column blocks
+    for col_start in range(0, in_features, block_size):
+        col_end = min(col_start + block_size, in_features)
+        count = col_end - col_start
+
+        # Sub-block weights and sub-block Hessian inverse
+        W_block = W[:, col_start:col_end].clone()
+        Q_block = torch.zeros_like(W_block)
+        err_block = torch.zeros_like(W_block)
+
+        H_inv_block = H_inv[col_start:col_end, col_start:col_end]
+
+        # Column-by-column inside the block
+        for i in range(count):
+            w_col = W_block[:, i]
+            d = H_inv_block[i, i]
+
+            # Quantize individual column
+            q_col = quantize_dequantize_cols(w_col.unsqueeze(1)).squeeze(1)
+            Q_block[:, i] = q_col
+
+            # OBS column error
+            err = (w_col - q_col) / d.clamp(min=1e-7)
+            err_block[:, i] = err
+
+            # Propagate error to remaining columns in the current block
+            if i + 1 < count:
+                W_block[:, i + 1 :] -= err.unsqueeze(1) * H_inv_block[i, i + 1 :].unsqueeze(0)
+
+        Q[:, col_start:col_end] = Q_block
+
+        # Propagate block error to all remaining columns outside the block (GEMM)
+        if col_end < in_features:
+            # W[:, col_end:] -= err_block @ H_inv[col_start:col_end, col_end:]
+            W[:, col_end:] -= err_block @ H_inv[col_start:col_end, col_end:]
+
+    return Q.to(dtype=dtype)