llm-inference-explained
← /learn · 05

Memory management

PagedAttention, block tables, prefix sharing and preemption.

Concept

Continuous batching wants as many sequences resident as memory allows. The KV cache decides that, and caches grow one token at a time to a length nobody knows in advance.

The naive answer is to reserve, per request, a contiguous slab big enough for the maximum length it could reach. Most requests stop far short, so most of each slab sits empty: memory reserved but never used, and fragmentation between slabs as requests come and go. The vLLM paper profiled existing systems storing actual token states in only 20.4% to 38.2% of their KV-cache memory.

PagedAttention borrows the operating system's answer to the same problem: virtual memory with pages.

  • The KV memory is cut into fixed-size blocks (say 16 tokens each).
  • Each sequence has a block table: logical block 0, 1, 2… → physical block numbers, which can be anywhere.
  • A block is allocated only when a sequence's last block fills up. The only waste is the unfilled tail of each sequence's last block.
  • The attention kernel follows the block table to gather keys and values, so the blocks don't need to be contiguous.

Paging makes three more things easy:

  • Prefix sharing. Two requests with the same system prompt, or several samples of one prompt, can point their block tables at the same physical blocks, with a reference count. When one of them writes into a shared, partly full block, it gets its own copy first (copy-on-write). Prefix caching keeps such blocks around after a request ends, so the next request with that prefix skips its prefill. SGLang's RadixAttention organises cached prefixes as a radix tree.
  • Preemption. When a running sequence needs a block and none is free, the scheduler evicts a sequence. In the vLLM paper the latest-arrived requests are preempted first, and either swap their blocks to CPU memory or drop them and recompute the cache later (a prefill over prompt plus output so far).
  • Admission control. A request is admitted only when its blocks fit.

Step through the scenario: watch the shared blocks (S) appear when request 2 joins with request 0's prompt, and shrink the pool to force a preemption.

Interactive

Blocks, block tables and preemption

Six requests share a pool of 4-token blocks. Each step every running sequence appends one token. Request 2 has the same prompt as request 0, so it reuses its full blocks. When a sequence needs a block and none is free, the latest-arrived running request is preempted and recomputed later.

0
0
0
0

Numbers: the request owning each block. S: a block shared by several sequences (reference count > 1). Grey: free.

Running
#0
Waiting
–
Finished
–
Blocks used
4 / 28
Paged slack
2 slots
last-block gaps
Contiguous waste
50 slots
64-token slabs
Shared tokens
0
stored once
Preempted, to recompute
–

Latest event: #0 admitted.

Maths

With block size BB, a sequence of nn tokens occupies ⌈n/B⌉\lceil n/B \rceil blocks and wastes

B⌈n/B⌉−n<BB\lceil n/B \rceil - n < B

slots, against nmax⁡−nn_{\max} - n for a contiguous reservation of nmax⁡n_{\max}. Over a batch, paged waste is under one block per sequence; contiguous waste grows with how far each sequence is from the maximum.

A prefix of pp tokens shared by kk sequences is stored once instead of kk times: ⌊p/B⌋\lfloor p/B \rfloor full blocks are shared outright and the partial tail block is copied on the first write.

Code

// src/lib/inference/paging.ts — append with copy-on-write (excerpt)
const needed = blocksFor(p, len + n) - table.length;
const last = table[table.length - 1];
const cow = last !== undefined && len % p.blockSize !== 0 && p.refs[last]! > 1;
if (needed + (cow ? 1 : 0) > p.free.length) return false; // caller preempts
if (cow) {
  p.refs[last]!--;
  table[table.length - 1] = takeBlock(p)!;
}
for (let i = 0; i < needed; i++) table.push(takeBlock(p)!);

The unit tests check that blocks are never double-booked, that sharing and copy-on-write keep the reference counts right, and that every request in the scenario finishes, with or without preemption.