FlashAttention 3: How Exact Attention Became 2x Faster Than FlashAttention 2

Why Attention Is Slow: The Memory Wall

Standard attention materializes the full N×N attention matrix in High Bandwidth Memory (HBM). For a sequence of 8192 tokens with 128-dim heads, that is 8192² = 67M entries — several GB just for intermediate results. The bottleneck is not arithmetic (FLOPs) but memory bandwidth: reading and writing to HBM is orders of magnitude slower than on-chip SRAM computation.

FlashAttention 1: IO-Aware Tiling

The original FlashAttention (arXiv:2205.14135) introduced IO-aware attention: tile the Q, K, V matrices into blocks that fit in SRAM, compute attention incrementally using the online softmax trick, and recompute activations during the backward pass instead of storing them. The full N×N matrix is never materialized. This reduced HBM reads/writes from O(N²) to O(N) and enabled 2-4x speedup with exact (not approximate) attention.

FlashAttention 2: Better Parallelism

FlashAttention 2 (arXiv:2307.08691) reduced the number of non-matmul FLOPs (which have lower throughput than matmul FLOPs on modern GPUs), parallelized across the sequence length dimension (not just batch and head), and improved thread block partitioning for better warp occupancy. Result: 2-4x speedup over FA1, reaching 50-73% of theoretical peak FLOPs/s on A100.

FlashAttention 3: H100-Specific Optimizations

FA3 (arXiv:2407.08608) targets the NVIDIA H100 Tensor Core architecture specifically. Three key advances:

1. Asynchrony and Overlapping

H100 has a dedicated asynchronous copy engine (TMA — Tensor Memory Accelerator) that can transfer data between HBM and SRAM independently of CUDA cores. FA3 overlaps compute (matmul) with data loading using software pipelining, so the GPU never stalls waiting for memory transfers.

2. FP8 Support

H100 has dedicated FP8 Tensor Cores with 2x the throughput of BF16 cores. FA3 supports FP8 attention computation with careful handling of the softmax normalization to maintain numerical stability. This enables up to 1.2 PFLOPs/s on H100 for attention.

3. Intra-Warp Parallelism

FA3 restructures the warp-level computation to better utilize H100's warp-group asynchronous execution, pipelining the two GEMMs (QK^T and attention × V) together.

# Using FlashAttention 3 via the flash-attn package
from flash_attn import flash_attn_func
import torch

# All on H100, with FP8 quantization for maximum throughput
q = torch.randn(2, 8192, 32, 128, dtype=torch.bfloat16, device="cuda")
k = torch.randn(2, 8192, 32, 128, dtype=torch.bfloat16, device="cuda")
v = torch.randn(2, 8192, 32, 128, dtype=torch.bfloat16, device="cuda")

# causal=True for autoregressive decoding
output = flash_attn_func(q, k, v, causal=True)

Impact on Context Length

FA3 makes 128k-context models practical for production inference. Each doubling of context length was previously 4x more expensive (O(N²) memory). With FA3's IO efficiency, context can be extended to 1M+ tokens in research settings without running out of memory or being bandwidth-bound.

Speedup Numbers

On H100 SXM5, FA3 achieves 1.5-2.0x speedup over FA2 in the forward pass and similar gains in the backward pass. At BF16 with sequence length 8192, FA3 reaches ~1200 TFLOPs/s versus ~700 for FA2.

Further Reading

• FlashAttention 3 arXiv:2407.08608
• FlashAttention 2 arXiv:2307.08691
• flash-attention GitHub