Inference scaling checklist drawn from the PaLM TPU v4 paper
- Source
- wafer_ai
- 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…
- The thread warns that multiquery with a head-sharded layout can replicate the on every device, so vLLM users with grouped-query attention should check whether tensor parallelism exceeds the KV-head count before estimating per-GPU capacity.
- KV cache size and prefill attention scratch memory are separate problems: sharing KV heads does not shrink the full attention-score tensor, which grows with the square of sequence length. The paper points to microbatching and FlashAttention for the latter.
- PaLM's setup stored weights in int8 but kept bfloat16 matrix math, which cuts weight-loading traffic yet gives no speedup for compute-bound batches, so the thread says to pick what to quantize based on the measured bottleneck.
- The authors overlapped collectives with matrix multiplication using Looped CollectiveEinsum, but the thread cautions that an asynchronous collective does not prove its cost is hidden; check a GPU timeline for exposed communication and stalls.
- For serving benchmarks, the thread recommends tracking time to first token, inter-token latency and throughput under target load, keeping queueing time separate from prefill time, and remeasuring at the same latency target after each change.
- inference — Running a trained model to get answers — the phase where AI is actually used, as opposed to trained.
- context window — The maximum amount of text a model can consider at once — its working memory for the current conversation or task.
- 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.
Gives a checklist for locating bandwidth versus compute limits in LLM serving, and for choosing tensor-parallel layouts based on interconnect and matrix dimensions before scaling chip count.
postwe launched the most comprehensive ai performance engineering repo in the world follow and save to keep up with the seriwafer_ai
postWafer maps a classic GPU textbook onto real AI performance engineering workwafer
articleHow to Size GPUs for AI Inference and TCO Without OverspendingElizabeth Goodman
Checking sign-in…
Loading comments…