PagedAttention and KV memory management

Continuous batching only helps if the scheduler can admit another request when a slot frees up. The limit is usually not compute. It is the KV cache. Naive servers reserve a slab sized for the longest allowed sequence when a request arrives. Most chats never use that length, so the reserved space sits empty and fragments the memory.

PagedAttention, from the vLLM project, borrows an idea from operating systems. It splits KV memory into fixed-size blocks, maps each sequence to its blocks with a block table, and allocates blocks as they are needed. Wasted memory drops to roughly one partly filled block per sequence. Larger batches fit. Sharing a prompt between samples becomes a matter of pointing at the same blocks.

Concepts in this section

Read them in order. Each one builds on the one before.

The kernel side

Standard attention code assumes K and V sit in one contiguous tensor. Paged KV is a pool of blocks. The attention kernel must:

1. Read the block table for the sequence.
2. Load each physical block's K and V tiles in turn.
3. Compute scores and weighted values without first gathering everything into one buffer. A gather would use up the memory bandwidth that paging is meant to save.

Early numbers put the kernel overhead at roughly 5% to 10% compared with ideal contiguous attention. Fitting two to four times more sequences more than pays for it. Serving engines route this through tuned backends, so you should not write it yourself unless you are building a new engine.

Myths and failure modes

• "Paging makes attention slow." The indirection costs a few percent in the kernel. Turning away half your batch because of fragmentation costs far more. Measure end-to-end tokens per second, not only kernel microbenchmarks.
• "I set max_model_len to 128k, so I am future-proof." Paging does not reserve 128k per user. But the scheduler may still plan headroom for the maximum, and huge caps make capacity planning harder. Set the cap to what the product needs.
• "GPUs do not fragment." They do, whenever variable-sized allocations are freed out of order. Paging avoids that for KV.
• Ignoring copy-on-write. Running n = 8 parallel samples without sharing multiplies KV by about 8 after the common prefix.

In practice

When something runs out of memory on an LLM server:

1. Check that the engine pages its KV. A custom stack with contiguous caches should be fixed before you buy GPUs.
2. Lower max_model_len to what the product needs. The scheduler plans around it even with paging.
3. Watch preemption counts and time-to-first-token spikes. Constant preemption means you are oversubscribed.
4. For parallel sampling and beam search, confirm that prompt blocks are shared.
5. Compare KV block usage with concurrency. High concurrency and low block usage often means a configuration bug, not a memory shortage.

Interview angle: map OS paging terms to KV blocks, explain why contiguous reservation wastes most of the memory, and say why kernels have to understand block tables.

Key takeaways

• Contiguous max-length KV reservations waste most of the GPU memory for typical requests.
• PagedAttention stores KV in fixed-size blocks, mapped by per-sequence block tables.
• Less waste means a bigger batch, and continuous batching turns that into throughput.
• Reference-counted blocks let samples share a prompt without copying it.
• Under memory pressure, the scheduler preempts sequences by swapping or recomputing, and paging makes both finer-grained.

Further learning

Read the PagedAttention paper and the vLLM announcement post. For a walkthrough, see the MLSys seminar talk. Next, read prefix caching to see the same block idea used across requests.