FlashAttention & Efficient Attention

GPU Memory Architecture: HBM vs SRAM

To understand FlashAttention (Tri Dao et al., 2022 / Stanford), you must understand how GPUs work.

Modern GPUs have two main memory levels:

1. High Bandwidth Memory (HBM): Main GPU VRAM (for example, 80GB on an NVIDIA A100). Large capacity, but slow bandwidth (2 TB/s).
2. Static RAM (SRAM): Ultra fast on chip cache directly next to Tensor Cores (192KB per Streaming Multiprocessor). Tiny capacity, but ultra fast bandwidth (19 TB/s)!

  GPU HARDWARE MEMORY HIERARCHY
  ┌─────────────────────────────────────────────────────────────┐
  │ HBM (Main GPU VRAM 80GB) ──► Slow Transfer (2 TB/s)         │
  └──────────────────────────────┬──────────────────────────────┘
                                 │ Load Blocks
                                 ▼
  ┌─────────────────────────────────────────────────────────────┐
  │ SRAM (On-Chip Cache 192KB) ──► Ultra-Fast Transfer (19 TB/s)│
  │ Tensor Cores Compute Matrix Math Here!                      │
  └─────────────────────────────────────────────────────────────┘

Standard Attention Memory Bottleneck

Standard attention performs three separate GPU memory steps:

1. Load $Q$ and $K$ from HBM $\to$ Compute $S = Q K^T / \sqrt{d_k} \to$ Write $S$ ($N \times N$) back to HBM.
2. Load $S$ from HBM $\to$ Compute $P = \text{Softmax}(S) \to$ Write $P$ ($N \times N$) back to HBM.
3. Load $P$ and $V$ from HBM $\to$ Compute $O = P V \to$ Write $O$ ($N \times d$) back to HBM.

Notice that the massive $N \times N$ matrix is written to and read from slow HBM memory three times in a row!

Standard attention is not bottlenecked by math compute; it is bottlenecked by slow HBM Memory Read Write IO.

How FlashAttention Solves the Bottleneck

FlashAttention computes exact attention in a single fused GPU kernel using two ideas:

  FlashAttention IO Aware Tiling Pipeline:
  1. Load small Block of Q, K, V from HBM into ultra-fast SRAM (e.g. 128 x 128 elements).
  2. Compute Partial Attention using Online Softmax inside SRAM.
  3. Incrementally update running Output accumulator in SRAM.
  4. Write FINAL Output O directly to HBM! (NEVER write N x N matrix to HBM!)

Online Softmax Trick

Standard Softmax requires knowing the global sum of exponents across all $N$ tokens before dividing.

Online Softmax computes partial exponential sums tile by tile, rescaling the running output vector smoothly whenever a new larger local maximum is found in a new block.

Key Performance Results

1. 2x to 4x Faster Training: Accelerates wall clock Transformer training speed by avoiding HBM memory bottlenecks.
2. Linear Memory Scaling $O(N)$: HBM memory usage scales linearly with sequence length $N$, enabling training on 32k to 128k context windows on standard GPUs.
3. 100 Percent Exact Math: Unlike approximate sparse attention methods, FlashAttention outputs the exact same floating point numbers as standard attention.

Say this out loud

FlashAttention is an IO aware exact attention algorithm. Standard attention is memory bound, repeatedly writing massive N x N matrices to slow GPU HBM VRAM. FlashAttention tiles Query, Key, and Value matrices into small SRAM blocks, using Online Softmax to compute exact attention in SRAM without ever storing the N x N matrix in HBM. This reduces VRAM memory to linear O(N) and accelerates training by 2 to 4 times.

Followups to expect

1. What is FlashAttention 2 (Dao 2023)? Further optimizes work partitioning across GPU Streaming Multiprocessors, improving GPU Occupancy to reach 73 percent of theoretical peak A100 FLOP performance.
2. What is FlashAttention 3 (Dao 2024)? Exploits FP8 low precision tensor cores and asynchronous memory overlapping on NVIDIA Hopper (H100) GPUs, achieving nearly 1 PFLOPS execution speed.