FlashAttention

The Problem

The core of the Transformer is the self-attention mechanism. Mathematically, it requires calculating an $N \times N$ attention matrix (where $N$ is the sequence length). If you double the context window, the memory required quadruples ($O(N^2)$). But the real bottleneck wasn't just FLOPs (compute); it was memory bandwidth. The GPU had to constantly read and write this massive $N \times N$ matrix to its slow High Bandwidth Memory (HBM), leaving the super-fast compute cores (SRAM) sitting idle waiting for data.

The Idea

Tri Dao and the Stanford team realized that if you could calculate attention without ever writing the massive intermediate $N \times N$ matrix to slow memory, you could drastically speed up the model. They developed FlashAttention, an IO-aware algorithm that computes exact attention using "tiling." It loads blocks of the Query, Key, and Value matrices from slow HBM into fast SRAM, performs the attention math on that block, and writes only the final output back to HBM.

How It Works

Standard attention computes the entire $Q K^T$ matrix, writes it to HBM, reads it back to apply softmax, writes it to HBM, reads it back to multiply by $V$, and writes the result.

FlashAttention relies on two techniques to avoid this:

1. Tiling: It breaks the input matrices into blocks that fit perfectly into the GPU's ultra-fast SRAM.
2. Recomputation: During the backward pass (training), instead of reading the massive intermediate matrices from HBM to calculate gradients (which standard attention does), FlashAttention just recomputes them on the fly in SRAM. Recomputing the math is actually faster than reading the data from slow memory.

Why It Mattered

FlashAttention is arguably the most important engineering breakthrough in the LLM era. It sped up Transformer training by 2x to 4x and slashed memory usage from $O(N^2)$ to $O(N)$. Because memory was no longer a quadratic bottleneck, researchers could suddenly train models with massive context windows (jumping from 2k tokens to 32k, 100k, and eventually 1M+ tokens).

What Came After

FlashAttention was universally adopted. It was integrated natively into PyTorch 2.0. The authors released FlashAttention-2 and 3, further optimizing the CUDA kernels for newer GPU architectures (like Hopper), ensuring it remains the absolute foundation of modern AI compute.

What to Read Next

• Flow Matching
• FID (Fréchet Inception Distance)