AWQ Deep Dive: Activation-Aware Weight Quantization — The Algorithm Step by Step

In 2023, researchers at MIT discovered something remarkable: not all weights in a neural network matter equally, and the difference is extreme—about 1% of weight channels, when quantized poorly, cause 100x more output error than the other 99% combined. These “salient” channels correspond to weights that multiply large activation values, amplifying even small quantization errors into catastrophic output degradation. AWQ (Activation-Aware Weight Quantization) exploits this asymmetry by scaling up those critical 1% of weights before quantization, giving them finer INT4 precision while the rest get coarser bins. The result: INT4 models that match GPTQ quality while being faster to quantize and slightly better on perplexity benchmarks.

AWQ is a weight-only quantization method that achieves near-lossless INT4 quality by protecting salient weight channels — those that correspond to large activation magnitudes. The key insight is that not all weights are equally important: weights multiplied by large activations contribute more to the output and should be quantized more carefully.

AWQ does not modify the quantization grid (it still uses uniform INT4). Instead, it applies per-channel scaling to the weights before quantization, enlarging the salient channels so they occupy more of the INT4 range. This is mathematically equivalent to SmoothQuant’s scaling migration, but applied with a different objective: minimizing quantization error on the output rather than balancing activation and weight ranges.

This post implements AWQ from scratch, step by step.

The Core Problem

Consider a linear layer $Y = XW^T$ where $W \in \mathbb{R}^{N \times K}$ and $X \in \mathbb{R}^{M \times K}$. After quantizing $W$ to $\hat{W}$, the output error is:

$\Delta Y = X(W - \hat{W})^T = X \cdot E^T$

where $E = W - \hat{W}$ is the weight quantization error matrix. The output error for a single output neuron $i$ is:

$\Delta Y_i = \sum_j X_j \cdot E_{ij}$

The expected squared error is:

$\mathbb{E}[\Delta Y_i^2] = \sum_j \mathbb{E}[X_j^2] \cdot E_{ij}^2$

assuming independence between channels. This shows that the output error contribution of weight $W_{ij}$ is proportional to $\mathbb{E}[X_j^2]$ — the mean squared activation of channel $j$.

Channels with large $\mathbb{E}[X_j^2]$ are “salient”: their quantization errors are amplified by the activation magnitude. AWQ reduces the quantization error on these salient channels by scaling them up before quantization.

The AWQ Scaling Trick

For each input channel $j$, AWQ applies a scaling factor $s_j$:

$\hat{W}_{:,j} = \text{Quantize}(W_{:,j} \cdot s_j) / s_j$

The weight is multiplied by $s_j$ before quantization and divided by $s_j$ after dequantization. The mathematical value is preserved (if there were no quantization), but the quantization grid has shifted: channel $j$ now occupies $s_j$ times more of the integer range.

To preserve the output, the activation must be divided by $s_j$:

$Y = X W^T = (X \cdot S^{-1})(S \cdot W)^T$

This is identical to SmoothQuant’s transformation. The difference is how $s_j$ is chosen:

• SmoothQuant: $s_j = (\max |X_j|)^\alpha / (\max |W_j|)^{1-\alpha}$ (balances activation/weight ranges)
• AWQ: $s_j$ is chosen to minimize the quantization error on the output (grid search)

import torch
import numpy as np

def awq_scaling_intuition(W, X_calibration, bits=4):
    """Demonstrate why scaling helps salient channels.

    Without scaling: all channels get the same quantization grid spacing.
    With scaling: salient channels get finer grid spacing.
    """
    K = W.shape[1]
    qmax = 2 ** (bits - 1) - 1

    # Compute per-channel activation importance
    act_importance = (X_calibration ** 2).mean(dim=0)  # (K,)

    # Without scaling: per-channel weight quantization
    w_max = W.abs().amax(dim=0)  # (K,)
    step_size = w_max / qmax     # Quantization step per channel

    # Output error per channel (proportional to importance * step^2)
    error_no_scale = act_importance * (step_size ** 2 / 12)

    # With scaling (s=2 for top channels): step size halved for important channels
    s = torch.ones(K)
    top_channels = act_importance.argsort(descending=True)[:int(K * 0.01)]
    s[top_channels] = 2.0

    # After scaling: weights are multiplied by s, so max increases
    # But the group scale absorbs this -- the key is within a group,
    # the scaled channel gets more levels relative to unscaled channels
    scaled_step = step_size / s
    error_with_scale = act_importance * (scaled_step ** 2 / 12)

    improvement = error_no_scale.sum() / error_with_scale.sum()
    print(f"Total output error reduction: {improvement:.2f}x")

    return improvement

Step 1: Compute Channel Importance from Calibration Data

AWQ uses calibration data to estimate $\mathbb{E}[X_j^2]$ for each input channel:

def compute_channel_importance(model, calibration_dataloader, num_samples=128):
    """Compute per-channel activation importance for each linear layer.

    Returns dict mapping layer_name -> importance tensor of shape (K,)
    where K is the input dimension.
    """
    importance = {}
    hooks = []

    def make_hook(name):
        def hook(module, input_data, output):
            x = input_data[0].detach().float()
            x_flat = x.reshape(-1, x.shape[-1])  # (tokens, K)

            # Mean squared activation per channel
            batch_importance = (x_flat ** 2).mean(dim=0)  # (K,)

            if name not in importance:
                importance[name] = {'sum': batch_importance, 'count': 1}
            else:
                importance[name]['sum'] += batch_importance
                importance[name]['count'] += 1
        return hook

    for name, mod in model.named_modules():
        if isinstance(mod, torch.nn.Linear):
            hooks.append(mod.register_forward_hook(make_hook(name)))

    model.eval()
    count = 0
    with torch.no_grad():
        for batch in calibration_dataloader:
            if count >= num_samples:
                break
            model(batch['input_ids'].cuda())
            count += batch['input_ids'].shape[0]

    for h in hooks:
        h.remove()

    # Average importance
    channel_importance = {}
    for name, data in importance.items():
        channel_importance[name] = data['sum'] / data['count']

    return channel_importance

Step 2: The Per-Group Scaling Search

AWQ operates within per-group quantization. For each group of $g$ weights (e.g., $g = 128$), it finds optimal per-channel scales within that group.

The search is performed per-group because the quantization scale factor is computed per-group. Scaling a channel within a group affects that group’s scale factor, which in turn affects all other channels in the same group.

def awq_search_scales_per_group(
    W_group,           # (N, g) weights for one group
    X_group,           # (num_tokens, g) calibration activations for this group
    bits=4,
    n_grid=20,         # Number of grid search points
):
    """Search for optimal per-channel scales within one group.

    For each channel j in the group, we search over candidate scales
    s_j and pick the one that minimizes the output reconstruction error.

    The output for this group is: Y_group = X_group @ W_group^T  (N_tokens, N)

    After quantization with scales s:
    Y_hat = (X_group / s) @ Quant(W_group * s)^T

    We search s to minimize ||Y_group - Y_hat||^2.
    """
    N, g = W_group.shape
    qmax = 2 ** (bits - 1) - 1

    # Target output (FP32 ground truth for this group)
    Y_target = X_group @ W_group.T  # (tokens, N)

    # Channel importance for this group
    importance = (X_group ** 2).mean(dim=0)  # (g,)
    importance = importance / importance.max()  # Normalize to [0, 1]

    best_scales = torch.ones(g, device=W_group.device)
    best_error = float('inf')

    # Grid search: try scaling important channels by different factors
    for ratio in torch.linspace(0, 1, n_grid + 1, device=W_group.device)[1:]:
        # Scale = importance ^ ratio
        # ratio=0: no scaling (all s=1)
        # ratio=1: full importance scaling (s proportional to sqrt(importance))
        scales = importance.pow(ratio).clamp(min=1e-4)

        # Normalize so geometric mean is 1 (preserve overall magnitude)
        scales = scales / scales.pow(1.0 / g).prod().pow(1.0 / g)

        # Apply scaling: W_scaled = W * s (per channel)
        W_scaled = W_group * scales.unsqueeze(0)

        # Quantize the scaled weights (per-group: one scale for all g channels)
        w_abs_max = W_scaled.abs().amax(dim=1, keepdim=True)  # (N, 1)
        w_scale = w_abs_max / qmax
        w_scale = w_scale.clamp(min=1e-10)

        W_q = (W_scaled / w_scale).round().clamp(-qmax - 1, qmax)
        W_deq = W_q * w_scale  # Dequantized (still scaled)

        # Undo the channel scaling in dequantized weights
        W_deq_unscaled = W_deq / scales.unsqueeze(0)

        # Compute output with quantized weights
        Y_hat = X_group @ W_deq_unscaled.T

        # Compute error
        error = ((Y_target - Y_hat) ** 2).mean().item()

        if error < best_error:
            best_error = error
            best_scales = scales.clone()

    return best_scales, best_error

Step 3: Apply Scales and Quantize

def awq_quantize_layer(
    linear_layer,
    channel_importance,  # (K,) importance scores
    bits=4,
    group_size=128,
    n_grid=20,
    calibration_X=None,  # (num_tokens, K) calibration activations
):
    """Full AWQ quantization of a single linear layer.

    Args:
        linear_layer: nn.Linear with weight (N, K)
        channel_importance: per-channel importance, shape (K,)
        bits: quantization bits
        group_size: per-group quantization group size
        n_grid: grid search resolution
        calibration_X: calibration activations for error evaluation

    Returns:
        W_q: quantized weights, shape (N, K), dtype int8
        scales: per-group scale factors, shape (N, num_groups)
        channel_scales: AWQ per-channel scales, shape (K,)
    """
    W = linear_layer.weight.data.float()
    N, K = W.shape
    num_groups = K // group_size
    qmax = 2 ** (bits - 1) - 1

    all_channel_scales = torch.ones(K, device=W.device)

    # Process each group independently
    for gi in range(num_groups):
        start = gi * group_size
        end = start + group_size

        W_group = W[:, start:end]  # (N, group_size)
        imp_group = channel_importance[start:end]

        if calibration_X is not None:
            X_group = calibration_X[:, start:end]
        else:
            # Fallback: use importance as proxy
            X_group = torch.randn(128, group_size, device=W.device)
            X_group *= imp_group.sqrt().unsqueeze(0)

        group_scales, _ = awq_search_scales_per_group(
            W_group, X_group,
            bits=bits, n_grid=n_grid
        )
        all_channel_scales[start:end] = group_scales

    # Apply scales to weights
    W_scaled = W * all_channel_scales.unsqueeze(0)  # (N, K)

    # Per-group quantization of scaled weights
    W_grouped = W_scaled.reshape(N, num_groups, group_size)
    group_abs_max = W_grouped.abs().amax(dim=2)  # (N, num_groups)
    group_scales = group_abs_max / qmax
    group_scales = group_scales.clamp(min=1e-10)

    W_q = (W_grouped / group_scales.unsqueeze(2)).round().clamp(
        -(qmax + 1), qmax
    )
    W_q = W_q.reshape(N, K).to(torch.int8)

    return W_q, group_scales, all_channel_scales

def awq_dequantize(W_q, group_scales, channel_scales, group_size):
    """Dequantize AWQ-quantized weights."""
    N, K = W_q.shape
    num_groups = K // group_size

    W_grouped = W_q.float().reshape(N, num_groups, group_size)
    W_deq = W_grouped * group_scales.unsqueeze(2)
    W_deq = W_deq.reshape(N, K)

    # Undo channel scaling
    W_deq = W_deq / channel_scales.unsqueeze(0)

    return W_deq

Step 4: Fuse Scales into the Model

The per-channel scales must be absorbed into the model to avoid runtime overhead. AWQ fuses the scale division into the preceding LayerNorm (for the activation path) and absorbs the scale multiplication into the weights (already done during quantization).

def fuse_awq_scales(model, layer_scales):
    """Fuse AWQ channel scales into the model.

    For each linear layer with AWQ scales s:
    - Weight: already scaled (W * s is quantized)
    - Activation: must be divided by s at runtime

    The activation division X / s is fused into the preceding LayerNorm:
    - LayerNorm: y = gamma * (x - mean) / std + beta
    - After fusion: y = (gamma / s) * (x - mean) / std + (beta / s)
    """
    for layer_name, scales in layer_scales.items():
        # Find the preceding LayerNorm
        parts = layer_name.split('.')
        # Example: model.layers.0.self_attn.q_proj
        # Preceding LN: model.layers.0.input_layernorm (for attn)
        #                model.layers.0.post_attention_layernorm (for MLP)

        layer_idx = None
        for i, part in enumerate(parts):
            if part == 'layers':
                layer_idx = int(parts[i + 1])
                break

        if layer_idx is None:
            continue

        if 'attn' in layer_name:
            ln_name = f"model.layers.{layer_idx}.input_layernorm"
        elif 'mlp' in layer_name:
            ln_name = f"model.layers.{layer_idx}.post_attention_layernorm"
        else:
            continue

        # Get the LayerNorm module
        ln_module = dict(model.named_modules()).get(ln_name)
        if ln_module is not None:
            ln_module.weight.data /= scales
            if ln_module.bias is not None:
                ln_module.bias.data /= scales

Step 5: Full AWQ Pipeline

class AWQQuantizer:
    """Complete AWQ quantization pipeline."""

    def __init__(self, model, bits=4, group_size=128, n_grid=20):
        self.model = model
        self.bits = bits
        self.group_size = group_size
        self.n_grid = n_grid

    def quantize(self, calibration_dataloader, num_samples=128):
        """Full AWQ pipeline: calibrate, search, quantize, fuse."""

        # Step 1: Collect activation statistics
        print("Step 1: Collecting activation statistics...")
        importance = {}
        activations = {}
        hooks = []

        def make_hook(name):
            def hook(module, input_data, output):
                x = input_data[0].detach().float()
                x_flat = x.reshape(-1, x.shape[-1])

                if name not in importance:
                    importance[name] = (x_flat ** 2).mean(dim=0)
                    activations[name] = [x_flat[:64]]  # Store subset
                else:
                    importance[name] += (x_flat ** 2).mean(dim=0)
                    if len(activations[name]) < 4:
                        activations[name].append(x_flat[:64])
            return hook

        for name, mod in self.model.named_modules():
            if isinstance(mod, torch.nn.Linear):
                hooks.append(mod.register_forward_hook(make_hook(name)))

        self.model.eval()
        count = 0
        with torch.no_grad():
            for batch in calibration_dataloader:
                if count >= num_samples:
                    break
                self.model(batch['input_ids'].cuda())
                count += batch['input_ids'].shape[0]

        for h in hooks:
            h.remove()

        # Normalize importance
        for name in importance:
            importance[name] /= count

        # Step 2: Search and quantize each layer
        print("Step 2: Searching optimal scales and quantizing...")
        layer_scales = {}
        quantized_layers = {}

        for name, mod in self.model.named_modules():
            if not isinstance(mod, torch.nn.Linear):
                continue
            if name not in importance:
                continue

            cal_X = torch.cat(activations.get(name, []), dim=0)
            if cal_X.shape[0] == 0:
                cal_X = None

            W_q, group_sc, ch_sc = awq_quantize_layer(
                mod, importance[name],
                bits=self.bits,
                group_size=self.group_size,
                n_grid=self.n_grid,
                calibration_X=cal_X,
            )

            layer_scales[name] = ch_sc
            quantized_layers[name] = {
                'W_q': W_q,
                'group_scales': group_sc,
                'channel_scales': ch_sc,
            }

        # Step 3: Fuse scales into LayerNorms
        print("Step 3: Fusing scales into model...")
        fuse_awq_scales(self.model, layer_scales)

        # Step 4: Replace linear layers with quantized versions
        print("Step 4: Replacing layers...")
        for name, qdata in quantized_layers.items():
            # Create quantized module and replace in model
            pass  # Implementation specific to model architecture

        return quantized_layers

AWQ vs RTN vs GPTQ: The Quality Difference

Why does AWQ outperform RTN? The answer is in the error weighting. RTN minimizes uniform weight MSE. AWQ minimizes activation-weighted output MSE.

def compare_awq_rtn_error(W, X, bits=4, group_size=128):
    """Compare AWQ and RTN on a single group.

    Shows that AWQ reduces output error even though it may increase
    weight MSE.
    """
    N, K = W.shape
    qmax = 2 ** (bits - 1) - 1

    # Ground truth output
    Y_true = X @ W.T

    # RTN: simple round-to-nearest per-group
    num_groups = K // group_size
    W_rtn = torch.zeros_like(W)
    for gi in range(num_groups):
        start = gi * group_size
        end = start + group_size
        g_w = W[:, start:end]
        g_max = g_w.abs().amax(dim=1, keepdim=True)
        g_scale = g_max / qmax
        g_scale = g_scale.clamp(min=1e-10)
        g_q = (g_w / g_scale).round().clamp(-(qmax+1), qmax)
        W_rtn[:, start:end] = g_q * g_scale

    Y_rtn = X @ W_rtn.T
    rtn_output_mse = ((Y_true - Y_rtn) ** 2).mean().item()
    rtn_weight_mse = ((W - W_rtn) ** 2).mean().item()

    # AWQ: activation-aware scaling
    importance = (X ** 2).mean(dim=0)
    W_q_awq, g_scales, ch_scales = awq_quantize_layer(
        torch.nn.Linear(K, N, bias=False).requires_grad_(False),
        importance, bits=bits, group_size=group_size,
        calibration_X=X
    )
    # Manually set weight for the function
    # (simplified -- actual implementation operates on module)

    # ... compute AWQ output error ...

    return {
        'rtn_output_mse': rtn_output_mse,
        'rtn_weight_mse': rtn_weight_mse,
    }

AWQ with Clipping

A further optimization clips the weight range before quantization, shrinking the scale factor to reduce rounding error at the cost of clipping error on outlier weights. AWQ searches for the optimal clipping ratio:

def awq_clip_search(W_group, X_group, bits=4, n_clip_grid=20):
    """Search for optimal weight clipping ratio within a group.

    Instead of using max(|w|) to set the scale, clip at ratio * max(|w|)
    where ratio < 1. This reduces the step size but clips extreme weights.

    The optimal ratio balances clipping error (on large weights) vs
    rounding error (on all weights).
    """
    N, g = W_group.shape
    qmax = 2 ** (bits - 1) - 1

    Y_target = X_group @ W_group.T

    best_ratio = 1.0
    best_error = float('inf')

    for ratio in torch.linspace(0.5, 1.0, n_clip_grid, device=W_group.device):
        # Clip weights
        clip_val = W_group.abs().amax(dim=1, keepdim=True) * ratio
        W_clipped = W_group.clamp(-clip_val, clip_val)

        # Quantize clipped weights
        w_max = W_clipped.abs().amax(dim=1, keepdim=True)
        w_scale = w_max / qmax
        w_scale = w_scale.clamp(min=1e-10)

        W_q = (W_clipped / w_scale).round().clamp(-(qmax + 1), qmax)
        W_deq = W_q * w_scale

        Y_hat = X_group @ W_deq.T
        error = ((Y_target - Y_hat) ** 2).mean().item()

        if error < best_error:
            best_error = error
            best_ratio = ratio.item()

    return best_ratio, best_error

The Relationship Between AWQ and SmoothQuant

AWQ and SmoothQuant share the same mathematical transformation: $Y = (XS^{-1})(SW^T)$. The differences are:

def demonstrate_awq_smoothquant_equivalence():
    """Show that AWQ's scaling is mathematically identical to SmoothQuant."""

    torch.manual_seed(42)
    N, K = 256, 256
    W = torch.randn(N, K) * 0.02
    X = torch.randn(64, K) * 0.5

    # SmoothQuant scaling
    act_max = X.abs().amax(dim=0)
    weight_max = W.abs().amax(dim=0)
    alpha = 0.5
    sq_scales = act_max.pow(alpha) / weight_max.pow(1 - alpha)
    sq_scales = sq_scales.clamp(min=1e-5)

    # SmoothQuant: X_smooth = X / s, W_smooth = W * s
    X_sq = X / sq_scales.unsqueeze(0)
    W_sq = W * sq_scales.unsqueeze(0)
    Y_sq = X_sq @ W_sq.T

    # AWQ scaling (using importance-based s)
    importance = (X ** 2).mean(dim=0)
    awq_scales = importance.pow(0.5)
    awq_scales = awq_scales / awq_scales.mean()  # Normalize

    X_awq = X / awq_scales.unsqueeze(0)
    W_awq = W * awq_scales.unsqueeze(0)
    Y_awq = X_awq @ W_awq.T

    # Both produce identical outputs to the original
    Y_orig = X @ W.T
    print(f"SQ output diff: {(Y_orig - Y_sq).abs().max():.2e}")
    print(f"AWQ output diff: {(Y_orig - Y_awq).abs().max():.2e}")
    # Both should be < 1e-5 (floating point only)

Integration with Inference Kernels

AWQ-quantized models are compatible with the same W4A16 kernels as GPTQ:

# AWQ model loading in vLLM
# The AWQ format stores:
# Quantized INT4 weights (packed)
# Per-group scale factors (FP16)
# Per-group zero points (optional, for asymmetric)
# AWQ channel scales are already fused into LayerNorm and weights

AWQ_MODEL_FORMAT = {
    'qweight': 'packed INT4, shape (K//8, N) or (N, K//8)',
    'qzeros': 'packed INT4 zero points, shape (K//g, N) or (N, K//g)',
    'scales': 'FP16 per-group scales, shape (K//g, N) or (N, K//g)',
}

# vLLM automatically detects AWQ format and routes to Marlin kernel
# if the model is compatible (symmetric, group_size=128, no act_order)

# AutoAWQ library quantization command:
# from awq import AutoAWQForCausalLM
# model = AutoAWQForCausalLM.from_pretrained(model_path)
# model.quantize(
#     tokenizer,
#     quant_config={
#         'zero_point': True,
#         'q_group_size': 128,
#         'w_bit': 4,
#         'version': 'GEMM',
#     }
# )

Practical AWQ Quantization with AutoAWQ

# Using the AutoAWQ library (production implementation)

# Step 1: Install
# pip install autoawq

# Step 2: Quantize
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = "meta-llama/Llama-2-7b-hf"
quant_path = "llama2-7b-awq-w4-g128"

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoAWQForCausalLM.from_pretrained(model_path)

quant_config = {
    "zero_point": True,      # Asymmetric quantization
    "q_group_size": 128,     # Group size
    "w_bit": 4,              # 4-bit weights
    "version": "GEMM",       # GEMM-compatible layout
}

# Quantize (uses calibration data internally)
model.quantize(tokenizer, quant_config=quant_config)

# Save in safetensors format
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

# Step 3: Load for inference (auto-selects best kernel)
# vLLM: vllm serve llama2-7b-awq-w4-g128
# Marlin kernel is used automatically if compatible

Scaling Factor Analysis on Real Models

def analyze_awq_scales(quantized_model):
    """Analyze the AWQ scaling factors across layers."""
    for name, mod in quantized_model.named_modules():
        if not hasattr(mod, 'awq_channel_scales'):
            continue

        scales = mod.awq_channel_scales
        print(f"\n{name}:")
        print(f"  Mean scale: {scales.mean():.4f}")
        print(f"  Max scale: {scales.max():.4f}")
        print(f"  Min scale: {scales.min():.4f}")
        print(f"  Std: {scales.std():.4f}")
        print(f"  Channels > 2x mean: {(scales > 2 * scales.mean()).sum()}")
        print(f"  Channels < 0.5x mean: {(scales < 0.5 * scales.mean()).sum()}")

# Typical findings on Llama-2 7B:
# - 1-3% of channels have scales > 5x the mean (salient channels)
# - Early layers have higher scale variance (more outliers)
# - MLP gate projections have the most uniform scales
# - Attention V projections have the highest scale variance