FlashAttention: How IO-Aware Tiling and Online Softmax Solved Transformer Memory Bottlenecks

Standard multi-head attention is the fundamental computational primitive of modern autoregressive language models. While mathematically straightforward, the operation introduces a severe operational bottleneck as context windows scale. Naive implementations of scaled dot-product attention exhibit quadratic memory complexity $O(N^2)$ and quadratic memory access costs, bounding sequence lengths and leaving modern GPU tensor cores severely underutilized.

FlashAttention, introduced by Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré in FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022), resolved this bottleneck without approximations. By restructuring the attention algorithm around the physical hierarchy of GPU memory, FlashAttention reduces memory reads and writes between high-bandwidth memory (HBM) and on-chip SRAM, cutting memory footprints from quadratic $O(N^2)$ to linear $O(N)$ while delivering two to four times faster execution.

FlashAttention GPU memory hierarchy and tiling architecture

The GPU Memory Wall in Standard Attention

To understand why standard attention stalls modern hardware, one must examine the physical memory hierarchy of modern accelerator chips like the NVIDIA A100 and H100.

The Accelerator Memory Hierarchy

A modern datacenter GPU consists of two primary storage tiers:

High-Bandwidth Memory (HBM): Off-chip dynamic RAM (such as 80 GB HBM2e on A100 or 80 GB HBM3 on H100). HBM offers large capacity but bounded bandwidth (typically 1.5 to 3.35 TB/s).
On-Chip SRAM (Shared Memory / L1 Cache): Fast, localized static memory residing directly inside each Streaming Multiprocessor (SM). While offering massive bandwidth (roughly 19 TB/s aggregate on A100), SRAM capacity is strictly limited (typically 192 KB to 228 KB per SM).

Compute engines like Tensor Cores execute matrix operations at hundreds of teraFLOPs, but they can only operate directly on data loaded into registers and local SRAM.

Naive Attention Execution Profile

Given input matrices $Q, K, V \in \mathbb{R}^{N \times d}$ (where $N$ is sequence length and $d$ is head dimension), standard attention computes:

$S = Q K^T \in \mathbb{R}^{N \times N}$ $P = \text{softmax}(S) \in \mathbb{R}^{N \times N}$ $O = P V \in \mathbb{R}^{N \times d}$

In a naive implementation (such as standard PyTorch execution):

$Q$ and $K$ are loaded from HBM into SRAM to compute $S = Q K^T$ . The resulting intermediate score matrix $S$ ( $N \times N$ ) is written back to HBM.
$S$ is read back from HBM to SRAM to compute row-wise softmax: $P = \text{softmax}(S)$ . The full probability matrix $P$ ( $N \times N$ ) is written back to HBM.
$P$ and $V$ are read from HBM to SRAM to compute the output projection $O = P V$ . The final tensor $O$ ( $N \times d$ ) is written to HBM.

For a sequence length of $N = 16,384$ tokens and 32 attention heads, the intermediate matrix $P$ alone requires several gigabytes of memory per layer. Because the kernel repeatedly reads and writes $N \times N$ matrices to slow HBM, execution time is dominated by memory bandwidth bottlenecks rather than floating-point math.

Standard Attention (Memory Bandwidth Bound):
[HBM] --(Load Q, K)--> [SRAM: Q*K^T] --(Write S: N x N)--> [HBM]
[HBM] --(Load S)-----> [SRAM: Softmax] -(Write P: N x N)--> [HBM]
[HBM] --(Load P, V)--> [SRAM: P*V] ----(Write O: N x d)--> [HBM]

FlashAttention (Compute Bound via SRAM Tiling):
[HBM] --(Load Blocks of Q, K, V)--> [SRAM: Tiled Dot-Product + Online Softmax] --(Write Block O)--> [HBM]

Core Algorithmic Mechanics of FlashAttention

FlashAttention achieves exact mathematical equivalence to standard attention while eliminating intermediate $N \times N$ reads and writes through three core techniques: tiling, online softmax, and backward recomputation.

1. IO-Aware Tiling

Instead of calculating the full $N \times N$ score matrix at once, FlashAttention splits the input matrices $Q, K, V$ into smaller blocks of size $B_r \times d$ and $B_c \times d$ , chosen specifically to fit within on-chip SRAM.

The kernel loads a block of $Q$ and iterates over blocks of $K$ and $V$ , performing local matrix multiplications entirely inside SRAM:

Load query block $Q_i \in \mathbb{R}^{B_r \times d}$ into SRAM.
Load key block $K_j \in \mathbb{R}^{B_c \times d}$ and value block $V_j \in \mathbb{R}^{B_c \times d}$ into SRAM.
Compute local attention scores $S_{ij} = Q_i K_j^T \in \mathbb{R}^{B_r \times B_c}$ .
Update the running output accumulator directly in SRAM.

Only the final accumulated result $O_i \in \mathbb{R}^{B_r \times d}$ is written out to HBM once the inner loop completes.

2. Online Softmax Reduction

The challenge with block-by-block processing is that standard softmax normalization requires global information across the entire sequence row:

$\text{softmax}(x)_i = \frac{e^{x_i - m}}{\sum_{j=1}^N e^{x_j - m}}, \quad \text{where } m = \max_{j} x_j$

To evaluate this incrementally without buffering all blocks, FlashAttention builds on the online softmax algorithm formulated by Milakov and Gimelshein in Online Normalizer Calculation for Softmax (2018) and analyzed by Rabe and Staats in Self-attention Does Not Need $O(n^2)$ Memory (2021).

When moving from block $j-1$ to block $j$ , the algorithm tracks running row maximums and normalization sums:

Let $m^{(j-1)}$ be the row maximum computed through block $j-1$ .
Let $\ell^{(j-1)}$ be the unnormalized sum of exponentials computed through block $j-1$ .
Compute the new local block maximum $\tilde{m} = \max(S_{ij})$ .
Update the global maximum: $m^{(j)} = \max(m^{(j-1)}, \tilde{m})$ .
Scale and update the running normalizer:

$\ell^{(j)} = e^{m^{(j-1)} - m^{(j)}} \ell^{(j-1)} + \sum e^{S_{ij} - m^{(j)}}$

Rescale the running output accumulator $O_i^{(j-1)}$ to match the new global maximum before adding the new block contribution:

$O_i^{(j)} = \text{diag}\left(e^{m^{(j-1)} - m^{(j)}}\right) O_i^{(j-1)} + e^{S_{ij} - m^{(j)}} V_j$

Once all key-value blocks have been processed, the final block output is normalized:

$O_i = \text{diag}\left(\ell^{(\text{final})}\right)^{-1} O_i^{(\text{final})}$

This formulation guarantees exact numerical results while ensuring no $N \times N$ matrix is ever written to or read from HBM.

Online Softmax Rescaling Step:
Previous Partial Output: O_prev (based on max m_prev)
New Block Scores: S_curr (local max m_curr)
New Global Max: m_new = max(m_prev, m_curr)

Scale Factor Alpha = exp(m_prev - m_new)
Updated Output: O_new = Alpha * O_prev + exp(S_curr - m_new) * V_curr

3. Recomputation in the Backward Pass

During neural network backpropagation, computing gradients for $Q, K, V$ requires access to the attention weight matrix $P$ . In standard backpropagation, $P$ is saved during the forward pass, consuming $O(N^2)$ HBM storage.

FlashAttention discards $P$ entirely during the forward pass. Instead, it saves only:

The output tensor $O \in \mathbb{R}^{N \times d}$
The row statistics: max vector $m \in \mathbb{R}^N$ and normalizer vector $\ell \in \mathbb{R}^N$

In the backward pass, FlashAttention reloads blocks of $Q, K, V$ from HBM, recomputes the tile scores $S_{ij}$ in fast SRAM using the saved $(m, \ell)$ vectors, and immediately evaluates the gradients $\partial Q, \partial K, \partial V$ .

Although recomputation incurs additional floating-point operations, it eliminates $O(N^2)$ memory reads from HBM. Because GPUs are memory-bandwidth constrained, avoiding HBM IO yields an overall speedup during training while reducing activation memory from quadratic $O(N^2)$ to linear $O(N)$ .

Architecture Iterations: FlashAttention-2 and FlashAttention-3

The FlashAttention kernel architecture has evolved across successive GPU hardware generations to extract higher fractions of theoretical peak hardware throughput.

FlashAttention-2: Parallelism Across the Sequence Dimension

In 2023, Tri Dao released FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. The second iteration introduced three major microarchitectural enhancements:

Outer Loop Query Parallelization: FlashAttention-1 parallelized over batch size and attention heads, scheduling key-value blocks in the outer loop. When batch sizes were small or sequence lengths large, this led to low occupancy across GPU Streaming Multiprocessors. FlashAttention-2 moved the query blocks to the outer loop and parallelized across the sequence length dimension, maximizing thread block occupancy.
Reduced Non-Matmul FLOPs: Softmax scaling and normalizer division were restructured to minimize scalar arithmetic operations on GPU CUDA cores, keeping execution concentrated on high-throughput Tensor Cores.
Warp Work Partitioning: FlashAttention-2 optimized how individual warps (groups of 32 threads) cooperate within an SM thread block, splitting matrix blocks across warps to eliminate shared-memory bank conflicts and synchronization barriers.

These optimizations raised execution efficiency from roughly 30-50% of theoretical peak FLOPs in FlashAttention-1 to 50-73% on NVIDIA A100 GPUs, achieving a 2x speedup over the original kernel.

FlashAttention-3: Asynchronous Execution on Hopper Architectures

In 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao introduced FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision, specifically tailored for NVIDIA Hopper (H100) microarchitecture:

Tensor Memory Accelerator (TMA) Integration: FlashAttention-3 utilizes Hopper's dedicated hardware TMA engine to execute asynchronous multi-dimensional tensor copies directly between HBM and SRAM without consuming register file bandwidth or ALU cycles.
Warp-Group Matrix Multiply-Accumulate (WGMMA): Leverages Hopper's specialized instructions executing across 128-thread warp groups, allowing matrix multiplications to proceed directly from shared memory into registers asynchronously.
Producer-Consumer Warp Specialization: Warps within a thread block are specialized into distinct roles: producer warps issue TMA memory transfers, while consumer warps execute math pipelines. This completely overlaps data movement latency with tensor computation.
Low-Precision FP8 Support: FlashAttention-3 incorporates 8-bit floating point (FP8) precision with block quantization and incoherent processing, doubling computational throughput while controlling numerical error accumulation.

On NVIDIA H100 SXM5 systems, FlashAttention-3 achieves up to 1.2 PFLOPs/s in FP16 (over 75% SM utilization) and reaches near 2.0 PFLOPs/s in FP8, outperforming FlashAttention-2 by 1.5 to 2.0 times.

System-Level Impact on Production Serving and Training

The transition to IO-aware attention fundamentally reshaped how frontier models are trained and deployed:

Context Window Expansion: Prior to FlashAttention, training transformer models beyond 2,048 or 4,096 tokens incurred unsustainable memory overheads. By linearizing memory consumption with sequence length, FlashAttention provided the underlying execution engine required to scale contexts to 32K, 128K, and 1M+ tokens.
Standardized Framework Integration: The core concepts of IO-aware tiling and online softmax are now ubiquitous. They are natively integrated into PyTorch as torch.nn.functional.scaled_dot_product_attention, Hugging Face Transformers, DeepSpeed, Megatron-LM, and Triton.
Inference Optimization: High-throughput LLM serving systems, including vLLM, TensorRT-LLM, and SGLang, combine FlashAttention kernels with PagedAttention and prompt caching to maximize continuous batching concurrency during multi-tenant token generation.