Inference rooflines

Inference is two workloads in one server [SB7]. Prefill processes the whole prompt in one pass and writes the KV cache; it looks like a training forward pass and is compute-bound. Decode generates one token per sequence per step, reading every weight and every active KV cache to do so; it is memory-bound. Time to first token (TTFT) is set by prefill. Time per output token is set by decode. Every serving design in Chapter 8 starts from this split.

The decode step

In decode, each linear layer multiplies a $[B, D]$ activation by a $[D, F]$ weight, where $B$ is the number of sequences in the batch. From Chapter 1, it is compute-bound when [SB7]

Weight quantization lowers the batch needed; faster low-precision math raises it again. This is formula 1.4; the Scaling Book writes the bf16-math case as β · C/W.

Computed by scripts/ch07_inference.py from scripts/specs.py.

Attention is different. From Chapter 4, decode attention has an arithmetic intensity of about $G$, the number of query heads per KV head, at any batch size. It is always bandwidth-bound. Put together, the Scaling Book's general step time is [SB7]

Every step reads every sequence's KV cache, and either reads every weight or does the matmul FLOPs, whichever takes longer. With N chips, C and W are summed over the chips.

For small batches, the lower bound reduces to bytes over bandwidth: $(B \cdot \text{KV} + \text{params}) / W_\text{hbm}$. A decode step can never be faster than reading the weights once.

KV growth makes the breakeven unreachable

At long context, two things stop the batch from reaching $B_\text{crit}$. The KV cache fills HBM before the batch gets there. And even if it fit, every extra sequence adds its whole cache to the step, so throughput saturates. The Scaling Book's LLaMA-2-13B tables on 8 v5e chips show both [SB7]:

• With standard multi-head attention, at 8k context, the model runs out of memory beyond batch 16.
• With a KV cache 5× smaller (GQA), batch 64 fits, and throughput keeps rising toward batch 240.

Prefill

Prefill pushes $T$ prompt tokens through every layer at once. Its linear layers have an effective batch of $T$, usually far above $B_\text{crit}$. Its attention has intensity about $T/2$, compute-bound above a few hundred tokens [SB7]. Prefill time is therefore FLOPs over achieved FLOP/s:

Prefill is a training forward pass without the backward: time scales with prompt tokens.

An 8,192-token prompt on LLaMA-3-70B takes about 0.91 s on 16 v5e chips at 40% MFU [SB8], and about 157 ms on 8 Ironwood chips derived. Because prefill is compute-bound, batching several prompts adds latency with little gain in throughput. The book recommends prefilling at batch 1 unless prompts are very short [SB7].

Disaggregated serving

Prefill wants compute and little memory. Decode wants memory bandwidth and HBM for many caches. They also want different sharding. Running them on separate servers lets each be sized for its job, at the cost of sending each request's KV cache from the prefill server to the decode server [SB7].

The server ratio follows from throughput. The Scaling Book's example: 8,192-token prompts take 0.91 s to prefill, and decode runs batch 32 at 19 ms per step for 512 tokens. Prefill servers then complete $P / 0.91$ requests per second and decode servers $32G / (0.019 \cdot 512)$. Matching them gives $P \approx 3G$: three prefill servers per decode server [SB8, Q12].

Sharding for decode

The training strategies of Chapter 6 mostly do not apply [SB7]:

• No FSDP. FSDP moves weights over ICI every step. In decode, a step's weight bytes over ICI take far longer than over HBM, so it costs a large multiple of the step. The book calls it "an easy and common way to have order of magnitude regressions".
• No data parallelism within a server. It just runs independent replicas; use separate replicas.
• No sequence parallelism. There is one new token per sequence.
• Model parallelism, further than in training. Each extra chip adds HBM bandwidth, which is what decode needs. ICI becomes the bottleneck only when

Shard the model until the per-chip weight read becomes shorter than the activation collectives.

On v5e, $\beta_\text{ici} \approx 8$, so a model with $F = 16{,}384$ at batch 32 can be split about 64 ways [SB7]. On Ironwood, 7,380 GB/s of HBM against 200 GB/s per axis gives $\beta_\text{ici} \approx 37$. LLaMA-3-70B at batch 32 can be split about 24 ways before ICI binds derived.

The KV cache. Shard it by KV heads first, which is free under head-sharded attention but limited to $K$ ways (8 for LLaMA-3). Beyond that, shard it by batch. That costs two AllToAlls per attention layer: one to move queries from head sharding to batch sharding, one to move outputs back [SB7].

Latency-bound collectives. Decode activations are small. A batch of 16 at $D = 8{,}192$ in int8 is 131 kB per collective. An 8-way collective is latency-bound below about 360 kB on v5e [SB7], and below about 800 kB on Ironwood derived. Decode collectives are almost always in this regime, where each costs about hops × 1 µs. This is the traffic TPU 8i's 7-hop Boardfly network and its Collectives Acceleration Engine target (Chapter 3).

Quantization on TPU

Fewer bytes per weight and per KV entry is the most direct lever on decode.

Where the low-precision rate equals the BF16 rate, quantization still saves HBM bytes and capacity; it just does not speed up the arithmetic.

Qwix is Google's JAX quantization library and the one tpu-inference uses for online quantization of JAX models [R-qwix] [R-tpuinf].

• It works by interception: while the model runs, it intercepts jax.lax.dot_general, jnp.einsum, conv_general_dilated and ragged_dot, and substitutes quantized versions according to rules.
• Each rule matches modules by regex and sets a weight and an activation dtype, for example QuantizationRule(module_path='.*', weight_qtype='int8', act_qtype='int8').
• It supports weight-only, dynamic-range and static-range schemes, native int4, int8 and fp8, and emulated narrower formats.
• Interception is disabled inside Pallas kernels, which have their own qwix.pallas helpers.
• In tpu-inference, an fp8 configuration is passed as --additional_config with weight_qtype and act_qtype set to float8_e4m3fn.

KV-cache quantization halves the attention term of formula 7.2. On a long-context workload that is most of the step: the 70B example above went from 14.0 ms to 7.0 ms at batch 64 when both weights and KV moved to fp8.

Speculative decoding

A small draft model proposes $k$ tokens; the large model scores all $k$ in one forward pass and accepts the longest correct prefix. Because a decode step is memory-bound, scoring $k$ tokens costs nearly the same as scoring one: the extra FLOPs come from idle matrix units, and the weights and KV are read once. With greedy decoding the output is unchanged, and probabilistic acceptance preserves the sampling distribution. The book's example of Chinchilla-70B with a 4B drafter finds the best lookahead at 3–4 tokens for summarization [SB7]. Speculation trades throughput for per-token latency when the batch is small. At long context it can help throughput too, because KV reads are shared across the speculated tokens.

MoE inference

From Chapter 4, an MoE needs $E/k$ times the batch to keep its expert matmuls compute-bound. The Scaling Book's chapter 7 model as an MoE with 16 experts and 2 active needs 1,920 tokens at bf16 on v5e [SB7, Q5]. Its KV cache per token is unchanged, and its weight bytes grow 12×. MoE decode is weight-bandwidth-bound and capacity-bound at almost any batch. Serving it pays two AllToAlls per layer for expert routing, at decode sizes that are latency-bound. All three are pressures behind TPU 8i's design.