Inference and reasoning / Inference scaling and serving / The KV cache, and why decoding is memory-bound

The KV cache, and why decoding is memory-bound

A language model generates one token at a time, and each new token attends to every token before it. The KV cache is the optimisation that stops the model recomputing the past at every step. It removes most of the arithmetic, and in doing so it changes what the hardware is waiting on: not the tensor cores, but the memory bus. This note explains both halves, puts numbers on the cache, and lists what architectures do about it.

The short version

Keys and values for a past token never change, so compute them once and keep them. That turns generation from quadratic to linear in the sequence length. But a decode step then does very little arithmetic per byte it reads: the whole weight set and the whole cache stream from memory to produce a single token. Throughput is set by bandwidth, not FLOPs, and the cache grows linearly with context until it is larger than the model. Grouped-query attention, paged allocation and lower-precision caches are the standard responses.

01What is cached, and why it can be

In every attention layer, each token is projected into a query, a key and a value (the previous note covers why those are three different things). The new token's query is dotted with the keys of all earlier tokens, and the resulting weights average their values.

The point that makes caching possible: in a decoder, the key and value of token t depend only on tokens up to t, and those tokens do not change once generated. So the key and value computed for token 5 while generating token 6 are exactly the ones needed while generating token 7, 8 and 500. Recomputing them is pure waste.

The KV cache is therefore, per layer, two tensors that grow by one row per generated token: the stack of keys and the stack of values for everything so far. At each step the model projects only the newest token, appends its key and value, and runs attention for its single query against the stored rows.

step n, each layer
  q_n, k_n, v_n  =  x_n W_Q,  x_n W_K,  x_n W_V        ← one token projected
  K_cache  ←  [K_cache ; k_n]                          ← append
  V_cache  ←  [V_cache ; v_n]
  out_n    =  softmax(q_n K_cacheT / √d_k) V_cache          ← one query, n keys

The query is not cached, because it is used once and discarded.

02What the cache saves

Consider generating N tokens after a prompt.

  • Without the cache, producing token n means re-running every earlier token through WK and WV to rebuild the keys and values: n projections at step n. Summed over the sequence that is 1 + 2 + … + N, so O(N2) projections per layer, and the per-step latency grows as the answer gets longer.
  • With the cache, step n projects one token, so N tokens cost O(N) projections per layer. The attention itself still reads n keys at step n, which is unavoidable, but that is a dot product per key rather than a full matrix projection per key.

For any realistic sequence the saving is enormous, and no production decoder runs without it. The interesting part is what happens next.

03Prefill and decode are different workloads

Serving a request has two phases, and they stress the hardware in opposite ways.

PrefillDecode
What happensThe whole prompt is processed in one pass, filling the cacheTokens are produced one at a time, each appending to the cache
Shape of the workMatrix × matrix: thousands of tokens against every weightVector × matrix: one token against every weight
Each weight loaded is usedOnce per prompt token, thousands of timesOnce
Bound byCompute (FLOPs)Memory bandwidth (bytes per second)
Shows up asTime to first tokenTime per output token

Prefill is what GPUs are built for. Every weight matrix is read from memory once and multiplied against a block of prompt tokens, so the tensor cores stay busy. Decode is the opposite, for two reasons that compound.

04Reason 1: low arithmetic intensity

Arithmetic intensity is the number of floating-point operations a kernel performs per byte it moves from high-bandwidth memory into the chip. Every accelerator has a ridge point: the intensity at which compute and memory traffic take equal time. Below it, the chip is waiting on memory and its FLOPs rating is irrelevant.

intensity = FLOPs / bytes moved,    ridge = peak FLOPs/s / bandwidth bytes/s

For a current datacentre GPU, peak dense 16-bit throughput is on the order of 1,000 TFLOPs/s and memory bandwidth on the order of 3 TB/s, so the ridge is roughly 300 FLOPs per byte. A kernel needs to do about 300 operations on every byte it loads to keep the chip busy.

Now look at a decode step at batch size 1. Each weight is a 16-bit value, two bytes. It is loaded and used in exactly one multiply-add for the one token being generated, which is two FLOPs. The intensity is about 1 FLOP per byte, two orders of magnitude below the ridge. The same is true of the cache: every stored key is loaded, dotted once against the one query, and discarded until the next step. The GPU streams gigabytes through the memory bus and performs almost no arithmetic on them.

This is why batching helps decode so much. Eight concurrent requests share one read of the weights, so the intensity rises eight-fold for the weight matmuls. It does not help the attention part, because each request has its own cache, which brings us to the second reason.

05Reason 2: the cache grows with context

Every token in the context contributes one key and one value per layer, per key-value head. The total footprint is:

bytes = 2 × nlayers × nkv heads × dhead × s × b × p

  • 2, one for keys and one for values
  • nlayers, the number of Transformer layers
  • nkv heads × dhead, the width of the key and value projections; for standard multi-head attention this equals dmodel
  • s, the sequence length in tokens, prompt plus generated
  • b, the batch size
  • p, bytes per element, 2 for FP16 or BF16

Two things about this formula matter. It is linear in s, so doubling the context doubles the cache. And it is linear in b, so the cache, unlike the weights, is paid per request. A server holding 32 concurrent long conversations holds 32 caches.

At every decode step the whole cache for the request is read to compute attention. So a longer context does not only take more memory; it makes each subsequent token slower, because more bytes have to stream through the same bus to produce it.

06A 70B model at 2k and 128k

Take a 70-billion-parameter model with 80 layers, a hidden size of 8,192 and 64 heads of 128 dimensions, served in FP16 at batch size 1.

ContextMulti-head attention (64 KV heads)Grouped-query attention (8 KV heads)
2,048 tokens5.4 GB0.67 GB
8,192 tokens21.5 GB2.7 GB
32,768 tokens85.9 GB10.7 GB
131,072 tokens343.6 GB42.9 GB

The left column is the formula with nkv heads × dhead = dmodel: 2 × 80 × 8192 × 2048 × 2 bytes is 5.4 GB at 2k, and 343.6 GB at 128k, which is more than twice the 140 GB the weights themselves occupy in FP16. A single request's cache would not fit on any GPU, and it is read in full for every token generated.

The right column is the same model with grouped-query attention and 8 key-value heads, which is how Llama 2 70B and Llama 3 70B are actually built. The cache is eight times smaller, and 128k of context costs 43 GB. That is still a third of the weights, per request, and it still has to be streamed for every token. The architecture change moved the problem; it did not remove it.

A rule of thumb for decode latency

At batch size 1, a decode step cannot be faster than (weight bytes + cache bytes) divided by memory bandwidth. For the GQA model above at 32k context on hardware with 3 TB/s, that is about (140 + 10.7) GB / 3 TB/s, roughly 50 ms per token, before any compute or overhead. Spread across several GPUs with tensor parallelism the bytes per device fall, which is the main reason large models are served that way even when they would fit on fewer devices.

07Architectural mitigations

Each term in the size formula is a lever, and the serving stack pulls several at once.

Multi-query attention

All query heads share a single key head and a single value head. The nkv heads term becomes 1, so the cache shrinks by the number of heads, 64× in the example. The cost is some model quality, since every head now has to find what it needs through the same key space. PaLM and Falcon used it.

Grouped-query attention

The compromise that most current models adopt: query heads are divided into groups, and each group shares one key-value head. With 64 query heads and 8 KV heads, the cache is 8× smaller than multi-head and quality is close to it. Llama 2 and 3, Mistral and most models since use GQA. A further step in the same direction is multi-head latent attention, which stores a compressed latent per token and reconstructs keys and values from it.

Paged attention

The cache is a tensor that grows one row at a time, and nobody knows in advance how long a response will be. Naive serving reserves the maximum context for every request up front, so most of the reserved memory is empty, and freed regions of different sizes fragment. vLLM's PagedAttention manages the cache like an operating system manages virtual memory: fixed-size blocks allocated on demand and indexed through a block table. Waste falls from a large fraction of memory to a few percent, so far more requests fit on a device, and shared prefixes such as a system prompt can be stored once and referenced by many sequences.

Cache quantisation

The p term: storing keys and values in 8-bit integers or FP8 instead of 16-bit halves the cache and halves the bytes streamed per step. Keys are more sensitive than values, so some schemes keep keys at higher precision or quantise per channel. Four-bit caches exist and cost some accuracy on long-context tasks.

Two more that are not about the architecture

  • Prefix caching. Many requests share a long prefix: a system prompt, a document, a tool list. Keeping that prefix's cache across requests turns its prefill into a lookup, which is what the hosted providers' prompt caching exposes.
  • Sliding windows and eviction. Limit attention to the last w tokens, or drop cache rows that attention weights show are never used. Both bound the s term at the cost of forgetting.

08Takeaways

  1. Keys and values of past tokens never change, so cache them. Generation goes from quadratic to linear in sequence length.
  2. Decode is a vector against every weight. About one FLOP per byte loaded, against a ridge of hundreds. The memory bus, not the tensor cores, sets the pace.
  3. The cache is linear in context and in batch. It is paid per request and streamed in full for every token.
  4. At 128k, a full multi-head cache is larger than the weights. GQA divides it by the head-group ratio; it does not make it small.
  5. Batching helps the weights, not the attention. Concurrent requests share one weight read but each brings its own cache.
  6. The fixes map onto the formula's terms. MQA and GQA shrink the head term, quantisation the precision term, paging the waste around the sequence term, prefix caching the prefill that fills it.