KV Cache Compression: Eviction and Quantization

The KV Cache sized the KV cache, and How LLM Architecture Evolved, 2017–2026 and Anatomy of a Frontier Model showed the architectural ways to shrink it: share key-value heads (GQA), or store a compressed latent (MLA). For a model that is already trained, two levers remain, both applied at serving time: keep fewer tokens in the cache, or keep each token in fewer bits. Both lose something. This post measures what.

Table of contents

The budget

# The KV cache's budget, worked for Llama-2-13B (40 layers, 40 heads of 128,
# full multi-head attention): bytes per token of context, and for one
# 32,768-token sequence, under each lever and their combinations. Group
# scales for quantized caches are counted (a 16-bit scale and zero-point
# per 32 numbers adds 1 bit per number).
L, H, D, CTX = 40, 40, 128, 32_768
full = 2 * L * H * D * 2                                     # K and V, BF16
rows = [("BF16, every token (baseline)", 16, 1.0, H),
        ("if it had 8 KV heads (GQA)", 16, 1.0, 8),
        ("4-bit K and V", 4 + 1, 1.0, H),
        ("2-bit K and V (KIVI)", 2 + 1, 1.0, H),
        ("keep a quarter of the tokens", 16, 0.25, H),
        ("GQA-8 + 4-bit + a quarter of the tokens", 4 + 1, 0.25, 8)]
print(f"{'':42s} {'bytes per token':>16s} {'one 32k sequence':>17s} {'vs baseline':>12s}")
for name, bits, frac, heads in rows:
    b = 2 * L * heads * D * bits / 8 * frac
    print(f"{name:42s} {b:14,.0f} B {b * CTX / 2**30:13.2f} GiB {b / full:11.1%}")

Full script: budget.py

                                            bytes per token  one 32k sequence  vs baseline
BF16, every token (baseline)                      819,200 B         25.00 GiB      100.0%
if it had 8 KV heads (GQA)                        163,840 B          5.00 GiB       20.0%
4-bit K and V                                     256,000 B          7.81 GiB       31.2%
2-bit K and V (KIVI)                              153,600 B          4.69 GiB       18.8%
keep a quarter of the tokens                      204,800 B          6.25 GiB       25.0%
GQA-8 + 4-bit + a quarter of the tokens            12,800 B          0.39 GiB        1.6%

A 13B model with full multi-head attention holds 800 KiB of cache per token: a single 32,768-token sequence is 25 GiB, slightly more than the model’s own 16-bit weights (26 GB, 24 GiB). (Llama-2 itself stops at 4,096 tokens; the arithmetic is the point.) The levers multiply. Five times fewer key-value heads, about three times fewer bits (4 bits plus scales), four times fewer tokens: together 1.6% of the original, 0.4 GiB. The question is what each costs in quality.

Eviction: keep fewer tokens

Eviction decides, as the sequence grows, which cached tokens to drop:

• A sliding window keeps the last $W$ tokens. Simple, bounded, and it forgets everything older.
• Attention sinks (Xiao et al., 2024, StreamingLLM; Long Context, and Where It Breaks): models put a large share of attention on the first few tokens whatever they say, as a place to park attention that has nowhere better to go. Evict those, and the softmax has to put that mass somewhere else, and the model breaks. Keep the first four tokens plus a window, and perplexity stays stable over millions of tokens (4 million in the paper), though nothing outside the window is remembered.
• Heavy hitters (Zhang et al., 2023, H2O): a small share of tokens receive most of the attention over time. Keep the tokens with the largest accumulated attention plus a recent window, and evict the rest.
• SnapKV (Li et al., 2024): after reading a long prompt, look at the attention its last few tokens pay to the rest, per head, and keep only the positions they attend to, compressing the prompt’s cache once before generation starts.

All eviction shares one weakness: an evicted token is gone. If a later question needs a detail from a part of the document no earlier token attended to, no policy that decided before the question was asked can have kept it. SnapKV picks after the first question, which sits at the end of the prompt; a follow-up in the next turn may need exactly what it dropped. H2O also needs the attention scores, which FlashAttention never materializes, which is part of why SnapKV’s one-time selection is easier to deploy.

Quantization: keep each token in fewer bits

Quantization from Scratch’s scales and zero-points apply to the cache too, with one difference: what to group. Keys and values have different shapes of error. KIVI (Liu et al., 2024) found that keys have a few channels with large values throughout a sequence, like the activation outlier channels whose weights AWQ protects, while values don’t. So it quantizes keys per channel, a scale for each channel over a group of tokens, and values per token, a scale for each token over its channels, keeping the most recent tokens in full precision until a group fills (the last 128 tokens, in the paper). With that, 2-bit caches hold up where per-token 2-bit keys don’t. In deployed servers today, the common form is simpler, an FP8 cache; lower-bit caches need attention kernels that read and dequantize the packed cache directly, or they save memory but not bandwidth.

Measured on a real model

Both levers on Qwen2.5-0.5B, on held-out text from this blog: the loss on the last 1,024 tokens of eight 2,048-token windows, with the model allowed to attend only to the tokens a policy keeps, or with its cached keys (after RoPE, as caches store them) and values rounded to fewer bits. The script runs on the CPU of an Apple M4 laptop and takes about 6 minutes after the model download. measure.py evicts tokens with an attention mask that hides them, and quantizes by registering its own attention function with transformers, which rounds the cached keys and values before it reads them. The rounding and that function:

def round_to(x, bits, per, group):
    """Round x (..., tokens, channels) to `bits`, with a scale and zero-point
    for each group: per="token", groups of `group` channels of one token;
    per="channel", one channel over groups of `group` tokens."""
    if per == "token":
        g = x.reshape(*x.shape[:-1], -1, group)
        lo, hi = g.amin(-1, keepdim=True), g.amax(-1, keepdim=True)
    else:
        g = x.reshape(*x.shape[:-2], x.shape[-2] // group, group, x.shape[-1])
        lo, hi = g.amin(-2, keepdim=True), g.amax(-2, keepdim=True)
    s = (hi - lo).clamp(min=1e-8) / (2 ** bits - 1)
    return (((g - lo) / s).round() * s + lo).reshape(x.shape)

# The cache as the attention reads it: keys after RoPE, as real caches store
# them, rounded by CACHE's settings; the last R tokens before each query
# kept in full precision (KIVI's residual window).
CACHE = {}
def attention(module, query, key, value, attention_mask, scaling, dropout=0.0, **kw):
    kq = round_to(key, *CACHE["k"]) if "k" in CACHE else key
    vq = round_to(value, *CACHE["v"]) if "v" in CACHE else value
    K, V, Kq, Vq = (repeat_kv(t, module.num_key_value_groups) for t in (key, value, kq, vq))
    scores = query @ Kq.transpose(2, 3) * scaling
    R = CACHE.get("R", 0)
    if R:
        near = (torch.arange(T)[:, None] - torch.arange(T)[None]) < R
        scores = torch.where(near, query @ K.transpose(2, 3) * scaling, scores)
    w = torch.softmax(scores + attention_mask, -1)
    out = (w * near) @ V + (w * ~near) @ Vq if R else w @ Vq
    return out.transpose(1, 2).contiguous(), None

Full script: measure.py

full cache: 3.063 bits per token, 8 windows
eviction: keep
  the last 512 tokens                          4.957  (+1.894 ± 0.031)
  4 sink tokens + the last 508                 3.235  (+0.172 ± 0.011)
  the last 256 tokens                          5.254  (+2.191 ± 0.034)
  4 sink tokens + the last 252                 3.437  (+0.374 ± 0.016)
  the last 128 tokens                          5.735  (+2.672 ± 0.039)
  4 sink tokens + the last 124                 3.689  (+0.626 ± 0.021)
quantization of the cache (keys after RoPE):
  8-bit, per token                             3.112  (+0.049 ± 0.006)
  4-bit, per token                             5.912  (+2.849 ± 0.036)
  4-bit, only values, per token                3.074  (+0.011 ± 0.002)
  4-bit, only keys, per token                  5.937  (+2.874 ± 0.036)
  4-bit, keys per channel                      3.087  (+0.024 ± 0.003)
  2-bit, per token                             11.506  (+8.443 ± 0.059)
  2-bit, keys per channel                      4.592  (+1.529 ± 0.024)
  2-bit, keys per channel, last 128 exact (KIVI) 3.388  (+0.325 ± 0.013)
control: keys per token with the key bias taken out first (before RoPE)
  4-bit keys, per token, bias removed          3.136  (+0.073 ± 0.006)
  2-bit keys, per token, bias removed          8.048  (+4.985 ± 0.047)
layer  0 keys: largest bias 130.0; a token's 64 channels span 243.7 on average,  4.4 without the bias; a channel varies from token to token by 0.40
layer 12 keys: largest bias   4.4; a token's 64 channels span  11.8 on average,  7.4 without the bias; a channel varies from token to token by 0.75

• A window alone is a disaster; four sink tokens rescue it. Keeping only the last 512 tokens costs 1.89 bits per token, more than half again the full cache’s loss. Keeping the first 4 tokens as well, and 508 recent ones, costs 0.17. At 128 tokens, 2.67 against 0.63. The four tokens at the start say almost nothing; they matter because the model parks attention on them.
• Eviction still costs. Even with sinks, keeping 512 of up to 2,048 tokens costs 0.17 bits per token, and 128 costs 0.63: this text does refer back further than a window reaches. Smarter selection (H2O, SnapKV) keeps more of what is needed, but no selection made before a question is asked keeps everything it might need.
• 8 bits is nearly free; 4 bits per token fails, and the keys are why. Rounding keys and values per token costs 0.05 bits at 8 bits and 2.85 at 4 bits. Rounding only the values at 4 bits costs 0.011; only the keys, 2.87.
• In Qwen, the key bias makes it extreme. Qwen’s key projections add a bias, a fixed offset per channel, up to 130 in the first layer, and Llama has none. A token’s 64 key channels span 244 on average there, while any one channel varies from token to token by 0.4, so a 4-bit grid stretched over each token, with a step of about 16, erases everything that differs between tokens. In a control that takes the bias out and rounds the keys before RoPE, 4-bit per-token keys cost only 0.07. But at 2 bits, per-token keys still fail without the bias (+5.0): keys also carry outlier channels, which KIVI found in bias-free Llama and Mistral. The bias exaggerates the per-channel structure; it doesn’t create it.
• Per-channel keys, and a full-precision window, fix it. A scale per channel over groups of 32 tokens absorbs constant offsets and outlier channels alike: 4-bit keys per channel cost 0.024 bits per token. At 2 bits, per-channel keys cost 1.53, against 8.44 per token; keeping the last 128 tokens in full precision, KIVI’s full scheme, brings 2 bits down to 0.33.

Three simplifications. Eviction is simulated by masking, so kept tokens keep their original positions, while StreamingLLM renumbers the positions inside its cache; at 2,048 tokens, well inside Qwen2.5’s trained context, that should matter little (it was not measured here), and renumbering matters for streams longer than the model was trained on. Per-channel keys are rounded a whole group of 32 tokens at a time, so a query in the middle of a group reads keys whose scale was set partly by tokens after it; a real cache rounds a group only once it is full, which is what KIVI’s full-precision window is for. And quantization here measures quality only: rounded values are stored back as 32-bit floats, so it says nothing about speed, which depends on kernels that read the packed cache. KVQuant (Hooper et al., 2024) goes further than KIVI: it rounds keys before RoPE, where a channel’s offsets haven’t yet been rotated across positions, uses non-uniform levels, and keeps a few outliers and the sink token in full precision.

Where this leads

• Paged Attention and vLLM Architecture manage the cache that remains, in blocks.
• Speculative Decoding, Derived makes each decode step, and so each read of the cache, produce more than one token.

Takeaways

• The KV cache can outweigh the model: 25 GiB for one 32k-token sequence with Llama-2-13B’s shapes. Fewer heads, fewer bits and fewer tokens multiply: together they cut it to under 2%.
• Eviction needs the attention sinks: keeping the last 512 tokens alone cost 1.89 bits per token on held-out text, and adding the first 4 tokens cut that to 0.17. Anything evicted is gone for good.
• Quantizing the cache per token failed at 4 bits (+2.85), because of the keys: values alone cost 0.011. Keys per channel at 4 bits cost 0.024.
• At 2 bits, per-channel keys cost 1.53 against 8.44 per token, and KIVI’s full-precision window of the last 128 tokens brought it to 0.33.
• Group by where the structure is: keys vary by channel, values by token. In Qwen a large key bias makes the per-token failure extreme, but the per-channel structure is there without it.

References

Start with StreamingLLM for attention sinks and KIVI for per-channel keys, the two results this post’s measurements reproduce; SnapKV and H2O are the smarter eviction policies.

• Xiao, Tian, Chen, Han and Lewis, Efficient Streaming Language Models with Attention Sinks (ICLR, 2024): StreamingLLM, attention sinks and the sinks-plus-window policy.
• Zhang et al., H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models (NeurIPS, 2023): heavy-hitter eviction by accumulated attention.
• Li et al., SnapKV: LLM Knows What You are Looking for Before Generation (NeurIPS, 2024): one-time compression of a prompt’s cache from what its last tokens attend to.
• Hooper et al., KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization (NeurIPS, 2024): keys rounded before RoPE, non-uniform levels and full-precision outliers.
• Liu et al., KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache (ICML, 2024): per-channel keys, per-token values and the full-precision recent window.