Serving

Every generated token costs a pass through the model, and a popular model generates trillions. Serving is the engineering of doing that cheaply for many users at once. This module starts from the KV cache's memory, explains why generation is limited by memory traffic rather than arithmetic, and covers the three techniques every modern serving system uses: continuous batching, paged attention and speculative decoding.

23.1 The Memory WallRun this section

The KV cache of section 12.2 is large at production scale. Its size in bytes is the product of every stored dimension:

• The $2$ counts keys and values.
• Each of the $n_{\text{layers}}$ layers has its own cache, with $n_{\text{kv\_heads}}$ key/value heads of width $d_{\text{head}}$.
• $L$ is the context length in tokens (not a layer count) and $B$ the batch size.
• $n_{\text{bytes}}$ is bytes per number, 2 for fp16/bf16.

Worked example: LLaMA-2 7B.

• $n_{\text{layers}} = 32$, $n_{\text{kv\_heads}} = 32$, $d_{\text{head}} = 128$, $L = 4096$, $B = 1$, $n_{\text{bytes}} = 2$ (fp16):

A GiB is $2^{30}$ bytes, so 2 GiB is about 2.15 GB.

Measured (lab 3a, experiments/kv_cache.py). The lab recomputes the total from the parameters and checks every factor. Two factors are easy to leave out, and leaving out either one halves the answer: the leading $2$ for K and V, and the precision term $n_{\text{bytes}} = 2$ at the end.

With 32 sequences, the cache alone needs $64.0 \text{ GiB}$, about 68.7 GB: almost three times an NVIDIA RTX 4090's 24 GB before a single weight is loaded, and far more than the model's own 13.5 GB of fp16 weights.

Worked example. LLaMA-3 8B: 32 layers, 8 KV heads (grouped-query attention, section 15.5), $d_{\text{head}} = 128$, $L = 8192$, fp16, $B = 1$ gives $2 \times 32 \times 8 \times 128 \times 8192 \times 2 = 1{,}073{,}741{,}824$ bytes $= 1.0$ GiB: twice the context of LLaMA-2 7B in half the cache, entirely from having 4 times fewer key/value heads.

In words. A calculator from layers, key-value heads, head width, context, batch and precision to the KV cache in gigabytes, beside the memory of common GPUs.

23.2 Prefill and DecodeRun this section

The two phases of section 12.1 stress the hardware in opposite ways:

• Prefill reads each weight once and uses it for every prompt token, so each byte feeds hundreds or thousands of multiply-adds: compute-bound.
• Decode reads every weight and the whole KV cache for one token per sequence, so each byte feeds one or two multiply-adds and the arithmetic units wait: memory-bound.

Arithmetic intensity is operations per byte read. A modern GPU does a few hundred operations in the time it reads one byte (section 16.1), so lower intensity is memory-limited. With 2-byte weights and $B$ tokens at once, a matrix multiplication's intensity is roughly $B$: about 1 for single-sequence decode.

Prefill mostly sets time to first token and decode the time per output token; systems often schedule the phases differently, some on separate machines. Since decode is limited by reading weights and cache, serving is about getting more tokens from each read. Sections 12.2 and 15.5 shrank the reads; the rest of this module shares them.

23.3 Continuous Batching

Applying the weights to many sequences per read costs little more than to one, so servers batch requests. In a fixed batch, everything waits for the longest sequence while finished slots sit idle. Continuous batching (Yu et al., 2022) re-forms the batch every step, so a waiting request takes a finished one's place at once, multiplying a busy server's throughput several times over.

23.4 Paged Attention

Many sequences mean many KV caches of unknown final length, and reserving each one's maximum wasted 60–80% of cache memory in earlier systems (Kwon et al., 2023, the vLLM paper). PagedAttention borrows virtual memory from operating systems: the cache lives in small fixed-size blocks, say 16 tokens, allocated as a sequence grows, with a per-sequence table mapping positions to blocks. Waste falls below 4%, so the same memory holds several times as many sequences, and sequences sharing a prefix, such as a long system prompt, can share its blocks.

In words. A pool of 96 blocks serving eight requests that arrive over time and grow token by token. Reserving each maximum fills the pool with empty blocks, so only four run. Paging hands out 16-token blocks as needed, so every request runs on arrival.

23.5 Speculative DecodingRun this section

Every token costs a full pass of the large model, however easy. Speculative decoding (Leviathan et al., 2023; Chen et al., 2023) lets a small draft model guess the next few tokens, and the large target model checks them all in one pass, which costs about the same as one token since decode is limited by reading weights. Each drafted token $x$ is accepted with probability

where $p$ is the target model's probability for that token and $q$ the draft model's. At the first rejection, a replacement is drawn from the leftover distribution

and the remaining drafts are discarded. The output distribution is then exactly the target's: tokens the draft over-predicts are thinned by the acceptance test, and under-predicted ones are made up by $p'$. A good draft gets two or three tokens accepted per pass, a two- to three-fold speedup with identical output.

In words. Target and draft probabilities over candidate tokens at each position. Stepping through four drafted tokens, each is accepted with probability the smaller of one and p over q; at a rejection a replacement is drawn from the leftover distribution and the rest are discarded. The readout counts tokens per large-model pass.

Exercises

Core Academic References

• Pope et al. (2023): Efficiently Scaling Transformer Inference on TPU v4. MLSys 2023. arXiv:2211.05102
• Yu et al. (2022): Orca: A Distributed Serving System for Transformer-Based Generative Models (continuous batching). OSDI 2022.
• Kwon et al. (2023): Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM). SOSP 2023. arXiv:2309.06180
• Leviathan, Kalman & Matias (2023): Fast Inference from Transformers via Speculative Decoding. ICML 2023. arXiv:2211.17192
• Chen et al. (2023): Accelerating Large Language Model Decoding with Speculative Sampling. arXiv:2302.01318
• Zhong et al. (2024): DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving. OSDI 2024. arXiv:2401.09670