FlashAttention: IO-Aware Tiling, Online Softmax Mathematics, Memory Hierarchy Dynamics, and Kernel Evolution

Standard multi-head attention is the computational core of transformer language models. While mathematically straightforward, standard attention implementations suffer from severe memory and throughput bottlenecks as sequence lengths grow. The bottleneck is not raw arithmetic compute, but memory traffic: repeatedly reading and writing intermediate attention matrices across GPU memory tiers.

FlashAttention, introduced by Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré in 2022, restructured attention computation around hardware memory hierarchy. By fusing the entire attention operation into a single GPU kernel, computing softmax incrementally via online normalizers, and recomputing intermediate matrices during the backward pass, FlashAttention reduced memory footprint from quadratic $O(N^2)$ to linear $O(N)$ while delivering 2x to 4x wall-clock speedups without approximating attention weights.

Subsequent iterations, FlashAttention-2 and FlashAttention-3, refined work partitioning, warp-level specialization, and asynchronous hardware primitives, scaling attention throughput to over 75% of theoretical peak compute on modern hardware.

The Hardware Memory Hierarchy and the IO Bottleneck

To understand why attention required re-engineering, consider the memory hierarchy of modern accelerator hardware like NVIDIA A100 and H100 GPUs.

GPU Memory Tiers: SRAM vs. High-Bandwidth Memory (HBM)

A modern GPU contains two primary memory tiers:

High-Bandwidth Memory (HBM / VRAM): Large capacity (80 GB to 192 GB), but bounded bandwidth (1.5 TB/s to 3.35 TB/s on A100 and H100 SXM).
On-Chip SRAM (Shared Memory and Register File): Small capacity (164 KB to 228 KB shared memory per Streaming Multiprocessor, or roughly 20 MB to 50 MB total across the chip), but massive aggregate bandwidth (over 19 TB/s on A100 and over 33 TB/s on H100) with single-cycle access latency.

Operations on GPUs fall into two regimes:

Compute-Bound: Execution time is dominated by arithmetic operations on Tensor Cores. High arithmetic intensity (FLOPs per byte transferred).
Memory-Bound: Execution time is dominated by reading from or writing to HBM. Low arithmetic intensity, leaving arithmetic units idle while waiting for memory transfers.

Standard Multi-Head Attention Execution Flow

Given sequence length $N$ and head dimension $d$ , multi-head attention projects input tokens into Query ( $Q$ ), Key ( $K$ ), and Value ( $V$ ) matrices in $\mathbb{R}^{N \times d}$ . The core mathematical formulation is:

S = Q * K^T / sqrt(d)       (Score matrix in R^{N x N})
P = softmax(S)              (Attention probabilities in R^{N x N})
O = P * V                   (Output matrix in R^{N x d})

In a standard PyTorch implementation (prior to kernel fusion), each operation is executed as a separate kernel:

Load $Q$ and $K$ from HBM into SRAM, compute $S = Q K^T / \sqrt{d}$ , and write $S \in \mathbb{R}^{N \times N}$ back to HBM.
Read $S$ from HBM into SRAM, compute $P = \text{softmax}(S)$ , and write $P \in \mathbb{R}^{N \times N}$ back to HBM.
Read $P$ and $V$ from HBM into SRAM, compute $O = P V$ , and write $O \in \mathbb{R}^{N \times d}$ back to HBM.

For a sequence length of $N = 16,384$ and 32 attention heads, materializing $S$ and $P$ in 16-bit precision requires:

Memory per matrix per head = 16,384^2 * 2 bytes = 536.87 MB
Memory per matrix across 32 heads = 536.87 MB * 32 = 17.18 GB

Materializing both $S$ and $P$ across a single forward pass requires transferring tens of gigabytes to and from HBM per transformer layer. The operation consumes $O(N^2)$ memory and is heavily memory-bound, running at a fraction of theoretical GPU compute capacity.

The FlashAttention Formulation: Tiling and Online Softmax

FlashAttention eliminates the materialization of $S$ and $P$ in HBM entirely. The algorithm computes attention in a single fused GPU kernel by splitting $Q, K, V$ into blocks that fit within fast on-chip SRAM.

FlashAttention Online Softmax and Tiling Schematic

The Softmax Tiling Challenge

Matrix multiplication $Q K^T$ and $P V$ can be trivially partitioned into independent blocks. However, row-wise softmax presents a structural dependency:

softmax(x)_i = exp(x_i - m) / sum_{j=1}^N exp(x_j - m)
where m = max_{j=1...N}(x_j)

To compute any element in a row of $\text{softmax}(S)$ , standard implementations must first inspect all $N$ elements of that row to determine the maximum $m$ (for numerical stability against exponential overflow) and calculate the normalization denominator $\sum \exp(x_j - m)$ . Naively, this prevents processing a row in small tiles.

Online Softmax Mathematics

FlashAttention solves this dependency by implementing online softmax normalization, adapting principles established by Milakov and Gimelshein (2018) and Rabe and Staats (2021).

Consider a row vector $x \in \mathbb{R}^N$ partitioned into two contiguous blocks $x = [x^{(1)}, x^{(2)}]$ of size $B$ .

For the first block $x^{(1)}$ :

m^{(1)} = max_j(x^{(1)}_j)
l^{(1)} = sum_j exp(x^{(1)}_j - m^{(1)})

When evaluating the second block $x^{(2)}$ :

m^{(2)} = max_j(x^{(2)}_j)
l^{(2)} = sum_j exp(x^{(2)}_j - m^{(2)})

The combined global maximum $m$ and global normalizer $l$ across both blocks are computed dynamically without re-reading $x^{(1)}$ :

m = max(m^{(1)}, m^{(2)})
l = exp(m^{(1)} - m) * l^{(1)} + exp(m^{(2)} - m) * l^{(2)}

Dynamic Accumulation of Attention Output

When computing the output tile $O = P V$ , each intermediate block output $O^{(1)}$ computed with local normalizer $l^{(1)}$ and local max $m^{(1)}$ is rescaled when transitioning to the global statistics:

O_new = (l^{(1)} * exp(m^{(1)} - m) / l) * O^{(1)} + (exp(m^{(2)} - m) / l) * (P^{(2)} * V^{(2)})

By tracking two scalar statistics per row, the running maximum $m$ and the running normalizer $l$ , FlashAttention updates the accumulated output vector $O$ in SRAM as new blocks of $K$ and $V$ are streamed through. The intermediate attention matrices $S$ and $P$ are never written to HBM.

IO Complexity and Backward Pass Recomputation

The primary theoretical contribution of FlashAttention is establishing IO-awareness for deep learning primitives.

IO Complexity Analysis

Let $N$ be sequence length, $d$ be head dimension, and $M$ be the size of SRAM (where $d \le M \le N d$ ).

Standard Attention IO Complexity: Standard attention transfers $\Theta(N d + N^2)$ words between HBM and SRAM due to reading and writing $S$ and $P$ .
FlashAttention IO Complexity: By selecting block sizes $B_c, B_r = \Theta(M / d)$ , FlashAttention requires $\Theta(N^2 d^2 / M)$ HBM memory transfers.

For standard configurations where $M \gg d$ , FlashAttention reduces memory traffic by several multiples, converting a memory-bound kernel into a compute-dense operation.

Backward Pass Recomputation

During model training, the backward pass requires the attention probabilities $P$ to calculate gradients with respect to $Q, K,$ and $V$ :

dV = P^T * dO
dP = dO * V^T
dS = P * (dP - diag(dP * P^T))
dQ = (dS * K) / sqrt(d)
dK = (dS^T * Q) / sqrt(d)

Standard attention caches the full $N \times N$ matrix $P$ in HBM during the forward pass, consuming $O(N^2)$ memory per layer.

FlashAttention introduces selective recomputation:

Forward pass: Stores only the final output $O \in \mathbb{R}^{N \times d}$ and the softmax normalization statistics $L = m + \log(l) \in \mathbb{R}^N$ in HBM ( $O(N)$ storage).
Backward pass: Loads $Q, K, V$ , and the vector $L$ into SRAM in blocks. The kernel recomputes the tile of $P$ on-the-fly in fast SRAM from $Q, K$ and $L$ , evaluates $dQ, dK, dV$ , and discards $P$ .

Because recomputing $P$ in SRAM is fast and avoids HBM memory reads, the backward pass runs faster despite performing additional arithmetic FLOPs.

The Architectural Evolution: FlashAttention-1, 2, and 3

The FlashAttention kernel architecture has undergone three major design iterations.

FlashAttention-1: Fused IO-Aware Kernel

Published in Dao et al. (2022), the original FlashAttention structured the outer loop over key-value blocks ( $K, V$ ) and the inner loop over query blocks ( $Q$ ).

Limitations:

The outer loop over $K, V$ required thread blocks to update output accumulators concurrently, requiring atomic additions or inter-thread synchronization in shared memory.
Non-matmul operations (exponential calculations, rescaling) occupied significant execution cycles.
Kernel achieved roughly 25% to 40% of theoretical peak FLOPs on NVIDIA A100 GPUs.

FlashAttention-2: Inverted Loops and Warp Partitioning

Published in Dao (2023), FlashAttention-2 introduced three major structural refinements:

Outer Loop Inversion: Swapped loop order so the outer loop iterates over Query blocks ( $Q$ ) and the inner loop iterates over Key-Value blocks ( $K, V$ ). Each thread block processes a fixed row of $Q$ and updates its own local accumulator without atomic operations or cross-block synchronization.
Parallelization Over Sequence Length: In addition to parallelizing across batch size and attention heads, FlashAttention-2 parallelizes across the sequence length dimension (outer loop blocks), saturating GPU compute resources even with small batch sizes.
Rescaling Optimization: Eliminated intermediate scalar division during the inner loop. The accumulator maintains unnormalized dot-products, applying normalizer division only once at the conclusion of the row.
Warp-Level Work Partitioning: Split the $K$ and $V$ dimensions among warps within a thread block using efficient tensor core matrix-multiply-accumulate (MMA) layouts, eliminating shared memory bank conflicts.

Performance: FlashAttention-2 achieved 50% to 73% of theoretical peak FP16/BF16 compute on A100 (reaching up to 225 TFLOPS).

FlashAttention-3: Asynchrony and Low-Precision on Hopper

Published by Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao (2024), FlashAttention-3 re-architected attention for NVIDIA Hopper (H100) hardware.

FlashAttention-3 leverages specialized architectural features:

Tensor Memory Accelerator (TMA): Uses dedicated asynchronous hardware units on Hopper to transfer multi-dimensional tensor tiles directly between HBM and Shared Memory, bypassing general-purpose registers and integer ALU instructions.
Warp-Specialization (Producer-Consumer Warps): Dedicated producer warps issue TMA transfer instructions and manage hardware barriers, while consumer warps execute matrix arithmetic on Tensor Cores (using Hopper WGMMA instructions).
Overlapping Softmax and GEMM: Overlaps non-matmul softmax operations (exponentiation, reduction) with asynchronous matrix multiply instructions across pipeline stages.
FP8 Low-Precision Support: Implements attention in 8-bit floating point (FP8 E4M3/E5M2). To prevent quantization error accumulation during softmax, FlashAttention-3 uses block-level quantization and dynamic scaling factors.

Performance: FlashAttention-3 achieves up to 75% to 85% utilization on H100 SXM5, achieving between 650 and 850 TFLOPS in FP16 and exceeding 1.2 PFLOPS in FP8.

Architectural Comparison of Attention Implementations

The following summary outlines the operational and hardware characteristics across implementations:

Standard Attention (PyTorch baseline):
Memory Complexity: $O(N^2)$
Memory Access (IO): $\Theta(N d + N^2)$ HBM reads/writes
Backward Pass Storage: Full $N \times N$ attention probability matrix $P$
GPU Compute Utilization: Typically 15% to 25% of theoretical peak FLOPs
Hardware Bottleneck: HBM bandwidth bound
FlashAttention-1:
Memory Complexity: $O(N)$
Memory Access (IO): $\Theta(N^2 d^2 / M)$
Backward Pass Storage: Output $O$ and softmax statistics vector $L \in \mathbb{R}^N$
GPU Compute Utilization: 25% to 40% on A100
Hardware Bottleneck: Shared memory synchronization and non-matmul instruction overhead
FlashAttention-2:
Memory Complexity: $O(N)$
Memory Access (IO): $\Theta(N^2 d^2 / M)$ with minimal synchronization
Work Partitioning: Outer loop over queries, warp partitioning across inner dimension
GPU Compute Utilization: 50% to 73% on A100
Hardware Bottleneck: Register pressure and Tensor Core latency
FlashAttention-3:
Memory Complexity: $O(N)$
Hardware Acceleration: TMA asynchronous copies, Hopper WGMMA instructions, warp specialization
Numeric Formats: FP16, BF16, and FP8 with block scaling
GPU Compute Utilization: 75% to 85% on H100
Hardware Bottleneck: Instruction issue bandwidth and pipeline occupancy

Ecosystem Integration and Production Impact

The techniques introduced by FlashAttention have become foundational across the deep learning ecosystem:

PyTorch Core: Integrated natively as the backend for torch.nn.functional.scaled_dot_product_attention (SDPA), which automatically routes to FlashAttention when hardware and tensor layouts match.
Serving Frameworks: Integrated into inference engines including vLLM, TensorRT-LLM, SGLang, and FlashInfer for high-throughput prefill and prompt processing.
Long-Context Scaling: Enabled scaling sequence lengths from 2,048 tokens in original GPT-3 architectures to 32k, 128k, and 1M+ tokens in modern models like Llama 3, Qwen 2.5, and Gemini by eliminating the $O(N^2)$ VRAM wall during training and inference.