The KV cache

Problem

A decoder-only model generates by running itself over and over, appending the chosen token each time (see decoding). Naively, step t runs the whole network on all t tokens, so writing n tokens costs about n²/2 token-positions of work. On the page's four-token example that is 10 key/value computations where only 4 were really new.

The waste is avoidable because of causality: the keys and values of token 3 depend only on tokens 1 to 3, so they never change when token 4 arrives. Keep them and the cost of a step no longer grows with the work already done. The price is memory, and for long conversations that price is large.

Intuition

Imagine taking notes in a meeting. To answer a question about what was said, you do not replay the recording from the start; you look at your notes. Each attention layer keeps two notes per past token: a key (what this token is about, used for matching) and a value (what it contributes if matched). A new token writes its own two notes, then reads all the notes to decide what to attend to. Nothing already written is edited.

Generation therefore has two phases. Prefill reads the whole prompt at once and fills the cache, with many tokens processed in parallel, which keeps the hardware busy. Decode then adds one token at a time, and each step must read the entire cache but does little arithmetic with it. That makes decode limited by memory bandwidth, not by compute.

Prefill is compute-heavy and parallel over the prompt; decode is a sequential loop over the cache.

Mechanics

What is cached, and why only K and V

In each layer, token i produces a query q_i, key k_i and value v_i (see attention). Its output is softmax(q_i·K_≤i^T/√d) V_≤i. For a new token only its own query is needed: past queries were used once, for the past tokens' outputs, and are never read again. Past keys and values are read by every later token, so they are the things to store. The causal mask is also free: the cache only contains earlier positions, so nothing needs masking at decode time. The cache has one slot per layer, per token, per K/V head: shape [layers, 2, B, Hkv, S, D] with S the length so far.

How big the cache is

Per token the cache holds 2 (K and V) × layers × H_kv × head width × bytes per value. For an 8B-class configuration with 32 layers, 32 query heads, 8 key/value heads, head width 128 and 2-byte values (fp16 or bf16) that is 131,072 bytes, 128 KiB, per token. An 8,192-token context costs 1 GiB for one sequence, and a batch of 32 such sequences costs 32 GiB, comparable to the weights of the model itself.

Multi-query and grouped-query attention

Cache size is proportional to the number of key/value heads, so share them. Multi-query attention (MQA) keeps all the query heads but a single K/V head. Grouped-query attention (GQA) is the middle: query heads are split into groups, each group sharing one K/V head. With 32 query heads and 8 K/V heads, every K/V head serves 4 queries. In PyTorch, F.scaled_dot_product_attention(..., enable_gqa=True) does this by repeating each key and value head to match the query heads (lines 6442-6444 of the reference code in its docstring), so the saving is in what you store, not in the attention arithmetic.

Sharing the cache across layers, and storing it in 4 bits

The per-token size is a product, 2 × layers × H_kv × head width × bits, and GQA attacks only one factor. The others can shrink too. Cross-layer attention (CLA, Brandon et al. 2024) lets a group of adjacent layers share one set of keys and values: only the first layer in the group projects and caches K and V, and the others attend to them with their own queries. Sharing across pairs halves the cache again, on top of GQA. You Only Cache Once (YOCO, Sun et al. 2024) takes it to the limit: the bottom half of the network is a self-decoder with cheap local attention, and the middle layer's output is projected once into a single global K/V that every upper layer cross-attends to. The upper layers no longer need the prompt run through them to build their own caches, so prefill does roughly half the work as well.

DeepSeek-V4.1-Flash (2026) combines all three axes. Its report frames them as entry size (GQA, latent attention), sequence (keeping one compressed entry per m tokens) and layer (cache reuse). It splits its 40 layers into a causal encoder-decoder after YOCO: the 20 decoder layers project their global K/V from the hidden state at layer 20, so a prompt only passes through the encoder half to fill them. Inside each half, most layers reuse a neighbour's work: a Full layer computes its own compressed K/V, sparse-attention indexer keys and top-k selection, a Reindex layer keeps the K/V but picks its own top-k, and a Reuse layer takes both. Finally the main cache is stored in FP4: E2M1 values (sign, 2 exponent bits, 1 mantissa bit, so magnitudes 0 to 6) with one E4M3 scale shared by every 16 channels, 4.5 bits per value. The values are quantized after rotary embeddings and the model is trained for it with quantization-aware training; the small sliding-window cache stays in FP8. The report gives about 890 bytes of global K/V per token, a quarter of its predecessor's.

Paged attention

A serving system holds many sequences of different, unknown final lengths. Reserving one contiguous region of the maximum length for each wastes most of it, and the leftovers fragment memory. Paged attention (the idea behind vLLM) stores each sequence's cache in fixed-size blocks of tokens that need not be adjacent, with a per-sequence block table mapping positions to blocks, as an operating system pages virtual memory. A sequence wastes at most the unused tail of its last block, and blocks can be shared between sequences with a common prefix.

Worked example

Implementation

import torch
import torch.nn.functional as F

class CachedAttention(torch.nn.Module):
    def __init__(self, d, q_heads, kv_heads):
        super().__init__()
        self.hq, self.hkv, self.dh = q_heads, kv_heads, d // q_heads
        self.q = torch.nn.Linear(d, q_heads * self.dh, bias=False)
        self.k = torch.nn.Linear(d, kv_heads * self.dh, bias=False)
        self.v = torch.nn.Linear(d, kv_heads * self.dh, bias=False)
        self.o = torch.nn.Linear(d, d, bias=False)

    def forward(self, x, cache=None):          # x: [B, T, d]; T = prompt length, then 1
        B, T, _ = x.shape
        q = self.q(x).view(B, T, self.hq, self.dh).transpose(1, 2)
        k = self.k(x).view(B, T, self.hkv, self.dh).transpose(1, 2)
        v = self.v(x).view(B, T, self.hkv, self.dh).transpose(1, 2)
        if cache is not None:                  # append, never rewrite
            k = torch.cat([cache[0], k], dim=2)
            v = torch.cat([cache[1], v], dim=2)
        out = F.scaled_dot_product_attention(
            q, k, v, is_causal=(cache is None),   # prefill: causal; decode: T == 1
            enable_gqa=self.hq != self.hkv,
        )
        out = out.transpose(1, 2).reshape(B, T, -1)
        return self.o(out), (k, v)

Prefill calls it with the whole prompt and no cache, then each decode step passes the previous (k, v) back with one new token. torch.cat copies the whole cache every step; real systems preallocate a buffer and write into it (or use blocks, as above). Rotary positions, if the model uses them, are applied to q and k before k is cached, so the cache holds already-rotated keys. A decode step with T = 1 needs no mask, because the cache holds only earlier positions.

Tradeoffs

• Speed versus memory. The cache turns quadratic recompute into linear work but its size is linear in context and batch, and often decides how many requests fit on a GPU.
• MQA and GQA shrink the cache by H_q/H_kv at some loss in quality that grows as the K/V heads shrink; GQA is the usual compromise. A model must be trained with the choice (or converted and fine-tuned), since the projections differ.
• Precision. Storing the cache in 8-bit or lower cuts memory again, trading accuracy; see quantization. At 4 bits the per-block scales are essential and the model generally has to be trained for it.
• Sharing across layers (CLA, YOCO-style encoder-decoders) divides the cache by the sharing factor and, for YOCO, roughly halves prefill, but the layers that read a shared cache lose their own view of the prefix, which costs quality, and the architecture must be chosen before pre-training.
• Paging improves utilisation and allows prefix sharing, at the cost of a more complex attention kernel that follows the block table.
• Pitfalls. Passing is_causal=True with a one-token query and a longer key sequence aligns the mask to the wrong corner (the attention page covers the two alignments). Off-by-one position ids when resuming from a cache silently corrupt outputs. And a cache is tied to the exact prefix: change an earlier token and every later entry is stale.

Check yourself

Connections

References

• Shazeer. Fast Transformer Decoding: One Write-Head is All You Need (multi-query attention). 2019 (arXiv 1911.02150).
• Ainslie et al. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP, 2023 (arXiv 2305.13245).
• Kwon et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP, 2023 (arXiv 2309.06180).
• Brandon, Mishra, Nrusimha, Panda and Ragan-Kelley. Reducing Transformer Key-Value Cache Size with Cross-Layer Attention. NeurIPS, 2024 (arXiv 2405.12981).
• Sun, Dong, Zhu, Huang, Wang, Ma, Zhang, Wang and Wei. You Only Cache Once: Decoder-Decoder Architectures for Language Models. NeurIPS, 2024 (arXiv 2405.05254).
• DeepSeek-AI. DeepSeek-V4.1-Flash: Pushing the Limits of KV Cache Compression. 2026 (technical report).
• Vaswani et al. Attention Is All You Need. NeurIPS, 2017 (arXiv 1706.03762).
• PyTorch v2.15.0: scaled_dot_product_attention, MultiheadAttention.