Tutorial 13: Paged Attention and Memory Management¶
The dense KV cache from Tutorial 12 works but wastes memory: every request allocates enough K/V storage for its maximum possible sequence length, even if the actual generation is much shorter. Paged attention solves this by splitting K/V storage into fixed-size pages and allocating them on demand — just like virtual memory in an operating system.
1. The Memory Problem¶
With the dense cache, if you set max_seq_len=4096 and serve 4 concurrent
requests:
Per-layer memory: 4 requests × 4096 tokens × 2 (K+V) × 8 heads × 128 dim × 2 bytes
= 4 × 4096 × 2 × 8 × 128 × 2
= 67 MB per layer
Total for 32 layers: 67 × 32 = 2.1 GB
But if most requests are only 200 tokens long, you wasted 4096 - 200 = 96% of
the allocated memory. With paged attention, you only allocate pages for the
tokens actually generated.
2. The Page Analogy¶
Paged attention borrows the concept of virtual memory from operating systems:
| OS Concept | Paged Attention Analog |
|---|---|
| Virtual page | Logical position in the K/V sequence |
| Physical page frame | Fixed-size slot in the page array |
| Page table (block table) | Maps logical page → physical page |
| Page fault | Need to allocate a new page |
| Page size | Typically 128 tokens |
Each request maintains a block table — a small array that maps logical page indices to physical page IDs:
Request A's block table: [3, 7, 1, ...]
→ logical page 0 is stored at physical page 3
→ logical page 1 is stored at physical page 7
→ logical page 2 is stored at physical page 1
The physical pages do not need to be contiguous — the block table hides the fragmentation.
3. The Page Pool¶
Each transformer layer has one page pool that manages physical storage:
class TinyKvPagedPool:
def __init__(self, page_size=128):
"""
Args:
page_size: Number of tokens per page. Larger pages = less overhead,
smaller pages = less waste. 128 is a common choice.
"""
self.page_size = page_size
self._key_pages = None # Physical page storage for keys
self._value_pages = None # Physical page storage for values
self.free_page_ids = [] # Pages available for reuse
self.used_page_ids = set() # Pages currently in use
self.num_allocated_pages = 0
Allocation¶
When a request needs a new page, the pool either reuses a freed page or allocates fresh storage:
def allocate_page(self):
"""
Get a page ID. Reuses freed pages if available,
otherwise grows the storage array.
"""
if self.free_page_ids:
# Reuse a previously freed page — no memory allocation needed
page_id = self.free_page_ids.pop()
else:
# Allocate a new page — grow the storage if needed
page_id = self.num_pages
self.num_allocated_pages += 1
self.used_page_ids.add(page_id)
return page_id
Deallocation¶
When a request finishes, its pages are returned to the free list:
def free_page(self, page_id):
"""Return a page to the free list for reuse."""
self.used_page_ids.remove(page_id)
self.free_page_ids.append(page_id)
The stale data in the freed page is not erased — it is simply ignored because the block table no longer references it.
4. The Request Cache¶
Each request gets one paged cache per layer. It manages a list of page IDs:
class TinyKvPagedCache:
def __init__(self, pool):
self.pool = pool # Shared pool for this layer
self.page_ids = [] # Which pages this request owns
self.page_lens = [] # How many valid tokens in each page
self.offset = 0 # Total tokens stored
def _append_chunk(self, key, value):
"""
Write new K/V tokens into pages.
This handles two cases:
1. The last page has free slots → fill them first
2. Remaining tokens need new pages
"""
S = key.shape[2] # number of new tokens
start = 0
# Case 1: Fill the partial tail page
if self.page_ids and self.page_lens[-1] < self.pool.page_size:
page_id = self.page_ids[-1]
page_start = self.page_lens[-1]
take = min(self.pool.page_size - page_start, S)
self.pool.write_page_slice(page_id, page_start,
key[:, :, :take, :],
value[:, :, :take, :])
self.page_lens[-1] += take
start += take
# Case 2: Allocate new pages for remaining tokens
while start < S:
end = min(start + self.pool.page_size, S)
page_id = self.pool.allocate_page()
self.pool.write_page_slice(page_id, 0,
key[:, :, start:end, :],
value[:, :, start:end, :])
self.page_ids.append(page_id)
self.page_lens.append(end - start)
start = end
self.offset += S
Example Walkthrough¶
page_size = 4 (for simplicity)
Write 3 tokens: pages = [0], page_lens = [3], offset = 3
Write 2 tokens: pages = [0, 1], page_lens = [3, 2], offset = 5
Write 4 tokens: pages = [0, 1, 2], page_lens = [3, 2, 4], offset = 9
Page 0: [t0, t1, t2, _] ← 3 valid tokens, 1 empty slot
Page 1: [t3, t4, _, _] ← 2 valid tokens, 2 empty slots
Page 2: [t5, t6, t7, t8] ← 4 valid tokens, full
5. Reading Pages for Attention¶
To compute attention, we need the full K/V history. With paged storage, we gather pages into a contiguous view:
def gather_dense(self):
"""Convert paged storage back to a contiguous tensor (for testing)."""
key_chunks = []
value_chunks = []
for page_id, page_len in zip(self.page_ids, self.page_lens):
key_page, value_page = self.pool.read_page(page_id)
# Only take the valid prefix of each page
key_chunks.append(key_page[:, :, :page_len, :])
value_chunks.append(value_page[:, :, :page_len, :])
# Concatenate all pages into a single contiguous tensor
return mx.concat(key_chunks, axis=2), mx.concat(value_chunks, axis=2)
In production, the Metal kernel reads pages directly using the block table — no gathering needed. This is the key performance advantage of paged attention: the kernel knows which physical pages to read without materializing a contiguous tensor.
6. Paged Attention Metadata¶
Instead of passing contiguous K/V tensors to the attention kernel, we pass metadata that describes the page layout:
@dataclass
class PagedKvMetadata:
key_pages: mx.array # (num_pages, num_kv_heads, page_size, head_dim)
value_pages: mx.array # same shape as key_pages
block_table: mx.array # (batch, max_pages) — maps logical → physical page
context_lens: mx.array # (batch,) — total tokens per request
page_size: int
The Metal attention kernel uses this metadata to: 1. For each query token, compute which physical pages it needs 2. Read only the relevant portions of each page 3. Compute attention scores and weighted sums
7. Transactional Safety¶
When multiple requests share a page pool, operations must be atomic. If writing new K/V fails (e.g., out of memory), the pool must be restored to its previous state. Our engine uses snapshot/restore:
def _snapshot_state(self):
"""Save the current state of the pool for rollback."""
return (
self._key_pages,
self._value_pages,
list(self.free_page_ids),
set(self.used_page_ids),
self.num_allocated_pages,
# ... other bookkeeping fields
)
def _restore_state(self, state):
"""Restore the pool to a previous snapshot."""
(self._key_pages, self._value_pages, self.free_page_ids,
self.used_page_ids, self.num_allocated_pages, ...) = state
The batch engine wraps every page operation in a try/except:
pool_state = pool._snapshot_state()
try:
# ... allocate pages and write K/V ...
except Exception:
pool._restore_state(pool_state) # rollback on failure
raise
This prevents partial writes from corrupting the page pool.
8. Memory Efficiency Comparison¶
| Scenario | Dense Cache | Paged Cache |
|---|---|---|
| 4 requests, max 4096, actual avg 200 | 4 × 4096 pages allocated | ~7 pages allocated (200/128 × 4) |
| Memory waste | ~96% | ~2% (one partial page per request) |
| Fragmentation | None (contiguous) | Managed by block table |
| Cross-request sharing | Impossible | Pages freed and reused |
Paged attention typically saves 80-95% of KV cache memory for mixed workloads.
9. Summary¶
- Pages are fixed-size blocks of K/V storage (typically 128 tokens each).
- Block tables map logical pages to physical pages (like OS page tables).
- Page pools manage allocation and deallocation with free-list recycling.
- Transactional safety via snapshot/restore prevents corruption on errors.
- The Metal kernel reads pages directly via the block table, avoiding costly gathering into contiguous tensors.