Knowledge 01 — LLM Inference Internals

Goal: by the end you can derive, from memory, the FLOPs and memory cost of an LLM forward pass, the exact size of the KV cache, why prefill and decode have opposite bottlenecks, and the math that makes batching profitable. This is the layer where the roofline meets the actual transformer. Everything in serving (Knowledge 04) and distributed inference (Knowledge 05) is built on these facts.

Table of Contents

• 1. The transformer forward pass, just enough
• 2. Autoregressive generation: prefill vs decode
• 3. The KV cache — what it is, why it exists, how big it is
• 4. Why prefill is compute-bound and decode is memory-bound
• 5. Attention variants: MHA → MQA → GQA → MLA
• 6. FlashAttention: the IO-aware kernel
• 7. The batching math (the money slide)
• 8. Latency metrics: TTFT, TPOT, and the SLO
• 9. Speculative decoding and other decode accelerators
• 10. Sampling, long context, and other practical knobs
• 11. Worked example: size Llama-3-70B serving end to end
• 12. References

1. The transformer forward pass, just enough

A decoder-only transformer (GPT/Llama family) maps a sequence of token IDs to a probability distribution over the next token. One pass through one layer, for hidden size d, does:

1. Embedding (once at the input): token ID → vector of size d.
2. Per layer (repeated L times):
   • Attention block: project the input to Q, K, V (three matmuls, each d×d); compute attention scores softmax(QKᵀ/√d_head)·V; project the result back (d×d).
   • MLP block: two (or three, for gated/SwiGLU) matmuls expanding to an intermediate size d_ff (typically 4·d or ~2.7·d for SwiGLU) and back.
   • Each block is wrapped in a residual add and a normalization (RMSNorm/LayerNorm).
3. Final norm + LM head: project the last hidden state to vocabulary logits (d × vocab), then softmax/sample.

Where the FLOPs go: ~⅓ in attention projections, ~⅔ in the MLP (it's wider). The actual softmax(QKᵀ)V attention math is small in FLOPs for short sequences but its memory behavior dominates for long contexts (§6). The famous approximation 2P FLOPs per token (from Knowledge 00 §4) covers the weight matmuls, which are the bulk.

You do not need to derive transformers from scratch for this role, but you must know which matmuls exist because that's what compilers fuse (Knowledge 03), what quantizers quantize (Knowledge 02), and what parallelism splits (Knowledge 05).

2. Autoregressive generation: prefill vs decode

LLMs generate one token at a time, feeding each output back in. This splits inference into two phases with completely different performance characters — the single most important operational distinction in LLM serving.

Prefill (a.k.a. "context encoding" / "prompt processing")

You feed the entire prompt (say 2000 tokens) through the model in parallel, in one forward pass. All 2000 positions are computed at once. This is a big, dense matmul — lots of independent work — and produces the first output token plus the KV cache for every prompt position.

Decode (a.k.a. "generation" / "token-by-token")

Now you generate output tokens one at a time. Each step:

• takes the single token you just produced,
• runs it through the model (a tiny matmul: batch of 1 token),
• attends over all previous tokens (using the cached K and V — see §3),
• produces one new token.

You repeat decode once per output token. A 500-token answer = 1 prefill over the prompt + 500 sequential decode steps.

Prompt: "Explain the roofline model in one sentence."   (8 tokens)
 ┌─ PREFILL ─────────────────────────┐  ┌────────── DECODE (×N, sequential) ──────────┐
 │ all 8 prompt tokens at once →      │  │ t1→ t2→ t3→ ... →tN  (1 token per pass each)│
 │ first output token + KV for 8 pos  │  │ each attends to all prior tokens via KV cache│
 └────────────────────────────────────┘  └──────────────────────────────────────────────┘
   compute-bound, parallel, fast/token     memory-bound, sequential, slow/token

Consequences you must internalize:

• Prefill cost scales with prompt length (and ~quadratically in attention for very long prompts). It sets TTFT (time to first token).
• Decode cost scales with number of output tokens, each step roughly constant time. It sets TPOT (time per output token).
• A long prompt + short answer is prefill-dominated; a short prompt + long answer (chat, agents, reasoning) is decode-dominated. Customers' workloads differ wildly on this axis, and it changes the right hardware and optimization. Always ask a customer for their input/output length distribution.

3. The KV cache — what it is, why it exists, how big it is

Why it exists

In decode step t, the new token must attend to all previous tokens. The attention needs the Key (K) and Value (V) vectors of every prior position. Those don't change once computed. So instead of recomputing K and V for all prior tokens every step (O(t²) waste), we cache them. Each decode step computes K,V for only the new token and appends to the cache. This is what makes decode O(t) per step instead of O(t²).

The KV cache is therefore the memory that makes autoregressive decoding affordable — and it's also the thing that eats your VRAM and bounds your batch size.

The size formula (memorize it)

KV cache bytes = 2 · n_layers · n_kv_heads · head_dim · seq_len · batch · bytes_per_element
                 ↑
                 2 = one K and one V

• 2 for K and V.
• n_layers × (n_kv_heads × head_dim) is the per-token KV size per layer. Note n_kv_heads (not query heads) — this is where GQA/MQA save memory (§5).
• seq_len = prompt + generated so far. batch = concurrent sequences. bytes_per_element = 2 for FP16, 1 for INT8/FP8 KV-cache quantization.

Worked size: Llama-3-70B

Llama-3-70B: n_layers=80, n_kv_heads=8 (GQA), head_dim=128, FP16. Per token: 2 · 80 · 8 · 128 · 2 bytes = 327,680 bytes ≈ 320 KB/token.

• At 8k context, 1 sequence = 320 KB × 8192 ≈ 2.56 GB of KV cache.
• At batch 32, 8k context = ~82 GB — more than a single 80 GB H100 holds, on top of the 140 GB of FP16 weights. This is why 70B serving needs multiple GPUs and why KV cache, not weights, is often the binding memory constraint at high concurrency.

Principal instinct: when a customer asks "how many concurrent users fit on this box?", the limiting equation is usually VRAM = weights + KV_cache(batch, context) + activations + overhead. You solve for batch. Halving KV bytes (GQA, INT8 KV quant, shorter context) directly doubles concurrency. This is one of the highest-leverage knobs in serving and the reason PagedAttention (Knowledge 04) was invented — to stop wasting KV memory to fragmentation.

4. Why prefill is compute-bound and decode is memory-bound

This follows directly from arithmetic intensity (Knowledge 00 §4–5).

Prefill: you process S prompt tokens at once against weights of size P. You read the weights once (≈2P bytes) but do ≈2·P·S FLOPs (S tokens through every weight). Arithmetic intensity ≈ S → for any non-trivial prompt, high AI → compute-bound. The tensor cores are busy; this is "good" GPU work. Prefill throughput is limited by FLOPs → lower precision (FP8/INT8) and tensor-core utilization help.

Decode: you process 1 token against the same weights. Read weights once (≈2P bytes), do ≈2P FLOPs. Arithmetic intensity ≈ 1 → memory-bound (far below the ~295 ridge point). The tensor cores starve waiting on HBM. Decode throughput is limited by bandwidth → batching, KV/weight quantization, and bandwidth-rich chips help; raw FLOPs don't.

This table is a whiteboard staple. When someone says "our LLM is slow," your first question is which phase — because the fix for one does little for the other. Mislabeling this is the most common junior mistake.

5. Attention variants: MHA → MQA → GQA → MLA

Attention's KV cache is the memory hog, so architects shrink it. You must know all four because they change the KV formula in §3 and a customer's model could use any.

• MHA (Multi-Head Attention) — the original. Every query head has its own K and V head. n_kv_heads = n_heads. Maximum quality, maximum KV cache.
• MQA (Multi-Query Attention) — all query heads share a single K and V head. n_kv_heads = 1. Shrinks KV cache by n_heads× (e.g., 32×!) → huge memory/bandwidth win for decode, small quality cost. Used by PaLM, Falcon.
• GQA (Grouped-Query Attention) — the compromise that won. Query heads are split into G groups; each group shares one K/V head. n_kv_heads = G (e.g., 8). Most modern models (Llama-2-70B, Llama-3, Mistral) use GQA — it captures most of MQA's memory savings with near-MHA quality.
• MLA (Multi-head Latent Attention) — DeepSeek's approach: compress K/V into a low-rank latent vector that's cached, then reconstruct per head. Dramatically smaller KV cache than even GQA, with strong quality. The current frontier for KV efficiency. (MLA is half the DeepSeek architecture story — the sparse-MoE other half, and what it does to your sizing math, is Knowledge 09 §2.)

Why you care operationally: KV cache size ∝ n_kv_heads. MHA→GQA(8) on a 64-head model is an 8× KV reduction → 8× the concurrency at the same memory, or 8× longer context. When sizing a deployment, read the model config (config.json → num_key_value_heads) — don't assume MHA.

Customer translation: "Their model uses GQA with 8 KV heads, so per-token KV is small; we can fit ~X concurrent 8k-context sessions on one chip. If they were on an old MHA model, we'd need ~8× the KV memory and the box would hold a fraction of the users."

6. FlashAttention: the IO-aware kernel

Naive attention computes the full S×S score matrix QKᵀ, applies softmax, then multiplies by V. For long sequences that S×S matrix is huge (S=8k → 64M entries per head) and gets written to and read from HBM — the classic memory-wall disaster. The math is cheap; the memory traffic kills you.

FlashAttention (Dao et al., 2022; v2 2023; v3 2024) computes the identical result without ever materializing the full score matrix in HBM. It tiles Q, K, V into blocks that fit in on-chip SRAM, and uses the online softmax trick (running max + running sum) to combine block results correctly without seeing all scores at once. Net effect: attention memory traffic drops from O(S²) to O(S), it stays in fast SRAM, and it runs many times faster and with far less memory — with bit-exact-equivalent math (it's an IO optimization, not an approximation).

You won't write FlashAttention, but you must:

• Recognize it in stacks (every serious LLM server uses it or a variant — FlashAttention-2/3, xFormers, FlashInfer).
• Know when it applies: it's the reason long-context inference is feasible at all, and the reason "we enabled flash attention" gives big speedups.
• Know its successors: FlashAttention-3 exploits Hopper FP8 and async; FlashInfer/PagedAttention kernels integrate it with paged KV cache for serving (see Knowledge 04).

The lesson FlashAttention teaches (and that you repeat to customers): the same math can be 5× faster purely by changing how it touches memory. This is the soul of inference optimization — it's usually not about doing less math, it's about moving less data.

7. The batching math (the money slide)

Decode is memory-bound: you pay ~2P bytes of weight reads per step regardless of how many sequences you decode together. So if you decode B sequences in one step (a batch), you read the weights once but produce B tokens. The weight-read cost is amortized across the batch.

Per-step time ≈ max( weight_read_time , B · compute_time_per_token )
              ≈ weight_read_time         while memory-bound (small B)
Throughput  ≈ B / per-step-time  →  grows ~linearly with B until you hit the compute roof

Concretely with the Knowledge 00 Example A numbers (8B model, H100):

• Batch 1: ~210 tok/s total, ~4.8 ms/token. Tensor cores ~0.3% utilized.
• Batch 32: still ~one weight read per step (mostly), but 32 tokens out → ~6000+ tok/s aggregate, dramatically better $/token. Per-user latency barely changes until you approach the compute roof or run out of KV memory.

This is the economic argument for batched serving and the reason continuous batching (Knowledge 04) — refilling batch slots the instant a sequence finishes — can yield 5–20× throughput over naive serving.

But batching has two ceilings:

1. Compute roof — once AI ≈ ridge point (~295), you're compute-bound and more batch stops helping latency (throughput plateaus).
2. KV memory — each sequence in the batch needs its own KV cache (§3). At long context, KV memory caps B before the compute roof. This is why KV-cache efficiency (GQA, paging, quantization) directly buys throughput.

The fundamental tension you'll articulate forever: throughput vs latency. Bigger batches = higher throughput (better $/token, the operator's friend) but each individual request waits longer in the batch and TPOT can rise (worse UX, the user's enemy). The serving knob max_num_seqs/batch policy is this trade-off. A principal sets it to the customer's SLO, not to a vanity throughput number.

8. Latency metrics: TTFT, TPOT, and the SLO

Customers buy on these. Define them precisely or you'll talk past each other.

• TTFT — Time To First Token: from request arrival to the first output token. Dominated by prefill (and queuing). Critical for interactive feel. Long prompts → high TTFT.
• TPOT — Time Per Output Token (a.k.a. ITL, inter-token latency): average time between subsequent tokens during decode. Sets the "typing speed" the user perceives. Dominated by decode (memory bandwidth).
• End-to-end latency = TTFT + (num_output_tokens − 1) · TPOT. For a 500-token answer, decode dominates total time.
• Throughput: tokens/s or requests/s aggregate across all users.
• Goodput: throughput counting only requests that met the SLO. The honest metric — a server "doing 10k tok/s" while violating every latency SLO has near-zero goodput. (Term popularized by the DistServe/Sarathi line of work.)
• Percentiles, not averages: SLOs are p50/p95/p99 ("p99 TTFT < 500 ms"). Averages hide the tail that users actually complain about. Always benchmark percentiles (Knowledge 06).

The SLO conversation (pure JD2): turn "make it fast" into "p95 TTFT < X ms and p95 TPOT < Y ms at Z concurrent requests, with task quality ≥ Q." That four-number spec is what you then size hardware against (Knowledge 08). Without it, "fast" is unfalsifiable and you can't be successful.

9. Speculative decoding and other decode accelerators

Decode is sequential and memory-bound — the worst case. Several techniques attack it:

Speculative decoding (the big one)

Idea: a small, cheap draft model proposes the next k tokens quickly; the big target model then verifies all k in a single parallel forward pass (like a mini-prefill) and accepts the longest correct prefix. Because verification is parallel (compute-bound, cheap on a memory-bound regime) you get multiple tokens per expensive target pass when the draft is right.

• Speedup: typically 2–3× on decode with no quality change (the target model's distribution is preserved exactly via the acceptance rule — Leviathan et al., Chen et al. 2023).
• Cost/risk: you run a second model; acceptance rate depends on draft↔target agreement; gains shrink for hard/creative text.

Variants you should name-drop correctly

• Self-speculation / Medusa (2023): bolt extra "heads" onto the target model to predict several future tokens — no separate draft model.
• EAGLE / EAGLE-2 (2024): predict at the feature level for higher acceptance; current SOTA-ish for spec decode.
• EAGLE-3 and MTP-as-draft (2025+): the frontier — draft heads trained into the base model — covered in Knowledge 09 §5.
• Lookahead decoding: n-gram-based parallel decoding, no draft model.
• Prompt/prefix caching: reuse KV cache across requests that share a prefix (system prompt, few-shot examples) — skips re-prefilling shared context. Massive win for agent/RAG workloads with big fixed system prompts. (vLLM "automatic prefix caching".)
• Chunked prefill / piggybacking: interleave prefill chunks with ongoing decode steps so a long prompt doesn't stall everyone's decode (Sarathi-Serve). Improves TPOT tail under mixed load (Knowledge 04).

When to reach for spec decode: low-concurrency, latency-critical decode (e.g., a single-user coding assistant) where the GPU is otherwise underutilized — spec decode spends spare compute to buy latency. At high batch, the GPU is already compute-busy and spec decode helps less (sometimes hurts). Knowing that nuance is the difference between a buzzword and an engineer.

10. Sampling, long context, and other practical knobs

• Sampling (greedy / temperature / top-k / top-p / min-p): cheap relative to the forward pass, but affects determinism and reproducibility of benchmarks. Hold sampling fixed when comparing systems.
• Long context (32k–1M tokens): attention cost and KV cache grow with sequence length; this is where MLA, KV quantization, ring/sequence parallelism (Knowledge 05), and context-extension tricks (RoPE scaling, YaRN) matter. Long context is often prefill-bound (huge prompt) and KV-memory-bound (huge cache) simultaneously.
• Quantizing the KV cache (FP8/INT8 KV): halves the §3 KV bytes → directly doubles concurrency or context. A favorite high-leverage knob; small quality cost if done well (Knowledge 02).
• Stop conditions / max_tokens: bound decode length; runaway generations destroy tail latency and cost.

11. Worked example: size Llama-3-70B serving end to end

A customer wants to serve Llama-3-70B, prompts ~1k tokens, answers ~500 tokens, target p95 TPOT ≤ 40 ms, as many concurrent users as possible on H100 80GB GPUs.

1. Weights: 70B params. FP16 = 140 GB → needs ≥2 H100s just for weights → tensor parallelism TP=2 (or quantize). With INT8 weights ≈ 70 GB → fits on... still tight with KV; realistically TP=2 or TP=4.
2. KV cache (from §3): GQA 8 KV heads, 80 layers, head_dim 128, FP16 → ~320 KB/token. Context ≈ 1k+500 ≈ 1.5k tokens → ~480 MB per sequence.
3. Memory budget: say TP=4 across 4×80 GB = 320 GB. Weights 140 GB. Overhead/activations ~30 GB. Leaves ~150 GB for KV → /0.48 GB ≈ ~310 concurrent sequences of KV capacity (before fragmentation; PagedAttention recovers most of it).
4. Decode throughput: memory-bound. Aggregate decode tok/s scales with batch until compute roof; TP=4 splits both compute and KV across 4 chips. Tune max_num_seqs so that p95 TPOT stays ≤ 40 ms (smaller batch = lower TPOT, fewer users; bigger batch = more users, higher TPOT). This is the throughput↔latency dial set to the SLO.
5. Levers if it doesn't fit the SLO/budget: INT8/FP8 weights (2× capacity + decode speed), INT8 KV (2× concurrency), spec decode (lower TPOT at low load), or a bandwidth-richer chip. Each is a Knowledge 02/04/05 topic.

This single example exercises every concept in the module and is exactly the system-design interview you'll face. Practice it cold.

12. References

• Pope et al., "Efficiently Scaling Transformer Inference" (2022) — the canonical prefill/decode + KV + parallelism math.
• Dao et al., FlashAttention (NeurIPS 2022), FlashAttention-2 (2023), FlashAttention-3 (2024).
• Shazeer, "Fast Transformer Decoding: One Write-Head is All You Need" (2019) — MQA. Ainslie et al., GQA (2023). DeepSeek-V2/V3 reports — MLA.
• Kwon et al., PagedAttention / vLLM (SOSP 2023) — KV memory management.
• Leviathan et al. and Chen et al., speculative decoding (2023); Medusa (2023); EAGLE/EAGLE-2 (2024).
• Agrawal et al., Sarathi-Serve (chunked prefill, 2024); Zhong et al., DistServe (prefill/decode disaggregation, goodput, 2024).
• Kaplan et al. / Hoffmann et al. scaling-law papers for the 2P/6P FLOP approximations.

➡ Next: Knowledge 02 — Quantization — how to make all of the above 2–4× smaller and faster without breaking accuracy.