FlashAttention 1, 2 & 3: Memory Tiling & Eliminating the O(N²) HBM Bottleneck

1. The Standard Attention Memory Bottleneck

Standard Self-Attention calculates:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V$$

In standard PyTorch implementations:

1. $QK^T$ materializes an $N \times N$ intermediate matrix in GPU High-Bandwidth Memory (HBM).
2. For an 8,000-token sequence, $N \times N = 64,000,000$ elements.
3. For a 64,000-token sequence, $N \times N = 4,096,000,000$ elements → Out of Memory (OOM) crash!

Standard Attention: Slow HBM Thrashing
GPU SRAM (Fast, 19 TB/s, ~100 KB) ◄────┐
                                        ├─► Read/Write huge N x N matrix repeatedly!
GPU HBM  (Slow, 3 TB/s, ~80 GB)   ◄────┘

2. FlashAttention Mechanics: Tiling & Online Softmax

Pioneered by Tri Dao (Stanford), FlashAttention never materializes the full $N \times N$ matrix in HBM:

• Tiling: Splits $Q, K, V$ matrices into small blocks that fit entirely inside ultra-fast on-chip SRAM (L1 cache).
• Online Softmax: Computes incremental softmax normalization factors dynamically across tiles without needing the global maximum upfront.
• Kernel Fusion: Combines matrix multiply, mask, softmax, and value accumulation into a single fused GPU kernel.

3. FlashAttention Evolution