Lesson 6 of 7 · 49 min

Attention & serving at scale

The production techniques that make transformers servable: MQA/GQA (cut KV bytes at the source), FlashAttention 1/2/3 (IO-aware exact attention), sliding-window/sparse attention, PagedAttention + continuous batching, and quantization. Real measured impact, the tradeoffs, and how vLLM/TensorRT-LLM/SGLang differ.

Where the textbook transformer meets the GPU budget

Everything so far was the model. This lesson is the serving stack — the techniques that decide whether your transformer costs $0.50 or $5 per million tokens, fits one GPU or eight, hits a 200ms or 2s p95. Three cost pressures drive every technique here: KV-cache bandwidth, attention-compute arithmetic, and serving-pool fragmentation. Each lever trades against the others and against quality. Senior LLM-engineer interviews live in this lesson, because this is where you prove you’ve actually run models in production, not just read the papers.

Recall the two binding constraints from Lessons 2 and 5: attention’s quadratic cost (compute and HBM I/O) and the KV cache’s linear-but-heavy memory. The serving toolkit attacks both. Fewer KV heads (MQA/GQA) and lower-precision cache shrink KV bytes; IO-aware kernels (FlashAttention) and sparse/windowed attention shrink attention cost; PagedAttention + continuous batching kill pool fragmentation and pack the GPU. Let’s take them in order of leverage.

vLLM: Easy, Fast, and Cheap LLM Serving (PagedAttention)Woosuk Kwon, UC Berkeley

MQA & GQA — cut the cache at the source

The highest-leverage architectural choice for serving is the KV head count. Standard MHA gives every query head its own K and V head. Multi-Query Attention (MQA) keeps all H query heads but shares a single K/V head — so the cache shrinks by H×. Grouped-Query Attention (GQA) interpolates: G groups, each group of H/G query heads sharing one K/V head (MQA is G=1, MHA is G=H). Llama 2 70B and Mistral 7B both use 8 KV heads, cutting cache ~4–8× vs MHA at near-MHA quality. GQA’s killer trick: you can uptrain an existing MHA checkpoint into GQA by mean-pooling the per-group K/V weights and continuing for only ~5% of original pre-training compute.

code

1MHA -> MQA -> GQA  (the KV-head lever; cache scales with H_kv)23  variant   query heads   KV heads   cache vs MHA   quality       used in4  -------   -----------   --------   ------------   -----------   -------------------5  MHA       H             H          1x (baseline)  baseline      orig Transformer, BERT6  MQA       H             1          1/H            slight drop   PaLM, Falcon7  GQA       H             G (e.g 8)  G/H            ~= MHA        Llama 2 70B, Mistral8  MLA       latent        compressed lowest         ~= MHA        DeepSeek-V2910  GQA-8 on a 64-query-head model = 8x smaller KV cache at near-MHA quality.11  Uptrain MHA -> GQA for ~5% of pretraining compute. Decide BEFORE fine-tuning.

Interview angle. “MQA vs GQA — when each?” GQA-8 recovers essentially all of MHA’s quality at near-MQA speed, so it’s the default for new long-context models; reserve MQA (G=1) for extreme-memory or constrained-device deployment where a 1–2% quality hit is acceptable. The tradeoff vendors hide: many “GQA helps” benchmarks also quantize the cache, double-counting the savings — isolate the lever when you report. And don’t retroactively convert a deployed MHA model unless you can retrain ~5% and re-run your full SLO suite; the conversion must be redone per fine-tune, so GQA is most attractive adopted before any domain tuning.

FlashAttention — IO-awareness, not smarts

FlashAttention (Dao et al., 2022) is the most important attention kernel in production, and the key insight is counterintuitive: it’s not an approximation and not cleverer math — it’s the same exact attention, reorganized to minimize HBM traffic. A naive kernel materializes the T×T score matrix in HBM, applies softmax, reads it back. FlashAttention tiles Q/K/V into blocks that fit in fast SRAM, computes softmax incrementally (online softmax with rescaling), and never writes the full T×T matrix to HBM. The output is bit-for-bit the standard result. It reduces HBM I/O from O(T²·d) to O(T²·d²/M) where M is SRAM size — and on the backward pass it recomputes the matrix from cheap stored stats instead of storing it, trading FLOPs for memory.

code

1FLASHATTENTION FAMILY (exact attention; the win is HBM I/O, not approximation)23  version   key idea                                    measured impact4  -------   -----------------------------------------   ----------------------------------5  FA-1      tile Q,K,V in SRAM; never write T x T HBM    3x on GPT-2 (1k); 2.4x LRA (1k-4k)6  FA-2      better work partitioning across warps        2x over FA-1; 50-73% of A100 peak7  FA-3      warp specialization + TMA + FP8 (Hopper)     740 TFLOPS FP16; ~1.2 PFLOPS FP889  A100: HBM ~2 TB/s vs SRAM ~19 TB/s -> moving the T x T matrix IS the bottleneck.10  Treat naive 3-loop attention as a bug, not a baseline.

The numbers: FA-1 gave a 3× wall-clock speedup on GPT-2 at 1k context and 2.4× on Long-Range Arena; FA-2 raised hardware utilization to 50–73% of A100 peak FLOPs (~2× over v1) by better warp-level work partitioning; FA-3 hits 740 TFLOPS in FP16 and ~1.2 PFLOPS in FP8 on H100 (~75% of theoretical max) via warp specialization and the Tensor Memory Accelerator. Interview angle. “How does FlashAttention save memory — derive it.” Don’t say “it’s faster”; say it keeps the score tiles SRAM-resident, never materializes the T×T matrix in HBM, and recomputes in the backward pass — IO-awareness, exact output. The pitfall worth naming: FP16 attention can have numerical-stability issues at long context, an active discussion on the project’s issue tracker.

Sliding-window & sparse attention

If you can’t afford full O(T²) attention, restrict which pairs attend. Sliding-window attention (SWA) lets each token attend only to its W nearest neighbors — linear, not quadratic. The clever part is stacking: Mistral 7B uses W=4096 across 32 layers, so through layered composition a token at the top can reach back ~32×4096 ≈ 131k tokens of effective context, while a rotating buffer halves cache memory above 8k. Longformer adds a few global tokens that attend to everything (carrying long-range state); BigBird combines window + global + random for linear complexity (~8× less time/memory than vanilla at 4k). Mistral 7B is competitive with Llama 2 13B across benchmarks on this design.

The failure mode: windowing alone breaks bare retrieval — if the answer sits outside the window and there’s no global token or layered reach, the model simply can’t find it. Interview angle. “Sliding-window vs full attention for long context?” → stacked SWA (Mistral-style) preserves long-range reach via composition and is production-friendly; isolated single-layer windows lose retrieval; for multi-million-token contexts, combine SWA with retrieval augmentation rather than ever-larger windows. Global tokens are the right tool for encoders/summarization heads, not general decoders.

PagedAttention & continuous batching — the serving breakthrough

Naive KV-cache management pre-allocates a contiguous block of max_seq_len per request, leaving most slots empty — typical fragmentation waste is 60–80%. PagedAttention (vLLM) borrows OS virtual-memory paging: store K/V in fixed-size physical blocks (e.g. 16 tokens) and keep a per-request block table mapping logical → physical blocks. Waste drops to <4%, and prefixes can be shared safely across requests (system prompts, parallel samples). The measured win: 2–4× throughput over FasterTransformer/Orca at matched latency, and up to 24× over HuggingFace Transformers on high-batch offline inference.

The other half — and the most underrated lever — is continuous (iteration-level) batching: instead of waiting for a whole batch to finish, the engine evicts completed sequences and admits new ones after every forward pass. Orca reported 36.9× throughput over FasterTransformer on GPT-3 175B at matched latency. Interview angle. Many “vLLM is 4× faster” claims actually measure PagedAttention and continuous batching combined — isolate the lever when you report. SGLang’s RadixAttention extends paging to cross-request prefix reuse via a radix tree, which is the win for RAG/agent workloads with shared system prompts.

Quantization — and how it interacts with internals

Quantization is where naive reporting most often hides the real lift, because three axes are independent: what you quantize (weights, activations, KV cache), how many bits, and where the conversion lives. SmoothQuant (W8A8) migrates activation outliers into the (smooth) weights so both fit INT8, enabling INT8 GEMM end-to-end — ~1.56× speedup, ~2× memory, training-free. AWQ (W4) protects the ~1% salient weight channels and 4-bit-quantizes the rest; great memory, hardware-friendly. GPTQ (W3/W4) uses one-shot Hessian-based rounding; the historical weight-only baseline. FP8 (Hopper/Blackwell via Transformer Engine) ~doubles matmul throughput with minimal accuracy work — and pairs cleanly with FA-3.

code

1QUANTIZATION: PICK BY HARDWARE + WORKLOAD23  method        precision   helps        cost           best on4  -----------   ---------   ----------   ------------   ----------------------5  GPTQ          W3/W4       memory       visible @3-bit  weight-only baseline6  AWQ           W4          memory       small           single-GPU budget infer7  SmoothQuant   W8A8        compute+mem  small           Ampere (INT8 GEMM)8  FP8 (TE)      W/A/KV FP8  both         small           Hopper/Blackwell only910  W4 wins MEMORY but barely speeds compute (activations still FP16).11  W8A8 wins SPEED on Ampere. FP8 wins BOTH on Hopper -- and only there.12  KV-cache quant (INT8/FP8) ~2x cache, but watch perplexity drift past ~64k.

How it touches the internals you learned: weight-only W4 saves memory but not compute (activations stay FP16); W8A8 speeds compute via INT8 GEMM; FP8 does both, but only on Hopper+. KV-cache quantization (INT8/FP8) roughly halves cache but needs careful per-tensor scaling and can drift on retrieval-heavy tasks past ~64k tokens. RoPE interacts badly with naive integer K-quantization (the rotation operates on bounded magnitudes), so FP8 K-cache with per-token scaling is the production-safe config. Interview angle. “4-bit vs 8-bit vs FP8 — which and why?” → name the hardware (Ampere → W8A8; Hopper → FP8; tight single-GPU memory → AWQ W4) and always ablate perplexity at your longest context, where the failure hides.

The engines: vLLM, TensorRT-LLM, SGLang

All three ship FlashAttention, continuous batching, and PagedAttention-style pooling by default — the differences are in fusion depth, FP8 maturity, quantization breadth, and prefix reuse. vLLM is the general-purpose default: broadest model/quant coverage (GPTQ, AWQ, SmoothQuant, FP8), simplest to ship. TensorRT-LLM is the throughput leader on NVIDIA hardware (compiled kernels, best FP8/INT8), at the cost of a heavy engine-compile/calibration cycle. SGLang wins prefix-heavy RAG/agent workloads via RadixAttention and structured-output primitives.

code

1H100, Llama-3.3-70B FP8, concurrency 100 (Spheron benchmark, Mar 2026)23  engine          throughput (tok/s)   p95 TTFT (ms)   pick when4  -------------   ------------------   -------------   ------------------------------5  TensorRT-LLM    2,780                1,280           max throughput on NVIDIA fleet6  SGLang          2,460                1,380           prefix-heavy RAG / agents7  vLLM            2,400                1,450           broadest support, fastest iter89  TensorRT-LLM ~8-16% ahead on throughput, ~12-15% on TTFT -- for a heavier deploy.10  SGLang's real edge (prefix reuse) does NOT show in this raw-throughput table.

Interview prep

This is the lesson that decides senior LLM-engineer serving rounds. Interviewers test whether you bind techniques to measured numbers, know which lever attacks which constraint (bandwidth vs arithmetic vs fragmentation), and pick the stack by traffic shape — not by absolute tokens/sec. Lead with the constraint, the lever, and a number.

01“MQA vs GQA vs MHA?” → KV-head lever: MHA H heads, MQA 1, GQA G; GQA-8 ≈ MHA quality at ~MQA speed; uptrain MHA→GQA for ~5% compute.
02“How does FlashAttention save memory?” → tile Q/K/V in SRAM, never materialize the T×T matrix in HBM, recompute in backward; exact, not approximate; 2–3×.
03“Sliding-window vs full attention?” → stacked SWA (Mistral W=4096×32 layers ≈ 131k reach) is linear and production-friendly; isolated windows lose retrieval.
04“What is PagedAttention and why does vLLM use it?” → OS-style paging of the KV cache → waste 60–80%→<4%, safe prefix sharing, 2–4× throughput.
05“PagedAttention vs continuous batching?” → different levers; continuous (iteration-level) batching alone was Orca’s 36.9×; they compose.
06“W4 vs W8A8 vs FP8?” → W4 saves memory only; W8A8 speeds compute on Ampere; FP8 does both on Hopper — choose by hardware, ablate at longest context.
07“Which serving engine?” → vLLM (general/iteration speed), TensorRT-LLM (max throughput on NVIDIA, heavy compile), SGLang (prefix-heavy RAG via RadixAttention).
08“KV cache quantization risks?” → ~2× cache saving but per-tensor scaling needed; RoPE+naive INT8 K is unstable; FP8 K with per-token scaling, watch >64k drift.

Going deeper: “how do you size a serving box for Llama-7B FP16 at 4k, batch 8?” (weights ~14 GB + KV 16 GB + framework ~2 GB ≈ 32 GB → right up against an A100 40 GB; a 24 GB 4090 maxes ~batch 2–3; moving to GQA is the single biggest lever); “TTFT-dominated vs throughput-dominated SLO?” (TTFT → FlashAttention + continuous batching + keep K in FP16; throughput → GQA + FP8 + TensorRT-LLM or vLLM); “prefix-reuse-dominated?” (SGLang RadixAttention). And always: re-test needle-in-a-haystack after any change — quantization, RoPE scaling, GQA conversion — because perplexity hides structural-retrieval regressions.

paperFlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessDao et al. (arXiv)paperGQA: Training Generalized Multi-Query Transformer ModelsAinslie et al. (arXiv)paperEfficient Memory Management for LLM Serving with PagedAttention (vLLM)Kwon et al. (arXiv)videoFlashAttention — Tri Dao | Stanford MLSys #67Stanford MLSys Seminars paperSmoothQuant: Accurate and Efficient Post-Training Quantization for LLMsXiao et al. (arXiv)

Checkpoint

A 70B MHA model is too KV-cache-heavy to serve at your target context and concurrency. You have ~5% of original pre-training compute available. Best move?

AUptrain the MHA checkpoint into GQA (e.g. 8 KV heads) by mean-pooling per-group K/V weights, then re-benchmark the SLO suiteBSwitch to MQA (1 KV head) for the maximum cache cutCIncrease the GPU count and keep MHA unchanged

Checkpoint

At 8k+ context your attention is the latency bottleneck on an A100, and you’re running a hand-written three-loop attention kernel. Highest-impact change, and why it’s safe?

AApproximate attention with a low-rank kernel to cut FLOPsBAdopt FlashAttention: it tiles Q/K/V in SRAM and never materializes the T×T matrix in HBM — exact output, 2–3× faster because the bottleneck was HBM I/OCLower the batch size so each attention call is smaller

Checkpoint

You deploy a single-layer sliding-window model and users report it can answer about recent text but fails whenever the needed fact is far earlier in a long document. Best explanation?

AThe window is too small only for short promptsBWindowing alone breaks bare retrieval outside the window; stack SWA across layers (Mistral-style composition) or add retrieval augmentation / global tokensCTemperature is too low, so the model ignores early context

Checkpoint

Your serving GPU has plenty of KV-cache memory free, yet throughput is poor and many requests sit idle waiting for a batch to finish. Which lever helps most?

APagedAttention — it will reclaim the wasted cache memoryBContinuous (iteration-level) batching — evict finished sequences and admit new ones after every forward pass, so the GPU stays saturatedCQuantize the weights to 4-bit to free compute

Checkpoint

On H100s you want both lower memory and higher matmul throughput for a Llama-class model, with minimal accuracy work. Best quantization choice?

AGPTQ W4 weight-onlyBSmoothQuant W8A8CFP8 weights/activations (and FP8 KV cache) via Transformer Engine — Hopper roughly doubles matmul throughput and cuts memory with minimal accuracy work

Could you choose GQA vs MQA, justify FlashAttention by IO, design a PagedAttention + continuous-batching stack, and pick quantization by hardware?

New to itGetting thereConfident

Takeaways

KV-head count is the top serving lever: GQA-8 ≈ MHA quality at ~MQA speed; uptrain MHA→GQA for ~5% compute, before fine-tuning.
FlashAttention is exact, not approximate — tile in SRAM, never materialize T×T in HBM; 2–3× (FA-1), 50–73% A100 peak (FA-2), FP8 on Hopper (FA-3).
Stacked sliding-window (Mistral W=4096×32 ≈ 131k reach) is linear; isolated windows lose retrieval — combine with RAG for million-token contexts.
PagedAttention (waste 60–80%→<4%) + continuous batching (Orca 36.9×) compose — don’t credit one for the combined lift.
Quantization is hardware-bound: W4 saves memory, W8A8 speeds Ampere, FP8 does both on Hopper; always ablate at your longest context.

Finally: the capstone — annotate a forward pass end to end, then size a real serving deployment.

Sources

Free to read · better with Enzo

Learn it with Enzo

Save your progress, answer the checkpoints, and let Enzo quiz you on what you just read.

Continue with Enzo Sign up free