Phase 08 — Inference Optimization at Depth
Difficulty: ⭐⭐⭐⭐⭐
Estimated Time: 3 weeks (70–90 hours)
Roles Supported: Inference engineers, model deployment engineers — FlashAttention and speculative decoding are now baseline expectations for Senior Staff
Why This Phase Exists
FlashAttention is the single most important algorithm to have implemented from scratch for this role. Every Senior Staff inference engineer is expected to understand it. Not just "it's IO-efficient" — but: what is the tiling strategy, how does the softmax normalization work across tiles, what are the backward pass gradients, and how does it extend to multi-query attention.
This phase covers the three most impactful inference optimizations in the last three years: FlashAttention (reduces attention memory from O(n²) to O(n)), speculative decoding (increases generation throughput 2–4×), and continuous batching (increases server-side utilization 5–10×). If you have implemented all three from scratch, you can speak to any inference optimization discussion at depth.
Concepts
- Standard attention memory: stores full $n \times n$ attention matrix — $O(n^2)$ memory, becomes prohibitive at $n > 8192$
- FlashAttention v1/v2/v3: tiled SRAM-efficient attention; reads/writes HBM $O(n)$ times instead of $O(n^2)$; no intermediate attention matrix materialized
- Online softmax: compute softmax in tiles without seeing all logits first; requires tracking running max and normalizing factor per tile; enables tiling
- FlashAttention tiling: split Q into blocks of $B_r$ rows; for each Q block, iterate over all K/V blocks of $B_c$ columns; update running softmax state; final output is accumulated correctly
- FlashAttention backward: requires storing only the output $O$ and logsumexp $L$ (not the attention matrix) — enables $O(n)$ memory backward
- Causal masking in FlashAttention: only certain tiles are "causal" (span the diagonal); others are fully masked or fully unmasked — optimized separately
- KV cache: store K and V for all previously generated tokens; avoid recomputing; grows linearly with sequence length; enables $O(1)$ compute per decode step
- Paged attention (vLLM): KV cache stored in non-contiguous "pages" (blocks of tokens); allows dynamic allocation without pre-allocating max-sequence-length contiguous memory; enables higher server utilization
- Speculative decoding: use a small draft model to generate k tokens; verify all k tokens with the large target model in one forward pass; accept tokens that match target model's distribution
- Acceptance rate: fraction of draft tokens accepted by verifier; acceptance rate $\alpha$ → expected tokens per target model call = $(1 - \alpha^{k+1}) / (1 - \alpha)$; higher $\alpha$ → more speedup
- Continuous batching: serve multiple requests simultaneously; when one request finishes, immediately insert a new one without waiting for the entire batch to complete; critical for high-throughput inference servers
- Prefix caching: cache KV states for common prompt prefixes; next request with the same prefix reuses cached KVs; important for chat history, system prompts
- Chunked prefill: split long prefill (prompt processing) across multiple decode steps; prevents long prefills from blocking short decode steps in a batched server
Labs
Lab 01 — FlashAttention v2 from Scratch in Triton
| Field | Value |
|---|---|
| Goal | Implement FlashAttention v2 forward (and optionally backward) in Triton; verify numerical correctness against standard attention within 1e-4; demonstrate memory savings and speedup at $n=4096, 8192$ |
| Concepts | Tiled attention, online softmax with running max/normalizer, SRAM vs DRAM bandwidth, Triton tile sizes, causal masking optimization, tl.dot, tl.load/store with masks |
| Steps | 1. Implement standard attention as reference; 2. Implement flash_attn_forward_kernel in Triton: tile Q in blocks of BLOCK_Q; iterate over K/V blocks; use online softmax to accumulate output; 3. Validate output matches standard attention within 1e-4; 4. Implement causal masking (skip future tiles entirely); 5. Benchmark: standard attention vs FlashAttention at seq_len=[512, 1024, 2048, 4096, 8192]; plot memory and latency; 6. (Bonus) implement backward pass storing only O and L |
| Stack | Triton 2.3+, PyTorch 2.3+, CUDA GPU (≥16 GB VRAM for seq_len=8192 comparison) |
| Datasets | Synthetic |
| Output | flash_attn_forward.py Triton kernel; numerical correctness report; memory comparison plot (quadratic vs linear); speedup table |
| How to Test | pytest test_lab.py — checks: output matches standard attention within 1e-4 for 5 test cases, memory usage is O(n) not O(n²), causal version passes autoregressive generation test, speedup > 1.5× at seq_len=4096 |
| Talking Points | "Derive the online softmax update formula." (you can write the 3-line recurrence: $m_i = \max(m_{i-1}, \max(q_i k^T))$, $l_i = e^{m_{i-1}-m_i} l_{i-1} + \sum e^{q_i k^T - m_i}$, $O_i = \text{diag}(e^{m_{i-1}-m_i}) O_{i-1} + e^{q_i k^T - m_i} v^T$) |
| Resume Bullet | Implemented FlashAttention v2 from scratch in Triton; validated within 1e-4 of standard attention; demonstrated 3.2× memory reduction and 2.1× speedup at seq_len=4096 |
| Extensions | Add multi-query attention (MQA) variant; implement sliding window attention (Mistral); add ALiBi bias support; implement backward pass |
Lab 02 — Speculative Decoding Engine
| Field | Value |
|---|---|
| Goal | Implement speculative decoding from scratch: a draft model generates $k=5$ candidate tokens; the target model verifies all in one forward pass; acceptance/rejection sampling produces lossless output; benchmark throughput vs standard decoding |
| Concepts | Token acceptance sampling, rejection sampling proof (lossless = same distribution as target), acceptance rate measurement, draft model selection criteria, adaptive $k$ |
| Steps | 1. Implement SpeculativeDecoder class with draft_model and target_model fields; 2. Implement draft_k_tokens(input_ids, k) using the draft model greedily; 3. Implement verify_and_accept(draft_tokens, draft_probs, target_logits) — compare distributions and accept/reject using sampling; 4. Implement the rejection sampling correction: if token rejected at position $i$, sample corrected token from adjusted distribution; 5. Prove (in comments) that the output distribution equals target distribution; 6. Benchmark: measure acceptance rate $\alpha$ for (LLaMA-3.2-3B target, LLaMA-3.2-1B draft) on WikiText-2; measure tokens/sec vs standard greedy decoding |
| Stack | PyTorch 2.3+, transformers |
| Datasets | WikiText-2, ShareGPT (for realistic prompt distribution) |
| Output | SpeculativeDecoder class; acceptance rate vs prompt type analysis; throughput comparison table |
| How to Test | pytest test_lab.py — checks: output distribution matches target-only sampling (KL < 0.01 on 1000 samples), acceptance rate > 60% for matched draft/target family, throughput > 1.5× vs target-only |
| Talking Points | "Prove speculative decoding is lossless." — use the rejection sampling argument; "When does speculative decoding hurt?" — when acceptance rate < threshold, or when target model is not the bottleneck; "What is the optimal k?" — function of target model latency and acceptance rate |
| Resume Bullet | Implemented speculative decoding engine with 74% acceptance rate using LLaMA-3.2-1B/3B pair; achieved 2.3× throughput improvement over target-only decoding on ShareGPT prompts |
| Extensions | Implement tree-based speculative decoding (speculative tree instead of chain); add adaptive k selection; implement batch speculative decoding |
Lab 03 — Continuous Batching Engine with KV Cache Management
| Field | Value |
|---|---|
| Goal | Build a mini LLM serving engine with continuous batching, paged KV cache, and request scheduling; demonstrate higher utilization vs static batching at the same latency target |
| Concepts | Continuous batching (iteration-level scheduling), paged attention (non-contiguous KV cache pages), preemption (evict a request's KV cache to make room), FCFS vs priority scheduling, throughput vs latency trade-off |
| Steps | 1. Implement KVCacheManager with fixed-size pages (e.g., 16 tokens per page) and free-list allocation; 2. Implement Request dataclass with state (waiting/prefilling/decoding/done); 3. Implement ContinuousBatchingScheduler — at each step, pack as many decode requests as possible into a batch; 4. Implement PagedAttentionKernel — attention that handles non-contiguous KV pages via an indirection table; 5. Run simulation: 100 requests with varying prompt/output lengths; compare: static batching vs continuous batching (throughput tokens/sec at P99 latency < 500ms) |
| Stack | PyTorch 2.3+, transformers, asyncio |
| Datasets | ShareGPT request trace (simulated arrival times) |
| Output | ServingEngine class; throughput vs latency plot for static vs continuous batching; KV cache utilization over time |
| How to Test | pytest test_lab.py — checks: all requests produce correct output (same as sequential decoding), KV cache never overflows, continuous batching achieves ≥2× throughput at same P99 latency vs static |
| Talking Points | "How does paged attention handle variable-length KV caches?" — explain the page table indirection; "What is the scheduling problem in continuous batching?" — explain the TTFT (time to first token) vs TBT (time between tokens) trade-off |
| Resume Bullet | Built mini LLM serving engine with continuous batching and paged KV cache; demonstrated 3.1× throughput improvement vs static batching at equivalent P99 latency on simulated ShareGPT traffic |
| Extensions | Add chunked prefill; add speculative decoding from Lab 02; add request preemption (swap out to CPU); benchmark against vLLM baseline |
Deliverables Checklist
- FlashAttention Triton kernel with correctness proof and benchmark plots
- Speculative decoding engine with acceptance rate analysis
- Continuous batching engine with throughput comparison
-
All
test_lab.pysuites pass
Interview Relevance
- "Implement FlashAttention on a whiteboard." — you derive the online softmax update, draw the tiling diagram, and explain why no intermediate attention matrix is needed
- "How does speculative decoding change the compute pattern on an NPU?" — you explain batched verification vs sequential generation, the model size ratio requirement
- "Design a serving system for a 70B LLM on 4× A100s." — you architect tensor parallelism + continuous batching + paged attention from first principles