PagedAttention

All steps

1. Reserving for the worst case

   Not knowing how long each answer will be, a simple server reserves room for the longest possible one in a single stretch of memory. Most of it is never used, and new requests wait while memory sits empty; vLLM measured only 20 to 38% of KV memory holding tokens in such systems.

   reserve max length per request

2. Blocks on demand

   PagedAttention splits the cache into fixed-size blocks (16 tokens in vLLM, 4 here) and gives a request a new block, from anywhere in memory, only when its last one is full. A block table per request lists its blocks in order, like virtual-memory pages.

   logical block → physical block

3. Attention across blocks

   The attention kernel follows the block table: for each of the request’s blocks it fetches the keys and values from wherever that block lives, and computes the same attention as before. Only the memory layout changed.

   for each block: K, V ← memory[table[b]]

4. Sharing and copy-on-write

   Two samples from one prompt can share the prompt’s blocks; a reference count says how many sequences use each block. When a sample must write into a shared block that is partly filled, it first copies it (copy-on-write), just like forked processes in an operating system.

   shared blocks · ref counts · copy-on-write

5. More requests at once

   The same 16 requests, the same memory. Paging fits more requests into the batch at once, so the whole queue finishes sooner; in the paper vLLM served 2 to 4 times more requests per second than earlier systems at the same latency.

   more sequences per batch → more tokens per second

Code

# the KV cache as a pool of fixed-size blocks
free = list(range(num_blocks)); block_table = {seq: [] for seq in running}
# before writing token t of a sequence, make sure its block exists
if t % BLOCK == 0: block_table[seq].append(free.pop())
b = block_table[seq][t // BLOCK]; kv[b, t % BLOCK] = (k, v)
# attention gathers the sequence's blocks through its table
K = torch.cat([kv_k[b] for b in block_table[seq]])
# sharing: forked sequences copy the table, not the blocks
block_table[child] = list(block_table[parent]); ref[b] += 1 for each b
if ref[b] > 1 and writing: b = copy_block(b); ref[b_old] -= 1

Go deeper

• Kwon et al. 2023, Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM)