II · THE IDEA · ARTIFICIAL INTELLIGENCE
Why Long Context Eats Your RAM
▶ Listen · narrated
The context window keeps growing, but your GPU keeps running out of memory long before you hit the advertised limit. The culprit is not the model weights.
At a glance
- What it stores
- One key vector and one value vector for every token, at every layer
- Llama 2 70B at 32k
- Roughly 160 GB in float16, more than twice the weight footprint
- Why it exists
- Avoids recomputing attention scores from scratch at every generation step
- Growth rate
- Linear in context length, linear in layer count, linear in hidden dimension
Imagine a model generating a story one word at a time. To choose the next word, it must look back at every word written so far and decide how much attention to pay to each one. If it had to re-read and re-process the entire story from the beginning at every step, writing a thousand-word story would mean reading the story a thousand times, and the cost would spiral out of control. Instead, the model keeps a cheat sheet: for every word already written, it stores two lists of numbers that summarise what that word contributed. These lists are called keys and values, and together they form the key-value cache. When generating the next word, the model consults the cache instead of re-reading everything, which makes generation fast. But the cache itself takes up space. Each word adds a fixed amount of data to the cache, so a longer story means a bigger cache. For a large model handling a very long document, the cache can grow so large that it fills up the computer's fast memory, even though the model itself would have fit comfortably. The cache is what makes long context possible, but it is also what makes long context expensive.
The key-value cache stores the output of the key and value projections in each transformer layer for every token in the context. During generation, the model computes a query vector for the new token, retrieves the cached keys for all previous tokens, computes attention scores as the dot product of the query with each key, applies softmax to get attention weights, and finally computes the output as a weighted sum of the cached values. This avoids recomputing keys and values for tokens that have already been processed, reducing generation cost from O(n²) to O(n) in the number of output tokens. Cache size is 2 × L × T × D × B, where L is the number of layers, T is the number of tokens, D is the hidden dimension, and B is bytes per parameter (2 for float16, 1 for int8). For a model with 32 layers, hidden dimension 4096, and float16 precision, each token costs 524 KB. At 32k tokens that is 16 GB; at 128k it is 64 GB. Grouped-query attention reduces D by a factor equal to the number of heads per group, because key-value pairs are shared across heads within a group rather than computed independently. The cache must reside in GPU memory during generation because attention requires reading the entire cache at every step, and PCIe bandwidth is two orders of magnitude lower than GPU memory bandwidth. Quantising the cache to int8 or lower precision halves memory use but introduces approximation error into every attention computation, which can degrade output quality, particularly at very long context where small errors accumulate. PagedAttention manages the cache in fixed-size blocks and allows non-contiguous allocation, which reduces fragmentation and enables prefix sharing between requests that begin with identical token sequences.
Look closer
The cache size formula is straightforward but unforgiving
For each token, each layer stores two vectors: the key and the value produced by that layer's attention mechanism. Each vector has dimension equal to the model's hidden size. In float16 precision, that is two bytes per number. So the formula is: 2 (key and value) × layers × tokens × hidden_dimension × 2 bytes. For Llama 2 7B with 32 layers and hidden dimension 4096, a single token costs 2 × 32 × 4096 × 2 = 524,288 bytes, or half a megabyte. An 8,192-token context therefore needs 4 GB. At 32,768 tokens it is 16 GB. At 131,072 tokens it is 64 GB — and that is before the model weights, the activations, or the optimizer state.
Multi-query and grouped-query attention reduce the cache, not the computation
Standard multi-head attention gives every attention head its own key and value projection, so a 32-head model stores 32 separate key vectors per token per layer. Multi-query attention (MQA) shares a single key and value across all heads, cutting cache size by the number of heads. Grouped-query attention (GQA), introduced by Ainslie and colleagues in 2023, is a compromise: heads are divided into groups, each group sharing one key-value pair. Llama 2 70B uses GQA with 8 groups across 64 heads, reducing cache size by a factor of eight compared to full multi-head. The attention computation itself remains almost unchanged; this is purely a memory optimization, and it works because the value of having fully independent keys per head is smaller than the engineering cost of storing them all.
The cache must stay in fast memory during generation
Every new token requires attending to every cached token, which means reading the entire cache once per generation step. If the cache lives in system RAM rather than GPU memory, that read happens over the PCIe bus — typically 16 or 32 GB/s, compared to GPU memory bandwidth in the terabytes per second. The result is that generation becomes bottlenecked by data transfer rather than computation. This is why long-context serving on consumer hardware often means choosing between a smaller model that fits entirely in VRAM, or a larger model with most requests truncated well below the advertised context limit. Offloading the cache to disk is technically possible but renders generation so slow as to be impractical for interactive use.
The story
When a transformer generates text, it does so one token at a time. At each step, the new token attends to every token that came before it — the prompt and all previously generated tokens. Attention requires computing a score between the new token's query vector and the key vector of every earlier token, then using those scores to mix their value vectors. If the model had to recompute every key and value from scratch at every step, generation would require passing the entire context through every layer every time, and a thousand-token generation would mean a thousand full forward passes through the model. The cost would be quadratic in output length.
The key-value cache solves this by storing the key and value vectors produced at each layer for each token that has already been processed. When generating token number 500, the model does not recompute the keys and values for tokens 1 through 499. It retrieves them from the cache, computes only the new token's query, key and value, appends the new key and value to the cache, and proceeds. This turns quadratic cost into linear cost, which is why generation is feasible at all.
But the cache is not free. Each token contributes two vectors per layer, and those vectors are the same width as the model's hidden dimension — 4096 for Llama 2 7B, 8192 for Llama 2 70B. In float16, each number occupies two bytes. For Llama 2 7B with 32 layers, one token adds 2 × 32 × 4096 × 2 bytes to the cache: half a megabyte. That sounds modest until you multiply by context length. An 8,192-token context needs 4 gigabytes. A 32,768-token context needs 16 gigabytes. A 131,072-token context needs 64 gigabytes, and the model weights themselves are only 13 gigabytes in float16.
For Llama 2 70B, the situation is worse. With 80 layers and hidden dimension 8192, each token costs 2 × 80 × 8192 × 2 = 2,621,440 bytes, or 2.5 megabytes. An 8k context is 20 GB. A 32k context is 80 GB. A 128k context is 320 GB. The weights are 140 GB, so at 128k tokens the cache is more than twice the size of the model. On a node with 80 GB of GPU memory, you can load the weights and serve perhaps 24,000 tokens of context before memory is exhausted — and that leaves no room for batching multiple requests, which is how production serving achieves acceptable cost per token.
This is why grouped-query attention matters. By sharing key-value pairs across groups of attention heads rather than giving each head its own, GQA cuts cache size by the number of heads per group. Llama 2 70B uses 8 groups for 64 heads, an eightfold reduction. That brings the per-token cost down from 2.5 MB to 320 KB, and a 32k context from 80 GB to 10 GB. The model becomes servable on hardware that would otherwise be unable to approach the advertised context limit.
PagedAttention, introduced by Kwon and colleagues, goes further by managing the cache in small blocks rather than contiguous arrays, allowing the system to allocate and free memory at fine granularity and to share cached prefixes between requests. But the fundamental constraint remains: the cache grows linearly with context length, and at some point it will dominate memory use no matter how cleverly it is managed. Long context is not free. It is a trade-off between capability and cost, and the cost is paid in gigabytes of fast memory that must be read in full at every generation step.
Why it mattered then
The key-value cache has been part of transformer inference since the architecture was introduced, but it became a bottleneck only when context windows began to grow beyond a few thousand tokens. Early models like GPT-2 had a 1,024-token limit, where cache size was a small fraction of overall memory use. As models scaled to tens of billions of parameters and context windows reached 8,192, 32,768 and eventually 100,000 tokens or more, the cache crossed a threshold: it stopped being an implementation detail and became the dominant cost of serving. The shift from multi-head to grouped-query attention in models like Llama 2 was a direct response to this pressure. Without it, long-context models would have been unservable on the hardware available to most researchers and companies, and the advertised context limits would have been theoretical rather than practical.
Why it matters now
Long context is now a headline feature. Every major model release advertises a larger window, and users expect to be able to paste entire codebases, legal documents or transcripts into a single prompt. But the memory cost of the cache means that the effective context limit on real hardware is often much lower than the number in the marketing material. A 128k-token model may truncate most requests to 32k or less, not because the model cannot handle longer sequences, but because serving them would require memory that is not available or would make batching impossible. Understanding cache size is therefore essential for capacity planning, cost estimation and choosing which models to deploy. It is also why quantisation and other memory-reduction techniques are applied not just to weights but to the cache itself, even though doing so introduces error into every attention computation. The trade-off is considered worth it, because the alternative is not serving long context at all.
The surprising detail
PagedAttention allows multiple requests to share the same cached prefix if they begin with identical tokens, which means that a system prompt or few-shot examples that appear in every request can be computed once and reused across thousands of conversations. The memory saving is proportional to the length of the shared prefix and the number of concurrent requests, and in production serving it can be substantial. But it depends on exact token-level matching: a single changed character invalidates the entire shared block, because the tokeniser may split the text differently and the token ids no longer align. This makes prefix sharing fragile in ways that are not always obvious. A trailing space, a different tokeniser version or a change in prompt template can silently disable the optimisation, and the system will allocate separate cache blocks for what a human would consider the same text.
Remember this
The cache grows linearly with context length and layer count, and at long context it will use more memory than the model weights. This is not a bug. It is the cost of making generation linear instead of quadratic.
Test yourself
You are serving a 7B model with 32 layers and hidden dimension 4096, using float16 precision. You have 24 GB of GPU memory. Roughly 14 GB is occupied by the model weights and overhead. How long a context can you serve for a single request before the key-value cache alone exhausts the remaining 10 GB? Show the arithmetic.
Each token adds 2 (key and value) × 32 layers × 4096 dimension × 2 bytes = 524,288 bytes, or 0.5 MB. With 10 GB available, that is 10,000 MB ÷ 0.5 MB per token = 20,000 tokens. So the practical limit is around 20k tokens for a single request, well below the 32k or 128k that the model might theoretically support. If you want to batch multiple requests or leave headroom for activation memory, the effective limit drops further. This is why consumer hardware often truncates context aggressively despite the model being advertised as long-context capable.
Go deeper
- Efficient Memory Management for Large Language Model Serving with PagedAttention · arXiv · Woosuk Kwon et al. · 2023-09-12
- GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints · arXiv · Joshua Ainslie et al. · 2023-05-22
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.