The KV cache: memory formula, GQA, MQA, MLA and long-context sizing

The KV cache is the per-sequence memory a transformer keeps during generation: for every layer and every token already processed, the key and value vectors that later tokens will attend to. Caching them turns autoregressive decoding from quadratic recomputation into one new token's worth of work per step, which is why every serving engine does it. The price is memory that grows linearly with context length and batch size and, for long contexts, outgrows the model weights themselves. The cache size decides how many users fit on a GPU, how long a context you can offer, and, through the roofline, how fast each token is generated. This article derives the size formula, shows how multi-query, grouped-query and multi-head latent attention change it, sizes real models for long context, verifies a cache implementation against full recomputation, and covers what fine-tuning does to the cache.

Why a cache exists#

In a decoder-only transformer, layer $\ell$ maps the hidden state $h_t$ of token $t$ to a query, key and value per head, $q_t=W_Q{h}_t$, $k_t=W_K{h}_t$, $v_t=W_V{h}_t$, and computes causal attention

Because of the causal mask, $k_s$ and $v_s$ depend only on tokens up to $s$. Once computed they never change, no matter what comes later. Generating token $t+1$ needs $q_{t+1}$ and every earlier $k_s,v_s$; recomputing the earlier ones means running the whole prefix through every layer again. Without a cache, producing $n$ tokens after a prompt of length $p$ processes ${\sum }_{i=0}^{n-1}(p+i)=np+n(n-1)/2$ token positions. With a cache it processes $p+n-1$: the prompt once, then one position per new token. For $p=64$ and $n=1024$, that is 589,312 positions against 1,087, a factor of 542.

Caching keys and values was implicit in the incremental decoding procedure of Vaswani et al. (2017) and is standard in every inference implementation. Its cost became the central problem of LLM serving once models and contexts grew. Shazeer (2019) noted that the repeated loading of cached keys and values makes decoding memory-bandwidth bound and introduced multi-query attention to shrink them. Pope et al. (2022) showed that with multi-query attention and the right partitioning, PaLM 540B could serve contexts up to 32 times longer. Kwon et al. (2023) measured that serving systems of the time used only 20.4 to 38.2% of their cache memory for actual token states, the rest lost to reservation and fragmentation, which motivated PagedAttention.

The memory formula#

Let $L$ be the number of layers, $n_{kv}$ the number of key-value heads per layer, $d_h$ the head dimension and $b$ the bytes per stored element. Each token stores one key and one value vector per key-value head per layer:

A batch of sequences with context lengths $S_1,\dots ,S_B$ needs

and the inequality is the whole capacity problem: the right side is fixed by the GPU and model, so the left side caps concurrency times context. Every decode step also reads all $M_{\text{KV}}$ bytes once. At bandwidth $\beta$, each cached token in the batch adds $\kappa /\beta$ to the time per output token; for Llama 3 8B on an H100 that is $131,072/3.35\times 10^{12}\approx 39$ ns per cached token, so a batch holding 100,000 tokens of context pays about 3.9 ms per step for cache reads alone, comparable to the 4.8 ms for reading the weights.

All Llama 3 and Mistral 7B models use $d_h=128$. The OPT-13B row matches the 800 KB per token Kwon et al. compute in the PagedAttention paper. Note the crossover: the cache for one 128k-token Llama 3 8B sequence (17.2 GB) is larger than the model's own BF16 weights (16.1 GB). For the 8B model the crossover sits at about 122,500 tokens of total cached context, summed over the batch; serving 32 users at 4k tokens each already puts you there.

How MQA, GQA and MLA shrink the cache#

The formula has four factors and the architecture controls two of them. Multi-head attention (MHA) caches a key and value for every query head, $n_{kv}=n_h$. The variants differ only in how many distinct key-value vectors each token keeps per layer:

Multi-query attention (Shazeer, 2019) shares one key head and one value head across all query heads. It divides the cache by $n_h$, a factor of 32 for a model with 32 heads, and Shazeer reported much faster decoding with minor quality loss. Grouped-query attention (Ainslie et al., 2023) interpolates: $n_h$ query heads are split into $n_g$ groups, each sharing one key-value head. GQA-1 is MQA and GQA-$n_h$ is MHA. Ainslie et al. also showed that an existing MHA checkpoint can be converted by mean-pooling the key and value projections within each group and continuing pre-training for about 5% of the original steps, reaching quality close to MHA at speed close to MQA. Llama 3 and Mistral 7B use 8 key-value heads, cutting their caches by 4× (8B) to 16× (405B) relative to MHA at the same width.

Multi-head latent attention (DeepSeek-AI, 2024, in DeepSeek-V2) takes a different route. Instead of caching keys and values, it caches a low-rank latent $c_t$ per token and reconstructs per-head keys and values from it:

with $d_c\ll n_h{d}_h$. At inference the up-projections can be folded into neighbouring matrices: since $q_t^{\top }{k}_s^C=(W^{UQ}{c}_t^Q{)}^{\top }{W}^{UK}{c}_s=c_t^{Q\top }(W^{UQ\top }{W}^{UK})c_s$, the product $W^{UQ\top }{W}^{UK}$ is precomputed and attention runs directly against the cached latents; $W^{UV}$ folds into the output projection in the same way. Rotary position embeddings break this, because a position-dependent rotation would sit between $W^{UQ}$ and $W^{UK}$, so MLA adds a small decoupled key $k_t^R\in {\mathbb{R}}^{d_h^R}$ that carries the rotary part and is cached alongside $c_t$. DeepSeek-V2 uses $n_h=128$, $d_h=128$, $d_c=512=4d_h$ and $d_h^R=64=d_h/2$, so each token caches $576$ elements per layer, which the paper notes equals GQA with 2.25 groups, while reporting quality stronger than MHA. Against an MHA model of the same shape, which would cache $2\times 128\times 128=32,768$ elements per layer, that is 1.8% of the size. Against its predecessor DeepSeek 67B the paper reports a 93.3% smaller cache and 5.76× higher maximum generation throughput.

Sliding-window attention bounds $S$ instead of $\kappa$. Mistral 7B attends only to the previous $W=4,096$ tokens in each layer and stores the cache in a rolling buffer, writing position $i$ to slot $i\mathrm{mod}W$. The cache stops growing at $W$ tokens, which the Mistral paper reports cuts cache memory by 8× on a 32k-token sequence. Information from further back still reaches the output through the stack of layers, up to $W$ times the number of layers in theory. The trade is that exact retrieval from beyond the window is no longer guaranteed. The KV cache compression article covers the inference-time versions of this idea, which evict tokens from a model that was trained with full attention.

Worked example: sizing a long-context deployment#

Suppose you fine-tuned Llama 3 8B for document question answering and want to serve 32k-token contexts on one 80 GB H100.

Per token. $\kappa =2\times 32\times 8\times 128\times 2=131,072$ bytes = 128 KiB.

Per sequence. $32,768\times 131,072=4.29$ GB.

Budget. 80 GB minus 16.1 GB of weights minus a few GB for activations, CUDA graphs and the runtime leaves roughly 60 GB, so at most $\lfloor 60/4.29\rfloor =13$ full-length sequences at once.

Throughput consequence. Each decode step reads 16.1 GB of weights plus up to $13\times 4.29=55.8$ GB of cache. At 3.35 TB/s that is at least 21.5 ms per step, so at most $13/0.0215\approx 605$ output tokens per second across all users, and about 47 tokens per second per user. At 2k-token contexts the same GPU holds over 200 sequences and can exceed 8,000 tokens per second (worked through in the roofline article). Long context is roughly an order of magnitude more expensive per output token on the same hardware.

Levers. Storing the cache in FP8 halves $\kappa$ to 64 KiB, fitting 27 sequences and roughly halving the cache term of each step. Moving to the 70B model raises $\kappa$ to 320 KiB: 10.7 GB per 32k sequence, and 141 GB of weights that need at least two GPUs before any cache is allocated. A model with MLA at DeepSeek-V2's ratio would need 2.3 GB per 32k sequence despite having 60 layers and 128 heads.

A toy run: cached decoding is exact#

A cache is an optimisation that must not change the output. We checked that directly with a small decoder in PyTorch: 4 layers, width 512, 8 query heads sharing 2 key-value heads (GQA with group size 4), random weights, seed 0, FP32 on CPU. From a 64-token random prompt we generated greedily twice, once by re-running the full sequence every step and once with a cache that stores each layer's keys and values and appends one row per step. For 64, 128, 256, 512 and 1,024 generated tokens the two runs produced identical token sequences, and the largest absolute difference in any logit at any step was $3.6\times 10^{-6}$ to $5.0\times 10^{-6}$, floating-point reordering noise. The cache held $2\times 4\times 2\times 64\times 4=4,096$ bytes per token, so 4,452,352 bytes after 1,087 positions, exactly the formula's prediction.

The work saved is deterministic and grows linearly with output length: the uncached run pushed 589,312 token positions through the model for 1,024 new tokens, the cached run 1,087. We also timed both, but the machine was shared with other jobs and the wall-clock speedups (3.6× to 9.6×) were too noisy to be worth more than confirming the direction, so the chart reports the exact work counts instead.

What fine-tuning does to the cache#

Adapters on key and value projections change the cache. A LoRA adapter on $W_K$ or $W_V$ changes every cached key or value, and one on any earlier layer changes the hidden states that produce them. Cached entries computed under one adapter are therefore invalid for another, even for an identical prompt. Engines that reuse cache across requests must key it by adapter: vLLM's prefix-cache block hash includes the LoRA ID for this reason. In a multi-LoRA deployment, a shared system prompt is cached once per adapter, not once overall. Adapters on only the query projection or the MLP of the last layers would, strictly, leave earlier caches valid, but no common engine exploits that and it is rarely worth designing around.

Merged weights behave like the base. A merged fine-tune has the same $\kappa$, and its cache is valid only for itself.

The architecture is fixed at fine-tuning time. You cannot turn an MHA model into a GQA model with LoRA. Ainslie et al.'s conversion needs mean-pooled projections plus continued pre-training at about 5% of the original budget. If serving memory matters, choose a GQA or MLA base before post-training, because switching later means redoing the whole post-training pipeline.

Long-context fine-tuning raises the bill linearly. Extending a model from 8k to 128k tokens with rotary scaling and long-document SFT leaves $\kappa$ unchanged but multiplies the worst-case $S$ by 16. Set the engine's maximum context to what the product needs, not to what the model supports.

Quantised caches need the fine-tuned model's statistics. FP8 or INT4 caches use scales calibrated on activations. Fine-tuning shifts key and value distributions, especially for outlier channels, so scales calibrated on the base model can clip the fine-tune's values. Recalibrate on the fine-tuned model's data, as discussed in KV cache compression.

Templates and reasoning traces are cache too. Every token of chat template, tool schema and system prompt sits in the cache for the whole request. A reasoning model trained with GRPO that thinks for 8,000 tokens holds an 8,000-token cache even if its visible answer is 200 tokens.

Engine knobs that decide cache capacity#

One subtlety with tensor parallelism: key-value heads are split across GPUs, so with 8 key-value heads and 16-way parallelism each head must be replicated, and the per-GPU cache stops shrinking beyond 8 GPUs. With MQA this happens at 2.

Failure modes and how to spot them#

Preemption storms. When the cache fills, engines preempt running sequences, either swapping their blocks to CPU memory or discarding and recomputing them later. Kwon et al. found recomputation never cost more than 20% of swapping's latency in their measurements. Symptom: throughput collapses under load, logs report preempted requests, and tail latency rises sharply. Fix: lower concurrency, lower maximum length, quantise the cache, or add GPUs.

Capacity planned at average context. Memory needs follow the sum of live contexts, which at peak can be several times the average. Size for the 95th percentile of total live tokens, not mean prompt length times mean concurrency.

Over-reservation. Systems that reserve a contiguous region for the maximum length per request waste most of it; this is the 20.4 to 38.2% utilisation Kwon et al. measured. Symptom: few concurrent sequences although actual contexts are short. Paged allocation fixes it.

Stale or mismatched cache. Reusing cache entries across a weight update, a different adapter, a different tokenizer or a changed chat template produces silently wrong outputs rather than errors. Treat model version, adapter ID and template as part of any cache key.

Quantised-cache drift. FP8 or lower caches that pass perplexity checks can still degrade long-range retrieval. Evaluate needle-in-a-haystack style retrieval at the context lengths you serve, on the fine-tuned model.

Choosing a strategy#

Minimal implementation#

A production cache is paged and fused into the attention kernel, but the logic is a preallocated buffer per layer and a write pointer. This version supports GQA and mirrors the toy run above.

import torch, math

class KVCache:
    def __init__(self, layers, batch, max_len, kv_heads, head_dim, dtype=torch.bfloat16, device="cuda"):
        shape = (layers, 2, batch, max_len, kv_heads, head_dim)
        self.buf = torch.empty(shape, dtype=dtype, device=device)
        self.len = 0
    def bytes_per_token(self):
        L, _, B, _, H, D = self.buf.shape
        return 2 * L * H * D * self.buf.element_size()       # = kappa
    def append(self, layer, k, v):                            # k, v: (batch, new, kv_heads, head_dim)
        n = k.shape[1]
        self.buf[layer, 0, :, self.len:self.len + n] = k
        self.buf[layer, 1, :, self.len:self.len + n] = v
        return self.buf[layer, 0, :, :self.len + n], self.buf[layer, 1, :, :self.len + n]
    def advance(self, n):                                     # call once per forward pass, after all layers
        self.len += n

def gqa_attention(q, k, v, pos0):
    # q: (B, T, Hq, D); k, v: (B, S, Hkv, D); query t sits at absolute position pos0 + t
    B, T, Hq, D = q.shape; S, Hkv = k.shape[1], k.shape[2]
    k = k.repeat_interleave(Hq // Hkv, dim=2); v = v.repeat_interleave(Hq // Hkv, dim=2)
    scores = torch.einsum("bthd,bshd->bhts", q, k) / math.sqrt(D)
    mask = torch.arange(S, device=q.device)[None, :] > (pos0 + torch.arange(T, device=q.device))[:, None]
    scores = scores.masked_fill(mask, float("-inf"))
    return torch.einsum("bhts,bshd->bthd", scores.softmax(-1), v)

# prefill: append the whole prompt (T = prompt length, pos0 = 0)
# decode:  append one token per step  (T = 1, pos0 = cache.len before the step)

In Hugging Face Transformers the same object is the past_key_values cache returned by the model and passed back each step; in serving engines such as vLLM, SGLang and TensorRT-LLM it is a pool of fixed-size blocks managed by the scheduler. Kernels never materialise the repeated heads: GQA-aware attention kernels index the shared key-value head directly, which is where the bandwidth saving actually comes from. In POFE the per-token cache size of the chosen base is shown next to its quality numbers, because it fixes long-context serving cost before a single training step is run.

Further reading#

Vaswani et al. (2017) defines the attention being cached. Shazeer (2019) is the first paper to treat cache bandwidth as the decode bottleneck and introduces MQA; Ainslie et al. (2023) generalises it to GQA and shows how to convert existing checkpoints. Pope et al. (2022) analyses cache memory and partitioning at scale. The DeepSeek-V2 report introduces MLA with its absorption trick and decoupled rotary key. The Mistral 7B report describes sliding-window attention with a rolling buffer. Kwon et al. (2023) measures how badly naive allocation wastes cache memory and is the starting point for the PagedAttention article. The Llama 3 report gives the shapes used in the tables here.

Frequently asked questions#

References#

1. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin (2017). Attention Is All You Need. arxiv.org/abs/1706.03762
2. Noam Shazeer (2019). Fast Transformer Decoding: One Write-Head is All You Need. arxiv.org/abs/1911.02150
3. Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, Sumit Sanghai (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arxiv.org/abs/2305.13245
4. DeepSeek-AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arxiv.org/abs/2405.04434
5. Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, Jeff Dean (2022). Efficiently Scaling Transformer Inference. arxiv.org/abs/2211.05102
6. Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, Ion Stoica (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arxiv.org/abs/2309.06180
7. Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, et al. (2023). Mistral 7B. arxiv.org/abs/2310.06825
8. Llama Team, AI @ Meta (Aaron Grattafiori, Abhimanyu Dubey, et al.) (2024). The Llama 3 Herd of Models. arxiv.org/abs/2407.21783
9. vLLM project (2025). vLLM design documentation: Automatic Prefix Caching. docs.vllm.ai/en/latest/design/prefix_caching.html

Related reading