Inference Optimization: The Bandwidth Wall, PagedAttention, and FlashAttention
Why is autoregressive token generation bottlenecked by memory bandwidth rather than GPU compute, and how does the KV cache scale?
How this lesson unfolds.
Training an LLM is only the first step. Serving it to millions of users in real time requires overcoming physical hardware bottlenecks. In this capstone chapter, we master modern inference engineering.
Why is autoregressive token generation bottlenecked by memory bandwidth rather than GPU compute, and how does the KV cache scale?
Imagine a super-fast chef who can chop an onion in 1 millisecond. But the pantry is 100 meters away, and the assistant can only bring one onion at a time. The chef spends 99.9% of their time waiting for the assistant to walk back and forth! GPUs are the chef, HBM is the pantry, and memory bandwidth is the walking speed.
- $$\text{KV}_{\text{bytes}} = 2 \cdot L \cdot B \cdot T \cdot d \cdot \text{bytes\_per\_elem}$$
- Learning goalCalculate operational intensity and quantify KV cache memory growth across batch sizes and sequence lengths.
- Helpful foundationsAutoregressive decoding · GPU memory hierarchy
- Common trapBelieving GPU FLOP capacity determines token generation speed; single-token autoregressive decoding speed is dictated by memory bandwidth.
- Next connectionEliminate memory fragmentation in KV caches using PagedAttention in Screen 2.
More context & mechanics
In the roofline model of computer architecture:
Operational Intensity = (FLOPs) / (Bytes Transferred)
- Prefill Phase (Prompt Ingestion): Compute-bound. Matrix multiplications process B × T tokens simultaneously, achieving high operational intensity (>100 FLOPs/byte).
- Decode Phase (Token Generation): Memory-bandwidth bound. Generating 1 token requires loading all Φ model parameters:
Time per token ≥ (2 × Φ bytes) / (HBM Bandwidth (Bytes/s))
The KV Cache Memory Footprint:
To avoid recomputing keys and values for past tokens, the KV cache stores tensors across layers:
KV Cache Size = 2 × L × B × T × d × precision (bytes)
As context length T and concurrency B grow, KV cache memory quickly dwarfs the model weights themselves!
Consider serving LLaMA-3 (8B parameters, FP16 = 16 GB weights) on an NVIDIA A100 GPU (2,000 GB/s HBM bandwidth):
- Theoretical Maximum Decoding Speed (Batch Size 1):
Time per token = (16 GB) / (2,000 GB/s) = 0.008 seconds = 8.0 ms ⟹ 125 tokens/sec
Even if the GPU has 312 TFLOPs of tensor core compute, it sits idle >95\% of the time waiting for weights to arrive from HBM!
- KV Cache Size for 64 Concurrent Users at T=4096 tokens (L=32, d=4096):
KV Size = 2 × 32 × 64 × 4096 × 4096 × 2 bytes = 137.4 GB VRAM
The KV cache alone requires two entire 80GB A100 GPUs!
The primary design constraint governing all production LLM serving architectures.
Williams, Waterman, and Patterson (Berkeley, 2009) established the Roofline Model, formalizing memory bandwidth bottlenecks in deep learning.
Prefill is compute bound, decode is bandwidth bound: the memory bus is the true clock of AI.
How does PagedAttention eliminate memory fragmentation and dynamic allocation waste in LLM serving?
In old operating systems, if an app wanted 100MB of RAM, it needed one unbroken, contiguous 100MB slab. If RAM was scattered in small gaps, the app crashed with an Out-of-Memory error! OS virtual paging solved this by chopping RAM into 4KB pages. PagedAttention brings this exact OS revolution to the KV cache.
- $$\text{WastedMemory} \le \frac{B_{\text{size}} - 1}{T} \ll \text{ContiguousPaddingLoss}$$
- Learning goalExplain the architectural principles of PagedAttention and calculate memory savings from eliminating KV cache fragmentation.
- Helpful foundationsThe Memory Bandwidth Wall · Operating system virtual memory paging
- Common trapBelieving KV cache tensors must reside in contiguous physical VRAM; PagedAttention gathers non-contiguous physical blocks on the fly inside GPU kernels.
- Next connectionReduce the total number of KV heads using MQA and GQA in Screen 3.
More context & mechanics
In traditional serving (Hugging Face Transformers), allocating contiguous tensors for the maximum sequence length (e.g. T=4096) causes severe waste:
- Reserved Memory Waste (60-80\%): Pre-allocating memory for tokens that haven't been generated yet.
- Internal Fragmentation: Padding sequences to identical batch dimensions.
- External Fragmentation: Memory holes between deallocated requests.
The PagedAttention Architecture (vLLM):
- Fixed-Size Physical Blocks: The KV cache is divided into blocks holding a fixed number of tokens (e.g. B_(size) = 16 tokens).
- Block Tables (Page Tables): Logical sequence positions map to arbitrary, non-contiguous physical blocks in GPU VRAM.
- Copy-on-Write for Parallel Sampling: Multiple sampled paths share identical prompt KV blocks; a new physical block is allocated only when a branch diverges!
A server handles 10 requests with max length T=2048, but average request length is only T=512:
- Traditional Contiguous Allocation:
- Must reserve memory for 10 × 2048 = 20,480 tokens.
- Active tokens used: 10 × 512 = 5,120 tokens.
- Memory Waste: (20480 - 5120) / (20480) = 75.0\% waste!
- PagedAttention (Block size 16):
- Allocates exactly ⌈ 512 / 16 ⌉ = 32 blocks per request.
- Total allocated: 10 × 32 × 16 = 5,120 tokens.
- Memory Waste: Near zero (<4\% for the final partially filled block).
- Throughput Leap: The exact same GPU can now serve 4× more concurrent requests!
Core technology powering vLLM, TensorRT-LLM, Hugging Face TGI, and modern cloud inference platforms.
Woosuk Kwon et al. (UC Berkeley LMSYS, 2023) published 'Efficient Memory Management for Large Language Model Serving with PagedAttention'.
Chop KV into pages, map through a table: zero wasted memory, maximum serving capacity.
How do Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) dramatically compress KV cache memory while preserving multi-head expressiveness?
Imagine a classroom of 32 students (query heads). In MHA, every student has their own private textbook (KV head)—that's 32 textbooks! In MQA, all 32 students share one single textbook, saving paper but causing fighting. In GQA, the class is divided into 4 study groups of 8 students each, sharing 4 textbooks: comfortable reading and 87.5% less paper!
- $$\text{CompressionRatio} = \frac{G}{H} \quad (\text{where } G \le H)$$
- Learning goalCompare MHA, MQA, and GQA architectures and calculate KV cache compression ratios across head grouping factors.
- Helpful foundationsMulti-head attention projections · The Memory Bandwidth Wall
- Common trapBelieving GQA reduces model parameter count significantly; GQA reduces KV projection parameters slightly, but its massive breakthrough is slashing runtime KV cache memory by 8x.
- Next connectionScale context lengths smoothly using Rotary Position Embeddings (RoPE) in Screen 4.
More context & mechanics
Comparing the three attention architectures for H query heads:
- Multi-Head Attention (MHA): H query heads, H key heads, H value heads.
KV Cache Ratio = 1.0×
- Multi-Query Attention (MQA): H query heads, 1 key head, 1 value head.
KV Cache Ratio = 1/H (e.g. 1/32 ≈ 3.1\% of MHA size)
Tradeoff: Up to 30× KV compression, but suffers slight degradation in complex multi-task reasoning.
- Grouped-Query Attention (GQA): H query heads grouped into G groups (G ∈ 4, 8), each sharing 1 key-value head.
KV Cache Ratio = G/H (e.g. 8/32 = 25\% of MHA size)
Tradeoff: Matches 99\% of MHA quality while slashing KV cache memory by 75-87.5\%.
Consider LLaMA-3 70B (H = 64 query heads, G = 8 KV heads, d_(head) = 128, L = 80 layers) at context T = 8192:
- Under Traditional MHA (64 KV heads):
KV_(MHA) = 2 × 80 × 1 × 8192 × (64 × 128) × 2 bytes = 21.47 GB per user!
A single user at 8k context eats an entire 24GB GPU!
- Under GQA (G = 8 KV heads — 8:1 ratio):
KV_(GQA) = 2 × 80 × 1 × 8192 × (8 × 128) × 2 bytes = 2.68 GB per user!
The KV cache footprint shrinks by 8×, allowing the server to host 8 concurrent users on the exact same hardware.
Standard architecture in all modern open and proprietary LLMs (Llama 3, Mistral, Gemma 2, DeepSeek-V2).
Noam Shazeer (Google, 2019) introduced MQA. Ainslie et al. (Google Research, 2023) developed GQA, bridging the quality gap.
MHA gives each query its own key; MQA shares one for all; GQA groups them into clusters: eight-fold memory savings with zero quality fall.
How does Rotary Position Embedding (RoPE) encode relative token distance, and how do frequency scaling techniques extend context length?
Imagine clock hands ticking. A minute hand rotates once every hour; an hour hand rotates once every 12 hours. RoPE rotates mathematical vectors around clock faces based on token position. To stretch a 4,000-page book onto a 128,000-page book without confusing the hands, you slow down the clock gears smoothly so the hands never spin out of control.
- $$(\mathbf{R}_m\mathbf{q})^T (\mathbf{R}_n\mathbf{k}) = \mathbf{q}^T \mathbf{R}_{n-m} \mathbf{k}$$
- Learning goalDerive RoPE relative position invariance and evaluate NTK-aware frequency scaling for context extension.
- Helpful foundationsPositional embeddings · 2D rotation matrices in linear algebra
- Common trapBelieving RoPE adds positional vectors to embeddings; RoPE performs multiplicative orthogonal 2D rotations on queries and keys.
- Next connectionEliminate the quadratic attention memory bottleneck using FlashAttention hardware tiling in Screen 5.
More context & mechanics
Traditional absolute position embeddings add a vector pₘ to token embedding xₘ.
Rotary Position Embedding (RoPE):
RoPE (Su et al., 2021) multiplies 2D slices of query and key vectors by a 2D rotation matrix:
R_(\Theta, m)^((i)) = cos(m θᵢ) & -sin(m θᵢ) \ sin(m θᵢ) & cos(m θᵢ) , θᵢ = 10000^(-2(i-1)/d)
When computing the inner product qₘ^T kₙ, the rotation angles subtract:
qₘ^T kₙ = (Rₘ q)^T (Rₙ k) = q^T R_(n-m) k
The attention score depends strictly on the relative distance n - m!
Frequency Scaling for Long Context:
When evaluating at context L_(new) = s · L_(train):
- Linear Position Interpolation (PI): Scale positions m' = m / s. (Causes high-frequency loss).
- NTK-Aware Scaling: Scales base frequency θ_(base) = 10000 · s^(d/(d-2)), preserving high frequencies while stretching low frequencies.
Let dimension d = 2 and base frequency θ = 0.1 rad/token.
Consider token at position m = 3 and token at position n = 5 (relative distance n - m = 2):
- Rotation angles: φₘ = 3 × 0.1 = 0.3 rad, φₙ = 5 × 0.1 = 0.5 rad.
- Input vectors: q = [1, 0]^T, k = [1, 0]^T.
- Rotated vectors:
q₃ = [cos(0.3), sin(0.3)] = [0.9553, 0.2955]
k₅ = [cos(0.5), sin(0.5)] = [0.8776, 0.4794]
- Dot Product:
q₃^T k₅ = (0.9553 × 0.8776) + (0.2955 × 0.4794) = 0.8384 + 0.1417 = 0.9801
- Direct Relative Formula Check:
cos((5 - 3) × 0.1) = cos(0.2) = 0.9801
The attention score is identical and depends solely on the distance of 2 tokens!
Universal positional backbone in LLaMA, Mistral, Qwen, and Gemma models supporting 128k+ contexts.
Jianlin Su introduced RoPE in 2021. Peng et al. (2023) developed YaRN and NTK-aware scaling.
Rotate 2D pairs by angle m theta; dot product captures relative distance: scale frequencies to stretch context to infinity.
Why is standard attention bottlenecked by GPU High-Bandwidth Memory (HBM), and how does FlashAttention's on-chip SRAM block tiling solve it?
Imagine baking 1,000 cookies. The standard attention baker has a tiny kitchen counter and a giant pantry down three flights of stairs in the basement. Every time they mix two ingredients, they walk down to the basement pantry, write down the intermediate batter weight in a ledger, walk back upstairs, mix one more egg, walk back down to the pantry, and store the bowl in the basement. They spend 95% of their day walking up and down stairs. FlashAttention is a master baker who rolls a mobile stainless steel prep cart (SRAM) right next to the oven: all mixing, cracking, and whisking occurs directly on the cart in small batches, walking to the basement pantry only once to load ingredients and once to store the finished cookies.
- Learning goalExplain the GPU memory hierarchy bottleneck in standard attention and illustrate how FlashAttention's SRAM block tiling avoids materializing T × T matrices in HBM.
- Helpful foundationsScaled dot-product attention formulation · FlashAttention: Fast and Memory-Efficient Exact Attention (Dao et al., 2022)
- Common trapBelieving that FlashAttention is an approximation like Linformer or Performer. FlashAttention is an exact mathematical computation of standard attention; it produces mathematically identical numerical outputs with zero approximation error.
- Next connectionNext, derive the online softmax rescaling calculus and SRAM tile update equations in Part B.
More context & mechanics
Modern GPU performance is governed by the Memory Wall:
- GPU Memory Hierarchy Reality:
- SRAM (On-Chip Cache): Lightning fast (≈ 19 TB/s bandwidth), but tiny (100 KB per Streaming Multiprocessor, ≈ 20 MB total).
- HBM (High-Bandwidth Memory / VRAM): Massive (80 GB), but relatively slow (≈ 2 to 3 TB/s bandwidth).
- The Standard Attention Flaw (Memory-Bound):
Standard attention executes as separate kernel launches:
- Compute S = Q K^T ⟹ Write T × T matrix to slow HBM (O(T²) memory).
- Read S from HBM ⟹ Compute A = softmax(S) ⟹ Write A to HBM.
- Read A from HBM ⟹ Compute O = A V ⟹ Write O to HBM.
Even though the GPU has teraflops of compute power, execution is stalled waiting for memory transfers over the HBM bus.
- FlashAttention Solution (Tri Dao et al., 2022):
FlashAttention loads blocks of Q, K, V into fast SRAM (e.g. 128 × 128 tiles), computes attention locally in cache, and updates output O incrementally. The massive T × T attention matrix is never written to HBM.
FlashAttention-1, 2, and 3 are integrated directly into PyTorch (`torch.nn.functional.scaled_dot_product_attention`). It delivers a 2× to 4× wall-clock speedup and reduces memory footprint from O(T²) down to O(T), enabling 128k+ context windows in Llama-3.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré published FlashAttention at Stanford in 2022. By designing an IO-aware algorithm that prioritized hardware memory hierarchies over raw FLOP counts, they revolutionized transformer inference and training.
- SRAM is gold; HBM is slow: Don't waste time running to the basement; compute in on-chip cache.
- IO-aware over FLOP-aware: Faster memory access beats fewer arithmetic operations.
How does the online softmax algorithm compute exact softmax normalization incrementally across tiled memory blocks?
Imagine calculating the grade average for a class of 10,000 students. In the standard method, you must read all 10,000 grades to find the highest score, write it down, read all 10,000 grades again to subtract the max and compute exponentials, and read them a third time to divide by the sum. In online softmax, you process students in small groups of 10. If group 2 has a higher student than group 1, you don't start over; you simply multiply group 1's previous running total by a correction factor (e^m_(old) - m_(new)) to align it with the new higher champion. The math remains 100% exact in a single pass.
- $$Online softmax updates running maximum $m_{\text{new}} = \max(m_{\text{old}}, m_{\text{block}})$ and rescales previous partial sums by $e^{m_{\text{old}} - m_{\text{new}}}$, enabling exact block-by-block accumulation.$$
- Learning goalDerive the online softmax rescaling formulas and trace block-by-block numerical normalization without global memory materialization.
- Helpful foundationsFlashAttention hardware concepts · Softmax max-subtraction stabilization calculus
- Common trapAssuming that online softmax introduces cumulative numerical drift. Because the rescaling factor e^m_(old) - m_(new) is algebraically exact, online softmax is mathematically identical to global softmax.
- Next connectionNext, transition to foundational neural network mechanics with the anatomy of a single artificial neuron in Neural Networks Screen 1.
More context & mechanics
#### 1. The Block Softmax Challenge
Standard softmax requires computing the maximum over the entire sequence:
m = max_(1 ≤ j ≤ T) xⱼ, d = ∑_(j=1)^T e^(xⱼ - m), pᵢ = (e^(xᵢ - m)) / (d)
If you tile the sequence into blocks B₁, B₂, ṡ, you cannot compute the global denominator d until all blocks are processed.
When processing new tile block Bₖ:
- Compute local block maximum:
m_(block) = max_(j ∈ Bₖ) xⱼ
- Update global running maximum:
m_(new) = max(m_(old), m_(block))
- Compute local block unnormalized denominator:
d_(block) = ∑_(j ∈ Bₖ) e^xⱼ - m_(new)
- The Rescaling Correction Factor:
To incorporate the previous accumulated sum d_(old) (which was normalized by m_(old)), scale it by the exponent of the maximum shift:
d_(new) = d_(old) · e^m_(old) - m_(new) + d_(block)
- Output Vector Update:
Similarly, the accumulated output representation vector O is rescaled and updated in-place:
O_(new) = O_(old) · ( (d_(old) · e^m_(old) - m_(new)) / (d_(new)) ) + V_(block) · ( (e^x_(block) - m_(new)) / (d_(new)) )
Let input be split into two 2-element blocks: B₁ = [2.0, 4.0] and B₂ = [1.0, 5.0].
- Block 1 Processing:
- m₁ = max(2.0, 4.0) = 4.0.
- Exponentials: [e^(2-4), e^(4-4)] = [e^(−2), e⁰] = [0.1353, 1.0].
- Sum d₁ = 0.1353 + 1.0 = 1.1353.
- Block 2 Processing:
- Local max: m_(block) = max(1.0, 5.0) = 5.0.
- New global max: m₂ = max(m₁=4.0, 5.0) = 5.0.
- Rescale factor for Block 1: e^(m₁ - m₂) = e^(4.0 - 5.0) = e^(−1) ≈ 0.3679.
- Corrected previous sum: d₁' = 1.1353 × 0.3679 = 0.4177.
- Block 2 exponentials with new max 5.0: [e^(1-5), e^(5-5)] = [e^(−4), e⁰] = [0.0183, 1.0].
- Block 2 sum: d_(block) = 0.0183 + 1.0 = 1.0183.
- Combined global sum: d₂ = d₁' + d_(block) = 0.4177 + 1.0183 = 1.4360.
- Verification against standard global softmax on [2, 4, 1, 5] (m=5.0):
e^(−3) + e^(−1) + e^(−4) + e⁰ = 0.0498 + 0.3679 + 0.0183 + 1.0 = 1.4360.
The online formula matches the full global calculation to the exact last decimal place! \blacksquare
FlashAttention-2 and 3 implement this online rescaling loop directly in GPU registers. By avoiding all intermediate HBM writes, it achieves up to 73% of theoretical peak FP16 TFLOPs on NVIDIA H100 GPUs.
Maxim Milakov and Natalia Gimelshein at NVIDIA developed the online softmax algorithm in 2018. Tri Dao et al. recognized that combining Milakov's online softmax with GPU shared memory tiling was the missing key to conquering the attention memory wall.
- e^m_(old) - m_(new): The magic rescaling bridge that links past blocks to new champions.
- Single-pass exactness: Calculate global softmax in local tiles with zero memory bloat.
How does Key-Value (KV) Cache: Eliminating Redundant Attention Recomputation operate mathematically at scale?
Visual comparison of next-token generation:
- $$Title in white sans-serif: "The KV Cache: Eliminating O(L^2) Autoregressive Waste".$$
- $$Subtitle in muted cyan: "Caching past Key and Value projection tensors to reduce next-token generation from O(L^2) to O(L)."$$
- Learning goalMaster the algebraic and geometric derivation of Key-Value (KV) Cache: Eliminating Redundant Attention Recomputation.
- Helpful foundationsFoundations of Build Llm From Scratch · Coordinate Invariance & Geometric Projections · Vaswani et al. (2017), Attention Is All You Need, NeurIPS
- Common trapAssuming Key-Value (KV) Cache: Eliminating Redundant Attention Recomputation is an arbitrary computational trick rather than an invariant spatial law.
- Next connectionConnects upstream foundational lemmas to modern high-dimensional computational systems.
More context & mechanics
An inference computational cache diagram in square 1:1 format on dark slate (#0f172a) with modular memory buffers.
Top section: Title in white sans-serif: "The KV Cache: Eliminating O(L²) Autoregressive Waste". Subtitle in muted cyan: "Caching past Key and Value projection tensors to reduce next-token generation from O(L²) to O(L)."
Middle section: Visual comparison of next-token generation:
Top (Without Cache - Redundant Disaster):
To generate token t = 5, the model re-projects all tokens 0, 1, 2, 3, 4 through W_Q, W_K, W_V from scratch. Generates L² total projection operations across sequence!
Bottom (With KV-Cache - Efficient Standard):
- Keys K_(0:t-1) and Values V_(0:t-1) from previous tokens are stored in GPU VRAM memory buffers.
- For new incoming token x_t, compute ONLY single query q_t, single key k_t, and single value v_t.
- Append k_t and v_t to cache.
- Attention computed as q_t · [K_(cache), k_t]^\top.
Bottom section: Memory tradeoff: "Transforms computation from compute-bound to memory-bandwidth-bound: generation speed is governed by how fast the GPU reads the KV-cache from VRAM."
Clear buffer append diagrams, highlighted single-vector calculation, modern dark aesthetic.
Every transformation preserves structural invariants while mapping coordinate spaces. When operating in high dimensions, numerical precision and computational efficiency require adhering strictly to canonical mathematical laws.
Formal derivation and peer-reviewed analysis:
"Key-Value (KV) Cache: Eliminating Redundant Attention Recomputation: Invariance, symmetry, and structure govern mathematical truth across continuous and discrete domains."
How does PagedAttention & vLLM: Virtual Memory for Non-Contiguous KV-Caches operate mathematically at scale?
Visual comparison of GPU VRAM allocation:
- Learning goalMaster the algebraic and geometric derivation of PagedAttention & vLLM: Virtual Memory for Non-Contiguous KV-Caches.
- Helpful foundationsFoundations of Build Llm From Scratch · Coordinate Invariance & Geometric Projections · Vaswani et al. (2017), Attention Is All You Need, NeurIPS
- Common trapAssuming PagedAttention & vLLM: Virtual Memory for Non-Contiguous KV-Caches is an arbitrary computational trick rather than an invariant spatial law.
- Next connectionConnects upstream foundational lemmas to modern high-dimensional computational systems.
More context & mechanics
An operating systems virtual memory architecture diagram in square 1:1 format on dark slate (#0f172a) with memory tables.
Top section: Title in white sans-serif: "PagedAttention: OS Virtual Memory for LLM Serving". Subtitle in muted emerald: "Kwon et al. (vLLM 2023): Eliminating 60-80% KV-cache memory waste via non-contiguous block allocation."
Middle section: Visual comparison of GPU VRAM allocation:
Left (Traditional Static Allocation - Massive Fragmentation):
Pre-allocates contiguous memory for maximum possible sequence length (e.g., 2,048 tokens). Over 70% of memory sits completely idle as internal fragmentation!
Right (PagedAttention - Block Table Paging):
- KV-cache divided into fixed-size physical blocks (e.g., 16 tokens per block).
- Logical blocks map to physical blocks anywhere in VRAM via a Page Table (identical to OS virtual memory!).
- Blocks allocated on-demand as tokens generate.
- Multiple requests (e.g., parallel beam search branches) share identical prompt physical blocks via Copy-on-Write!
Bottom section: Throughput gain: "Shrinks KV-cache memory waste to under 4%, enabling 2x-4x higher concurrency on the same GPU cluster."
Clean virtual memory address mapping lines, distinct block cells, modern dark theme.
Every transformation preserves structural invariants while mapping coordinate spaces. When operating in high dimensions, numerical precision and computational efficiency require adhering strictly to canonical mathematical laws.
Formal derivation and peer-reviewed analysis:
"PagedAttention & vLLM: Virtual Memory for Non-Contiguous KV-Caches: Invariance, symmetry, and structure govern mathematical truth across continuous and discrete domains."