FlashAttention: IO-Aware Exact Attention, Tiling, Online Softmax, and the Evolution to FlashAttention-3

The attention mechanism is the computational bottleneck of every Transformer model. Standard implementations materialize the full $N \times N$ attention matrix in high-bandwidth memory (HBM), incurring $O(N^2)$ memory reads and writes that dominate runtime long before arithmetic intensity saturates the GPU. FlashAttention and its successors eliminate this bottleneck by restructuring the attention computation around the GPU memory hierarchy—trading redundant recomputation for dramatically reduced data movement.

This explainer covers the algorithmic foundations, the progression through FlashAttention-2 and FlashAttention-3, and the practical implications for training and serving long-context LLMs.

1. The IO Problem in Standard Attention

Given query, key, and value matrices $Q, K, V \in \mathbb{R}^{N \times d}$ (sequence length $N$, head dimension $d$), standard attention computes:

$$S = Q K^T \in \mathbb{R}^{N \times N}$$ $$P = \text{softmax}(S)$$ $$O = P V \in \mathbb{R}^{N \times d}$$

The $N \times N$ matrix $S$ is written to HBM, read back for softmax, and written again as $P$; then $P$ is read for the $PV$ multiplication. On an A100 with 1.5 TB/s HBM bandwidth and 19.5 TFLOPs/s FP16 peak, attention at $N=8192, d=128$ spends >90% of cycles moving data rather than computing. The arithmetic intensity (FLOPs/byte) is far below the hardware ridge point.

2. FlashAttention (v1): Tiling, Recomputation, and Online Softmax

2.1 Tiling by Blocks

FlashAttention partitions $Q, K, V$ into blocks that fit in on-chip SRAM (shared memory). For block sizes $B_r, B_c$ (rows of $Q$, columns of $K/V$), the outer loop iterates over blocks of $K$ and $V$, the inner loop over blocks of $Q$. Each block loads $Q_i \in \mathbb{R}^{B_r \times d}$, $K_j \in \mathbb{R}^{B_c \times d}$, $V_j \in \mathbb{R}^{B_c \times d}$ into SRAM, computes the local $S_{ij} = Q_i K_j^T$, and updates the output block $O_i$ incrementally.

This reduces HBM accesses from $O(N^2)$ to $O(N^2 \cdot d / M)$ where $M$ is SRAM capacity—linear in $N$ for fixed $M$.

2.2 Online Softmax (No $N \times N$ Materialization)

Standard softmax requires the full row of $S$ to compute the normalization constant $\sum_j e^{S_{ij}}$. FlashAttention computes softmax online using the classic stable formulation:

Initialize $m_i = -\infty$, $l_i = 0$, $o_i = 0$ for each output row $i$.

For each $j$ block: $$S_{ij} = Q_i K_j^T$$ $$\tilde{m}_{ij} = \max(m_i, \text{rowmax}(S_{ij}))$$ $$\tilde{l}_{ij} = l_i \cdot e^{m_i - \tilde{m}_{ij}} + \text{rowsum}(e^{S_{ij} - \tilde{m}_{ij}})$$ $$o_i \leftarrow o_i \cdot e^{m_i - \tilde{m}_{ij}} + \text{diag}(e^{S_{ij} - \tilde{m}_{ij}}) \cdot V_j$$ Update $m_i \leftarrow \tilde{m}_{ij}$, $l_i \leftarrow \tilde{l}_{ij}$.

After the $K/V$ loop, normalize: $O_i = o_i / l_i$.

No $N \times N$ matrix ever leaves SRAM. The backward pass recomputes $S_{ij}$ from saved $Q, K, V$ blocks—trading extra FLOPs for saved bandwidth.

2.3 Block-Sparse Attention

FlashAttention extends naturally to block-sparse patterns (local window, strided, random). Only selected $K_j, V_j$ blocks are loaded, and the same online softmax accumulates over the chosen sparsity pattern. This yields an approximate attention algorithm faster than any dense baseline while retaining exact computation within attended blocks.

3. FlashAttention-2: Better Parallelism and Work Partitioning

FlashAttention-1 parallelized over batch and heads but kept sequence-length loops sequential within a block. FlashAttention-2 (Dao, 2023) introduces three key changes:

1. Parallelism over sequence length: The outer loop over $K/V$ blocks is parallelized across thread blocks. Each thread block computes a disjoint chunk of the output $O_i$ rows. This requires a reduction across thread blocks for the softmax statistics $(m_i, l_i)$, implemented via atomicAdd in shared memory followed by a cross-block sync.
2. Work partitioning for load balance: For causal attention, the number of $K/V$ blocks per $Q$ row varies (first row attends to 1 block, last to $N/B_c$). FlashAttention-2 assigns work by backward block index so each thread block handles a contiguous range of $Q$ rows with roughly equal $K/V$ iterations.
3. Tuning for Hopper/Ampere: Block sizes are chosen to maximize occupancy—$B_r=64, B_c=64$ for $d=64$; $B_r=128, B_c=64$ for $d=128$—and kernel fusion eliminates intermediate stores.

Results: FlashAttention-2 reaches 230 TFLOPs/s (73% MFU) on A100, 1.7–3.0× faster than v1, and 2.8× end-to-end training speedup on GPT-style models at 8K context.

4. FlashAttention-3: Asynchrony, Warp Specialization, and FP8

FlashAttention-3 (Shah et al., 2024) targets the H100 (Hopper) architecture and introduces:

4.1 Asynchronous Tensor Memory Accelerator (TMA) + Warp Specialization

Hopper's TMA enables asynchronous copies between HBM and shared memory without SM involvement. FlashAttention-3 dedicates warps to:

• Producer warps: Issue TMA loads for $K_j, V_j$ blocks.
• Consumer warps: Execute MMA (tensor core) matrix multiplies on data already in shared memory.
• Scheduler warps: Overlap softmax reduction with the next TMA load.

This pipeline hides HBM latency behind computation, sustaining 740 TFLOPs/s (75% MFU) on H100 FP16.

4.2 FP8 with Block Quantization and Incoherent Processing

FlashAttention-3 supports FP8 E4M3/E5M2 with per-block quantization scales. Crucially, it uses incoherent processing: the softmax exponent is computed in FP32 from FP8-dequantized values, avoiding the catastrophic error accumulation of naive FP8 softmax. With this scheme, FP8 FlashAttention-3 achieves 2.6× lower error than standard per-tensor FP8 attention under outlier features, and approaches 1.2 PFLOPs/s on H100.

4.3 Variable-Length Sequences and Packed Kernels

A single kernel handles variable-length sequences via a packed representation with cuSeqlens, eliminating padding overhead and enabling efficient training on datasets with diverse sequence lengths.

5. Practical Impact

| Metric | Standard Attention | FlashAttention-2 (A100) | FlashAttention-3 (H100 FP16) | FlashAttention-3 (H100 FP8) | |--------|-------------------|------------------------|------------------------------|----------------------------| | Peak TFLOPs/s | ~30 | 230 | 740 | ~1200 | | MFU | ~15% | 73% | 75% | ~75% | | Context length feasible | 4K–8K | 128K+ | 128K–1M+ | 1M+ | | Memory (N=128K, d=128) | OOM | ~20 GB | ~20 GB | ~10 GB |

FlashAttention is now the default attention kernel in PyTorch (torch.nn.functional.scaled_dot_product_attention with torch.backends.cuda.enable_flash_sdp), xFormers, Hugging Face Transformers, vLLM, and TensorRT-LLM. It directly enabled the context-length explosion from GPT-3's 2K to Llama-3's 128K and beyond.

6. When Not to Use FlashAttention

• Very short sequences ($N < 512$): Launch overhead and SRAM pressure can make standard attention faster.
• Non-GPU accelerators without shared memory / tensor core analogues: The algorithm assumes a fast on-chip scratchpad and warp-level MMA.
• Custom attention variants (e.g., rotary embeddings computed outside the kernel, attention sinks, ALiBi) may require kernel modifications or fall back to the standard path.

Sources

• Dao, Fu, Ermon, Rudra, Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. https://arxiv.org/abs/2205.14135
• Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICLR 2024. https://arxiv.org/abs/2307.08691
• Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv 2024. https://arxiv.org/abs/2407.08608
• PyTorch Blog. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. https://pytorch.org/blog/flashattention-3
• Dao AILab. FlashAttention GitHub Repository. https://github.com/dao-ailab/flash-attention