flashattention-2
**FlashAttention-2** is an **optimized implementation of the attention mechanism that achieves 2× the speed of the original FlashAttention** — reaching up to 230 TFLOPS/s on NVIDIA A100 GPUs (73% of theoretical peak), through better work partitioning across GPU thread blocks, improved parallelism along the sequence length dimension, and elimination of redundant floating-point operations in the online softmax computation, making it the standard attention implementation for all production LLM training and inference.
**What Is FlashAttention-2?**
- **Definition**: The second-generation IO-aware attention algorithm (Dao, 2023) that computes exact attention without materializing the full N×N attention matrix in GPU high-bandwidth memory (HBM) — using tiling and online softmax to keep all intermediate computations in fast GPU SRAM (~20MB) rather than slow HBM (~80GB).
- **Why "Flash"**: Standard attention writes the full N×N attention matrix to GPU HBM (slow memory), then reads it back for softmax and value multiplication. FlashAttention keeps computations in fast on-chip SRAM, avoiding these slow memory round-trips. It's not an approximation — it computes the exact same result, just faster.
- **v2 Improvements**: FlashAttention-1 was already 2-4× faster than standard attention. FlashAttention-2 adds another 2× by optimizing thread block scheduling, reducing non-matmul FLOPs, and improving sequence-length parallelism.
**FlashAttention-2 Improvements Over v1**
| Optimization | v1 Problem | v2 Solution | Speedup |
|-------------|-----------|------------|---------|
| **Non-matmul FLOPs** | Rescaling operations during online softmax | Eliminate rescaling by restructuring the algorithm | ~15% |
| **Sequence Parallelism** | Parallelized only across batch and heads | Also parallelize across sequence length dimension | ~50% on long sequences |
| **Warp Partitioning** | Suboptimal work distribution between warps | Better partition between thread warps, reducing shared memory reads/writes | ~20% |
| **Causal Masking** | Applied mask to all tiles | Skip computation for fully masked tiles | ~2× for causal (autoregressive) |
**Performance Comparison**
| Implementation | TFLOPS/s (A100) | % of Peak | Memory | Exact? |
|---------------|----------------|-----------|--------|--------|
| **Standard PyTorch** | ~30 | 10% | O(N²) | Yes |
| **FlashAttention v1** | ~120 | 39% | O(N) | Yes |
| **FlashAttention-2** | ~230 | 73% | O(N) | Yes |
| **FlashAttention-3** | ~300+ (H100) | 75%+ | O(N) | Yes |
| **Theoretical Peak** | 312 (A100 BF16) | 100% | — | — |
**How FlashAttention-2 Works (Tiled Algorithm)**
| Step | Action | Memory Level |
|------|--------|-------------|
| 1. Load Q tile from HBM to SRAM | Load Q block (Br × d) | HBM → SRAM |
| 2. Load K, V tiles sequentially | Load K, V blocks (Bc × d) | HBM → SRAM |
| 3. Compute S = Q × K^T (tile) | Matrix multiply in SRAM | SRAM only |
| 4. Online softmax (no rescaling in v2) | Compute softmax incrementally | SRAM only |
| 5. Compute O = softmax(S) × V | Accumulate output tile | SRAM only |
| 6. Write output tile to HBM | Store final result | SRAM → HBM |
| 7. Repeat for all K, V tiles | Iterate through sequence | Overlapped loads |
**Adoption**
| Framework | Integration Status |
|-----------|-------------------|
| **PyTorch 2.0+** | Built-in via `torch.nn.functional.scaled_dot_product_attention` |
| **Hugging Face Transformers** | Default for supported models (`attn_implementation="flash_attention_2"`) |
| **vLLM** | Default attention backend for LLM serving |
| **DeepSpeed** | Integrated for training |
**FlashAttention-2 is the standard attention implementation for modern LLMs** — delivering exact attention computation at 73% of GPU peak throughput through IO-aware tiling, optimized warp scheduling, and sequence-length parallelism, enabling 2-4× faster training and longer context lengths without any approximation or quality loss compared to standard attention.