← AI Terminology
Flash Attention
Flash Attention is an IO-aware, memory-efficient implementation of the attention mechanism that computes exact attention without materialising the full attention matrix in GPU HBM — achieving 2–4× speedup and O(n) memory instead of O(n²) for sequence length n.
Introduced by Tri Dao et al. (Stanford, 2022), Flash Attention 2 and 3 further improved performance.
Introduced by Tri Dao et al. (Stanford, 2022), Flash Attention 2 and 3 further improved performance.
Why It Matters in AI
Standard attention is the training and inference bottleneck at long context lengths: computing a 100K×100K attention matrix requires petabytes of memory bandwidth. Flash Attention tiles the computation in GPU SRAM (fast on-chip memory), reads/writes HBM (slow off-chip memory) only once, and produces bit-identical results to standard attention. It enabled long-context LLMs economically and is now the default attention implementation in PyTorch, HuggingFace, and every frontier model.
Key Points
| Aspect | Description |
|---|---|
| Exact | Mathematically identical to standard attention — not an approximation |
| Speed | 2–4× faster than standard PyTorch attention; Flash Attention 3 on H100 achieves >500 TFLOP/s |
| Impact | Made 128K–1M token context windows economically viable; adopted by all frontier LLM labs |
| Memory | O(N) HBM memory (vs O(N²)) — linear in sequence length |
| Availability | PyTorch 2.0+ SDPA, flash-attn package, cuDNN 9 native kernels |
| Key innovation | Fused CUDA kernel: tiles QK^T softmax V in SRAM — avoids materialising full N×N matrix in HBM |
Simple Analogy
Standard attention is like running every calculation through a large external hard drive (HBM) — accurate but slow due to data transfer. Flash Attention does all the heavy arithmetic on the CPU's fast cache (SRAM) and only reads/writes the slow drive for final results — same answer, but far fewer expensive trips to slow memory.
Common Usage Examples
pip install flash-attn— install CUDA kernels;from flash_attn import flash_attn_func- PyTorch 2.0:
with torch.backends.cuda.sdp_kernel(enable_flash=True): F.scaled_dot_product_attention(...) - HuggingFace:
model = AutoModelForCausalLM.from_pretrained(..., attn_implementation="flash_attention_2") - Flash Attention 3: H100-specific optimisations — asynchronous GEMM + softmax pipelining
- Context length scaling: Llama 3.1 128K context enabled by Flash Attention without quadratic memory cost
Summary
In short: Flash Attention computes exact attention in linear memory by tiling in fast SRAM instead of materialising the N×N matrix in slow GPU memory — the key enabler of long-context LLMs.