← Writing

Machine learning Explainer · 2026 · 17 min read

The KV cache and the long-context race

A transformer re-reads everything it has written so far to produce each next token. The KV cache is the trick that makes that affordable, and the reason a long prompt is expensive. Here is the history of that one data structure: where it came from, the papers that shrank it, and how context windows grew from about 2,000 tokens to a million.

Generate text with a transformer and it works one token at a time, attending to every token before it. Do that literally and you redo a mountain of work: at position 1,000 the model recomputes what it already computed at positions 1 through 999. The KV cache is the fix for that waste. A few years on, it is also the main thing standing between you and a million-token prompt. Almost every efficiency paper of the last few years is, underneath, an attack on this one structure.

What the cache actually is

Attention turns each token into three vectors: a query, a key and a value. To produce the next token, its query is compared against the keys of all earlier tokens, and those scores are used to mix the earlier tokens’ values. The point that matters for cost: once a token is in the past, its key and value never change. So you compute them once and keep them. That store of past keys and values is the KV cache, and it is a direct consequence of the attention mechanism from Attention Is All You Need (2017), the paper that dropped recurrence and built the model out of attention alone.

The catch is how the cache grows. It holds one key and one value for every token, in every attention head, in every layer. So it scales with sequence length times model depth times head count, and for a large model at long context it runs to many gigabytes. Worse, every single new token has to stream that whole cache out of memory to attend over it. The 2019 multi-query attention paper, Fast Transformer Decoding, named the real problem: during generation the bottleneck is not arithmetic, it is the memory bandwidth spent reloading those keys and values. That framing set the agenda for everything that followed.

Making attention itself cheap: FlashAttention

Before you shrink the cache, there is the attention computation, which is quadratic in sequence length in both time and memory. FlashAttention (2022) made a sharp observation: the dominant cost was moving data between the GPU’s slow main memory and its fast on-chip memory, not the math. By tiling the computation to stay on-chip and never writing the full attention matrix out, it computes exact attention, no approximation, with memory that grows linearly instead of quadratically, and ran about 3x faster than a standard implementation on GPT-2.

FlashAttention-2 (2023) reworked the parallelism and cut the non-matmul work, reaching around 225 TFLOPs/s on an A100, about 72% of what the hardware can do and roughly double the first version. FlashAttention-3 (2024) rebuilt it again for Hopper H100s, overlapping compute with data movement and adding FP8, up to 740 TFLOPs/s in FP16 and close to 1.2 PFLOPs/s in FP8. None of this changes the numbers the model produces. It just makes long context cheap enough to train and serve.

Shrinking the cache: fewer keys and values

The cache size is tokens times layers times heads times head dimension. Architecture cannot change how many tokens you have, but it can change the rest. The head count went first.

Multi-query attention, from that same 2019 paper, shares a single key and value head across all the query heads. The cache shrinks by the head count, decoding gets much faster, and quality drops only a little. Grouped-query attention (2023) found the middle ground: a handful of key-value groups rather than one or many. It lands close to full multi-head quality at close to multi-query speed, and it can be uptrained from an existing checkpoint with about 5% of the original pretraining compute. That cheapness is why Llama 2 and 3 use it.

DeepSeek took a different angle. Multi-head latent attention in DeepSeek-V2 (2024) does not drop heads; it compresses the keys and values into a small shared low-rank vector and caches that instead. They reported the KV cache down 93.3% and up to 5.76x the generation throughput of their earlier 67B model.

The cache also grows with depth, so the next move was to share it across layers. Cross-layer attention (2024) reuses key and value heads between neighbouring layers for roughly another 2x cut at nearly the same accuracy. YOCO (2024) goes further with a decoder-decoder design that computes the global cache once and lets later layers read it, holding a 1M-token context with near-perfect retrieval.

Serving the cache: paging and reuse

Even a right-sized cache is wasted if the serving system manages memory badly. Requests arrive and finish at different times, each with its own cache that grows as it generates, and plain contiguous allocation leaves the GPU full of unusable gaps. PagedAttention (2023), the idea behind vLLM, borrows virtual memory from operating systems: store the cache in fixed-size blocks that need not sit next to each other, and share identical blocks across requests. Near-zero waste, bigger batches, and 2 to 4x the throughput of the serving stacks before it.

The other easy win is reuse. A system prompt, a set of tool descriptions, a long document you ask five questions about: all of it repeats across calls, and recomputing its attention state every time is pure waste. Prompt Cache (2023) precomputes those reusable chunks once and splices them into new prompts, cutting time-to-first-token by about 8x on GPU and up to 60x on CPU. This is the machinery behind the “prompt caching” you now see in commercial APIs.

Keeping less: windows, sinks, eviction, quantization

The other way to shrink the cache is to not store all of it.

Sliding windows. Only attend to the last N tokens. Longformer (2020) did this for long documents with a local window plus a few global tokens, turning quadratic attention linear. Mistral 7B (2023) shipped sliding-window attention in a general-purpose model, stacking windows across layers so information still travels further than any one window.

Attention sinks. Naive windowing breaks a model the moment the text runs past its training length. StreamingLLM (2023) found the reason: models dump a large share of their attention onto the first few tokens no matter what those tokens say. Keep those few “sink” tokens plus a recent window and the model stays stable over effectively endless streams, with a 22.2x speedup over recomputing from scratch.

Eviction by importance. H2O (2023) noticed that a small set of “heavy hitter” tokens carry most of the attention weight. Keep those plus the recent tokens, drop the rest, and throughput rises up to 29x.

Fewer bits. Store the same tokens at lower precision. KVQuant (2024) gets the cache to about 3 bits with under 0.1 perplexity cost, by quantizing keys per channel and before the rotary embedding, enough to fit a 1M-token context for a 7B model on one 80GB A100. KIVI (2024) is a tuning-free 2-bit scheme, keys per channel and values per token, for 2.6x less peak memory and up to 4x larger batches.

Teaching a model a longer context

Shrinking the cache lets you afford a long context. Position encoding decides whether the model can use it. Most current LLMs use RoPE (2021), which encodes position by rotating the query and key vectors, so attention ends up depending on the distance between two tokens. The problem is that a RoPE model falls apart if you simply feed it sequences longer than it trained on: the extrapolated positions produce runaway attention scores.

Position Interpolation (2023) fixed this by squeezing the position indices back into the trained range rather than extrapolating past it, extending LLaMA to 32,768 tokens with about a thousand fine-tuning steps. YaRN (2023) refined the idea, building on “NTK-aware” interpolation by treating different frequency bands differently and scaling the attention temperature, reaching 128k context with far less training. LongRoPE (2024) searched for the best non-uniform interpolation and pushed past 2 million tokens.

Scaling across machines, and the numbers everyone watched

Some contexts are too long for a single device’s memory no matter how hard you compress. Ring Attention (2023) splits the sequence across a ring of GPUs and overlaps passing key-value blocks around the ring with the attention compute, so the context you can hold scales with the number of devices, with no approximation. Put all of these tricks together and you get the context windows that shipped in real models:

Gemini 1.5 is the high-water mark on paper: its report holds over 99% retrieval out to 10 million tokens across text, video and audio, with a 1M-token production window. Two decisions made that possible: enough of the cache tricks above to fit the context, and enough position work to make the model attend across it.

Leaving the cache behind

The cache exists because attention looks back at every token. Some architectures reject that premise. Linear attention (2020) rewrites attention with kernel feature maps so it costs O(n) instead of O(n²), and at inference it runs like an RNN, carrying a fixed-size state rather than a cache that grows with every token, up to thousands of times faster on very long sequences. The price was quality, which is why linear models lagged for years.

Mamba (2023) closed much of that gap by making the state selective, letting the model decide what to keep or forget based on the input rather than compressing everything the same way. It scales linearly, generates about 5x faster than a comparable transformer, and a 3B Mamba matched transformers twice its size. The frontier now is hybrids that keep a few full-attention layers for precise recall and fill the rest with cheap linear or state-space layers, buying attention’s memory where it earns its keep and constant memory everywhere else.

The through-line is short. Attention made models good and made them expensive. The KV cache made generation affordable and became the new ceiling. Almost everything since has been chipping at that one structure: cache fewer heads, fewer layers, fewer bits, fewer tokens, or refuse to store them at all. Context went from 2,000 tokens to a million in about four years on the back of those moves, and the next order of magnitude will come from the same place.