Reference — the answer first

What is the KV cache?

Direct answer

A Transformer language model generates text one token at a time. At each step it keeps a key and a value for every earlier token, in every attention layer and every key-value head. That store is the KV cache. It exists so the model does not recompute the keys and values of earlier tokens at every step. Its size grows linearly with the number of tokens stored and with the number of key-value heads each layer keeps. Grouped-query attention (GQA) and multi-query attention (MQA) shrink it by letting several query heads share the same stored key-value heads. With long contexts and large batches, the cache can take more memory than the model weights themselves.

The memory formula

Plain text: Mem_KV = B * N_layers * T * H_kv * d_head * 2 * bytes

B
batch — sequences decoded together
N_layers
model layers that each store keys and values
T
earlier token positions kept in the cache
H_kv
stored key-value heads per layer — the term head sharing changes
d_head
dimensions per head
2
one tensor for keys, one for values
bytes
bytes per stored number, such as 2 for fp16/bf16

This counts the bytes of the keys and values themselves. A real serving system also spends memory on allocation overhead, paging metadata and padding, and changing the attention architecture can change model quality; the formula covers neither.

A worked example

Take one configuration, the same defaults as the KV memory step in the Lab: batch 1, 32 layers, 4,096 stored positions, 32 query heads, head dimension 128, fp16/bf16 (2 bytes per number). All results are in decimal GB (10⁹ bytes), computed with the formula above.

  • Full multi-head attention stores H_kv = 32 key-value heads: 2.15 GB.
  • GQA with H_kv = 8 at the same context: 0.54 GB — 4.0× smaller, because 8/32 of the heads remain.
  • The same GQA cache at 32,768 positions: 4.29 GB, 2.0× the full multi-head cache at 4,096 positions, because multiplying the context by 8 outweighs dividing the stored heads by 4.

With B, N_layers, d_head and precision fixed, KV-cache memory is proportional to T × H_kv. Sharing heads makes the cache narrower; a longer context makes it longer.

MHA vs GQA vs MQA

Cache memory for the worked configuration at T = 4,096, with 32 query heads throughout.
VariantStored K/V heads (H_kv)Cache memoryvs MHA
MHA — no sharing322.15 GB1.0×
GQA — grouped sharing80.54 GB4.0× smaller
MQA — one shared head10.07 GB32.0× smaller

The formula gives cache size only. Quality and speed have to be measured: MQA can lower quality, and Ainslie et al. (2023) report that GQA models converted from multi-head checkpoints come close to multi-head quality while decoding nearly as fast as MQA.

Three common misconceptions

  1. “Sharing shrinks the query heads.” GQA keeps every query head — 32 stay 32 in the worked model. What shrinks is the number of stored key-value heads (H_kv), from 32 to 8. Query-head count and stored-head count are different quantities.
  2. “Sharing forgets earlier words.” The same earlier positions stay in the cache — T does not change. Each position simply stores fewer distinct key-value sets for the heads to read.
  3. “A smaller cache proves faster, better answers.” The formula estimates stored bytes only. Allocator behavior, paging, batching, scheduling, and answer quality sit outside it; speed and quality claims need separate evidence.

Equation to code

The same calculation in Python, term for term, as used in the efficient-attention notebook:

def kv_cache_gb(
    batch=1,
    layers=32,
    tokens=4096,
    kv_heads=32,
    head_dim=128,
    bytes_per_value=2,
):
    scalars = batch * layers * tokens * kv_heads * head_dim * 2
    return scalars * bytes_per_value / 1e9
Every equation factor and its variable in the code.
Equation factorCode termMeaning
Bbatchsequences decoded together
N_layerslayerslayers that store keys and values
Ttokensearlier positions kept in the cache
H_kvkv_headsstored key-value heads per layer
d_headhead_dimdimensions per head
2literal 2one key tensor plus one value tensor
bytesbytes_per_valuebytes per stored number (2 for fp16/bf16)
÷ 1e9/ 1e9decimal gigabytes, not GiB

Sources

Test your understanding

To test yourself, answer the two questions in the KV-cache challenge, or open the KV memory step in the Lab to change one input at a time and watch this formula respond. Further reading: grouped-query attention and efficient attention at scale.