Efficiently Scaling Transformer Inference
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, et al.2022MLSys 2023
Read it on arxiv.org(opens in a new tab)Why this one
Read this after Attention Is All You Need and before Efficient Memory Management for Large Language Model Serving with PagedAttention. The useful idea is not a trick kernel, it is a cost model for deciding how to split a huge decoder across devices when batch size, sequence length, and latency target all pull in different directions. People often talk about tensor parallelism, pipeline parallelism, and data parallelism as fixed architectural choices. This paper treats them as knobs whose best setting changes between prefill and decode, and between one impatient user and a full batch. That is the mental model you need when an LLM serving stack is slow but the GPUs are not obviously idle. It is worth your evening because it turns inference from folklore into arithmetic: bytes moved, FLOPs used, cache size, and the shape of attention.
What to take away
- Decode is often memory bandwidth bound, while prefill can use much higher compute utilization.
- The right partitioning changes with batch size, so one serving layout will not fit every latency target.
- Multi-query attention matters operationally because it shrinks KV cache traffic and unlocks longer contexts.
Reads with
- Attention Is All You Need
the decoder cost model only lands once self attention and kv caching are familiar
- Efficient Memory Management for Large Language Model Serving with PagedAttention
turns this paper's kv cache arithmetic into a serving memory manager
- Fast Transformer Decoding: One Write-Head is All You Need
attacks the same decode bottleneck by shrinking kv cache bandwidth instead of repartitioning work