LLM inference serving: KV cache, batching, paging and speculative decoding

Before this lesson: the roofline, the ridge point and the matmul threshold $B^*$ (lesson 2); merging partial attention results and split-KV decoding (lesson 4); parameter counts (lesson 6); weight-only (W4A16) and W8A8 quantization (lesson 11). Text generation with a KV cache and grouped-query attention (deep learning 11, deep learning 9). For the proof in section 3: discrete distributions and expectation (probability 2). The model view of the KV cache and of decoding will be in the LLMs track (in preparation). Time: about 85 minutes, plus the notebook.

---

Why this matters #

For a widely used model, the total cost of serving can be larger than the cost of training. Each generated token reads every weight and the whole KV cache of its sequence from HBM, so the cost of serving is set by bytes, and by how many sequences share one read of the weights. Paged KV memory raised the throughput of LLM servers by 2–4× at the same latency, compared with FasterTransformer and Orca (Kwon et al., 2023). Speculative decoding gave 2–3× at small batch, with the same output distribution (Leviathan et al., 2023).

The open-source engines vLLM, SGLang and TensorRT-LLM combine all the methods of this lesson. The same numbers decide how many GPUs a chat service needs, how fast RL rollouts (many samples per prompt) run, and which quantization format to choose.

1. Predict first #

2. The picture #

Two phases. Generation starts with a prefill: one forward pass over the whole prompt. It writes the keys and values (K and V) of every prompt token into the KV cache and gives the first output token. Then decode runs one step per output token. Each step feeds one new token per sequence, appends its K and V to the cache, and attends over everything cached so far.

Simplest case: one sequence. One decode step multiplies one token vector by every weight matrix. Each weight is read from HBM once and used for one multiply-add: 2 FLOPs per 2 bytes in bf16, an intensity of 1 FLOP/byte. The ridge point of an H100 is about 295 FLOP/byte (lesson 2), so the tensor cores are idle more than 99% of the time. The step time is the time to read the weights: $16.1\ \text{GB} / 3.35\ \text{TB/s} \approx 4.8$ ms for an 8B model.

Batching shares the read. With $B$ sequences in one step, each weight that is read is used $B$ times: the same bytes, $B$ times the useful work. But each sequence brings its own cache. Its K and V serve only that sequence, so their bytes grow with the batch, and they need memory. The batch is limited by the free memory divided by the cache per sequence.

Interactive figure: loading…

A server is a memory manager plus a scheduler. The scheduler (continuous batching) decides at every step which sequences run. The memory manager (PagedAttention) gives KV memory to sequences as they grow. Speculative decoding attacks a different limit: the number of sequential steps per token.

3. Definitions and main results #

Notation (as in the whole track). $L$ layers; $a$ query heads; $n_{kv}$ key/value heads ($n_{kv} = a$ for multi-head attention, MHA; fewer for grouped-query attention, GQA; $n_{kv} = 1$ for multi-query attention, MQA); $g = a / n_{kv}$ query heads per KV head; $d_h$ head dimension; $h$ model width; $e$ bytes per stored number (2 for bf16); $\Psi$ parameters; $W = e\Psi$ bytes of weights; $B$ sequences in a decode step; $C$ cached tokens per sequence (the context); $P$ peak FLOP/s; BW the HBM bandwidth in bytes/s; $R = P/\mathrm{BW}$ the ridge point. H100 SXM: $P = 989$ TFLOP/s (dense bf16), BW $= 3.35$ TB/s, $R \approx 295$ FLOP/byte, 80 GB of HBM.

KV-cache memory and the batch that fits #

Definition. The KV cache stores, for every past token, one key vector and one value vector of length $d_h$ per KV head and per layer. So one token needs

The 2 counts K and V. Let $M_{kv}$ be the memory left for the cache: the GPU memory minus the weights, the activations and the runtime buffers. Then $M_{kv}/k$ tokens fit, and at context $C$ the batch is at most

GQA divides $k$ by $g$ (compared with MHA of the same shape); an fp8 cache halves $e$.

Example. Llama-3-8B has $L = 32$, $n_{kv} = 8$, $d_h = 128$. In bf16, $k = 2 \cdot 32 \cdot 8 \cdot 128 \cdot 2 = 131{,}072$ B $= 128$ KiB per token. Llama-2-7B is about the same size but uses MHA with 32 KV heads: 512 KiB per token, 4× more.

Proposition (decode step: bytes, FLOPs, intensity and time) #

A decode step with $B$ sequences of context $C$, in bf16 ($e = 2$ for the weights and the cache), reads at least $W + BCk$ bytes and does about $B \cdot W + g \cdot BCk$ FLOPs. Its intensity is a byte-weighted average of $B$ and $g$:

Since $g \ll R \approx 295$ in common models ($g = 4$–8 with GQA, $g = a$ with MQA), the step is memory-bound unless $B > R$ and the weights are most of the bytes.

Why this is true, in one sentence. The weights are shared by the whole batch, so each weight byte does $B$ units of work; the cache is private, so each cache byte does only $g$ units, whatever the batch is.

Proof.

1. Weights. Every weight matrix is read once per step, whatever $B$ is: $W$ bytes. Each weight takes part in one multiply and one add per sequence: $2\Psi B$ FLOPs. In bf16, $W = 2\Psi$, so this is $B \cdot W$.
2. Cache. Each sequence reads its own K and V, $Ck$ bytes, so $BCk$ bytes in total. Nothing is shared across the batch.
3. Attention FLOPs. Per sequence, layer and query head, the scores $q^\top K$ cost $2Cd_h$ FLOPs and the weighted sum of the values another $2Cd_h$. In total $4Cd_h \cdot a \cdot L$ per sequence. The cache of that sequence has $Ck = 4CLn_{kv}d_h$ bytes (with $e = 2$). The ratio is $a/n_{kv} = g$: attention does $g$ FLOPs per cache byte.
4. Intensity. Divide the FLOPs of steps 1 and 3 by the bytes of steps 1 and 2. The weights contribute intensity $B$ and the cache contributes $g$, each weighted by its share of the bytes.
5. Bound. The roofline gives $t \ge \max(\text{bytes}/\mathrm{BW},\ \text{FLOPs}/P)$, and the step is memory-bound exactly when $I < R$. Since $g \ll R$, the average reaches $R$ only if $B > R$ and $W$ is most of the bytes. $\blacksquare$

Left out: the query and output vectors, norms and activations ($O(Bh)$ bytes, small) and the KV writes of the new tokens. With other formats, the weight term becomes $2B/e_w$ FLOP per weight byte and the cache term $2g/e_{kv}$: int4 weights ($e_w = 1/2$) make weight bytes 4× fewer and their intensity 4× higher.

When are the weight matmuls compute-bound? By the matmul result of lesson 2, a $(B \times D) \cdot (D \times F)$ matmul is compute-bound from $B^* = R\,DF/(DF - R(D + F))$, about 300 rows for large $D, F$. For Llama-3-8B: $B^* = 344.9$ for the $4096 \times 4096$ projections and $B^* = 325.4$ for the $4096 \times 14{,}336$ MLP matrices. A batch of 326–345 sequences fits in memory only with short contexts (section 4).

Prefill is the opposite. A prompt of $T$ tokens goes through each weight matrix as one $(T \times D) \cdot (D \times F)$ matmul, so its intensity is about $T$ (for $T \ll D, F$). A prompt of more than a few hundred tokens is compute-bound. Its time is about $2\Psi T / P$ plus the attention FLOPs, about $2 L a d_h T^2$ with a causal mask.

Latency and throughput: TTFT, TPOT, goodput #

• TTFT (time to first token) = the time in the queue + the prefill. It grows with the prompt length.
• TPOT (time per output token; also called inter-token latency) = the time between two later tokens of one sequence, about one decode step.
• Throughput = output tokens per second over all sequences, $B / t_{\text{step}}$. It sets the cost per token.
• Goodput = the number of requests per second that finish within their latency targets (service-level objectives, SLOs), for example TTFT ≤ 1 s and TPOT ≤ 50 ms for 99% of requests.

A larger $B$ raises throughput, but each step reads more cache, so TPOT grows too: the curve in the figure above. A server chooses the largest batch whose TPOT still meets the SLO.

Continuous batching #

Static (request-level) batching starts a batch together and runs it until its longest request ends. Finished sequences keep their slot and idle. New requests wait for the whole batch (head-of-line blocking), so TTFT suffers.

Continuous batching (iteration-level scheduling; Orca, Yu et al., OSDI 2022) lets the scheduler decide again before every step. A finished sequence leaves at once; a waiting request joins at the next step, and its prefill runs in that step. The non-attention layers treat the batch as one flat list of tokens, so prefill and decode tokens of different sequences share one matmul (Orca's "selective batching"). Attention runs per sequence, so its kernel must accept ragged lengths, and the KV memory of any sequence must be able to grow at any step. Paging gives that.

Interactive figure: loading…

PagedAttention #

Earlier servers stored the cache of each request as one contiguous tensor, sized for its maximum length. The vLLM paper (Kwon et al., 2023, §3) names three kinds of waste:

• reserved slots for future tokens, held for the whole life of the request;
• internal fragmentation: slots up to the maximum length that the request never uses;
• external fragmentation: gaps between buffers of different sizes, left by the allocator.

In their measurements, only 20.4–38.2% of the KV memory of existing systems held actual token states; vLLM reached 96.3%.

PagedAttention (§4.1–4.3) applies the idea of virtual memory:

• The KV memory is a pool of fixed-size physical blocks of $n_\text{blk}$ tokens each (vLLM's default is $n_\text{blk} = 16$).
• Each sequence has a block table: logical block $i$ (tokens $i \cdot n_\text{blk}$ to $(i+1) \cdot n_\text{blk} - 1$) maps to a physical block id anywhere in the pool.
• A sequence gets a new block only when its last block is full.
• The attention kernel follows the block table and reads K and V one block at a time. The partial results of the blocks are merged with the online-softmax update of lesson 4, so no contiguous copy is needed.

Waste bound. A sequence of $n$ tokens uses $\lceil n/n_\text{blk} \rceil$ blocks, so it wastes at most $n_\text{blk} - 1$ slots, all in its last block. All blocks have the same size, so there is no external fragmentation.

Sharing (§4.4). A physical block can belong to several sequences; a reference count says how many. Parallel sampling (several answers to one prompt) and beam search map their prompt blocks to the same physical blocks.

• Copy-on-write: before a sequence writes into a block whose reference count is above 1, it gets a private copy, and the count of the shared block drops by 1. The last sharer finds a count of 1 and writes in place.
• Prefix caching: a request whose prompt starts with the same tokens as cached blocks (a system prompt, earlier chat turns) maps its first logical blocks to them and skips their prefill. The paper reserves blocks for known shared prefixes; later vLLM versions find a cached full block by a hash of its tokens and of all tokens before it.

Preemption (§4.5). When no block is free, some sequence must stop. vLLM serves requests first-come-first-served, so it preempts the newest request first: the oldest requests keep their progress. It evicts all blocks of a sequence or none, because a sequence needs all its blocks at every step. To restore the sequence later, it either swaps the blocks back from CPU memory or recomputes them with one prefill over the prompt plus the tokens generated so far.

Interactive figure: loading…

What paging costs. Reading through a block table made the attention kernel 20–26% slower (vLLM §7.1). The larger batch more than pays for it.

Theorem (speculative sampling) #

Setup. A cheap draft model gives next-token distributions $q$; the target model gives $p$. In one round, the draft proposes $\gamma$ tokens $x_1, \dots, x_\gamma$, one after the other, each sampled from $q$. The target scores all $\gamma + 1$ positions in one forward pass, which gives $p$ at each position. For $i = 1, \dots, \gamma$: accept $x_i$ with probability $\min(1,\ p(x_i)/q(x_i))$. At the first rejection, sample a replacement token from the residual distribution

and end the round. If all $\gamma$ tokens are accepted, sample one bonus token from the target's distribution at position $\gamma + 1$.

(a) Every emitted token has exactly the distribution $p$, given the same prefix.

(b) If each draft token is accepted with probability $\alpha$, independently, the expected number of tokens per target pass is

In this lesson, $\alpha$ is the acceptance rate and $\gamma$ the draft length, as in the papers. This $\alpha$ is not the message latency of lesson 7.

Small example first. Take two tokens, target $p = (0.6, 0.4)$ and draft $q = (0.9, 0.1)$.

• The draft draws $x_1$ with probability 0.9 and keeps it with probability $0.6/0.9 = 2/3$. So $x_1$ is drawn and kept with probability $0.9 \cdot 2/3 = 0.6$.
• The draft draws $x_2$ with probability 0.1 and always keeps it ($p/q = 4 > 1$): probability 0.1.
• A rejection happens with probability $0.9 \cdot 1/3 = 0.3$. The residual is $\max(0, p - q) = (0, 0.3)$, normalised to $(0, 1)$: after a rejection the output is always $x_2$.
• In total: $P(x_1) = 0.6$ and $P(x_2) = 0.1 + 0.3 = 0.4$. This is exactly $p$.

The draft over-proposes $x_1$; the rule throws away exactly the surplus and gives it to the token the draft under-proposes.

Interactive figure: loading…

Proof of (a) (one position; $p$ and $q$ are the distributions at that position).

1. Token $x$ can be emitted in two ways: the draft draws $x$ and it is accepted, or a draft token is rejected and $x$ is drawn from $p'$.
2. Drawn and accepted:

3. The probability of a rejection is one minus the probability of a kept draft. Because $\sum_y p(y) = 1$,

4. This is exactly the normaliser of $p'$, so it cancels:

5. Add the two ways:

(If $p(x) \le q(x)$, the sum is $p(x) + 0$; if $p(x) > q(x)$, it is $q(x) + p(x) - q(x)$.)

A rejection ends the round, and the later draft tokens are discarded, because their prefix has changed. Each position is handled given the accepted prefix, so the whole sequence has the target's joint distribution. (If the residual is all zero, then $p = q$ and no rejection can happen.)

Proof of (b).

1. One round gives (number $A$ of draft tokens accepted before the first rejection) + 1: the replacement token, or the bonus token if all $\gamma$ are accepted.
2. At least $j$ draft tokens are accepted exactly when the first $j$ all pass: $P(A \ge j) = \alpha^j$ for $j \le \gamma$.
3. For a count $A \in \{0, \dots, \gamma\}$, the tail-sum formula for expectations gives

(Each outcome $A = m$ is counted once in each of the terms $j = 1, \dots, m$.) 4. Add the one extra token and use the geometric sum:

By step 3 of (a), the acceptance probability at one position is $\sum_y \min(p(y), q(y)) = 1 - \mathrm{TV}(p, q)$, where $\mathrm{TV}(p, q) = \tfrac12 \sum_y |p(y) - q(y)|$ is the total-variation distance. In the example: $0.6 + 0.1 = 0.7 = 1 - 0.3$. The rate $\alpha$ is the average of this over positions; the independence in (b) is a simplifying assumption (Leviathan et al. make the same one).

Why it is faster: the roofline. Decode is memory-bound, so a target pass over $\gamma + 1$ positions costs about the same as a pass over one: the weights and the cache are read once, and the extra FLOPs stay below the ridge point. Let $c$ be the time of one draft step divided by the time of one target step. A round costs $1 + \gamma c$ target steps and gives $\mathbb{E}[\text{tokens}]$ tokens, so the speed-up is

(Leviathan et al., Theorem 3.8). It needs a high $\alpha$ and a cheap draft. It shrinks at large batch: the verify pass processes $B(\gamma + 1)$ tokens per weight read; above the ridge point it is compute-bound, and the rejected tokens then cost real compute.

Interactive figure: loading…

Prefill and decode interference #

A long prompt that enters the running batch makes that iteration long, and every decoding sequence waits: a TPOT spike. If the scheduler delays prefills instead, TTFT grows.

• Chunked prefill (Sarathi-Serve, Agrawal et al., OSDI 2024): split a long prompt into chunks of a few hundred tokens and put one chunk into each iteration, next to the decode tokens. A decode step is memory-bound, so its unused compute absorbs a chunk almost for free (section 4), and TPOT stays smooth.
• Prefill/decode disaggregation (DistServe, Zhong et al., OSDI 2024; Splitwise, Patel et al., ISCA 2024): separate GPU pools for prefill and decode; after the prefill, the KV cache is sent to a decode GPU. Each pool gets the batch size and parallelism for its own target (TTFT or TPOT), without interference. The cost: the KV transfer, two pools to size and balance, a more complex system. It pays off at large scale with strict latency targets and a fast interconnect.

Interactive figure: loading…

Answers to section 1. (1) About 6 ms, not 150. The 16.1 GB of weights are read once for all 32 sequences (4.8 ms). The 32 caches add $32 \times 1024 \times 128$ KiB $= 4.3$ GB, another 1.3 ms. The step does 32 times the work in 1.27 times the time. (2) No: speculative sampling keeps the target distribution exactly (part (a)). On average you get $(1 - 0.8^5)/(1 - 0.8) = 3.36$ tokens per target pass.

Smaller caches: GQA, MQA, MLA and quantized caches

The KV cache per token is $2 L n_{kv} d_h e$. Each factor has its own fix. GQA and MQA lower $n_{kv}$ (Ainslie et al., 2023, show that MHA checkpoints can be converted to GQA with a little extra training). Multi-head latent attention (MLA, DeepSeek-V2, 2024) stores one compressed latent vector per token and layer (512 numbers, plus 64 for the rotary part of the key) instead of full K and V per head, and recovers K and V from it with matrices that are merged into the attention computation. An fp8 or int8 cache halves $e$; this doubles $B_{\max}$ and halves the cache read time, at a small cost in accuracy (lesson 11). Sliding-window attention bounds $C$ for some layers. All of them raise either the batch that fits or the intensity $g$ of decode attention.

4. Worked example #

Serve Llama-3-8B ($\Psi = 8.03 \cdot 10^9$, $L = 32$, $a = 32$, $n_{kv} = 8$, so $g = 4$; $d_h = 128$, $h = 4096$, MLP width 14,336) in bf16 on one H100 SXM (80 GB).

Step 1: memory. Weights $W = 2\Psi = 16.1$ GB. Assume $M_{kv} = 50$ GB remain for the cache after activations, CUDA graphs and runtime buffers. With $k = 128$ KiB: $50 \cdot 10^9 / 131{,}072 = 381{,}469$ tokens, that is 23,841 blocks of 16 tokens.

Step 2: the batch that fits. At $C = 2048$: $B_{\max} = \lfloor 381{,}469 / 2048 \rfloor = 186$. Contiguous allocation for a maximum length of 8,192 tokens fits only $\lfloor 381{,}469 / 8192 \rfloor = 46$ sequences, whatever their real length. Paging counts only the live tokens, plus less than one block per sequence.

Step 3: step time at $C = 2048$, from the proposition:

$B$	bytes	FLOPs	$I$ (FLOP/byte)	bytes/BW	FLOPs/$P$	tokens/s
1	16.3 GB	0.017 T	1.0	4.87 ms	0.02 ms	205
8	18.2 GB	0.14 T	7.5	5.44 ms	0.14 ms	1,472
64	33.2 GB	1.10 T	33.0	9.92 ms	1.11 ms	6,450
186	66.0 GB	3.19 T	48.3	19.70 ms	3.22 ms	9,442

Every row is memory-bound. From $B = 1$ to $B = 64$ the throughput grows 31×; from 64 to 186 only 1.5×, because the cache reads now dominate the bytes. TPOT grows from 4.9 to 19.7 ms.

Step 4: when is decode compute-bound? The MLP matmuls need $B \ge 326$, the projections $B \ge 345$. At $C = 2048$ only 186 sequences fit, so never. With very short contexts, $C = 128$ and $B = 512$, the cache is only 8.6 GB, $I = 335 > 295$, and the compute time (8.35 ms) exceeds the memory time (7.36 ms). Compute-bound decode needs large batches of short sequences.

Step 5: prefill and chunking. A 2,048-token prompt costs $2\Psi T + 2 L a d_h T^2 = 3.29 \cdot 10^{13} + 0.11 \cdot 10^{13} = 3.40 \cdot 10^{13}$ FLOPs (3% of it causal attention): 34 ms at peak, about 69 ms at 50% of peak. Take 64 decoding sequences with $C = 1024$: a step needs 7.36 ms for its bytes but only 1.07 ms of compute. If the whole prompt joins one step at 50% of peak, that step takes about 71 ms: a 10× TPOT spike for all 64 sequences.

A chunk of $T_c$ prompt tokens adds about $2\Psi T_c$ FLOPs and almost no bytes (the weights are read anyway). So it is free while the compute time stays below the memory time:

With chunks of 256 tokens, the prompt enters over 8 steps of 7.4 ms (at peak) or about 11 ms (at 50% of peak), instead of one step of 71 ms.

Step 6: disaggregation cost. The cache of a 2,048-token prompt is $2048 \times 131{,}072$ B $= 268$ MB: 5.4 ms over a 50 GB/s network link, 0.6 ms over NVLink (450 GB/s). This is small next to a 34–69 ms prefill.

Step 7: speculative decoding at batch 1. Take a draft of about 1.2B parameters with the same tokenizer (for example Llama-3.2-1B). Decode time follows the weight bytes, so $c \approx 1.24/8.03 \approx 0.15$. With $c = 0.15$ and $\alpha = 0.8$:

$\gamma$	2	3	4	5	6	8
tokens per target pass	2.44	2.95	3.36	3.69	3.95	4.33
speed-up $S$	1.88	2.04	2.10	2.11	2.08	1.97

The best $\gamma$ is 4–5. The verify pass scores $\gamma + 1 = 5$ tokens per weight read, an intensity of about 5, far below 295: it costs one memory-bound step, as the formula assumes. At $B = 64$ the same pass has $64 \times 5 = 320$ tokens per weight read, near the compute-bound thresholds of the matmuls (326–345), and the gain shrinks.

5. In code #

NumPy checks the numbers of section 4 and the theorem. The notebook has these cells and the exercises.

import numpy as np

# Llama-3-8B shape, bf16 weights and cache (e = 2 bytes), one H100 SXM
L, a, n_kv, d_h, e = 32, 32, 8, 128, 2
Psi = 8.03e9                                 # parameters
P, BW = 989e12, 3.35e12                      # dense bf16 FLOP/s, HBM bytes/s
W = e * Psi                                  # weight bytes
g = a // n_kv                                # query heads per KV head

k = 2 * L * n_kv * d_h * e                   # KV bytes per token (K and V, every layer)
print(k, k // 1024, "KiB")                   # 131072 128 KiB

M_kv = 50e9                                  # bytes left for the KV cache
print(int(M_kv // k), int(M_kv // (k * 2048)))   # 381469 186

# Decode step at context C = 2048: bytes, FLOPs, intensity, roofline times
C = 2048
B = np.array([1, 8, 64, 186])
kv_bytes = B * C * k
n_bytes = W + kv_bytes
flops = 2 * Psi * B + 4 * B * C * L * a * d_h    # weight matmuls + attention (q.K and p.V)
t_mem, t_comp = n_bytes / BW, flops / P
print((n_bytes / 1e9).round(1))              # [16.3 18.2 33.2 66. ]
print((flops / n_bytes).round(1))            # [ 1.   7.5 33.  48.3]  far below 295
print((t_mem * 1e3).round(2), (t_comp * 1e3).round(2))
# [ 4.87  5.44  9.92 19.7 ] [0.02 0.14 1.11 3.22]
print((B / np.maximum(t_mem, t_comp)).round())   # [ 205. 1472. 6450. 9442.] tokens/s

# The intensity is a byte-weighted average of B (weights) and g (cache)
print(np.allclose(flops / n_bytes, (B * W + g * kv_bytes) / n_bytes))   # True

The speculative sampling theorem, checked step by step:

import numpy as np

rng = np.random.default_rng(0)

# 1. Part (a), step by step, on random target/draft pairs over 50 tokens
for _ in range(1000):
    p = rng.dirichlet(np.full(50, 0.3))          # target distribution
    q = rng.dirichlet(np.full(50, 0.3))          # draft distribution
    kept = q * np.minimum(1.0, p / q)            # step 2: equals min(p, q)
    residual = np.maximum(p - q, 0.0)            # unnormalised residual
    p_reject = 1.0 - kept.sum()
    assert np.isclose(p_reject, residual.sum())  # step 3
    emitted = kept + p_reject * residual / residual.sum()
    assert np.allclose(emitted, p)               # step 5: exactly the target
print("exact on 1000 random pairs")

# 2. Part (b): Monte Carlo of one round with alpha = 0.8, gamma = 4
alpha, gamma = 0.8, 4
acc = rng.random((200_000, gamma)) < alpha
n_acc = np.where(acc.all(axis=1), gamma, acc.argmin(axis=1))   # accepted before the first rejection
print(round((n_acc + 1).mean(), 2), round((1 - alpha**(gamma + 1)) / (1 - alpha), 4))   # 3.36 3.3616

And the memory that paging saves, on 200 requests with live lengths between 100 and 3,000 tokens:

import numpy as np

# 200 running requests: live lengths (prompt + tokens so far), maximum length 8192
rng = np.random.default_rng(1)
live = rng.integers(100, 3000, size=200)
block = 16
reserved = 8192 * len(live)                         # contiguous: maximum length per request
paged = block * np.ceil(live / block).sum()         # paged: whole blocks, on demand
print(live.sum(), reserved, int(paged))             # 322055 1638400 323568
print(f"used share: reserved {live.sum() / reserved:.1%}, paged {live.sum() / paged:.1%}")
# used share: reserved 19.7%, paged 99.5%
print("waste per request < one block:", bool((block * np.ceil(live / block) - live).max() < block))
# waste per request < one block: True

With reservation, about 20% of the cache memory holds tokens, close to the low end of the vLLM measurements; with paging, 99.5%.

6. Where it shows up #

• LLM engines (vLLM, SGLang, TensorRT-LLM) combine everything here: paged KV memory, continuous batching, prefix caching, chunked prefill, speculative decoding, and CUDA graphs for the decode step (lesson 13).
• Long contexts use split-KV decoding (lesson 4): at batch 1 the cache read must still be spread over all SMs.
• RL rollouts (as in GRPO) need throughput only, with many samples per prompt that share the prompt blocks. Agents and chat reuse long prefixes, so prefix caching saves most of their prefill.
• Quantization (lesson 11): weight-only int4 shrinks $W$ (shorter steps at small batch); an fp8 KV cache halves $k$ (twice the sequences, shorter long-context steps).
• ML for science. An autoregressive emulator (weather, Earth system) processes the whole state at every step: thousands of tokens per weight read, like a prefill. It is compute-bound even at batch 1 when its tokens per step exceed $B^*$, which needs wide enough weights (lesson 2: $DF/(D+F) > R$). It needs no KV cache unless it attends over past steps. A latent diffusion or rectified-flow sampler (for example for seismic ground motion) also runs the whole latent at every step; its lever is fewer steps, not a cache.

7. Common mistakes #

• Using $a$ instead of $n_{kv}$, or forgetting the factor 2 (K and V), in the KV size. For Llama-3-70B this is the difference between 320 KiB and 2.5 MiB per token.
• "Decode is slow because it needs much compute." It needs little compute. It is limited by reading the weights and the caches, so a GPU with more FLOP/s and the same bandwidth does not help.
• Expecting batching to make decode compute-bound. Batching helps the weight matmuls only. Each sequence reads its own cache (intensity $g$), so at long contexts the step stays memory-bound at any batch.
• Calling speculative decoding approximate. It is exact in distribution. (Outputs can still differ bit by bit, because different kernel shapes give different floating-point sums.) Its gain disappears at large batch, where the verify pass is compute-bound.
• Reporting only mean throughput. A server with high throughput and TPOT spikes fails its SLO. Report TTFT and TPOT percentiles, and goodput.
• Sizing the batch from the prompt length only. The cache grows with every generated token. A batch that fits at the start can run out of blocks later; this is why a server needs preemption.

8. Check yourself #

1. Llama-2-7B has 32 layers, 32 KV heads and $d_h = 128$. Give its KV bytes per token in bf16, and the number of 4,096-token sequences that fit in 50 GB. Compare with Llama-3-8B.
   Answer

   $k = 2 \cdot 32 \cdot 32 \cdot 128 \cdot 2 = 524{,}288$ B $= 512$ KiB. One sequence needs $4096 \times 524{,}288$ B $= 2.15$ GB, so $\lfloor 50/2.147 \rfloor = 23$ sequences. Llama-3-8B ($n_{kv} = 8$) stores 4× less per token: 93 sequences.
2. Llama-3-8B on one H100 with $B = 64$ and $C = 4096$: step time, tokens/s and compute time. Which limit binds?
   Answer

   Bytes: 16.1 GB of weights + $64 \times 4096 \times 131{,}072$ B $= 34.4$ GB of cache, 50.4 GB in total. So the step takes at least $50.4 / 3.35 = 15.1$ ms: about 4,250 tokens/s. Compute: $2 \cdot 8.03 \cdot 10^9 \cdot 64 / (989 \cdot 10^{12}) = 1.04$ ms plus 0.14 ms of attention. Memory-bound by more than 10×.
3. Parallel sampling: 3 answers to one 100-token prompt, blocks of 16 tokens. How many blocks without and with sharing? What happens when each sample writes its first new token?
   Answer

   $\lceil 100/16 \rceil = 7$ blocks (6 full, the last with 4 tokens). Without sharing $3 \times 7 = 21$; with sharing 7, the shared ones with reference count 3. The first new token goes into the shared, partly filled last block: the first two writers trigger copy-on-write (each gets a private copy, and the count drops to 2, then 1), and the third writes in place. Total 9 blocks.
4. Speculative decoding with $\alpha = 0.6$, $\gamma = 3$ and $c = 0.1$: expected tokens per target pass and speed-up. Does the output distribution change?
   Answer

   $(1 - 0.6^4)/(1 - 0.6) = 2.176$ tokens; $S = 2.176/(1 + 0.3) = 1.67$. No: by the theorem, every token has the target distribution.
5. Why does vLLM preempt the newest request first, and evict all of its blocks? For a 1,500-token Llama-3-8B sequence, compare the cost of swapping and of recomputation.
   Answer

   First-come-first-served: the oldest requests keep their progress. All blocks, because a sequence needs all of them at every step. Swap: $1500 \times 128$ KiB $= 197$ MB each way, about 4 ms at 50 GB/s over PCIe, plus CPU memory. Recompute: one prefill of $2 \cdot 8.03 \cdot 10^9 \cdot 1500 = 2.4 \cdot 10^{13}$ FLOPs, about 24 ms at peak, with no CPU memory and no transfer. Here swapping takes less time; recomputation is simpler.
6. A 32,768-token prompt arrives while 64 sequences decode on the Llama-3-8B server (step about 7.4 ms). What happens to TPOT without chunked prefill, and what does chunked prefill change?
   Answer

   The prefill costs $2\Psi T + 2 L a d_h T^2 = 5.3 \cdot 10^{14} + 2.8 \cdot 10^{14} = 8.1 \cdot 10^{14}$ FLOPs: 0.8 s at peak, about 1.6 s at 50%. In one piece, all 64 sequences wait that long for one token: TPOT jumps more than 100×. With 64 chunks of 512 tokens, each step carries 1/64 of the prompt. The steps take 9.5 to 18 ms at peak (19 to 36 ms at 50%); the later chunks are slower because they attend over more earlier tokens. So TPOT rises for 64 steps instead of one long stall. The prompt's TTFT grows a little (0.89 s instead of 0.82 s at peak).

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Write the KV-cache size per token.

   Show answer

2. Compute the KV cache per token of a model with $L = 80$, $n_{kv} = 8$, $d_h = 128$ in bf16.

   Show answer

3. Why is a decode step memory-bound?

   Show answer

4. In the decode intensity, which part grows with the batch and which does not?

   Show answer

5. What is the difference between TTFT and TPOT?

   Show answer

6. What is the idea of continuous batching?

   Show answer

7. What does a PagedAttention block table store, and what is the waste per sequence?

   Show answer

8. When does copy-on-write happen in PagedAttention?

   Show answer

9. Give the acceptance rule of speculative sampling.

   Show answer

10. Expected tokens per target pass with acceptance rate $\alpha$ and draft length $\gamma$?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Why is decode memory-bound, and where does batching stop helping?" The weights are shared (intensity about $B$); each cache is not (intensity $g$). KV memory caps $B$, and TPOT grows with $BCk$.
• "Estimate the step time of model X at batch B and context C." Bytes $W + BCk$ over the bandwidth; check that FLOPs/$P$ is smaller. Then tokens/s $= B / t$.
• "How does vLLM work?" Paged KV memory (block tables, blocks on demand, reference counts, copy-on-write, prefix caching), continuous batching, and preemption of the newest request with recompute or swap.
• "Is speculative decoding exact? When does it not help?" Exact by the accept/resample rule. It does not help with a low $\alpha$, a slow draft, or a large batch.
• "TPOT spikes when long prompts arrive." Chunked prefill with a token budget per step, or separate prefill and decode pools.
• "Swap or recompute after preemption?" It depends. Recompute: one compute-bound prefill, no CPU memory. Swap: PCIe copies both ways. The vLLM paper found recomputation faster for small blocks and swapping faster for large blocks (§7.3).

10. Practice #

The notebook for this lesson runs in Colab, with NumPy only. It has five exercises, each with a check cell and a folded solution:

1. kv_bytes_per_token(cfg) and max_batch(m_kv, C, cfg), checked on Llama-3-8B, Llama-2-7B and Llama-3-70B.
2. decode_step(B, C, cfg, gpu): the roofline step time, tokens/s and the bound; find the largest batch whose TPOT meets a target.
3. A BlockAllocator with block tables, reference counts and copy-on-write; compare its memory use with contiguous reservation on a synthetic request trace.
4. speculative_round(p, q, rng): show by Monte Carlo that the outputs follow $p$, and that the mean tokens per round match $(1 - \alpha^{\gamma+1})/(1 - \alpha)$.
5. A small simulator of static and continuous batching: steps, slot use and mean latency.

11. Watch and read #

Fast LLM Serving with vLLM and PagedAttentionAnyscale · the vLLM team (UC Berkeley), Ray Summit 2023 · 32 min

The authors of vLLM explain the memory waste of earlier servers, paging, sharing and the measured gains.

Speculative Decoding: When Two LLMs are Faster than OneEfficient NLP · 13 min

The draft-and-verify loop, the acceptance rule, and why the residual is (p − q)₊ (7:52); the video writes p for the draft, so its letters are swapped.

How LLMs Get Faster Without Changing Their Outputs | Speculative DecodingJia-Bin Huang · 25 min

The rejection-sampling argument (2:08–7:08) in animation, then newer drafting methods (draft trees, Medusa, multi-token prediction, EAGLE).

vLLM: Easy, Fast, and Cheap LLM Serving for Everyone - Woosuk Kwon & Xiaoxuan Liu, UC BerkeleyPyTorch Foundation · 24 min

How the vLLM engine adopts the optimizations of this lesson in practice, from its authors.

Lecture 58: Disaggregated LLM InferenceGPU MODE · Junda Chen · 75 min

Why prefill and decode interfere, and how separate prefill and decode pools are built and balanced.

• Kwon et al. (2023), Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM, SOSP 2023): §3 (the three kinds of waste), §4.1–4.5 (paging, sharing, preemption), §7.1–7.3 (kernel overhead, block size, recompute vs swap).
• Leviathan, Kalman, Matias (2023), Fast Inference from Transformers via Speculative Decoding: §2.3 (the rule), §3.1 (the expected number of tokens), Theorem 3.8 (the speed-up), Appendix A.1 (the proof).
• Chen et al. (2023), Accelerating Large Language Model Decoding with Speculative Sampling: the same rule, found independently, on a 70B model.
• Yu et al. (2022), Orca: A Distributed Serving System for Transformer-Based Generative Models (OSDI 2022): iteration-level scheduling and selective batching.
• Agrawal et al. (2024), Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve: chunked prefill and stall-free scheduling.
• Zhong et al. (2024), DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving; Patel et al. (2024), Splitwise.
• Stanford CS336, Language Modeling from Scratch (Spring 2025), lecture 10 "Inference": arithmetic intensity of prefill and decode, KV-cache size, batching, speculative sampling and paging (the lecture video is on YouTube, Stanford Online).
• Austin et al., How to Scale Your Model, All About Transformer Inference: latency and throughput estimates, the design of an inference engine, and speculative sampling.
• nano-vllm: a readable vLLM in about 1,200 lines of Python (the scheduler and the block manager with reference counts and prefix hashing).