Inference scaling checklist drawn from the PaLM TPU v4 paper
- Source
- x.com
- Date
we gathered all the resources you'll ever need to become the most cracked ai performance engineer follow and save to keep up with the series. links in thread 🧵 part 7: Efficiently Scaling Transformer Inference Reiner Pope and coauthors studied PaLM inference on TPU v4. the paper gives performance engineers a way to reason about where weights, activations, and KV caches should live, what must move each step, and which transfers limit latency. the paper's mechanisms suggest these checks for an inference deployment: - prefill and decode can have different arithmetic intensity, meaning computation per byte transferred. prefill can process many prompt tokens in one pass; standard autoregressive decode produces one new token per sequence per step. larger decode batches amortize weight reads, but add sequence-specific KV cache traffic. sweep batch size with context length and profile the linear layers and attention to locate the bandwidth or compute limit. - tensor parallelism can reduce local compute without removing the communication bottleneck. in the paper's 1D layout, activation-aggregation time stays roughly constant as chip count grows for a fixed workload. 2D partitioning…
Context
Running a large language model for users has two phases that strain hardware differently, which is why more chips or bigger batches do not always speed things up. In the PaLM inference paper, Google researchers analyze the phases separately: reading the prompt processes many input tokens in parallel, while generating the answer produces one token at a time for each sequence.
At small batch sizes the time to load the model's weights dominates, so processing more sequences together spreads that cost across more work. At larger batches and longer contexts, each sequence's stored state (the ) must also be read, and the time to load it becomes the bottleneck. Splitting the model across more chips does not help with communication, which the paper says stays roughly constant regardless of chip count and so matters more as chips are added. Batch size, , and chip layout therefore interact, and adding hardware alone may leave the limiting data transfer in place.
- inference — Running a trained model to get answers — the phase where AI is actually used, as opposed to trained.
- attention — The mechanism that lets a model weigh which earlier words matter for the word it's currently processing — the core operation of a transformer.
- KV cache — The memory a model keeps about text it has already read, so generating each new token doesn't require reprocessing the whole conversation.
- context window — The maximum amount of text a model can consider at once — its working memory for the current conversation or task.
Checking sign-in…
Loading comments…