KV cache explained: the formula, a diagram, and a memory example

What the KV cache stores

In transformer attention, a token produces a query, a key, and a value in each attention layer. The query is compared with available keys to determine how to combine their values. These are numerical vectors, not a database of facts or a copy of the model’s weights. The original Transformer paper defines the attention operation.

For a causal language model, previously processed positions cannot depend on tokens that come later. Their keys and values can therefore be retained and reused during generation. At a decode step, the model computes attention state for the newly processed token, appends its key and value, and uses its query to attend to the available history. That retained state is the key-value cache, usually shortened to KV cache.

Caching avoids repeating previous tokens’ computation, but attention still reads relevant cached state. A longer full-attention context can therefore increase the work for a new token. Cached generation is not constant-cost at every sequence length. Hugging Face’s cache explanation describes the tensors and update step.

Calculate the raw cache bytes

For a standard full-attention model with equal key and value head dimensions, the uncompressed cache payload for one sequence is:

bytes = 2 × layers × kv_heads × head_dimension × cached_tokens × bytes_per_element

Grouped-query attention (GQA) lets multiple query heads share a key/value head. Multi-head attention has separate KV heads for the query heads; multi-query attention uses one KV head. GQA sits between these arrangements. Count num_key_value_heads, when the model config exposes it, rather than automatically substituting num_attention_heads. The GQA paper explains this architectural choice.

To derive the formula, count kv_heads × head_dimension scalars in one token’s key tensor, double it for the value tensor, then multiply by tokens and layers. Finally, convert scalars into bytes. If dimensions or attention types vary by layer, calculate each layer separately and sum the results.

Work through a 512 MiB example

Take an illustrative configuration with 32 attention layers, 8 KV heads, a head dimension of 128, and two-byte cache elements. Each cached token occupies 2 × 32 × 8 × 128 × 2 = 131,072 bytes, or 128 KiB. At 4,096 cached tokens, that becomes 512 MiB per sequence. These are calculated tensor sizes, not measurements of an inference engine.

layers = 32
kv_heads = 8
head_dimension = 128
bytes_per_element = 2

bytes_per_token = (
    2 * layers * kv_heads * head_dimension * bytes_per_element
)

print(f"Per token: {bytes_per_token / 1024:.0f} KiB")
for cached_tokens in (1024, 4096, 8192):
    payload_bytes = bytes_per_token * cached_tokens
    print(f"{cached_tokens:,} tokens: {payload_bytes / 1024**2:.0f} MiB")

active_lengths = [4096, 4096, 1024]
batch_bytes = bytes_per_token * sum(active_lengths)
print(f"Three sequences: {batch_bytes / 1024**2:.0f} MiB")

assert bytes_per_token == 131072
assert bytes_per_token * 4096 == 512 * 1024**2
assert batch_bytes == 1152 * 1024**2

Per token: 128 KiB
1,024 tokens: 128 MiB
4,096 tokens: 512 MiB
8,192 tokens: 1024 MiB
Three sequences: 1152 MiB

The three-sequence example adds the actual cached lengths. Multiplying by batch size works when all sequences have equal lengths. A densely padded allocation might reserve more; a paged allocator can allocate blocks as needed. For capacity planning, include the prompt and the generation budget. A sampled final token may not yet have been fed through the model, so exact cached length at a particular instant can differ by one from the visible sequence length.

Why real GPU usage differs from the formula

The calculation covers the logical KV payload. It excludes weights, intermediate activations, attention workspaces, allocator overhead, reserved memory, and execution-graph allocations. Treat it as one line in a memory budget, not as a GPU sizing guarantee. Loading the model and observing free memory is also insufficient: some allocations appear only when work begins.

Cache layout is another variable. A static cache can reserve capacity beyond the current sequence length. Sliding-window or chunked-attention layers need not retain an ever-growing full history; mixed architectures require a layer-by-layer calculation. Offloading moves some state to another memory tier, introducing transfer costs. Hugging Face’s cache strategy guide documents these alternatives.

Multi-head Latent Attention (MLA) uses a different representation, so the ordinary expanded-KV formula is not the right model for every implementation. DeepSeek-V2, for example, describes caching a compressed latent representation. Inspect the architecture and the serving backend’s actual storage format. DeepSeek-V2 paper.

Likewise, do not divide the total by the number of GPUs without checking the parallel layout. Some dimensions may be sharded while others are replicated. Use per-device allocation information from the engine to validate a deployment budget.

Paging, prefix reuse, and quantization solve different problems

Paged KV storage divides sequences into blocks and maps their logical positions to physical allocations. This reduces the need for one large contiguous reservation and makes memory sharing practical. It does not make every token’s key and value disappear: partly filled blocks and metadata still matter. The PagedAttention paper explains the memory-management approach behind vLLM.

Prefix caching reuses compatible cached state for an identical token prefix. Two requests with the same long system prompt may reuse work; two differently worded prompts about the same topic generally cannot. Arrange genuinely shared content consistently, then measure cache hits with realistic traffic. Reuse primarily saves repeated prefill, not the new answer’s decode work. vLLM’s prefix caching documentation.

KV quantization reduces the precision of stored cache elements. A one-byte representation halves the main element payload relative to two-byte storage, but scaling information and implementation details affect the actual allocation. Quantization can change accuracy or add conversion overhead. Verify backend support and compare quality and latency on your workload before treating the smaller arithmetic estimate as a win. vLLM’s quantized KV-cache guide.

Investigate cache pressure with controlled changes

1. Read the model configuration and identify its attention types, KV-head count, head dimensions, and cache dtype. Compute a baseline for the intended sequence lengths.
2. Run one short request, then a longer one with the same output cap. Compare active cache tokens or blocks with the expected growth. A reserved pool may stay the same size even while its occupancy changes.
3. Increase concurrency separately from length. Watch waiting requests, cache occupancy, preemptions, errors, and latency so that queueing is not mistaken for faster execution.
4. Change one constraint, such as maximum context or admitted concurrent requests. Re-run the same workload and check that answers still meet the application’s requirements.

For this example, doubling the cached length doubles the logical payload. Replacing eight KV heads with thirty-two would quadruple it if all other factors stayed equal. You can check those predictions without a GPU using the script, then compare them with the engine’s behavior. Any gap becomes a specific question about allocation or architecture instead of an unexplained out-of-memory error.