Unit 5.5 · about 3 h
322 / 324Two regimes: prefill and decode
Latency and throughput numbers are meaningless without the regime behind them. Prefill scales with prompt length in parallel while decode pays per token in serial, so each needs its own bottleneck story.
- explain
- derive
You can now
- Explain prefill parallelism versus decode serialization and why one is compute-bound and the other bandwidth-bound.
- Derive MLP and attention intensities and predict the limiting regime for short prefill, long prefill, single decode and batched decode.
- Explain batch effects and phase interference between concurrent prefills and decodes.
Covers
- Prefill parallelism and prompt-length scaling
- Decode serialization and per-token work
- Compute, memory-bandwidth, launch and capacity limits
Study, in this order
3 steps
- CS336 Lecture 10: arithmetic intensity of inferenceStanford CS336 - Lecture 10 Inference
arithmetic_intensity_of_inference(naive \(O(T^3)\) without cache, MLP intensity \(B \times T\) with prefill large-\(T\) versus generation \(T=1\), attention intensity \(S \times T/(S+T)\) with prefill \(S/2\) versus generation below 1, batching helps MLP but not attention) - JAX scaling book inference: linear ops, attention and the step-time formulaHow to Scale Your Model - inference chapter
Linear operations (critical batch 240 on TPU v5e, about 280 on H100), What about attention (prefill \(T/2\) compute-bound, generation about 1 and always bandwidth-bound) and Theoretical estimates (step time as batch x KV plus params over bandwidth, with the batch 4 versus 256 Pop Quiz)
Depth · Read the three named subsections only. Skip sharding, the serving engine and worked problems; they belong to the later tracks.
- Modern LLM Notebook 21 Section 9: prefill versus decodeModern LLM Notebook - KV cache and inference acceleration
Section 9 table (prefill parallel and compute-bound with FlashAttention versus decode serial and bandwidth-bound with KV cache plus PagedAttention) with the 50-token prefill versus 20-step decode timing cell
Focus on
- Prefill processes the whole prompt at once like training, so long prompts saturate compute. Decode handles one token per request, so it streams weights and KV from memory.
- MLP intensity is \(B \times T\): batching requests helps decode. Attention intensity is about 1 at decode because each sequence owns its KV, so batching cannot fix it.
- Short prefills and single-token decodes can be launch-bound: tiny kernels cost more to start than to run. Long contexts and big batches turn capacity-bound when the KV no longer fits.
- A prefill running beside decodes steals the compute the decodes share. That interference is why serving separates or interleaves the phases.
Reproduce closed-book
Traps
- Expecting bigger batches to fix decode attention. Batching helps MLP reuse of weights, but per-sequence KV keeps attention bandwidth-bound.
- Calling decode compute-bound because FLOPs look small. Small FLOPs with large streaming reads is the definition of bandwidth-bound.
- Treating every small step as compute-bound. A 32-token prefill sits far below the roughly 240-token critical batch, so streaming the weights bounds it; for a tiny model, kernel-launch overhead can dominate even that.
Practice
JAX scaling book inference Pop Quiz: step-time bounds at batch 4 and 256
Pop Quiz under Theoretical estimates (30B dense, TPU v5e \(4 \times 4\) slice, int8 params with bf16 FLOPs, 8192 context, 100 kB per token KV; batch 4 bounds 2.5 ms memory-bound, batch 256 bounds 21 ms compute-bound)
Derive both bounds from the step-time formula by hand first, then check against the given answer.
Self-test
-
MLP intensity is \(B \times T\). Compare \(B=8\), \(T=1\) decode with \(B=1\), \(T=1024\) prefill against a critical batch near 240.
8 versus 1024. The decode step is memory-bound and the prefill is compute-bound inbf16. -
Attention intensity is \(S \times T/(S+T)\). Compare prefill with \(S=2048\) to decode with \(S=2048\) and \(T=1\).
Prefill gives \(S/2=1024\) (compute-bound) while decode gives about 1 (always bandwidth-bound). Batching does not raise the decode value because each sequence has its own KV. -
The JAX Pop Quiz bounds a 30B step at 2.5 ms for batch 4 and 21 ms for batch 256. Why does latency rise while throughput improves?
Batch 4 is memory-bound (batch x KV plus params over bandwidth). Batch 256 pushes the MLP into the compute-bound term, so each step costs more but carries 64x the tokens. -
A 32-token single-stream decode feels slow despite tiny FLOPs. Name two likely limiters besides FLOPs.
Kernel-launch overhead on tiny ops and memory bandwidth from streaming weights plus KV each step. Capacity is not the issue at this size.
Completion criterion
Stop when
You derive both Pop Quiz bounds within rounding and classify all four regimes with the intensity formulas.
Mark passed, then open 5.6 Only mark this when you can do it without notes.Leads to
- 5.6 Measuring time to first token and inter-token latency Measure the two regimes you just analyzed.
- 5.7 Beyond one cache: prefix sharing, paging and speculation See how paging and speculation change the picture.
Comes after
- 4.7 Memory, arithmetic intensity and reconciling estimatesCarry intensity regimes into separate prefill and decode measurement.
- 5.1 Why cache: recomputation in autoregressive decodingSee how the two phases stress hardware differently.
- 5.4 KV memory arithmeticUse the sizes to predict compute versus bandwidth limits.
Resources for this unit
- Modern LLM Notebook - KV cache and inference acceleration Use the KV-cache portions of English notebooks 11-mla-kv-cache and 21-inference-acceleration. Skip MLA and advanced acceleration as core work. Recalculate cache sizes and verify timing after warmup and explicit synchronization.
- Stanford CS336 - Lecture 10 Inference Use the Spring 2026 inference lecture to connect KV arithmetic, phase behavior and batching vocabulary. Keep kernels, distributed serving and expensive performance experiments for the later tracks.
- How to Scale Your Model - inference chapter After cache equivalence and phase measurement pass, read the inference chapter for arithmetic intensity, parallelism and production-scale reasoning. Treat it as the handoff to inference and performance tracks, not required foundation work.
Also in this section
- LLMs from Scratch companion - KV cache and GQA Use ch04/03_kv-cache and ch04/04_gqa with the chapter model code. Add explicit per-layer ownership, position offsets, variable lengths, EOS compaction and MHA-versus-GQA cache arithmetic to the runnable comparison.
- AI Performance Engineering - inference-engine bridge Use Sections 4 and 5 only after this gate to choose the next inference, KV-cache, scheduling, quantization or distributed-serving resource. Do not treat its GPU and kernel material as unfinished LLM Foundations work.
- No Magic - micro KV or paged-cache simulation Use one small CPU simulation such as microkv.py or micropaged.py only to visualize state ownership. Do not present Python timing as GPU or serving-performance evidence.