FlashAttention: Mathematical Foundations, IO-Aware Tiling, Online Softmax Scaling, and Memory-Hierarchy Optimization in Transformer Architectures

Standard multi-head self-attention represents the core computational bottleneck in scaling modern Transformer models to long context windows. While the theoretical arithmetic operations of self-attention scale quadratically with sequence length $O(N^2)$ , modern GPU execution profiles reveal that standard implementations are constrained not by floating-point arithmetic throughput (FLOPs), but by memory access latency and bandwidth between High Bandwidth Memory (HBM) and on-chip Static Random-Access Memory (SRAM).

FlashAttention, introduced by Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Ré in 2022, re-engineers attention as an exact, IO-aware algorithm. By restructuring the computation to tile input matrices into on-chip SRAM blocks, computing softmax incrementally via the online softmax trick, and recomputing intermediate attention matrices during the backward pass, FlashAttention reduces memory access complexity from $\Theta(N^2)$ to $\Theta(N^2 d^2 / M)$ , where $M$ is the SRAM capacity and $d$ is the head dimension. This eliminates the quadratic memory footprint without approximating attention weights, yielding substantial end-to-end training and inference acceleration.

The Hardware Memory Hierarchy and the Roofline Model

Understanding why standard attention bottlenecks on modern GPUs requires analyzing GPU memory hierarchies through the lens of the Roofline performance model.

Modern accelerator hardware architectures, such as NVIDIA Ampere (A100) and Hopper (H100), partition memory into distinct physical layers with contrasting capacity and bandwidth characteristics:

High-Bandwidth Memory (HBM / Global Memory): Large capacity (40 GB to 80 GB on A100; 80 GB to 144 GB on H100), but constrained bandwidth (1.55 TB/s to 2.0 TB/s on A100 SXM4; 3.35 TB/s on H100 SXM5).
On-Chip SRAM (Shared Memory / L1 Cache): Ultra-fast bandwidth (approximately 19 TB/s aggregate on A100; over 33 TB/s on H100), but constrained physical capacity (192 KB per Streaming Multiprocessor on A100; 228 KB per SM on H100, totaling roughly 20 MB to 50 MB across the entire chip).
Tensor Core Register Files: Immediate single-cycle register storage directly wired into specialized matrix-multiply-accumulate (MMA) processing units.

The Roofline model defines the theoretical maximum performance $P$ (FLOP/s) as a function of arithmetic intensity $I$ (FLOPs per byte of memory transferred):

$P = \min(P_{\text{peak}}, I \times B_{\text{peak}})$

Where $P_{\text{peak}}$ is peak compute throughput and $B_{\text{peak}}$ is peak memory bandwidth.

When arithmetic intensity falls below the machine balance threshold $I_{\text{threshold}} = P_{\text{peak}} / B_{\text{peak}}$ , the operation is strictly memory-bandwidth bound (IO-bound). On an NVIDIA A100 (312 TFLOP/s BF16 Tensor Core compute, 2.0 TB/s HBM bandwidth), $I_{\text{threshold}} \approx 156 \text{ FLOPs/byte}$ . Pointwise and reduction operations (such as dropout, masking, and elementwise softmax) operate at $I \le 1 \text{ FLOPs/byte}$ , spending the majority of their execution cycles stalling on HBM memory transfers.

The Standard Attention IO Bottleneck: Quadratic Memory Traffic

Given input sequence length $N$ , head dimension $d$ , and batch size $B$ (analyzed below for a single attention head), the standard attention mechanism evaluates:

$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}$

Where $Q, K, V \in \mathbb{R}^{N \times d}$ are query, key, and value matrices.

In standard PyTorch or CUDA implementations, each distinct operation is executed as a separate GPU kernel call:

Algorithm 1: Standard Multi-Head Attention Execution Flow
-----------------------------------------------------------
1. Load Q, K from HBM -> Compute S = Q K^T in SRAM -> Write S to HBM (size N x N)
2. Load S from HBM -> Compute P = softmax(S) in SRAM -> Write P to HBM (size N x N)
3. Load P, V from HBM -> Compute O = P V in SRAM -> Write O to HBM (size N x d)

Even though the matrix multiplications $Q K^T$ ( $2 N^2 d$ FLOPs) and $P V$ ( $2 N^2 d$ FLOPs) exhibit high theoretical arithmetic intensity when $N$ and $d$ are large, the intermediate materialization of $S$ and $P$ in HBM creates severe IO overhead:

HBM Reads: $Q, K$ ( $2Nd$ ), $S$ ( $N^2$ ), $P$ ( $N^2$ ), $V$ ( $Nd$ ) elements.
HBM Writes: $S$ ( $N^2$ ), $P$ ( $N^2$ ), $O$ ( $Nd$ ) elements.
Total HBM Access: $\Theta(N d + N^2)$ memory reads and writes.

When $N = 8,192$ and $d = 128$ , the intermediate matrices $S$ and $P$ each require $67.1 \times 10^6$ elements ( $134.2 \text{ MB}$ in 16-bit precision per head). Across 32 attention heads, materializing $S$ and $P$ requires over $8.5 \text{ GB}$ of intermediate HBM allocations per layer. At sequence lengths of $32\text{k}$ or $128\text{k}$ , standard attention causes immediate out-of-memory (OOM) failures or throttles GPU compute utilization to under 15% of peak theoretical capacity.

FlashAttention Tiling and Memory Flow Architecture

The IO-Awareness Theorem and Theoretical Lower Bounds

Dao et al. formalized the memory communication model for attention by analyzing data movement between fast memory (SRAM of capacity $M$ ) and slow memory (HBM).

Theorem (IO Complexity of Standard Attention vs. FlashAttention)

Let $N$ be the sequence length, $d$ be the head dimension, and $M$ be the fast memory size where $d \le M < N d$ .

Standard Attention requires $\Theta(N d + N^2)$ HBM memory accesses.
FlashAttention computes exact attention using $\Theta\left(\frac{N^2 d^2}{M}\right)$ HBM memory accesses.
For any algorithm computing attention within the class of matrix-multiplication based operations where the output $O = \text{softmax}(Q K^T) V$ , any computation must perform $\Omega\left(\frac{N^2 d^2}{M}\right)$ memory accesses between HBM and SRAM.

Because FlashAttention matches the lower bound $\Omega(N^2 d^2 / M)$ , it is asymptotically optimal in IO complexity. By choosing block sizes $B_r, B_c \approx \Theta(M / d)$ , FlashAttention reduces total memory access traffic by a factor of $\frac{M}{d}$ , which in practice translates to a $4\times$ to $20\times$ reduction in HBM bandwidth consumption.

Online Softmax Mathematics and Numerical Stability

The mathematical barrier to tiling the attention computation across blocks is the non-local normalization term in the standard three-pass softmax function.

Given a row vector $x \in \mathbb{R}^N$ , standard numerically stable softmax computes:

$m = \max_{j=1 \dots N} x_j$

$f(x) = \left[\exp(x_1 - m), \exp(x_2 - m), \dots, \exp(x_N - m)\right]$

$l = \sum_{j=1}^N f(x)_j$

$\text{softmax}(x) = \frac{f(x)}{l}$

Computing $m$ and $l$ requires observing all $N$ elements before normalizing any individual component, which traditionally prevents executing softmax in a single pass over sub-blocks.

FlashAttention overcomes this limitation using the online softmax formulation (pioneered by Milakov and Gimelshein and analyzed for attention by Rabe and Staats).

Suppose a vector $x$ is partitioned into two consecutive concatenated blocks $x = [x^{(1)}, x^{(2)}]$ where $x^{(1)} \in \mathbb{R}^{B_1}$ and $x^{(2)} \in \mathbb{R}^{B_2}$ .

Let the local statistics for block 1 be:

$m^{(1)} = \max_{j} x^{(1)}_j, \quad l^{(1)} = \sum_{j} \exp(x^{(1)}_j - m^{(1)})$

When block 2 arrives, local statistics are calculated:

$m^{(2)} = \max_{j} x^{(2)}_j, \quad l^{(2)} = \sum_{j} \exp(x^{(2)}_j - m^{(2)})$

The global maximum $m^{\text{new}}$ and global normalizer $l^{\text{new}}$ combining both blocks are computed analytically without revisiting individual elements of $x^{(1)}$ :

$m^{\text{new}} = \max\left(m^{(1)}, m^{(2)}\right)$

$l^{\text{new}} = l^{(1)} \exp\left(m^{(1)} - m^{\text{new}}\right) + l^{(2)} \exp\left(m^{(2)} - m^{\text{new}}\right)$

Output Vector Rescaling Dynamics

To compute the attention output $O = \text{softmax}(Q K^T) V$ incrementally, let $O^{(1)} \in \mathbb{R}^{B_1 \times d}$ be the intermediate output accumulator computed up to block 1:

$O^{(1)} = \frac{1}{l^{(1)}} \sum_{j=1}^{B_1} \exp\left(x^{(1)}_j - m^{(1)}\right) v^{(1)}_j$

Upon processing block 2, the updated unnormalized accumulated vector is scaled to reflect the new global maximum $m^{\text{new}}$ and summed with the new block contribution:

$O^{\text{new}} = \frac{l^{(1)} \exp\left(m^{(1)} - m^{\text{new}}\right) O^{(1)} + \exp\left(x^{(2)} - m^{\text{new}}\right) V^{(2)}}{l^{\text{new}}}$

This recursive formulation generalizes inductively across arbitrary numbers of blocks $T = \lceil N / B_c \rceil$ , ensuring mathematical equivalence to full-sequence softmax while maintaining numeric stability against floating-point overflow.

Tiling Algorithm: SRAM Block Execution and Forward Pass Mechanics

FlashAttention maps the mathematical recurrence directly to GPU execution blocks.

Given SRAM capacity $M$ , block sizes are configured as:

$B_c = \left\lfloor \frac{M}{4d} \right\rfloor, \quad B_r = \min\left(\left\lfloor \frac{M}{4d} \right\rfloor, d\right)$

Let $T_r = \lceil N / B_r \rceil$ and $T_c = \lceil N / B_c \rceil$ . The matrices $Q, K, V$ are partitioned into row and column blocks:

$Q = [Q_1, Q_2, \dots, Q_{T_r}]$ where $Q_i \in \mathbb{R}^{B_r \times d}$
$K = [K_1, K_2, \dots, K_{T_c}]$ where $K_j \in \mathbb{R}^{B_c \times d}$
$V = [V_1, V_2, \dots, V_{T_c}]$ where $V_j \in \mathbb{R}^{B_c \times d}$

Algorithm 2: FlashAttention Forward Pass
-----------------------------------------------------------
Require: Q, K, V in HBM of dimensions N x d; SRAM memory size M.
1. Initialize O = (0)_{N x d} in HBM, l = (0)_N in HBM, m = (-inf)_N in HBM.
2. For j = 1 to T_c do:
3.     Load K_j, V_j from HBM to SRAM
4.     For i = 1 to T_r do:
5.         Load Q_i, O_i, l_i, m_i from HBM to SRAM
6.         Compute S_ij = (Q_i K_j^T) / sqrt(d) in SRAM
7.         Compute m_ij = rowmax(S_ij), p_tilde_ij = exp(S_ij - m_ij), l_ij = rowsum(p_tilde_ij)
8.         Compute m_i_new = max(m_i, m_ij)
9.         Compute l_i_new = exp(m_i - m_i_new) * l_i + exp(m_ij - m_i_new) * l_ij
10.        Compute O_i = diag(l_i * exp(m_i - m_i_new) / l_i_new) * O_i 
                         + diag(exp(m_ij - m_i_new) / l_i_new) * p_tilde_ij * V_j
11.        Write O_i, l_i_new, m_i_new back to HBM
12. Return O

By fusing matrix multiplication, masking, softmax reduction, and value multiplication into a single fused GPU kernel, the $N \times N$ attention matrix is never written to or read from HBM.

Exact Gradient Backpropagation with On-Chip Recomputation

During standard deep learning backpropagation, intermediate activations from the forward pass must be retained in memory to compute backward gradients. For standard attention, this requires retaining the $N \times N$ matrix $P = \text{softmax}(S)$ in HBM for each attention head across every transformer layer.

FlashAttention completely eliminates the $O(N^2)$ backward memory storage requirement through activation recomputation.

Instead of storing $P \in \mathbb{R}^{N \times N}$ , the forward pass writes only the output $O \in \mathbb{R}^{N \times d}$ and the compact softmax statistics $(m, l) \in \mathbb{R}^N$ to HBM.

Activation Storage Footprint Comparison per Layer (Single Head):
-----------------------------------------------------------------
Standard Attention:   O(N^2) elements  (Matrix P stored in HBM)
FlashAttention:       O(N d + N) elements  (Q, K, V and statistics m, l)

During the backward pass:

Blocks of $Q_i, K_j, V_j$ and statistics $m_i, l_i$ are reloaded into SRAM from HBM.
The attention matrix block $S_{ij} = Q_i K_j^T / \sqrt{d}$ and normalized probability block $P_{ij} = \text{diag}(l_i)^{-1} \exp(S_{ij} - m_i)$ are recomputed on-the-fly directly inside fast SRAM registers.
Gradients with respect to queries $d Q$ , keys $d K$ , and values $d V$ are accumulated:

$d V_j = \sum_{i=1}^{T_r} P_{ij}^T d O_i$

$d P_{ij} = d O_i V_j^T$

$d S_{ij} = P_{ij} \circ \left( d P_{ij} - \text{rowsum}(d P_{ij} \circ P_{ij}) \right)$

$d Q_i = \frac{1}{\sqrt{d}} \sum_{j=1}^{T_c} d S_{ij} K_j, \quad d K_j = \frac{1}{\sqrt{d}} \sum_{i=1}^{T_r} d S_{ij}^T Q_i$

Although recomputing $S_{ij}$ and $P_{ij}$ requires an additional $2 N^2 d$ FLOPs (a theoretical 33% increase in backward FLOPs), execution speed increases by $2\times$ to $3\times$ because the kernel avoids reading gigabytes of attention weights over the slow HBM interconnect.

Evolution Across Generations: FlashAttention-1, FlashAttention-2, and FlashAttention-3

The FlashAttention architecture has progressed through three major generational refinements, adapting to modern GPU microarchitectures:

FlashAttention-1 (2022)

Target Hardware: NVIDIA Ampere (A100) and Turing architectures.
Outer Loop Structure: Outer loop iterates over $K, V$ blocks; inner loop iterates over $Q$ blocks.
Parallelization Scheme: Parallelized across Batch and Attention Heads dimensions.
Accumulator Rescaling: Rescaled output $O_i$ at each step in SRAM, incurring non-matmul FLOPs.
Peak Compute Efficiency: Achieved 25% to 40% of theoretical peak device FLOPs.

FlashAttention-2 (2023)

Target Hardware: NVIDIA Ampere, Ada Lovelace, and Hopper architectures.
Outer Loop Structure: Inverted loop ordering: outer loop iterates over $Q$ blocks; inner loop iterates over $K, V$ blocks.
Parallelization Scheme: Parallelized across Batch, Attention Heads, and Sequence Length ( $Q$ blocks), maximizing Streaming Multiprocessor occupancy even for batch size 1.
Accumulator Rescaling: Deferring division by the normalizer until final loop exit, accumulating unnormalized values directly in registers.
Peak Compute Efficiency: Achieved 50% to 73% of theoretical peak device FLOPs.

FlashAttention-3 (2024)

Target Hardware: NVIDIA Hopper architecture (H100, H200).
Asynchronous Data Movement: Leverages Tensor Memory Accelerator (TMA) hardware units to transfer tensor tiles directly between HBM and SRAM bypassing register files.
Warp Specialization: Partitions warps into dedicated producer and consumer groups, achieving complete hardware overlap between memory loading and matrix multiply operations.
Low-Precision FP8 Attention: Supports FP8 (E4M3 and E5M2) attention pipelines with block quantization and GEMM scaling.
Peak Compute Efficiency: Achieved 75% to 85% of theoretical peak device compute (over 650 to 800 TFLOP/s FP16 and up to 1.2 PFLOP/s FP8).

Quantitative Hardware Utilization and Serving Benchmarks

The operational benefits of IO-aware attention extend across both foundation model pre-training and high-concurrency inference serving.

GPU Execution Profile on NVIDIA A100 (BF16, d=128, Forward Pass):
-----------------------------------------------------------------
Sequence Length (N)   Standard PyTorch Attention    FlashAttention-2    Compute Speedup
2,048                 0.95 ms                       0.32 ms             3.0x
8,192                 12.80 ms                      2.10 ms             6.1x
16,384                51.40 ms                      6.40 ms             8.0x
64,512                Out Of Memory (OOM)           52.30 ms            Enables Long Context

Memory Footprint Reduction

By reducing peak activation memory from $\Theta(N^2 \cdot L \cdot H)$ to $\Theta(N \cdot d \cdot L \cdot H)$ (where $L$ is model depth and $H$ is head count), FlashAttention allows training context windows to scale from 2,048 tokens to 32,768, 131,072, and beyond on standard cluster topologies without requiring activation checkpointing to disk or multi-node tensor parallel communication overhead.

Impact on Modern LLM Architectures

FlashAttention has been adopted as the standard attention kernel across major inference and training runtimes, including vLLM, TensorRT-LLM, Hugging Face Transformers, Megatron-LM, and PyTorch (torch.nn.functional.scaled_dot_product_attention). Its principles have also been extended to adjacent paradigms:

PagedAttention: Applying FlashAttention block tiling over non-contiguous virtual memory page tables to eliminate KV cache fragmentation in inference engines.
FlashDecoding / FlashDecoding++: Partitioning the key-value sequence length across multiple thread blocks and performing parallel reductions to accelerate single-token autoregressive generation at long context lengths.
Multi-Head Latent Attention (MLA): Combining low-rank key-value projections with fused FlashAttention kernels to compress KV cache memory bandwidth during autoregressive decoding.

By identifying the memory hierarchy as the true performance limiter of deep sequence models, IO-aware computing has transformed the algorithmic foundation of large-scale artificial intelligence.