Module 14: Quantization

Introduction

You have trained a model (Modules 06–08). Every weight in it is a 32-bit float — 4 bytes each. A 7-billion-parameter model is therefore 28 GB just to store, and every matmul reads all of it from memory. Quantization is the single most effective way to shrink that: round each weight onto a small grid of integers and keep one floating-point scale to undo the rounding at run time.

Store 8-bit integers instead of 32-bit floats and the model is 4× smaller; store 4-bit integers and it is 8× smaller — often the difference between a model fitting on your GPU or not. The cost is a small, measurable rounding error, and this module builds the whole machinery from scratch so you can see exactly what that error is and how to keep it tiny.

What You’ll Learn

  • Why a trained weight has far more precision than inference needs
  • Symmetric (absmax) and affine (zero-point) quantization, from the formulas
  • How to quantize a tensor to int8 / int4 and measure the error in SNR
  • Why per-channel scales beat a single per-tensor scale (the outlier problem)
  • How to build a QuantizedLinear that replaces nn.Linear at a fraction of the memory
  • How GPTQ pushes past int4 by minimizing output error with second-order (Hessian) error feedback
  • How AWQ protects the salient weight channels using activation magnitudes — no Hessian, no backprop
  • How KIVI shrinks the KV cache — the other half of inference memory — by quantizing Keys per-channel and Values per-token, with a full-precision residual window
  • How the Microscaling (MX) formats — the block-floats Blackwell runs in hardware — fuse fp8 elements with a shared E8M0 power-of-two scale (MXFP4 at 4.25 bits)
  • How NVFP4 trades E8M0’s exactness for an E4M3 block scale that lands between the powers of two, held in range by a second per-tensor fp32 scale (two-level scaling)
  • How SmoothQuant quantizes the activations too — migrating their outliers into the weights with an equivalent per-channel rescale — to reach per-tensor int8 W8A8
  • Why the shipping activation granularity is per-token — a scale per row of X — and how it and per-channel weights factor out of the int8 GEMM as an outer product (the LLM.int8 vector-wise recipe)

Prerequisites

Intuition: A Weight Is a Number With Too Many Digits

A trained weight might be 0.0417328.... During inference, does the model really need all those digits? Almost never. The weights in a given tensor span some range — say roughly [-0.2, 0.2] — and if you chop that range into a few hundred evenly spaced buckets, rounding each weight to the nearest bucket barely changes the model’s output.

That is the whole idea. Pick a grid of integer levels, store which bucket each weight fell into (a small integer), and remember one scale — the width of a bucket — to turn buckets back into real numbers:

w \approx \text{scale} \times q, \qquad q \in \{-127, \dots, 127\}\ \text{(int8)}

Eight bits gives 2^8 = 256 buckets; four bits gives just 16. Fewer buckets = a coarser grid = more rounding error. The art of quantization is choosing the grid — how wide, how centred, and how many independent grids to keep — so that the error stays negligible while the integers stay tiny.

NoteKey Insight

Quantization trades precision for size. A float carries ~7 decimal digits of precision that a trained, redundant network mostly does not use. Rounding weights onto a 256-level (int8) grid typically changes outputs by well under a percent, while cutting the model’s memory — and its memory-bandwidth cost, which usually dominates inference — by 4×.

The Math: Symmetric and Affine Quantization

Symmetric (absmax). The simplest scheme centres the grid on zero. For b bits we use the signed range [-q_{\max}, q_{\max}] with q_{\max} = 2^{b-1}-1 (so 127 for int8). One scale covers the whole tensor:

s = \frac{\max_i |w_i|}{q_{\max}}, \qquad q_i = \text{clamp}\!\left(\text{round}\!\left(\frac{w_i}{s}\right),\, -q_{\max},\, q_{\max}\right), \qquad \hat{w}_i = s\, q_i

Because the grid is symmetric, 0.0 maps exactly to the integer 0 — no offset needed. This is the default for weights, which are roughly zero-centred.

Affine (zero-point). When the values are one-sided — think the output of a ReLU, all \geq 0 — a symmetric grid wastes half its levels on negatives that never occur. Affine quantization maps the true range [w_{\min}, w_{\max}] onto the unsigned range [0, 2^b - 1] using an integer zero-point z (the code that represents 0.0):

s = \frac{w_{\max} - w_{\min}}{2^b - 1}, \qquad z = \text{round}\!\left(-\frac{w_{\min}}{s}\right), \qquad q_i = \text{clamp}\!\left(\text{round}\!\left(\tfrac{w_i}{s}\right) + z,\, 0,\, 2^b - 1\right)

and we recover \hat{w}_i = s\,(q_i - z). No levels are wasted, so affine is the better choice for activations.

NoteKey Insight

Symmetric needs only a scale; affine needs a scale and a zero-point but uses the full code range. Weights → symmetric (zero-centred, and a zero-point would slow the matmul). Activations → affine (often one-sided). Both are just a linear map between reals and integers; the only lossy step is round.

Code: Quantize a Tensor from Scratch

Everything lives in quantization.py. Load it and quantize a weight tensor to int8, then read back the reconstruction error.

import importlib.util
import sys
from pathlib import Path

import torch
import torch.nn as nn

spec = importlib.util.spec_from_file_location("quantization", Path("quantization.py").resolve())
q = importlib.util.module_from_spec(spec)
sys.modules["quantization"] = q
spec.loader.exec_module(q)

torch.manual_seed(0)
w = torch.randn(4, 4) * 0.1                       # a small weight matrix
qi, scale = q.absmax_quantize(w, num_bits=8)

print("original[0]:   ", w[0].tolist())
print("integers[0]:   ", qi[0].tolist(), f"  (dtype {qi.dtype})")
print(f"scale:          {scale.item():.6f}")
print("dequantized[0]:", q.dequantize(qi, scale)[0].round(decimals=4).tolist())
original[0]:    [-0.11258398741483688, -0.11523602157831192, -0.025057857856154442, -0.04338788613677025]
integers[0]:    [-68, -69, -15, -26]   (dtype torch.int8)
scale:          0.001666
dequantized[0]: [-0.11330000311136246, -0.11490000039339066, -0.02500000037252903, -0.043299999088048935]

The four floats became four small integers plus one shared scale. Now measure how much precision that cost, in int8 versus int4:

w = torch.randn(512, 512) * 0.05
for bits in (8, 4):
    qi, s = q.absmax_quantize(w, num_bits=bits)
    err = q.quantization_error(w, q.dequantize(qi, s))
    print(f"int{bits}: max error {err['max_abs']:.5f}   SNR {err['snr_db']:5.1f} dB")
int8: max error 0.00092   SNR  39.5 dB
int4: max error 0.01664   SNR  14.3 dB

int8 reconstructs the tensor at ~40 dB (rounding error ~10 000× smaller than the signal); int4, with only 16 levels, drops to the mid-teens. Signal-to-noise ratio is the honest way to compare: every extra bit adds ~6 dB, so int8 is worth about 24 dB more than int4.

Per-Tensor vs Per-Channel: The Outlier Problem

One scale for a whole matrix has a weakness. Neural-network weight matrices routinely contain a few outlier channels whose magnitudes dwarf the rest. The absmax scale is set by that single largest value, so every other channel is forced onto a needlessly coarse grid.

The fix is per-channel quantization: give each output channel (row) its own scale. One extra float per row buys back the precision the outlier stole. Watch it on a matrix where channel 0 is 20× louder than the rest:

torch.manual_seed(0)
w = torch.randn(24, 64) * 0.05
w[0] *= 20.0                                       # one loud output channel

qt, st = q.absmax_quantize(w, num_bits=8)          # one scale for everything
qc, sc = q.quantize_per_channel(w, num_bits=8, axis=0)   # one scale per row

# Per-channel mean-squared error under each scheme (for the visual below).
mse_tensor = ((w - q.dequantize(qt, st)) ** 2).mean(dim=1)
mse_chan = ((w - q.dequantize(qc, sc)) ** 2).mean(dim=1)

ojs_define(quantChan = {
    "perTensor": mse_tensor.tolist(),
    "perChannel": mse_chan.tolist(),
    "outlier": 0,
})
et = q.quantization_error(w, q.dequantize(qt, st))
ec = q.quantization_error(w, q.dequantize(qc, sc))
print(f"per-tensor  : SNR {et['snr_db']:.1f} dB")
print(f"per-channel : SNR {ec['snr_db']:.1f} dB   ← the quiet channels keep their precision")
per-tensor  : SNR 29.1 dB
per-channel : SNR 42.8 dB   ← the quiet channels keep their precision

Per-tensor lets one loud channel drag every quiet channel’s error up by the same factor it is louder — here ~20×. Per-channel absorbs the outlier into its own row and leaves the rest untouched, recovering well over 10 dB of SNR. This is exactly why production weight quantization is always per-channel (and why methods like LLM.int8() go further and pull genuine outlier features out into fp16).

Code: A Quantized Linear Layer

Now assemble the piece that actually saves memory in a transformer: a drop-in for nn.Linear whose weight is stored per-channel int8. QuantizedLinear.from_linear quantizes a trained layer; its forward dequantizes the weight and runs an ordinary linear — the faithful, readable stand-in for what an integer kernel does (integer matmul, then a per-channel rescale).

torch.manual_seed(0)
linear = nn.Linear(512, 512)
qlinear = q.QuantizedLinear.from_linear(linear, num_bits=8)

x = torch.randn(8, 512)
ref = linear(x)
out = qlinear(x)

rel = (out - ref).norm() / ref.norm()
print(f"output relative error: {rel.item():.4%}")     # well under 1%

mem = qlinear.memory_bytes()
print(f"fp32 weight bytes: {mem['fp32_weight_bytes']:,}")
print(f"int8 stored bytes: {mem['weight_bytes'] + mem['scale_bytes']:,.0f}"
      f"  ({mem['compression']:.2f}× smaller)")
output relative error: 0.3824%
fp32 weight bytes: 1,048,576
int8 stored bytes: 264,192  (3.97× smaller)

The quantized layer’s output is within a fraction of a percent of the original, at ~4× less memory. Scale that up to a whole model:

for params, name in [(7e9, "7B"), (13e9, "13B"), (70e9, "70B")]:
    line = f"{name:>4}:  "
    for bits in (32, 16, 8, 4):
        line += f"int{bits}={q.model_footprint(int(params), bits)/1e9:6.1f}GB  "
    print(line)
  7B:  int32=  28.0GB  int16=  14.0GB  int8=   7.0GB  int4=   3.5GB  
 13B:  int32=  52.0GB  int16=  26.0GB  int8=  13.0GB  int4=   6.5GB  
 70B:  int32= 280.0GB  int16= 140.0GB  int8=  70.0GB  int4=  35.0GB

A 70B model is 280 GB in fp32 — many GPUs — but 35 GB in int4, within reach of a single high-memory card. That table is why quantization is not optional at the frontier.

Group-wise Quantization: One Scale Per Block

Per-channel fixed the outlier channel. But it left a subtler leak: a single loud weight inside a row still sets that whole row’s scale, so every other weight in the row — dozens or thousands of them — is forced onto a needlessly coarse grid. The fix is to stop sharing a scale across the whole row: split each row into groups of group_size consecutive weights and give every group its own absmax scale. This is the granularity every production int4 model actually uses. When a checkpoint says “int4,” it almost always means int4, group size 128 — GPTQ-g128, AWQ-g128, and llama.cpp’s k-quants are all group-wise.

The key idea is that group size is a dial spanning the two granularities you just built:

\underbrace{g = d_{\text{in}}}_{\text{per-channel (one scale/row)}} \;\;\longleftrightarrow\;\; \underbrace{g = 1}_{\text{one scale/weight (lossless)}}

Everything useful lives in between. grouped.py holds the whole thing — load it and confirm the two endpoints are exactly the granularities the module already knows:

spec_g = importlib.util.spec_from_file_location("grouped", Path("grouped.py").resolve())
gq = importlib.util.module_from_spec(spec_g)
sys.modules["grouped"] = gq
spec_g.loader.exec_module(gq)

torch.manual_seed(0)
w = torch.randn(32, 128) * 0.05
w[3, 5] *= 40.0                                    # loud weights buried in
w[17, 100] *= 40.0                                 # otherwise-quiet rows

# Endpoint 1: a group as wide as the row IS per-channel — bit for bit.
qg, sg = gq.quantize_grouped(w, num_bits=4, group_size=128, axis=1)
qc, sc = q.quantize_per_channel(w, num_bits=4, axis=0)
print("g = d_in  ==  per-channel:", bool(torch.equal(qg, qc) and torch.equal(sg, sc)))

# Endpoint 2: a group of one gives every weight its own scale — exact.
q1, s1 = gq.quantize_grouped(w, num_bits=4, group_size=1, axis=1)
recon = gq.dequantize_grouped(q1, s1, group_size=1, axis=1)
print("g = 1     ==  exact reconstruction:", torch.allclose(recon, w, atol=1e-6))
g = d_in  ==  per-channel: True
g = 1     ==  exact reconstruction: True
NoteKey Insight

Group-wise quantization is not a new idea bolted on — it is the continuous interpolation between the two granularities you already have. Per-channel is one endpoint (g = d_in), a scale per weight is the other (g = 1). Choosing a group size in between is choosing a point on the accuracy–memory curve.

The Math: Group Size Is a Dial

Split a weight row into n_g = d_{\text{in}} / g groups. Each group G carries its own symmetric absmax scale, then rounds against it:

s_G = \frac{\max_{i \in G} |w_i|}{q_{\max}}, \qquad \hat{w}_i = s_G \cdot \operatorname{clamp}\!\Big(\operatorname{round}\big(w_i / s_G\big),\, -q_{\max},\, q_{\max}\Big)

Why does a smaller group help? Because refining a group into sub-groups can only lower each sub-group’s absmax (a subset’s maximum never exceeds the whole set’s). So as g shrinks, the scale seen by any given weight is non-increasing, and with it the worst-case rounding error:

|w_i - \hat{w}_i| \;\le\; \frac{s_G}{2} \quad\text{is monotone non-increasing in decreasing } g.

The price is storage. Each group stores one extra scale (a 16-bit float), so the scale amortizes to \text{scale\_bits} / g extra bits on every weight:

b_{\text{eff}} = b + \frac{\text{scale\_bits}}{g}.

At the industry-default int4, g = 128: b_{\text{eff}} = 4 + 16/128 = 4.125 bits — about 0.4% over raw int4, for a large slice of the accuracy back. The GPTQ paper measures this exact trade in perplexity: “group-size 128 (≈ 0.15 extra bits) improves perplexities by another 0.1.”

Step by Step

Drive the dial on a real weight row below. Each step halves the group size, subdividing the strip into finer blocks — watch each block’s scale tighten to its own contents and the worst-case error bound shrink, while the bit cost climbs.

# One real, outlier-bearing weight row for the step-through (64 values).
torch.manual_seed(7)
row = torch.randn(64) * 0.06
row[19] *= 30.0                                    # a loud weight mid-row
row_list = row.tolist()

q_max4 = 2 ** (4 - 1) - 1
group_scales = {}
for g in (64, 32, 16, 8, 4, 2, 1):
    absmax = gq.group_absmax(row.unsqueeze(0), group_size=g, axis=1)[0]  # (n_g,)
    group_scales[str(g)] = (absmax / q_max4).tolist()

ojs_define(groupRow = {
    "values": row_list,
    "scales": group_scales,
    "sizes": [64, 32, 16, 8, 4, 2, 1],
    "qMax": q_max4,
})
NoteKey Insight

Notice the loud weight (red) sits in one block. At group_size = 64 its scale is imposed on the whole row; as you halve the group, the outlier gets quarantined into an ever-smaller block and every other block collapses onto a tight scale of its own. The outlier’s own block never gets cheaper — that residual is exactly what the error-correcting methods below (GPTQ, AWQ) attack from a different angle.

The Accuracy–Memory Dial

Turning the dial has a measurable price and a measurable payoff. Sweep the group size on the outlier matrix and read both axes at once — reconstruction SNR (higher is better) against the effective bits per weight (4 + 16/g):

sizes = [128, 64, 32, 16, 8, 4, 2, 1]
sweep = []
for g in sizes:
    qi, s = gq.quantize_grouped(w, num_bits=4, group_size=g, axis=1)
    err = q.quantization_error(w, gq.dequantize_grouped(qi, s, group_size=g, axis=1))
    sweep.append({
        "g": g,
        "snr": None if err["snr_db"] == float("inf") else round(err["snr_db"], 2),
        "bits": round(gq.grouped_effective_bits(4, g), 4),
    })

ojs_define(groupSweep = sweep)
print(f"{'group':>6}  {'eff.bits':>8}  {'SNR(dB)':>8}")
for r in sweep:
    snr = "  exact" if r["snr"] is None else f"{r['snr']:8.1f}"
    tag = "  ← per-channel" if r["g"] == 128 else ("  ← per-weight" if r["g"] == 1 else "")
    print(f"{r['g']:>6}  {r['bits']:>8.3f}  {snr}{tag}")
 group  eff.bits   SNR(dB)
   128     4.125      14.2  ← per-channel
    64     4.250      16.6
    32     4.500      18.5
    16     5.000      20.7
     8     6.000      23.9
     4     8.000      27.0
     2    12.000      31.1
     1    20.000     157.2  ← per-weight

The two curves cross the story: shrinking the group buys SNR fast at first, then flattens — while the bit cost stays nearly flat until it explodes past g ≈ 16. The sweet spot sits where they diverge, which is exactly why g = 128 (and, for even lower bit-widths, g = 64) is the near-universal default.

Now assemble the payoff — a GroupQuantizedLinear, the sibling of the per-channel QuantizedLinear, storing one fp16 scale per input block:

torch.manual_seed(0)
linear = nn.Linear(1024, 1024)
qlin = gq.GroupQuantizedLinear.from_linear(linear, num_bits=4, group_size=128)

x = torch.randn(8, 1024)
rel = (qlin(x) - linear(x)).norm() / linear(x).norm()
mem = qlin.memory_bytes()
print(f"int4 g=128 output error: {rel.item():.3%}")
print(f"effective bits/weight:   {mem['effective_bits']:.3f}")
print(f"compression vs fp32:     {mem['compression']:.2f}×")
int4 g=128 output error: 7.111%
effective bits/weight:   4.125
compression vs fp32:     7.76×

That is genuine 4-bit storage (~7.8× smaller than fp32) whose output still tracks the fp32 layer — the honest starting point the error-correcting methods below refine.

Where the Precision Goes

The per-group scale map makes the localization visible. Each cell is one group’s scale; group-wise gives the outlier its own bright block and lets every other block stay dark (precise), where per-channel would smear one column and per-tensor a single value across the whole matrix.

g_map = 16
_, s_grp = gq.quantize_grouped(w, num_bits=4, group_size=g_map, axis=1)
_, s_chan = q.quantize_per_channel(w, num_bits=4, axis=0)
qt, s_tensor = q.absmax_quantize(w, num_bits=4)

ojs_define(scaleMap = {
    "grouped": s_grp.tolist(),                       # (rows, n_groups)
    "perChannel": s_chan.squeeze(-1).unsqueeze(-1).tolist(),  # (rows, 1)
    "perTensor": float(s_tensor),
    "rows": s_grp.shape[0],
    "groupSize": g_map,
})
TipTry This
  1. Flip to per-tensor. The whole map turns one uniform bright colour — the single outlier has coarsened every weight.
  2. Flip to per-channel. Only the two outlier rows light up; the rest go dark. Better — but those two rows are still coarse across all 128 of their weights.
  3. Back to group-wise. Now only the blocks holding an outlier are bright. The other seven blocks in each outlier row keep their own fine scale. That reclaimed precision is the SNR the curve above measured.

Beyond Round-to-Nearest: GPTQ

Everything above rounds each weight in a vacuum — nearest grid point, one weight at a time, blind to how the layer will actually be used. That is round-to-nearest (RTN), and it minimizes the weight error \lVert W - \hat{W}\rVert. But a layer does not ship its weights to the user; it ships \hat{W}x. What matters is the output error over the inputs the layer really sees:

\arg\min_{\hat{W}}\ \lVert W X - \hat{W} X \rVert_2^2

where X is a small calibration batch — a few hundred real activations, one per column. When the input features are correlated (they always are in a trained network), the weight that minimizes rounding error is not the one that minimizes output error. RTN leaves accuracy on the table, and at 3–4 bits that gap is the difference between a usable model and a broken one.

GPTQ (Frantar et al., 2022) fixes this with one idea: error feedback. Quantize the weight matrix one column at a time; after you snap a column to the grid, look at the output error you just created and nudge the columns you have not quantized yet to cancel it. The later weights absorb the earlier weights’ rounding error. By the time you reach the last column, the accumulated damage to the layer’s output has been continuously paid down instead of piling up.

NoteKey Insight

RTN asks “what is the nearest grid point for this weight?” GPTQ asks “given that I must round this weight, how should I adjust the remaining weights so the layer’s output barely moves?” The first ignores the calibration data; the second is steered by it. That is the whole difference — and it is why GPTQ took large models to 4 bits when RTN could not.

The Math: Second-Order Error Compensation

How much should the remaining weights move? The objective \lVert W X - \hat{W} X \rVert^2 is a quadratic in the weights, and a quadratic’s shape is its Hessian

H = 2 X X^\top \qquad (\text{shape } d_\text{in} \times d_\text{in}),

the correlation matrix of the input features. This is the Optimal Brain Surgeon (OBS) setup from network pruning: if we fix one weight w_q to its quantized value \text{quant}(w_q), the update to the other free weights that keeps the loss lowest, and the loss it costs, are both closed-form:

\boldsymbol{\delta} = -\frac{w_q - \text{quant}(w_q)}{[H^{-1}]_{qq}}\, H^{-1}_{:,\,q}, \qquad \Delta L = \frac{\bigl(w_q - \text{quant}(w_q)\bigr)^2}{[H^{-1}]_{qq}}.

The off-diagonal of H^{-1} is exactly the “how correlated is weight q with the others” signal that tells the remaining weights how to compensate. Applying this greedily, column by column, is OBQ. GPTQ’s contribution is making it fast: quantize the columns in a fixed left-to-right order (the order barely matters for a big layer), and precompute every conditional H^{-1} quantity with a single Cholesky factorization C (upper-triangular, C^\top C = H^{-1}) instead of re-inverting a shrinking matrix at each step. The per-column update becomes:

q_j       = quant(W[:, j])          # round column j on its fixed per-row grid
err       = (W[:, j] − q_j) / C[j, j]
W[:, j:] −= err ⊗ C[j, j:]          # push the error onto the not-yet-done columns

A tiny constant — 1% of the mean diagonal — is added to H before inverting so it stays invertible. That is the entire algorithm.

NoteKey Insight

If the features are uncorrelated, H is diagonal, C has no off-diagonal terms, and C[j, j:] touches only column j — no error propagates and GPTQ becomes RTN exactly. Every bit of GPTQ’s advantage is the correlation correction. We prove this bit-for-bit in the code below.

Step by Step

Walk one GPTQ pass over an 6\times 8 layer at 2 bits. Each step quantizes one column (it turns solid), then propagates its error rightward onto the columns still to come. Watch the running output loss track below the RTN line the whole way — that gap is the compensation working.

Code: GPTQ from Scratch

The algorithm lives in gptq.py: layer_hessian builds 2XX^\top, round_to_nearest is the RTN baseline, and gptq_quantize runs the Cholesky Algorithm 1 above. Build a small layer, a correlated calibration batch, and compare the two on output loss.

gptq_spec = importlib.util.spec_from_file_location("gptq", Path("gptq.py").resolve())
gptq = importlib.util.module_from_spec(gptq_spec)
sys.modules["gptq"] = gptq
gptq_spec.loader.exec_module(gptq)

torch.manual_seed(0)
W = torch.randn(64, 64) * 0.1                    # a layer's weight matrix
mix = torch.randn(64, 64)
X = mix @ torch.randn(64, 128)                   # 128 correlated calibration inputs
H = gptq.layer_hessian(X)                        # 2 X Xᵀ, damped

for bits in (4, 3, 2):
    rtn, _ = gptq.round_to_nearest(W, num_bits=bits)
    gp, _ = gptq.gptq_quantize(W, H, num_bits=bits)
    lr = gptq.output_loss(W, rtn, X)
    lg = gptq.output_loss(W, gp, X)
    print(f"int{bits}:  RTN loss {lr:9.1f}   GPTQ loss {lg:9.1f}   "
          f"({1 - lg / lr:.0%} lower)")
int4:  RTN loss    4105.7   GPTQ loss    1915.3   (53% lower)
int3:  RTN loss   21976.2   GPTQ loss   10072.2   (54% lower)
int2:  RTN loss  175683.4   GPTQ loss   92849.2   (47% lower)

GPTQ roughly halves the output error at every bit width, and the win is largest where it matters most — the aggressive low-bit regime. Now the two facts that pin the algorithm down. First, uncorrelated features collapse GPTQ to RTN, exactly:

# Orthogonal calibration features → diagonal Hessian → no error to propagate.
H_diag = gptq.layer_hessian(torch.eye(64) * 3.0, damp=0.0)
gp_diag, _ = gptq.gptq_quantize(W, H_diag, num_bits=3)
rtn3, _ = gptq.round_to_nearest(W, num_bits=3)
print("diagonal H → GPTQ == RTN, bit-for-bit:", torch.equal(gp_diag, rtn3))
diagonal H → GPTQ == RTN, bit-for-bit: True

Second, GPTQ’s Cholesky reformulation is not an approximation — it reproduces the slow, explicit Optimal Brain Surgeon recursion (obq_reference, which shrinks H^{-1} by a Schur complement at every step) to floating-point precision:

gp, _ = gptq.gptq_quantize(W, H, num_bits=3)
ref = gptq.obq_reference(W, H, num_bits=3)
print("GPTQ (Cholesky) == OBS reference:", torch.allclose(gp, ref, atol=1e-4))
GPTQ (Cholesky) == OBS reference: True

RTN vs GPTQ at a Glance

The bars below come straight from demonstrate_gptq. Each bit width is normalized so RTN is a full bar; GPTQ’s shorter bar is the output loss that survives — the rest is what the error feedback recovered.

demo = gptq.demonstrate_gptq(verbose=False)
ojs_define(gptqBars = {
    "bits": [4, 3, 2],
    "reduction": [demo["results"][b]["reduction"] for b in (4, 3, 2)],
})

# The stepper trace: a compact 6×8 layer at 2 bits where the win is vivid.
torch.manual_seed(0)
W_trace = torch.randn(6, 8) * 0.1
X_trace = torch.randn(8, 8) @ torch.randn(8, 64)
ojs_define(gptqTrace = gptq.gptq_trace(W_trace, X_trace, num_bits=2))
TipTry This
  1. Kill the correlation. In the code, replace X = mix @ torch.randn(64, 128) with plain X = torch.randn(64, 128). The features are now near-orthogonal, H is nearly diagonal, and GPTQ’s advantage almost vanishes — you can see that its whole edge is the correlation correction.
  2. Push to the extreme. Add 1 to the bit sweep. At the same 1-bit ternary edge where RTN is hopeless, GPTQ still recovers a large fraction of the loss — this is why the paper reports usable 2-bit and even ternary models.

A Different Lever: AWQ

GPTQ fixes rounding after the fact: round a column, then use the Hessian to push the error onto the columns you haven’t touched yet. AWQ (Lin et al., 2023) reshapes the problem before you quantize at all — with no backprop, no Hessian, and no reconstruction. It is the method most int4 open weights actually ship with.

The idea starts from a fact GPTQ and RTN both ignore: not all weights are equally important, and importance lives in the activations, not the weights. A weight column multiplied by a channel whose activations are large contributes far more to the output than one multiplied by a near-silent channel. Protect those salient channels and you protect the output.

How do you protect a channel without spending extra bits on it? An equivalent transformation. For a linear y = Wx, scale one input channel’s weight column up by s > 1 and divide that channel’s activation by the same s. In full precision the product is unchanged — but now the salient column occupies more of the integer grid, so quantizing it costs less.

NoteKey Insight

GPTQ asks “given this rounding, how do I cancel its output error?” AWQ asks “which channels can I afford to round?” — and rescales the layer so the answer is “all of them.” One uses second-order weight statistics; the other uses first-order activation statistics. They attack the same \lVert WX - \hat W X\rVert from opposite ends.

Step through the equivalent transform — watch the salient column grow, its activation shrink, and the output stay put:

The Math: Protect the Salient Channels

Recall the symmetric grid: \Delta = \max(|w|) / 2^{N-1} is the width of one quantization bucket, and rounding a weight introduces an error of about \text{RoundErr} \approx 0.25 buckets (uniform on [0, 0.5]). The output error a single weight contributes is that rounding error times the activation it feeds:

\text{Err}\big(Q(w)\,x\big) = \Delta \cdot \text{RoundErr} \cdot x .

Now apply the equivalent transform — scale the weight by s and the activation by 1/s:

\text{Err}\big(Q(w\,s)\,(x / s)\big) = \Delta' \cdot \text{RoundErr} \cdot x \cdot \tfrac{1}{s}, \qquad \frac{\text{Err}'}{\text{Err}} = \frac{\Delta'}{\Delta}\cdot\frac{1}{s} .

Here is the whole trick. Scaling one channel up barely moves the group’s maximum, so \Delta' \approx \Delta and the ratio is \approx 1/s < 1: the salient channel’s error shrinks. Push s too far, though, and that column becomes the new maximum — \Delta' jumps, coarsening every other weight in the group. So s is a trade-off, and AWQ tunes it with a single knob \alpha tied to activation salience:

s_X = \text{mean}(|X|) \ \text{(per input channel)}, \qquad s = s_X^{\alpha},

\alpha^{*} = \arg\min_{\alpha \in [0,1]} \ \big\lVert \, Q\!\big(W\,\text{diag}(s)\big)\,\text{diag}(s)^{-1} X \;-\; W X \, \big\rVert .

At \alpha = 0 every scale is s = s_X^{0} = 1 — no rescaling, i.e. plain round-to-nearest. At \alpha = 1 the scale is the raw activation magnitude, which over-protects the loud channels and wrecks the quiet ones. The best \alpha sits between, found by a cheap grid search (the paper uses 20 points over [0,1]).

NoteKey Insight

Because \alpha = 0 is round-to-nearest and is one of the grid points, the search can never return something worse than RTN — it takes the arg-min. AWQ’s output loss is \le RTN’s, guaranteed, exactly the way GPTQ’s is. Two very different mechanisms, the same safety net.

Code: AWQ from Scratch

The method lives in awq.py. activation_scale reads salience straight off a calibration batch; apply_awq_scale is the equivalent transform W\,\text{diag}(s); awq_effective quantizes the reshaped weights and folds the inverse scale back so the result is comparable to the original W; and awq_search grid-searches \alpha. First, prove the transform is free until you quantize:

awq_spec = importlib.util.spec_from_file_location("awq", Path("awq.py").resolve())
awq = importlib.util.module_from_spec(awq_spec)
sys.modules["awq"] = awq
awq_spec.loader.exec_module(awq)

torch.manual_seed(0)
W_awq = torch.randn(64, 64) * 0.1                # a layer's weights
X_awq = torch.randn(64, 128)                     # calibration inputs
salient_ch = torch.randperm(64)[:4]
X_awq[salient_ch] *= 12.0                        # a few loud (salient) channels

s_X = awq.activation_scale(X_awq)                # per-channel salience
s = s_X ** 0.5                                   # some scale
same = torch.allclose(awq.equivalent_output(W_awq, s, X_awq), W_awq @ X_awq, atol=1e-4)
print(f"(W·diag(s))(x/s) == W·x in full precision:  {same}")
(W·diag(s))(x/s) == W·x in full precision:  True

Now quantize. RTN rounds blindly; AWQ searches \alpha, scales the salient columns up, and quantizes the reshaped matrix:

for bits in (4, 3, 2):
    rtn = awq.round_to_nearest(W_awq, num_bits=bits)
    awq_w, alpha, _ = awq.awq_quantize(W_awq, X_awq, num_bits=bits)
    rtn_loss = awq.output_loss(W_awq, rtn, X_awq)
    awq_loss = awq.output_loss(W_awq, awq_w, X_awq)
    reduction = 1 - awq_loss / rtn_loss
    print(f"int{bits}: RTN {rtn_loss:8.1f}   AWQ {awq_loss:8.1f}   "
          f"α*={alpha:.2f}   −{reduction:.0%}")
int4: RTN    642.9   AWQ    328.6   α*=0.30   −49%
int3: RTN   3680.4   AWQ   1785.0   α*=0.30   −52%
int2: RTN  28191.1   AWQ  12606.1   α*=0.45   −55%

AWQ roughly halves the output loss here, and — like GPTQ — the gap widens as bits shrink: fewer bits mean more rounding error, so protecting the salient channels matters more. The alpha=0 collapse is exact:

s_one = torch.ones(64)                           # s = s_X ** 0
awq_rtn = awq.awq_effective(W_awq, s_one, num_bits=3)
rtn_only = awq.round_to_nearest(W_awq, num_bits=3)
print(f"AWQ at α=0 is RTN, bit-for-bit:  {torch.equal(awq_rtn, rtn_only)}")
AWQ at α=0 is RTN, bit-for-bit:  True

RTN vs AWQ, Channel by Channel

Where does the win actually come from? Break the output loss down per input channel — each bar is a channel’s weight error weighted by how loud its activations are. RTN leaves the salient channels towering. AWQ pulls exactly those down.

W_awq_hat, _, _ = awq.awq_quantize(W_awq, X_awq, num_bits=3)
rtn_hat = awq.round_to_nearest(W_awq, num_bits=3)
rtn_ch = awq.channel_output_error(W_awq, rtn_hat, X_awq)
awq_ch = awq.channel_output_error(W_awq, W_awq_hat, X_awq)
# show the 12 loudest channels so the salient ones are visible
order = torch.argsort(rtn_ch, descending=True)[:12]
ojs_define(awqChannels = {
    "idx": order.tolist(),
    "rtn": rtn_ch[order].tolist(),
    "awq": awq_ch[order].tolist(),
    "salient": salient_ch.tolist(),
})

AWQ and GPTQ are complementary, not rivals: GPTQ steers on the weight-side Hessian, AWQ on activation salience, and production quantizers often stack them — AWQ’s scaling first, GPTQ’s error feedback second. Both take large models to honest int4.

Activations Fight Back: the W8A8 Problem

Everything so far quantizes a weight. But the matmul that actually runs is Y = XW — an activation X times a weight W — and to make that matmul fast on int8 tensor cores you must quantize both operands. That is a W8A8 GEMM: 8-bit weights and 8-bit activations, multiplied as integers and accumulated in int32. Keep X in fp16 and the tensor core has to convert on the fly; the int8 speedup evaporates.

Here is the catch, and it is the whole reason this section exists. Weights are flat — a well-behaved bell curve you have quantized five different ways already. Activations are not. Run real text through a transformer and a few input channels light up with magnitudes 10–100× everything else, and — crucially — they are the same channels for every token. A single per-tensor int8 scale on X has to reach that loud channel, so it spends almost all of its 255 codes covering one column and leaves ordinary channels with a code or two. The activations are, in a word, unquantizable by the plain per-tensor recipe.

You cannot fix this the way you fixed weights. A per-channel weight scale folds cleanly into the matmul (that is the whole QuantizedLinear story). But activation channels run along the contraction axis of XW — a per-channel activation scale does not commute through the sum, so you can’t just fold it away. This is exactly the wall LLM.int8() hit and answered with slow mixed precision: pull the outlier channels out into a separate fp16 matmul. SmoothQuant answers it without any mixed precision at all.

NoteKey Insight

Activation outliers are hard but systematic — the same channels are loud for every token. Weights are easy and have headroom. So don’t fight the outlier where it lives; move it. Rescale each input channel so some of its magnitude slides out of the activation and into the weight, using a transform that leaves the product XW exactly unchanged.

The Math: Migrate the Difficulty

Insert a per-input-channel diagonal \operatorname{diag}(s) and its inverse between the two operands. Because they cancel, the product is untouched:

Y = X W = \big(X\,\operatorname{diag}(s)^{-1}\big)\big(\operatorname{diag}(s)\,W\big) = \hat X \,\hat W .

\hat X = X\,\operatorname{diag}(s)^{-1} divides the columns of X (input channels) by s; \hat W = \operatorname{diag}(s)\,W multiplies the rows of W (same input channels) by s. Choose s per channel j to balance how hard each side is to quantize:

s_j = \frac{\max_t |X_{t,j}|^{\,\alpha}}{\max_i |W_{j,i}|^{\,1-\alpha}}, \qquad \alpha \in [0, 1].

The exponent \alpha is a migration dial:

  • \alpha = 1 \Rightarrow s_j = \max_t|X_{t,j}|: every activation channel is divided by its own peak, so \hat X flattens to magnitude 1 everywhere — but \hat W absorbs the entire outlier and may itself become hard.
  • \alpha = 0 \Rightarrow s_j = 1/\max_i|W_{j,i}|: the weights flatten and the activations keep all their range — you are back where you started.
  • \alpha = 0.5 (the default) is the geometric split s_j = \sqrt{\max|X_j|/\max|W_j|}: push half the difficulty across. For most layers that is the sweet spot.

And because s is fixed offline (from a calibration pass over a little text), the \operatorname{diag}(s)^{-1} factor folds into the previous layer’s LayerNorm weights — exactly like AWQ’s scale folds backward. So at run time there is no extra op: you ship \hat W and a LayerNorm whose scale already carries 1/s, and both operands are plain per-tensor int8.

NoteThe same trick, a different goal

Look back at AWQ: it uses the identical equivalent transform X\operatorname{diag}(s)^{-1}\cdot\operatorname{diag}(s)W. The difference is the purpose. AWQ keeps activations in fp16 (W4A16) and scales to protect salient weights. SmoothQuant scales to make the activations themselves quantizable (W8A8). One algebraic identity, two quantization regimes.

Step by Step

Watch one smoothing pass on a small layer. Channel 2 is the loud one; drag the control and see its magnitude slide out of the activations and into the weights while the product stays put.

spec_sq = importlib.util.spec_from_file_location("smoothquant", Path("smoothquant.py").resolve())
sq = importlib.util.module_from_spec(spec_sq)
sys.modules["smoothquant"] = sq
spec_sq.loader.exec_module(sq)

torch.manual_seed(3)
_Xs = torch.randn(8, 6)
_Ws = torch.randn(6, 6)
_Xs[:, 2] *= 40.0                                    # the loud input channel

_xmax = sq.channel_absmax(_Xs, dim=0)
_wmax = sq.channel_absmax(_Ws, dim=1)
_Xh, _Wh, _s = sq.smooth(_Xs, _Ws, alpha=0.5)
_xmax_after = sq.channel_absmax(_Xh, dim=0)
_wmax_after = sq.channel_absmax(_Wh, dim=1)
_inv = (_Xh @ _Wh - _Xs @ _Ws).abs().max().item()

ojs_define(smoothWalk = {
    "actMax": _xmax.tolist(),
    "wtMax": _wmax.tolist(),
    "scale": _s.tolist(),
    "actMaxAfter": _xmax_after.tolist(),
    "wtMaxAfter": _wmax_after.tolist(),
    "outlier": 2,
    "invErr": _inv,
})
TipTry This
  1. Follow one channel. Step to the end and watch channel 2 only: its activation bar shrinks by s and its weight bar grows by s — the product of the two is invariant, which is why the output never moves.
  2. Read the invariance tag. At the last step the max output difference is ~10^{-6}: pure float round-off. Smoothing adds no error of its own; every bit of the eventual error comes from the int8 rounding that follows.

Code: SmoothQuant from Scratch

smoothquant.py is short because the idea is one transform. First, the invariance that makes it safe — for any \alpha, the smoothed product equals the original:

torch.manual_seed(0)
X = torch.randn(32, 16)
W = torch.randn(16, 16)
X[:, 3] *= 100.0                                     # a persistent ~100× outlier channel

X_hat, W_hat, s = sq.smooth(X, W, alpha=0.5)
print(f"max |X̂Ŵ − XW| = {(X_hat @ W_hat - X @ W).abs().max():.2e}   (float round-off only)")
print(f"activation channel 3 max: {sq.channel_absmax(X, 0)[3]:7.1f}  →  {sq.channel_absmax(X_hat, 0)[3]:6.2f}")
print(f"weight     channel 3 max: {sq.channel_absmax(W, 1)[3]:7.2f}  →  {sq.channel_absmax(W_hat, 1)[3]:6.2f}")
max |X̂Ŵ − XW| = 6.10e-05   (float round-off only)
activation channel 3 max:   223.2  →   22.77
weight     channel 3 max:    2.32  →   22.77

The outlier moved out of the activations and into the weights. Now quantize both operands to per-tensor int8 — naively, then after smoothing — and read the output SNR:

ref = X @ W

naive = sq.quantize_w8a8(X, W)                       # per-tensor int8 both operands
smart = sq.smoothquant_linear(X, W, alpha=0.5)       # smooth, then W8A8

naive_snr = q.quantization_error(ref, naive)["snr_db"]
smart_snr = q.quantization_error(ref, smart)["snr_db"]
print(f"naive  W8A8 : {naive_snr:6.2f} dB")
print(f"smooth W8A8 : {smart_snr:6.2f} dB   (+{smart_snr - naive_snr:.1f} dB)")
naive  W8A8 :  31.60 dB
smooth W8A8 :  41.42 dB   (+9.8 dB)

Roughly +9–10 dB — about a bit and a half of accuracy — bought by a per-channel rescale that costs nothing at run time.

Drive the Migration Dial

\alpha is a dial, and the picture is a trade. Slide it right and the activation range collapses while the weight range climbs; the W8A8 output SNR traces a U that peaks where the two are balanced — near \alpha = 0.5.

_demo = sq.demonstrate_smoothquant(verbose=False)
ojs_define(sqSweep = {
    "alphas": [r["alpha"] for r in _demo["sweep"]],
    "snr": [r["snr_db"] for r in _demo["sweep"]],
    "actRange": [r["act_range"] for r in _demo["sweep"]],
    "wtRange": [r["wt_range"] for r in _demo["sweep"]],
    "best": _demo["best_alpha"],
    "naiveSnr": _demo["naive_snr_db"],
})
TipTry This
  1. Find the peak. Sweep \alpha from 0 to 1. At \alpha = 0 nothing moves and you sit on the naive baseline; push toward 1 and you over-migrate — the weights become the outlier and SNR falls again. The best point is in the middle.
  2. Watch the two bars cross. As \alpha climbs, the activation range dives and the weight range rises. The output is happiest where neither dominates — the geometric balance SmoothQuant defaults to.
WarningPer-channel means the contraction axis — and only there

SmoothQuant’s scale is per input channel because that is the one axis you can fold away: \operatorname{diag}(s)^{-1} merges into the previous LayerNorm, s into this layer’s weights, so smoothing costs nothing at run time. A per-input-channel scale on the activations themselves has no such escape — it sits on the contraction axis k, buried inside the \sum_k, and cannot be pulled out of the matmul at all. That is the axis SmoothQuant refuses to quantize directly. The token axis is a different story: you cannot fold a per-token scale, but you can factor it — and the next section does exactly that.

A Scale Per Token: Granularity on the Axis That Survives

Weights in this module climbed a granularity ladder — per-tensor → per-channel → group — and every rung bought accuracy by keeping a loud channel from setting the step for the rest. Activations never left the bottom rung: SmoothQuant and fp8 both quantize X per-tensor, one scale for the whole matrix. The callout above is why the obvious upgrade looked blocked — a per-channel activation scale is trapped on the contraction axis.

But X has shape (\text{tokens} \times C_\text{in}), and only one of its two axes is the contraction axis. The token axis is an output dimension: it indexes the rows of Y, exactly the way an output channel of W indexes the columns. A scale on an output axis survives the matmul. So the activation granularity that ships is not per-channel — it is per-token.

Intuition: A Different Outlier, One Axis Over

SmoothQuant’s villain was a loud column — the same input channel large for every token. There is a second, equally common pattern: a loud row — a token whose activations are large across all channels. Attention sinks, the beginning-of-sequence token, and the “massive activation” tokens transformers park huge norms on all look like this. Lay X out as a grid and the two outliers are perpendicular:

  • SmoothQuant fixed the loud column by migrating it into the weights (a per-channel move on the contraction axis).
  • A loud row is fixed by giving that row its own scale — a per-token move on the output axis.

One per-tensor scale is set by the single loudest token, so every quiet token is quantized on a step far too coarse for it and collapses toward zero. Per-token hands each row the tightest step it can use.

NoteKey Insight

Per-token quantization is per-channel quantization on the other axis. Per-channel weights keep the output-column axis of W independent; per-token activations keep the row axis of X independent. In code it is the same function (quantize_per_channel) with axis=0 instead of axis=1. Both axes are outer dimensions of the GEMM, which is the whole reason the scales factor out.

The Math: Two Outer Axes, One Outer Product

Per-token symmetric int8 gives row t its own absmax scale:

s^x_t = \frac{\max_j \lvert X_{t,j}\rvert}{2^{b-1}-1}, \qquad X_{t,j} \approx s^x_t \, X^q_{t,j}.

Pair it with per-channel weights (s^w_o for output column o) and watch the two scales walk straight out of the accumulation:

Y_{t,o} = \sum_k X_{t,k} W_{k,o} \approx \sum_k \big(s^x_t X^q_{t,k}\big)\big(s^w_o W^q_{k,o}\big) = s^x_t \, s^w_o \sum_k X^q_{t,k} W^q_{k,o} = \big(s^x \otimes s^w\big)_{t,o}\,\big(X^q W^q\big)_{t,o}.

The hardware accumulates \sum_k X^q_{t,k} W^q_{k,o} in int32, then a single outer-product rescale s^x \otimes s^w — one multiply per output element, after the matmul — turns it back into Y. This is LLM.int8()’s vector-wise quantization: a scale per row of X, a scale per column of W.

Contrast the forbidden move. A per-channel activation scale s^x_k would sit inside the sum, \sum_k s^x_k X^q_{t,k} W^q_{k,o}, glued to the contraction index k; no factoring pulls it out, so it would have to be applied during accumulation and breaks the int8 GEMM. The token scale s^x_t escapes only because t is an outer index.

NoteKey Insight

You could not fold a per-token scale offline — row t’s range depends on the actual activation, unknown until you see it. So per-token is dynamic: computed at run time from each row’s own absmax. No calibration set, no offline pass — just one cheap max-reduction per token, and every token keeps its own dynamic range.

Code: Vector-Wise W8A8 from Scratch

per_token.py is thin because per-token quantization reuses the module’s own quantize_per_channel — with axis=0. First, the scales: a loud token gets its own larger step instead of imposing it on everyone.

spec_pt = importlib.util.spec_from_file_location("per_token", Path("per_token.py").resolve())
pt = importlib.util.module_from_spec(spec_pt)
spec_pt.loader.exec_module(pt)

torch.manual_seed(0)
X = torch.randn(24, 16)
W = torch.randn(16, 16)
for r in (3, 11, 19):
    X[r] *= 40.0                                     # a few loud *tokens* (rows)

_, s_tok = pt.quantize_per_token(X)                  # one scale per row of X
print(f"per-token scale shape: {tuple(s_tok.shape)}   (one per token)")
print(f"quiet row 0 scale : {s_tok[0].item():.4f}")
print(f"loud  row 3 scale : {s_tok[3].item():.4f}   (its own, larger step)")
per-token scale shape: (24, 1)   (one per token)
quiet row 0 scale : 0.0167
loud  row 3 scale : 1.0742   (its own, larger step)

Now the two W8A8 recipes — naive per-tensor vs vector-wise — and their output SNR:

ref = X @ W

naive  = pt.w8a8_per_tensor(X, W)                    # one scale each for X and W
vector = pt.w8a8_vector_wise(X, W)                   # per-token X, per-channel W

naive_snr  = q.quantization_error(ref, naive)["snr_db"]
vector_snr = q.quantization_error(ref, vector)["snr_db"]
print(f"naive       W8A8 : {naive_snr:6.2f} dB")
print(f"vector-wise W8A8 : {vector_snr:6.2f} dB   (+{vector_snr - naive_snr:.1f} dB)")
naive       W8A8 :  31.01 dB
vector-wise W8A8 :  40.14 dB   (+9.1 dB)

And the factoring is exact — rescaling the int32 accumulation by the outer product equals dequantizing both operands first:

Xq, sx = pt.quantize_per_token(X)
Wq, sw = pt.quantize_weight_per_channel(W)
dequant_first = q.dequantize(Xq, sx) @ q.dequantize(Wq, sw)
print(f"max |vector-wise − dequant-then-matmul| = {(vector - dequant_first).abs().max():.2e}")
print("(the scales pull out of the sum — the identity is algebraic, not approximate)")
max |vector-wise − dequant-then-matmul| = 3.05e-05
(the scales pull out of the sum — the identity is algebraic, not approximate)

Drive the Loud-Token Magnitude

The left panel is the per-row reconstruction error at this layer: per-tensor (one shared scale) spikes on the quiet rows it crushes, while per-token stays flat. The right panel sweeps how loud the outlier tokens are — slide it and watch per-tensor’s output SNR dive while vector-wise barely moves.

_demo = pt.demonstrate_per_token(verbose=False)
ojs_define(ptDemo = {
    "loudRows": _demo["loud_rows"],
    "errTensor": _demo["err_per_tensor"],
    "errToken": _demo["err_per_token"],
    "factors": [r["loud_factor"] for r in _demo["sweep"]],
    "snrTensor": [r["per_tensor_snr_db"] for r in _demo["sweep"]],
    "snrVector": [r["vector_wise_snr_db"] for r in _demo["sweep"]],
})
TipTry This
  1. Push the magnitude right. At 1\times the two curves nearly touch — with no loud token, per-tensor is already fine. Past \sim\!20\times per-tensor dives while vector-wise holds flat: that gap is the value of per-token granularity.
  2. Read the left bars. The per-tensor error (red) is worst on the quiet rows, not the loud ones — those are the tokens whose real signal fell below the shared coarse step. Per-token (green) gives each row a step its own size, so the bars stay level.

Where the Two Axes Meet

Per-token and SmoothQuant are not rivals — they tame perpendicular outliers, so production stacks use both: SmoothQuant migrates the loud channels into the weights offline, then a per-token dynamic + per-channel weight W8A8 GEMM handles whatever per-token range remains at run time. Weights get per-channel (or group); activations get per-token. Every scale that ships lives on an outer axis — the one place a scale can sit and still let the int8 tensor cores do their job.

FP8: The Same Bits, Spent on Range

SmoothQuant paid for a per-channel transform to make int8 activations survive. There is another way to pay — change the grid instead of the data. Keep eight bits, but arrange the 256 codes as a tiny floating-point format (a minifloat) rather than an integer. This is fp8, the native 8-bit type of H100-class tensor cores, and it is Module 07’s precision.py put to work at inference.

The whole difference is where the codes sit:

  • int8 spaces its codes uniformly — a fixed step \Delta = \text{scale}/127 everywhere. Constant absolute error. But one loud value sets \Delta for the entire tensor, so every quiet value is left with a handful of codes across the empty span up to the outlier — and anything below \Delta/2 rounds to zero.
  • fp8 spaces its codes exponentially — dense near zero, coarse for large values, the step doubling every binade. Constant relative error across its normal range. A quiet value keeps its relative precision no matter how loud its neighbour.

A per-tensor scale over an outlier tensor is exactly the case where constant-relative beats constant-absolute. Drive a value through both grids and watch what each does with it:

spec_fp8 = importlib.util.spec_from_file_location("fp8_gemm", Path("fp8_gemm.py").resolve())
fg = importlib.util.module_from_spec(spec_fp8)
sys.modules["fp8_gemm"] = fg
spec_fp8.loader.exec_module(fg)

_e4m3_grid = fg.fp8_grid(fg.E4M3)
_peak = _e4m3_grid[-1]                                # 448 — e4m3's max_normal
_int8_step = _peak / 127.0
_int8_grid = [k * _int8_step for k in range(128)]     # a uniform grid at the same peak

ojs_define(
    fp8Grid = [v for v in _e4m3_grid if v > 0],
    int8Grid = [v for v in _int8_grid if v > 0],
    fp8Peak = _peak,
    int8Step = _int8_step,
)

Slide down toward the quiet channels’ scale (near 0.1) and int8’s story falls apart: its first nonzero code sits at \Delta \approx 3\.5, so any value below \Delta/2 snaps to zero — the signal is simply gone. fp8 keeps a few percent relative error the whole way down, because its ticks stay evenly spaced in log.

NoteKey Insight

int8 gives you constant absolute precision; fp8 gives you constant relative precision. When a single per-tensor scale must cover a wide dynamic range, absolute precision is spent almost entirely on the loud end and starves the quiet end — which is why fp8 tolerates the activation outlier that int8 needs SmoothQuant to remove.

The Two Formats, and the fp8 GEMM

Module 07 built both fp8 formats; we reuse them, never re-derive them:

format exp mantissa max note
e4m3 4 3 448 no infinity, one NaN slot — the forward format (weights, activations)
e5m2 5 2 57344 IEEE-like inf + NaN — more range, the gradient format

Both operands of Y = XW get a per-tensor scale that maps their largest magnitude onto the format max M, are rounded to fp8, and the matmul rescales once at the end — because the two scalars factor cleanly out of the sum:

s_X = \frac{M}{\max|X|}, \quad s_W = \frac{M}{\max|W|}, \qquad Y \;\approx\; \frac{Q(s_X X)\,Q(s_W W)}{s_X\, s_W}

where Q rounds to fp8. That factoring is exactly what fp8 tensor cores do: fp8 inputs, an fp32 accumulator, a single rescale. fp8_gemm.py is that one line.

Now run all three W8A8 paths on the same layer — but with the realistic massive-activation shape (Sun et al., 2024): the loud channel is loud yet downstream-weak (its weight row is small — the model has learned to mostly ignore it), so the signal the output needs lives in the quiet channels:

torch.manual_seed(0)
Xf = torch.randn(32, 16)
Wf = torch.randn(16, 16)
Xf[:, 3] *= 100.0                                   # a loud channel...
Wf[3, :] *= 0.01                                    # ...that barely couples downstream
ref_f = Xf @ Wf

int8_naive = fg.quantize_w8a8_int8(Xf, Wf)          # per-tensor int8, no help
int8_smooth = fg.smoothquant_linear(Xf, Wf, alpha=0.5)  # int8 after SmoothQuant
fp8_e4m3   = fg.fp8_gemm(Xf, Wf, fg.E4M3)           # fp8, no smoothing at all

snr = lambda a: q.quantization_error(ref_f, a)["snr_db"]
print(f"int8  W8A8, naive        : {snr(int8_naive):6.2f} dB   (loud channel crushes the scale)")
print(f"int8  W8A8 + SmoothQuant : {snr(int8_smooth):6.2f} dB   (needs a per-channel transform)")
print(f"fp8   W8A8, e4m3         : {snr(fp8_e4m3):6.2f} dB   (no transform — the grid does it)")
int8  W8A8, naive        :   5.64 dB   (loud channel crushes the scale)
int8  W8A8 + SmoothQuant :  38.35 dB   (needs a per-channel transform)
fp8   W8A8, e4m3         :  27.89 dB   (no transform — the grid does it)

fp8 rides out the outlier for free, landing far above collapsed int8 — though a few dB short of smoothed int8’s peak (e4m3 spends only 3 bits on the mantissa, so where int8 is well-scaled it is genuinely more precise). And the scale-factoring is exact, so an fp8 GEMM is just two rounded operands and one rescale:

s_x = fg.fp8_per_tensor_scale(Xf, fg.E4M3)
s_w = fg.fp8_per_tensor_scale(Wf, fg.E4M3)
factored = fg.fp8_round_trip(Xf, fg.E4M3, s_x) @ fg.fp8_round_trip(Wf, fg.E4M3, s_w)
print(f"max |fp8_gemm − roundtrip(X)@roundtrip(W)| = {(factored - fp8_e4m3).abs().max():.2e}")
max |fp8_gemm − roundtrip(X)@roundtrip(W)| = 2.38e-06

Where the Grid Wins — and Where It Doesn’t

The honest picture is a crossover. With no real outlier, int8’s uniform grid is the more precise of the two — seven effective bits beat three mantissa bits. Grow the outlier and int8’s per-tensor scale is dragged up until the quiet channels vanish, while fp8 holds flat. Drive the outlier magnitude and watch the two curves swap places:

_fp8_demo = fg.demonstrate_fp8_gemm(verbose=False)
ojs_define(fp8Sweep = {
    "outlier": [r["outlier"] for r in _fp8_demo["sweep"]],
    "int8": [r["int8_snr_db"] for r in _fp8_demo["sweep"]],
    "fp8": [r["fp8_snr_db"] for r in _fp8_demo["sweep"]],
})
TipTry This
  1. Find the crossover. At 1–3\times (no real outlier) int8 sits above fp8 — its seven bits of precision beat fp8’s three mantissa bits. Push past \sim10\times and int8 dives while fp8 barely moves. That swap is the entire fp8 value proposition in one picture.
  2. Read the grid widget alongside it. The int8 curve falls off exactly when the quiet channels drop below its step \Delta and snap to zero — the same underflow you can watch in the grid explorer above.
Warningfp8 is not a free lunch — and per-tensor is not the whole story

fp8’s constant-relative precision holds only in its normal range: near zero it degrades into subnormals, and above 448 (e4m3) values saturate — a truly enormous activation still clips. And per-tensor fp8 is not universally enough: frontier fp8 training (DeepSeek-V3) keeps a finer per-tile / per-block scale precisely because a single scalar can still lose the tail. The lesson is not “fp8 beats int8” but “fp8 trades peak precision for dynamic range” — choose the grid that matches your data.

The Other Half of Memory: the KV Cache

Everything so far shrinks the weights. But quantized weights are only half the story of a running model. During autoregressive decode (Module 08) each weight is read once per step, while the KV cache — the stored keys and values of every past token — grows linearly with the sequence and is re-read on every step. At long context and large batch, the KV cache, not the weights, is the memory that runs you out of GPU.

The size is easy to write down. For one token the cache holds a key and a value vector in every head of every layer:

\text{KV bytes} = \underbrace{2}_{K,\,V}\times n_\text{layers}\times n_\text{heads} \times d_\text{head}\times L \times \frac{b}{8}

For a Llama-2-7B-shaped model (32 layers, 32 heads, d_\text{head}=128) at L=4096 tokens in fp16, that is over 2 GB — for the cache alone, on top of the weights. Double the context and it doubles. This is the lever quantization has not pulled yet, and it is the single biggest one for long-context inference.

We build the fix the way KIVI does it, and the whole lesson pays off in one picture: the KV cache forces you to quantize its two halves along opposite axes.

Intuition: Keys and Values Want Opposite Axes

Lay one head’s cache out as a matrix of shape (tokens \times channels). The Key cache and the Value cache look nothing alike:

The Key cache has a few fixed channels whose magnitude is far larger than the rest, and they stay large across every token — a bright vertical stripe. (This is the same emergent-outlier-channel phenomenon that broke naive int8 in LLM.int8(), now in the cache.) The Value cache has no such structure: its magnitude is spread evenly, with no channel or token standing out.

That single difference dictates the axis. Recall the outlier lesson from the top of this module: a scale is only as fine as its widest value. So put each scale where it isolates the wide values:

  • Key → per-channel. Give every channel its own scale (take min/max over the token axis). The one outlier channel gets its own coarse scale; every other channel keeps a fine one. Quantize the Key per-token instead and that outlier sets a coarse scale for the whole token, wrecking the informative channels.
  • Value → per-token. Give every token its own scale (min/max over the channel axis). There is no outlier channel to isolate, and the attention output is a weighted sum over tokens — a contraction along the token axis — so per-token error stays confined to a token and averages out under the softmax weights.
NoteKey Insight

The Key and the Value cache are quantized along opposite axes — per-channel Keys, per-token Values — and it is not a convention, it is forced by the data: the Key has persistent outlier channels to isolate; the Value does not, and its token-axis contraction makes per-token error benign.

The Math: Per-Channel Keys, Per-Token Values

The quantizer itself is the affine (asymmetric) min/max scheme from the very top of this module — the one built for one-sided or offset ranges — now pointed at activations (the cache) instead of weights. Per group of values:

z = \min(X), \qquad s = \frac{\max(X) - \min(X)}{2^{B}-1}, \qquad Q = \Big\lfloor \tfrac{X - z}{s} \Big\rceil, \qquad \hat X = s\,Q + z

“Per-channel” and “per-token” are just which axis the min/max is taken over, and KIVI adds two practical wrinkles — a group size and a residual window:

  • Group size G. Instead of one scale over the whole token axis, split it into groups of G (KIVI uses G=32) so a scale stays local. This is the same block-quantization dial as int4-g128 for weights — here it also means a per-channel Key scale needs G tokens to exist before you can compute it.
  • Residual length R. Keep the most recent R tokens (KIVI uses R=128) in full precision — a small fp16 sliding window — and quantize only the older tokens. This solves the streaming problem the group size creates: tokens arrive one at a time, but a per-channel scale spans a whole group, so new tokens sit in the fp residual until a full group has accumulated behind them.

Step through quantizing one cache, group by group:

Code: Quantize the KV Cache from Scratch

Everything lives in kv_cache_quant.py. The one primitive is quantize_along — asymmetric group-wise quantization along a chosen axis — and Keys and Values differ only in which axis they reduce over:

import torch
import kv_cache_quant as kvq

torch.manual_seed(0)
T, C = 256, 64
K = torch.randn(T, C)
K[:, 7] = 40.0 + 0.5 * torch.randn(T)   # a persistent outlier channel
V = torch.randn(T, C)

qk = kvq.quantize_key_cache(K, num_bits=2, group_size=32)     # per-channel
qv = kvq.quantize_value_cache(V, num_bits=2, group_size=32)   # per-token

print(f"Key   scale shape {tuple(qk.scale.shape)}  -> one per channel, grouped over tokens")
print(f"Value scale shape {tuple(qv.scale.shape)}  -> one per token, grouped over channels")
Key   scale shape (8, 64)  -> one per channel, grouped over tokens
Value scale shape (256, 2)  -> one per token, grouped over channels

The axis is not cosmetic. Quantize the Key the wrong way — per-token, like a Value — and the outlier channel sets a coarse scale for every token, destroying the informative channels. Measure the reconstruction error on the non-outlier channels:

per_channel = kvq.dequantize_along(kvq.quantize_key_cache(K, num_bits=2, group_size=32))
per_token   = kvq.dequantize_along(kvq.quantize_value_cache(K, num_bits=2, group_size=32))

mask = torch.ones(C, dtype=torch.bool); mask[7] = False   # informative channels only
err_pc = (K[:, mask] - per_channel[:, mask]).pow(2).mean()
err_pt = (K[:, mask] - per_token[:, mask]).pow(2).mean()
print(f"per-channel (correct)  MSE on informative channels: {err_pc:.4f}")
print(f"per-token   (wrong)    MSE on informative channels: {err_pt:.4f}")
print(f"the wrong axis is {err_pt / err_pc:.1f}x worse")
per-channel (correct)  MSE on informative channels: 0.1563
per-token   (wrong)    MSE on informative channels: 2.7507
the wrong axis is 17.6x worse

What ultimately matters is the attention output. attention_output_error runs one attention step against the true cache and against the 2-bit cache, both ways round:

query = torch.randn(C)
right = kvq.attention_output_error(query, K, V, num_bits=2, group_size=32)
wrong = kvq.attention_output_error(query, K, V, num_bits=2, group_size=32, swap_axes=True)
print(f"per-channel K / per-token V (KIVI): {right['snr_db']:5.2f} dB SNR")
print(f"swapped axes (wrong):               {wrong['snr_db']:5.2f} dB SNR")
per-channel K / per-token V (KIVI):  7.42 dB SNR
swapped axes (wrong):                2.23 dB SNR

Finally, QuantizedKVCache is the streaming object: it keeps the last residual_length tokens in fp16 and quantizes the rest in groups as they age out — so the recent, most-attended tokens are always exact.

cache = kvq.QuantizedKVCache(num_bits=2, group_size=32, residual_length=128)
true_k = []
for _ in range(600):
    k, v = torch.randn(C), torch.randn(C)
    true_k.append(k)
    cache.append(k, v)

true_k = torch.stack(true_k)
recent_exact = torch.allclose(cache.keys()[-128:], true_k[-128:], atol=1e-6)
oldest_lossy = not torch.allclose(cache.keys()[0], true_k[0], atol=1e-4)
print(f"tokens cached: {cache.n_tokens}, key cache shape: {tuple(cache.keys().shape)}")
print(f"last 128 tokens bit-exact (residual): {recent_exact}")
print(f"oldest token quantized (lossy):       {oldest_lossy}")
tokens cached: 600, key cache shape: (600, 64)
last 128 tokens bit-exact (residual): True
oldest token quantized (lossy):       True

Interactive: The Memory–Error Dial

Drive the two knobs quantization gives you on the cache — the bit width and the residual length — and watch the two things they trade: how much GPU the cache costs, and how much the attention output degrades. The bit width sets the error; the residual buys back exactness on recent tokens at a memory cost.

TipTry This
  1. Ride the bit width down. Start at B=16 (the fp16 baseline: perfect output, full memory) and step to 2. Memory collapses ~5\times while the correct-axis output stays usable — and the swapped-axis line falls apart much faster.
  2. Buy back the recent tokens. Raise R from 0 to 512. Memory creeps up because those tokens are fp16 again — the price of keeping the most-attended tokens exact. KIVI’s default R=128 is the knee of that trade.
  3. Delete the outlier. In the bridge cell, drop the _Kd[:, 7] = ... line. With no outlier channel the two SNR lines converge — the axis only matters because the Key has outlier channels.

Microscaling (MX): The Format the Hardware Speaks

So far every scheme in this module has stored integers and paid for a floating-point scale to undo the rounding — int4 weights with an fp16 group scale, a KV cache with an fp16 min/max per group. Two of those choices are quietly expensive. The scale is a 16-bit float even though all it ever does is stretch a block, and the elements are integers even though a trained weight is really a float with a wide dynamic range.

Microscaling (MX) — an Open Compute Project standard (2023) that NVIDIA’s Blackwell tensor cores run natively — flips both. It fuses the two ideas this module already built:

  • the fp8 minifloats of Module 07 (a number is a tiny float, not an integer), and
  • the block scale of the group-wise section above (one scale per block of 32).

An MX block is just those together:

\underbrace{P_0, P_1, \dots, P_{31}}_{\text{32 tiny floats (E2M1 / E2M3 / E4M3)}} \quad + \quad \underbrace{X}_{\text{one shared scale (E8M0)}}, \qquad v_i = X \cdot P_i .

The magic is in what the shared scale is.

The Math: An E8M0 Scale and an E2M1 Element

The scale X is E8M0 — eight bits, all exponent. No sign, no mantissa. It encodes exactly one power of two, X = 2^{\,e-127} for e \in \{0,\dots,254\} (bias 127), with e = 255 reserved for NaN. Because it is a pure power of two, multiplying by X never rounds anything — it only shifts the binary point. Its entire job is to slide a block into the element format’s range. That is why one 8-bit exponent can replace a 16-bit fp16 scale with no loss: an fp16 scale spends 10 mantissa bits it did not need.

The element is a tiny float. MXFP4 uses E2M1: 1 sign, 2 exponent, 1 mantissa bit. Its whole non-negative vocabulary is eight numbers:

\text{E2M1} = \{\,0,\ 0.5,\ 1,\ 1.5,\ 2,\ 3,\ 4,\ 6\,\}

— note the widening gaps (0.5 steps near zero, then 1, then 2): a float spends its precision where the values are, unlike an integer’s even ladder. There are no infinities and no NaNs; all sixteen bit patterns are finite numbers, and the largest is 6.0. (MXFP6 uses E2M3/E3M2; MXFP8 reuses m07’s E4M3/E5M2.)

Quantizing a block is the OCP Algorithm 1 — line the loudest value up with the top of the element range, then snap:

\text{shared\_exp} = \Big\lfloor \log_2 \max_i |v_i| \Big\rfloor - e_{\max}^{\text{elem}}, \qquad X = 2^{\text{shared\_exp}}, \qquad P_i = \operatorname{round\_to\_grid}\!\left(\frac{v_i}{X}\right),

where e_{\max}^{\text{elem}} = 2 for E2M1 (its largest normal exponent). Drive a value onto the E2M1 grid below and watch it snap — the resolution is dense near zero and coarse out at 6.

spec_mx = importlib.util.spec_from_file_location("mx", Path("mx.py").resolve())
mx = importlib.util.module_from_spec(spec_mx)
sys.modules["mx"] = mx
spec_mx.loader.exec_module(mx)

_grid = mx.element_grid(2, 1)                       # the E2M1 vocabulary
_probes = [0.3, 0.7, 1.7, 2.6, 4.9, 9.0]           # values to snap, one per step
_snapped = [mx.round_to_grid(v, _grid) for v in _probes]

ojs_define(mxGridData = {
    "grid": _grid,
    "probes": _probes,
    "snapped": _snapped,
})
NoteKey Insight

An integer grid is evenly spaced; a float grid is dense near zero and coarse far out. That is exactly the shape a trained weight wants — most weights are small, a few are large — which is why 4-bit floats often beat 4-bit integers for the same bit budget once each block gets its own power-of-two exponent.

Code: MX from Scratch

mx.py builds the whole format. First the element vocabulary, enumerated directly from the bit layout — no lookup table:

print("E2M1 grid:", mx.element_grid(2, 1))          # MXFP4
print("E2M3 max: ", mx.element_grid(2, 3)[-1])      # MXFP6
print("E4M3 max: ", mx.element_grid(4, 3)[-1])      # MXFP8 element
E2M1 grid: [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
E2M3 max:  7.5
E4M3 max:  480.0

Now quantize one real block. These weights are tiny (~10^{-2}) — far below E2M1’s smallest step of 0.5 — so casting them straight to fp4 would round almost everything to zero. The shared exponent rescues them:

torch.manual_seed(0)
block = torch.randn(32) * 0.015
block[7] = 0.3                                       # the block's loudest value

fmt = mx.MX_FORMATS["mxfp4"]
X, P = mx.quantize_mx(block, fmt)
print(f"shared exponent : 2^{int(round(__import__('math').log2(X)))}  →  X = {X}")
print(f"element values  : {sorted(set(P))[:6]} ...  (all on the E2M1 grid)")

recon = torch.tensor(mx.dequantize_mx(X, P))
mx_snr = q.quantization_error(block, recon)["snr_db"]

# Baseline: no block scale (X = 1), snap straight to E2M1.
naive = torch.tensor([mx.round_to_grid(v, mx.element_grid(2, 1)) for v in block.tolist()])
naive_snr = q.quantization_error(block, naive)["snr_db"]

print(f"\nMXFP4 (shared E8M0 exponent) : {mx_snr:6.2f} dB")
print(f"fp4 with no block scale      : {naive_snr:6.2f} dB")
shared exponent : 2^-4  →  X = 0.0625
element values  : [-0.5, -0.0, 0.5, 4.0] ...  (all on the E2M1 grid)

MXFP4 (shared E8M0 exponent) :  12.71 dB
fp4 with no block scale      :   3.18 dB

The shared power-of-two exponent buys roughly +9 dB on this block — for the price of a single 8-bit number shared across all 32 elements.

Watch the algorithm run, step by step, on that same block:

_amax = block.abs().max().item()
_sexp = mx.shared_exponent(block, fmt)
_X = mx.e8m0_scale(_sexp)
_scaled = (block / _X).tolist()
_snap = [mx.round_to_grid(v, mx.element_grid(2, 1)) for v in _scaled]
_recon = mx.dequantize_mx(_X, _snap)

ojs_define(mxBlock = {
    "values": block.tolist(),
    "amax": _amax,
    "emax": 2,
    "sharedExp": _sexp,
    "scale": _X,
    "scaled": _scaled,
    "snapped": _snap,
    "recon": _recon,
    "gridMax": 6.0,
})

The Bit Budget: Why 4.25 Bits

The whole trade is one number. An E8M0 scale is 8 bits, shared across a block of 32, so every value costs its element size plus 8/32 = 0.25 bits:

\text{bits}_{\text{MXFP4}} = 4 + \frac{8}{32} = 4.25, \qquad \text{bits}_{\text{MXFP8}} = 8 + \frac{8}{32} = 8.25 .

Compare that to the group-wise section’s int4 block with an fp16 scale at the same block of 32: 4 + 16/32 = 4.5 bits. MX is cheaper per block and its scale is exact. Drive the block size below and watch the two levers move — a smaller block fits each shared exponent more tightly (higher SNR) but pays more scale overhead per value (more bits). MXFP4 at block 32 sits right at the knee.

torch.manual_seed(3)
_w = torch.randn(256) * 0.05
_w[100] = 1.4                                        # one loud weight
_sizes = [1, 2, 4, 8, 16, 32, 64, 128]
_curve = {"sizes": _sizes, "snr4": [], "snr8": [], "bits4": [], "bits8": []}
for s in _sizes:
    f4 = mx.MXFormat("mxfp4", 2, 1, block_size=s)
    f8 = mx.MXFormat("mxfp8", 4, 3, block_size=s)
    r4 = mx.mx_quantize_dequantize(_w, f4)
    r8 = mx.mx_quantize_dequantize(_w, f8)
    _curve["snr4"].append(q.quantization_error(_w, r4)["snr_db"])
    _curve["snr8"].append(q.quantization_error(_w, r8)["snr_db"])
    _curve["bits4"].append(mx.mx_effective_bits(f4))
    _curve["bits8"].append(mx.mx_effective_bits(f8))

ojs_define(mxDial = _curve)
TipTry This

Step the block size from 128 down to 1. The SNR climbs the whole way — a smaller block means each shared exponent hugs its own values more tightly — but so does the bit cost, from 4.06 at block 128 to 12 at block 1 (one scale per weight). The industry picks 32: nearly all the SNR of a tiny block for a quarter-bit of overhead.

NVFP4: A Scale That Lands Between the Powers of Two

MX is not a paper curiosity — Blackwell’s 5th-generation tensor cores execute MXFP8, MXFP6, and MXFP4 in silicon. But they also execute a different 4-bit format, NVIDIA’s own NVFP4, and the difference is the one thing E8M0 can’t do.

E8M0 is exact — a power of two only shifts the binary point, so the scale is never itself rounded — but that is also its cage: it can line a block’s peak up only to the nearest binade. A block whose ideal scale is 0.15 is stuck choosing between 0.125 and 0.25; it cannot sit in between. NVFP4 spends that exactness to buy a tighter fit. It changes two dials and adds a level:

  • a 16-element microblock (half of MX’s 32) — twice as many scales, each hugging a smaller slice of the row;
  • an E4M3 block scale instead of E8M0 — a real fp8 float (max 448), so the scale can land between the powers of two, at the cost that the scale is now itself rounded (E4M3 has 3 mantissa bits);
  • a second, per-tensor fp32 scale on top — two-level scaling, the piece that makes the E4M3 block scale expressible at all (next).

Step across a handful of blocks and watch the two scale formats choose. E8M0’s dots are pinned to the power-of-two gridlines; E4M3’s dots float free, landing right next to the ideal scale wherever it happens to fall:

import sys
from pathlib import Path

spec_nv = importlib.util.spec_from_file_location("nvfp4", Path("nvfp4.py").resolve())
nv = importlib.util.module_from_spec(spec_nv)
sys.modules["nvfp4"] = nv
spec_nv.loader.exec_module(nv)

ojs_define(nvSnaps = nv.scale_snap_demo())
NoteKey Insight

E8M0 and E4M3 are both 8-bit scales. E8M0 spends all 8 bits on exponent, so it is exact but can only be a power of two. E4M3 spends 3 of them on a mantissa, so it rounds — but it can sit between the powers of two, right where the block’s ideal scale usually is. NVFP4 bets the finer fit is worth the rounded scale.

The Two-Level Recipe

There is a catch. An E4M3 scale maxes out at 448, but a block’s peak can be anything — a scale of, say, 5000 simply isn’t representable. NVFP4’s fix is the second level: first divide the whole tensor by one fp32 global scale,

s_{\text{global}} = \frac{\text{amax}_{\text{tensor}}}{448 \times 6} = \frac{\text{amax}_{\text{tensor}}}{M_{\text{E4M3}} \cdot M_{\text{E2M1}}},

then give each 16-block its own E4M3 scale on top. Why this value? Because it is exactly what makes every block’s ideal scale representable. A block’s ideal multiplier (peak lands on the E2M1 grid’s top, 6) is \text{amax}_b / 6; divide it by s_{\text{global}} and

\frac{\text{amax}_b / 6}{s_{\text{global}}} = 448 \cdot \frac{\text{amax}_b}{\text{amax}_{\text{tensor}}} \le 448,

because \text{amax}_b \le \text{amax}_{\text{tensor}} for every block. The inequality is the whole reason the global level exists: it guarantees each local E4M3 scale fits, with the block holding the tensor’s own peak landing at exactly 448. Every value then decodes as the product of both levels:

v \approx s_{\text{global}} \cdot s_{\text{block}} \cdot P, \qquad s_{\text{block}} \in \text{E4M3}, \quad P \in \text{E2M1}.

The global scale’s only job is to make the local scale expressible; the local E4M3 scale does the accuracy work.

Code: NVFP4 from Scratch

nvfp4.py builds the two levels on top of the pieces you already have — the E2M1 grid from mx.py, and Module 07’s E4M3 rounder for the block scale (NVFP4’s scale is an fp8 number, so it snaps to the exact same grid). First the global scale and the range guarantee it buys:

torch.manual_seed(0)
weights = torch.randn(256) * 0.02
weights[42] = 0.7                                   # the tensor's loud outlier

s_global = nv.per_tensor_scale(weights)             # amax / (448 * 6)
headroom = nv.block_scale_headroom(weights)         # largest ideal E4M3 scale over blocks

print(f"s_global (fp32)                : {s_global:.3e}")
print(f"loudest block's ideal E4M3 scale: {headroom:.1f}   (E4M3 max = {nv.E4M3_MAX})")
print(f"range guarantee holds          : {headroom <= nv.E4M3_MAX}")
s_global (fp32)                : 2.604e-04
loudest block's ideal E4M3 scale: 448.0   (E4M3 max = 448.0)
range guarantee holds          : True

The loudest block lands on 448 exactly — the global scale is calibrated to that edge. Now the full two-level round-trip, and the payoff versus a naive fp4 cast:

recon = nv.nvfp4_quantize_dequantize(weights)       # 16-blocks, E4M3 scale, fp32 global
nvfp4_snr = q.quantization_error(weights, recon)["snr_db"]

naive = torch.tensor([mx.round_to_grid(v, mx.element_grid(2, 1)) for v in weights.tolist()])
naive_snr = q.quantization_error(weights, naive)["snr_db"]

print(f"NVFP4 (two-level)     : {nvfp4_snr:6.2f} dB")
print(f"fp4, no scale at all  : {naive_snr:6.2f} dB")
NVFP4 (two-level)     :  22.34 dB
fp4, no scale at all  :   6.49 dB
WarningNVFP4’s scale is not a power of two

Because the E4M3 block scale can be any fp8 value, the effective scale s_{\text{global}}\cdot s_{\text{block}} is generally not a power of two — so a plain value like 0.5 is no longer guaranteed to round-trip exactly, the way it always does under MX’s E8M0 scale. That is the price of landing between the powers of two; the tighter fit more than pays it back on real weights.

MXFP4 vs NVFP4

Put the two 4-bit formats on the same weights. To separate the block-size win from the scale-format win, we cast four ways: naive fp4 (no scale), standard MXFP4 (block 32, E8M0), the same E8M0 scale at block 16, and full NVFP4 (block 16, E4M3, two-level). The jump from block-16 E8M0 to NVFP4 is the E4M3 scale alone:

def _casts(x):
    c = nv.compare_mxfp4_nvfp4(x)
    order = ["naive_fp4", "mxfp4", "e8m0_block16", "nvfp4"]
    labels = {"naive_fp4": "fp4 (no scale)", "mxfp4": "MXFP4 · blk 32 · E8M0",
              "e8m0_block16": "blk 16 · E8M0", "nvfp4": "NVFP4 · blk 16 · E4M3"}
    return [{"name": labels[k], "snr": c[k]["snr_db"], "bpw": c[k]["bpw"]} for k in order]

torch.manual_seed(1)
_normal = torch.randn(512) * 0.03; _normal[77] = 0.9          # weights + one outlier
_heavy = (torch.randn(512) ** 3) * 0.1                        # heavy-tailed
_uniform = (torch.rand(512) - 0.5) * 0.4                      # flat, no outliers

ojs_define(nvCompare = {
    "weights + outlier": _casts(_normal),
    "heavy-tailed": _casts(_heavy),
    "uniform": _casts(_uniform),
})
TipTry This

Switch the distribution. On weights + outlier, NVFP4 opens a clear gap over standard MXFP4 — part from the smaller block (see block-16 E8M0), part from the E4M3 scale on top. On uniform weights, with no outliers to fit around, the block-scale format barely matters and the four casts nearly converge: NVFP4’s edge is a dynamic-range tool, and it earns its keep exactly where the range is wide.

WarningNot every “block-scaled fp8” is MX

DeepSeek-V3 famously trained end-to-end in fp8 with tile/block-wise scaling (1×128 activations, 128×128 weights). That is the same family of idea — group the numbers, scale each group — but it is not the OCP MX format: the elements are ordinary E4M3/E5M2, the block shapes are DeepSeek’s own, and the scales are fp32, not E8M0. When you read “fp8 with fine-grained scaling,” check whether it means MX (a ratified standard with an E8M0 scale over 32 elements) or a bespoke block-scaling recipe. They are cousins, not the same format.

Interactive Exploration

Step Through the Pipeline

Every quantization is the same six steps: measure the range, form a scale, divide, round, store the integers, and dequantize on the way back out. Step through them and watch where the lossy round sits.

Drive the Grid Yourself

Slide the bit-width and switch schemes. Fewer bits coarsen the grid; the arrows show each weight snapping to its nearest level, and the readouts track the error. Below it, the per-channel chart (built from the tested code above) shows the outlier channel’s damage — and how per-channel scales contain it.

TipTry This
  1. Watch int4 break. In the grid explorer, drag bits from 8 down to 2. At int2 there are only 4 levels — the snapped dots barely resemble the weights, and the SNR collapses. This is why naive int4 needs help (grouping, GPTQ).
  2. Re-centre with affine. At low bit-widths, switch from symmetric to affine. The grid slides to hug the data’s actual range instead of straddling zero, and the error drops — the payoff of spending a zero-point.
  3. Read the per-channel chart. The red bars (per-tensor) tower over the green (per-channel) for every quiet channel, because the outlier set their shared scale. Only the outlier channel itself is a near-tie.

Common Pitfalls

  1. Forgetting to clamp. After round(w/s) a value can land at \pm(q_{\max}+1) from floating-point error. Without clamp it overflows the integer type and corrupts the weight. Always clamp to the representable range.

  2. Symmetric quantization on skewed data. A symmetric grid centred on zero wastes half its levels on a range the data never visits (e.g. post-ReLU activations). Use affine — with a zero-point — whenever the values are one-sided.

  3. Per-tensor scales on weights. One outlier channel forces a coarse grid on the whole matrix. Per-channel (per-row) scales are nearly free — one float per channel — and are the standard for weight quantization.

  4. Expecting int4 to just work. Four bits is 16 levels; plain per-channel int4 usually loses too much. Real 4-bit uses smaller groups (a scale per 64 or 128 weights, built above) and error-correcting methods like GPTQ and AWQ.

  5. Feeding GPTQ the wrong calibration data. GPTQ steers on the Hessian H = 2XX^\top of a calibration batch. Too few samples, or samples off the model’s real distribution, give a poor H and the compensation can hurt — the method is only as good as the activations you show it. A few hundred in-distribution sequences is the usual recipe.

  6. Confusing storage bits with compute. Quantizing to int8 shrinks storage and memory bandwidth immediately. Getting the speedup also needs an integer matmul kernel; the QuantizedLinear here dequantizes-then-matmuls to stay readable, which is correct but not itself faster.

  7. Over-scaling in AWQ. Bigger s is not always better. Scaling a salient column past the group maximum raises \Delta for every weight in the group, trading the salient channel’s error for everyone else’s — this is the right wall of the \alpha curve. Search \alpha; don’t hand-pick a large scale. And read salience from activations, not weight magnitude: the two disagree, and it is the activation-heavy channels that dominate the output.

  8. Forgetting to count the group scales. Group-wise storage is not num_bits per weight — it is num_bits + scale_bits/group_size. Quote int4-g128 as 4.125 bits, not 4. Shrink the group without accounting for the scales and you can end up larger than you think: at g = 8 those fp16 scales already add two full bits per weight. Use grouped_footprint, not model_footprint, once groups are small.

  9. A group size that doesn’t divide the row. group_size must partition the quantized axis evenly. Real kernels pad the last partial group (llama.cpp uses fixed 32-wide blocks precisely so every tensor tiles cleanly); quantize_grouped raises rather than silently truncate. Pick a group_size that divides in_features (32/64/128 are the standard choices) or pad first.

  10. Reaching for fp8 as a strictly-better int8. It isn’t. fp8’s constant-relative precision only helps when the data’s dynamic range is wide; on a well-scaled tensor with no outlier, int8’s uniform grid is more precise (seven bits vs e4m3’s three mantissa bits — see the crossover widget). And fp8 still saturates: an activation above 448 clips in e4m3. Pick fp8 for range, int8 for peak precision — and remember per-tensor fp8 can still need a finer per-tile scale on the hardest layers.

  11. Trying to quantize activations per channel. It is the intuitive mirror of per-channel weights, and it is a trap: an activation channel is the contraction axis of XW, so its scale is buried inside \sum_k and cannot be factored out of the int8 GEMM — the kernel would have to rescale mid-accumulation. Activations quantize per-token (an output axis) or per-tensor; the channel axis belongs to the weights, or to SmoothQuant’s fold. When you do go per-token, remember it must be dynamic (each row’s absmax at run time) — a single calibrated activation scale cannot know a future token’s range.

  12. Treating NVFP4’s scale like E8M0’s. MX’s E8M0 scale is a power of two, so it never rounds and a plain value like 0.5 always round-trips exactly. NVFP4’s E4M3 block scale does round (3 mantissa bits), and its effective scale s_{\text{global}}\cdot s_{\text{block}} is generally not a power of two — so don’t assume nice round weights survive untouched. And don’t drop the per-tensor fp32 level: without it a loud block’s ideal scale can exceed E4M3’s 448 and become unrepresentable. Two levels, not one.

Exercises

Exercise 1: int4 vs int8 error

Quantize the same tensor to int8 and int4 and confirm int4’s SNR is roughly 24 dB lower (about 6 dB per bit).

x = torch.randn(10_000) * 0.1
q8, s8 = q.absmax_quantize(x, num_bits=8)
q4, s4 = q.absmax_quantize(x, num_bits=4)
snr8 = q.quantization_error(x, q.dequantize(q8, s8))["snr_db"]
snr4 = q.quantization_error(x, q.dequantize(q4, s4))["snr_db"]
print(f"int8 {snr8:.1f} dB,  int4 {snr4:.1f} dB,  gap {snr8 - snr4:.1f} dB")
int8 40.4 dB,  int4 15.3 dB,  gap 25.1 dB

Exercise 2: Build the per-channel win

Take a matrix, make one channel an outlier, and show per-channel MSE is far below per-tensor MSE. (Hint: quantize_per_channel(..., axis=0) vs absmax_quantize.)

# Your implementation here

Exercise 3: Footprint of a model that fits

What is the largest model (in billions of parameters) that fits in 24 GB of GPU memory at int4, leaving 4 GB for activations? Use model_footprint.

# budget = 20 GB for weights; solve params from model_footprint(params, 4)
budget_bytes = 20e9
params = budget_bytes / q.bytes_per_param(4)
print(f"~{params/1e9:.0f}B parameters fit in int4 within a 20 GB weight budget")
~40B parameters fit in int4 within a 20 GB weight budget

Exercise 4: GPTQ beats RTN where it counts

On a correlated layer, confirm GPTQ’s output loss is below RTN’s at int3, and that the gap widens as you drop to int2 (fewer bits ⇒ more rounding error ⇒ more for the compensation to recover).

torch.manual_seed(1)
W_ex = torch.randn(48, 48) * 0.1
X_ex = torch.randn(48, 48) @ torch.randn(48, 128)     # correlated inputs
H_ex = gptq.layer_hessian(X_ex)
for bits in (3, 2):
    rtn, _ = gptq.round_to_nearest(W_ex, num_bits=bits)
    gp, _ = gptq.gptq_quantize(W_ex, H_ex, num_bits=bits)
    red = 1 - gptq.output_loss(W_ex, gp, X_ex) / gptq.output_loss(W_ex, rtn, X_ex)
    print(f"int{bits}: GPTQ recovers {red:.0%} of RTN's output loss")
int3: GPTQ recovers 56% of RTN's output loss
int2: GPTQ recovers 48% of RTN's output loss

Exercise 5: AWQ finds the salient channel

Give a layer one loud activation channel, quantize with AWQ, and confirm two things: the searched \alpha^{*} is greater than 0 (the search left RTN behind), and the loud channel’s own per-channel error is smaller than under RTN.

torch.manual_seed(2)
W_a = torch.randn(32, 32) * 0.1
X_a = torch.randn(32, 128)
loud = 7
X_a[loud] *= 20.0                                     # one salient channel
W_a_hat, alpha_star, _ = awq.awq_quantize(W_a, X_a, num_bits=3)
rtn_a = awq.round_to_nearest(W_a, num_bits=3)
rtn_e = awq.channel_output_error(W_a, rtn_a, X_a)[loud]
awq_e = awq.channel_output_error(W_a, W_a_hat, X_a)[loud]
print(f"α* = {alpha_star:.2f}  (> 0: {alpha_star > 0})")
print(f"channel {loud} error:  RTN {rtn_e:.1f} → AWQ {awq_e:.1f}  ({awq_e < rtn_e})")
α* = 0.35  (> 0: True)
channel 7 error:  RTN 784.7 → AWQ 57.8  (True)

Exercise 6: The group-size dial

Confirm the two endpoints and the monotone middle: quantize_grouped with group_size == in_features reproduces quantize_per_channel bit-for-bit, and each halving of the group can only lower a weight’s scale. Then read off the bit cost of the accuracy you buy.

torch.manual_seed(1)
Wg = torch.randn(16, 64) * 0.05
Wg[4, 9] *= 40.0                                      # a weight buried in a quiet row

# Endpoint: a full-width group is exactly per-channel.
qg, sg = gq.quantize_grouped(Wg, num_bits=4, group_size=64, axis=1)
qc, sc = q.quantize_per_channel(Wg, num_bits=4, axis=0)
print("full group == per-channel:", bool(torch.equal(qg, qc) and torch.equal(sg, sc)))

prev = None
for g in (64, 32, 16, 8, 4, 2, 1):
    qi, s = gq.quantize_grouped(Wg, num_bits=4, group_size=g, axis=1)
    err = q.quantization_error(Wg, gq.dequantize_grouped(qi, s, group_size=g, axis=1))
    bits = gq.grouped_effective_bits(4, g)
    snr = "exact " if err["snr_db"] == float("inf") else f"{err['snr_db']:5.1f}"
    print(f"g={g:>3}:  {bits:5.3f} bits   SNR {snr} dB")
full group == per-channel: True
g= 64:  4.250 bits   SNR  17.8 dB
g= 32:  4.500 bits   SNR  20.0 dB
g= 16:  5.000 bits   SNR  21.5 dB
g=  8:  6.000 bits   SNR  24.0 dB
g=  4:  8.000 bits   SNR  28.8 dB
g=  2:  12.000 bits   SNR  32.3 dB
g=  1:  20.000 bits   SNR 161.8 dB

Exercise 7: The KV cache wants opposite axes

Plant a persistent outlier channel in a Key cache and confirm the KIVI rule: quantize Keys per-channel and Values per-token, and swapping the axes is much worse. Measure the reconstruction error on the informative (non-outlier) channels.

torch.manual_seed(2)
Kx = torch.randn(128, 32)
Kx[:, 6] = 30.0 + 0.4 * torch.randn(128)             # a persistent outlier channel

per_channel = kvq.dequantize_along(kvq.quantize_key_cache(Kx, num_bits=2, group_size=32))
per_token   = kvq.dequantize_along(kvq.quantize_value_cache(Kx, num_bits=2, group_size=32))

mask = torch.ones(32, dtype=torch.bool); mask[6] = False
err_pc = (Kx[:, mask] - per_channel[:, mask]).pow(2).mean()
err_pt = (Kx[:, mask] - per_token[:, mask]).pow(2).mean()
print(f"per-channel (correct): {err_pc:.4f}")
print(f"per-token   (wrong):   {err_pt:.4f}   -> {err_pt / err_pc:.0f}x worse")
# Your turn: sweep num_bits in (2, 3, 4, 8) and watch the gap shrink as bits rise.
per-channel (correct): 0.1558
per-token   (wrong):   5.4012   -> 35x worse

Exercise 8: MX turns a coarse fp4 into a usable one

Take a block of small weights, cast it straight to E2M1 (no scale), then quantize it as MXFP4 with its shared power-of-two exponent. Confirm the block scale buys a large SNR jump — and that the exact scale means a grid-aligned block reconstructs with no error at all.

torch.manual_seed(5)
blk = torch.randn(32) * 0.02
blk[11] = 0.4                                         # the block's loudest value

grid = mx.element_grid(2, 1)
naive = torch.tensor([mx.round_to_grid(v, grid) for v in blk.tolist()])
Xb, Pb = mx.quantize_mx(blk, mx.MX_FORMATS["mxfp4"])
mxrec = torch.tensor(mx.dequantize_mx(Xb, Pb))
print(f"no scale : {q.quantization_error(blk, naive)['snr_db']:5.1f} dB")
print(f"MXFP4    : {q.quantization_error(blk, mxrec)['snr_db']:5.1f} dB")

# The E8M0 scale is exact: a block that already lies on grid x 2^k is lossless.
onGrid = torch.tensor([0.5, 1.0, 1.5, 3.0]) * 2 ** -3
print("grid-aligned block exact:",
      bool(torch.equal(mx.mx_quantize_dequantize(onGrid, mx.MX_FORMATS["mxfp4"]), onGrid)))
# Your turn: sweep the block size (MXFormat(..., block_size=s)) and plot SNR vs bits.
no scale :   8.2 dB
MXFP4    :  18.4 dB
grid-aligned block exact: True

Exercise 9: SmoothQuant migrates the outlier

Give a layer one ~100× activation outlier channel. Confirm the smoothing transform is exactly invariant, then that per-tensor int8 W8A8 on the smoothed operands beats the naive W8A8 by several dB — and that the balanced \alpha = 0.5 beats over-migrating to \alpha = 1.

torch.manual_seed(4)
Xs = torch.randn(32, 24)
Ws = torch.randn(24, 24)
Xs[:, 5] *= 100.0                                    # the persistent loud channel
ref = Xs @ Ws

Xh, Wh, _ = sq.smooth(Xs, Ws, alpha=0.5)
print("invariant:", bool(torch.allclose(Xh @ Wh, ref, atol=1e-3)))

snr = lambda y: q.quantization_error(ref, y)["snr_db"]
print(f"naive  W8A8 : {snr(sq.quantize_w8a8(Xs, Ws)):5.1f} dB")
print(f"α=0.5  W8A8 : {snr(sq.smoothquant_linear(Xs, Ws, 0.5)):5.1f} dB")
print(f"α=1.0  W8A8 : {snr(sq.smoothquant_linear(Xs, Ws, 1.0)):5.1f} dB")
# Your turn: sweep alpha in 0..1 and confirm the SNR peaks in the middle (the U-curve).
invariant: True
naive  W8A8 :  32.1 dB
α=0.5  W8A8 :  41.6 dB
α=1.0  W8A8 :  31.8 dB

Exercise 10: fp8 rides out the outlier that int8 can’t

Take the massive-activation layer — a loud channel that is downstream-weak — and confirm fp8 W8A8 (no smoothing) clears naive int8 by a wide margin, yet loses to int8 once you remove the outlier. That crossover is the point of fp8.

torch.manual_seed(7)
Xe = torch.randn(32, 24)
We = torch.randn(24, 24)

snr = lambda ref, y: q.quantization_error(ref, y)["snr_db"]

# With a loud-but-weak outlier: fp8 wins for free.
Xo, Wo = Xe.clone(), We.clone()
Xo[:, 5] *= 100.0; Wo[5, :] *= 0.01
refo = Xo @ Wo
print(f"outlier layer   int8 {snr(refo, fg.quantize_w8a8_int8(Xo, Wo)):5.1f} dB   "
      f"fp8 {snr(refo, fg.fp8_gemm(Xo, Wo, fg.E4M3)):5.1f} dB")

# Without it: int8's uniform grid is more precise.
refc = Xe @ We
print(f"clean layer     int8 {snr(refc, fg.quantize_w8a8_int8(Xe, We)):5.1f} dB   "
      f"fp8 {snr(refc, fg.fp8_gemm(Xe, We, fg.E4M3)):5.1f} dB")
# Your turn: try e5m2 instead of e4m3 — more range, one fewer mantissa bit. When does it help?
outlier layer   int8   6.8 dB   fp8  28.8 dB
clean layer     int8  39.1 dB   fp8  29.2 dB

Exercise 11: Per-token rescues the quiet tokens

Build a layer with one very loud token (a whole row scaled up) and confirm that per-tensor W8A8 crushes the quiet rows while per-token (vector-wise) keeps them. Then verify the outer-product factoring is exact.

torch.manual_seed(11)
Xt = torch.randn(16, 32)
Wt = torch.randn(32, 32)
Xt[7] *= 80.0                                        # one loud token

reft = Xt @ Wt
snr = lambda ref, y: q.quantization_error(ref, y)["snr_db"]
print(f"per-tensor  W8A8 : {snr(reft, pt.w8a8_per_tensor(Xt, Wt)):5.1f} dB")
print(f"vector-wise W8A8 : {snr(reft, pt.w8a8_vector_wise(Xt, Wt)):5.1f} dB")

# The factoring identity: rescale-the-int32 == dequant-then-matmul.
Xq, sx = pt.quantize_per_token(Xt)
Wq, sw = pt.quantize_weight_per_channel(Wt)
err = (pt.w8a8_vector_wise(Xt, Wt) - q.dequantize(Xq, sx) @ q.dequantize(Wq, sw)).abs().max()
print(f"factoring max error : {err:.2e}   (algebraic — the scales pull out of the sum)")
# Your turn: which rows have the largest per-tensor error — the loud one, or the quiet ones?
# (Use pt.per_token_row_error(Xt) and look at the quiet rows.)
per-tensor  W8A8 :  33.4 dB
vector-wise W8A8 :  43.0 dB
factoring max error : 1.83e-04   (algebraic — the scales pull out of the sum)

Exercise 12: NVFP4’s E4M3 scale lands between the powers of two

Take a block of weights and compare the block scale two ways at the same block size: MX’s power-of-two E8M0 scale versus NVFP4’s fp8 E4M3 scale. Confirm the E4M3 scale sits closer to the ideal \text{amax}/6, and that the full two-level NVFP4 round-trip beats the E8M0 scale on the same 16-element blocks.

torch.manual_seed(9)
w = torch.randn(256) * 0.03
w[100] = 0.8                                          # a loud outlier

# One block: the ideal scale vs the two formats' snaps (s_global = 1).
snap = nv.scale_snap_demo([w[:16].abs().max().item()])[0]
print(f"ideal scale {snap['ideal']:.3f}   E8M0 {snap['e8m0']:.3f} (off {snap['e8m0_gap']:.0%})   "
      f"E4M3 {snap['e4m3']:.3f} (off {snap['e4m3_gap']:.0%})")

# Whole tensor, block 16: E8M0 scale vs NVFP4's two-level E4M3 scale.
e8m0_16 = mx.mx_quantize_dequantize(w, mx.MX_FORMATS["nvfp4"])   # block 16, E8M0
nvfp4 = nv.nvfp4_quantize_dequantize(w)                          # block 16, E4M3, 2-level
print(f"blk16 E8M0 : {q.quantization_error(w, e8m0_16)['snr_db']:5.1f} dB")
print(f"NVFP4      : {q.quantization_error(w, nvfp4)['snr_db']:5.1f} dB")
# Your turn: verify block_scale_headroom(w) == 448 — the global scale calibrates the amax block.
ideal scale 0.012   E8M0 0.016 (off 27%)   E4M3 0.012 (off 5%)
blk16 E8M0 :  19.2 dB
NVFP4      :  21.1 dB

Summary

Key takeaways:

  1. Quantization stores integers, not floats. Round each weight onto a grid of 2^b levels and keep a scale to undo it: w \approx s\,q. int8 is 4× smaller than fp32, int4 is 8×.

  2. Two schemes. Symmetric (absmax) centres the grid on zero — best for weights. Affine adds a zero-point to use the full code range — best for one-sided activations.

  3. SNR is the honest metric. Every bit is worth ~6 dB; int8 reconstructs weights near 40 dB, int4 near the mid-teens.

  4. Per-channel beats per-tensor. A single outlier channel wrecks a per-tensor scale; one scale per row costs almost nothing and recovers the lost precision.

  5. QuantizedLinear is the payoff. Stored per-channel int8, it matches the fp32 layer to a fraction of a percent at ~4× less memory — the building block of every quantized transformer.

  6. Group size is a dial. Split each row into blocks with their own scales: g = d_in is per-channel, g = 1 is lossless per-weight, and the useful middle costs b + scale_bits/g bits — 4.125 for the near-universal int4-g128. It is the granularity real int4 checkpoints ship.

  7. GPTQ minimizes output error, not weight error. Round-to-nearest rounds each weight blindly; GPTQ quantizes column by column and uses the calibration Hessian H = 2XX^\top to push each rounding error onto the not-yet-quantized weights. On correlated layers it roughly halves the int3/int2 output loss — and collapses exactly to RTN when the features are uncorrelated.

  8. AWQ protects the channels the activations care about. No Hessian, no backprop: measure per-channel salience s_X = \text{mean}(|X|), scale each weight column by s = s_X^{\alpha} (and the activation by 1/s — an equivalent transform), and grid-search \alpha \in [0,1] to minimize the same output loss. Since \alpha = 0 is RTN, AWQ is never worse; with real salient channels it roughly halves the int3 output loss and stacks cleanly with GPTQ.

  9. The KV cache is the other half of memory — quantize it on opposite axes. Weights are read once per step; the KV cache grows with context and is re-read every step, so at long context it dominates. KIVI quantizes the Key cache per-channel (isolating its persistent outlier channels) and the Value cache per-token (matching attention’s token-axis contraction) with the same affine min/max scheme, a group size G, and a full-precision residual window of the most recent tokens. Swap the axes and the Key’s outlier channel wrecks every token — the choice is forced by the data, and it cuts the cache to ~2-bit at near-baseline quality.

  10. Microscaling (MX) is fp8 elements + a shared power-of-two scale. The format Blackwell runs in hardware fuses this module’s two threads: tiny float elements (E2M1’s {0, .5, 1, 1.5, 2, 3, 4, 6} for MXFP4) sharing one E8M0 scale — 8 bits, all exponent — over a block of 32, with v_i = X\,P_i. Because the scale is a pure power of two it never rounds, and 8 bits shared over 32 elements is just 8/32 = 0.25 extra bits: MXFP4 = 4.25 bits, cheaper and more exact than int4 with an fp16 scale (4.5) at the same block. NVFP4 varies it with two-level scaling: a 16-element block whose scale is a real E4M3 fp8 value (so it can land between the powers of two, unlike E8M0), held in range by one per-tensor fp32 scale s_{\text{global}} = \text{amax}/(448\cdot 6) — decode is v \approx s_{\text{global}}\,s_{\text{block}}\,P, ~4.5 bits. Both ship on current hardware.

  11. SmoothQuant quantizes the activations, not just the weights. A real int8 GEMM is W8A8 — both operands int8 — but activations carry a few outlier channels 10–100× louder than the rest that wreck a per-tensor scale. Since those channels are systematic and the weights have headroom, migrate the difficulty: the equivalent transform XW = (X\operatorname{diag}(s)^{-1})(\operatorname{diag}(s)W) with s_j = \max|X_j|^{\alpha}/\max|W_j|^{1-\alpha} (default \alpha = 0.5) slides the outlier out of X and into W, and \operatorname{diag}(s)^{-1} folds into the previous LayerNorm — free at run time. Both operands are then plain per-tensor int8, worth several dB of output SNR and no mixed precision.

  12. fp8 spends its bits on range, not peak precision. Swap int8’s uniform grid for an 8-bit minifloat (e4m3, max 448 / e5m2, max 57344 — from Module 07) and the codes become exponential: constant relative error instead of constant absolute. That rides out the activation outlier with no SmoothQuant transform, because the quiet channels keep their relative precision. The per-tensor scales factor cleanly out of the GEMM (Y \approx Q(s_X X)\,Q(s_W W)/(s_X s_W) — the fp8-tensor-core recipe). But it is a trade: with no outlier, int8’s seven bits beat e4m3’s three mantissa bits — a genuine crossover — and fp8 still saturates past its max.

  13. Activations quantize per-token, on the axis that survives the GEMM. A per-channel activation scale is trapped on the contraction axis and cannot leave the sum — so the activation granularity that ships is per-token: one scale per row of X (an output axis). Paired with per-channel weights, the two scales factor out as an outer product s^x\!\otimes s^w applied to the int32 result — LLM.int8’s vector-wise recipe. It is dynamic (each row’s absmax at run time, no calibration), it is literally per-channel quantization with axis=0, and it complements SmoothQuant: SmoothQuant tames loud channels offline, per-token tames loud tokens at run time.

What’s Next

Quantization is the first lever of fast inference. You have now built the integer round-trip, group-wise storage (the int4-g128 layout every checkpoint ships), GPTQ’s Hessian-guided error feedback, and AWQ’s activation-aware scaling — the two dominant int4 weight-quant methods, from opposite directions. With SmoothQuant you crossed to the other operand — quantizing the activations for a genuine per-tensor int8 W8A8 GEMM — and with fp8 you saw the same W8A8 problem dissolve by changing the grid instead of the data: an exponential minifloat (E4M3/E5M2, from Module 07) rides out the outlier for free, at the cost of peak precision. The frontier goes further with grouped int4 threaded through GPTQ/AWQ (the error-correcting methods choosing per-group scales), AWQ’s tuned scale search, and fine-grained per-tile fp8 (DeepSeek-V3’s training recipe). You have also quantized the KV cache itself — the long-context memory lever — which hands off to the rest of fast inference: speculative decoding (a small draft model proposes, the big model verifies) and paged attention (vLLM’s KV-cache memory manager), where a quantized cache means even more tokens per block.

Going Deeper

Core Papers:

  • Jacob et al., Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference (2017) — the affine/zero-point scheme. https://arxiv.org/abs/1712.05877
  • Dettmers et al., LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (2022) — outlier features and why plain int8 breaks at scale; the mixed-precision baseline SmoothQuant replaces. https://arxiv.org/abs/2208.07339
  • Xiao et al., SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models (ICML 2023) — migrate activation outliers into the weights with the equivalent transform s = \max|X|^{\alpha}/\max|W|^{1-\alpha} for per-tensor int8 W8A8; the activation quantizer built above. Up to 1.56× faster, 2× smaller, 530B on one node. https://arxiv.org/abs/2211.10438
  • Frantar et al., GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers (2022) — accurate one-shot 4-bit weight quantization; the Cholesky Algorithm 1 built above. https://arxiv.org/abs/2210.17323
  • Frantar & Alistarh, Optimal Brain Compression (2022) — the OBQ greedy quantization/pruning framework GPTQ accelerates. https://arxiv.org/abs/2208.11580
  • Hassibi & Stork, Second Order Derivatives for Network Pruning: Optimal Brain Surgeon (1993) — the \delta = -(w_q - \hat w_q)/[H^{-1}]_{qq}\,H^{-1}_{:,q} update at the heart of it.
  • Lin et al., AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration (2023) — protect the salient channels with an activation-derived scale s = s_X^{\alpha}; the \alpha-search built above. MLSys 2024 best paper. https://arxiv.org/abs/2306.00978
  • Dettmers & Zettlemoyer, The case for 4-bit precision: k-bit Inference Scaling Laws (2022) — the block-size vs. bit-width Pareto frontier; why a small group plus 4 bits beats other budgets. https://arxiv.org/abs/2212.09720
  • Liu et al., KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache (2024, ICML) — per-channel Keys, per-token Values, group size G{=}32, residual R{=}128; ~2.6× less peak memory and up to 4× larger batch at near-baseline accuracy. The KV-cache quantizer built above. https://arxiv.org/abs/2402.02750
  • Hooper et al., KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization (2024) — 3-bit non-uniform quant with pre-RoPE per-channel Keys, per-token Values, and dense-and-sparse outlier isolation. https://arxiv.org/abs/2401.18079
  • Micikevicius et al., FP8 Formats for Deep Learning (2022) — the E4M3/E5M2 minifloats the fp8 W8A8 section and MXFP8 reuse; the format story of Module 07. https://arxiv.org/abs/2209.05433
  • Peng et al., FP8-LM: Training FP8 Large Language Models (2023) — per-tensor fp8 end to end, the recipe the fp8-GEMM section builds toward. https://arxiv.org/abs/2310.18313
  • Rouhani et al., Microscaling Data Formats for Deep Learning (2023) — the MX family built above: a block of 32 elements sharing one E8M0 power-of-two scale, v_i = X\,P_i, and the OCP Algorithm 1 quantizer. https://arxiv.org/abs/2310.10537
  • Open Compute Project, Microscaling Formats (MX) Specification v1.0 (2023) — the ratified standard: E8M0 scales, E2M1/E2M3/E3M2 elements, MXINT8. https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
  • NVIDIA, Introducing NVFP4 for Efficient and Accurate Low-Precision Inference (2024) — the Blackwell 4-bit variant: 16-element microblocks, an E4M3 block scale, and a second per-tensor fp32 scale. https://developer.nvidia.com/blog/introducing-nvfp4-for-efficient-and-accurate-low-precision-inference/
  • NVIDIA et al., Pretraining Large Language Models with NVFP4 (2025) — the two-level recipe in full: the global s_{\text{global}} = \text{amax}/(448\cdot 6) encode scale, per-block E4M3, and stability tricks for training in 4-bit. https://arxiv.org/abs/2509.25149
  • DeepSeek-AI, DeepSeek-V3 Technical Report (2024) — end-to-end fp8 training with fine-grained tile/block scaling (a cousin of MX, not the MX format). https://arxiv.org/abs/2412.19437

Practical Resources: