Part 6 — Inference Fundamentals
Transformer and attention intuition
Connect the Transformer's actual math — Q/K/V attention, multi-head, the MLP, the LM head — to the inference vocabulary you keep hearing: KV cache, FlashAttention, PagedAttention, MQA/GQA, and continuous batching. Understand why each serving trick exists by seeing the computation it optimizes.
The words you hear in inference engineering — KV cache, FlashAttention, PagedAttention, GQA, continuous batching — are not arbitrary jargon. Each one exists to make a specific piece of the Transformer's math cheaper to run for many users at once. This chapter walks the computation once, and hangs each serving technique on the step it optimizes.
What you will understand by the end
- What Q, K, and V actually compute, and why that creates the KV cache.
- Why fewer key/value heads (MQA/GQA) directly shrink decode's memory pressure.
- What FlashAttention and PagedAttention each change — and what they don't.
- How iteration-level scheduling (ORCA) became today's continuous batching.
The computation, in one pass
A token is first turned into a vector by an embedding lookup (a memory read from a learned table, not heavy math). Then it flows through a stack of Transformer layers, each with two parts: attention (tokens exchange information) and an MLP / feed-forward block (each token's vector is transformed). At the very end, the LM head projects the final vector onto the vocabulary to produce logits, and a sampler picks the next token.
Attention mixes information across tokens; the MLP transforms each token on its own; the LM head turns the result into a next-token distribution. Attention is where the interesting inference costs live, because it is the part that must remember every previous token.
Q, K, V — and why they force a cache
For each token, the model computes three vectors from its hidden state, using learned weight matrices:
- Q (query): what am I looking for?
- K (key): what do I contain?
- V (value): what should I pass on if attended to?
Attention scores every token's query against every token's key, softmaxes the scores into weights, and blends the values:
Attention(Q, K, V) = softmax(Q Kᵀ / √d) · V
Here is the inference-critical consequence. When the model generates token 501, its query must attend to the keys and values of all 500 previous tokens. Recomputing those every step would be quadratic and wasteful — so the serving engine stores them:
The KV cache is exactly the stored keys and values of every prior token. It is built during prefill and read (and appended to) on every decode step. This is what turns attention from "re-read the whole history each token" into a single new query scanning a growing cache — and it is the thing consuming decode's memory bandwidth and capacity.
So the two runtime facts from Prefill and decode — decode is bandwidth-bound and capacity is the batch ceiling — both trace back to this cache.
Multi-head, and the MQA/GQA lever
Attention runs in several heads in parallel — different learned "views" of context (one might track syntax, another long-range references) — whose outputs are concatenated and projected back. More heads means more expressive attention, but also more K/V tensors to store, i.e. a bigger KV cache and more decode bandwidth.
That trade is where two architecture choices come in:
- MQA (Multi-Query Attention): many query heads share a single key/value head.
- GQA (Grouped-Query Attention): query heads share K/V in groups — a middle ground.
MQA and GQA reduce the number of key/value heads, which directly shrinks the KV cache. Smaller KV cache → less HBM bandwidth per token and more requests fit in memory → higher decode throughput. They are a model-architecture answer to a serving bottleneck.
Two smaller pieces round out a layer: RoPE (rotary positional embeddings) rotates Q and K so the cached keys stay position-aware — cheap elementwise math, not a tensor-core cost — and the MLP is big dense matmuls that the tensor cores handle but that can still stall on weight bandwidth at low batch.
Modern models also swap the original block components for lighter equivalents — RMSNorm in place of LayerNorm, a SwiGLU-gated MLP in place of a plain ReLU one (RoPE is itself one of these upgrades). These are implementation refinements; the inference story is unchanged: attention is the only step that mixes across tokens, and everything else — normalisation, the MLP, the LM head — runs independently per token. That is exactly why the KV cache and attention, not the MLP, are where serving work concentrates.
FlashAttention and PagedAttention — two different fixes
These are constantly confused. They optimize different things and neither changes the model's output.
- FlashAttention is an IO-aware attention kernel. Naive attention materialises the big score matrix in slow HBM; FlashAttention tiles the computation and keeps chunks in fast on-chip memory, cutting memory traffic. Same math, far fewer memory reads/writes.
- PagedAttention is KV-cache memory management. Requests have different prompt and output lengths, so naive contiguous KV allocation fragments memory and wastes it. PagedAttention stores the cache in fixed pages (like virtual memory), so more requests pack into the same VRAM. It grows the batch you can physically hold.
FlashAttention makes each attention computation cheaper; PagedAttention lets you fit more concurrent requests. One is a compute-kernel win, the other a memory-packing win. Saying "FlashAttention increases batch size" (or vice versa) mixes up the two levers.
ORCA → continuous batching
The last bridge is scheduling. Early serving ran request-at-a-time: finish one user's whole answer, then the next. The ORCA system introduced iteration-level scheduling — run one decode step for everyone, rebuild the active set, run the next step — plus selective batching (batch the ops that benefit, don't force a rigid shape).
That idea is exactly today's continuous batching, and this project measured what it's worth: once requests could join the batch every step instead of waiting, throughput rose ~22×. The ancestor paper's insight, made real. See the experiment →
The connection map
| Runtime word | The math it optimizes |
|---|---|
| KV cache | storing K/V so a new query needn't recompute all history |
| MQA / GQA | fewer K/V heads → smaller cache, less decode bandwidth |
| FlashAttention | the Q Kᵀ softmax V kernel, with less memory traffic |
| PagedAttention | packing the KV cache so more requests fit |
| Continuous batching | ORCA-style per-iteration scheduling of the decode loop |
This is the minimum bridge from architecture to serving, not a full Transformer course. Training-time ideas (dropout, weight decay) are irrelevant to inference and shouldn't be confused with these serving techniques — they shape how the weights were learned, not how they're served.
Common mistakes
- Confusing FlashAttention and PagedAttention. Kernel efficiency vs memory packing.
- Thinking MQA/GQA are quality tricks. They are primarily serving optimizations that cut KV-cache cost.
- Treating the KV cache as an implementation detail. It is the origin of decode's bandwidth and capacity limits — the thing most serving work is fighting.
Practical guidance
- When a model card advertises GQA, read it as "cheaper decode / bigger batch," and factor that into your KV-cache budget from Prefill and decode.
- Expect any serious engine (vLLM, SGLang, TRT-LLM) to give you FlashAttention-class kernels, PagedAttention-class memory management, and continuous batching — their presence is table stakes, and Inference runtimes compares how they differ.
- Reason about a new "attention variant" by asking which axis it touches: compute kernel, KV-cache size, or scheduling.
Summary
- Attention's Q/K/V math is what forces a KV cache, and the KV cache is the source of decode's bandwidth and capacity limits.
- MQA/GQA shrink the cache; FlashAttention cuts attention's memory traffic; PagedAttention packs the cache to grow the batch; continuous batching (from ORCA) schedules the decode loop per iteration.
- Each serving technique maps to a specific step of the computation — which is how to reason about the next one you meet.
Knowledge check
A colleague says "we enabled FlashAttention, so we can now fit more concurrent requests." What's the correction?
FlashAttention makes the attention kernel cheaper (less HBM traffic, same math); it does not change how much KV cache fits in memory. Fitting more concurrent requests is what PagedAttention does, by packing the KV cache into pages. Different lever.
Why does Grouped-Query Attention (GQA) raise decode throughput?
Decode's cost is dominated by reading the KV cache from memory each token. GQA has query heads share fewer key/value heads, so the KV cache is smaller — less bandwidth per token and more requests fit in VRAM, both of which lift decode throughput.