Vol. VI · Positions and Long Context Plate 9 of 10KV cache, long context, memory, grouped-query attention, 128k tokens, serving costC

What a Long Context Costs.

Every token kept in context costs a key and a value in every layer. Multiply by 128 000 and the cache can be larger than the model.

To generate the next token, a model needs the keys and values of every earlier token in every layer — that is what the key–value cache holds (Vol. XXIII). Per token it stores 2 × L × (number of key–value heads) × (head dimension) numbers: for the published Llama 3 8B configuration, 2 × 32 × 8 × 128 = 65 536 numbers, 128 kilobytes in 16-bit precision. The cache grows linearly with context, one such slice per token, and it must sit in fast accelerator memory beside the weights.

Context length also costs compute — attention’s 4 T d per layer per token (plate IV.7) — but at serving time memory is usually the binding constraint: it limits how many conversations fit on one accelerator at once, and so the cost per token. The workshop computes the cache for published configurations at any length and compares it with the weights.

Below: the key–value cache of one sequence against context length for the chosen configuration, with the size of the model’s 16-bit weights for comparison; and the same configuration without grouped-query attention, as if it had one key–value head per query head.

What you will learn
Workshop
Key–value cache of one sequence (16-bit) against context length (log scales) for the chosen configuration (blue) and the same configuration with as many key–value heads as query heads (red); the dashed line is the size of its 16-bit weights.

The cache outgrows the model

The published Llama 3 8B configuration keeps 131.1 KB of keys and values per token. At its extended context of 131 072 tokens one sequence’s cache is 17.2 GB, 1.07 times the 16.1 GB of its 16-bit weights; the two are equal at about 123 thousand tokens. Grouped-query attention — 8 key–value heads shared by 32 query heads (plate VII.2) — is what keeps this manageable: with a key–value head per query head the cache would be 4 times larger. GPT-3, with full multi-head attention over 96 layers, kept 4.7 MB per token, about 9.7 GB for its 2 048-token context.

This arithmetic shapes long-context engineering. Serving systems page the cache in blocks and share it between requests with common prefixes (Vol. XXIII); the cache can be quantised to 8 or 4 bits (plate XXIV.3); architectures shrink it with grouped and latent attention (Vol. VII) or windows (plate VI.7); and the price of a long prompt reflects all of it. A 128 000-token context is a capability bought with memory, and the next volumes count the bill.

These figures are for one sequence. A server holding 32 conversations of 32 000 tokens each in Llama 3 8B needs 32 × 32 000 × 128 KB ≈ 134 GB of cache — more than one accelerator — which is why batching and cache management dominate serving design.

Terms

Key–value cache
The stored keys and values of every earlier token, in every layer, reused at each generation step.
Per-token cache size
2 × layers × key–value heads × head dimension × bytes.
Grouped-query attention
Sharing each key–value head among several query heads to shrink the cache.
Context length
The number of tokens a model reads at once.
Each token in context costs 2 L H_KV d_h numbers of cache; at 128 000 tokens the cache of one sequence can outweigh the model’s weights.

Sources

  1. Llama 3 8B and its 128 000-token extension: Llama Team, AI @ Meta, arXiv:2407.21783, 2024.
  2. Mistral 7B: Jiang et al., arXiv:2310.06825, 2023.
  3. GPT-3: Brown et al., NeurIPS 2020.