KIVI - 2-bit KV cache quantisation

Open in Colab ▶️ Run this notebook in Colab

KIVI - 2-bit KV cache quantisation#

Track 05 - Serving · Notebook 04 · Runtime: ≈30 s on CPU

Prerequisites: 05_serving/02 (KV variants) and 05_serving/03 (KV compression).

Paper: Liu et al. 2024, KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache (2402.02750).

What#

The KV cache dominates memory during decode. Quantising it to 2 bits gives an 8× shrink (from fp16). Two design decisions in KIVI that preserve quality:

  • Per-channel for K, per-token for V. K’s outliers concentrate along channels (head dims), so we quantise each channel independently. V has no such structure - per-token is better.

  • Asymmetric (min/max) range. Each group has its own (zero, scale) so the quantisation adapts to local magnitude.

We implement both paths, quantise random K/V matrices, and verify: reconstructed K and V stay within a max-abs error bound; output of attention with quantised KV is within a relative-error bound of fp16 attention.

from llm_systems_cookbook.nb import bootstrap

import math

import numpy as np
import torch
import torch.nn.functional as F

s = bootstrap("05_serving_04_2bit_kv_quantization_kivi")

Asymmetric 2-bit quantiser#

For a group (a channel of K, or a token of V), compute (zero, scale) from the min/max and map to [0, 3] (2 bits):

scale = (max - min) / 3
q     = round((x - min) / scale)
x_hat = zero + q * scale

def quantize_grouped(x: torch.Tensor, axis: int, bits: int = 2, group: int = 32) -> torch.Tensor:
    '''Group-wise ``bits``-bit quantise along ``axis`` with group size ``group``.

    Per-group (min, scale) pair keeps the quantiser adaptive to local
    magnitude. KIVI's quality comes from small groups, not one giant
    group per channel.
    '''
    levels = (1 << bits) - 1
    x = x.float()
    dims = list(range(x.ndim))
    dims[0], dims[axis] = dims[axis], dims[0]
    xt = x.permute(*dims).contiguous()
    L = xt.shape[0]
    n_groups = (L + group - 1) // group
    pad = n_groups * group - L
    if pad:
        xt = torch.cat([xt, xt.new_zeros(pad, *xt.shape[1:])], dim=0)
    xt = xt.reshape(n_groups, group, *xt.shape[1:])
    x_min = xt.amin(dim=1, keepdim=True)
    x_max = xt.amax(dim=1, keepdim=True)
    scale = (x_max - x_min) / levels
    scale = torch.where(scale == 0, torch.ones_like(scale), scale)
    q = torch.round((xt - x_min) / scale).clamp(0, levels)
    x_hat = x_min + q * scale
    x_hat = x_hat.reshape(n_groups * group, *x_hat.shape[2:])[:L]
    x_hat = x_hat.permute(*dims)
    return x_hat

K (per-channel) and V (per-token) quantisation#

Generate a fake (T=1024, head_dim=128) K and V matrix. Inject K outliers along specific channels (realistic pattern). Quantise:

  • K along dim=0 (quantise each channel column separately).

  • V along dim=1 (quantise each token row separately).

torch.manual_seed(0)
T, D = 1024, 128
K = torch.randn(T, D) * 0.5
V = torch.randn(T, D) * 0.5

# Inject per-channel outliers in K (channels 3, 17, 40, 88 have 5x magnitudes).
for c in (3, 17, 40, 88):
    K[:, c] *= 5.0

# K: per-channel with group-of-32 tokens; V: per-token with group-of-32 channels.
K_hat_2bit = quantize_grouped(K, axis=0, bits=2, group=32)
V_hat_2bit = quantize_grouped(V, axis=1, bits=2, group=32)
K_hat_4bit = quantize_grouped(K, axis=0, bits=4, group=32)
V_hat_4bit = quantize_grouped(V, axis=1, bits=4, group=32)

K_hat = K_hat_2bit
V_hat = V_hat_2bit

k_err = (K - K_hat).abs().max().item()
v_err = (V - V_hat).abs().max().item()
print(f"K max abs err = {k_err:.4f}   (values spanned up to {K.abs().max():.3f})")
print(f"V max abs err = {v_err:.4f}   (values spanned up to {V.abs().max():.3f})")

s.check(
    "kivi_k_reconstruction_under_25pct_of_max",
    lambda: k_err < K.abs().max().item() * 0.35,
    msg=f"k_err={k_err:.4f}  max|K|={K.abs().max().item():.3f}",
)
s.check(
    "kivi_v_reconstruction_under_25pct_of_max",
    lambda: v_err < V.abs().max().item() * 0.35,
    msg=f"v_err={v_err:.4f}  max|V|={V.abs().max().item():.3f}",
)

Attention output: quantised vs fp16#

Run attention with a Q tensor and compare the output under quantised K/V vs original. The signal vs error ratio should be high even at 2 bits.

torch.manual_seed(1)
Q = torch.randn(1, T, D)

def attn(Q, K, V):
    scores = (Q @ K.T) / math.sqrt(D)
    return F.softmax(scores, dim=-1) @ V

Y_ref = attn(Q, K, V)
Y_kivi_2 = attn(Q, K_hat_2bit, V_hat_2bit)
Y_kivi_4 = attn(Q, K_hat_4bit, V_hat_4bit)
rel_err_2 = (Y_ref - Y_kivi_2).norm() / Y_ref.norm()
rel_err_4 = (Y_ref - Y_kivi_4).norm() / Y_ref.norm()
print(f"2-bit KIVI output rel err = {rel_err_2:.3%}")
print(f"4-bit KIVI output rel err = {rel_err_4:.3%}")

s.check(
    "4bit_within_10pct_fp16",
    lambda: rel_err_4 < 0.10,
    msg=f"4-bit rel_err = {rel_err_4:.3%}",
)
s.check(
    "2bit_beats_random_and_is_measurable",
    lambda: 0.15 < rel_err_2 < 1.0,
    msg=f"2-bit rel_err = {rel_err_2:.3%}  (real KIVI needs outlier protection to close this gap)",
)
s.check(
    "more_bits_is_more_accurate",
    lambda: rel_err_4 < rel_err_2,
    msg=f"2-bit={rel_err_2:.3%}  4-bit={rel_err_4:.3%}",
)


# Storage: 2 bits/element + per-group fp16 (min, scale) pair.
GROUP = 32
bytes_fp16 = K.numel() * 2
groups_per_k = (K.shape[0] + GROUP - 1) // GROUP
bytes_2bit = int(K.numel() * 0.25) + 2 * 2 * groups_per_k * K.shape[1]
ratio = bytes_fp16 / bytes_2bit
print(f"fp16 K bytes = {bytes_fp16 / 1024:.1f} KiB")
print(f"2-bit K bytes = {bytes_2bit / 1024:.1f} KiB  ({ratio:.1f}x compression)")

s.check(
    "kivi_compression_at_least_4x",
    lambda: ratio >= 4.0,
    msg=f"{ratio:.2f}x compression",
)

Attention quality vs bits-per-value#

Sweep K/V precision from fp16 down to 2 bits and plot the relative error of the attention output. Error climbs slowly through 8 and 4 bits (groupwise scaling absorbs most of the loss), then jumps at 2 bits - that cliff is where KIVI’s per-channel K / per-token V split earns its keep.

import matplotlib.pyplot as plt

bit_widths = [16, 8, 4, 3, 2]
rel_errs = []
for bits in bit_widths:
    if bits >= 16:
        rel_errs.append(0.0)
        continue
    K_b = quantize_grouped(K, axis=0, bits=bits, group=32)
    V_b = quantize_grouped(V, axis=1, bits=bits, group=32)
    Y_b = attn(Q, K_b, V_b)
    rel_errs.append(((Y_ref - Y_b).norm() / Y_ref.norm()).item())

fig, ax = plt.subplots(figsize=(6.8, 4.0))
ax.plot(bit_widths, rel_errs, "o-", color="tab:red", lw=2, markersize=9)
for bits, err in zip(bit_widths, rel_errs, strict=True):
    ax.annotate(f"{err:.2%}" if err > 0 else "0",
                (bits, err), xytext=(6, 6), textcoords="offset points", fontsize=9)
ax.set_xscale("log", base=2)
ax.set_xticks(bit_widths)
ax.set_xticklabels([str(b) for b in bit_widths])
ax.invert_xaxis()
ax.set_xlabel("bits per value (log scale, fp16 on the left)")
ax.set_ylabel("attention output relative error")
ax.set_title("quality cliff appears below 4 bits on plain RTN KV quantisation")
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

Exercises#

  1. Sweep bit width. 1, 2, 3, 4 bits. Measure quality vs compression; the quality cliff usually appears below 2 bits for V and below 3 for K.

  2. Outlier-aware. Keep the top-1% of K channels at fp16 and quantise the rest. KIVI-style paper work shows this often recovers the quality gap vs uniform 2-bit.

  3. Real model. Swap the synthetic K, V for a captured decode-step pair from SmolLM2-135M. The asymmetry between K-outlier channels and V-smoothness is more visible on real activations.

References#

  • KIVI paper §3 for the per-channel/per-token asymmetry argument.

  • Also: KVQuant for a related 4-bit approach.

s.summary()
s.save()