FlashAttention

FlashAttention is a family of hardware-aware algorithms and software implementations for computing the attention mechanism used in the Transformer architecture. Its central innovation is to reorganize attention around the memory hierarchy of a graphics processing unit (GPU), rather than merely reduce arithmetic operations. Introduced in May 2022 by Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré, the original algorithm computes exact dense attention while avoiding storage of the complete attention matrix in external GPU memory. Subsequent versions improve work partitioning and exploit newer hardware capabilities. (arxiv.org)

Attention and the memory bottleneck

For a single attention head, scaled dot-product attention takes query, key, and value matrices $Q$ , $K$ , and $V$ , and computes

O=\operatorname{softmax}\left(\frac{QK^{\mathsf T}}{\sqrt{d}}\right)V,

where $d$ is the query/key dimension and the softmax function operates across each score row. Masks can restrict permitted query–key interactions. In self-attention over a sequence of length $N$ , the score matrix has $N^2$ entries. Conventional implementations commonly write this matrix, and then normalized attention weights, to device memory between separate operations. (arxiv.org)

FlashAttention addresses the resulting data movement. GPU high-bandwidth memory has much greater capacity than on-chip memory, but accessing it is more expensive. By dividing attention into tiles that fit on chip and combining operations, FlashAttention avoids repeatedly transferring large intermediate matrices. Its performance advantage therefore comes from reduced memory traffic, not from discarding query–key comparisons. (arxiv.org)

Tiling and online normalization

The apparent obstacle to tiled attention is that softmax requires a normalization factor spanning an entire score row. Computing softmax independently within each tile would change the result. FlashAttention instead uses online normalization, maintaining a running maximum, an exponential sum, and an accumulated weighted output for each query row. (tridao.me)

For a newly encountered score block $s_b$ , these quantities can be updated as

m'=\max(m,\max s_b),

\ell'=e^{m-m'}\ell+\sum_j e^{s_{b,j}-m'},

u'=e^{m-m'}u+\sum_j e^{s_{b,j}-m'}v_{b,j}.

The final output is $u/\ell$ . Rescaling the previous accumulator incorporates the new block without retaining earlier scores. Subtracting the running maximum provides numerical stability for the exponential function. This formulation preserves global normalization while allowing blockwise computation. (tridao.me)

“Exact” distinguishes the dense algorithm from methods that approximate attention through sparsity or reduced-rank representations. It does not imply bitwise identity with another implementation: reordered floating-point arithmetic can produce small numerical differences. The official implementation tests outputs and gradients against reference implementations within numerical tolerances. (arxiv.org)

Training and computational complexity

During backpropagation, attention derivatives ordinarily require intermediate probabilities. FlashAttention retains compact normalization information and recomputes attention blocks as needed, rather than storing the full probability matrix. This is a form of activation checkpointing applied within the attention operation. Although recomputation adds arithmetic, reduced memory access can also make the backward pass faster. (arxiv.org)

For fixed head dimension, attention’s auxiliary space complexity becomes linear rather than quadratic in sequence length. Dense attention’s arithmetic time complexity nevertheless remains quadratic: the permitted query–key pairs still have to be evaluated. FlashAttention therefore separates storage efficiency from arithmetic scaling, an important distinction in computational complexity. (arxiv.org)

Development across generations

FlashAttention-2, introduced in July 2023, improves parallel computing and work partitioning. It reduces non-matrix-multiplication operations, distributes a single attention head across multiple thread blocks, and reorganizes work among GPU warps to reduce communication through shared memory. The paper reported approximately twice the speed of the original implementation and 50–73% of theoretical peak arithmetic throughput on A100 GPUs under its benchmarks. (arxiv.org)

FlashAttention-3, introduced in July 2024, targets NVIDIA Hopper GPUs. It overlaps data transfers with computation, interleaves matrix multiplication with softmax, and uses hardware support for low-precision FP8 computation. These techniques address hardware utilization beyond the original memory-traffic problem. Its FP8 path also introduces block quantization and processing intended to reduce quantization error; such low-precision execution should be distinguished from exact dense attention performed at higher precision. (arxiv.org)

FlashAttention-4, described in a March 2026 paper, addresses increasingly asymmetric hardware scaling: matrix-multiplication throughput has grown faster than some other GPU resources. It combines algorithm changes with kernel pipelining and is implemented in the Python-embedded CuTe-DSL. The paper reports up to 1,613 TFLOPs/s for BF16 attention on B200 GPUs in its benchmark configurations. These are workload-specific kernel measurements, not universal application speedups. (arxiv.org)

Software integration and limitations

The official open-source software repository provides interfaces for multi-head attention, causal and sliding-window masks, variable-length sequences, and inference with cached keys and values. Support varies across implementation paths, hardware generations, and numerical formats. These interfaces operate on tensors and can integrate attention into training or neural-network inference workflows. (github.com)

PyTorch integrated FlashAttention-2 into its scaled-dot-product-attention operation in version 2.2. Framework dispatch can select optimized attention kernels when their constraints are satisfied, or use alternative implementations otherwise. Actual gains depend on sequence length, head dimension, masking, precision, and hardware; attention-kernel acceleration does not translate directly into an equivalent speedup for an entire large language model. (pytorch.org)

Attention and the memory bottleneck

Tiling and online normalization

Training and computational complexity

Development across generations

Software integration and limitations

References