Chapter 18 Speculative Decoding and Serving Optimizations
18.1 Speculative Sampling
18.1.1 The Draft/Target Algorithm and Its Math
Speculative sampling comes up in nearly every inference-focused system design interview because it is the rare technique that cuts latency without any approximation: the output distribution is provably identical to that of the target model. Interviewers at frontier labs ask about it both to probe whether a candidate understands why generation is slow (it is memory-bound, not compute-bound) and to test whether they can reason about probabilistic correction schemes under pressure. Expect it at the senior to staff level, where the follow-up is almost always “prove the acceptance step is unbiased” or “how would you choose the draft model?”
Autoregressive generation produces one token per forward pass, and each pass is dominated by loading the model weights from HBM-not by arithmetic. For a 70 B parameter model in bf16 on an H100 with 3.35 TB/s memory bandwidth, a single generation step requires moving roughly GB of weights, which takes about 42 ms regardless of batch size. Prefill, by contrast, processes all prompt tokens in parallel and is compute-bound: you pay the weight-load cost once and amortise it over the whole sequence. Speculative sampling exploits this asymmetry directly.
The algorithm uses two models: a small, fast draft model and the large, accurate target model . Given a prefix , the draft model autoregressively produces candidate tokens under its own distribution. The target model then scores all tokens in a single forward pass-this is just a prefill over the extended sequence, so it runs in parallel and is fast. For each position the candidate is accepted with probability
If the token is accepted the algorithm moves to position . If it is rejected, a corrected token is sampled from the residual distribution
and the loop terminates. Finally, if all tokens are accepted, the algorithm samples one additional token directly from for free. The net effect is that every call to the expensive target model potentially yields more than one new token.
The central insight is that verifying a proposed token under the target model is identical to prefill, which is cheap; only generating the next token requires a fresh, expensive decode step.
The acceptance-rejection correction is not a heuristic-it is an exact importance-sampling step borrowed from the Metropolis-Hastings framework. The proposal distribution is ; the target is . When the draft agrees with the target (), the token is accepted unconditionally. When the draft is overconfident relative to the target, the excess probability mass is redistributed via the residual , ensuring the marginal distribution of every emitted token matches exactly. No temperature adjustment, no post-hoc filtering-the samples are mathematically identical to what you would get from greedy-free sampling the target model alone.
The practical speed gain is governed by the expected number of tokens produced per target call, which equals where is the average per-token acceptance rate. With a well-matched draft-target pair, routinely reaches 0.7-0.85 on natural language tasks, giving 2-3 useful tokens per expensive decode step. The latency wall from weight-loading is amortised across those tokens, reducing wall-clock time per output token proportionally.
The key trade-off is draft quality versus draft cost. A stronger draft model pushes higher but also consumes more memory and FLOPs, narrowing the gap to the target model. A weaker draft is cheaper but drops acceptance rates; if falls below roughly the overhead of running the draft outweighs the savings. The sweet spot in production is typically a draft 5-10 smaller than the target, from the same model family and ideally distilled from the target.
A Worked Example
Consider a setup with Llama 3 70B as the target and a 7B draft, both in bf16 on a single H100 (80 GB HBM, 3.35 TB/s bandwidth).
Baseline latency per token (target only). Weight size: GB. Time to load weights: ms per step. At lookahead tokens the target model is called once for every 5-step window.
Draft cost. The 7B draft weighs GB. Loading time: ms per step. Five draft steps cost ms total.
Verification cost. The target scores all 5 candidates in one prefill call. Prefill is compute-bound, not bandwidth-bound; at reasonable sequence lengths this is roughly equivalent to one generation step: ms.
Effective throughput with . Expected accepted tokens per round: . Total time per round: ms. Effective latency per output token: ms, versus the ms baseline. That is a wall-clock speed-up with zero change in output quality.
Whiteboard summary:
For each round of K draft steps:
1. Draft model: generate x_{t+1} ... x_{t+K} [K * draft_cost ms]
2. Target model: one prefill over the K tokens [~1 decode step ms]
3. Accept/reject each token, emit 1 + K*alpha tokens on average
Speedup = (1 + K*alpha) / (1 + K * draft_cost / target_cost)
What You Actually Decide in Practice
-
Draft model size. Aim for a 7-10 parameter ratio between target and draft. Below that the draft is too slow; above it acceptance rates collapse. When the draft comes from the same family and is distilled from the target, acceptance rates are 5-10 percentage points higher than a generic model of the same size.
-
Lookahead length . Start at -. Increasing raises the ceiling on speedup but also raises the cost of a bad draft run-if is low, long lookaheads waste draft compute. Profile on representative prompts before tuning upward.
-
Memory budget. The draft model lives alongside the target model in GPU memory. On an 80 GB H100 a 70B target leaves roughly 10 GB for the draft after KV caches-enough for a 4-5B model in
bf16or a 7B model quantized to INT4. If you cannot fit both models, consider offloading the draft weights to CPU and streaming them in; the CPU-GPU transfer cost is often less than the savings on generation steps. -
When speculative decoding hurts. If the target and draft distributions are poorly aligned (cross-domain prompts, highly stochastic sampling with temperature ) drops below 0.5 and the overhead dominates. Monitor the acceptance rate per request type in production and fall back to direct decoding when
for a given prompt class.
- Batch size interaction. Speculative decoding is most valuable at batch size 1, where the target model is entirely bandwidth-bound. At large batch sizes ( on an H100) the target model starts to become compute-bound during the verification prefill, and the relative advantage of speculative decoding shrinks. For high-throughput serving at large batch sizes, continuous batching and KV cache management usually matter more than speculation.
How This Shows Up in Interviews
-
“Walk me through how speculative decoding works and why it speeds things up.” A strong answer names the draft-target pair, explains that the target’s verification step is a prefill (parallelisable, cheap) rather than a decode (sequential, bandwidth-bound), and quantifies the speed-up in terms of and . Weak answers say “the small model helps the big model” without connecting it to the memory-bandwidth bottleneck.
-
“How do you guarantee the output is drawn from the target distribution?” The key is the acceptance probability and the residual distribution for rejected tokens. A strong answer sketches why these two together give an unbiased sample: accepted tokens come with weight , and the correction redistributes the rejected mass exactly as requires.
-
“What would you do if the acceptance rate on production traffic is lower than expected?” Concretely: (1) profile which prompt categories have low ; (2) fine-tune or distill the draft on that domain; (3) fall back to direct decoding for those categories; (4) lower to reduce wasted draft compute.
-
Senior follow-up: “Can you modify speculative decoding to work without a separate draft model, and what are the accuracy implications?” Yes-Medusa adds multiple independent heads on top of the target model’s residual stream, one per lookahead position. Each head predicts the token steps ahead without autoregressive rollout, so there is no separate model cost. The trade-off is that these heads see only the last hidden state, not the full draft autoregressive context, so acceptance rates are typically 10-20 percentage points lower than a strong standalone draft model of comparable parameter count. EAGLE improves on this by feeding the target model’s top-layer features into a lightweight autoregressive draft, recovering most of the acceptance rate while keeping the draft small and co-located.
Key Takeaways
-
Speculative sampling exploits the fact that verifying a token (prefill) is much cheaper than generating one (decode), and the output is exactly drawn from the target distribution-no approximation.
-
The per-round speed-up is ; with a well-matched draft, - gives 2-3 latency reduction at batch size 1.
-
Choose the draft model from the same family as the target, target a 7-10 size ratio, and prefer a distilled draft over a generic pretrained model.
-
The technique is most effective at small batch sizes where generation is purely bandwidth-bound; at high batch sizes continuous batching yields bigger wins.
-
Monitor acceptance rates per prompt domain in production; fall back to direct decoding when .
18.1.2 Medusa, EAGLE, and Draft Model Design
When you ship speculative decoding to production, the single hardest engineering decision is the draft model itself-how to build it, where to host it, and how tightly to couple it to the target. This question shows up in system design interviews at exactly the point where a candidate has explained the basic algorithm and the interviewer pivots to: “OK, but how would you actually build the draft model for a 70B serving system you own?” Getting this right separates candidates who have thought about the operational reality from those who have only memorized the algorithm.
The draft model’s sole job is to propose tokens fast enough that the target model’s verification step-a single prefill-recovers more wall-clock time than the draft cost. Three architectures dominate in practice.
Standalone draft model. A fully independent, smaller model from the same family as the target. It has its own weights, its own KV cache, and runs its own autoregressive decode loop to produce tokens. Acceptance rates are highest here because the draft sees the full autoregressive context at each step. The cost is memory: on an 80 GB H100 with a 70B target, a standalone 7B draft in bf16 consumes 14 GB, leaving barely enough room for KV caches at non-trivial sequence lengths.
Medusa. Medusa attaches independent prediction heads directly on top of the target model’s last hidden state. Head predicts the token steps ahead of the current position without any autoregressive rollout. There is no separate model and no additional memory for weights beyond the heads themselves (each head is a single linear projection of size , roughly 0.5 GB per head for a 4096-dimensional model and 128 K vocabulary). The draft cost per round collapses to a single forward pass through the heads in parallel. The downside is expressive power: head sees only the hidden state at position , not the predicted tokens at positions , so the heads are conditionally independent. In practice this cuts acceptance rates by 10-20 percentage points relative to a well-tuned standalone draft.
EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency). EAGLE threads a middle path. It trains a lightweight autoregressive draft module-typically a single transformer layer-that takes as input the feature vectors from the target model’s final layer concatenated with the token embeddings of the previously accepted tokens. The draft module is therefore not standalone; it is a thin attachment that rides on the target’s residual stream. Because the draft sees the target’s own high-level representations, it has far more signal than Medusa’s static hidden state, and acceptance rates recover to within 5-10 percentage points of a standalone draft. The draft module is also small-one transformer layer for a LLaMA-sized target is roughly 500 M parameters or 1 GB in bf16-so memory pressure is minimal.
The design spectrum runs from standalone (highest acceptance rate, highest memory cost) through EAGLE (near-standalone quality, low memory cost) to Medusa (lowest cost, lowest acceptance rate); the right choice depends on how much GPU memory the target leaves unclaimed after KV caches.
The memory argument is decisive at the 70B scale. A standalone 7B draft eats 14 GB; EAGLE’s single-layer module eats 1 GB; Medusa’s heads eat 2.5 GB. For a 70B target in bf16 (140 GB) sharded across two H100s (160 GB total), the KV cache budget for a batch of 16 sequences at sequence length 2,048 is roughly 10 GB, leaving only 10 GB spare. EAGLE or Medusa fit easily; a standalone 7B draft does not without aggressive quantisation.
The co-location advantage of EAGLE over a standalone draft is not just memory. Because the draft module ingests the target model’s feature vectors, those vectors only need to be computed once per token position-they are produced by the target during verification and cached. The draft therefore uses information that the standalone model would have to independently re-derive, which is why EAGLE’s acceptance rates are high despite the module being tiny.
The residual between Medusa and EAGLE is conditional independence. When generating code or structured output where token strongly depends on token , Medusa’s heads-each conditioned only on position -systematically underestimate the entropy of later positions and generate plausible-looking but locally inconsistent proposals. EAGLE’s autoregressive rollout within the draft module preserves this conditional structure.
A Worked Example
Consider LLaMA 3 70B (bf16) served on two H100s via tensor parallelism, with lookahead tokens, batch size 1, and a 2,048-token context.
Memory budget per GPU (80 GB each, 160 GB total):
-
Target weights: GB across both GPUs, so GB per GPU.
-
KV cache (batch 1, 2048 tokens, 80 layers, 8 KV heads, head dim 128,
bf16): GB per GPU. -
Remaining per GPU: GB.
Option A: Standalone 7B draft. 7B model in bf16: 14 GB total, GB per GPU if sharded. Fits, but leaves only 2 GB per GPU for activations and the draft’s own KV cache-tight. Expected acceptance rate: (same family, distilled). Tokens per target call: .
Option B: EAGLE single-layer module. Module size: GB total, negligible per GPU. Fits comfortably. Expected acceptance rate: (empirically 5-10 pp below a good standalone). Tokens per target call: .
Latency comparison (target decode ms, draft step ms for standalone, ms for EAGLE’s single layer):
| Strategy | Draft cost | Tokens/call | ms/token |
|---|---|---|---|
| No speculation | - | 1.0 | 42 |
| Standalone 7B | ms | 5.0 | |
| EAGLE | ms | 4.6 | |
| Medusa | ms (parallel) | 3.5 |
EAGLE wins on effective ms/token in this configuration because its near-zero draft overhead and nearstandalone acceptance rate combine better than either extreme. Note that Medusa assumes , which is typical for code-generation tasks where later tokens are highly context-dependent.
What You Actually Decide in Practice
-
Choose Medusa only when memory is the hard constraint. If the target already occupies 95 % of device memory and there is no room for even a 1 GB module, Medusa is your only option. In any other scenario EAGLE dominates on latency.
-
Distill the draft from the target, not from general pretraining. A draft trained only on next-token prediction from public data will have a vocabulary distribution mismatch vs. the target after RLHF or instruction tuning. Even a lightweight supervised fine-tuning run on 1-5 B tokens of target-model outputs raises by 8-15 percentage points. This is the single highest-ROI investment in a speculative decoding deployment.
-
Profile acceptance rate by prompt domain before setting . Open-ended creative tasks typically hit
; multi-step reasoning and code with long logical dependencies can fall to . For low-acceptance domains, reduce to 2-3 or disable speculation entirely; a run with and is slower than no speculation because draft cost exceeds the gain.
-
Quantise the draft, not the target. Users tolerate no accuracy loss from the target. The draft’s sole quality metric is acceptance rate. INT4 quantisation of a standalone 7B draft typically costs 3-5 percentage points in but halves its memory footprint and speeds up each draft step by 1.5-2-usually a net positive for the overall tokens-per-second figure.
-
Tree-based speculative decoding for Medusa. Because Medusa’s heads are independent, you can sample multiple candidates per position and evaluate them as a tree in a single target forward pass rather than a linear chain of tokens. A tree with branching factor 2 at each of levels evaluates 15 nodes (a complete binary tree) in one prefill call, and the best accepted path is selected. This partially compensates for Medusa’s lower per-head acceptance rate without increasing target calls; implement this if you adopt Medusa.
How This Shows Up in Interviews
-
“You are designing speculative decoding for a 70B chat model on two H100s. The memory budget is tight. What draft strategy do you pick and why?” A strong answer works through the memory budget explicitly (target weights, KV cache, headroom), rules out a standalone 7B draft as too large without quantisation, lands on EAGLE as the best acceptance-rate-per-GB option, and mentions distilling the EAGLE module from target outputs.
-
“What is the fundamental quality difference between Medusa and EAGLE?” The answer is conditional independence: Medusa’s head does not see what heads predicted, so proposals are locally inconsistent. EAGLE’s autoregressive draft rollout preserves the conditional structure, recovering most of a standalone draft’s acceptance rate at a fraction of the memory cost.
-
“How do you measure whether speculative decoding is actually helping in production?” Track three metrics: (1) mean acceptance rate per request class; (2) effective tokens-per-second compared to the direct-decode baseline; (3) p99 latency to first token, which should be unchanged since speculation only affects the generation phase. Alert on dropping below 0.55 for any significant traffic slice-this usually signals a distribution shift between the draft and target.
-
Senior follow-up: “An EAGLE draft module trained on general chat data is deployed with a target that was recently fine-tuned on a new coding dataset. Acceptance rates have dropped from 0.75 to 0.52. What do you do?” The gap is a distribution mismatch introduced by the target fine-tune. The immediate fix is to fine-tune the EAGLE module on a sample of outputs from the updated target-even a 200 M token run at low learning rate is usually enough to close most of the gap. Longer term, bake the EAGLE module fine-tune into the target update pipeline so both stay in sync. While the fix is in flight, fall back to direct decoding for the coding traffic slice to avoid negative latency impact.
Key Takeaways
-
Standalone drafts give the highest acceptance rates but consume the most memory; EAGLE recovers near-standalone quality at GB by conditioning on target features; Medusa is cheapest but suffers from conditional independence across heads.
-
EAGLE dominates for most 70B deployments: the memory saving over a standalone draft is decisive, and the acceptance rate gap relative to Medusa is worth the marginal extra cost.
-
Distilling the draft from target-model outputs is the single most impactful tuning lever-expect an 8-15 percentage point gain in over a generic pretrained draft.
-
Quantise the draft aggressively (INT4 is fine); keep the target at
bf16orfp8. -
Monitor per traffic slice; disable speculation when it falls below and retrain the draft whenever the target is fine-tuned.
18.2 Serving System Design
18.2.1 Continuous Batching and Selective Batching
Production serving of LLMs looks nothing like the clean batched inference your benchmarks describe. Real traffic is chaotic: requests arrive at random intervals, prompts span 50 to 50,000 tokens, and generations finish at unpredictable times. An interview question like “how would you serve a 70B model at high throughput with low latency?” is really asking whether you understand the scheduling and batching machinery that separates a demo from a production system. This topic appears consistently at senior and staff levels wherever teams own inference infrastructure, and the two mechanisms here-continuous batching and selective batching-are the architectural foundation of every serious open-source serving system today.
The naive approach is static batching: collect requests until the batch is full, run all sequences in lockstep until every one of them finishes, then start the next batch. The problem is obvious once stated: a request that generates 5 tokens holds a batch slot hostage for the full 512-step budget, and the GPU sits idle between batch arrivals while you wait for enough requests to form a full batch.
Continuous batching (also called iteration-level scheduling) fixes this by making the generation step the atomic scheduling unit. The worker generates one token for every active sequence, then yields control back to the scheduler. The scheduler inspects the waiting queue: any sequence that just generated an end-of-sequence token frees its slot, and a new request can be inserted immediately. The GPU is never blocked waiting for stragglers or for a batch to fill.
Selective batching addresses the second problem: variable-length sequences cannot be naively padded into a rectangular tensor when the length variance is large, because attention compute scales quadratically with sequence length and most of that padding is wasted work. The key observation is that attention and the feed-forward (MLP) layers behave differently. Attention is inherently per-sequence-each sequence has its own KV cache-so it must be dispatched separately for each sequence (or handled with variable-length kernels like FlashAttention- 2’s cu_seqlens interface). MLP layers, however, apply the same weight matrices to every token regardless of which sequence the token belongs to. Given tokens from sequences of different lengths, you can simply flatten the batch dimension: reshape inputs from into and process all tokens together as one large matrix multiply. Because the tokens from different sequences do not interact inside an MLP, this is mathematically identical to processing them separately but with much higher arithmetic intensity.
The key principle: continuously schedule at the granularity of individual generation steps, and batch MLP computation across all active tokens regardless of their sequence boundaries.
The throughput gains from continuous batching are compound. First, you eliminate idle GPU time from batch formation. Second, you eliminate the slot-wasting effect of stragglers: a 20-token request that finishes early frees its KV cache and batch slot after exactly 20 steps, not 512. Third, continuous batching naturally increases the average number of concurrent sequences in steady-state, which raises the arithmetic intensity of MLP layers and drives the GPU toward its compute-bound regime.
Selective batching amplifies this by ensuring that the bulk of the computation-MLP layers constitute roughly two-thirds of total FLOPs in a standard transformer-is batched at maximum efficiency even when sequence lengths vary wildly. Without selective batching, a variable-length batch either suffers quadratic padding waste in attention or requires an inefficient loop over individual sequences in the MLP.
The alternative-waiting for a fixed batch and running to completion-gives lower memory pressure per step because fewer sequences are active at once, but it produces dramatically worse GPU utilization and higher latency variance. In practice, static batching delivers 20-30% GPU utilization for typical interactive workloads; continuous batching pushes that to 70-85%.
A Worked Example
Consider a 7B model on an H100, with a KV cache budget of 512 concurrent sequence slots, each pre-allocated to 2048 tokens maximum.
Under static batching with batch size : the batch is padded to the longest sequence in the batch. If the distribution of actual generation lengths is geometric with mean 256 tokens but maximum 2048, then at least one sequence in every batch of 16 will be near the maximum, so all 16 are padded to 2048 tokens. A request that actually generates 30 tokens occupies a slot for the full 2048 steps. Average slot utilization: .
Under continuous batching with the same memory budget: slots free as sequences finish. In steady state, the scheduler keeps 16 sequences running simultaneously. When sequence generates its EOS at step 30, slot is immediately assigned to the next waiting request. Effective throughput is:
where is the per-step latency (roughly 8 ms per token for a 7B model with on an H100, scaling with batch size). At , ms, giving a steady-state throughput of tokens/s, compared to tokens/s with static batching at 12.5% slot utilization.
For selective batching, suppose active sequences have lengths tokens in their KV caches. The MLP forward pass flattens these into a single tensor of shape and runs two large matrix multiplies of size and . Arithmetic intensity is , comfortably compute-bound on an H100 (threshold for bf16). Attention is dispatched per-sequence with per-sequence KV lengths, using variable-length flash attention kernels.
What You Actually Decide in Practice
-
Schedule in token budget, not sequence count. Set a maximum of, say, 8,192 active tokens across all sequences rather than a fixed batch size. This bounds both memory usage and per-step latency regardless of sequence length distribution. A batch of 64 short sequences and a batch of 4 long sequences both consume the same compute.
-
Chunk long prefills. A single request with a 32,768-token prompt will monopolize the GPU for the entire prefill, stalling all ongoing generation steps and spiking inter-token latency for other users. Chunked prefill breaks the prompt into pieces of, say, 512 tokens per step, interleaving them with generation steps. This bounds prefill latency per step at the cost of slightly slower time-to-first-token for the long-prompt request. Use chunk sizes of 512-2048 tokens; smaller chunks hurt MLP efficiency.
-
Preemption when memory fills. If a high-priority request arrives but there are no free KV cache slots, preempt the lowest-priority active sequence: save its KV cache to CPU DRAM (swap) or simply discard it and restart the sequence from scratch when a slot frees. Swap is preferable for long sequences (large KV cache already built); restart is preferable for short sequences (cheap to redo). The crossover is typically around 500-1000 tokens of context.
-
Priority tiers for multi-tenant serving. Continuous batching makes priority trivially composable: each step, the scheduler fills available token budget starting from the highest-priority queue. Latency-sensitive interactive requests get first pick; batch-processing jobs fill remaining capacity. This is simpler to implement than priority in static batching because slots turn over every step.
-
Flash attention with variable-length support is non-negotiable. Without FlashAttention-2’s
varlenkernels (or equivalent), selective batching for attention requires either a Python loop over sequences (ter-rible) or padding to the max length in the batch (defeats the purpose). Verify your serving stack usesflash_attn_varlen_funcor equivalent before deploying.
How This Shows Up in Interviews
-
“How does vLLM achieve higher throughput than naive batching?” Strong answer: iteration-level continuous batching reclaims slots as soon as a sequence finishes, eliminating straggler waste. Selective batching ensures MLP computation is always batched at full efficiency across all active tokens. PagedAttention (the next subsection) then eliminates KV cache fragmentation to maximize the number of concurrent sequences.
-
“Why can’t you just pad all sequences to the same length and run a standard batch?” Strong answer: padding wastes compute quadratically in attention (padded positions still participate in ), wastes memory for the padded KV cache, and lets long stragglers hold batch slots for all shorter sequences. At mean generation length of 256 and max 2048, effective utilization is .
-
Senior follow-up: “If prefill is compute-limited and generation is memory-limited, what happens to generation latency when a large-batch prefill lands?” Senior answer: a large prefill dominates the GPU for many steps, stalling all generation steps and causing inter-token latency (ITL) spikes for other users. The fix is chunked prefill: bound the prefill tokens per step to, say, 512, so that each step is a mix of prefill and generation. FlashAttention-2 supports this via the
varleninterface, which handles sequences of mixed prefill and generation tokens in a single kernel call. The trade-off is slightly higher TTFT for the long-prompt request, which is acceptable if TTFT SLOs are looser than ITL SLOs.
Key Takeaways
-
Continuous batching reclaims GPU slots every generation step, eliminating straggler waste and idle time between batches.
-
Selective batching separates attention (per-sequence, variable-length kernels) from MLP (all tokens flattened into one large matrix multiply), keeping MLP compute-bound.
-
Schedule by token budget, not sequence count, to bound per-step latency regardless of sequence length distribution.
-
Chunked prefill prevents long prompts from monopolizing the GPU and spiking inter-token latency for concurrent generation requests.
-
FlashAttention-2 variable-length kernels are the prerequisite for practical selective batching; confirm your stack supports them before benchmarking.
18.2.2 Paged Attention and Memory Fragmentation
The KV cache is the dominant consumer of GPU memory in inference, often eclipsing model parameter storage once you scale to reasonable batch sizes. A naive allocator that pre-reserves a contiguous chunk for each request wastes 20-40% of KV cache capacity to fragmentation-directly capping the number of concurrent sequences you can run and therefore your throughput ceiling. If you are interviewing for any role that touches inference infrastructure, expect to explain exactly how PagedAttention works and why it matters: this is the core mechanism behind vLLM and the design pattern that every serious serving system has since adopted or reimplemented. The concept is worth understanding deeply not just for interviews but because it reveals how much of systems work is rediscovering operating-system ideas in new contexts.
The fundamental problem is that you do not know how many tokens a request will generate before it starts. A pre-emptive allocator that reserves the maximum possible sequence length (say, 2048 tokens) wastes GPU memory proportional to the gap between that maximum and the actual generation. An allocator that starts small and grows dynamically requires copying the KV cache each time it outgrows its current allocation-expensive and disruptive.
PagedAttention borrows the solution that operating systems developed for process memory: virtual memory with paging. Physical DRAM is divided into fixed-size pages; a page table maps virtual addresses to scattered physical pages. A process sees a contiguous virtual address space, but the underlying physical layout is fragmented.
Applied to the KV cache, the translation is direct. The KV cache is divided into fixed-size blocks, each storing tokens of KV state for all layers and all KV heads. A block table maps each sequence to an ordered list of physical block indices. When a sequence needs more KV storage, the allocator assigns the next free physical block from a pool-no contiguity requirement. When a sequence finishes, all its blocks are returned to the free pool immediately.
The memory cost of a single block is:
where is the block size in tokens, is the number of layers, is the number of KV heads, is the head dimension, and the final factor of 2 accounts for bf16 precision. The factor before doubles for key and value.
Copy-on-write enables the second major feature: prefix caching. Many requests share a common prefix-a system prompt, a few-shot example block, or a retrieved document chunk. Without sharing, every request independently computes and stores the KV cache for that prefix, paying both compute and memory costs that scale linearly with the number of concurrent requests. PagedAttention gives each physical block a reference count. Blocks in a shared prefix region have reference count equal to the number of sequences using them; those sequences read from the block but never write to it. When a sequence generates its first token past the shared prefix, the scheduler forks it: new blocks are allocated for the diverging suffix, and the shared prefix blocks remain intact and untouched. This is copy-on-write, and it makes the KV cache cost of a repeated system prompt a one-time allocation regardless of how many requests share it.
Eliminate fragmentation first, then exploit prefix sharing: together these two mechanisms double or triple the effective batch size you can sustain within a fixed GPU memory budget.
Pre-PagedAttention serving systems suffered two kinds of fragmentation. Internal fragmentation: you reserve 2048 tokens for a request that generates 80, wasting of its allocation. External fragmentation: the gaps between allocations that no single request can use, analogous to heap fragmentation in a long-running process. The vLLM paper measured 20-40% total KV cache waste from these effects on realistic workloads.
With block-based allocation, internal fragmentation is bounded by at most tokens per sequence-at a block size of 16 tokens, that is at most 15 wasted token slots per sequence, negligible at any reasonable sequence length. External fragmentation is eliminated because blocks are fungible: any free block can be assigned to any sequence.
Prefix caching compounds the benefit. A 1,000-token system prompt on a 7B model (32 layers, 8 KV heads, head dimension 128, bf16) costs:
per request. With 100 concurrent requests all using the same system prompt, the naive cost is 13.1 GB-a significant fraction of the entire KV budget on an 80 GB GPU. Prefix caching reduces this to 131 MB, a 100 reduction in KV memory for the shared prefix. This frees space for more concurrent sequences, which directly improves throughput.
The alternative-a bump allocator with padding-is simpler to implement but puts a hard ceiling on batch size that degrades with sequence length variance. A slab allocator (fixed-size objects) partially addresses external fragmentation but does not support prefix sharing.
A Worked Example
Setup: 7B model, H100 80 GB GPU, model weights in bf16: approximately 14 GB. KV cache budget: roughly 60 GB. Block size tokens, layers, KV heads, .
Bytes per block:
Total physical blocks: GB MB blocks, representing token slots of KV capacity.
Without PagedAttention (pre-allocation): Reserve blocks per sequence at request start. Maximum concurrent sequences: . If actual generation length averages 256 tokens (needing 16 blocks), then blocks per sequence are reserved but unused: internal fragmentation . Effective concurrent sequences at the 256-token average: sequences actually needed simultaneously (22 waste factor).
With PagedAttention (on-demand allocation): Each sequence holds exactly blocks at any given step. At 256-token average generation: 16 blocks per sequence. Maximum concurrent sequences given 28,500 blocks: . Practical limit is set by other memory constraints and scheduling, but the point is that fragmentation no longer caps batch size at sequences.
With prefix caching: If all 100 concurrent requests share a 512-token system prompt: blocks, reference count 100. Those 32 blocks serve all 100 sequences. Without prefix caching: blocks. Savings: blocks GB freed for generation context.
What You Actually Decide in Practice
-
Block size is a tunable constant with real trade-offs. Smaller blocks (8 tokens) reduce internal fragmentation but multiply the size of the block table and increase pointer-chasing overhead during attention, which must now gather non-contiguous blocks. Larger blocks (32 tokens) are more cache-friendly and reduce metadata overhead but waste up to 31 token slots per sequence. In practice, 16 tokens per block is the right default for mixed-length workloads; adjust toward 32 if your sequence length distribution is long-tailed and toward 8 if you serve many very short requests.
-
Prefix cache eviction needs a policy. Shared prefix blocks are held in cache indefinitely under the assumption that the prefix is fixed. For multi-tenant serving where each tenant has a different system prompt, you need LRU eviction on prefix blocks, triggered when the free block pool drops below a watermark (e.g., 10% of total blocks). Evicting a prefix block invalidates all sequences that depend on it-force a KV recompute for those sequences or preempt them. Radix attention (used in SGLang) automates this with a prefix trie and LRU eviction at the subtree level.
-
CPU swap is a last resort, not a primary strategy. When the free block pool is exhausted, the scheduler can swap a paused sequence’s KV blocks to CPU DRAM over PCIe or NVLink. NVLink bandwidth is roughly 900 GB/s bidirectional (H100 SXM); PCIe 4.0 is roughly 32 GB/s. At MB per block and 16 blocks for a 256-token context: MB transferred over PCIe takes ms-tolerable once but not repeatedly. Prefer preemption-and-restart for short sequences (under 500 tokens) and swap for long sequences where the recompute cost exceeds the transfer cost.
-
Block table lookup is on the critical path in the attention kernel. The attention kernel must translate logical block indices to physical addresses for each sequence. This adds one level of indirection per block boundary. Verify that your attention kernel handles non-contiguous blocks natively (as vLLM’s
paged_attentionkernel does) rather than copying KV blocks into a contiguous staging buffer before attention-the latter negates the memory savings. -
Chunked prefill interacts with block allocation. When a 1,024-token prompt arrives and is chunked into 512-token pieces, the first chunk allocates blocks at step one, and the second chunk allocates another 32 blocks at step two, before any generation blocks are needed. Account for this prefill block reservation in your watermark-based eviction policy, otherwise a burst of long-prompt requests can exhaust the block pool before generation even starts.
How This Shows Up in Interviews
-
“Explain how PagedAttention works and why it matters.” Strong answer: describe block-based KV cache allocation, the block table mapping logical to physical blocks, copy-on-write for prefix sharing, and quantify the fragmentation reduction. Concrete numbers matter: pre-PagedAttention systems waste 20-40% of KV cache; with block size 16, internal fragmentation is at most 15 tokens per sequence.
-
“Why does KV cache fragmentation hurt throughput?” Strong answer: fragmentation reduces the number of concurrent sequences you can hold in GPU memory. Fewer concurrent sequences means lower arithmetic intensity for MLP layers, pushing the GPU toward its memory-bound regime. The throughput loss is multiplicative: halving the effective batch size roughly halves throughput at the operating point where the system is memory-bound.
-
Senior follow-up: “How does SGLang’s radix attention extend PagedAttention, and when does it help most?” Senior answer: radix attention maintains a prefix trie over all KV-cached sequences across all current and recent requests, not just those that share the same system prompt. Any common prefix subtree-including multi-turn conversation history and RAG-retrieved context-is automatically deduplicated. LRU eviction at the subtree level keeps the trie bounded. The gain is largest in multi-turn chat (where conversation history grows and is shared across turns) and RAG applications (where the same retrieved documents recur across users). The trade-off is trie management overhead and the risk of cache thrashing if prefix diversity is high relative to the block pool size.
Key Takeaways
-
PagedAttention divides the KV cache into fixed-size blocks (typically 16 tokens), allocated on demand and returned immediately when a sequence finishes-eliminating fragmentation that previously wasted 20-40% of KV cache capacity.
-
Copy-on-write lets multiple sequences share KV blocks for a common prefix at no extra memory cost; prefix caching turns a repeated system prompt from a per-request allocation to a single shared allocation.
-
The block table lookup adds one level of indirection in the attention kernel; use a kernel that handles non-contiguous blocks natively rather than copying to a staging buffer.
-
Block size, prefix eviction policy, and CPU-swap thresholds are the three key tunable parameters in production; defaults of 16 tokens/block, LRU eviction at 10% free watermark, and swap-for-long/restart-for-short cover most workloads.
-
Radix attention (SGLang) extends prefix caching to a full prefix trie, maximizing cache hits for multi-turn conversations and RAG pipelines where prefix overlap spans multiple users.