Lab 01 — Paged KV-Cache Allocator + Prefix Sharing (Python)

Phase: 07 — LLM Serving | Difficulty: ⭐⭐⭐⭐☆ | Time: 5–8 hours Language: Python (stdlib) | Hardware: none

Concept primer: ../WARMUP.md Ch. 2–4.

Run

python solution.py        # correctness tests + memory comparison + prefix-share demo

0. The mission

Build PagedAttention's memory manager (not the attention math — the allocator, which is the actual innovation, Phase 05 Ch. 5):

  • BlockAllocator — a pool of fixed-size KV blocks (16 tokens each).
  • BlockTable — a sequence's logical→physical block map (OS page table).
  • Copy-on-write prefix sharing — multiple sequences referencing the same physical blocks for a shared system prompt, with ref-counting and copy-on- divergence.

Then measure the three regimes the WARMUP describes.

1. What the output shows

== memory efficiency: sequences that fit in a 4096-block pool ==
naive contiguous (reserve max_len=2048)  :   32 sequences
paged (allocate-on-demand, actual=768)    :   85 sequences   (2.7x)

== prefix sharing: 50 requests, shared 1024-token system prompt ==
without sharing :  3350 blocks used
with sharing    :   214 blocks used   (15.7x less; shared prefix stored once)
copy-on-write   : divergence forked block; refcounts correct

(The prefix-sharing ratio is large here because the shared 1024-token prompt dwarfs each request's 40-token generation — exactly the chat-serving regime where prefix caching pays off most.)

  • Naive contiguous reserves max_len per sequence → most blocks unused (internal fragmentation), so few sequences fit. Paged allocates only what each sequence actually uses → 2–4× more concurrent sequences (WARMUP Ch. 3).
  • Prefix sharing stores the shared system prompt's blocks once; 50 requests reference them instead of copying — the chat-serving win (WARMUP Ch. 4).
  • Copy-on-write forks a shared block only when a sequence diverges, ref-counts staying correct.

2. Reading order (solution.py)

  1. BlockAllocator — allocate/free, the free list, ref-counts.
  2. BlockTable / Sequence — append tokens, grow blocks on demand.
  3. share_prefix + the COW path — ref-count increment on share, copy on write to a shared (refcount>1) block.
  4. The three demos in main().

3. Extensions

  1. Eviction under pressure: when the pool is full, evict an LRU sequence; implement both recovery policies — recompute (drop KV) and swap (move blocks to a host-side pool) — and compare their cost (WARMUP Ch. 7).
  2. Beam search / parallel sampling: fork a sequence into N continuations sharing the prompt's blocks; show COW only on the diverging tokens.
  3. Block-size sweep: vary block size (4/8/16/32 tokens); plot internal fragmentation vs block-table size — find the sweet spot (WARMUP Ch. 3).
  4. Hash-based prefix cache: auto-detect shared prefixes by hashing block contents (RadixAttention-lite) instead of explicit sharing.

4. Common pitfalls

  1. Freeing a shared block while refcount > 1 — corrupts other sequences; free must decrement and only reclaim at zero.
  2. Writing into a shared block without COW — silently corrupts every sharer (the test test_cow_on_divergence catches it).
  3. Off-by-one at block boundaries — token T lives in block T // block_size at offset T % block_size; the append logic must roll to a new block exactly at the boundary.
  4. Counting "blocks used" as allocated-not-shared — the memory win is in physical blocks; shared blocks count once.

5. What this lab proves about you

You've implemented the core of vLLM's memory manager and can explain it as the allocator it is — fixed-size paging + COW sharing for the dominant memory consumer. Reading vllm/core/block_manager_v2.py is now a code review; you can reason about block size, prefix-cache hit rate, and eviction policy as the platform knobs they are.