The Papers Behind Modern Inference Optimisation

These three papers are unusually satisfying to read because their claims are checkable, their mechanisms are comprehensible, and their results are running in production behind essentially every hosted model. They also share a property that makes them a set rather than a list: none of them changes what the model computes.

One bottleneck, three attacks

Generation is memory-bound, not compute-bound. Modern accelerators can do arithmetic far faster than they can move data between high-bandwidth memory and the small fast memory on the chip, so for token-by-token decoding the limiting resource is memory rather than FLOPs. Once you hold that, the three papers stop looking like separate tricks:

• FlashAttention attacks memory traffic — how much data moves per attention operation.
• PagedAttention attacks memory capacity — how many requests fit at once, which is what throughput is made of.
• Speculative decoding attacks memory bandwidth per token — how many tokens you get per pass over the weights.

And all three are exact. FlashAttention computes the same attention. PagedAttention stores the same values. Speculative decoding provably preserves the target model’s output distribution. That is what makes them a different category from quantisation or distillation, which trade quality for speed and require you to evaluate what you lost.

FlashAttention: memory traffic

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022) observes that standard attention implementations materialise the full attention matrix in high-bandwidth memory. For a sequence of length N that matrix is N-by-N, and writing it out and reading it back is the dominant cost — not the multiplications.

The method tiles the computation so blocks of queries, keys and values are loaded into on-chip memory, the attention for that tile is computed there, and only the output is written back. Making this work requires computing the softmax without ever having the whole row, which is done with a running normalisation, and requires recomputing some intermediate values during the backward pass rather than storing them — deliberately doing more arithmetic to move less data.

The framing is the contribution as much as the kernel. The paper argues that attention implementations should be analysed for their IO cost against the memory hierarchy, not for their FLOP count, and that argument reorganised how the field thinks about efficiency. Later versions improved work partitioning across the parallel units of the accelerator and adapted to newer hardware generations. Note also what it does not do: it does not change the quadratic scaling of attention with sequence length, only the constant. Papers claiming to fix quadratic attention are a separate and much more contested literature.

PagedAttention: memory capacity

Efficient Memory Management for Large Language Model Serving with PagedAttention (2023), the vLLM paper, is a systems paper and reads like one. The observation is that serving systems allocated the KV cache as one contiguous block per request, sized for the maximum possible output length. A request that generates fifty tokens against a two-thousand-token reservation wastes the difference, and fragmentation wastes more.

The fix is borrowed wholesale from operating systems: split the cache into fixed-size blocks, keep a block table mapping a request’s logical positions to physical blocks, and allocate on demand. Blocks need not be contiguous. Internal waste is bounded by one block per request instead of by the length you guessed.

Two consequences follow, and the second is the one people miss. Better packing means more concurrent requests in the same memory, and concurrency is throughput — the paper reports large throughput gains over prior serving systems at comparable latency. And because blocks are indirected through a table, they can be shared: several sequences generated from one prompt can point at the same prefix blocks, with copy-on-write when they diverge. That mechanism is the direct ancestor of prefix reuse across requests, which is what you are exploiting when prompt caching discounts a shared prefix.

Speculative decoding: bandwidth per token

Two 2023 papers, from different groups, describe the same idea: speculative decoding, or speculative sampling. Generating one token requires reading the entire model from memory. Reading the model for one token and reading it for five costs nearly the same, because the bottleneck is the read, not the arithmetic.

So: have a small fast draft model propose several tokens; run the large model once over the whole proposed continuation, scoring all positions in parallel; accept the prefix that the large model agrees with, and correct at the first disagreement. The subtle part is the acceptance rule, which uses a modified rejection sampling scheme so that the distribution of accepted tokens is provably identical to sampling from the target model directly. You are not approximating the big model. You are getting its exact output distribution with fewer passes over its weights.

The speedup depends entirely on the acceptance rate, which depends on how well the draft model imitates the target on your traffic. That is why the technique helps enormously on predictable text and much less on text where the models disagree, and why later work explored drafting with lightweight heads attached to the target model rather than a separate model. See how speculative decoding behaves in practice.

How to read this literature

Systems papers reward a different reading procedure from method papers, and it is a simpler one.

• Find the resource being conserved. Every one of these papers has a single sentence identifying the bottleneck. Locate it and the rest of the paper is a consequence.
• Check whether the result is exact. An exact optimisation needs no quality evaluation. An approximate one does, and if a paper offering a large speedup does not evaluate output quality, find out why.
• Read the hardware assumptions. Results are stated for a specific accelerator generation with a specific memory hierarchy. Ratios shift with hardware, and a technique tuned for one generation may be less compelling on the next.
• Distinguish throughput from latency. Serving papers usually optimise aggregate throughput; a change that improves tokens per second across a busy server can leave a single request no faster. Both matter and they are different quantities.