All About Transformer Inference — Chapter 7 of How To Scale Your Model

TL;DR

Chapter 7 of the JAX/Google “How To Scale Your Model” series — a first-principles walkthrough of transformer inference math from a single sampled token through multi-slice serving. Establishes prefill vs generation as two structurally distinct workloads (compute-bound vs memory-bandwidth-bound), derives the critical batch size $B_\text{crit} = \beta \cdot \alpha_\text{hbm}$ (~240 tokens on TPU v5e, ~280 on H100) at which the linear/FFW layers cross into compute-bound territory, and shows attention during generation is stuck at intensity $\approx 1$ regardless of batch — so it’s always bandwidth-bound. Ends with sharding rules for both phases (prefill ≈ training; generation is model-parallel-only), KV-cache reduction levers (GQA/MQA, local attention, layer sharing, quantization, PagedAttention), and inference-engine designs (naive → interleaved → disaggregated prefill/decode with continuous batching).

Key claims

Naive sampling is $O(n^2)$ on FFW and $O(n^3)$ on attention to generate $n$ tokens; the KV cache reduces this to $O(n)$ FFW and $O(n^2)$ attention by never reprocessing prior tokens [§“Sampling is conceptually simple”].
Prefill and generation are two different workloads: prefill sees compute-bound linear layers (long sequences batch tokens along the sequence dim), generation only produces one token per request per step and must batch across requests to reach $B_\text{crit}$ [§“Prefill vs generation”].
Critical batch size for the linear/FFW blocks: $B_\text{crit} = (\text{bits per param}/\text{bits per activation}) \cdot (C/W_\text{hbm})$ ; ≈ 240 tokens on TPU v5e, ≈ 280 on H100 for bf16 activations; int8 weights halve it, int8 activations double it back [§“Takeaway: Transformer matmuls are compute-bound”].
Attention arithmetic intensity is $ST/(S+T)$ ; during prefill $S=T$ so intensity is $\Theta(T)$ (compute-bound above ~480 tokens), during generation $T=1$ so intensity $\approx 1$ regardless of $S$ or $B$ — attention is always memory-bandwidth-bound during generation [§“Multiheaded Attention Arithmetic Intensity”].
Bandwidth-bound decode lower bound: step time ≈ $(B \cdot \text{KV bytes} + \text{param bytes}) / \text{HBM bandwidth}$ ; general form adds a compute-bound term for the MLP when batch is large enough [§“Theoretical Step Time (General)”].
KV cache dominates memory quickly: for LLaMA 2-13B (40 layers, 40 KV heads, H=128) a single 8192-token cache at bf16 is ~6.7 GiB, so four caches exceed the parameter footprint; batch 240 needs 1.6 TiB [§“KV cache table”].
Reducing KV cache 5× (via 1:5 GMQA — 8 KV heads shared across 40 Q heads) turns theoretical per-chip throughput at batch=240 from 963 tok/s → 4529 tok/s on 8×TPU v5e and unlocks batch=64 within 128 GiB HBM; LLaMA-3 8B adopts exactly this ratio [§“With a smaller KV cache”].
KV-cache reduction levers, from small to large impact: GMQA/GQA (share KV heads across Q heads), local/sliding-window attention (cap the KV per local layer), cross-layer KV sharing (reduce storage but not necessarily bandwidth), quantization (int8/int4/fp8 on both weights and KV), ragged HBM reads / PagedAttention (don’t read padding); together give order-of-magnitude reductions [§“Tricks for improving generation throughput and latency”].
Sharding rules diverge sharply between phases: prefill allows essentially any training-time sharding (model parallelism up to ICI bound, then sequence parallelism); generation rules out FSDP (parameters are memory-bandwidth-bound, moving them via ICI is orders slower than HBM), rules out pure data parallelism (replicating params doesn’t help load them faster), rules out sequence sharding, and leaves only variants of model parallelism [§“General rules for sharding prefill”, “Sharding generation”].
Because generation is memory-bandwidth-bound, model parallelism can go beyond the training-era ICI-bound (up to $F/(B\alpha)$ ≈ 64-way for typical F=16384, B=32, α=8 on TPU v5e) to improve latency at minimal throughput cost — the FLOPs roofline stops binding [§“Note on ICI bounds for generation”].
KV caches are sharded along both head dim (Megatron, up to K-way) and batch dim; the cost is two AllToAlls per attention layer on Q and O activations, which is cheap relative to loading the stationary KVs [§“Sharding the KV cache”].
Inference engine designs on a Pareto: naive (batched-prefill-then-batched-generate) is simple but wastes TTFT and interleaves poorly; interleaved does batch-1 prefill + batched generation, high throughput but generation stalls during prefill; disaggregated prefill/decode on separate replicas gives predictable latency and lets each replica be sharded/quantized differently [§“Inference engine design”].

Method

This is Part 7 of the “How To Scale Your Model” scaling-book. The presentation is derivational: start with an FFW matmul roofline ( $T_\text{math} \geq T_\text{comms}$ → $B \geq C/W_\text{hbm}$ ), extend to attention using flash-attention-style intensity accounting ( $4BSTD$ flops over $4BSD + 4BTD$ bytes), then plug in real LLaMA 2-13B dimensions to concretize the memory picture. Sharding math follows the same skeleton used in the training chapters (Section 5), specialized to the memory-bound decode regime. A worked “pop quiz” computes step-time lower bounds for a 30B int8/bf16 model at batch 4 vs 256 on a TPU v5e 4×4 slice (2.5 ms vs 21 ms). The chapter references JetStream (github.com/google/JetStream) as the reference implementation of the interleaved and disaggregated designs.

Results

Not applicable — this is a tutorial, not an experimental paper. The reference numbers it produces (all analytical):

$B_\text{crit}$ ≈ 240 tokens (TPU v5e bf16), ≈ 280 tokens (H100 bf16).
LLaMA 2-13B: 13B params → 26 GB in bf16 → 32.7 GiB total memory at batch=1 with 6.7 GiB KV; 133.2 GiB at batch=16 (OOMs on 8×v5e which has 128 GiB HBM).
With 1:5 GMQA on the same architecture: 47.4 GiB total at batch=16, 4529 tok/s theoretical at batch=240.
30B int8/bf16 on 4×4 v5e: 2.5 ms/step at batch=4, 21 ms/step at batch=256.

Why it’s interesting

This chapter is the canonical entry point for the transformer-inference mental model that most of the LLM Inference Efficiency entries on the wiki either build on or push against. The $B_\text{crit}$ derivation is the same math Speculative Decoding: Performance or Illusion? uses to explain why speculative-decoding gains decay monotonically with batch size (once you’re compute-bound in the MLPs, drafting stops paying for itself), and the “attention is always bandwidth-bound during generation” observation is exactly what the KV-compression / GQA / MLA / paged-attention line of work is designed to attack — DeepSeek-V4.1-Flash — Pushing the Limits of KV Cache Compression‘s CED architecture, Kimi Linear: An Expressive, Efficient Attention Architecture, and Qwen3.8-Flash-Next — Hybrid Attention with QSA, Gated Residual, N-gram Embedding are all downstream of this diagnosis. The disaggregated-serving design point sketched here is what production engines like Mooncake and DistServe operationalize; the chapter’s JetStream reference is the JAX-side counterpart to vLLM/SGLang. Reiner Pope, one of the authors, has a companion talk filed as How GPT, Claude, and Gemini are actually trained and served – Reiner Pope.