flash attention io aware
**Flash Attention** is the **IO-aware exact attention algorithm that computes self-attention with O(N) memory instead of O(N²) and 2-4× faster wall-clock time than standard attention — by restructuring the attention computation into tiles that fit in GPU SRAM (shared memory), minimizing expensive reads and writes to GPU HBM (global memory), without any approximation or change to the mathematical output**.
**The Memory Bandwidth Bottleneck**
Standard attention implementation:
1. Compute S = QK^T (N×N matrix) — write to HBM: O(N²) writes.
2. Compute P = softmax(S) — read S from HBM, write P: O(N²) reads + writes.
3. Compute O = PV — read P and V from HBM: O(N²) reads.
Total HBM accesses: O(N²). For N=8K with d=128 at FP16: the N×N attention matrix is 128 MB — larger than the GPU's SRAM (~20 MB per SM), forcing multiple round-trips to/from HBM (bandwidth: ~2 TB/s on A100 vs. ~19 TB/s SRAM bandwidth).
**Flash Attention Algorithm**
Key insight: compute attention in tiles without ever materializing the full N×N attention matrix in HBM.
1. **Outer loop**: Iterate over blocks of K and V (block size B_c).
2. **Inner loop**: For each K/V block, iterate over blocks of Q (block size B_r).
3. **In SRAM**: Load a Q block and K/V block into SRAM. Compute the local attention scores S_ij = Q_i · K_j^T. Compute local softmax (with online softmax tracking running max and sum). Compute local output O_ij = softmax(S_ij) · V_j. Accumulate into the output using the running softmax denominator.
4. **Write only O**: The final output O is written to HBM once. The N×N attention matrix never exists in HBM.
**Online Softmax**
The mathematical challenge: softmax requires the max and sum across all K positions, but we process K in blocks. The online softmax algorithm (Milakov & Gimelshein, 2018) maintains running max and exponential sum, updating them as each new K block is processed. This allows correct softmax computation without a second pass over the data.
**Performance Impact**
| Metric | Standard Attention | Flash Attention |
|--------|-------------------|-----------------|
| Memory | O(N²) | O(N) |
| HBM reads/writes | O(N²) | O(N²/M) where M = SRAM size |
| Wall clock (A100, N=4096) | ~15 ms | ~5 ms |
| Max sequence length (40 GB) | ~16K | ~64K+ |
**Flash Attention 2 and 3**
- **Flash Attention 2**: Better work partitioning across GPU warps and thread blocks. Non-matmul FLOPs reduced 2-4×. Achieves 50-73% of A100 peak FLOPS (vs. 25-40% for FA1).
- **Flash Attention 3**: Optimized for Hopper architecture (H100). Uses asynchronous TMA (Tensor Memory Accelerator) for hardware-accelerated data movement, FP8 support, and warp-specialized pipelining. Up to 1.5-2× speedup over FA2 on H100.
**Integration**
Flash Attention is integrated into PyTorch (torch.nn.functional.scaled_dot_product_attention), HuggingFace Transformers, and all major LLM training/inference frameworks. It is the default attention implementation for virtually all modern LLM training.
Flash Attention is **the systems optimization that made long-context Transformers practical** — demonstrating that the attention bottleneck was not the O(N²) computation but the O(N²) memory traffic, and that restructuring the computation to respect the GPU memory hierarchy eliminates the bottleneck without changing a single mathematical operation.