PagedAttention
A KV cache in fixed-size blocks: how much memory contiguous allocation wastes, what block tables, reference counts and copy-on-write buy, and what greedy admission costs.
A server with room for 32,768 tokens of KV cache can hold exactly 32 sequences if it sets aside 1,024 tokens for each. In my simulation of a plausible chat workload, those 32 sequences used 33 per cent of the memory they had been given. The other two thirds sat reserved for tokens that, on average, would never be generated.
Give the same budget to an allocator that hands out 16-token blocks as sequences grow, and the server holds between 92 and 96 sequences at a time, using 98 per cent of the memory, and finishes the same 3,000 requests in 2.6 to 2.7 times fewer steps.
Then let it admit requests greedily, as fast as blocks come free, and it thrashes. Memory runs out mid-sequence, the newest request is thrown away and recomputed, and the run suffers 2,094 preemptions for 3,000 requests. The same allocator that was 2.7 times better is now, once recomputation is charged at four steps per preemption, only 10 per cent better than reserving the maximum and clearly worse than reserving the declared limit (13,754 steps against 15,260 and 10,637).
Paging is a good idea. It is not a free one, and the part that decides whether it works is the scheduler’s policy and not the page table. This chapter builds a real paged cache for GPT-2, measures what it costs when attention has to gather its blocks, and then uses a simulation to show where the memory goes and where the policy bites.
Where this sits
Chapter 10 built a cache that grew by concatenation and found that, for eight sequences generating to 960 tokens, the allocator held 2.9 GiB to store a cache of 540 MiB. It also found that a preallocated buffer fixed the fragmentation by paying for the maximum up front. Chapter 8 computed that at a batch of about seven the cache, not the weights, dominates the memory traffic of a step. Chapter 9 introduced a scheduler that interleaves work from different requests. Paging is the answer to the question those chapters left open: how should a set of caches of unpredictable, growing sizes share one pool?
The Inference Lab piece added here is the block layer: a reference-counted BlockAllocator, BlockTables with copy-on-write and fork, a real PagedKV cache for GPT-2 with a ragged-batch decode, and an allocation simulator with three policies. The rule it adds to the cache rule of Chapter 10 is that an allocator’s quality is measured in sequences held, not bytes saved: capacity is what the scheduler can spend.
For the Pi series, the relevant fact is that an agent’s context is long and variable, and a shared prefix, such as a system prompt, is common to many of its calls. The block layer is what makes both cheap to hold. Chapter 16 builds prefix caching on exactly the sharing mechanism measured here.
Where contiguous allocation wastes memory
A cache stored as one contiguous region per sequence can waste memory in three ways.
Reservation. The server does not know how long a sequence will become. It either reserves the model’s maximum context, or the limit the request declared, so that the region never needs to grow. Almost all of that reservation is empty most of the time.
Internal fragmentation. Even a perfectly sized region has an unused tail while the sequence is still growing.
External fragmentation. Regions of different sizes are created and destroyed in different orders, and free memory ends up in pieces too small to use. Chapter 10 measured this effect in the allocator: reserved memory 5.4 times the cache.
I modelled the first two. The simulation is deliberately simple and I will state what it is: a saturated server (a request is admitted the moment there is room), a fixed KV budget of 32,768 tokens, 3,000 requests, and sequence lengths drawn from log-normal distributions. Each resident sequence gains one token per step, and prefill is ignored. It is a model of an allocator under a synthetic workload, not a measurement of any server. Three workloads:
| Workload | Mean prompt | Mean output | Mean final length |
|---|---|---|---|
| Chat-like | 217 | 154 | 371 |
| Long prompt, short answer | 627 | 47 | 674 |
| Short prompt, long answer | 47 | 359 | 406 |
(The distributions were specified by their medians; lognormal means are higher.) The maximum context is 1,024, and a policy that “reserves the declared limit” is given a limit of 512 new tokens, which is a common default and is generous for these workloads. Memory utilisation is the tokens actually stored divided by the tokens set aside:
| Policy | Chat-like | Long prompt | Long answer |
|---|---|---|---|
| Reserve the maximum (1,024) | 33% | 61% | 30% |
| Reserve the declared limit | 48% | 64% | 55% |
| 16-token blocks, on demand | 98% | 99% | 97% |
Reserving the maximum wastes two thirds of memory for the workloads with short sequences. The declared-limit policy recovers some but not most, and it relies on clients declaring honest limits. Paging reaches 97 to 99 per cent for all three.
Paging
The paged cache decouples a sequence’s logical positions from where they sit in memory.
- Physical memory is one pool of equal blocks, each holding the keys and values for
block_sizeconsecutive positions of every layer. - Each sequence has a block table: a list mapping its logical block \(i\) (positions \(ib \ldots ib + b - 1\)) to a physical block number.
- A block is allocated when a sequence’s next position falls into a new logical block, and returned when the sequence finishes.
- Blocks are reference-counted. A block with more than one owner is shared and read-only; writing into it requires a private copy first (copy on write).
This is the structure of an operating system’s virtual memory, applied to a cache, and it has the same properties. Allocation is constant time and never fails for want of a contiguous region, because every block is the same size and any free block will do. External fragmentation disappears. The cost is the block table (64 entries of 4 bytes for a 1,024-token sequence in 16-token blocks, 256 bytes against 72 MiB of cache) and the indirection when reading.
What it costs in the last block
What remains is internal fragmentation in each sequence’s last block, which on average holds about half its capacity. The expected waste per sequence is about \((b - 1)/2\) tokens, so the waste fraction is
| Block size \(b\) | Predicted waste | Measured waste |
|---|---|---|
| 1 | 0% | 0.0% |
| 16 | 2.0% | 2.2% |
| 64 | 8.5% | 8.7% |
| 128 | 17.1% | 16.0% |
| 256 | 34% | 27.3% |
The formula tracks the simulation until the blocks are large relative to the sequences, where the finite spread of lengths starts to matter. Small blocks waste little. Large blocks recreate the problem they were meant to solve.
Why not use \(b = 1\)? It wastes nothing, but everything else gets worse: the block table has one entry per token, and an attention kernel reading through it makes one indirect read per token. The measurement below shows that cost does not appear in the cache copy, but a fused kernel reading through the table would meet it. Production engines settle in the range of tens of tokens, and the right value depends on the kernel.
The real cache
PagedKV stores two tensors of shape [layers, blocks, block_size, heads, head_dim], one for keys and one for values, plus the block tables. Prefill for one sequence goes through the same forward_cached as Chapter 10, with a thin view that writes new keys and values into blocks and gathers the sequence’s blocks for attention. Decode for a batch of sequences of different lengths is a separate function: it computes each sequence’s query, key and value, writes each new key and value to the sequence’s block and offset, gathers every sequence’s blocks into a padded tensor, and attends with a mask that hides each sequence’s padding.
It is the same function
Four sequences of 5, 19, 40 and 77 tokens were decoded together for 20 steps with 16-token blocks and compared against four independent Chapter 10 caches. The generated tokens were identical. The largest logit difference was \(1.1 \times 10^{-3}\), on logits that reach magnitudes of a few hundred, so about 5 parts per million relative; the same rounding-level disagreement Chapter 10 reported. The tests repeat the comparison on a small model, along with two others: blocks freed by one sequence and reused by another do not leak the old values, and a failed multi-block reservation leaks no blocks.
Sharing and copy-on-write
A 40-token prefix occupies three blocks (two full, one holding 8 tokens). Forking a second sequence from it copied nothing: the block count stayed at three, with each block’s reference count at two. Then each branch decoded a different token. The first branch to write found the half-full block shared and made a private copy; the second then found itself the sole owner and wrote in place. The pool held four blocks instead of the six that two independent caches would need, and both branches matched an independent computation to \(3 \times 10^{-5}\) and \(6 \times 10^{-5}\).
The saving is the prefix: two of the three blocks stay shared for as long as both sequences live. This is a toy of one prefix and two branches. Chapter 16 asks what happens when thousands of requests share thousands of prefixes, and when it is correct to share.
What the gather costs
My attention gathers a sequence’s blocks into a contiguous tensor before running. A paged-attention kernel avoids that by reading blocks in place through the table. None exists on this machine, so this cost is the price of not having one. The paged decode step against the static cache of Chapter 10, for the same batch and context, in two runs:
| Batch | Context | Static cache | Paged (16-token blocks) | Ratio |
|---|---|---|---|---|
| 1 | 128 | 18.1 – 19.6 ms | 18.8 – 21.1 ms | 1.04 – 1.08 |
| 8 | 128 | 22.0 ms | 24.3 – 24.6 ms | 1.12 |
| 32 | 128 | 18.5 – 21.3 ms | 20.4 – 21.4 ms | 0.96 – 1.16 |
| 64 | 128 | 20.3 – 22.8 ms | 20.9 – 22.0 ms | 0.96 – 1.03 |
| 32 | 512 | 18.3 – 18.7 ms | 33.6 – 37.5 ms | 1.8 – 2.0 |
| 64 | 512 | 19.4 – 19.9 ms | 59.2 – 63.8 ms | 3.0 – 3.3 |
At small batches or short contexts the gather is invisible, because the step is launch-bound (an absolute 18 to 25 ms here, about twice what earlier chapters saw: the host-dependent floor of Chapter 11). At 64 sequences with 512 tokens each it triples the step. The arithmetic says why: that cache is 2.4 GB in fp32. The gather reads and writes all of it each step, and attention reads it again, and the layout changes after the gather probably cost further passes. A static cache reads it once. The extra work is a few passes over a cache of that size at a few hundred gigabytes a second, which is the 40 ms or so added.
That is the argument for the kernel half of PagedAttention: without a fused kernel that reads through the block table, the memory-capacity win is paid for in bandwidth exactly where it matters, at large batch and long context.
I also measured the gather alone: eight sequences of 512 tokens (6 MiB per layer, fp16), blocks scattered randomly through the pool or taken in order.
| Block size | Gather time, scattered | In order |
|---|---|---|
| 1 | 81 – 82 µs | 80 – 81 µs |
| 16 | 78 µs | 77 – 78 µs |
| 256 | 74 µs | 74 µs |
The block size made essentially no difference, from one token to 256, and neither did the physical order of the blocks. At this size the copy is limited by the bytes moved, 155 to 170 GB/s counting the read and the write. The scattered layout costs nothing here because each block is a whole row of the pool and a gather with an index is a plain copy. A kernel reading through the table would see the fragmentation differently, and I could not measure it.
The policy is where it succeeds or fails
Paging makes it possible to admit a sequence knowing only its prompt. What it cannot do is make memory appear when a growing sequence needs a block and none is free. The server must then pick a victim, throw its blocks away and recompute it later. I simulated the simplest rule: admit whenever the prompt’s blocks fit, and preempt the newest sequence on shortage. Then I added one parameter, a watermark: admit a request only if, after its blocks are taken, a stated fraction of the pool would still be free as headroom for growth.
Steps to complete the same 3,000 requests (lower is better), with 16-token blocks and a recomputation charged at four extra steps per preemption:
| Policy | Chat-like | Long prompt | Long answer |
|---|---|---|---|
| Reserve the maximum | 15,260 | 4,665 | 35,223 |
| Reserve the declared limit | 10,637 | 4,536 | 19,711 |
| Paged, watermark 0 | 13,754 | 3,877 | 94,960 |
| Paged, watermark 5% | 5,611 | 3,106 | 17,752 |
| Paged, watermark 10% | 5,858 | 3,273 | 14,599 |
| Paged, watermark 20% | 6,550 | 3,671 | 14,489 |
| Paged, watermark 30% | 7,446 | 4,186 | 15,650 |
(The number of preemptions at watermark 0 is 2,094, 228 and 20,809.)
With no headroom, paging is better than reserving the maximum in the first two workloads and much worse than anything in the third: the sequences there are long, they all grow together, and the server repeatedly admits as many as fit and then loses the newest of them. A modest watermark, 5 to 10 per cent, removes almost all of the thrash and keeps most of the capacity. A large one gives the memory back as idle headroom: at 30 per cent the paged server holds 67 sequences instead of 96 in the chat-like workload.
With a sensible watermark, the speed-ups against reserving the maximum are 2.6 to 2.7 times (chat-like), 1.4 to 1.5 times (long prompt) and 2.0 to 2.4 times (long answer), all with recomputation charged at four steps per preemption. They are largest where the maximum is most over-generous. The PagedAttention paper’s abstract reports throughput of 2 to 4 times that of two earlier systems at the same latency, with near-zero memory waste. My numbers are of the same order. They are not a replication: the paper measured real models on real traces and a full system, and I simulated an allocator under assumed length distributions.
The simulation is crude in a few ways that I should state. Recomputation is charged as a constant. A real preempted sequence costs a prefill of its prompt and everything it had generated, which from Chapter 9 is 1 to 5 decode steps at these lengths. Real systems can also swap blocks to host memory instead of recomputing, or choose victims by policy; I tried neither. And the admission rule I used is the first thing anyone would write. The aim is to show that the choice of rule changes the result by more than the choice of data structure.
A projection, by arithmetic
GPT-2 is small enough that memory capacity is not its bottleneck on this card (Chapter 10 estimated room for several hundred full-length sequences). For a model of 7 billion parameters with 32 layers and 32 heads of 128 dimensions in fp16, the cache costs \(2 \times 32 \times 4096 \times 2 = 524{,}288\) bytes per token, half a mebibyte. A full 1,024-token sequence is 512 MiB. In 10 GiB of cache memory that is 20 sequences with the maximum reserved, or about 54 sequences if the cache is stored at the chat-like workload’s mean length of 371 tokens with 98 per cent utilisation.
That figure is arithmetic and not a measurement: such a model does not fit this GPU, and many such models use grouped-query attention, which shrinks the cache several-fold (Chapter 21). It shows the shape of the benefit. The more memory a token costs, the more a policy that stops reserving unused tokens is worth.
What the numbers do not show
- A simulation, not a server. Length distributions are assumed, the queue is saturated, prefill is ignored and recomputation is a flat charge.
- A gather, not a paged kernel. The measured paging overhead is the overhead of this implementation, and it could be removed or could be replaced by different overheads in a fused kernel.
- GPT-2 on one GPU in a noisy session. Absolute step times were about twice those of earlier chapters; I use ratios within a run.
- One prefix and two branches for sharing.
- No swapping to host memory, no victim-selection policies other than “newest”, no priorities.
- Exactness checked on four sequences for 20 steps.
Inference Lab increment
The Lab now has:
BlockAllocator(reference counts, out-of-blocks and double-free errors) andBlockTables(on-demand reservation that leaks nothing on failure,fork, copy on write).PagedKV, a real paged cache for GPT-2;prefill_paged, andforward_paged_decodefor ragged batches.simulate(policy, workload, budget, block_size, watermark), with three policies, deterministic given a seed, and the three workloads.
The rule added: a cache allocator is judged by how many sequences it lets the scheduler run, at what recomputation cost, and the number that governs it is the admission headroom. Chapter 15 takes up the scheduler that spends this capacity.
What you should now be able to do
- Name the three sources of waste in a contiguous KV allocation and estimate each.
- Describe a block table, a reference count and copy-on-write, and say what each is for.
- Predict internal fragmentation from block size and mean sequence length.
- Say what an attention kernel must do to read a paged cache without an extra copy, and what an implementation without that kernel costs.
- Explain why admission policy decides whether paging helps, and what a watermark trades.
- Estimate how many sequences fit in a given memory for a given model dimension.
Key terms
- Block / page: a fixed number of consecutive positions’ keys and values.
- Block table: the map from a sequence’s logical blocks to physical ones.
- Reference count: how many sequences use a block.
- Copy on write: copying a shared block before modifying it.
- Internal fragmentation: unused space in a sequence’s last block.
- External fragmentation: free memory split into unusable pieces.
- Preemption: evicting a running sequence to free blocks; it must be recomputed or restored.
- Watermark: free-memory headroom required before admitting a request.
Reader exercises
- Reproduce the waste table for a workload with a mean length of 100 tokens. At what block size does paging waste more than 20 per cent?
- Replace “preempt the newest” with “preempt the sequence with the fewest generated tokens”. How many preemptions does the long-answer workload have?
- Add swapping: a preempted sequence’s blocks are copied to a host pool and restored later, charged at the PCIe bandwidth. When does that beat recomputation?
- Make
forward_paged_decodeavoid the gather for batch 1 by using the sequence’s blocks directly. Does the step change? - Measure a gather where each block is a row of 16 positions read through an index, against a single contiguous copy of the same bytes. What does the difference say about the cost of a table?
- Share a 512-token prefix among 16 sequences. How many blocks does the pool hold, and how many bytes does it save against 16 independent caches?
- Compute the cache per token and per full sequence for a model with 8 key/value heads instead of 32. How many more sequences fit?
References
- Efficient Memory Management for Large Language Model Serving with PagedAttention — the paper behind this chapter. Its abstract reports “near-zero waste in KV cache memory” and 2 to 4 times the throughput of earlier systems at the same latency; I have not reproduced those figures.
- SGLang: Efficient Execution of Structured Language Model Programs — organises prefix reuse as a radix tree; the subject of Chapter 16.
- Cost-Efficient Large Language Model Serving for Multi-turn Conversations with CachedAttention — retaining KV state across turns and a storage hierarchy; not implemented here.
Evidence in this chapter
| Claim | Category |
|---|---|
| Fixed-size blocks, block tables, reference counts and copy-on-write | Established background (the cited paper) |
| Reserving the maximum used 30–61% of memory; paging used 97–99% | Book result of a simulation with a synthetic workload |
| Internal fragmentation ≈ \((b-1)/(2\bar{L})\) | Derived; matched the simulation to within a few points up to \(b=128\) |
| The paged cache reproduces independent decoding for ragged batches | Book result, four sequences, 20 steps |
| Forking a prefix saves two of three blocks and branches are exact | Book result, one case |
| The gather path triples the decode step at batch 64 and 512 tokens | Book result, two runs; implementation-specific |
| Gather cost does not depend on block size or block order | Book result, one shape |
| Greedy admission thrashes; a 5–10% watermark removes most preemptions | Book result of the simulation |
| Paging gives 1.4–2.8× fewer steps in the simulation | Simulation; not a replication of the paper’s 2–4× |
| 55 sequences instead of 20 for a 7B-class model in 10 GiB | Arithmetic, not measured |
| A fused paged kernel, swapping, real traces, prefix hit rates | Not measured |
Implementation status
- Manuscript:
content/books/inference/14-chapter.md - Runnable examples:
examples/ch14/paged_cache.py(allocator, paged cache, forward and simulator); the original toy allocator remains inexamples/ch14/ - Tests:
tests/test_ch14_paged.py,tests/test_ch14_paged_cache.py(allocator, tables, copy-on-write, no leaks on failure, ragged-batch exactness, fork, block reuse, simulation behaviour; all on CPU with a small model) - Evidence:
research/ch14/sources.md,research/ch14/results.json,research/evidence-register.md - Verification command:
python -m unittest discover -s tests -v
The chapter is a draft with draft = false. It makes no claim about a fused paged-attention kernel, real traces or production schedulers. Chapter 15 takes the capacity this chapter creates and asks how a scheduler should spend it: continuous batching.