Warmup Guide — Inference Optimization at Depth

Zero-to-expert primer for Phase 08. Builds the three pillars of modern LLM serving from first principles: FlashAttention (online softmax + tiling), speculative decoding (rejection sampling over a draft model), and continuous batching with PagedAttention.

Table of Contents

Chapter 1: Why Attention Is Memory-Bound

Standard attention materializes the score matrix:

  1. $S = QK^\top$ — an $n \times n$ matrix, written to DRAM ($O(n^2)$ bytes).
  2. $P = \text{softmax}(S)$ — read $n^2$, write $n^2$.
  3. $O = PV$ — read $n^2$ again.

FLOPs are $O(n^2 d)$, but DRAM traffic is $O(n^2)$ with a small constant of reuse — arithmetic intensity is low (Phase 07's framework), so for realistic $d$ the kernel runs on the bandwidth roof, far below compute peak. Worse, at $n = 32K$, the FP16 score matrix is $32768^2 \times 2 \approx 2$ GB per head per batch element — it doesn't even fit.

The fix is not approximation (sparse/linear attention change the math). FlashAttention computes the exact same output while never materializing $S$ — pure memory choreography. Two ideas make it possible: safe softmax (Ch. 2) and its online form (Ch. 3).

Chapter 2: Numerically Safe Softmax

$\text{softmax}(x)_i = e^{x_i} / \sum_j e^{x_j}$ overflows FP16 at $x \approx 11$ ($e^{11} > 65504$). The standard fix subtracts the max (mathematically identity — numerator and denominator share the factor $e^{-m}$):

$$m = \max_j x_j, \qquad \text{softmax}(x)_i = \frac{e^{x_i - m}}{\sum_j e^{x_j - m}}$$

Now every exponent is ≤ 0, every $e^{x_i - m} \in (0, 1]$. The cost: softmax now needs two reduction passes (max, then sum) before the normalize pass — three sweeps over data you wanted to touch once. That structure is what Chapter 3 dismantles.

Chapter 3: Online Softmax — The Enabling Trick

The question: can you compute a safe softmax in one pass, seeing elements block-by-block, never knowing the global max in advance?

Yes — keep running statistics and rescale when the max changes. Process blocks $b = 1, 2, \ldots$; maintain the running max $m$ and running denominator $\ell$:

$$m^{new} = \max(m, \tilde{m}_b), \qquad \ell^{new} = \ell \cdot e^{m - m^{new}} + \tilde{\ell}_b \cdot e^{\tilde{m}_b - m^{new}}$$

where $\tilde{m}_b, \tilde{\ell}_b$ are the block's local max and local exp-sum. The factor $e^{m - m^{new}}$ retroactively corrects everything accumulated under the old max — exactness preserved by exponent algebra, not approximation.

The same correction extends to the attention output accumulator: maintain $O = \sum p_j v_j$ un-normalized, rescale by $e^{m - m^{new}}$ whenever the max moves, divide by $\ell$ once at the very end. This is the entire mathematical content of FlashAttention; everything else is engineering.

Chapter 4: FlashAttention — Tiling the Whole Thing

The kernel (Lab 01 implements it tile-by-tile):

for each block Q_i of queries (loaded once into SRAM):
    m = -inf; l = 0; O_acc = 0
    for each block (K_j, V_j) of keys/values (streamed through SRAM):
        S_ij = Q_i @ K_j.T / sqrt(d)          # tile of scores — lives only in SRAM
        (apply causal mask for the diagonal block; skip blocks above it)
        update m, l, O_acc with the online-softmax rescaling (Ch. 3)
    O_i = O_acc / l                           # normalize once
    write O_i to DRAM                         # the ONLY n×d-sized write

Accounting: DRAM traffic falls from $O(n^2)$ to $O(n^2 d / M)$ where $M$ is SRAM size — in practice 5–10× less traffic, kernel moves toward the compute roof, and memory footprint is $O(n)$, unlocking long context. FLOPs slightly increase (the rescaling) — a deliberately losing trade on compute to win on bandwidth, the signature move of memory-bound optimization.

Backward pass: don't store $P$; recompute tiles from $Q, K$ in the backward sweep, keeping only the per-row $m, \ell$ statistics — same recompute-over-store logic as activation checkpointing. FA2 vs FA1 (worth knowing): FA2 reorders loops (parallelize over query blocks AND heads/sequence), reduces non-matmul FLOPs, and improves work partitioning between warps — ~2× over FA1 on the same math.

Chapter 5: Speculative Decoding — The Bandwidth Arbitrage

The setup (Phase 07 Ch. 5): batch-1 decode reads all weights to produce one token. Verifying $k$ tokens at once costs nearly the same wall-clock as generating one — the weight read amortizes over $k$ positions (it's a small prefill).

The arbitrage: let a small, fast draft model propose $\gamma$ tokens autoregressively (cheap — its weights are small), then run the big target model once over all $\gamma$ positions in parallel and accept the longest valid prefix. Big-model bandwidth cost per accepted token drops by the mean accepted length.

Expected speedup: with per-token acceptance rate $\alpha$ and draft length $\gamma$, the expected tokens per target-model pass is $\frac{1 - \alpha^{\gamma+1}}{1 - \alpha}$; speedup follows after subtracting draft cost. Practical: α ≈ 0.6–0.8 for a well-matched draft → 2–3×. Tuning $\gamma$ trades wasted draft work (rejections) against amortization — Lab 02 sweeps it and profiles α per domain (code drafts accept far more than creative prose).

Chapter 6: The Acceptance Rule — Why the Distribution Is Exact

The non-obvious part: the output distribution is provably identical to sampling the target model alone. For draft distribution $q$ and target $p$ at each position:

  • Accept the drafted token $x$ with probability $\min(1,, p(x)/q(x))$.
  • On rejection, resample from the residual distribution $\text{norm}(\max(0,, p - q))$ and stop the speculation there.

Why it works (the one-line proof sketch): P(output = x) = q(x)·min(1, p(x)/q(x)) + P(reject)·residual(x) = min(q, p) + (p − min(q, p)) = p(x). The draft can be anything — a tiny LLM, an n-gram cache, Medusa heads — correctness never depends on draft quality, only speed does. Lab 02's test validates exactly this distributional identity, which is also the production guarantee you quote to stakeholders: speculative decoding is a pure latency optimization with zero accuracy risk.

Chapter 7: Serving Many Requests — Static vs Continuous Batching

Static batching: collect $B$ requests, run them in lockstep until all finish. Two pathologies: (1) head-of-line waiting — short generations idle while the longest finishes (padding waste scales with length variance); (2) batch-formation latency — requests wait for the batch to fill.

Continuous (in-flight) batching: scheduling at iteration granularity. Every decode step: finished sequences exit, queued requests join, the step runs on whoever is active. GPU utilization stays high regardless of length variance; new requests start immediately. Costs: a scheduler (admission, preemption), per-step bookkeeping, and the memory problem — admitting a request whose KV cache will grow means you must manage fragmentation and out-of-memory preemption, which is Chapter 8's subject.

Prefill/decode interference is the remaining subtlety: a long prefill stalls every decoding request that iteration — production schedulers chunk prefills or disaggregate prefill and decode onto different workers.

Chapter 8: PagedAttention — Virtual Memory for the KV Cache

The problem: KV caches grow token-by-token to unpredictable lengths. Contiguous pre-allocation forces reserving max-length per request — vLLM measured 60–80% of KV memory wasted on reservation + fragmentation.

The solution is literally OS paging: divide KV storage into fixed-size blocks (e.g., 16 tokens each); each sequence holds a block table (logical → physical mapping); allocate blocks on demand as generation proceeds. The attention kernel gathers K/V through the block table (one indirection per tile — a few % overhead).

What the indirection buys, beyond near-zero waste:

  • Prefix sharing: identical prompt prefixes (system prompts, few-shot headers) map to the same physical blocks with copy-on-write on divergence — beam search and parallel sampling share almost everything.
  • Preemption: evict a sequence's blocks (swap to CPU or recompute later) under pressure, by table manipulation.

Lab 03 builds the block manager + continuous scheduler and measures utilization vs static batching — the waste numbers are the lesson.

Lab Walkthrough Guidance

Order: Lab 01 (FlashAttention) → Lab 02 (speculative) → Lab 03 (continuous batching).

  • Lab 01: implement the non-tiled safe softmax attention as your reference oracle first. Then online softmax on a 1-D toy (verify the rescaling identity numerically per block). Then the tiled kernel; test against the oracle at several sizes including non-divisible tile boundaries and the causal mask's diagonal blocks (the two classic bug sites).
  • Lab 02: implement the acceptance rule and validate the distribution (histogram over many samples vs direct target sampling) before measuring speed. Then sweep $\gamma$, measure α, and reproduce the expected-tokens formula empirically.
  • Lab 03: block manager first (allocate/append/free with a fragmentation counter), then the iteration-level scheduler, then the comparison: static vs continuous on a synthetic workload with high length variance. Report utilization + p50/p99 latency.

Success Criteria

You are ready for Phase 09 when you can, from memory:

  1. Write safe softmax and the online update for $(m, \ell, O)$, and explain the rescaling factor's role.
  2. State FlashAttention's traffic result ($O(n^2) \to O(n^2 d/M)$) and why FLOPs go up.
  3. Derive the speculative acceptance rule and sketch the exactness proof.
  4. Compute expected accepted tokens for given $(\alpha, \gamma)$.
  5. Explain iteration-level scheduling and both static-batching pathologies.
  6. Describe the block table, prefix sharing, and what fragmentation PagedAttention eliminates.

Interview Q&A

Q: Why does FlashAttention speed things up while doing MORE FLOPs? Attention is bandwidth-bound (Ch. 1); the binding resource is DRAM traffic, not compute. Tiling + online softmax cut traffic ~$d/M$-fold at the cost of rescaling arithmetic — spending the abundant resource (FLOPs) to save the scarce one (bytes). On the roofline: the point moves right (higher intensity) along the bandwidth roof toward the ridge.

Q: Where does speculative decoding NOT help? (1) Compute-bound regimes — large batches/prefill: verification no longer rides free bandwidth; (2) poorly matched draft (low α) — wasted draft work exceeds amortization; (3) when the draft model itself competes for the same memory/compute budget (edge devices may not fit both); (4) sampling at high temperature over flat distributions — α drops because q and p diverge more after sharpening differences.

Q: Your continuous-batching server shows great throughput but terrible p99 latency. Why? Prefill interference: long prompts entering the batch stall every in-flight decode for those iterations. Fixes: chunked prefill (split long prompts across iterations), prefill/decode disaggregation, admission control on prompt length per iteration, priority scheduling. Also check preemption storms under memory pressure — evicted-and-recomputed sequences show up purely in the tail.

Q: How do FlashAttention and PagedAttention interact? They compose: FlashAttention defines the compute pattern (tiled streaming over K/V); PagedAttention defines the storage layout (blocks via table indirection). The fused kernel streams K/V tiles by walking the block table — block size is chosen to align with the kernel's tile size so the indirection happens once per tile, not per element.

References

  • Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022) — arXiv:2205.14135
  • Dao, FlashAttention-2 (2023) — arXiv:2307.08691
  • Milakov & Gimelshein, Online normalizer calculation for softmax (2018) — arXiv:1805.02867
  • Leviathan et al., Fast Inference from Transformers via Speculative Decoding (2022) — arXiv:2211.17192
  • Chen et al., Accelerating LLM Decoding with Speculative Sampling (2023) — arXiv:2302.01318
  • Kwon et al., Efficient Memory Management for LLM Serving with PagedAttention (vLLM) (SOSP 2023) — arXiv:2309.06180
  • Yu et al., Orca: A Distributed Serving System for Transformer-Based Generative Models (OSDI 2022) — continuous batching's origin
  • vLLM documentation — read the scheduler and block-manager design docs