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 splits both weight dimensions so each chip communicates smaller activation shards. evaluate this against the actual interconnect and matrix dimensions before increasing the parallelism degree. - the best tensor to move changes with the token batch. weight-stationary layouts exchange activations; weight-gathered layouts transfer more weights to reduce activation exchange. large prefills can justify that trade because their activation tensors are large. count the tokens processed per forward pass, including prompt tokens during prefill, when comparing those layouts. - multiquery attention shares one KV head across query heads, but a head-sharded layout can replicate that cache across devices. the paper shards attention across the batch during decode and exchanges current-step activations to reduce cache reads per chip. in vLLM's grouped-query attention deployments, check whether tensor parallelism exceeds the KV-head count and introduces cache replication before estimating per-GPU capacity. - KV cache capacity and prefill attention intermediates are separate memory problems. at a fixed batch size and query-head count, sharing KV heads does not shrink the full attention-score tensor. for dense prefill, that tensor's element count grows with the square of sequence length. Appendix G discusses microbatching to reduce temporary allocations and FlashAttention to avoid materializing the full score matrix. use that distinction to identify whether a memory failure comes from retained KV state or temporary attention allocations. - communication overlap depends on the execution schedule and tensor layout. the authors use Looped CollectiveEinsum to overlap collectives with matrix multiplication, including choosing output sharding that exposes overlap opportunities. inspect a GPU timeline for exposed communication and dependent compute stalls; an asynchronous collective does not establish that its cost is hidden. - quantization changes different bottlenecks depending on what is quantized. the paper stores weights in int8 while retaining bfloat16 matrix arithmetic. this reduces weight-loading traffic, but does not provide faster arithmetic for compute-bound batches. evaluate weight, activation, and KV cache precision against the measured bottleneck and benchmark the effect on both phases. - benchmark the operating point that the application needs. the paper compares latency with accelerator-time per token and model FLOPS utilization. in a serving system, track time to first token, inter-token latency, and throughput under the target load. keep queueing time separate from prefill time. the practical use is to narrow an optimization experiment: estimate compute, HBM traffic, and collective costs for the actual workload, inspect the bottleneck, change the relevant layout or precision, and remeasure at the same latency target.