Inference scaling checklist from the PaLM TPU v4 paper

Serving a large language model can be slowed by moving data as well as by doing arithmetic, and the right fix depends on which one is the limit. That is the lesson wafer_ai draws in part 7 of its series on AI performance engineering, a post that turns the paper "Efficiently Scaling Transformer Inference" by Reiner Pope and coauthors into a checklist. The paper studied PaLM inference on TPU v4 chips, and the post says it shows where weights, activations and KV caches should live, what must move at each step, and which transfers limit latency.

The first check is that serving has two phases. Prefill reads the prompt and can process many prompt tokens in one pass; decode then produces one new token per sequence per step. The post says the two phases can differ in arithmetic intensity, the amount of computation done per byte moved. Running more sequences together during decode spreads the cost of reading the model's weights across more work, but each sequence adds its own KV cache traffic. The KV cache is the stored attention state from earlier tokens, and it has to be read again at every step. The post therefore recommends varying batch size together with context length and profiling to see whether bandwidth or compute is the limit.

Adding chips has a similar catch. Splitting each layer across chips, called tensor parallelism, reduces the computation each chip does without removing the communication between them. In the paper's 1D layout, as the post describes it, the time spent aggregating activations stays roughly constant as chip count grows for a fixed workload; a 2D layout splits both weight dimensions so each chip sends smaller pieces. The post advises checking this against the actual interconnect and matrix sizes before raising the parallelism degree.

This explanation relies on wafer_ai's summary. The post says links are in the thread, but only this one post was collected, so the paper itself and the attached image were not read here. Some advice, such as checking KV cache replication in vLLM deployments, is the post's own extension of the paper to GPU serving.