Lensa ML
Lensa ML

KV Cache

Step 1 of 6: One Token at a Time

Each new query attends to every earlier key

Prompt only: the last token's query reads every keynext token → "Paris"ThecapitalofFranceispromptquery
Sequence length
5
Keys this query reads
5
Next token
"Paris"
q·k scores so far
5

The prompt has 5 tokens. The query of "is" attends to all 5 keys and values and predicts "Paris".

Press Play, or Step through one token at a time, and watch the arcs from the newest token. Each arc is one query-key score, q · k / √d, with d = 128 per head. The causal mask is what makes caching safe: appending a token never changes the keys and values of earlier positions, because they could not see it. In a bidirectional encoder such as BERT, adding a word changes every earlier representation, so nothing could be reused. The continuation after the prompt is illustrative, not the output of a real model.

Step 2 of 6: The Wasted Work

Recomputing the prefix at every step

Tokens generated16
Prompt length P5
K/V vectors computed per step (one row per step, per layer)No cacheposition → (20 total)With KV cacheposition → (20 total)computed, newrecomputedread from cache
K/V computed, no cache
200
K/V computed, cached
20
Ratio
10.0×

At 16 steps the waste grows with the square of the length: 200 pairs versus 20, 10.0 times the work, in each of the model's 32 layers.

Drag Tokens generated from 1 to 64 and Prompt length P from 1 to 16. Without a cache, step s recomputes keys and values for all P + s − 1 positions, so T steps cost T · P + T(T − 1)/2 projections per layer. With a cache the cost is P + T − 1. At P = 5 and T = 64 that is 2,336 against 68, about 34 times the work. Attention itself still reads every cached key, so the cache removes the recomputed projections, not the growing attention cost.

Step 3 of 6: Caching Keys and Values

Append one k and v, reuse the rest

Prefill: write 5 prompt tokens to the cachetokenK cacheV cachesoftmax(q·k), q of "is"The7%capital23%of33%France5%is31%Vector values and attention scores are illustrative
Cache length
5
New K/V computed
5
q·k dot products
5
Cache, this sequence
2.5 MiB

Prefill computes k and v for all 5 prompt tokens in one pass and stores them. The last prompt token's query then scores all 5 cached keys.

Press Play or Step and watch one k row and one v row join the cache each step. In practice the cache is a pair of tensors per layer, shaped [batch, kv_heads, tokens, head_dim], and every decode step writes one new slice along the token axis. Only the newest token passes through the Q, K, and V projections and the feed-forward layers. Position information is applied before storage: with rotary embeddings (RoPE), each key is rotated for its position once and cached in rotated form.

Step 4 of 6: Prefill vs Decode

Compute-bound versus memory-bound

Prompt length512
Output tokens256
TFLOP/s attainable (illustrative GPU: 300 TFLOP/s, 2 TB/s)0.11101001k10k110100FLOP per byte read from memory (log scale)memory-boundcompute-bounddecodeprefill
Tokens in prefill
512
Decode steps
255
Read per decode step
13.9 GB
Prefill FLOP/byte
502

Prefill handles 512 tokens in one pass at about 502 FLOP per byte (compute-bound, roughly 23 ms on this GPU). Each of the 255 decode steps reads up to 13.9 GB for one token, about 6.9 ms each (illustrative timings).

Drag Prompt length and Output tokens, then compare the two dots on the roofline. Prefill multiplies every weight matrix by all P prompt tokens at once, so each byte of weights read does about P floating-point operations. Decode reuses each byte for only one token, about 1 FLOP per byte, far below the ridge of roughly 150 on this illustrative GPU. That is why time per output token tracks memory bandwidth, and why batching many users raises throughput: one read of the weights serves every sequence in the batch.

Step 5 of 6: The Memory Bill

Cache size grows with context and batch

Context length (tokens)4,096
Batch size (sequences)1
2 × 32 layers × 32 heads × 128 × 2 B = 512 KiB per token× 4,096 tokens × 1 sequence = 2.1 GB24 GB GPU80 GB GPUweights 13.5 GBKV cache 2.1 GBfp16, 7B chat model; activations and runtime overhead not counted
KV cache
2.1 GB
Weights
13.5 GB
Total
15.6 GB
Fits on
24 GB GPU

The cache is 2.1 GB, smaller than the 13.5 GB of weights, and the total fits on a 24 GB GPU.

Drag Context length from 512 to 32,768 and Batch size from 1 to 64. The cache per token is 2 × layers × kv_heads × head_dim × bytes: 2 × 32 × 32 × 128 × 2 = 524,288 bytes for the 7B model, or 0.5 MiB. A 4,096-token sequence therefore needs 2 GiB (2.1 GB). Weights are paid once, while the cache is paid per sequence, so the batch size a server can run is roughly free memory divided by the cache for one sequence. Activations and runtime buffers come on top.

Step 6 of 6: Shrinking the Cache

Fewer heads, fewer bits, fewer tokens

Key/value heads
Cache precision
Sliding window
32 query heads32 key/value heads (MHA), 1 query head eachMHA, fp16, full context: 68.7 GBThis setup: 68.7 GB512 KiB/token × 16,384 tokens × 8 sequences (16,384-token context)
Per token
512 KiB
Cached tokens
16,384
Cache size
68.7 GB
Smaller by
1×

Full multi-head attention at fp16 keeps 68.7 GB for 8 sequences of 16,384 tokens. Pick fewer key/value heads, int8, or a window to shrink it. PagedAttention (vLLM) attacks a different waste: it stores the cache in small blocks allocated on demand, so no memory is reserved for tokens a sequence never reaches.

Toggle Key/value heads, Cache precision, and Sliding window and watch the cache size. Multi-query attention (Shazeer, 2019) shares one key/value head across all query heads. Grouped-query attention (Ainslie et al., 2023) uses a few groups, and Llama 2 70B uses 8 KV heads while the 7B keeps full multi-head attention. Both are architecture choices made before or during training, as is a sliding window, used by Mistral 7B with 4,096 tokens. Int8 or lower-precision caches can be applied at serving time, usually with a small accuracy cost.

What is a KV cache in LLM inference?

A KV cache is the memory where a decoder-only language model keeps the key and value vectors of every token it has already processed, so each new token only needs its own keys and values computed. During generation the model produces one token per forward pass, and each new token's query has to attend to all earlier tokens. Without the cache, every pass would recompute keys and values for the whole prefix.

Every production LLM server uses a KV cache, from Hugging Face Transformers to vLLM, TensorRT-LLM, and llama.cpp. It turns generation from quadratic recomputation into one short pass per token, but it costs memory that grows with context length and batch size. That memory, often larger than the model weights, limits how many users a GPU can serve at once and how long a context it can hold, which is why grouped-query attention, cache quantization, and paged memory exist.

Key ideas

  • A KV cache stores each token's keys and values so they are computed once, not at every generation step.
  • The causal mask guarantees that cached keys and values never change when new tokens are appended.
  • Prefill processes the prompt in parallel and is compute-bound, while decode emits one token per step and is memory-bound.
  • Cache size equals 2 × layers × kv_heads × head_dim × bytes × tokens × batch.
  • For a 7B model at fp16 the cache costs 0.5 MiB per token, so long contexts and large batches can outgrow the weights.
  • Grouped-query attention, multi-query attention, cache quantization, and sliding windows shrink the cache, and paged allocation stops memory from being wasted.

Common questions

Does a KV cache change the model's output?

No. With exact arithmetic the cached and uncached computations produce the same logits, since the cache only stores values that would otherwise be recomputed. Small differences can appear from floating-point ordering in different kernels. Quantizing the cache to int8 or lower is an approximation and can shift outputs slightly.

What is prompt caching?

Prompt caching, also called prefix caching, keeps the KV cache of a shared prefix, such as a long system prompt or document, and reuses it across requests. A new request that starts with the same tokens skips their prefill. API providers often bill cached input tokens at a lower price for this reason.

Why is the first token slower than the rest?

The first token waits for prefill, which runs the whole prompt through the model before any output appears. This delay is called time to first token, and it grows with prompt length. Later tokens each take one decode step, so their pace depends mainly on memory bandwidth and batch size.

What is PagedAttention?

PagedAttention, introduced with the vLLM serving system by Kwon et al. (2023), stores the KV cache in fixed-size blocks instead of one contiguous buffer per sequence. Blocks are allocated as a sequence grows and can be shared between sequences with a common prefix. This removes most of the memory lost to reserving space for the maximum length, which lets a server batch more requests.

Related explainers

References

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

Last updated