Inference optimization

Part V left you with a trained model, pre-trained, supervised-fine-tuned, preference-aligned, optionally PEFT-adapted, optionally distilled. The mathematical object exists. Part VI asks the engineering question that follows: how do you actually serve it? A transformer forward pass is conceptually simple; serving one at production scale is an entire discipline. Naive implementations are shockingly inefficient. They rerun every layer for every generated token, pad short sequences to match long ones in the batch, and materialize huge intermediate attention matrices that the GPU spends most of its time reading and writing rather than computing. A production stack does none of this.

This chapter covers the foundational optimizations that distinguish production serving from a naive forward pass. KV cache, the central optimization, an ~$700\times$ speedup over re-computing past tokens. Prefill vs decode: two phases with completely different performance characteristics, which dictates which optimizations help where. Continuous batching, Yu et al. (2022, Orca) showed that the batch should adapt request by request, not be fixed at startup. Flash Attention, Dao (2022) eliminated the $(N \times N)$ intermediate that made long contexts impractical. Speculative decoding, Leviathan et al. (2023) and Chen et al. (2023) amortize each big-model forward pass across multiple accepted tokens. PagedAttention, Kwon et al. (2023, vLLM) manages KV cache memory like an operating system manages RAM.

The goal of the chapter is not to make you re-implement vLLM. It is to give you the mental model: what’s expensive, what’s cheap, where the bottlenecks live, and which optimization attacks which bottleneck. With that frame, the next two chapters fall into place, quantization (Ch 18) reduces bits per parameter, sampling (Ch 19) decides what the decoder actually emits, and the modern serving stacks (vLLM, TGI, SGLang, TensorRT-LLM) start to read as recombinations of a small number of primitives.

The inference cost problem

The naive baseline

The naive way to generate text is the way the math is written down: at every step, push the entire so-far sequence through every layer of the model and read off the next-token distribution. Cost is $O(N^2)$ in attention per step and $O(N^3)$ to generate $N$ tokens in total, the outer loop over generated tokens multiplies into the per-step attention cost. For a single $1024$-token completion, the attention term alone contributes roughly $10^9$ redundant multiply-adds, on top of the far larger redundant weight-matmul cost from re-running every layer on every previous token at every step, even though those tokens have not changed since the previous step. Wall-clock throughput on an A100 lands somewhere around $5$ to $10$ tokens per second. That is fast enough to type out a response, but nowhere close to what production economics demand.

Why this matters

The same model, same hardware, with a fully optimized inference stack: $200$ to $500$ tokens per second. A $5$ to $50\times$ throughput gap that has nothing to do with model architecture and nothing to do with hardware. It is purely the difference between running the math as written and running the system as designed for serving. The economic consequence is decisive. A service handling one billion tokens per day on naive inference might need a hundred A100s; on a production stack, five. Margin per request flips from impossible to comfortable. At LLM scale the savings are not marginal; they are the difference between a viable product and one that bleeds compute.

These optimizations are not exotic research. Every major open-source serving system (vLLM, TGI, SGLang, TensorRT-LLM) uses some combination of them, and most managed APIs (Together, Anyscale, Fireworks, the model-provider endpoints) layer additional proprietary tuning on top. The pattern is industry standard, and the patterns are the topic of this chapter.

Prefill vs decode: two phases

Prefill

Generation begins with prefill: the model processes the entire prompt in a single forward pass. All prompt tokens flow through every layer together, queries and keys and values are computed for every position at once, and the attention matrix is filled in one shot. Because each matmul operates on a wide sequence dimension, prefill is compute-bound, high arithmetic intensity, lots of FLOPs per byte of weight loaded from HBM, GPU fully saturated. Cost scales as $O(N^2)$ in attention and $O(N)$ in the rest of the model, where $N$ is the prompt length. The user-facing metric this phase determines is TTFT (Time To First Token): the latency between submitting a request and seeing the first output. Long prompts move TTFT up roughly quadratically because of the attention term.

Decode

After prefill, the model enters decode: tokens come out one at a time, each requiring its own forward pass. The sequence dimension of those forward passes is exactly one, just the newly emitted token, with the rest of the context handled through the KV cache (next section). This makes decode memory-bound. Each layer’s matmul has the same big weight matrix as in prefill but with a tiny activation tensor, so the GPU is dominated by the time spent loading weights from HBM rather than the time spent multiplying. Arithmetic intensity collapses. The user-facing metric this phase determines is TPOT (Time Per Output Token): how fast the stream of tokens scrolls once it has started. TPOT is roughly constant per token after the KV cache amortizes the past work, and is dominated by memory bandwidth, not compute.

Why this matters

The two phases have different bottlenecks, which means they call for different optimizations. Prefill optimizations target compute and kernel efficiency: bigger matmuls, kernel fusion, Flash Attention. Decode optimizations target memory bandwidth and unused compute: cache layouts that minimize HBM traffic, Flash Decoding, and speculative decoding (which uses up the slack compute to verify many drafted tokens at once). Naming the phases makes the rest of the chapter readable: when you read about a new technique, the first question to ask is “does this help prefill or decode?”, and the answer almost always reveals which bottleneck the technique attacks. Productionally, the two phases also trade off against each other in a serving scheduler: long prefills can stall the decode of ongoing requests, so modern stacks chunk prefills to interleave them with decode steps and keep the TPOT smooth for already-streaming users.

KV cache: the central optimization

The redundancy

Consider what a naive decode step at position $i$ actually does. To compute attention at that step, the model needs the keys and values $K_1, V_1, \ldots, K_{i-1}, V_{i-1}$ for every previous token, plus the new ones $K_i, V_i$ for the token just generated. The naive way is to compute all of them from scratch, re-running the K and V projections for the entire prefix at every step. But those past keys and values are deterministic functions of the past tokens, which have not changed since the previous step. Recomputing them is pure waste. The KV cache removes this waste by storing keys and values from previous steps and only computing the new ones at each decode iteration.

The savings are dramatic. Per layer per token, the work drops from $O(N)$ matmul (compute K and V for all $N$ tokens) to $O(1)$ matmul (compute K and V for just the new token). The attention computation itself drops from $O(N^2)$ per step to $O(N)$ per step, each new query still attends to all $N$ cached keys, but the keys themselves are looked up rather than re-derived. Summed across a $1024$-token decode, the total attention work collapses from $\sum_{i=1}^{1024} i^2 \approx 3.6 \times 10^8$ to $\sum_{i=1}^{1024} i \approx 5.2 \times 10^5$, roughly a $700\times$ reduction in attention FLOPs, plus the savings from never re-projecting K and V. This is the central optimization of LLM inference; everything else stacks on top of it.

What the cache stores

For every layer, the cache stores the keys and values for every previous position. A typical layout is (layers, 2, batch, heads, seq_len, head_dim), where the 2 separates K from V. At each decode step the model:

1. Computes Q, K, V for the single new token,
2. Appends the new K and V into the cache slot at position $\text{seq\_len}$,
3. Computes attention between the new Q and the full cache of K and V.

The Q for the new token is the only fresh query in the layer; the keys and values for everything else were paid for once, on the step they were produced. The KV cache is the model’s working memory: the snapshot of what has been said and how it has been encoded. After prefill, decode is just lookups against that snapshot plus a small new write each step.

The memory cost

Speed comes at the cost of memory. The cache size is given by:

where $L$ is the number of layers, $H$ is the number of attention heads, $d_{\text{head}}$ is the per-head dimension, $N$ is the current sequence length, and batch is the number of concurrent sequences sharing the GPU. In BF16, each entry takes $2$ bytes.

For Llama-7B at $4$K context, plugging in $L = 32$, $H = 32$, $d_{\text{head}} = 128$, $N = 4096$ gives roughly $2$ GB per sequence, manageable. For Llama-70B at $128$K context, with $L = 80$, $d_{\text{head}} = 128$, $N = 131072$, and the true GQA count of $H = 8$ KV heads, the number jumps to roughly $40$ GB per sequence; with the full $64$ query heads (no GQA) it would be roughly $320$ GB. The KV cache dominates inference memory at long context. Model weights are fixed regardless of how many users you serve; the cache scales linearly with both context length and concurrent requests. This is the fundamental tension of long-context serving: the algorithmic speedup is enormous, but the price is paid in memory, and at $128$K and beyond the cache becomes the binding constraint on how many concurrent sequences a GPU can hold.

The runnable cell below simulates the speedup directly. It defines a toy “transformer”, four layers of $W \cdot x$ with a quadratic-attention term, and times $50$ decode steps both ways: recomputing the entire prefix each step versus carrying a KV cache forward. Even on a CPU, even with small dimensions, the recomputation tax is visible. Real transformers see a much larger gap because real attention is genuinely $O(N^2)$ per layer.

import numpy as np
import time

np.random.seed(0)
d_model = 256
n_layers = 4
W = [np.random.normal(0, 0.02, (d_model, d_model)) for _ in range(n_layers)]

def naive_decode_one_step(tokens, weights):
    """No KV cache, recompute everything for all tokens."""
    x = tokens                # (seq, d_model)
    for w in weights:
        # Recompute Q, K, V for ALL tokens; attention is O(seq^2) per layer
        x = np.tanh(x @ w + (x @ x.T) @ x / x.shape[0])
    return x

def cached_decode_one_step(new_token, kv_cache, weights):
    """With KV cache, process only the new token."""
    x = new_token.reshape(1, -1)
    for layer_idx, w in enumerate(weights):
        kv_cache[layer_idx].append(x @ w)
        keys = np.concatenate(kv_cache[layer_idx], axis=0)
        x = np.tanh(x @ w + (x @ keys.T) @ keys / len(kv_cache[layer_idx]))
    return x

# Setup: 100-token prompt, decode 50 more tokens
prompt = np.random.normal(0, 1, (100, d_model))

# Naive: every step recomputes everything
naive_start = time.time()
running = prompt
for _ in range(50):
    new = np.random.normal(0, 1, (1, d_model))
    running = np.concatenate([running, new], axis=0)
    _ = naive_decode_one_step(running, W)
naive_time = time.time() - naive_start

# Cached: only the new token at each step
kv_cache = [[] for _ in range(n_layers)]
for tok in prompt:
    cached_decode_one_step(tok, kv_cache, W)   # prefill
cached_start = time.time()
for _ in range(50):
    new = np.random.normal(0, 1, (d_model,))
    cached_decode_one_step(new, kv_cache, W)
cached_time = time.time() - cached_start

print(f"Naive decoding:  {naive_time*1000:.0f} ms")
print(f"Cached decoding: {cached_time*1000:.0f} ms")
print(f"Speedup: {naive_time/cached_time:.1f}×")
print(f"\nIn production: the speedup is much larger because real attention is O(N^2) per layer.")
print(f"For 1024-token sequences, KV cache provides ~700× speedup vs naive.")

Batching strategies

Static, dynamic, continuous

Three batching strategies have dominated, in roughly chronological order. Static batching is the simplest: pad every sequence in the batch to the same length and run them as one big matrix. This is easy to implement and trivially correct, but it is wasteful, short sequences burn compute on padding tokens, and the batch only advances at the rate of the longest sequence, so any request finishing early just sits idle until the batch is done. Dynamic batching improves utilization by grouping incoming requests by similar length and serving each length-bucket as its own static batch. Less padding waste, but the cost is latency: each request has to wait for its bucket to fill before processing begins.

Continuous batching is the modern default, introduced by Yu et al. (2022) in the Orca paper. The idea is iteration-level scheduling instead of request-level scheduling: at each decode step, the serving system re-decides which sequences are in the batch. When a sequence finishes (hits EOS or the max length), its slot is freed immediately; when a new request arrives, it joins the batch at the next step without waiting for anyone else. Prefill and decode for different requests can even interleave on the same batch: the prefill of a new request runs alongside decode of established ones, with the scheduler choosing chunk sizes that keep both classes of work flowing. The result is higher sustained GPU utilization than any pre-2022 approach could achieve, while keeping per-request latency low because nothing artificially stalls.

The latency vs throughput tradeoff

Even with continuous batching there is no free lunch. Higher batch size increases throughput, more requests share each weight-loading round, which is exactly what decode (memory-bound) needs, but per-request latency climbs as the batch crowds shared memory and compute. Lower batch size cuts per-request latency but leaves the GPU underutilized, especially during decode. The continuous-batching scheduler walks this tightrope dynamically: fill the batch when there are requests waiting, drain it when there are not, never block a finished sequence behind an unfinished one. The empirically observed sweet spot for a $7$B model on an A100 sits somewhere around batch size $16$ to $64$; for bigger models on bigger hardware the optimum slides around based on memory pressure and the prefill/decode mix.

Flash Attention

The memory problem

Naive attention computes the full $(N \times N)$ score matrix $QK^T$ as an explicit intermediate, applies softmax row-wise, and multiplies by $V$. For a sequence of length $N = 8192$, that intermediate has $67$M entries, occupying $134$ MB in BF16. The intermediate is wasteful: only the final $\text{softmax}(QK^T) V$ output ever leaves the attention block; the full score matrix exists for milliseconds before being thrown away. But on the way there, it has to be written to HBM and read back to apply softmax, then read again to multiply by $V$. At long context, those round trips to HBM eat almost all of the kernel’s wall-clock time; the matmul itself is fast, the memory traffic is not.

Tiling + online softmax

Flash Attention (Dao 2022) eliminates the intermediate. The trick has two pieces. First, tiling: process attention in blocks of size $B_q \times B_k$ that fit in the GPU’s on-chip SRAM rather than its HBM. Second, online softmax: maintain a running max and a running sum for each query row so that partial softmax outputs from each tile can be merged into the final answer without ever reconstructing the full row. The output $\text{softmax}(QK^T)V$ is built incrementally as the kernel sweeps over $K$ and $V$ tiles, never materializing the full $(N \times N)$ matrix anywhere outside the on-chip register file. Chapter 10 covers the kernel-engineering side of this in more depth, the tiling scheme, SRAM/HBM bandwidth accounting, and how the technique fits into the broader custom-kernel story of training infrastructure; the treatment here focuses on what it buys at serving time.

The effect is two-fold. Memory complexity for the intermediate drops from $O(N^2)$ to $O(N)$, removing the practical ceiling on context length. Wall-clock time drops $2$ to $4\times$ versus naive attention on long sequences because the kernel is no longer pinned by HBM bandwidth. Flash Attention 2 (Dao 2023) refined the parallelism, better work partitioning across thread blocks and tighter outer loops, for roughly another $2\times$ on top of FA-1. Production frameworks (PyTorch SDPA, vLLM, xFormers, the major training stacks) all default to Flash Attention or an equivalent fused kernel; the “naive” attention from textbooks is essentially never what runs on a real GPU today.

Flash Attention is the kernel-level lever. Speculative decoding is the second algorithmic lever, after KV cache.

Speculative decoding

The recipe

Decode is memory-bound. Each step loads the entire weight matrix from HBM, multiplies it against a sequence dimension of one, and spends most of its time waiting on memory. There is unused compute capacity in every decode step. Speculative decoding (Leviathan, Kalman & Matias 2023; Chen et al. 2023) is the optimization that converts that unused capacity into throughput. The recipe is short:

1. A small draft model, typically $30\times$ to $100\times$ smaller than the target model, generates $k$ candidate tokens autoregressively. Because it is small, this is cheap.
2. The big model runs a single forward pass with all $k$ candidates concatenated to the existing context. Same matmuls as a normal decode step but with a sequence dimension of $k$ instead of $1$.
3. For each position $i = 1, \ldots, k$, compare the big model’s distribution to the draft’s choice. Accept every draft token that matches; on the first rejection, replace it with the big model’s choice and stop.

The big model emits at least one token per round (the rejected position, replaced with its own pick) and possibly all $k$ draft tokens if every match holds. Net effect: one big-model forward pass yields between $1$ and $k+1$ committed tokens, instead of exactly $1$ as in standard decode.

Why it works

Two facts make the bookkeeping pay off. First, the big model’s forward pass with $k$ tokens of input is only modestly slower than with $1$ token; the matmuls are the same shape on the weight side, only the sequence dim widens. For a memory-bound regime, $k$ tokens cost almost the same as $1$. Second, each accepted draft token saves a full big-model forward pass that the system would otherwise have had to run to produce that same token. As long as the draft model agrees with the big model often enough, the speedup is real. Acceptance rates of $0.6$ to $0.8$ are typical for a well-chosen draft model; below $0.5$ the technique struggles to break even because the draft cost and the rejection probability conspire against you.

The version in Chen et al. (2023) goes one step further: a rejection-sampling correction guarantees that the resulting token distribution is exactly what the big model alone would produce. Speculative decoding under this scheme is not an approximation; it is mathematically equivalent to the original decoder, just faster. That property is what made the technique a deployment-ready optimization rather than a quality-throughput tradeoff.

The speedup formula

The expected speedup can be written:

Numerator: tokens committed per round. Denominator: cost per round in units of one big-model forward pass. With $k = 5$, acceptance rate $\alpha = 0.7$, and a draft model $50\times$ cheaper than the big one, the formula gives roughly $2.5\times$ speedup, large enough to matter, modest enough to be honest about the engineering work required to capture it.

def speculative_speedup(k, alpha, draft_relative_cost=0.02, overhead=0.05):
    """
    Estimate speculative decoding speedup.

    k: number of draft tokens proposed per round
    alpha: acceptance rate (fraction of draft tokens accepted)
    draft_relative_cost: draft inference cost / big inference cost
    overhead: framework overhead per round
    """
    # Expected accepted tokens per round
    if alpha < 1:
        expected_accepted = (1 - alpha**(k+1)) / (1 - alpha) - 1
    else:
        expected_accepted = k

    # Cost: 1 big pass + k draft passes + overhead
    cost_per_round = 1 + k * draft_relative_cost + overhead

    # Tokens per cost
    speedup = (1 + expected_accepted) / cost_per_round

    return expected_accepted, speedup

# Sweep over k and alpha
print(f"{'k':>3} | " + ' '.join(f'α={a:.1f}' for a in [0.3, 0.5, 0.7, 0.9]))
print("-" * 50)
for k in [1, 3, 5, 7, 10]:
    speeds = []
    for alpha in [0.3, 0.5, 0.7, 0.9]:
        _, s = speculative_speedup(k, alpha)
        speeds.append(f'{s:.2f}x')
    print(f"k={k:>3} | " + ' '.join(s.rjust(6) for s in speeds))

print("\nObservations:")
print("- Low acceptance (α=0.3): marginal speedup at any k")
print("- High acceptance (α=0.7-0.9): k=5-7 gives 2-3x speedup")
print("- Too-high k can hurt: draft cost accumulates, acceptance drops over horizon")

PagedAttention and modern stacks

The fragmentation problem

A naive KV cache allocator reserves a contiguous memory block per sequence, sized for the maximum context that sequence might reach. Sequences arrive and finish at different times with different actual lengths, so the freed-up blocks are different sizes than the incoming ones need, classic memory fragmentation, the same phenomenon that motivated paged virtual memory in operating systems. At low concurrency it is annoying; at production scale, with hundreds of overlapping sequences, the GPU ends up unable to admit new requests despite having plenty of total free memory.

Pages, like an OS

PagedAttention (Kwon et al. 2023, vLLM) borrows the OS solution directly. Allocate the KV cache in fixed-size pages, typically $16$ tokens each, and let a sequence’s cache be a list of pages, not necessarily contiguous in physical memory. A page table maps logical positions in the sequence to physical pages. When a sequence finishes, its pages return to a free pool and are reusable by anyone, regardless of the next request’s length. Fragmentation effectively vanishes.

Pages buy something the contiguous layout could not. Shared prefixes: when many requests share a common system prompt or few-shot prefix, the corresponding KV cache pages can be shared across sequences with copy-on-write semantics. A single copy of the system prompt’s cache serves all users; only the per-user suffix needs new pages. Higher concurrency: $2$ to $4\times$ more sequences fit on the same GPU than under contiguous allocation, depending on the workload mix. Dynamic memory pressure: when a long-running sequence demands more pages than expected, the scheduler can swap out other sequences’ pages, much like an OS swaps memory under pressure.

Modern inference stacks

A short, opinionated map of the production landscape:

• vLLM, open source; introduced PagedAttention; the default starting point for self-hosted serving in $2024$–$2026$.
• TGI (Hugging Face Text Generation Inference), production-focused, multi-LoRA serving, deep Hugging Face Hub integration.
• SGLang, adds a programming-language abstraction for complex multi-step LLM workflows on top of efficient inference.
• TensorRT-LLM (NVIDIA), aggressive kernel-level optimization tuned for NVIDIA hardware; the choice when you have NVIDIA-only deployments and want the last drop of throughput.

For most teams the answer is vLLM unless there is a specific reason to reach elsewhere. Larger organizations mix and match, or build proprietary stacks, but the underlying primitives (continuous batching, paged KV cache, Flash Attention, speculative decoding) are common across every serious option.

Exercises

Four exercises that lock in the chapter’s machinery. Each is a self-contained problem with a starting template; hints are collapsed by default; try the problem first.

Exercise 1 (easy): KV cache memory calculator

Implement the KV cache memory formula and compute the cache size for various model sizes and context lengths.

def kv_cache_size_gb(seq_len, n_layers, n_heads, head_dim, batch=1, dtype_bytes=2):
    total_bytes = 2 * n_layers * n_heads * head_dim * seq_len * batch * dtype_bytes
    return total_bytes / 1024**3

MODELS = {
    "Llama-7B":  {"layers": 32, "heads": 32, "head_dim": 128},
    "Llama-13B": {"layers": 40, "heads": 40, "head_dim": 128},
    "Llama-70B": {"layers": 80, "heads": 64, "head_dim": 128},
}

contexts = [1024, 4096, 32768, 131072]
print(f"{'Model':<12} | {'1K':>10} {'4K':>10} {'32K':>10} {'128K':>10}")
print("-" * 60)
for name, spec in MODELS.items():
    sizes = [kv_cache_size_gb(c, spec["layers"], spec["heads"], spec["head_dim"]) for c in contexts]
    row = ' '.join(f'{s:>7.2f} GB' for s in sizes)
    print(f"{name:<12} | {row}")

Model        |         1K         4K        32K       128K
------------------------------------------------------------
Llama-7B     |    0.50 GB    2.00 GB   16.00 GB   64.00 GB
Llama-13B    |    0.78 GB    3.12 GB   25.00 GB  100.00 GB
Llama-70B    |    2.50 GB   10.00 GB   80.00 GB  320.00 GB

Exercise 2 (medium): Naive vs cached decoding

Implement naive decoding (recompute everything per token) and cached decoding (reuse stored K, V). Compare timing.

def naive_decode_step(all_tokens, weights):
    x = all_tokens
    for w in weights:
        x = np.tanh(x @ w)
    return x

def cached_decode_step(new_token, kv_cache, weights):
    x = new_token.reshape(1, -1)
    for layer_idx, w in enumerate(weights):
        x = np.tanh(x @ w)
        kv_cache[layer_idx].append(x[0])
    return x

prompt = np.random.normal(0, 1, (100, d_model))

naive_start = time.time()
running = prompt
for _ in range(50):
    new = np.random.normal(0, 1, (1, d_model))
    running = np.concatenate([running, new], axis=0)
    naive_decode_step(running, W)
naive_time = time.time() - naive_start

kv_cache = [[] for _ in range(n_layers)]
for tok in prompt:
    cached_decode_step(tok, kv_cache, W)
cached_start = time.time()
for _ in range(50):
    new_token = np.random.normal(0, 1, (d_model,))
    cached_decode_step(new_token, kv_cache, W)
cached_time = time.time() - cached_start

print(f"Naive decoding:  {naive_time*1000:.0f} ms")
print(f"Cached decoding: {cached_time*1000:.0f} ms")
print(f"Speedup: {naive_time/cached_time:.1f}×")

Exercise 3 (medium): Speculative speedup formula

Implement the speculative decoding speedup formula and verify it matches the widget’s behavior.

def speculative_speedup(k, alpha, draft_cost=0.02, overhead=0.05):
    expected_accepted = sum(alpha**i for i in range(1, k + 1))
    expected_emitted = expected_accepted + 1
    cost_per_round = 1 + k * draft_cost + overhead
    speedup = expected_emitted / cost_per_round
    return expected_accepted, expected_emitted, cost_per_round, speedup

K_VALUES = [1, 3, 5, 7, 10]
ALPHA_VALUES = [0.3, 0.5, 0.7, 0.9]

print(f"{'k':>3} | " + ' '.join(f'α={a:.1f}' for a in ALPHA_VALUES))
print("-" * 50)
for k in K_VALUES:
    speedups = []
    for alpha in ALPHA_VALUES:
        _, _, _, sp = speculative_speedup(k, alpha)
        speedups.append(f'{sp:.2f}x')
    print(f"k={k:>3} | " + ' '.join(s.rjust(6) for s in speedups))

  k | α=0.3 α=0.5 α=0.7 α=0.9
--------------------------------------------------
k=  1 |  1.21x  1.40x  1.59x  1.78x
k=  3 |  1.28x  1.69x  2.28x  3.10x
k=  5 |  1.24x  1.71x  2.56x  4.07x
k=  7 |  1.20x  1.67x  2.64x  4.79x
k= 10 |  1.14x  1.60x  2.61x  5.49x

Exercise 4 (hard): Continuous batching simulator

Simulate continuous batching: a fixed-size batch of decode requests where each request finishes at a different time. Track GPU utilization vs naive batching (wait for the longest request to finish before starting a new batch).

def continuous_batching(requests):
    pending = list(requests)
    active = []
    completed = 0

    throughputs = []
    utilizations = []

    for step in range(SIMULATION_STEPS):
        while len(active) < BATCH_SIZE and pending:
            active.append(pending.pop(0))

        tokens_this_step = len(active)
        for r in active:
            r.remaining -= 1
        finished = [r for r in active if r.remaining == 0]
        completed += len(finished)
        active = [r for r in active if r.remaining > 0]

        throughputs.append(tokens_this_step)
        utilizations.append(len(active) / BATCH_SIZE)

    return throughputs, utilizations

def naive_batching(requests):
    pending = list(requests)
    active = []
    throughputs = []
    utilizations = []

    for step in range(SIMULATION_STEPS):
        if not active and pending:
            active = pending[:BATCH_SIZE]
            pending = pending[BATCH_SIZE:]

        tokens_this_step = len(active)
        for r in active:
            r.remaining -= 1
        active = [r for r in active if r.remaining > 0]

        throughputs.append(tokens_this_step)
        utilizations.append(len(active) / BATCH_SIZE)

    return throughputs, utilizations

# Run both
requests = generate_requests(200)
cont_throughput, cont_util = continuous_batching(requests)
naive_throughput, naive_util = naive_batching(generate_requests(200))   # fresh requests

print(f"Continuous batching:")
print(f"  Avg throughput: {np.mean(cont_throughput):.1f} tokens/step")
print(f"  Avg utilization: {np.mean(cont_util)*100:.0f}%")
print(f"\nNaive batching:")
print(f"  Avg throughput: {np.mean(naive_throughput):.1f} tokens/step")
print(f"  Avg utilization: {np.mean(naive_util)*100:.0f}%")

The full inference picture

The optimization stack

A modern production inference stack combines:

• A distilled or quantized model (Ch 16, Ch 18), smaller compute and memory base.
• KV cache with PagedAttention, efficient algorithmic decode and efficient memory layout.
• Continuous batching: high sustained GPU utilization across variable-length requests.
• Flash Attention, fused, memory-aware kernels that make long context practical.
• Speculative decoding: amortize big-model passes across multiple committed tokens.
• Custom sampling logic (Ch 19), temperature, top-p, top-k, constrained decoding, JSON-mode.

The combined effect on the same hardware is roughly $5\times$ to $10\times$ throughput versus a naive baseline, sometimes more, depending on workload. A $7$B model on a single A100 with a naive stack delivers around $5$ to $10$ tokens per second; the same hardware running a fully optimized stack delivers $200$ to $500$. The model is identical. The system around it is different. That is the lever.

What’s next

Chapter 18 narrows in on a specific category of optimization: quantization. Where this chapter targeted wasted computation, Ch 18 targets bits per parameter, INT8, INT4, NF4, AWQ, GPTQ. The compositional story matters: quantization stacks multiplicatively with the inference-stack optimizations here, often giving another $2$ to $4\times$ throughput improvement on top of what this chapter establishes. Chapter 19 turns to the sampling layer, how the decoder actually picks a token given the logits the model produced. Temperature, top-p, top-k, beam search, structured-output constraints. Together with this chapter, those three chapters make up the full Part VI serving stack.

After Part VI, the curriculum opens into the second half: capabilities (Ch 20–23), safety and evaluation (Ch 24–26), and agents (Ch 27–30). The lessons here will recur: an agent loop is a sequence of inference calls, and the cost economics of every later chapter depend on whether the serving layer underneath them is the naive $5$ tokens-per-second or the production $500$.

Inference optimization is what separates a research artifact from a production system. A trained model is a mathematical object; serving it efficiently is engineering. The optimizations covered here, KV cache, continuous batching, Flash Attention, speculative decoding, PagedAttention, are not optional flourishes. They are the difference between a $5$-token-per-second toy and a $500$-token-per-second production service on the same chip. Part VI is the chapter cluster where that difference gets paid for, and Ch 17 is its foundation.