S06Memory & KV Cache·283 lines

PagedAttention

Contiguous per-request KV buffers waste 60–80% of KV-cache memory to fragmentation. Paging the cache into fixed-size blocks recovers nearly all of it.

block table
block manager
internal fragmentation
copy-on-write
gather kernel

Motivation

The problem

The contiguous cache from S05 has a fatal flaw, and it lives in the allocator rather than in the kernel. To hand a request a contiguous buffer you must decide its size at admission time, before you know how long the generation will be. So you reserve the maximum.

Measured on real traffic, that produces three distinct kinds of waste:

Internal fragmentation: the reserved-but-unwritten tail of every request. A request that stops at 40 tokens with a 2,048-token reservation wastes 98% of it.
External fragmentation: free blocks exist, but not adjacent free blocks, so a new request cannot be placed even though total free memory is ample.
Unshareable duplicates: two requests with the same 2,000-token system prompt store two identical copies, because each owns a private buffer.

The vLLM paper found existing systems keeping only 20–40% of KV memory in useful tokens. The other 60–80% was fragmentation.

Core idea

The solution

Operating systems solved this in the 1960s. Stop giving processes contiguous physical memory; give them contiguous virtual memory backed by scattered physical pages, and keep a page table to translate.

PagedAttention is that idea applied to the KV cache. Chop the cache into fixed-size blocks; 16 tokens each is the vLLM default. Give each request a block table mapping its logical block i to whichever physical block happens to be free, and allocate one block at a time as the request grows.

the block managerpython

class BlockManager:
    def __init__(self, num_blocks, block_size):
        self.block_size = block_size
        self.free = list(range(num_blocks))     # a free list. That is all.
        self.tables: dict[int, list[int]] = {}  # request id -> physical blocks

    def append_token(self, req_id, n_tokens_now) -> bool:
        table = self.tables.setdefault(req_id, [])
        needed = ceil(n_tokens_now / self.block_size)
        while len(table) < needed:
            if not self.free:
                return False                    # out of memory -> preempt (S10)
            table.append(self.free.pop())       # ANY free block will do
        return True

    def free_request(self, req_id):
        self.free.extend(self.tables.pop(req_id, []))

That is the allocator in fifteen lines. The hard part is the other half: attention can no longer assume its keys and values live in one flat array.

Why this needed a custom CUDA kernel

A normal attention kernel walks memory with a stride. A paged kernel must look up each block's physical address in the block table before it can read it. That indirection has to happen inside the kernel: gathering the blocks into a contiguous buffer beforehand would cost exactly the copy you were trying to avoid. PagedAttention is therefore as much a kernel paper as a memory-management one.

Mechanics

How it works

Diagramcontiguous vs paged allocation

The same ten tokens. On top they occupy 36 reserved slots and block a fourth request. On the bottom they occupy 10 blocks and leave 14 free.

Choosing the block size

Block size is the one real tuning knob, and it trades two costs against each other:

Larger blocks → shorter block tables, less per-block indirection, better memory coalescing in the kernel. The cost is more waste in each request's final partial block, and coarser prefix sharing.
Smaller blocks → near-zero waste and finer-grained sharing, at the price of longer tables and more kernel overhead per token.

Expected internal waste is block_size / 2 tokens per request, independent of request length. Waste therefore stops scaling with context.

vLLM default block: 16; 8 tokens wasted per request, on average
typical utilisation: >96%; versus 20–40% with contiguous reservation
throughput gain: 2–4×; entirely from fitting more requests in the same GPU

Copy-on-write, for free

Because blocks are referenced through a table, two requests can point at the same physical block. Keep a reference count and you get copy-on-write: parallel samples from one prompt (n=4), or beam search branches, share every block of the common prefix and only copy the block where they diverge.

Sharing a 2,000-token prompt across four samples turns 8,000 tokens of cache into roughly 2,000. In the next chapter we push this further and share blocks across different requests.

Explore

Try it

A real block manager running a realistic workload: mostly short generations with a few long ones. Run it in paged mode, then switch the allocator to contiguous and run it again. Watch how many requests each mode keeps in flight with identical memory.

Simulatorblock managerrun paged, then switch to contiguous

allocator

block size (tokens)4workload seed5

step 0 / 90

running: 14
completed: 0
blocks: 58/64
wasted / preempted: 17% · 0

physical KV block pool · 64 blocks × 4 tokens

free
occupied by real tokens
reserved but empty (waste)

Blocks are handed out one at a time, on demand, from anywhere in the pool. A request's blocks need not be adjacent; the block table below makes them look contiguous to the attention kernel. Waste is capped at under one block per request.

block tables (logical → physical)

r017 tok01234
r116 tok5678
r218 tok910111213
r38 tok1415
r417 tok1617181920
r510 tok212223

This table is the entire indirection. The attention kernel walks it to find where logical block i physically lives, exactly as a CPU page table maps virtual pages to physical frames.

internal fragmentation

live tokens: 213
waste in used blocks: 8%

Raise block size to 16 and watch waste climb: the last block of every request is only partly full, so expected waste is blockSize/2 tokens per request. Lower it to 1 and waste vanishes, but the block table gets 16× longer and the kernel does 16× more indirection. vLLM defaults to 16 because that is where the two curves cross.

Then push block size to 16 and back to 1. At 16 the waste bar grows visibly; at 1 it vanishes and every block table becomes sixteen times longer. What the default buys you is a good position on that curve.

Build it

Implementation

Here is the gather step that makes paged attention work. In NumPy it is an index operation; in CUDA it is the inner loop of the kernel.

code/s06_paged_attention.py (excerpt)python

def paged_attention(q, block_table, k_cache, v_cache, seq_len, block_size):
    """
    k_cache: [num_blocks, block_size, n_kv_heads, head_dim]  — the pool
    block_table: [num_logical_blocks]                        — this request's map
    """
    out_scores = []
    for logical, physical in enumerate(block_table):
        start = logical * block_size
        n = min(block_size, seq_len - start)
        if n <= 0:
            break
        k_blk = k_cache[physical, :n]        # the indirection, in one line
        out_scores.append(q @ k_blk.transpose(0, 2, 1) / sqrt(head_dim))

    scores = np.concatenate(out_scores, axis=-1)   # looks contiguous again
    weights = softmax(scores)
    return gather_values(weights, block_table, v_cache, seq_len, block_size)

Reference counts must be exact

Free a block that another request still points at and you get silent corruption: attention reads someone else's keys and the model produces fluent nonsense with no error anywhere. Free too conservatively and you leak until the pool is exhausted. Test the allocator directly. After every operation, assert that every physical block is owned by exactly as many block tables as its refcount claims.

Run it locally

Implements a block manager and a paged attention function, asserts paged attention is numerically identical to contiguous attention, then replays a synthetic workload under both allocators and reports utilisation, waste and how many requests each could keep resident.

$ python code/s06_paged_attention.py

Expect: A numerical equivalence assertion, a utilisation comparison (roughly 26% contiguous vs 52% paged), and a block-size sweep showing waste tracking block_size/2.

Only NumPy is required — setup instructions.

Production notes

In production

vLLM — the original implementation. The v1 engine also supports a "cascade" path that batches the shared prefix separately from the divergent suffixes.
SGLang — replaces the flat block table with a RadixAttention tree, which subsumes paging and prefix sharing into one structure (S07).
baseRT — advertises "paged key-value cache" alongside continuous batching; the two go together, because paging is what makes a variable-size batch affordable.
FlashInfer — the attention-kernel library that engines such as SGLang build on for paged attention, so you get the indirection without writing CUDA yourself.

Exercises

1
Add reference counting and implement copy-on-write for n>1 sampling. Verify that four samples from a 1,000-token prompt use ~1,000 blocks of cache, not 4,000.
2
Sweep block size from 1 to 32 and plot utilisation against block table length. Find where the curves cross for your workload and compare to vLLM's default of 16.
3
Make the allocator fail: drive it to exhaustion and implement the two recovery options, recomputing the request from scratch or swapping its blocks to host memory. Measure which is cheaper at various context lengths; this is the S10 preemption decision.

Continue

What's next

Blocks are shareable, and so far we have used that only for branches of the same request. A chat server sends the same system prompt thousands of times an hour. S07 hashes block contents so that any two requests with a common prefix share the same physical blocks. Prefill cost often drops by an order of magnitude.

Check yourself

Questions

Answer before you look. A wrong answer tells you more than a right one, because the explanation points at the part of the chapter to reread.

Check yourselfQuestion 1 of 6

With a block size of 16, what is the expected internal fragmentation per request, and how does it scale with context length?

Score 0/6