KV cache: what do we stop computing, and start storing?

Last reviewed: 2026-10 · Prerequisites: Multi-head attention, Decoder-only

After generating one token, should a model recompute the whole conversation from scratch? It does not need to. Some earlier calculations can be kept: that is the starting point for a KV cache.

“Caching makes it fast” is not enough, though. Long conversations consume more memory, and concurrent requests can move the bottleneck elsewhere. Follow a short generation to separate computation, storage, and the optimizations that address each.

Start with three input tokens

Suppose a prompt is tokenized as [A, B, C]. These letters are placeholders, not actual tokenizer output. The model processes the three positions in parallel and uses the output at C to predict D.

Immediately after sampling D, its K/V have not been computed. They appear when D is fed back through the model. Reversing this order can make an implementation process a token twice or skip a position.

Why can earlier positions stay unchanged?

In a standard causal decoder, A attends only to A; B attends only to A and B. Appending D does not change what A, B, or C could see. With fixed weights, positional encoding, and prefix, inference can reuse their per-layer K/V. An ordinary bidirectional encoder does not satisfy this condition: appending a token can change previous representations. The Transformers caching explanation develops this causal argument.

At one layer, for the new position $t$:

$$ \begin{aligned} K_{\le t}&=[K_{<t};k_t],\\ V_{\le t}&=[V_{<t};v_t]. \end{aligned} $$

$$ o_t=\operatorname{softmax}\left(\frac{q_tK_{\le t}^{\top}}{\sqrt{d_h}}\right)V_{\le t}. $$

The query is only the current $q_t$, but keys and values include all visible history. Earlier queries are not needed for this new attention row, so they usually need not be cached.

This is not an answer cache, and it does not skip the current token's FFN or remaining layers. It trades storing per-layer K/V for avoiding repeated computation at old positions.

Cheaper steps do not make long context free

A single query head attending over history of length $t$ still computes roughly $t$ matching scores. It avoids rebuilding the entire $t\times t$ matrix for old queries, but the cost per step is not constant.

Starting from a prompt of length $P$ and processing $N$ more tokens, the cumulative query–key match count per head is approximately:

$$ \sum_{j=1}^{N}(P+j)=NP+\frac{N(N+1)}2. $$

This excludes prefill, projections, FFNs, and hardware parallelism. It simply shows that history grows during generation: caching does not remove the cost of attending to that history.

Calculate the memory yourself

For a decoder with identical layer structures, unsharded caches, equal sequence lengths, and no window eviction or cache quantization:

$$ M_{KV}=2\,B\,L\,T\,H_{KV}\,d_h\,b. $$

Take a hypothetical model with 32 layers, 32 query heads, 8 KV heads, head dimension 128, 8192 cached tokens, and BF16 at 2 bytes per element.

def kv_bytes(batch, layers, tokens, kv_heads, head_dim, bytes_per_value=2):
    return 2 * batch * layers * tokens * kv_heads * head_dim * bytes_per_value

one_sequence = kv_bytes(1, 32, 8192, 8, 128)
eight_sequences = kv_bytes(8, 32, 8192, 8, 128)
print(one_sequence / 2**30, eight_sequences / 2**30)

The results are 1 GiB and 8 GiB. These are raw KV sizes, excluding weights, workspaces, and allocator headroom. With 32 KV heads, one sequence would require 4 GiB. For unequal lengths, sum over actual sequence lengths; preallocation, padding, sharding, or replication can change physical usage.

The weights fit. Why does the first batch run out of memory?

“A 7B model in BF16 is about 14 GB. Shouldn't a 24 GiB device be enough?” It may be enough to load weights. Requests also need KV state, temporary tensors, and kernel workspaces. Longer inputs and higher concurrency expose that difference. (Counting weights alone, how many GPUs of each generation a 70B model needs is worked out in NVIDIA generations.)

Write a budget before choosing the GPU count. Suppose the model has exactly seven billion BF16 parameters and the device exposes 24 GiB. For this calculation, assume 2 GiB for runtime allocations and another 2 GiB of headroom:

If its KV configuration is the earlier 32-layer, eight-KV-head example, each independent 8192-position cache takes 1 GiB. The budget accommodates at most six, not eight. At 16384 positions, each takes 2 GiB and only three fit. Length counts prompt and generated positions already processed by the model; prompt length alone understates the cost of a long answer.

The 2 GiB runtime value is an input to this example, not a constant for every model. Measure with the intended model, batch, prefill chunks, and graph settings. Weight quantization need not quantize KV, and scales, zero-points, and unquantized modules can add storage. The vLLM memory configuration guide is a useful deployment checklist.

Unequal lengths, pages, and shared prefixes

The equal-length formula is convenient, but real requests differ. Round each cache length up to its allocation block, then sum; for how blocks are mapped and why tail blocks have empty slots, see the PagedAttention example. With four-token blocks and lengths 5, 3, and 8, allocation needs 2 + 1 + 2 = 5 blocks: 20 slots for 16 computed positions, leaving four unused slots.

def paged_kv_bytes(lengths, layers=32, kv_heads=8, head_dim=128,
                   bytes_per_value=2, block_tokens=16):
    lengths = list(lengths)
    dimensions = (layers, kv_heads, head_dim, bytes_per_value, block_tokens)
    if any(type(value) is not int or value < 1 for value in dimensions):
        raise ValueError("Dimensions must be positive integers")
    if any(type(length) is not int or length < 0 for length in lengths):
        raise ValueError("Lengths must be nonnegative integers")
    blocks = sum((length + block_tokens - 1) // block_tokens for length in lengths)
    per_token = 2 * layers * kv_heads * head_dim * bytes_per_value
    return blocks * block_tokens * per_token

kv_budget = 24 * 2**30 - 7_000_000_000 * 2 - 4 * 2**30
assert paged_kv_bytes([8192] * 6) <= kv_budget
assert paged_kv_bytes([8192] * 7) > kv_budget
assert paged_kv_bytes([16384] * 3) <= kv_budget
assert paged_kv_bytes([16384] * 4) > kv_budget

Sharing requires counting unique physical blocks, not adding logical request lengths. Two eight-position caches with four-token blocks can share the first block and have separate second blocks: 16 logical positions occupy three blocks, or 12 physical slots. Retained, unused prefixes also occupy the cache pool; they are not memory already returned to the system.

Full-attention state commonly grows with sequence length. Sliding-window layers can retain only the state their windows require, while hybrid models need per-layer-type accounting. Do not apply one uniform $L\times T$ formula to every architecture. The hybrid KV-cache design addresses these distinctions.

Multiple GPUs do not divide every allocation equally

Consider just the earlier 1 GiB cache. Assume even partitioning and omit communication buffers:

A less obvious case is an MQA model with one KV head. In a TP=8 implementation that replicates KV, each rank may retain that head rather than an eighth of it. Check the model implementation and backend; tensor_parallel_size alone does not determine cache placement.

These calculations describe KV, not proportional reductions in every weight or workspace allocation. Deployment feasibility depends on the most constrained rank, not the sum of free memory across devices.

MHA, GQA, MQA: what is shared?

MHA: Q1 → K1,V1    Q2 → K2,V2    Q3 → K3,V3    Q4 → K4,V4
GQA: Q1,Q2 → K1,V1              Q3,Q4 → K2,V2
MQA: Q1,Q2,Q3,Q4 → K1,V1

This sketch uses four query heads. GQA and MQA reduce the number of KV groups; each query head still has its own query and attention output. Smaller caches come with a change to the parameter sharing. Deleting KV heads from a trained MHA model is not automatically quality-preserving. The GQA paper describes and evaluates an uptraining approach from MHA checkpoints.

Return to the 32-query-head example. Moving from 32 to 8 KV heads leaves 8/32 = 1/4 of the raw KV storage, with layers, length, head dimension, and dtype fixed. A fourfold reduction means one quarter remains, not four times as much. It also does not mean total model memory falls by 75%.

KV head count is an architecture choice, not necessarily the GPU count. Tensor-parallel sharding or replication then affects per-device storage. Teaching code may use repeat_interleave to expand shared K/V to the query-head count, but explicitly copying them does not demonstrate the memory saving. An efficient implementation preserves shared storage and supports grouped access in the kernel.

What do FlashAttention and PagedAttention save?

Both often appear next to the KV cache, but they save different things. FlashAttention still computes the same attention; it just avoids writing the full score/probability matrix to GPU memory and reading it back. PagedAttention manages KV storage for many requests in blocks, improving fragmentation and sharing. Neither shrinks K/V itself; that takes architectural changes such as GQA / MQA / MLA. For worked examples and limits of each, see FlashAttention and PagedAttention: what does each save?.

Different results with caching: what should you check?

1. Positions: the fourth position cannot start at zero again. RoPE positions must continue rather than reset at each decode call. The token order of cached K/V must not change either, and every layer must write and read the same set of positions.
2. Masks: three cached positions plus two new inputs require 2 × 5 attention, not 2 × 2. The first new position must not see the second. An upper-left triangle loses valid history; this averaging example shows the resulting error.
3. Prefixes: changing the system prompt, image input, or adapter can invalidate the cache. Identical text under different weights is not an identical computational prefix.
4. Determinism: use eval() to disable module dropout; if calling functional SDPA directly, also pass dropout_p=0.0. Compare logits before sampled text. Different random samples do not establish a cache bug.
5. Lifecycle: release caches when requests end. Reordering batches must reorder sequence identities and their caches together, never attaching one person's history to another's request.

The useful test is not “can it generate a sentence?” Fix an input and compare corresponding logits from a full forward pass against prefill followed by incremental appends, within numerical tolerance. The site's decoder tests already include cache-equivalence checks.

Measure these separately

Hold model, input/output length distributions, concurrency, and hardware fixed. A short-prompt, single-request test is not a substitute for concurrent long conversations. Before optimizing, find out whether time goes into compute, weight movement, KV reads, or queueing (for how batch size shifts this, see the decode roofline example). A newer kernel may not address the actual bottleneck.

Return to attention for the tensor computation, or continue to LoRA and QLoRA for training rather than inference memory.