PagedAttention block allocator

Open in Colab ▶️ Run this notebook in Colab

PagedAttention block allocator#

Track 01 - Inference · Notebook 03 · Runtime: ≈15 min on CPU

Prerequisites: 01_inference/01 (autoregressive decoding and KV cache).

Paper: Kwon et al. 2023, Efficient Memory Management for Large Language Model Serving with PagedAttention.

Problem#

The baseline KV cache stores each sequence’s K and V as a dense tensor of shape (max_seq_len, heads, head_dim). If any sequence reaches max_seq_len, you’ve reserved that much memory for every sequence in the batch - even ones that finished ten tokens in. Kwon et al. measure 60-80 % of reserved KV memory sitting unused in real workloads.

PagedAttention#

Instead of one dense tensor per sequence, store KV in fixed-size physical blocks (16 tokens each by default) pulled from a shared pool. Each sequence keeps a block table mapping its logical token positions to physical block IDs, the same pattern as OS virtual memory to physical pages. That buys three things:

  1. No per-sequence pre-reservation: a 10-token sequence uses ceil(10 / 16) blocks, not max_seq_len worth of tensor.

  2. Prefix sharing via refcounting: sequences with a common system prompt share the same physical blocks until one writes, at which point copy-on-write produces a fresh block.

  3. Fragmentation-free reuse: blocks return to the pool when refcount hits zero.

This notebook builds the three data structures that implement those properties - BlockAllocator, SequenceTable, and fork_sequence - and verifies the waste numbers from the paper.

from llm_systems_cookbook.nb import bootstrap

import random
from collections import deque
from dataclasses import dataclass, field

s = bootstrap("01_inference_03_pagedattention_block_allocator")

Block allocator#

alloc() pops a free block ID; free(block_id) decrements the refcount and returns the block to the pool if it hits zero. incref is what enables prefix sharing.

class BlockAllocator:
    '''Hands out fixed-size physical blocks with reference counting.'''

    def __init__(self, num_blocks: int, block_size: int = 16) -> None:
        self.num_blocks = num_blocks
        self.block_size = block_size
        self._free: deque[int] = deque(range(num_blocks))
        self._refcount: list[int] = [0] * num_blocks

    def alloc(self) -> int:
        if not self._free:
            raise RuntimeError("out of blocks (would trigger preemption or CPU swap)")
        block_id = self._free.popleft()
        self._refcount[block_id] = 1
        return block_id

    def incref(self, block_id: int) -> None:
        self._refcount[block_id] += 1

    def free(self, block_id: int) -> None:
        self._refcount[block_id] -= 1
        if self._refcount[block_id] == 0:
            self._free.append(block_id)

    def refcount(self, block_id: int) -> int:
        return self._refcount[block_id]

    @property
    def free_blocks(self) -> int:
        return len(self._free)


def _pool_exhausts(a: BlockAllocator) -> bool:
    a.alloc()
    a.alloc()
    try:
        a.alloc()
    except RuntimeError:
        return True
    return False


alloc = BlockAllocator(num_blocks=8, block_size=16)
print(f"start:     {alloc.free_blocks} free")
b0 = alloc.alloc()
b1 = alloc.alloc()
print(f"after 2 alloc: {alloc.free_blocks} free, blocks=({b0}, {b1})")
alloc.incref(b0)
print(f"incref b0:   refcount[b0]={alloc.refcount(b0)}")
alloc.free(b0)
print(f"free b0 once: still held, refcount[b0]={alloc.refcount(b0)}")
alloc.free(b0)
print(f"free b0 again: {alloc.free_blocks} free")

s.check("alloc_returns_distinct_blocks", lambda: b0 != b1)
s.check("refcount_reaches_zero_frees_block", lambda: alloc.refcount(b0) == 0)
s.check(
    "alloc_fails_when_pool_empty",
    lambda: _pool_exhausts(BlockAllocator(num_blocks=2, block_size=16)),
)

Sequence block table#

Each sequence holds a list of physical block IDs. Translating a logical token position t to physical storage is

block_index = t // block_size
offset      = t % block_size
physical_id = block_table[block_index]

Appending a token writes into the current trailing block or (if it’s full) asks the allocator for a new one.

@dataclass
class SequenceTable:
    block_size: int
    block_ids: list[int] = field(default_factory=list)
    length: int = 0  # number of logical tokens stored

    def translate(self, logical_pos: int) -> tuple[int, int]:
        '''Map a logical token position to (physical_block_id, offset_in_block).'''
        if logical_pos < 0 or logical_pos >= self.length:
            raise IndexError(f"logical_pos {logical_pos} out of range [0, {self.length})")
        return self.block_ids[logical_pos // self.block_size], logical_pos % self.block_size

    def append(self, allocator: BlockAllocator) -> None:
        '''Append one token's worth of KV, allocating a new block if the last one is full.'''
        if self.length % self.block_size == 0:
            # New block needed (also handles the initial empty case).
            self.block_ids.append(allocator.alloc())
        self.length += 1

    def free(self, allocator: BlockAllocator) -> None:
        for bid in self.block_ids:
            allocator.free(bid)
        self.block_ids = []
        self.length = 0

Waste measurement#

Pre-PagedAttention waste = 1 − (used tokens / max_context_len × num_sequences). Paged waste = 1 − (used tokens / allocated-block bytes). The paged number is bounded above by 1 / block_size - only the trailing fractional block of each sequence wastes space.

BLOCK_SIZE = 16
MAX_CONTEXT = 4096
NUM_BLOCKS = 2048   # plenty to model the workload below

def simulate_workload(lengths: list[int]) -> tuple[float, float]:
    '''Return (pre_paged_waste_fraction, paged_waste_fraction) for ``lengths``.'''
    # Pre-paged: every sequence holds max_context_len worth of memory.
    pre_paged_reserved = sum(MAX_CONTEXT for _ in lengths)
    pre_paged_used     = sum(lengths)

    # Paged: each sequence rounds up to the nearest block.
    allocator = BlockAllocator(num_blocks=NUM_BLOCKS, block_size=BLOCK_SIZE)
    seqs = [SequenceTable(block_size=BLOCK_SIZE) for _ in lengths]
    for seq, L in zip(seqs, lengths, strict=True):
        for _ in range(L):
            seq.append(allocator)
    paged_reserved = sum(len(seq.block_ids) * BLOCK_SIZE for seq in seqs)
    paged_used     = sum(seq.length for seq in seqs)

    return (1 - pre_paged_used / pre_paged_reserved,
            1 - paged_used / paged_reserved)


rng = random.Random(0)
# 100 sequences drawn from a log-normal-ish length distribution.
LENGTHS = [max(4, int(rng.lognormvariate(mu=4.5, sigma=0.8))) for _ in range(100)]
pre_waste, paged_waste = simulate_workload(LENGTHS)
print(f"mean seq length       = {sum(LENGTHS) / len(LENGTHS):.1f}")
print(f"pre-paged waste      = {pre_waste:.1%}   (every seq reserves {MAX_CONTEXT})")
print(f"paged waste          = {paged_waste:.1%}   (block size {BLOCK_SIZE})")

s.check("pre_paged_waste_above_60pct", lambda: pre_waste >= 0.60, msg=f"{pre_waste:.1%}")
s.check(
    "paged_waste_bounded_by_block_size",
    lambda: paged_waste <= 1.0 / BLOCK_SIZE,
    msg=f"paged waste {paged_waste:.1%}, theoretical ceiling {1.0/BLOCK_SIZE:.1%}",
)
s.check(
    "paging_reduces_waste_by_at_least_10x",
    lambda: pre_waste / max(paged_waste, 1e-6) >= 10.0,
    msg=f"pre={pre_waste:.1%}  paged={paged_waste:.1%}  ratio={pre_waste / max(paged_waste, 1e-6):.1f}x",
)

Where the memory went#

One stacked bar per strategy: the blue slice is KV actually holding live tokens, the red slice is reserved-but-wasted slots. Pre-paged reservation dominates; paging reduces the waste slice to at most one block per sequence.

import matplotlib.pyplot as plt

used_tokens = sum(LENGTHS)
pre_reserved   = len(LENGTHS) * MAX_CONTEXT
alloc_play     = BlockAllocator(num_blocks=NUM_BLOCKS, block_size=BLOCK_SIZE)
seqs_play      = [SequenceTable(block_size=BLOCK_SIZE) for _ in LENGTHS]
for seq, L in zip(seqs_play, LENGTHS, strict=True):
    for _ in range(L):
        seq.append(alloc_play)
paged_reserved = sum(len(seq.block_ids) * BLOCK_SIZE for seq in seqs_play)

strategies = ["pre-paged (contiguous)", "paged (block=16)"]
used   = [used_tokens, used_tokens]
wasted = [pre_reserved - used_tokens, paged_reserved - used_tokens]

fig, ax = plt.subplots(figsize=(6.5, 3.6))
ax.bar(strategies, used,   label="live KV tokens",       color="tab:blue")
ax.bar(strategies, wasted, bottom=used, label="reserved but unused",
       color="tab:red", alpha=0.7)
for i, (u, w) in enumerate(zip(used, wasted, strict=True)):
    ax.text(i, u + w, f"{w / (u + w):.0%} wasted", ha="center", va="bottom", fontsize=9)
ax.set_ylabel("tokens of KV capacity")
ax.set_title("KV allocation: contiguous reservation vs paged blocks")
ax.legend()
ax.grid(True, axis="y", alpha=0.3)
fig.tight_layout()
plt.show()

Copy-on-write fork#

Two sequences sharing a prefix share the physical blocks holding that prefix. The refcount lets both read; when one appends a new token into a not-yet-full block, fork_sequence allocates a fresh block and copies so neither side clobbers the other.

Used in beam search, tree-of-thoughts, and speculative decoding’s draft/verify trees - anywhere one base sequence fans out into siblings with a shared history.

def fork_sequence(parent: SequenceTable, allocator: BlockAllocator) -> SequenceTable:
    '''Create a child that shares every full parent block via refcount bump.'''
    child = SequenceTable(block_size=parent.block_size, length=parent.length)
    full_blocks = parent.length // parent.block_size
    partial = parent.length % parent.block_size

    # Share every full block.
    for bid in parent.block_ids[:full_blocks]:
        allocator.incref(bid)
        child.block_ids.append(bid)

    # Copy the trailing partial block so future appends on either side
    # don't clobber the other's KV.
    if partial > 0:
        new_bid = allocator.alloc()
        # In a real system we'd memcpy the block's K/V tensors here.
        child.block_ids.append(new_bid)
        # Also freshen the parent's trailing block so it doesn't share
        # with the child going forward. (This is "eager COW"; lazy COW is
        # an optimisation.)
        old_trailing = parent.block_ids[-1]
        parent_new = allocator.alloc()
        allocator.free(old_trailing)
        parent.block_ids[-1] = parent_new

    return child


allocator = BlockAllocator(num_blocks=256, block_size=BLOCK_SIZE)
prompt = SequenceTable(block_size=BLOCK_SIZE)
for _ in range(100):
    prompt.append(allocator)
blocks_used_after_prompt = NUM_BLOCKS - allocator.free_blocks  # rough; we used a different allocator
blocks_before_fork = 256 - allocator.free_blocks
print(f"after 100-token prompt: {blocks_before_fork} blocks allocated")

child = fork_sequence(prompt, allocator)
blocks_after_fork = 256 - allocator.free_blocks
print(f"after fork:             {blocks_after_fork} blocks allocated")
print(f"fork cost in blocks:    {blocks_after_fork - blocks_before_fork}  "
      f"(vs {len(prompt.block_ids)} if we deep-copied)")

s.check(
    "fork_cost_at_most_two_extra_blocks",
    lambda: (blocks_after_fork - blocks_before_fork) <= 2,
    msg=f"fork cost={blocks_after_fork - blocks_before_fork}",
)
s.check(
    "full_blocks_are_shared_by_refcount",
    lambda: all(allocator.refcount(bid) == 2 for bid in prompt.block_ids[:100 // BLOCK_SIZE]),
    msg=f"refcounts={[allocator.refcount(bid) for bid in prompt.block_ids[:100 // BLOCK_SIZE]]}",
)
s.check(
    "child_length_matches_parent",
    lambda: child.length == prompt.length,
)

Exercises#

  1. Add preemption: when alloc() fails, evict the oldest sequence’s blocks instead of raising.

  2. Sweep block_size ∈ {1, 4, 8, 16, 32, 64} and plot paged waste vs block size for the workload above.

  3. Build a 1-parent 10-child tree-of-thought pattern using fork_sequence and count total blocks vs ten independent sequences.

  4. Wire a real torch.Tensor[(num_blocks, block_size, d)] as the physical storage and have SequenceTable.append write actual K/V values into it.

References#

  • vLLM source: block_manager.py for production preemption + swap.

  • SGLang RadixAttention: radix-tree-keyed prefix cache, covered in 01_inference/06_radix_prefix_cache.

s.summary()
s.save()