Quantization · AI infra interview lesson

Why It Matters

Representing 0.2 and 1.7 as integers requires choosing the value of one integer step: the scale. Compute reconstruction error before comparing granularities, integer formats and low-bit floats.

A concrete ledger: LLaMA-2-7B's 6,738,415,616 parameters take 13.48 GB in FP16. As a rough estimate, store all of them as INT4 with one FP16 scale per 128 weights: 4/8 + 2/128 = 0.515625 bytes per parameter, 3.47 GB in total, about 3.9x smaller. At an H100's 3.35 TB/s, the batch-1 weight-read lower bound per step drops from 4.0 ms to 1.04 ms. On the KV side, a 16-bit cache with 32 layers, 8 KV heads, head_dim=128 and 32768 tokens is 2*32*32768*8*128*2 = 4 GiB; FP8 is 2 GiB, and the 2-bit codes alone are 0.5 GiB (plus scales and the residual).

Walk the Shapes First

Two layouts show up: nn.Linear.weight is stored as (C_out, C_in), while the math is written Y = X W with W as (C_in, C_out). Below, C_in=10, C_out=4, T=3, group_size G=4.

Rules read off the table: the shape of the scale is just another way of saying "which elements share parameters", and with keepdim it broadcasts straight back; SmoothQuant's s runs along C_in while INT8 scales run along T and C_out, which are different things; groups that do not divide evenly and tokens that do not fill a group need separate handling.

Running the examples: execute this page in order in one Python environment. Later snippets may reuse earlier definitions. CPU examples validate calculations and state changes, not GPU timing. Download the notebook for the complete exercise workspace.

Core Implementations

Step 1: Symmetric / Asymmetric, and Rounding

What is new: a per-tensor quantize / dequantize.

• Symmetric: s = amax / (2^(b-1) - 1), q = clamp(round(x / s), -qmax, qmax), x_hat = s * q. INT8 uses [-127, 127], deliberately skipping -128 to stay symmetric.
• Asymmetric: first widen [lo, hi] to contain 0, then s = (hi - lo) / (2^b - 1), z = round(-lo / s), q = clamp(round(x / s) + z, 0, 2^b - 1), x_hat = s * (q - z).
• torch.round rounds half to even: 0.5 → 0, 1.5 → 2, 2.5 → 2, matching torch.fake_quantize_per_tensor_affine. Python's built-in round is also half-even; math.floor(x + 0.5) is not.
• All-zero tensor: amax = 0, and the scale must not be 0 or x / s becomes NaN.
• In-range error bound: |x - x_hat| <= s / 2.

Example: [-1, 3] with 4 bits gives s = 4/15, z = round(3.75) = 4; x = 1 encodes to 8 and reads back as 1.0667. Rounding the zero point shifts the ends too: -1 reads back as -1.0667.

import torch
torch.manual_seed(0)

def quantize_sym(x, bits=8):
    """对称、per-tensor：q ∈ [-(2^(b-1)-1), 2^(b-1)-1]，z = 0。返回 (q int32, scale)。"""
    qmax = 2 ** (bits - 1) - 1
    amax = x.abs().max()
    scale = amax / qmax if amax > 0 else torch.tensor(1.0)   # 全零张量：scale 随便给个正数
    q = torch.round(x / scale).clamp(-qmax, qmax).to(torch.int32)   # torch.round 是 half-to-even
    return q, scale

def quantize_asym(x, bits=8):
    """非对称：q ∈ [0, 2^b - 1]，区间 [min(x,0), max(x,0)] 必须包含 0，零点取整。"""
    qmax = 2 ** bits - 1
    lo, hi = x.min().clamp(max=0), x.max().clamp(min=0)
    scale = (hi - lo) / qmax if hi > lo else torch.tensor(1.0)
    zero = torch.round(-lo / scale).clamp(0, qmax)
    q = (torch.round(x / scale) + zero).clamp(0, qmax).to(torch.int32)
    return q, scale, zero

def dequantize(q, scale, zero=0):
    return (q.float() - zero) * scale

x = torch.randn(64) * 3
q, s = quantize_sym(x)
ref = torch.fake_quantize_per_tensor_affine(x, s.item(), 0, -127, 127)
assert torch.allclose(dequantize(q, s), ref)
assert (dequantize(q, s) - x).abs().max() <= s / 2 + 1e-6          # 范围内误差 ≤ scale/2

qa, sa, za = quantize_asym(x.relu(), bits=4)                         # 全非负：零点 0
ref = torch.fake_quantize_per_tensor_affine(x.relu(), sa.item(), int(za), 0, 15)
assert torch.allclose(dequantize(qa, sa, za), ref)
qb, sb, zb = quantize_asym(x + 2, bits=8)                            # 偏正的分布：零点不为 0
ref = torch.fake_quantize_per_tensor_affine(x + 2, sb.item(), int(zb), 0, 255)
assert torch.allclose(dequantize(qb, sb, zb), ref, atol=1e-6) and int(zb) > 0

# 手算例子：[-1, 3] 用 4-bit → s = 4/15, z = 4；x = 1 → q = 8 → 1.0667
q1, s1, z1 = quantize_asym(torch.tensor([-1.0, 1.0, 3.0]), bits=4)
assert q1.tolist() == [0, 8, 15] and int(z1) == 4
print("asym:", q1.tolist(), round(s1.item(), 4), dequantize(q1, s1, z1).tolist())

# half-to-even：0.5 → 0, 1.5 → 2, 2.5 → 2
q2, _ = quantize_sym(torch.tensor([0.5, 1.5, 2.5, -2.5, 127.0]))
assert q2.tolist() == [0, 2, 2, -2, 127]
# 全零张量
q3, s3 = quantize_sym(torch.zeros(4))
assert q3.tolist() == [0, 0, 0, 0] and not torch.isnan(dequantize(q3, s3)).any()

Step 2: Granularity: Per-tensor / Per-channel / Per-group

What is new: the scale is no longer a single number but an absmax along some dimension (keepdim).

• per-tensor: amax = x.abs().amax()
• per-channel (weights, one per output channel): in Linear layout amax(dim=1, keepdim=True), shape (C_out, 1)
• per-group: split C_in into chunks of G, view(C_out, n_groups, G), and take amax over the last dim. If C_in does not divide, F.pad with zeros first; padded zeros never change the absmax, and you slice them off at the end.

Finer granularity means one outlier contaminates fewer elements, but you store more scales: INT4 plus one FP16 scale per 128 is 0.515625 bytes per weight.

import torch
import torch.nn.functional as F
torch.manual_seed(0)

def qdq_sym(x, bits=8, dim=None):
    """对称 quantize→dequantize。dim=None: per-tensor；否则沿 dim 求 absmax（keepdim），即每个其余下标一组。"""
    qmax = 2 ** (bits - 1) - 1
    amax = x.abs().amax() if dim is None else x.abs().amax(dim=dim, keepdim=True)
    scale = torch.where(amax > 0, amax / qmax, torch.ones_like(amax))
    return torch.round(x / scale).clamp(-qmax, qmax) * scale, scale

def qdq_per_group(w, bits=4, group_size=128):
    """w: (C_out, C_in) PyTorch Linear 布局；沿 C_in 每 group_size 个一组。C_in 不整除时末尾补 0。"""
    C_out, C_in = w.shape
    pad = (-C_in) % group_size
    wp = F.pad(w, (0, pad))                                    # (C_out, C_in + pad)
    g = wp.view(C_out, -1, group_size)                         # (C_out, n_groups, G)
    deq, scale = qdq_sym(g, bits, dim=-1)                      # scale: (C_out, n_groups, 1)
    return deq.view(C_out, -1)[:, :C_in], scale.squeeze(-1)

W = torch.randn(4, 8)
W[:, 5] *= 20                                                  # 一个输入通道是离群值
_, s_t = qdq_sym(W)                                            # ()
deq_c, s_c = qdq_sym(W, dim=1)                                 # (4, 1)：每个输出通道一个
deq_g, s_g = qdq_per_group(W, bits=8, group_size=4)            # (4, 2)
print("scale shapes:", tuple(s_t.shape), tuple(s_c.shape), tuple(s_g.shape))

ref = torch.fake_quantize_per_channel_affine(W, s_c.flatten(), torch.zeros(4, dtype=torch.int32), 0, -127, 127)
assert torch.allclose(deq_c, ref)

err = {name: (W - d).pow(2).mean().item() for name, d in
       [("tensor", qdq_sym(W)[0]), ("channel", deq_c), ("group", deq_g)]}
print("MSE:", {k: f"{v:.2e}" for k, v in err.items()})
assert err["group"] <= err["channel"] <= err["tensor"]

# 不整除：C_in = 10，G = 4 → 3 组，最后一组只有 2 个真实元素
W2 = torch.randn(3, 10)
deq2, s2 = qdq_per_group(W2, bits=4, group_size=4)
assert deq2.shape == W2.shape and s2.shape == (3, 3)
for j in range(3):                                             # 朴素参考：逐组切片
    blk = W2[:, 4 * j: 4 * j + 4]
    ref_blk, _ = qdq_sym(blk, bits=4, dim=1)
    assert torch.allclose(deq2[:, 4 * j: 4 * j + 4], ref_blk)

# 存储账：INT4 + 每 128 个共享一个 FP16 scale
print("bytes/weight, G=128:", 4 / 8 + 2 / 128)                 # 0.515625

Step 3: W8A8 and SmoothQuant

What is new: activations are quantized too and the GEMM runs on INT8 inputs with INT32 accumulation; (Pro lesson)

Step 4: INT4 Pack/Unpack and Simplified GPTQ

What is new: 4-bit codes go two per byte; weight quantization uses calibration activations to compensate errors instead of rounding each wei… (Pro lesson)

Compensate error after quantizing a column

Direct rounding considers a weight locally. This simplified GPTQ model uses input statistics to estimate output error and compensate remaini… (Pro lesson)

Step 5: KV Cache and Low-bit Floats: KIVI, FP8 E4M3, FP4 E2M1

What is new: KV is written while generating, so quantization must be streaming; FP8 / FP4 grids are not uniform. (Pro lesson)

Low-bit floating-point rounding and overflow

This simulates saturating E4M3: clamp to [-448,448], then round. The reference must use the same clamping policy. (Pro lesson)

Low-bit values also need a shared scale

FP4 has few representable values, so groups share a scale. Compare how grouping changes maxima and reconstruction errors. (Pro lesson)

In Real Systems

nano-vllm : no quantization. ModelRunner.__init__ calls torch.set_default_dtype(hf_config.dtype) , so the whole model is BF16/FP16; (Pro lesson)

Common Pitfalls

Pattern: scale = x.abs().max() / 127 with no zero check. Problem: an all-zero tensor gives scale=0 and x / scale is NaN. (Pro lesson)

Say It in the Interview

"Quantization trades precision for bytes: in memory-bound decode, INT4 weights with group-128 FP16 scales are about 0.52 bytes per parameter… (Pro lesson)

Hands-on Notebook

Each exercise in the notebook gives a statement and a function skeleton; write your code, then run the tests below it. (Pro lesson)