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)
  • How the production recipe threads group scales through GPTQ and AWQ — AWQ’s scale outside, GPTQ’s feedback across groups, stacked — at int4-g128
  • Why the scales hardware runs are constrained — E8M0 floor-snaps and E4M3 nearest-snaps threaded through the same recipe (MXFP4-g32 at 4.25 bits, the gpt-oss layout)
  • Why stochastic rounding is the cast that removes RTN’s bias — round up with probability p = (x-lo)/(hi-lo), unbiased in expectation (E[Q_{SR}(x)] = x), and why that pays in training but not in inference
  • How an NVFP4 GEMM rescales each 16-value K-block’s partial dot product by both operands’ local scales before accumulation
  • How to wrap real m06 GPTModel linear layers with two-level NVFP4 weights, while preserving tied embeddings
  • How MXINT8 combines a six-fraction-bit integer grid with one shared power-of-two scale per 32-value block

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.04338788241147995]
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 — ~3.1% over bare int4 (0.125/4; the “0.4%” you may see quoted is vs fp32, 0.125/32), 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” — its ≈ 0.15 counts the asymmetric path too, one fp16 scale plus a 4-bit zero-point per group: b_{\text{eff}} = 4 + (16 + 4)/128 \approx 4.16.

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   92901.6   (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| = 3.05e-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| = 6.10e-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.86e-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 n \times \frac{b}{8}

For a Llama-2-7B-shaped model (32 layers, 32 heads, d_\text{head}=128) at n=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.

The Production Recipe: Group Scales Through GPTQ and AWQ

You have built three separate wins: groups localize the grid, AWQ reshapes salience before quantizing, GPTQ compensates rounding after it. A production int4 checkpoint uses all three at once — group scales threaded through the error-correcting methods, not beside them. When a checkpoint says “int4-g128,” it means this recipe: each group of 128 consecutive columns carries its own scale, AWQ picks an activation-derived scale outside the grouped quantizer, GPTQ feeds rounding error forward across the group grids, and the two corrections stack.

The three ideas compose because they attack different parts of the error. Groups shrink the grid where it matters: one loud weight coarsens only its own 128 columns instead of the whole row. AWQ reshapes the problem first: it reads per-channel salience off the activations (s_X = \text{mean}(|X|)), scales the loud weight columns up by s = s_X^{\alpha} — dividing the activation by the same s, so the product is unchanged — and grid-searches the single knob \alpha \in [0, 1]. That search happens outside the quantizer: for each \alpha, the grouped quantizer runs inside, on W\,\text{diag}(s). GPTQ then cleans up after: it quantizes column by column and steers each rounding error onto the not-yet-quantized columns through the calibration Hessian H = 2XX^{\top} (plus 1% damping, as before). Reshape first, compensate after — the levers are orthogonal, so they stack.

One rule makes the threading honest: recompute each group’s grid at its boundary. By the time GPTQ reaches a later group, earlier rounding errors have been pushed into it — its weights are no longer the original ones. So the scale for group G is fit to the most-current compensated W when G’s first column is reached, never frozen from the original weights upfront:

s_G = \frac{\max_{i \in G} |w_i^{\text{(compensated at boundary)}}|}{q_{\max}}.

Freezing every grid upfront is the natural shortcut and the wrong one — the two rules measurably differ, and the lesson’s pitfalls call it out.

The bit accounting grows one term. Each group still carries its fp16 scale (16 bits per group, the symmetric teaching simplification used throughout this module — no zero-point). The asymmetric path the papers ship also stores one b-bit integer zero-point per group:

b_{\text{eff}} = b + \frac{16}{g} \quad \text{(symmetric)}, \qquad b_{\text{eff}} = b + \frac{16 + b}{g} \quad \text{(asymmetric)}.

At int4-g128 that is 4.125 vs 4 + 20/128 \approx 4.16 — the GPTQ paper’s “≈ 0.15 extra bits” is the asymmetric count. recipe_bpw in grouped_recipe.py computes both; state which assumption you mean every time you quote a number.

WarningB = 128 Is Not g = 128

GPTQ processes columns in lazy blocks of B (default here: one group) — a throughput knob that only decides how many columns accumulate before the tail update. The group size g is an accuracy knob — how many columns share a scale. The paper’s B = 128 coinciding with g = 128 is a coincidence, not an identity: changing B regroups the same arithmetic (the loss barely moves), while changing g regrids every weight. grouped_gptq_quantize takes both arguments separately so you cannot confuse them.

Run the whole recipe below. The demo layer has both properties the recipe answers — loud weights buried in quiet rows plus correlated, salient-channel inputs — and reports the output loss at each stage, RTN → grouped → +AWQ → +GPTQ → stacked:

spec_gr = importlib.util.spec_from_file_location("grouped_recipe", Path("grouped_recipe.py").resolve())
gr = importlib.util.module_from_spec(spec_gr)
sys.modules["grouped_recipe"] = gr
spec_gr.loader.exec_module(gr)

recipe_demo = gr.demonstrate_recipe()
print(f"symmetric b_eff: {recipe_demo['bpw']:.3f} bits,   asymmetric: {recipe_demo['bpw_asym']:.5f} bits")
================================================================
THE PRODUCTION RECIPE  (output loss ||WX - What X||^2)
================================================================
Layer (64, 1024), int4-g128, 4 salient channels @ 12.0x

     rtn  8611364.0000
 grouped  3905020.7500
     awq  3035712.2500
    gptq   314380.6250
 stacked   304223.3125

Asymmetric b_eff @ g=128: 4.15625 bits
symmetric b_eff: 4.125 bits,   asymmetric: 4.15625 bits
ojs_define(recipeStages = [
    {"stage": name, "loss": loss} for name, loss in recipe_demo["stages"].items()
])
ojs_define(recipeSweep = recipe_demo["sweep"])
NoteKey Insight

Read the two charts together: the bars fall at every stage — grouping beats per-channel RTN, AWQ and GPTQ each beat grouping alone, and the stack wins overall — while the dial shows what the accuracy costs. From g = 1024 down to 128 you buy decibels almost for free (4.02 \to 4.125 bits); past 64 the bits climb steeply for little SNR. The knee at 64–128 is why the industry stops at 128.

TipTry This
  1. Move the dial left. Step g from 1024 down to 32 and watch the readout: SNR climbs fast, then flattens — while the bit cost accelerates. The knee is the whole economic argument for g128.
  2. Break the boundary rule. Re-run the demo cell with grid="upfront" in grouped_gptq_quantize (freeze every grid from the original weights) and compare the GPTQ loss — the gap is the price of ignoring the compensation the feedback already pushed downstream.
  3. Split the stack. Compare the awq and gptq bars: AWQ wins by reshaping salience, GPTQ by compensating rounding — then confirm stacked sits below both, the signature of orthogonal levers.

Constrained Scales: The Recipe Hardware Actually Runs

Every scale in the recipe above is a free fp32 number — s_G = max|W_G| / q_max, any positive real. In silicon that freedom does not exist. A constrained scale is a group scale snapped to one of the two grids the hardware GEMM understands, and the recipe above must run on the snapped grid, not the free one. Two forces demand it. First, storage: the scale ships with the weights, so an fp32 scale per group would double the scale tax the bit budget just counted. Second, factoring: the scale must pull out of the accumulation, and Blackwell’s tensor cores factor exactly two shapes — a power-of-two scale shared over 32 values (MXFP4) or an fp8 scale over 16 values (NVFP4). This is not hypothetical: gpt-oss ships its MoE weights — 90+% of its parameters — in MXFP4 at 4.25 bits each, which is what lets the 120B model fit on a single 80 GB GPU.

The two grids snap differently, and the difference is the whole section. E8M0 holds only an exponent, so its scales are pure powers of two and the spec default snaps down: the largest power of two at or below the ideal (floor). E4M3 is a real fp8 value, so it snaps to nearest and can land between the powers of two — at the cost of a scale that is itself rounded. Floor both coarsens every bucket and clips the group’s peak (any weight above 7\hat{s} clamps), which is why the E8M0 bar below towers over the other two — and why real MX fits the exponent to the block peak with element headroom instead of flooring a fitted scale.

\hat{s} = 2^{\lfloor \log_2 s \rfloor} \quad \text{(E8M0: floor, the spec default)}, \qquad \hat{s} = \text{round}_{\text{E4M3}}(\min(s, 448)) \quad \text{(E4M3: nearest, capped)}.

One subtlety the stepper calls out: rounding the ideal scale up instead of down is a documented implementation choice (Mishra et al., cited in the NVFP4 paper’s background) — it never clips the peak, at the cost of coarsening every bucket. The OCP MX v1.0 spec (§6.3) default is floor (“the largest power-of-two less than or equal to”), and this module teaches floor, with round-up shown as the contrast.

The bit arithmetic gets cheaper, because the scale itself shrinks from fp16 to 8 bits. One 8-bit scale amortized over the group:

b_{\text{eff}} = 4 + \frac{8}{32} = 4.25 \quad \text{(MXFP4-g32, symmetric — no zero-point)}.

That is the number gpt-oss quotes. And the NVFP4 half of the story reuses the two-level machinery you already built: decode is still v \approx s_{\text{global}}\,s_{\text{block}}\,P, and block_scale_headroom still reads 448.0 at the amax block — the range guarantee carries over unchanged. mx_recipe.py holds the whole path: snap_scale_e8m0, snap_scale_e4m3, constrained_grouped_quantize (the GPTQ/AWQ recipe with snapped boundary scales), constrained_vs_free, and demonstrate_mx_recipe.

Run it. The demo layer is the same honest shape as the g128 one — buried loud weights plus correlated, salient inputs — with the recipe run three ways:

spec_mxr = importlib.util.spec_from_file_location("mx_recipe", Path("mx_recipe.py").resolve())
mxr = importlib.util.module_from_spec(spec_mxr)
sys.modules["mx_recipe"] = mxr
spec_mxr.loader.exec_module(mxr)

mx_demo = mxr.demonstrate_mx_recipe()
print(f"MXFP4-g32 b_eff: {mx_demo['bpw']:.3f} bits (symmetric)")
group-32 4-bit output loss (lower is better):
  free: 41779.0156
  e8m0: 94056.8984
  e4m3: 51179.5039
MXFP4-g32 b_eff: 4.250 bits (symmetric)
MXFP4-g32 b_eff: 4.250 bits (symmetric)
ojs_define(mxStages = [
    {"stage": name, "loss": loss} for name, loss in mx_demo["stages"].items()
])
ojs_define(mxSweep = mx_demo["sweep"])
ojs_define(mxSnapRow = max(
    mx_demo["snaps"], key=lambda r: abs(r["e8m0"] - r["ideal"]) / r["ideal"]
))

Step through what happens to one illustrative group scale — the demo’s row with the largest E8M0 gap, where floor costs roughly half the scale and E4M3 stays within a few percent:

Now the price, measured. Three runs of the same GPTQ-threaded recipe — free, E8M0-snapped, E4M3-snapped — with the group-size dial re-cutting all three bars from the sweep:

NoteKey Insight

Read the dial against the bars: E4M3 hugs the free bar at every group size — the fractional scale recovers almost everything flooring throws away — while E8M0 towers over both, worst at small groups where each clipped peak poisons a larger share of the weights. The ordering at the staged g32 (free below both constrained, pinned in the tests) is a same-seed observation, not a theorem: snapping rewrites the feedback trajectory, so at some sizes a constrained run can dip under free. The gap is the price of hardware — the recipe pays it and ships anyway, at 4.25 bits.

TipTry This
  1. Price the clip. Step the snapper to step 1, then step 3: floor’s gap comes from clipping the peak (values above 7\hat{s} clamp), round-up’s from coarsening every bucket. Recompute the row in the demo cell with 2 ** math.ceil(math.log2(ideal)) and compare the two gaps.
  2. Shrink the groups. Move the dial from 128 down to 16 and watch the E8M0/free ratio: finer groups localize each clipped peak, so the floor tax falls — read the exact ratios off mx_demo["sweep"]["e8m0_gap"].
  3. Break the ordering. At g=128 the E8M0 bar can dip below free. That does not contradict the hardware price: the snapped grid steers GPTQ’s feedback down a different path, and output loss is not monotone in grid fineness once feedback is involved. Find the crossover, then confirm the staged g32 anchor (constrained_vs_free at the demo seed) still holds.

Stochastic Rounding: The Unbiased Cast

Every cast in this module has been round-to-nearest (RTN): GPTQ’s column-by-column snap, AWQ’s \alpha-search, the E8M0 floor, the E2M1 element snap. Each asks the same question — which grid point is closest? — and answers it the same way every time. That determinism is exactly what inference wants. It is also a bias, and a bias is what a training run cannot afford.

Put a value 0.49\Delta above a grid point and RTN always rounds it down: the error is -0.49\Delta, the same sign and the same size on every cast of that value, forever. Put it 0.51\Delta above and RTN always rounds up, +0.49\Delta. The error is a deterministic function of where the value sits, so it never averages out — a parameter that should drift upward gets nudged down on every single update. A zero-mean error cancels over a long enough run. A fixed-sign error accumulates.

Stochastic rounding (SR) is the cast that removes the bias. The grid and the scale do not change; only the decision does. Instead of taking the nearest neighbour, SR takes one of the two neighbours with probability proportional to how close it is. The rounding error becomes a random variable with mean zero, and the parameter moves in the right direction on average — which, for a quantity re-quantized once per update, is the only thing that matters.

One honest caveat before the machinery: SR is not a better cast. A single stochastic cast is worse than a nearest one — its worst case is twice as large. SR pays only where the same quantity is cast thousands of times, which is why it is a training cast and why inference leaves it alone.

Intuition: Flip a Coin Instead of Deciding

Here is the mental model. RTN stands at a value and asks a yes/no question — round up? — and always gives the same answer. SR replaces the answer with a coin weighted by distance: the closer neighbour gets the larger share of the probability mass, and the two shares add to one.

Sit a value 90% of the way up the interval and the coin is 90/10 in favour of rounding up; sit it 10% of the way up and it is 10/90. Sit it exactly at the midpoint and the coin is fair — the one position where SR is maximally noisy, and the one where it has the most room to be wrong in either direction.

Two things fall out of that picture immediately. First, the decision is randomized but the grid is not: SR never invents a value. Every cast lands on one of the same two neighbours RTN chooses between, so the reconstruction error is never worse than one full step. Second, the randomness is what makes the mean exact. RTN’s answer for a given value is always the same grid point — cast it again and you get the same answer, including the same error. SR’s answer changes from cast to cast, and the set of answers it produces has the original value at its centre. Repeat the cast and the errors cancel; repeat RTN and they do not.

The Math: Round Up With Probability p

Take a value x between two adjacent grid points lo < hi, with \Delta = hi - lo. SR rounds up with probability

p = \frac{x - lo}{hi - lo} = \frac{x - lo}{\Delta},

and down with 1 - p: the closer neighbour is the likelier one, and a value exactly on-grid has p = 0. That is the whole method. The NVFP4 paper describes it as rounding “a value probabilistically to one of its two nearest representable numbers, with probabilities inversely proportional to their distances” (§4.4), and Croci et al. give the same rule as their eq 7.3.

Now compute the expected value of the cast. There are only two outcomes, so the expectation is two terms — and they simplify:

E[Q_{SR}(x)] = p \cdot hi + (1 - p)\cdot lo = \frac{x - lo}{\Delta}\,hi + \frac{hi - x}{\Delta}\,lo = \frac{x(hi - lo)}{\Delta} = x.

That is not a sampling result, and it is not “true for enough draws”. It is an identity: whatever the value, whatever the grid, the expected cast is the value. The rounding error \delta = Q_{SR}(x) - x is therefore mean-zero. Croci et al. prove the stronger version (Thm 6.3) — the errors are not merely zero-mean but mean-independent, E(\delta_i \mid \delta_{i-1}, \dots, \delta_1) = E(\delta_i) = 0 — so a sum of stochastic roundings has expectation equal to the exact sum regardless of the order the sum is taken in (Thm 6.4).

RTN is the complement. Its error for a fixed x is one deterministic number, bounded by half a step:

\bigl|E[Q_{RTN}(x)] - x\bigr| \le \frac{\Delta}{2},

with equality whenever x sits exactly halfway between grid points. Read that carefully: the bias is not small “on average over a training run” — it is the same offset every time you cast the same value. Ten thousand updates at -0.49\Delta each is a drift of -4900\Delta, not zero.

The variance is where SR pays for itself. El Arar et al. (2023) give the exact variance of the cast, V(\hat{x}) = \varepsilon(x)^2\,\theta(x)\,(1 - \theta(x)) with \theta the up-probability, which is maximized at \theta = 1/2:

V \le \frac{\Delta^2}{4}.

The spread of the error is capped at a quarter of a step squared, worst at the midpoint where the coin is fair. The spread is the price; the zero mean is the point.

And the price is real in the per-cast worst case. Croci et al. (eq 6.4) bound the stochastic rounding error by |\delta| \le 2u against RN’s |\delta| \le u — “the worst-case error bounds for SR are a factor 2 larger than those for RN, however, and so SR does not benefit all computations.” In grid terms: a single SR cast can land a full step away (\Delta) where RTN never moves more than half of one (\Delta/2). A one-shot cast is the clearest case where SR does not pay.

Accumulation is the mirror image. In a sum of n terms the worst-case error grows like n under RTN but only like \sqrt{n} under SR: Croci et al.’s probabilistic constant is \tilde{\gamma}_n(\lambda) = \exp(\lambda\sqrt{n}\,u + n\,u^2)/(1 - u) - 1 with probability at least 1 - 2\exp(-\lambda^2/2), against the worst case \gamma_n = nu/(1 - nu). El Arar et al. (2026) state the contrast flatly: “In length n summation, the worst-case error of SR grows as sqrt(n) with high probability, unlike … round-to-nearest (RN), which grows as n.” Root-n noise is the good case — it averages.

NoteKey Insight

Unbiased is not the same as error-free. SR buys a smaller bias, not a smaller error: the expected cast is exactly x, but any single cast is still off by up to a full step — twice RTN’s half-step bound. What SR changes is the character of the error: RTN’s is a fixed offset that compounds, SR’s is zero-mean noise that cancels as 1/\sqrt{n} over repeated casts. Everything downstream follows from that trade: it is a win where repetition exists and a loss where it does not.

Code: Stochastic Rounding from Scratch

stochastic_round.py holds the whole path: sr_round_up_probability, stochastic_round_scalar, sr_absmax_quantize, sr_quantize_dequantize, sr_variance_bound, sr_bias_vs_rtn, accumulate, accumulation_error, rtn_stall_radius, quantized_sgd, mx_stochastic_cast, and the story layer demonstrate_stochastic_round.

The probability first — one line of arithmetic, and the coin it weights:

spec_sr = importlib.util.spec_from_file_location("stochastic_round", Path("stochastic_round.py").resolve())
sround = importlib.util.module_from_spec(spec_sr)
sys.modules["stochastic_round"] = sround
spec_sr.loader.exec_module(sround)

lo, hi = 0.5, 1.0                                    # one interval of the grid
for probe in (0.55, 0.75, 0.95):
    p = sround.sr_round_up_probability(probe, lo, hi)
    print(f"x={probe:.2f}  p_up={p:.3f}  →  {p:.0%} up / {1 - p:.0%} down")
x=0.55  p_up=0.100  →  10% up / 90% down
x=0.75  p_up=0.500  →  50% up / 50% down
x=0.95  p_up=0.900  →  90% up / 10% down

Same grid, same scale, two casts. Put SR and RTN on the identical symmetric absmax grid and compare:

torch.manual_seed(0)
sr_x = torch.randn(4096) * 0.05

sr_q, sr_scale = sround.sr_absmax_quantize(sr_x, num_bits=4, generator=torch.Generator().manual_seed(0))
rtn_q, rtn_scale = q.absmax_quantize(sr_x, num_bits=4)
sr_deq = q.dequantize(sr_q, sr_scale)
rtn_deq = q.dequantize(rtn_q, rtn_scale)

sr_levels = torch.unique(sr_deq)                     # the grid both casts share
sr_gaps = torch.diff(sr_levels)
sr_step = float(sr_gaps[sr_gaps > 1e-9].min())       # one grid interval

print(f"scale: SR {float(sr_scale):.6f}   RTN {float(rtn_scale):.6f}   (grid spacing Δ = {sr_step:.4f})")
print(f"mean error  : SR {float((sr_deq - sr_x).mean()):+.2e}   RTN {float((rtn_deq - sr_x).mean()):+.2e}")
print(f"max |error| : SR {float((sr_deq - sr_x).abs().max()):.4f}   RTN {float((rtn_deq - sr_x).abs().max()):.4f}")
print(f"bounds      : SR ≤ Δ = {sr_step:.4f}   RTN ≤ Δ/2 = {sr_step / 2:.4f}   "
      f"variance ≤ Δ²/4 = {sround.sr_variance_bound(sr_step):.4f}")
scale: SR 0.029296   RTN 0.029296   (grid spacing Δ = 0.0293)
mean error  : SR +2.10e-04   RTN -4.06e-05
max |error| : SR 0.0293   RTN 0.0146
bounds      : SR ≤ Δ = 0.0293   RTN ≤ Δ/2 = 0.0146   variance ≤ Δ²/4 = 0.0002

Both casts fit the same grid — the scales agree to the last digit — and the max |error| column is the price of the coin: RTN never moves more than half a step, while SR’s worst cast lands a full step away, exactly as the 2u bound warns. Over a symmetric ensemble both mean errors sit near zero, because RTN’s offsets cancel across values; what does not cancel is which value gets which error, and that is what the next cell isolates.

Repeat the cast on single values inside one grid interval, where RTN’s decision is fixed by position:

sr_probes = [0.05, 0.15, 0.25, 0.35, 0.45]           # positions inside one interval
sr_grid_step = 0.5

print(f"{'x':>6}  {'p_up':>6}  {'RTN error':>10}  {'SR mean':>9}  {'SR std':>8}")
sr_rows = []
for v in sr_probes:
    p = sround.sr_round_up_probability(v, 0.0, sr_grid_step)
    r = sround.sr_bias_vs_rtn(v, sr_grid_step, draws=4000,
                              generator=torch.Generator().manual_seed(0))
    sr_rows.append((p, float(r["rtn_bias"]), float(r["sr_bias"]), float(r["sr_std"])))
    print(f"{v:>6.2f}  {p:>6.3f}  {float(r['rtn_bias']):>+10.4f}  "
          f"{float(r['sr_bias']):>+9.4f}  {float(r['sr_std']):>8.4f}")

sr_n = len(sr_rows)
print(f"\nover {sr_n} positions: mean |RTN error| "
      f"{sum(abs(row[1]) for row in sr_rows) / sr_n:.4f}   "
      f"mean |SR mean| {sum(abs(row[2]) for row in sr_rows) / sr_n:.4f}   "
      f"max SR std {max(row[3] for row in sr_rows):.4f}   "
      f"(std ≤ Δ/2 = {sr_grid_step / 2:.4f})")
     x    p_up   RTN error    SR mean    SR std
  0.05   0.100     -0.0500    +0.0019    0.1525
  0.15   0.300     -0.1500    +0.0010    0.2296
  0.25   0.500     +0.2500    +0.0037    0.2500
  0.35   0.700     +0.1500    +0.0055    0.2267
  0.45   0.900     +0.0500    -0.0000    0.1500

over 5 positions: mean |RTN error| 0.1300   mean |SR mean| 0.0024   max SR std 0.2500   (std ≤ Δ/2 = 0.2500)

Read the columns. RTN error is a deterministic function of position — negative below the halfway point, positive at or above it, and the same number every time you cast that value, so no amount of repetition moves it. SR mean sits at zero to within its own noise at every position: the individual casts jump between lo and hi, but their average is the value itself. SR std is the price, and it peaks at the midpoint (p = 0.5), where the coin is fair.

The Absorbed Update

Now the fact that makes SR a training ingredient. A parameter being optimized is a large number that receives a small update every step — and if the update is smaller than half a step, RTN rounds it to zero. The update is not merely quantized, it is absorbed: the parameter does not move at all. Croci et al. name this stagnation — “unlike RN, it is immune to the phenomenon of stagnation, whereby a sequence of tiny updates to a relatively large quantity are lost” — and their §6.1 example is stark: starting from x = 1, ten additions of \varepsilon_M/4 leave RN’s answer at exactly 1, while SR lands on the exact sum. Note which half of that is exact: RTN’s lane needs no sampling to be frozen, because a value below \Delta/2 rounds to the same grid point every time.

The demo layer accumulates \delta = 0.1 two hundred times onto a grid whose spacing is \Delta = 0.5 — every update is well under \Delta/2 = 0.25, so RTN’s answer is fixed before the run starts:

sr_demo = sround.demonstrate_stochastic_round(verbose=False)

sr_rtn_losses = sr_demo["sgd"]["rtn"]["losses"]
sr_freeze = next((i for i in range(1, len(sr_rtn_losses))
                  if abs(sr_rtn_losses[i] - sr_rtn_losses[i - 1]) < 1e-12), None)
sr_len = len(sr_demo["rtn_traj"])                    # steps + 1: index n is after n updates
sr_checks = sorted({min(sr_len - 1, k) for k in (1, 2, 4, 8, 16, 32, 64, 128, sr_len - 1)})

ojs_define(srAbsorb = {
    "delta": float(sr_demo["delta"]),
    "step": float(sr_demo["step"]),
    "checkpoints": sr_checks,
    "exact": sr_demo["exact_traj"],
    "rtn": sr_demo["rtn_traj"],
    "sr": sr_demo["sr_traj"],
})
ojs_define(srSgd = {
    "rtn": sr_rtn_losses,
    "sr": sr_demo["sgd"]["sr"]["losses"],
    "rtnFinal": float(sr_demo["sgd"]["rtn"]["final_loss"]),
    "srFinal": float(sr_demo["sgd"]["sr"]["final_loss"]),
    "rtnStall": [float(r) for r in sr_demo["sgd"]["rtn"]["stall_radius"]],
    "freeze": sr_freeze,
})
ojs_define(srMxBias = {
    "rtn": float(sr_demo["mx"]["rtn_bias"]),
    "sr": float(sr_demo["mx"]["sr_bias"]),
})
print(f"p at the midpoint : {float(sr_demo['p_at_midpoint']):.3f}")
print(f"measured bias     : SR {float(sr_demo['sr_bias']):+.5f}   "
      f"RTN {float(sr_demo['rtn_bias']):+.5f}   "
      f"(SR std {float(sr_demo['sr_std']):.5f}, bound {float(sr_demo['variance_bound']):.5f})")
print(f"accumulate δ={float(sr_demo['delta'])} × {int(sr_demo['steps'])} on Δ={float(sr_demo['step'])}: "
      f"exact {float(sr_demo['exact_traj'][-1]):.2f}   "
      f"RTN {float(sr_demo['rtn_traj'][-1]):.2f} (error {float(sr_demo['rtn_final_error']):.2f})   "
      f"SR {float(sr_demo['sr_traj'][-1]):.2f} (error {float(sr_demo['sr_final_error']):.2f})")
p at the midpoint : 0.500
measured bias     : SR +0.00375   RTN +0.25000   (SR std 0.25000, bound 0.06250)
accumulate δ=0.1 × 200 on Δ=0.5: exact 20.00   RTN 0.00 (error -20.00)   SR 21.50 (error 1.50)

Step through the run and watch the two lanes separate:

RTN’s lane is a flat line at zero for the entire run — the update is absorbed before it ever moves the sum — while SR’s walks upward with the exact sum at its centre and an error that grows like \Delta\sqrt{n}/2, not like \Delta n. That is the whole argument for SR in one picture: the coarse grid does not lose the update, it just delivers it late.

TipTry This
  1. Watch the walk start. Step to 1, then 2, then 8 updates: SR moves in whole steps of \Delta (it can only land on grid points), so the lane jumps rather than creeps. Compare its error at each checkpoint to \Delta\sqrt{n}/2.
  2. Freeze the other lane. Rerun the accumulation with mode="rtn" in accumulate and confirm every entry is exactly 0.0 — the frozen lane is arithmetic, not a sampling accident. Then rerun it under SR with a different generator seed and watch the final value change while the mean does not.
  3. Cross the absorption boundary. Rerun with delta=0.3 — still under \Delta = 0.5, but now above \Delta/2 = 0.25. RTN’s lane leaves zero, because the update now rounds up to a full step. Absorption is a threshold at \Delta/2, not a rule about “small” updates.

Where It Pays: Training, Not Inference

The regime is the whole story, and it is narrow: SR pays where the same quantity is re-quantized thousands of times, and costs where it is not. A parameter accumulating gradient updates is the first case. It receives a small update every step, and under RTN that update can vanish entirely — Gupta et al. (2015, §4) put it plainly: “any parameter update in the range (-\varepsilon/2, \varepsilon/2) is always rounded to zero.” Zamirai et al. (2020) describe the same failure for bfloat16 weight updates, where nearest rounding “often cancels small updates, which degrades the convergence and model accuracy”, and their remedy (SR, or Kahan summation) bought “up to 7% absolute validation accuracy gain”.

A one-shot cast is the second case. An inference pass quantizes a weight once and then uses it; there is no second cast to average against, so SR would only add noise on top of the same grid — a strictly worse deal, since its worst case is twice RTN’s. Google Cloud’s guidance is blunt: “stochastic rounding is designed for training only. Because it is non-deterministic, you should stick to standard Round-to-Nearest for inference workloads where consistent outputs are required.” The same source states the training payoff in the language of this module: in a sum of N operations, RTN’s systematic error grows O(N) while SR’s grows O(\sqrt{N}).

Watch the two regimes run side by side. quantized_sgd does SGD on a two-parameter quadratic with a 4-bit grid — same learning rate, same start, the only difference is the cast:

sr_sgd = sr_demo["sgd"]
sr_radii = " / ".join(f"{float(r):.4f}" for r in sr_sgd["rtn"]["stall_radius"])
print(f"RTN: final loss {float(sr_sgd['rtn']['final_loss']):.4f}   "
      f"stalls at a Δ/2 radius of {sr_radii} per parameter")
print(f"SR : final loss {float(sr_sgd['sr']['final_loss']):.4f}   "
      f"no stall radius — the update never stops moving")
if sr_freeze is None:
    print("RTN's loss is still moving at the last step — the run ended before it froze.")
else:
    print(f"RTN's loss stops moving at step {sr_freeze} of {len(sr_rtn_losses) - 1}.")
RTN: final loss 0.6020   stalls at a Δ/2 radius of 1.4286 / 0.1429 per parameter
SR : final loss 0.0000   no stall radius — the update never stops moving
RTN's loss stops moving at step 4 of 300.

The two curves tell the same story as the lanes. RTN descends until its update falls under \Delta/2 — the stall radius the engine computes analytically in rtn_stall_radius(lr, step) — and then stops, because every remaining update rounds to zero. SR keeps descending past that point: the same updates that RTN discards still move the parameter, one grid step at a time, with the errors cancelling instead of accumulating.

Stochastic Rounding in the Hardware Formats

SR is not a simulation-only idea. Blackwell has native SR for FP4 conversion instructions — “Blackwell GPUs have native support for several rounding modes including round-to-nearest-even and stochastic rounding for FP4 conversion instructions” (NVFP4 paper §2) — and NVIDIA’s Transformer Engine documents the payoff directly: “each scaled value is probabilistically rounded to one of the two nearest representable NVFP4 values … which ensures that the expected value of the quantized tensor equals the original value, eliminating systematic quantization bias during training.”

Run it on a real MX block. Take a block of 32 and cast it twice — once with the deterministic E2M1 element snap, once with the stochastic one:

torch.manual_seed(7)
sr_block = torch.randn(32) * 0.02
sr_block[9] = 0.35                                   # one loud element sets the block peak
sr_fmt = mx.MX_FORMATS["mxfp4"]
sr_exp = mx.shared_exponent(sr_block, sr_fmt)
sr_scale = mx.e8m0_scale(sr_exp)
sr_rtn_block = mx.mx_quantize_dequantize(sr_block, sr_fmt)
sr_sr_block = sround.mx_stochastic_cast(sr_block, fmt=sr_fmt,
                                        generator=torch.Generator().manual_seed(0))
sr_draws = torch.stack([
    sround.mx_stochastic_cast(sr_block, fmt=sr_fmt, generator=torch.Generator().manual_seed(k))
    for k in range(128)
])

print(f"E8M0 shared scale: 2^{sr_exp} = {sr_scale:.6f}   "
      f"(block amax {float(sr_block.abs().max()):.4f}) — the same X in both casts")
print(f"elements that move between the RTN cast and one SR cast: "
      f"{int((sr_sr_block != sr_rtn_block).sum())} of 32")
print(f"|error|: RTN {float((sr_rtn_block - sr_block).abs().mean()):.5f}   "
      f"SR, mean of 128 draws {float((sr_draws.mean(0) - sr_block).abs().mean()):.5f}")
E8M0 shared scale: 2^-4 = 0.062500   (block amax 0.3500) — the same X in both casts
elements that move between the RTN cast and one SR cast: 7 of 32
|error|: RTN 0.00897   SR, mean of 128 draws 0.00102

The shared scale is untouched: E8M0 stores only an exponent, and the exponent is chosen from the block’s peak — the same amax in both casts — so X is the same number and the only thing that changes is which of the two neighbouring elements each value lands on. In the mean over draws, SR reconstructs the true block; RTN reconstructs the same approximation every time, bias included.

The implementation does not actually flip coins. Fasi & Mikaitis (2021) show the two formulations are equivalent — including for a finite number of random bits — because “most of the implementations … add bits from a random stream to the part of the number that will be truncated”, and the carry out of that addition decides the direction. PTX spells the same recipe out for cvt.rs (PTX ISA 8.7+, sm_100a/sm_103a/sm_107a): the value is “rounded in the direction toward zero or away from zero based on the carry out of the integer addition of the supplied random bits (rbits) to the truncated off (discarded) bits of mantissa” — 13 random bits for f32→f16, 16 for f32→fp8. AMD’s MI300 exposes the same idea as CVT_SR_FP8_F32 (20 random bits) and CV_SR_BF8_F32 (21), and Graphcore’s IPU and SpiNNaker2 carry it too. This is recent hardware: pre-Blackwell tensor cores did not have it — their additions used round-toward-zero, and CUDA’s fp16 arithmetic implemented round-to-nearest only (Fasi, Higham, Mikaitis & Pranesh, 2021).

The recipe is equally explicit about which tensors get it. The NVFP4 pretraining run — 12B parameters, 10T tokens, “the longest publicly documented training run in 4-bit precision to date”, MMLU-Pro 62.58% against FP8’s 62.62% — lists SR as one of four co-equal ingredients, applied to the gradients entering both Dgrad and Wgrad; weights and activations stay on round-to-nearest-even. The reason is stated in §4.4: “applying stochastic rounding to the forward pass tensors is detrimental, as it amplifies quantization error relative to nearest rounding.” Appendix E.3’s ablation is the demonstration — “applying stochastic rounding to gradients leads to stable convergence of the training loss for the 1.2B model, whereas using it on activations or weights causes divergence” — and the ablation that removes SR entirely degrades convergence (Fig. 4).

WarningTwo Ways to Get SR Wrong
  1. Applying it to the wrong tensors. SR is not a global “make 4-bit training work” switch. The NVFP4 ablations put it on gradients only: on activations or weights it causes divergence, because on a forward-pass tensor the noise is amplified rather than averaged away — a single pass has no repetition to cancel it. Enable SR where a quantity is re-cast thousands of times; leave weights, activations, and every inference cast on round-to-nearest.
  2. Assuming the coin is exact. With only r random bits you get limited-precision SR: the probabilities deviate from p = (x - lo)/(hi - lo) and the expectation is no longer exactly x (El Arar, Fasi, Filip & Mikaitis, 2026, §5). Hardware SR is bounded this way — 13 random bits for f32→f16, 16 for f32→fp8 in PTX cvt.rs — so the unbiasedness above is the ideal the implementation approximates, not a guarantee of the silicon.

NVFP4 GEMM: Scale Each Partial Before Adding

Quantizing the two matrices is only half the job. A GEMM multiplies them and sums along the shared K dimension. Because NVFP4’s E4M3 scale changes from one 16-value block to the next, each block contributes its own scaled partial dot product. The fp32 tensor scales are shared by the whole operand; the block scales are local. Both are needed when a block’s contribution joins the output.

The Math: Block Scales Inside the Reduction

For X \in \mathbb{R}^{M\times K} and W \in \mathbb{R}^{K\times N}, let P^X_b and P^W_b be the E2M1 grid values for the b-th 16-wide K block. Each block also has a local E4M3 scale, s^X_{i,b} for row i of X and s^W_{b,j} for column j of W. The tensor scales are g_X and g_W.

Y_{ij} \approx \sum_b (g_X s^X_{i,b})(g_W s^W_{b,j}) \left(P^X_b {P^W_b}^{\mathsf T}\right)_{ij}.

The scale product is an outer product over output rows and columns. The element grid dot product gives one partial result for the block; multiply that partial by its scale product before adding the next block. A single scale product after the full dot product would only be valid if the block scales did not change along K.

Code: A Two-Block NVFP4 GEMM

nvfp4_gemm.py keeps the E2M1 grid values unpacked as floats so you can see the arithmetic. It uses the existing NVFP4 quantizer for each K block, then accumulates the scaled partials in float32. Compare its answer with a separate decode-then-matmul reference:

gemm_spec = importlib.util.spec_from_file_location("nvfp4_gemm", Path("nvfp4_gemm.py").resolve())
nv_gemm = importlib.util.module_from_spec(gemm_spec)
sys.modules["nvfp4_gemm"] = nv_gemm
gemm_spec.loader.exec_module(nv_gemm)

nv_gemm_demo = nv_gemm.demonstrate_nvfp4_gemm()
print(f"X @ W output shape: {nv_gemm_demo['shape_y']}")
print(f"max error vs decoded-operand matmul: {nv_gemm_demo['reference_max_error']:.2e}")
print(f"error if local scales are omitted: {nv_gemm_demo['global_only_max_error']:.3f}")
X @ W output shape: [2, 2]
max error vs decoded-operand matmul: 7.45e-09
error if local scales are omitted: 0.101
ojs_define(nvGemmTrace=nv_gemm_demo["trace"])

Interactive: Follow One Partial Product

Move the step control to see each 16-wide piece of the reduction. The first matrix is still in E2M1 grid units; the highlighted result applies both operands’ scales before it is accumulated.

NoteKey Insight

The tensor-wide scales do not eliminate the local scales. During the reduction, each K block has a different pair of local scales, so its own matrix of partial products must be rescaled before accumulation. This lesson implements that algebra in float32; it does not pack fp4 values or run a GPU Tensor Core kernel.

MXINT8: A Fixed-Point Grid With a Block Scale

MXFP4 spends four bits on a small floating-point element. MXINT8 spends all eight element bits on a signed fixed-point number: one sign bit, one integer bit, and six fractional bits. Its element code q represents q/64, so the symmetric values range from -127/64 to 127/64. One E8M0 power-of-two scale X is shared by each block of 32 values, and decoding is simply \hat v_i = Xq_i/64.

That makes the two levels easy to separate: the integer element gives a uniform grid with a fixed step inside the block, while E8M0 shifts that grid to fit the block’s magnitude. The reference conversion picks the largest power of two no greater than the block’s maximum absolute value, rounds scaled values to nearest-even, then saturates at \pm127/64. The block maximum can approach 2X, while the largest positive decoded element is (127/64)X, leaving a small saturation band at the top of a power-of-two interval. The standard permits alternate conversion algorithms; this example follows its minimally supported scale rule.

Code: Quantize and Decode a Block

mxint8.py reuses m14’s E8M0 scale encoder and records each code beside the value it reconstructs. This is an unpacked arithmetic model, not a hardware kernel or a packed tensor representation.

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

mxint8_demo = mxint8.demonstrate_mxint8()
for block in mxint8_demo["trace"]:
    print(f"block {block['block']}: scale={block['scale']:g}, first codes={block['codes'][:8]}")
print(f"Maximum absolute reconstruction error: {mxint8_demo['max_abs_error']:.4f}")
block 0: scale=0.5, first codes=[23, -52, 93, -8, 37, -67, 14, 49]
block 1: scale=2, first codes=[35, -77, 102, -26, 54, -99, 16, 70]
Maximum absolute reconstruction error: 0.0125
ojs_define(mxint8Trace=mxint8_demo["trace"])

Interactive: Watch the Scale Move the Grid

Step through the two blocks. The code column stays an integer; changing the block scale changes the spacing between the real values those codes decode to.

NoteKey Insight

The block scale is a second exponent shared across 32 values. Inside the block, MXINT8’s quantization step is exactly X/64; the integer grid does not change shape when the block scale changes. OCP specifies a symmetric endpoint of \pm127/64, leaving the two’s-complement -128 code unused.

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.

Put NVFP4 Weights in a Real GPT

So far, NVFP4 has been a tensor conversion. Here we connect that two-level representation to the nn.Linear modules inside the m06 GPTModel. A linear layer computes Y = XW^T + b. We flatten its weight in the same order as nvfp4.py, quantize each group of 16 values, and save three pieces:

  1. One tensor-wide fp32 scale, s_{global}.
  2. One E4M3 block scale, s_{block(i)}, for every 16 consecutive weights.
  3. One E2M1 grid value, P_i, per weight.

At inference, the wrapper reconstructs \hat W_i = s_{global}s_{block(i)}P_i and passes it to ordinary torch.nn.functional.linear. The bias stays full precision. This lets us observe the effect of NVFP4 weights inside a real model; it does not pack four-bit values or run an accelerated NVFP4 kernel. The forward path still materializes a floating-point weight matrix before calling F.linear.

The default m06 model ties lm_head.weight to the token embedding matrix. That matrix is an nn.Embedding, not an ordinary linear weight to replace, so the converter deliberately leaves the tied output head alone. Set tie_weights=False to give lm_head its own weight; then it can be converted like the other linear layers.

from pathlib import Path
import importlib.util
import sys

import torch

repo_root = Path.cwd().resolve().parents[1]
sys.path.insert(0, str(repo_root / "modules"))
from m06_transformer.transformer import GPTModel

linear_path = repo_root / "modules" / "m14_quantization" / "nvfp4_linear.py"
linear_spec = importlib.util.spec_from_file_location("nvfp4_linear_demo", linear_path)
nvfp4_linear = importlib.util.module_from_spec(linear_spec)
linear_spec.loader.exec_module(nvfp4_linear)

torch.manual_seed(14)
model = GPTModel(
    vocab_size=128, embed_dim=32, num_heads=4, num_layers=2,
    max_seq_len=32,
).eval()
tokens = torch.randint(0, model.vocab_size, (2, 8))
with torch.no_grad():
    full_precision_logits = model(tokens)

converted = nvfp4_linear.quantize_linear_layers(model)
with torch.no_grad():
    nvfp4_logits = model(tokens)

print(f"Converted {len(converted)} layers: {', '.join(converted)}")
print(f"Logits shape: {tuple(nvfp4_logits.shape)}")
print(f"Largest logit change: {(nvfp4_logits - full_precision_logits).abs().max().item():.5f}")
print(f"Tied output head preserved: {model.lm_head.weight is model.token_embedding.weight}")
Converted 12 layers: blocks.0.attention.mha.q_proj, blocks.0.attention.mha.k_proj, blocks.0.attention.mha.v_proj, blocks.0.attention.mha.out_proj, blocks.0.ffn.linear1, blocks.0.ffn.linear2, blocks.1.attention.mha.q_proj, blocks.1.attention.mha.k_proj, blocks.1.attention.mha.v_proj, blocks.1.attention.mha.out_proj, blocks.1.ffn.linear1, blocks.1.ffn.linear2
Logits shape: (2, 8, 128)
Largest logit change: 0.11666
Tied output head preserved: True

This randomly initialized model is a wiring check, not a quality benchmark. The logit difference only shows that quantization changes the computation while the model still runs end to end. The unpacked int8 half-step codes are a teaching representation of the E2M1 grid; they do not demonstrate a four-bit memory footprint.

Note

The production NVFP4 training layout can group weight scales over two dimensions (for example, shared scales across a 16 × 16 tile). This wrapper intentionally reuses the earlier lesson’s flattened, consecutive 16-value quantizer. The representation connects the lesson’s scale arithmetic to a real GPT module, but is not a complete Blackwell training or inference implementation.

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.

  13. Freezing group grids upfront. Fitting every group’s scale to the original weights ignores what GPTQ’s feedback already pushed downstream — later groups no longer hold the weights you calibrated. Recompute each grid from the compensated W at its boundary (grid="boundary", the default); the grid="upfront" contrast in the recipe section measures exactly what the shortcut costs.

  14. Confusing B = 128 with g = 128. The lazy-update block B is a throughput knob (same arithmetic, regrouped work); the group size g is an accuracy knob (a new grid per group). Passing block_size=128 does not buy anything grouping buys — and shrinking B will not fix a too-coarse grid.

  15. Quoting 4.125 when you mean 4.16. The symmetric recipe stores one fp16 scale per group (4 + 16/128 = 4.125); the asymmetric checkpoints the papers report also store a zero-point (4 + 20/128 \approx 4.16). State the assumption with the number — use recipe_bpw(..., asymmetric=True) for the paper comparison.

  16. Assuming the recipe clips. The AWQ path built here is scale-only: it reshapes salience and rounds. The full paper recipe adds a per-group weight clipping search on top (tuning each group’s [\min, \max] before rounding), which this module does not implement — expect a small residual gap to reported numbers, concentrated in the loudest groups.

  17. Snapping E8M0 to nearest. The spec default is floor — the largest power of two at or below the ideal (OCP MX v1.0 §6.3) — not nearest, and the two differ on three of every four binades. Rounding up is a real implementation choice (Mishra et al.), but if you snap to nearest while the hardware floors, your simulated error underestimates the shipped error by up to half a scale. Use snap_scale_e8m0 and state the direction.

  18. Quoting unconstrained bits for constrained storage. The g128 recipe’s 4.125 counts an fp16 scale per group; once the scale snaps to 8 bits the honest number is 4 + 8/g (4.25 at g32) — and the error behind the number changes too, because flooring clips the group peak. Quote mxfp4_g32_effective_bits with the snapped loss beside it, never the free loss with the constrained bits.

  19. Treating stochastic rounding as a free win. It is not a better cast — it is a noisier one with a smaller bias. The worst-case per-cast error doubles (|\delta| \le 2u under SR against RN’s u, Croci et al. eq 6.4), and the variance it buys back is at best \Delta^2/4. It pays where the same quantity is re-quantized thousands of times — a parameter accumulating updates, a gradient summed over a batch — and costs everywhere else. Inference wants a deterministic answer, which is exactly why the NVFP4 recipe leaves weights and activations on round-to-nearest-even and enables SR on gradients only.

  20. Applying one scale after the whole GEMM. The E4M3 scale changes from block to block along K. Multiply every block’s partial dot product by the two effective scales for that block, then sum. nvfp4_gemm.py demonstrates the scale arithmetic; it is an educational emulation, not a fast kernel.

  21. Using all 256 signed codes for MXINT8. The format’s symmetric range is [-127,127], so -128 is intentionally unused. Also keep the implicit 1/64 element factor in the decode; the E8M0 block scale does not replace it.

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 41.1 dB,  int4 15.9 dB,  gap 25.2 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.2 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

Exercise 13: Price the recipe both ways, then predict the knee

Hand-compute the stored cost of int4-g128 under both accountings — symmetric (4 + 16/128) and asymmetric (4 + (16 + 4)/128) — and confirm recipe_bpw agrees. Then predict: at which group size does halving g stop buying SNR worth the bits? Check your prediction against the sweep chart’s knee band.

for g in (256, 128, 64, 32):
    print(f"g={g:>3}:  sym {gr.recipe_bpw(4, g):.5f}   asym {gr.recipe_bpw(4, g, asymmetric=True):.5f}")
# Your turn: read the SNR column off recipe_demo["sweep"] and mark where the
# dB-per-extra-bit trade turns bad — that is the knee the widget shades.
g=256:  sym 4.06250   asym 4.07812
g=128:  sym 4.12500   asym 4.15625
g= 64:  sym 4.25000   asym 4.31250
g= 32:  sym 4.50000   asym 4.62500

Exercise 14: Snap the scale by hand, then price the clip

Take one group scale (say s = 0.34), snap it both ways by hand — E8M0 floor (2^floor(log2(0.34))) and E4M3 nearest — and confirm snap_scale_e8m0 and snap_scale_e4m3 agree. Then run the GPTQ-threaded recipe free vs snapped on a small layer and confirm the staged ordering holds at the demo seed.

s = 0.34
print(f"free {s:.4f}   E8M0 {mxr.snap_scale_e8m0(s):.4f}   E4M3 {mxr.snap_scale_e4m3(s):.4f}")

torch.manual_seed(3)
Wm = torch.randn(8, 32) * 0.1
Wm[2, 5] *= 20.0
mix = torch.randn(32, 32)
Xm = mix @ torch.randn(32, 64)
Hm = gptq.layer_hessian(Xm)
cmp = mxr.constrained_vs_free(Wm, Hm, Xm, group_size=8, mode="gptq")
print(f"free {cmp['free']:.1f}   E8M0 {cmp['e8m0']:.1f}   E4M3 {cmp['e4m3']:.1f}")
# Your turn: flip scale="e4m3" to "e8m0" in a mode="awq" run and compare the gaps.
free 0.3400   E8M0 0.2500   E4M3 0.3438
free 28.1   E8M0 539.0   E4M3 29.9

Exercise 15: Round up with probability p

SR is unbiased by construction, not by luck. Take a value between two grid neighbours and verify the identity by hand: with p = (x - lo)/(hi - lo), the expected outcome p\cdot hi + (1-p)\cdot lo equals x exactly — no sampling involved, no “for large enough draws”. Then the RTN contrast: run the absorbed update under both casts and confirm RTN’s lane never leaves zero while SR’s mean lands on the exact sum.

lo, hi = 0.0, 0.5
x = 0.35
p = sround.sr_round_up_probability(x, lo, hi)
print(f"p = {p:.3f}   p·hi + (1-p)·lo = {p * hi + (1 - p) * lo:.6f}   x = {x:.6f}")

# Your turn: run accumulate(delta=0.1, steps=200, step=0.5, mode="rtn") and then
# mode="sr", and compare each lane's final value to the exact sum (delta * steps).
# RTN should read exactly 0.0 at every step — the update is absorbed, not merely
# rounded — while SR lands within about step·sqrt(n)/2 of the exact sum.
p = 0.700   p·hi + (1-p)·lo = 0.350000   x = 0.350000

Exercise 16: Restore the K-block scales

The GEMM emulation keeps E2M1 grid values and E4M3 scales separate. Rebuild the output from the trace in demonstrate_nvfp4_gemm: add each scaled_partial and compare the result with the demo’s output. Then try the tempting shortcut: sum the raw partials first and apply only the two tensor scales. Why does that shortcut fail when the E4M3 block scales differ?

demo = nv_gemm.demonstrate_nvfp4_gemm()
restored = torch.tensor(demo["trace"][0]["scaled_partial"])
for block in demo["trace"][1:]:
    restored += torch.tensor(block["scaled_partial"])
print(f"max reconstruction gap: {(restored - torch.tensor(demo['result'])).abs().max():.2e}")
max reconstruction gap: 0.00e+00

Exercise 17: Decode one MXINT8 block by hand

Use the block maximum to choose its E8M0 scale, then encode three values with the fixed q/64 element grid. Confirm that the returned codes use symmetric clamping and that decoding uses the product of the block scale and the implicit 2^{-6} element scale.

demo = mxint8.demonstrate_mxint8()
block = demo["trace"][1]
print(f"max abs input: {max(abs(v) for v in block['source']):.3f}")
print(f"E8M0 scale: {block['scale']}")
print(f"first inputs: {block['source'][:3]}")
print(f"first codes:  {block['codes'][:3]}")
print(f"first decoded:{block['decoded'][:3]}")
# Your turn: hand-compute q = round_even(v / scale * 64), then decode scale*q/64.
max abs input: 3.200
E8M0 scale: 2.0
first inputs: [1.100000023841858, -2.4000000953674316, 3.200000047683716]
first codes:  [35, -77, 102]
first decoded:[1.09375, -2.40625, 3.1875]

Exercise 18: Wire NVFP4 into GPTModel

Run the model example and inspect the paths it converts. Why does the tied lm_head remain an nn.Linear? Construct a second model with tie_weights=False, run quantize_linear_layers, and confirm that its output projection is converted. Then explain why a successful forward pass does not show a four-bit memory footprint or a speedup.

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.

  14. The production recipe threads group scales through GPTQ and AWQ. AWQ’s s = s_X^{\alpha} search runs outside the grouped quantizer (reshape salience first), GPTQ’s Cholesky feedback runs across the per-group grids with each scale recomputed at its group boundary (compensate after) — and the stack beats each alone, at 4.125 symmetric / \approx 4.16 asymmetric bits.

  15. The lazy block B is throughput, the group g is accuracy. Changing B regroups the same GPTQ arithmetic; changing g regrids every weight, with the SNR-vs-bits knee at 64–128. Quote B = 128 and g = 128 as the coincidence they are — never as one knob.

  16. Hardware runs constrained scales, not free ones. Snapping each group scale to E8M0 (floor to a power of two — the spec default) or E4M3 (nearest fp8, landing between the powers) threads through the same GPTQ feedback and AWQ search, costs error (floor also clips the group peak), and shrinks the scale to 8 bits: MXFP4-g32 at 4.25 bits, the number gpt-oss ships its MoE weights at.

  17. Finer scales recover what flooring loses. E4M3’s fractional grid hugs the free error at every group size, and shrinking the group localizes each clipped peak — while the two-level NVFP4 guarantee (block_scale_headroom at 448) carries over unchanged. The staged ordering (free below both constrained) is a pinned same-seed observation, not a theorem, because snapping rewrites the feedback path.

  18. Stochastic rounding is the cast that removes the bias. RTN always sends a value the same way, so its error is a fixed-sign offset that compounds over a training run; SR picks the upper neighbour with probability p = (x-lo)/(hi-lo), which makes the rounding error mean-zero — E[Q_{SR}(x)] = x is an identity, not a sampling result — with variance at most \Delta^2/4 and accumulation error growing like \sqrt{n} instead of n. The price is a worse single cast (|\delta| \le 2u against RN’s u — a factor 2), so SR pays only where the same quantity is re-quantized thousands of times: it is a training cast, applied to gradients, while inference keeps RTN.

  19. Block-scaled GEMM rescales before it adds. The two operand scales for each 16-wide K tile form an outer product over the output matrix. Scale that tile’s partial dot product, then accumulate; one final scale pair cannot replace the block-local products. nvfp4_gemm.py reproduces this arithmetic in float32, not the hardware instruction.

  20. MXINT8 pairs fixed point with a block scale. Each signed code represents q/64, and one E8M0 power of two per 32 elements sets the real step to X/64. Its symmetric code range is [-127,127].

  21. NVFP4 weights can run through a real GPTModel. NVFP4Linear stores the unpacked E2M1 grid values with their E4M3 block scales and fp32 tensor scale, then reconstructs a floating-point matrix for F.linear. The converter keeps a tied embedding/output head shared; the forward pass demonstrates wiring and numerical change, not packed memory savings or acceleration.

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 — done above as the production recipe (grouped_recipe.py: grouped_gptq_quantize with boundary-recomputed grids, grouped_awq_quantize with the \alpha-search outside, and stacked_awq_gptq_grouped combining them) — extended above to E8M0/E4M3-constrained scales through the same methods (mx_recipe.py: constrained_grouped_quantize with floor/nearest-snapped boundary scales at MXFP4-g32’s 4.25 bits, the layout gpt-oss ships) — and with stochastic rounding the cast itself becomes a choice rather than a constant (stochastic_round.py: the unbiased cast, its \Delta^2/4 variance bound, and the absorbed updates it rescues) — and an NVFP4 GEMM emulation (nvfp4_gemm.py: quantize both operands, rescale each K-block partial, then accumulate) — plus AWQ’s tuned scale search and per-group clipping, 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.

The MX family also includes MXINT8: six fractional bits give each element a uniform q/64 grid, and the shared E8M0 scale changes that grid’s real step once per 32-value block. mxint8.py builds the OCP conversion and decode rule; the scale and rounding path here is an educational reference, not a packed hardware implementation.

The final NVFP4 bridge (nvfp4_linear.py) puts the two-level weight conversion inside real m06 GPTModel linear layers. Its ordinary F.linear forward path reconstructs the floating-point matrix, and the converter leaves the default tied embedding/output head shared. This closes the m14 MX/NVFP4 implementation follow-ups while keeping the emulation’s scope clear.

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. Stochastic rounding enters as a gradient-only ingredient (§4.4): on forward-pass tensors it amplifies quantization error, and App. E.3’s ablation diverges when it is applied to activations or weights. https://arxiv.org/abs/2509.25149
  • NVIDIA Transformer Engine, FP8/FP4 Primer — contrasts 1D blocks of 16 values with a training recipe that shares weight scales across 16 × 16 tiles; the NVFP4Linear wrapper above follows the former flattened educational layout. https://docs.nvidia.com/deeplearning/transformer-engine/examples/fp8_primer.html
  • NVIDIA, Introducing NVFP4 for Efficient and Accurate Low-Precision Inference
    1. — the E2M1 values, 16-element E4M3-scaled microblocks, and fp32 tensor scale represented by the GEMM emulation. This lesson focuses on scale arithmetic and does not model a Tensor Core kernel. https://developer.nvidia.com/blog/introducing-nvfp4-for-efficient-and-accurate-low-precision-inference/
  • Croci, Fasi, Higham, Mary & Mikaitis, Stochastic rounding: implementation, error analysis and applications (2022) — unbiasedness, the √n bound, the stagnation theorem. https://doi.org/10.1098/rsos.211631
  • El Arar, Sohier, de Oliveira Castro & Petit, Stochastic rounding variance and probabilistic bounds: a new approach (2023) — the Δ²/4 variance bound. https://arxiv.org/abs/2207.10321
  • Gupta et al., Deep Learning with Limited Numerical Precision (2015) — where SR enters deep learning. https://arxiv.org/abs/1502.02551
  • Zamirai et al., Revisiting BFloat16 Training (2020) — RTN cancels small weight updates. https://arxiv.org/abs/2010.06192
  • 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
  • OpenAI, gpt-oss-120b & gpt-oss-20b Model Card (2025) — the existence proof for this section: post-trained MoE weights quantized to MXFP4 at 4.25 bits per parameter, fitting the 120B model on one 80 GB GPU. https://arxiv.org/abs/2508.10925

Practical Resources: