Paged attention applies the operating system idea of virtual memory and paging to the KV Cache. It is what lets a server pack many variable-length sequences into GPU memory with almost no waste; also enables Continuous Batching.
The problem it solves
Allocating one contiguous KV block per request forces you to reserve for the maximum possible length. Most requests are shorter, so a lot of memory is lost to internal fragmentation and over-reservation.
How it works
- KV cache is split into fixed-size blocks (pages), each holding the keys and values for a fixed number of tokens.
- A per-request block table maps logical token positions to physical blocks, exactly like an OS page table.
- Blocks are allocated on demand as a sequence grows, and need not be contiguous in memory.
- The attention kernel follows the block table to gather K and V.
Payoff
- Memory waste drops to under one block per sequence, so far more sequences fit and batch sizes rise.
- Copy-on-write sharing: several sequences can point at the same physical blocks until one diverges, which is how Prefix Caching and parallel sampling share a prompt’s KV cheaply.
from dataclasses import dataclass, field
from typing import Dict, List, Optional
import numpy as np
class KVCache:
"""
shape of K, V: (num_blocks, block_size, num_kv_heads, head_dim)
"""
def __init__(self, num_blocks: int, block_size: int, num_kv_heads: int, head_dim: int, dtype = np.float32):
self.num_blocks = num_blocks
self.block_size = block_size
self.num_kv_heads = num_kv_heads
self.head_dim = head_dim
shape = (num_blocks, block_size, num_kv_heads, head_dim)
self.k = np.zeros(shape, dtype=dtype)
self.v = np.zeros(shape, dtype=dtype)
def write(self, block_id: int, offset: int, k: np.ndarray, v: np.ndarray) -> None: