Lab 01 — PagedAttention, Continuous Batching & Speculative Decoding

Phase: 09 — Inference Serving Internals (vLLM/TensorRT-class) Difficulty: ⭐⭐⭐⭐☆ (the off-by-one block math and the acceptance-rule proof are where people fall) Time: 4–6 hours

A serving engine like vLLM or TensorRT-LLM is, at its core, three mechanisms bolted onto a transformer: a paged KV-cache allocator so memory doesn't fragment, an iteration-level scheduler so the batch stays full, and (optionally) speculative decoding so you do fewer sequential target-model steps. This lab builds runnable miniatures of all three — the BlockManager (KV in fixed pages, with a block table per sequence), the ContinuousBatchingScheduler (admit and retire requests every step, beating static batching on wasted slot-steps), and speculative_accept (the min(1, p/q) acceptance rule and the residual resample that provably preserves the target distribution).

What you build

• blocks_needed(num_tokens, block_size) — ceil division. The whole paging scheme rides on getting the exact-multiple boundary right (blocks_needed(8, 4) == 2, not 3).
• kv_cache_blocks_fragmentation(seq_lens, block_size) — the proof that paged allocation wastes far fewer slots than naive max_len-per-sequence reservation (the over-allocation PagedAttention removes).
• BlockManager(num_blocks, block_size) — allocate a prompt's blocks, append one token (a new block only when the last is exactly full — the soul of the lab), free to reclaim, num_free_blocks, and None on OOM. A per-sequence block table makes the cache non-contiguous.
• ContinuousBatchingScheduler(max_batch, block_manager) — a .step() that admits waiting requests if a slot + blocks are free, advances each running request one token, and retires finished ones in-flight; plus static_batching_steps(...) as the baseline to beat.
• speculative_accept(draft_tokens, target_probs, draft_probs, rng) — accept draft token i with prob min(1, p_target/p_draft); on rejection resample from the normalized residual max(0, p_target - p_draft); return (num_accepted, bonus_token).

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• You can explain why append must allocate a new block only when len % block_size == 0, and why getting it wrong silently leaks or over-allocates blocks (test_append_allocates_new_block_only_when_last_is_full is the soul test).
• You can explain why kv_cache_blocks_fragmentation shows paged waste << naive waste on a skewed batch, and why that gap is the motivation for PagedAttention.
• You can explain why test_continuous_batching_beats_static_on_wasted_slot_steps holds — what head-of-line blocking is and why in-flight retirement removes it.
• You can explain why speculative_accept preserves the target distribution, and reproduce the one-line argument behind test_spec_decode_preserves_target_distribution.
• You can size a paged KV-cache and reason about OOM/preemption from num_free_blocks.

How this maps to the real stack

Limits of the miniature (say these in the interview): there is no real attention kernel — the blocks here hold counts, not actual K/V tensors, so we model allocation, not the FlashAttention gather over a block table; the scheduler advances "one token = one step" and ignores prefill cost, chunked prefill, and preemption/swapping; speculative_accept verifies one sequence with given distributions and skips the parallel target forward pass and the draft model itself; and "slot-steps" is a throughput proxy, not measured latency (TTFT/TPOT live on real hardware).

Extensions (build these on real hardware)

• Prefix sharing / copy-on-write: ref-count blocks so two sequences sharing a prompt point at the same physical blocks; copy a block only when one diverges (the write). Show the KV savings for a fixed system prompt across 100 requests.
• Preemption & swapping: when allocate/append returns None, evict the lowest-priority running sequence's blocks (recompute or swap to host), then re-admit — the real OOM path.
• Chunked prefill: split a long prompt's prefill across steps so it interleaves with decode and doesn't spike TTFT for everyone else in the batch.
• Real KV blocks + a tiny attention kernel: store actual K/V vectors per block and implement the gather-over-block-table attention so PagedAttention is end-to-end.
• Tree/typical speculative acceptance (Medusa/EAGLE): verify a tree of draft tokens per step and accept the longest valid path; measure the acceptance-length distribution.
• Run a real vLLM server with --speculative-model and compare measured tokens/s and acceptance rate against your miniature's predictions.

Interview / resume

• Talking points: "Walk me through PagedAttention — why blocks, what's the block table, why near-zero fragmentation." "Static vs continuous batching — what's head-of-line blocking and why is continuous 2–4×?" "Speculative decoding — prove it preserves the target distribution." "When do you reach for a draft model vs Medusa/EAGLE vs n-gram?"
• Resume bullet: Built a from-scratch LLM-serving core — a PagedAttention KV-block allocator with per-sequence block tables, a continuous (in-flight) batching scheduler that eliminates head-of-line blocking, and a distribution-preserving speculative-decoding acceptance rule — verified with a 25-test suite including statistical checks that the speculative output matches the target distribution.