Attention is mathematically simple but expensive to implement efficiently. Given queries $Q$, keys $K$, and values $V$, standard scaled dot-product attention computes

where $S$ contains the attention scores, $P$ contains the normalized probabilities, and $O$ is the output. For a sequence of length $N$, the score and probability matrices have $N^2$ entries. Even when the final result fits in memory, repeatedly writing and reading these matrices from GPU HBM can dominate the runtime.

The IO problem

GPU computation uses several memory levels. Tensor Cores are fast, shared memory is smaller and faster than HBM, and HBM is large but comparatively slow. A naive attention implementation writes the score matrix to HBM, reads it to apply softmax, writes the probability matrix, and reads it again to multiply by $V$.

FlashAttention avoids materializing the $N\times N$ attention matrix. It divides $Q$, $K$, and $V$ into tiles, loads a small set of tiles into on-chip memory, performs the matrix multiplications and softmax updates there, and writes only the output and a small amount of per-row metadata. The algorithm is exact. It changes the order and location of the computation, not the attention function being computed.

The asymptotic amount of attention arithmetic is still $O(N^2d)$. The important improvement is the amount of data moved between HBM and on-chip memory. This is why FlashAttention is an IO-aware algorithm rather than a linear-attention approximation.

FlashAttention-1

FlashAttention introduced the tiled IO-aware algorithm and the online softmax reduction. It proved that exact attention can use linear additional memory in sequence length by avoiding the quadratic intermediate matrices. The main idea is algorithmic: reorganize the computation around the GPU memory hierarchy.

FlashAttention-2

FlashAttention-2 kept the same exact tiled computation but improved how the work is assigned to GPU thread blocks and warps. It reduced non-matrix-multiply work, parallelized attention across the sequence even for a single head, and reduced unnecessary communication through shared memory.

FlashAttention-3

FlashAttention-3 targets Hopper GPUs. It uses warp specialization and asynchronous operations to overlap Tensor Core matrix multiplication, data movement through Tensor Memory Accelerator, and softmax work. It also introduces low-precision techniques for FP8 attention while controlling the resulting numerical error.

FlashAttention-4

FlashAttention-4 targets Blackwell GPUs, where Tensor Core throughput grows faster than some other parts of the chip. Matrix multiplication is therefore less likely to be the only bottleneck. FA4 redesigns the pipeline around asynchronous MMA operations, larger tiles, software-emulated exponential operations, and reduced shared-memory traffic and atomic updates in backward.

References

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
• FlashAttention-4: Algorithm and Kernel Pipelining Co-design for Asymmetric Hardware Scaling
• FlashAttention lecture notes
• Tiling