The Inference Optimization Stack

01 The Big Picture

Docs 06–08 diagnosed the disease: inference is memory-bound, the KV cache grows without mercy, and decode starves the GPU. This doc is the pharmacy — every major treatment, organized by where in the system it intervenes.

The industry's optimizations look like a zoo: MoE, SSMs, GQA, quantization, PagedAttention, kernel fusion, tensor parallelism, Splitwise… But they snap into a clean four-layer stack, ordered by when the decision is made — from model design time down to datacenter deployment time. Higher layers change what must be computed and stored; lower layers change how efficiently it happens. Master the stack and any new technique you meet has an obvious slot — and an obvious question: which bytes does it eliminate, shrink, schedule, or distribute?

02 The Stack, Animated

The cheapest byte is the one that never exists. These techniques redesign the network so less must be computed or cached at all.

MoE — only wake the experts you need

Replace each FFN with many "expert" FFNs plus a router that activates 1–2 per token. A 8×7B model holds ~47B parameters but touches ~13B per token — decode reads fewer weight-bytes through the wall (doc 08's budget) while keeping big-model capacity. Cost: all experts must still fit in VRAM, and routing causes load imbalance. Covered in depth in doc 02.

GQA — stop caching what heads can share

In vanilla multi-head attention, 32 query heads each carry their own K and V — 32 sets of cache. The insight of Grouped-Query Attention: queries need diversity (they ask different questions), but keys/values are more redundant (the "address book" being queried can be shared). So: keep 32 Q heads, share K,V across groups — e.g. 8 KV heads, each serving 4 Q heads. KV cache shrinks 4× with minimal quality loss; the extreme (1 KV head, MQA) shrinks 32× but hurts. Plug it into doc 06's formula: n_kv_heads is the lever. This is why Llama-3 and Mistral serve long contexts affordably — the choice was made at training time and you inherit it.

SSMs — abolish the KV cache entirely

State Space Models (Mamba family) are the radical option. A transformer remembers by keeping everything — the KV cache IS its memory, growing O(n). An SSM instead carries a fixed-size hidden state updated recurrently each token, like an RNN with better math (selective state updates that train in parallel). Consequences:

Given a fixed architecture, store each number in fewer bits. Since decode speed ≈ bandwidth ÷ bytes (doc 08), shrinking bytes is a direct speed multiplier.

Quantization — fewer bits per weight

Training uses fp16/bf16 (16 bits). But inference weights are read-only — they can be compressed once, offline. Map each weight to an 8-bit or 4-bit integer plus a per-group scale factor: w ≈ scale × q. The matmul dequantizes on the fly — extra compute, but compute is free (doc 08's roofline) while bytes are precious:

Why it works at all: trained weights are redundant and roughly normally distributed — most of the information survives coarse rounding if outlier values are handled carefully (that's the difference between naive rounding and GPTQ/AWQ).

KV cache compression — fewer bits per remembered token

The same logic applied to the cache, which at long context outweighs the weights themselves (doc 06's 2.1 GB per sequence). The toolbox: quantize K,V to 8/4-bit (cache is also write-once-read-many per token); sliding windows — only keep the last W tokens' K,V (Mistral); eviction policies — drop K,V of tokens that stopped receiving attention (H2O's "heavy hitters": measure attention weights, evict the ignored — literally LFU cache eviction applied to tokens); and attention sinks — the odd discovery that the first few tokens must never be evicted or generation destabilizes.

The bytes that remain must be placed, scheduled, and moved without waste. This is the serving engine's job (vLLM, TensorRT-LLM, SGLang).

PagedAttention — virtual memory for the KV cache

Pre-vLLM servers allocated each request's cache as one contiguous region sized for max_tokens — like malloc'ing the worst case up front. Result: 60–80% of "used" VRAM was actually padding and fragmentation. PagedAttention copies the OS playbook wholesale: chop the cache into fixed-size blocks (~16 tokens), keep a per-sequence block table (page table) mapping logical position → physical block, allocate on demand. Fragmentation collapses to under 4%; bonus: two requests sharing a prefix can map the same physical blocks (copy-on-write), which is how prefix caching (doc 07) is implemented under the hood.

Continuous batching — no empty seats

Scheduler admits/retires requests every decode step, not every request lifetime. Covered in doc 06 — it lives at this layer because it's pure scheduling: same model, same bytes, just never an idle slot.

Kernel fusion — don't commute back to HBM between errands

A naive implementation of LayerNorm → matmul → bias → GeLU launches four kernels; each writes its output to HBM and the next reads it back. Three round trips of intermediate tensors through the slow lane, plus per-launch overhead. Fusion compiles them into one kernel: values stay in registers/SRAM from start to finish, touching HBM once on the way in and once on the way out. FlashAttention (doc 08) is exactly this idea executed on attention — fusion is the general principle, FlashAttention its most famous instance. This is what TensorRT, torch.compile, and Triton spend their lives doing.

When one GPU's 80 GB and 3.35 TB/s aren't enough, distribute — but every split buys capacity with communication.

Tensor parallelism — split the matrices themselves

A 70B fp16 model (140 GB) cannot fit one GPU. Tensor parallelism (TP) splits each weight matrix column- or row-wise across N GPUs: each computes its shard of every matmul, then an all-reduce stitches results — twice per transformer layer. Now the model fits, and N GPUs' bandwidth works in parallel (decode ceiling roughly ×N). The catch: those all-reduces happen every layer, every token, which is why TP lives inside an NVLink island (~900 GB/s); over PCIe or Ethernet it dies. Contrast with its sibling, pipeline parallelism: split by layers (GPU1 holds layers 1–40…), cheap point-to-point comms, but sequential decode hops through every stage — TP for latency, PP for capacity across slow links.

Splitwise — give each phase its own machines

The elegant endgame of doc 06. Prefill is compute-bound; decode is memory-bound — so why run them on the same GPU? Disaggregated serving (Splitwise, DistServe, and today's production stacks) routes each request to a prefill fleet (compute-heavy GPUs, maximally batched prompts), then ships the finished KV cache over the interconnect to a decode fleet (bandwidth/capacity-optimized GPUs, huge concurrent batches). Each fleet scales independently with its own bottleneck; the interference you saw in doc 06 — decode stuttering while someone's prompt prefills — disappears, because the phases no longer share silicon. The KV-cache transfer is the toll, paid once per request.

07 Mental Models

08 Common Misconceptions