# Module 23: Serving

> LLM serving: the KV-cache memory wall, prefill versus decode, continuous batching, paged attention as in vLLM, and speculative decoding.

- Part VI: Deployment · core track · Advanced · 4 hours
- By the end you can size a serving deployment and choose the techniques that raise throughput.
- Builds on: [Module 12: Generation](https://transformers-and-math.drmhse.com/md/12-generation.md), [Module 15: The LLaMA Block](https://transformers-and-math.drmhse.com/md/15-the-llama-block.md)
- Web page (interactive figures, labs and equation explainers): https://transformers-and-math.drmhse.com/modules/23-serving
- From *Transformers & Math* by Mike Chumba: https://transformers-and-math.drmhse.com

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 Wall

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

$$
\text{Memory}_{\text{KV}} = 2 \times n_{\text{layers}} \times n_{\text{kv\_heads}} \times d_{\text{head}} \times L \times B \times n_{\text{bytes}}
$$

- 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`):
$$
\text{Memory}_{\text{KV}} = 2 \times 32 \times 32 \times 128 \times 4096 \times 2 = 2{,}147{,}483{,}648 \text{ bytes} = 2.0 \text{ GiB per single sequence}
$$

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.

> **Interactive figure** `kv-cache-calculator`, on the web page: https://transformers-and-math.drmhse.com/modules/23-serving#figure-kv-cache-calculator
>
> **Try:** Set up LLaMA-3 8B at 8,192 tokens, then double the context. Before you look, predict what happens to the cache, and then what happens when you double the batch size instead.
>
> **What happens:** Both double it: the cache is a product of every dimension in the formula, so doubling any one factor doubles the total. Context length and batch size are the two that a serving system controls at run time, which is why they decide how many users fit on a GPU.

> **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 Decode

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.

> **Interactive figure** `paged-attention`, on the web page: https://transformers-and-math.drmhse.com/modules/23-serving#figure-paged-attention
>
> **Try:** Eight requests arrive with unknown final lengths. Before switching from reserving each request's maximum to paging, predict how many of them can run at once in each mode.
>
> **What happens:** With reservation, only four fit: each holds its maximum length of memory from the start, mostly empty. With paging, blocks are handed out as tokens are generated, so all eight run as soon as they arrive and almost no memory sits reserved but unused.

> **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 Decoding

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

$$
\min\left(1, \, \frac{p(x)}{q(x)}\right)
$$

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

$$
p'(x) = \frac{\max\left(0, \, p(x) - q(x)\right)}{\sum_{x'} \max\left(0, \, p(x') - q(x')\right)}
$$

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.

> **Interactive figure** `speculative-stepper`, on the web page: https://transformers-and-math.drmhse.com/modules/23-serving#figure-speculative-stepper
>
> **Try:** The draft model proposes four tokens. Before stepping, predict what happens to the remaining drafts once one is rejected, and whether the final text can ever differ in distribution from sampling the large model directly.
>
> **What happens:** Everything after the first rejection is discarded and replaced by one token drawn from the leftover distribution. The acceptance rule and the leftover distribution together make the output distribution exactly the large model's, whatever the draft model does; a bad draft only costs speed.

> **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.

> **Browser lab** `speculative-exact`: an experiment that runs in the reader's browser on the web page: https://transformers-and-math.drmhse.com/modules/23-serving#lab-speculative-exact

## Exercises

**Exercise (apply).** What is the KV cache for 32 sequences of 4,096 tokens on LLaMA-2 7B (32 layers, 32 KV heads, head width 128, fp16), in GiB?

*Answer:* 64 (within 0.01)

*Solution:* One sequence is $2 \times 32 \times 32 \times 128 \times 4096 \times 2 = 2^{31}$ bytes $= 2$ GiB, so 32 sequences need 64 GiB.

**Exercise (apply).** The large model gives a drafted token probability $p = 0.3$ and the draft model gave it $q = 0.5$. What is the chance it is accepted?

*Answer:* 0.6 (within 0.0001)

*Solution:* $\min(1, p/q) = \min(1, 0.6) = 0.6$. The draft over-predicted this token, so it is thinned out to the rate the large model wants.

**Exercise (apply).** If each drafted token is accepted independently with probability $\alpha = 0.8$ and the draft proposes $\gamma = 4$ tokens, the expected number of tokens produced per large-model pass is $\frac{1 - \alpha^{\gamma+1}}{1 - \alpha}$. Evaluate it.

*Answer:* 3.3616 (within 0.001)

*Solution:* $\frac{1 - 0.8^5}{0.2} = \frac{1 - 0.32768}{0.2} = 3.36$ tokens per pass, instead of 1.

**Exercise (check).** Why is decoding one sequence memory-bound while prefilling a long prompt is compute-bound, on the same hardware?

*Solution:* Both read every weight once per pass. Prefill uses each weight for every prompt token in the pass, hundreds or thousands of operations per byte read, so arithmetic is the limit. Decode uses each weight for just one token per sequence, about one operation per byte read, so the GPU waits on memory.

**Exercise (extend).** A system reserves each request's maximum length of KV cache up front. Explain two separate ways this wastes memory, and how paging fixes both.

*Solution:* First, internal waste: a request reserved for 4,096 tokens that finishes after 300 leaves most of its reservation unused for its whole life. Second, fragmentation: freed reservations leave gaps of different sizes that a new request may not fit into. Paging allocates small fixed-size blocks only as tokens are actually generated, so unused space is at most one partly filled block per request, and any free block can serve any request.

### Core Academic References
- **Pope et al. (2023)**: *Efficiently Scaling Transformer Inference on TPU v4*. MLSys 2023. [arXiv:2211.05102](https://arxiv.org/abs/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](https://arxiv.org/abs/2309.06180)
- **Leviathan, Kalman & Matias (2023)**: *Fast Inference from Transformers via Speculative Decoding*. ICML 2023. [arXiv:2211.17192](https://arxiv.org/abs/2211.17192)
- **Chen et al. (2023)**: *Accelerating Large Language Model Decoding with Speculative Sampling*. [arXiv:2302.01318](https://arxiv.org/abs/2302.01318)
- **Zhong et al. (2024)**: *DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving*. OSDI 2024. [arXiv:2401.09670](https://arxiv.org/abs/2401.09670)