Inference Optimization for LLMs
Introduction
Training a large language model is a one-time cost; serving it happens millions of times. Inference optimization is the art of making each generated token fast and cheap — maximizing the tokens-per-second a GPU can produce and minimizing how long a user waits for a reply. It is distinct from model optimization (quantization, pruning, distillation), which shrinks the model itself. Here the model stays the same; we make the execution smarter.
Almost every technique on this page exploits one structural fact: a decoder-only LLM generates text autoregressively — one token at a time, each token conditioned on all the tokens before it. That sequential dependency is what makes naive inference slow, and what the KV cache, batching, and speculative decoding all attack.
The headline: the KV cache turns generation from O(L²) work into O(L), continuous batching keeps the GPU busy, and speculative decoding produces several tokens per expensive forward pass — all without changing the model's output.
The two phases: prefill vs. decode
A single generation request runs in two very different phases, with very different performance characteristics:
1. Prefill (process the prompt)
The entire prompt is fed through the model in parallel in one forward pass. All prompt tokens are processed at once, producing the first output token and filling the KV cache for the whole prompt.
Because we have a big batch of tokens to crunch through dense matrix multiplies, prefill is compute-bound — limited by the GPU's raw FLOP/s.
2. Decode (generate the rest)
Tokens are produced one at a time. Each step runs a forward pass for a single new token, attending to all cached past tokens, then samples the next token and repeats.
Each step does little arithmetic but must read the entire model's weights (and the KV cache) from memory, so decode is memory-bandwidth-bound — limited by how fast we can move bytes, not by FLOPs.
Why it matters: the two phases have opposite bottlenecks, so they are optimized differently. Prefill wants big parallel matmuls; decode wants high memory bandwidth and tricks (batching, speculation) that amortize the per-step weight read across more useful work.
The KV cache: trading memory for recomputation
Self-attention computes, for each position, a query Q, a key K, and a value V, then mixes values by attention weights: Attention(Q, K, V) = softmax(QKᵀ / √dₖ) · V. The crucial observation: in a causal decoder, the keys and values of past tokens do not change as new tokens arrive — token 5's K and V are the same whether the sequence is 5 tokens long or 500.
So instead of recomputing every past K and V at each step, we cache them. At generation step t we compute only the new token's K and V, append them to the cache, and attend over the stored set:
Without a cache
Each new token recomputes K/V for all t positions. Total K/V vectors to reach length L: 1+2+…+L = L(L+1)/2 → O(L²) wasted recomputation.
With a cache
Each step computes only 1 new token's K/V. Total to reach length L: L → O(L). The cost is the memory to hold every past K and V.
The KV cache is a classic space-for-time trade: we spend memory so we never recompute. The demo below makes the saved work visible.
Interactive: KV-cache visualizer
Toggle the cache on and off and drag the slider to advance generation. Watch the per-step and cumulative key/value counts, and how the cumulative-cost curve bends from quadratic (red) to linear (green).
Key/Value computations
Each row is one generation step (t1…t8); each square is a key/value vector for one position. With the cache ON, past K/V are stored (blue) and only the new diagonal cell (green) is computed each step.
Drag to advance autoregressive generation step by step.
The red curve bends upward (quadratic); the green line stays straight (linear). The gap is the work the KV cache saves.
The memory cost of the KV cache
The cache is not free — for long contexts it becomes the dominant memory consumer, often larger than the model's activations and competing with the weights themselves for GPU memory. The size of the KV cache is:
The cache grows linearly with sequence length and with batch size. A long-context conversation or a big serving batch can need tens of gigabytes just for K/V — which is why long-context serving is fundamentally a memory problem, and why decode is memory-bandwidth-bound (every step re-reads this cache).
PagedAttention
Store the cache in fixed-size "pages" (like OS virtual memory) instead of one contiguous block, eliminating fragmentation and wasted reservation. The core idea behind vLLM.
GQA / MQA
Grouped-query and multi-query attention let many query heads share a few key/value heads, shrinking n_kv_heads and therefore the cache by a large factor.
Quantized cache
Storing K/V in 8-bit (or lower) instead of 16-bit roughly halves cache size at a small accuracy cost — directly cutting dtype_bytes.
Batching and continuous batching
Because decode is memory-bandwidth-bound, a single request leaves the GPU's compute mostly idle: we read all the weights to produce just one token. The fix is batching — process many sequences together so the same weight read serves many tokens, raising GPU utilization and total throughput (the per-request latency barely changes).
Naive (static) batching waits for the whole batch to finish before starting new requests, so a short request is stuck behind a long one and slots sit empty. Continuous batching (a.k.a. in-flight or iteration-level batching) instead schedules at the granularity of a single decode step: as soon as one sequence finishes, a waiting request takes its place in the very next iteration. This keeps the batch full and can multiply throughput several-fold on real traffic.
Intuition: batching converts a bandwidth-bound problem into a more compute-bound one by amortizing each weight read over many sequences. Continuous batching ensures those batch slots are never wasted waiting on stragglers.
Speculative decoding
Decode is sequential: normally one expensive forward pass of the big model yields exactly one token. Speculative decoding breaks this by using a small, cheap draft model to guess the next k tokens, then having the big target model verify all k guesses (plus one more position) in a single forward pass.
The target accepts the longest correct prefix of the draft — the leading tokens it would have produced anyway — and replaces the first wrong token with its own sample. Each big-model pass therefore yields between 1 and k+1 tokens. When the draft is good, this is a big speedup; when the draft is bad, you fall back to one token per pass (never worse than plain decoding, ignoring the small draft overhead).
Crucially, the output distribution is unchanged. The verification step uses a rejection-sampling rule so that the sequence of accepted-and-corrected tokens is distributed exactly as if the target model had generated every token on its own. Speculative decoding makes inference faster, not different — it is a pure latency optimization, not an approximation.
Interactive: speculative decoding
Step through rounds where the draft proposes 4 tokens. Green = accepted, red = the first rejected token, gray = tokens discarded after the mismatch, purple = the target's own correction. Notice each round still produces multiple tokens from one big-model pass.
Speculative decoding — one big-model pass
The target accepts the longest correct prefix (green), rejects the first token it disagrees with (red), then emits its own token there (purple). So one expensive forward pass yields 4 tokens instead of 1. Because every accepted token is exactly what the target would have produced — and the mismatch is replaced by the target's own sample — the final output distribution is identical to ordinary decoding.
FlashAttention: IO-aware exact attention
Standard attention materializes the full seq × seq score matrix QKᵀ in slow high-bandwidth memory (HBM), reading and writing it back for the softmax — a memory-traffic bottleneck. FlashAttention instead tiles the computation, keeps blocks in fast on-chip SRAM, and fuses the softmax using an online (running-max) formulation so the big score matrix is never written to HBM at all.
The result is the same exact attention output (it is not an approximation), but with far less memory traffic and a smaller memory footprint — speeding up both prefill and training. It is a textbook example of an IO-aware algorithm: the math is unchanged; the win comes entirely from respecting the memory hierarchy.
Metrics: what "fast" means
Inference speed is not a single number. Two metrics matter, and they map onto the two phases:
| Metric | What it measures | Dominated by | Matters for |
|---|---|---|---|
| Time to first token (TTFT) | Latency from request to the first generated token | Prefill (grows with prompt length) | Perceived responsiveness |
| Inter-token latency (ITL / TPOT) | Time between successive output tokens | Decode (memory bandwidth) | Streaming reading speed |
| Throughput (tokens/sec) | Aggregate tokens produced across all requests | Batch size & GPU utilization | Serving cost / capacity |
The fundamental tension: latency vs. throughput. Bigger batches improve throughput (and cost-per-token) but can raise per-request latency. Serving systems tune the batch size and scheduling to hit a target TTFT and inter-token latency while maximizing tokens/sec.
Putting it together
| Technique | What it optimizes | Changes output? |
|---|---|---|
| KV cache | Avoids O(L²) recomputation in decode (costs memory) | No |
| Continuous batching | GPU utilization & throughput | No |
| Speculative decoding | Multiple tokens per big-model pass (latency) | No (verified, exact) |
| FlashAttention | Attention memory traffic (prefill & training) | No (exact) |
| PagedAttention / GQA / quantized KV | Shrinks KV-cache memory → longer context, bigger batches | No / minimal |
Key Takeaways
- Inference optimization speeds up serving the model; it is separate from model compression (quantization/pruning/distillation).
- Inference has two phases: prefill processes the prompt in parallel and is compute-bound; decode generates one token at a time and is memory-bandwidth-bound.
- The KV cache stores past keys and values so each decode step computes only the new token's K/V — turning O(L²) recomputation into O(L), at the cost of memory.
- For long contexts and large batches the KV cache dominates memory; PagedAttention, GQA/MQA, and quantized caches shrink it.
- Continuous batching keeps the GPU full by swapping finished sequences out at every decode step, multiplying throughput.
- Speculative decoding verifies a small model's guesses in one big-model pass, yielding several tokens per step without changing the output distribution.
- FlashAttention computes exact attention with far less memory traffic by tiling and fusing the softmax — same result, less IO.
- Measure with TTFT (prefill-bound), inter-token latency (decode-bound), and throughput (batch/utilization-bound); latency and throughput trade off against each other.