KV cache 為何比模型權重更能決定 LLM 伺服器可承載的用戶數

A 7B model in fp16 takes 13.5 GB of GPU memory for its weights. On an 80 GB card that leaves about 66 GB free, which sounds like plenty. It runs out at around 30 users with 4,096-token contexts, and the weights aren't what fills it. The KV cache is.

This post works through the arithmetic. All the numbers are for a Llama-2-7B-shaped model: 32 layers, 32 attention heads, head dimension 128, fp16.

What the cache stores

A decoder-only model generates one token per forward pass. The new token's query has to attend to the keys and values of every earlier token. Without a cache, each pass recomputes those keys and values for the whole prefix.

The causal mask is what makes reuse safe. Earlier tokens can't see later ones, so appending a token never changes the keys and values already computed. The model can compute them once, store them, and on each step add one new key row and one new value row per layer.

The saving is large. With a 5-token prompt and 64 generated tokens, recomputing everything costs 2,336 key/value projections per layer. With a cache it costs 68, about 34 times less. Attention still reads every cached key at each step, so the cache removes the recomputed projections, not the attention cost itself.

What it costs

The cache for one token is:

2 (K and V) × layers × kv_heads × head_dim × bytes
= 2 × 32 × 32 × 128 × 2
= 524,288 bytes ≈ 0.5 MiB

Half a megabyte per token adds up quickly:

A single 32K-token conversation needs more memory for its cache than the model needs for its weights. And the weights are paid once, while the cache is paid per sequence. So the number of users a server can batch is roughly the free memory divided by the cache for one sequence: 66.5 GB / 2.1 GB ≈ 30 sequences of 4,096 tokens, before activations and runtime buffers.

Why batching matters so much

Generation has two phases. Prefill runs the whole prompt through the model at once, so each byte of weights read from memory does work for every prompt token. Decode produces one token per step, and each byte of weights read does work for just that one token. Decode is limited by memory bandwidth, not by compute.

Batching is the fix: one read of the weights serves every sequence in the batch. But every extra sequence brings its own KV cache, so cache size is what caps the batch, and with it the throughput.

How models shrink it

Each technique cuts one term of the formula:

• Fewer KV heads. Multi-query attention (Shazeer, 2019) shares a single key/value head across all query heads. Grouped-query attention (Ainslie et al., 2023) uses a few groups. With 8 KV heads instead of 32, the same 7B shape would fit about 123 sequences of 4,096 tokens instead of 30. Llama 2 70B uses 8 KV heads; the 7B keeps all 32.
• Fewer bytes. Storing the cache in int8 halves it again, to about 247 sequences with 8 KV heads, usually at a small accuracy cost.
• Fewer tokens. A sliding window keeps only the most recent tokens. Mistral 7B uses a 4,096-token window.
• Less waste. PagedAttention (Kwon et al., 2023), used in vLLM, stores the cache in fixed-size blocks instead of reserving the maximum length for every sequence, so less memory sits empty.

The first and third are architecture choices made before or during training. Cache quantization and paging can be applied at serving time.

Try it

I built an interactive version of this as part of Lensa ML, a free set of ML explainers. The KV cache explainer lets you step through generation token by token, drag the context length and batch size to watch the cache outgrow the weights, and toggle KV heads, cache precision and a sliding window to see each saving.

References

• Pope et al. (2022), Efficiently Scaling Transformer Inference
• Shazeer (2019), Fast Transformer Decoding: One Write-Head is All You Need
• Ainslie et al. (2023), GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
• Kwon et al. (2023), Efficient Memory Management for Large Language Model Serving with PagedAttention