KV-Cache and Inference Cost

The KV-cache is what makes long-context inference expensive

Quadratic cost is the training story. The inference story is the KV-cache: during autoregressive generation, you cache the key and value tensors for every token you've already processed so you don't recompute them on each new token. The cache grows linearly with context length, and at long context it dominates GPU memory.

The numbers are brutal. A 70B-parameter model with 64 attention heads, head dim 128, and 80 layers, serving fp16, uses roughly 2 × n_layers × n_heads × head_dim × 2 bytes per token of context. At 128K tokens that works out to about 40+ GB just for the cache — on top of the model weights. You can fit one user session on a single H100, maybe. Batching collapses.

GQA and MQA — buying back KV memory

Grouped Query Attention (GQA, used by Llama 2/3) shares K and V heads across groups of query heads. Llama 2 70B uses 8 KV heads for 64 query heads, an 8× KV-cache reduction with negligible quality loss. Multi-Query Attention (MQA) goes all the way: a single KV head, an N× reduction, with measurable quality loss on harder tasks. Most modern frontier models use GQA as a default.

Both still grow linearly with context. They reduce the constant — they don't change the slope. Which is why serving a 128K-context model still costs roughly 3–5× per request compared to 8K, even with GQA. That cost ratio shows up in API pricing tiers and in the long-context tax that every long-context provider passes through to customers.

Code

KV-cache size in GB for common model shapes·python

def kv_cache_gb(seq_len, n_layers, n_kv_heads, head_dim, dtype_bytes=2):
    # Per token: 2 (K and V) * layers * heads * head_dim * dtype_bytes
    bytes_per_tok = 2 * n_layers * n_kv_heads * head_dim * dtype_bytes
    return seq_len * bytes_per_tok / (1024 ** 3)

# Llama 3.1 70B uses GQA: 8 KV heads, 80 layers, head_dim 128
for n in [8_000, 32_000, 128_000, 1_000_000]:
    print(f'{n:>9} tok -> {kv_cache_gb(n, 80, 8, 128):>6.2f} GB KV-cache')
# 8000   -> 2.50 GB
# 32000  -> 10.00 GB
# 128000 -> 40.00 GB
# 1000000 -> 312.50 GB

Exercise

Compute the per-token KV-cache size for three model families you're likely to deploy: Llama 3.1 8B (GQA, 8 KV heads, 32 layers, head_dim 128), Llama 3.1 70B (8 KV heads, 80 layers, 128), and a hypothetical "MHA 70B" (64 KV heads, 80, 128). For each, plot the GPU memory needed for the KV-cache alone at context lengths from 4K to 1M. Then mark the line where you'd no longer fit on a single H100 80GB. This plot is the picture every serving team has on their wall.

The KV-cache is what makes long-context inference expensive

GQA and MQA — buying back KV memory

Code

External links

Exercise

Progress

Comments 0