KV Cache Quantization: Why Keys and Values Need Different Axes

Quantize your KV cache to 4 bits with a single symmetric scale per token and watch perplexity fall off a cliff — while the exact same scheme applied to the value cache costs you almost nothing. Same bit width, same model, same kernel. The keys break and the values don't.

That asymmetry is not a bug in your implementation. It is a structural property of what keys and values are, and it dictates the design of every serious KV cache quantization scheme shipping today.

TL;DR

KV cache quantization compresses the per-token key/value tensors that dominate memory at long context — for Llama-3.1-70B (80 layers, 8 KV heads, head_dim 128) the cache costs ~320 KiB per token in FP16, so 128K tokens is ~40 GiB, more than the weights of the model at 4-bit.
Keys have persistent outlier channels: a handful of head_dim positions carry magnitudes 10–100× the rest, and they're the same positions across every token. Values have no such structure.
Therefore quantize keys per-channel and values per-token. Per-token key quantization forces one scale to cover both the outlier channels and the normal ones, and the normal channels collapse into a couple of quantization levels.
RoPE smears the outliers. Rotation mixes channel pairs, so post-RoPE keys have less clean per-channel structure. Quantizing keys before RoPE preserves it, at the cost of applying rotation on dequantized values at attention time.
Keep a full-precision residual window of the most recent 32–128 tokens. It costs almost no memory and removes the worst of the error on the tokens attention weights most.

Why is the KV cache the thing that actually runs out of memory?

Because it grows linearly with sequence length and batch size, while weights are fixed. The size per token is:

bytes_per_token = 2 (K and V) × n_layers × n_kv_heads × head_dim × dtype_bytes

For Llama-3.1-70B with GQA (8 KV heads, head_dim 128, 80 layers) in FP16: 2 × 80 × 8 × 128 × 2 = 327,680 bytes ≈ 320 KiB per token. A single 128K-token request needs about 40 GiB of KV cache. Serve eight of them concurrently and you have exceeded an 8×H100 node on cache alone.

Halving that with FP8, or quartering it with INT4, is the difference between 8 concurrent long-context sessions and 32. That directly multiplies throughput, because on decode you are memory-bandwidth bound: fewer bytes read per attention step is fewer microseconds per token.

Why does per-token key quantization destroy accuracy?

Because key tensors have outlier channels — specific coordinates of the head dimension whose activations are consistently one to two orders of magnitude larger than their neighbors, for essentially every token in the sequence. This is the same phenomenon that motivated activation-aware weight quantization, but here it lives in the cache itself.

Group-wise quantization picks a scale per group. Where you draw the group boundary is everything:

Per-token (group = one token's whole key vector): the group spans both outlier and normal channels. The scale is set by the max, so scale ≈ outlier_max / 7 for INT4. Normal channels, two orders of magnitude smaller, all round to −1, 0, or 1. You have effectively quantized most of the key vector to under 2 bits.
Per-channel (group = one channel across tokens): each outlier channel gets its own large scale, each normal channel gets its own small scale. Nobody is dragged along.

Values have roughly isotropic magnitude across channels, so per-token grouping is fine there — and per-token is what you want for values, because it's the layout-friendly axis: a new token appends a complete, self-contained quantized block. Per-channel value quantization would force you to restate scales as the sequence grows.

Here's the effect, isolated:

import torch

torch.manual_seed(0)
T, D = 512, 128                       # tokens, head_dim
K = torch.randn(T, D)
outliers = [3, 17, 44, 91]            # persistent outlier channels
K[:, outliers] *= 60.0                # same channels for every token

def quant_dequant(x, dim, bits=4):
    """Asymmetric min-max quantization along `dim`."""
    qmax = 2 ** bits - 1
    mn = x.amin(dim=dim, keepdim=True)
    mx = x.amax(dim=dim, keepdim=True)
    scale = (mx - mn).clamp(min=1e-8) / qmax
    q = ((x - mn) / scale).round().clamp(0, qmax)
    return q * scale + mn

per_token   = quant_dequant(K, dim=1)   # scale shared across channels
per_channel = quant_dequant(K, dim=0)   # scale shared across tokens

def rel_err(a, b):
    return ((a - b).norm() / b.norm()).item()

print(f"per-token   : {rel_err(per_token, K):.4f}")
print(f"per-channel : {rel_err(per_channel, K):.4f}")

# What matters downstream: attention logits, not the raw tensor
q = torch.randn(8, D)
ref = q @ K.T
print("logit err (per-token)  :", rel_err(q @ per_token.T,   ref))
print("logit err (per-channel):", rel_err(q @ per_channel.T, ref))

Run it and per-channel wins by roughly an order of magnitude on both the tensor error and the attention logits. The logit error is the one that matters: attention applies softmax to those scores, and softmax is exponentially sensitive. An absolute logit error of ~1.0 can reorder which tokens the head attends to, which is how you get a model that reads fine for 200 tokens and then confidently cites the wrong retrieved chunk.

Note that the outlier pattern in the snippet is fixed across tokens. That's the load-bearing assumption. If outlier channels moved token to token, per-channel grouping would help far less.

Where does RoPE break this?

Rotary position embeddings rotate pairs of channels (2i, 2i+1) by an angle that depends on absolute position. That rotation mixes an outlier channel with its partner, and the mixing angle changes per token. The clean "channel 44 is always huge" structure gets smeared into "channels 44 and 45 trade magnitude as a function of position."

Two options:

Quantize post-RoPE. Simple: the cache holds exactly what attention consumes. But per-channel grouping is now fighting position-dependent rotation, so you need finer groups or more bits.
Quantize pre-RoPE. Store the raw keys, keep clean per-channel outlier structure, and apply RoPE after dequantization inside the attention kernel. This is what KVQuant-style schemes do, and it's why they hold up at 3–4 bits where naive post-RoPE schemes need 8.

Pre-RoPE isn't free. Your attention kernel must now fuse dequantize → rotate → dot-product, and RoPE application on the fly costs a little decode latency. It also breaks compatibility with any kernel that assumes the cache is directly consumable — which is most of them. Check what your serving stack actually implements before designing around it.

What should I actually turn on in production?

Start with FP8, not INT4. FP8 E4M3 has a wide dynamic range, which means the outlier-channel problem mostly stops mattering — the format absorbs a 100× spread that INT4's 16 uniform levels cannot. It halves cache memory and, on Hopper and later, has native hardware support.

In vLLM:

vllm serve meta-llama/Llama-3.1-70B-Instruct \
  --kv-cache-dtype fp8 \
  --max-model-len 131072 \
  --tensor-parallel-size 4

Without calibrated scales this uses a default scaling factor, which is usually fine for E4M3 but is the first thing to check if quality dips. Models published with pre-computed KV scales in the checkpoint will pick them up automatically.

Go below 8 bits only if you have measured that KV cache — not weights, not activations — is your binding constraint, and only with a scheme that gets the axes right. The checklist:

Keys per-channel, values per-token.
Group size on the key channel axis of 64–128, not "whole tensor."
A full-precision residual window of the most recent 32–128 tokens, appended unquantized and folded into the quantized cache in blocks. Recent tokens receive disproportionate attention mass, and this window is a rounding error in memory terms.
Evaluate on long-context retrieval tasks, not perplexity. Perplexity is dominated by short-range prediction and will look nearly unchanged while needle-in-haystack recall quietly degrades. That divergence is the single most common way teams ship a broken cache quantization config.

Why does this get worse with GQA and MLA?

Grouped-query attention already cut the KV cache by the query-to-KV head ratio — 8× for Llama-3.1-70B. Multi-head latent attention compresses further by caching a low-rank latent instead of full keys and values.

Both make quantization harder, not easier. Each cached number now carries more information: with 8 KV heads instead of 64, every cached key is read by eight query heads, so an error in it propagates eight ways. Compression schemes don't compose linearly. A 4-bit config validated on an MHA model is not validated on the GQA or MLA version of it, and stacking aggressive KV quantization on top of MLA's low-rank projection is where I'd expect trouble first. Re-measure per architecture.

The direct answer

KV cache quantization needs different axes for keys and values because key tensors contain persistent outlier channels — the same head-dimension coordinates carry 10–100× the magnitude of their neighbors across every token — while value tensors do not. Quantizing keys per-token forces one scale to span outlier and normal channels, collapsing the normal ones to a couple of levels and corrupting attention logits, which softmax amplifies. Quantize keys per-channel (ideally pre-RoPE, so rotation doesn't smear the outlier structure) and values per-token, keep the most recent 32–128 tokens in full precision, and validate on long-context retrieval rather than perplexity. If you just want the memory win with minimal risk, FP8 E4M3 sidesteps the whole problem: its dynamic range absorbs outliers that INT4's uniform levels cannot.