FlashAttention: Mathematical Foundations, IO-Aware Tiling, Online Softmax, and Recomputation Dynamics

Standard self-attention in the Transformer architecture scales quadratically with sequence length in both time and memory complexity. While algorithmic research historically focused on reducing FLOP counts via sparse or low-rank approximations, Tri Dao et al. (2022) demonstrated that the practical wall-clock bottleneck in multi-head attention is not compute capability, but memory IO.

FlashAttention reformulates exact scaled dot-product attention as an IO-aware algorithm. By leveraging the GPU memory hierarchy, fusing operations into a single kernel, and computing softmax incrementally via online statistics, FlashAttention eliminates the need to materialize the intermediate N x N attention matrix in GPU global memory.

The GPU Memory Hierarchy and the Attention Memory Wall

Modern GPU architectures, such as the NVIDIA A100 and H100, feature a hierarchical memory subsystem consisting of:

High Bandwidth Memory (HBM): Main GPU memory (40 GB to 80 GB on A100, up to 96 GB on H100) with bandwidth ranging from 1.5 TB/s to 3.35 TB/s.
On-Chip Static RAM (SRAM / Shared Memory): Small, low-latency memory local to each Streaming Multiprocessor (SM) (192 KB per SM on A100, 228 KB on H100), delivering over 19 TB/s of aggregate bandwidth.

In standard implementations of scaled dot-product attention according to Vaswani et al. (2017):

O = softmax(Q * K^T / sqrt(d)) * V

where Q, K, V are matrices of shape (N, d), sequence length is N, and head dimension is d.

The computation proceeds sequentially through discrete GPU kernels:

Compute Scores: Compute S = Q * K^T (shape N x N) and write S from SRAM to HBM (O(N^2) write traffic).
Compute Probabilities: Read S from HBM into SRAM, apply row-wise softmax P = softmax(S), and write P (shape N x N) back to HBM (O(N^2) read and write traffic).
Compute Output: Read P and V from HBM into SRAM, compute O = P * V (shape N x d), and write O back to HBM.

For modern context windows (N >= 4096), N^2 dwarfs N * d. Because the arithmetic intensity (FLOPs per byte of memory transfer) of pointwise softmax and memory-staging operations is low, the GPU compute engines (Tensor Cores) spend the majority of execution cycles stalled, waiting for memory bus transactions between HBM and SRAM.

Online Softmax: The Mathematical Core

The primary obstacle to fusing the matrix multiplication Q * K^T, the softmax normalization, and the output reduction P * V into a single kernel is the global reduction required by softmax.

Standard safe softmax avoids numerical overflow by subtracting the row maximum:

m(x) = max_j(x_j)
l(x) = sum_j(exp(x_j - m(x)))
softmax(x)_i = exp(x_i - m(x)) / l(x)

Traditionally, this requires three passes over the entire sequence:

Find row maximum m(x).
Compute unnormalized exponentials and sum them to obtain normalizer l(x).
Divide each exponential by l(x).

To compute attention without materializing all scores at once, FlashAttention adapts the online normalizer technique introduced by Milakov and Gimelshein (2018) and Rabe and Staats (2021).

Incremental Online Updates

Suppose an input vector x of length N is split into two contiguous blocks: x = [x^(1), x^(2)].

Let the local statistics for block 1 be:

m^(1) = max(x^(1))
l^(1) = sum_j(exp(x_j^(1) - m^(1)))

When processing block 2 with local statistics m^(2) = max(x^(2)) and l^(2) = sum_j(exp(x_j^(2) - m^(2))), the combined maximum and normalizer update incrementally:

m_new = max(m^(1), m^(2))
l_new = l^(1) * exp(m^(1) - m_new) + l^(2) * exp(m^(2) - m_new)

The running output vector O^(1) = sum_j (exp(x_j^(1) - m^(1)) / l^(1)) * V_j^(1) updates to incorporate block 2 without re-reading block 1:

O_new = (l^(1) * exp(m^(1) - m_new) / l_new) * O^(1) + (exp(-m_new) / l_new) * sum_j(exp(x_j^(2)) * V_j^(2))

By maintaining running scalars m and l alongside the partial output vector O, the exact softmax-weighted sum can be computed in a single sweep over blocks of K and V.

IO-Aware Tiling Algorithm

FlashAttention organizes the input matrices into blocks sized to fit directly inside the GPU on-chip SRAM capacity M:

Partitioning:

Q is partitioned into T_r = ceil(N / B_r) blocks Q_1, ..., Q_(T_r) of dimension B_r x d.
K is partitioned into T_c = ceil(N / B_c) blocks K_1, ..., K_(T_c) of dimension B_c x d.
V is partitioned into T_c = ceil(N / B_c) blocks V_1, ..., V_(T_c) of dimension B_c x d.
Block sizes B_r, B_c are chosen such that B_r * d + 2 * B_c * d <= M.

Execution Flow (FlashAttention-1):

Initialize output O = 0 (shape N x d), normalizers l = 0 (length N), and maximums m = -inf (length N) in HBM.
Outer loop: For each column block j in {1, ..., T_c}:
Load K_j, V_j from HBM into on-chip SRAM.
Inner loop: For each row block i in {1, ..., T_r}:
Load Q_i, O_i, l_i, m_i from HBM into SRAM.
Compute local scores S_ij = Q_i * K_j^T / sqrt(d) (shape B_r x B_c) on SRAM.
Compute local block row max m_tilde_ij and local normalizer l_tilde_ij.
Update global row statistics:
m_i_new = max(m_i, m_tilde_ij)
l_i_new = exp(m_i - m_i_new) * l_i + exp(m_tilde_ij - m_i_new) * l_tilde_ij
Update running output:
O_i = diag(exp(m_i - m_i_new))^-1 * O_i + exp(S_ij - m_i_new) * V_j
Write updated O_i, l_i, m_i back to HBM.

At completion, each row block O_i is normalized by diag(l_i)^-1 * O_i.

IO Complexity Reduction

According to the analysis by Dao et al. (2022):

Standard Attention IO: Theta(N * d + N^2) HBM memory accesses.
FlashAttention IO: Theta(N^2 * d^2 / M) HBM memory accesses, where M is the SRAM capacity.

For standard configurations (head dimension d = 64 to 128, SRAM M ~ 100 KB), FlashAttention reduces memory access operations by 4x to 20x, matching theoretical IO lower bounds for matrix computation.

Backward Pass and Activation Recomputation

In standard Transformer training, backpropagation through attention requires storing the full attention probability matrix P (shape N x N) in HBM during the forward pass. This creates an O(N^2) memory footprint per attention layer per head, causing out-of-memory errors at extended sequence lengths.

FlashAttention introduces a selective recomputation strategy:

Forward Storage: The forward pass saves zero intermediate N x N matrices to HBM. It stores only the final output O (shape N x d) and the summary statistics m, l (length N).
Backward Recomputation: During the backward pass, when computing gradients dQ, dK, dV, the attention matrix blocks S_ij and P_ij are recomputed dynamically in SRAM from Q_i, K_j and the stored statistics m_i, l_i.

The Compute-Memory Trade-off

Recomputing S_ij during the backward pass adds approximately 16% to 20% more theoretical arithmetic FLOPs to the backward pass. However, because these FLOPs execute entirely within fast on-chip SRAM without issuing HBM read or write instructions, the backward pass runs up to 3x faster than standard attention while reducing peak activation memory from O(N^2) to O(N).

Architectural Evolution: FlashAttention-2 and FlashAttention-3

The FlashAttention architecture has evolved across GPU hardware generations:

FlashAttention-2 (Dao, 2023)

Introduced by Tri Dao (2023), FlashAttention-2 resolved multiple scheduling bottlenecks:

Inverted Loop Nest: Swapped the loops to place Q in the outer loop and K, V in the inner loop. This eliminates repeated global HBM writes and reads for the output accumulator O_i, keeping it resident in registers throughout the inner loop.
Sequence-Length Parallelism: In addition to parallelizing across batch and head dimensions, FlashAttention-2 parallelizes across the sequence length dimension (row blocks of Q), achieving high SM utilization even with small batch sizes or single-head execution.
Warp-Level Matrix Partitioning: Split work across warps within each thread block to reduce shared memory bank conflicts and synchronization barriers.

FlashAttention-3 (Shah et al., 2024)

Designed for NVIDIA Hopper (H100) and Blackwell architectures by Shah et al. (2024):

Asynchronous Execution via TMA: Uses Hopper Tensor Memory Accelerator (TMA) hardware units to transfer tensor tiles directly between global memory and shared memory asynchronously, bypassing register files.
Warp Specialization and Ping-Pong Scheduling: Separates warps into dedicated producer (data loading and softmax) and consumer (Tensor Core matrix multiplication via WGMMA) roles, overlapping memory transfers with compute.
Low-Precision FP8 Support: Implements block-quantized FP8 matrix multiplication with dynamic scaling to maintain numerical accuracy while doubling throughput.

Implementation Comparison Across Generations

Standard PyTorch Attention: HBM access complexity O(N * d + N^2), activation memory O(N^2), multi-kernel execution, general GPU support.
FlashAttention (v1): HBM access complexity O(N^2 * d^2 / M), activation memory O(N), outer loop over K/V blocks, Ampere (A100) target, exact arithmetic.
FlashAttention-2: HBM access complexity O(N^2 * d^2 / M), activation memory O(N), outer loop over Q blocks with sequence parallelism, Ampere/Ada target, exact arithmetic.
FlashAttention-3: HBM access complexity O(N^2 * d^2 / M), activation memory O(N), warp-specialized TMA scheduling, Hopper (H100) target, exact FP16/BF16 and quantized FP8.

FlashAttention demonstrates that as hardware FLOP capabilities outpace memory bandwidth scaling, kernel design must optimize for data movement across the memory hierarchy rather than raw operation counts.