GPU Memory Architecture & FlashAttention-3: Solving the Memory Bandwidth Bottleneck in LLM Inference

As foundation models (LLMs) scale to $100\text{k}+$, $1\text{M}+$, and multi-million token context windows, AI infrastructure engineers face an unavoidable physical constraint: Large Language Model inference is fundamentally memory-bandwidth bound, not compute-bound.

Modern AI accelerators like the NVIDIA H100, H200, and Blackwell B200 boast monstrous compute capacities—reaching up to $4.5\text{ to }9\text{ PetaFLOPs}$ of FP8 Tensor Core arithmetic. Yet, during multi-head attention and token generation (decoding), these compute units sit idle for over $70\%–80\%$ of execution time, starving for data while waiting on memory buses.

The standard multi-head self-attention formulation: $$\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}}\right)\mathbf{V}$$ requires materializing a massive $N \times N$ intermediate attention matrix into High Bandwidth Memory (HBM). For a 128k token context window, this requires reading and writing gigabytes of intermediate memory on every single token step—a computational disaster known as the Memory Wall.

FlashAttention (developed by Tri Dao et al.) revolutionized AI systems by redesigning attention from an IO-aware hardware perspective. And with FlashAttention-3, the algorithm reaches near-theoretical hardware utilization on modern GPU architectures through SRAM Tiling, Online Softmax Normalization, and Hardware Warp Specialization.

In this technical deep dive, we explore the physical GPU memory hierarchy, dissect why naive attention collapses under the Roofline Model, analyze FlashAttention-3's mathematical and hardware breakthroughs, and explore PagedAttention & KV-Cache quantization in production serving engines like vLLM.

The GPU Memory Hierarchy: HBM vs. SRAM

To understand why IO-aware algorithms are critical, we must examine the physical memory hierarchy of modern AI GPUs (e.g. NVIDIA H100 SXM5 / B200):

Physical GPU Memory Hierarchy & Bandwidth Latencies:

┌─────────────────────────────────────────────────────────────┐
│ High Bandwidth Memory (HBM3 / HBM3e)                        │
│   • Capacity: 80 GB – 192 GB                                │
│   • Bandwidth: ~3.35 TB/sec                                 │
│   • Access Latency: ~500+ CPU/GPU clock cycles (SLOW!)      │
└─────────────────────────────┬───────────────────────────────┘
                              │ Memory Bus Bottleneck
┌─────────────────────────────▼───────────────────────────────┐
│ L2 Cache (Shared across all Streaming Multiprocessors)      │
│   • Capacity: 50 MB – 128 MB                                │
│   • Bandwidth: ~12 TB/sec                                   │
│   • Access Latency: ~100 clock cycles                       │
└─────────────────────────────┬───────────────────────────────┘
                              │
┌─────────────────────────────▼───────────────────────────────┐
│ On-Chip SRAM (Shared Memory / L1 Cache per SM)              │
│   • Capacity: ~228 KB per SM (Total: ~30 MB on-chip)        │
│   • Bandwidth: ~33 TB/sec (10x faster than HBM!)            │
│   • Access Latency: ~10 clock cycles (BLAZING FAST!)        │
└─────────────────────────────┬───────────────────────────────┘
                              │
┌─────────────────────────────▼───────────────────────────────┐
│ Tensor Core Registers (ALU Units)                           │
│   • Compute Capacity: Up to 4.5+ PetaFLOPs FP8              │
│   • Access Latency: 1 clock cycle                           │
└─────────────────────────────────────────────────────────────┘

The Standard Attention Bottleneck: The $O(N^2)$ Memory Wall

In standard PyTorch attention, the calculation executes as a sequence of discrete CUDA kernel launches, each writing intermediate results back to high-latency HBM:

Standard Attention Kernel Pipeline (High HBM I/O):

1. Kernel 1: Load Q, K from HBM -> Compute S = Q * K^T (N x N floats) -> WRITE S TO HBM!
                                                                                │
2. Kernel 2: Read S from HBM -> Compute P = softmax(S) -> WRITE P TO HBM! <────┘
                                                                │
3. Kernel 3: Read P, V from HBM -> Compute O = P * V -> Write O to HBM! <──────┘

Total HBM Memory Traffic: O(N^2) bytes written and read repeatedly from slow HBM!
At N = 64k tokens, the intermediate matrix P consumes 32 GB of RAM per layer!

FlashAttention Breakthrough 1: Tiling & Online Softmax

FlashAttention eliminates intermediate $N \times N$ memory materialization entirely by tiling the $\mathbf{Q}$, $\mathbf{K}$, and $\mathbf{V}$ matrices into small sub-blocks ($B_r \times d$) that fit completely within the on-chip SRAM ($228\text{ KB}$ per SM).

FlashAttention SRAM Tiling:

HBM (Global GPU RAM)
   Q (N x d)          K (N x d)          V (N x d)
       │                  │                  │
       ▼                  ▼                  ▼
┌────────────────────────────────────────────────────────┐
│ On-Chip SRAM (Shared Memory)                           │
│   Load Block Q_i (Br x d)                              │
│   Loop over Blocks K_j, V_j (Bc x d):                  │
│     1. Compute S_ij = Q_i * (K_j)^T (In SRAM)          │
│     2. Online Softmax Scaling (Update running max/sum) │
│     3. Accumulate Output O_i += P_ij * V_j (In SRAM)   │
└─────────────────────────────┬──────────────────────────┘
                              │
                              ▼ (Write FINAL output directly!)
                       O_i to HBM (N x d)

The Online Softmax Normalization Trick

Standard softmax requires knowing the global maximum value across the entire row ($\max(x)$) before computing the denominator $\sum e^{x_i - m}$. How can we compute softmax in incremental tiles without seeing future blocks?

FlashAttention uses the Online Softmax Rescaling Identity:

When moving from block $j-1$ to block $j$, given running row maximum $m_{\text{new}} = \max(m_{\text{prev}}, \max(S_j))$ and running denominator $d_{\text{new}} = e^{m_{\text{prev}} - m_{\text{new}}} d_{\text{prev}} + \sum e^{S_{j,k} - m_{\text{new}}}$, the previous partial output accumulator $O_{\text{prev}}$ is dynamically rescaled:

$$O_{\text{new}} = \text{diag}\left(e^{m_{\text{prev}} - m_{\text{new}}}\right) O_{\text{prev}} + P_j V_j$$

At the end of the block loop, a single division by $d_{\text{final}}$ yields the mathematically exact attention output—without ever saving a single $N \times N$ float to HBM!

FlashAttention-3: Harnessing Hopper & Blackwell Architecture

While FlashAttention-2 achieved ~70% of theoretical H100 FLOPs for FP16, it left substantial performance on the table for FP8 Tensor Core precision. FlashAttention-3 introduces three massive hardware-level innovations:

KV-Cache Optimization: PagedAttention in vLLM

During multi-turn conversational inference (decoding), the model must store past Key and Value vectors (the KV-Cache) so previous tokens don't need to be recomputed.

Traditionally, engines allocated contiguous virtual memory buffers for the maximum sequence length (e.g. 8k tokens), leading to $60\%–80\%$ memory waste due to internal and external memory fragmentation.

PagedAttention (Virtual Memory Paging for KV Cache):

Physical HBM Allocated into Non-Contiguous 16-Token Blocks:
[ Physical Block 0 ]   [ Physical Block 1 ]   [ Physical Block 2 ]   [ Physical Block 3 ]

Logical KV-Cache for Request A (Tokens 0 - 42):
• Page Table A:
    Logical Block 0 (Tokens 0 - 15)  ───> Maps to Physical Block 2
    Logical Block 1 (Tokens 16 - 31) ───> Maps to Physical Block 0
    Logical Block 2 (Tokens 32 - 42) ───> Maps to Physical Block 3

RESULT: Zero Memory Fragmentation, Dynamic Memory Growth, and Zero-Copy Forking!

PagedAttention (the core innovation of vLLM) applies the operating system's Virtual Memory Paging principle to LLM KV-caches. KV tensors are partitioned into small 16-token physical blocks. A page table maps logical token positions to physical blocks, allowing servers to run $2\times$ to $4\times$ higher batch sizes within the exact same GPU VRAM.

Performance Benchmarks: FlashAttention-3 vs. Standard Attention

Summary & Engineering Checklist

• Design for Memory IO, Not Just Math: In modern deep learning systems, the bottleneck is rarely raw floating-point operations; it is the cost of moving bytes between slow HBM and fast on-chip SRAM.
• Enable FlashAttention-3 in Inference Runtimes: On Hopper (H100/H200) and Blackwell GPUs, ensure your serving stack (vLLM, TensorRT-LLM, HuggingFace TGI) enables FlashAttention-3 with TMA warp specialization to unlock up to $1.2\text{ PetaFLOPs}$ per GPU.
• Adopt PagedAttention for Concurrency: Eliminate memory fragmentation in production inference servers by adopting block-paged KV caches, enabling dramatic increases in serving concurrency.
• Leverage FP8 KV-Cache Quantization: Compressing the KV-Cache from FP16 to FP8 cuts memory footprint by 50% with $< 0.1\%$ perplexity degradation, doubling the effective maximum context window length.