FlashAttention (FlashAttention-1, 2, and 3): Mathematical Foundations, IO-Awareness, Online Softmax Tiling, and Asynchronous Hardware Acceleration

The standard scaled dot-product attention mechanism defines the computational core of modern autoregressive large language models and vision transformers. While algorithmically straightforward, its practical execution on modern graphics processing units (GPUs) has historically suffered from severe hardware memory bottlenecks.

Standard self-attention incurs an $O(N^2)$ asymptotic computational and memory footprint with respect to sequence length $N$. On modern accelerators like NVIDIA A100 and H100 GPUs, the primary bottleneck is not raw arithmetic compute capacity (FLOPs), but memory access latency and bandwidth between slow off-chip High Bandwidth Memory (HBM) and fast on-chip Static Random-Access Memory (SRAM).

The FlashAttention family of algorithms, developed by Tri Dao and collaborators across Stanford University and Princeton University, fundamentally altered transformer execution by introducing IO-aware exact attention. By reformulating softmax computation through incremental online reduction, tiling operations across on-chip SRAM cache blocks, avoiding the materialization of the full $N \times N$ attention matrix in HBM, and leveraging asynchronous hardware primitives, FlashAttention achieves multi-fold wall-clock speedups while reducing peak memory footprint from quadratic $O(N^2)$ to linear $O(N)$ without numerical approximation.

The GPU Memory Wall in Standard Attention

To understand why standard attention stalls modern processors, one must analyze the physical hierarchy of GPU architectures:

• High Bandwidth Memory (HBM): Serves as main device memory (e.g., 80 GB on an NVIDIA A100 with 1.5 to 2.0 TB/s bandwidth; 80 GB or 96 GB on an NVIDIA H100 with 3.35 TB/s bandwidth).
• On-Chip SRAM (Shared Memory / L1 Cache): Integrated directly into each Streaming Multiprocessor (SM). An A100 SM contains 192 KB of combined shared memory/L1 cache per SM (totaling roughly 20 MB across 108 SMs with aggregate bandwidth exceeding 19 TB/s). An H100 SM contains 228 KB per SM (exceeding 33 TB/s aggregate bandwidth).
• Compute Units (Tensor Cores): Capable of dense matrix multiplications at hundreds of teraflops (e.g., 312 TFLOPS dense FP16 on A100; 989 TFLOPS dense FP16 on H100).

In standard attention, given query matrix $Q \in \mathbb{R}^{N \times d}$, key matrix $K \in \mathbb{R}^{N \times d}$, and value matrix $V \in \mathbb{R}^{N \times d}$ with sequence length $N$ and head dimension $d$:

1. Compute pre-softmax attention logits: $S = Q K^T \in \mathbb{R}^{N \times N}$, written to HBM.
2. Compute attention probabilities: $P = \text{softmax}(S / \sqrt{d}) \in \mathbb{R}^{N \times N}$, read from and written to HBM.
3. Compute output representation: $O = P V \in \mathbb{R}^{N \times d}$, reading $P$ and $V$ from HBM.

For a moderate sequence length of $N = 16,384$ tokens with 32 heads in 16-bit precision, materializing $S$ and $P$ requires allocating $2 \times (16,384 \times 16,384) \times 2 \text{ bytes} \approx 1.07 \text{ GB}$ per head per layer. Across 32 heads, a single transformer layer requires 34.3 GB of intermediate HBM allocations just for the attention matrix.

Under the Roofline Model of hardware execution, an operation is memory-bound when its arithmetic intensity (arithmetic operations per byte of memory transfer) falls below the hardware operational balance point. Standard attention operations (scaling, masking, element-wise exponentiation, row summations, and divisions) have an arithmetic intensity close to 1 FLOP/byte. Consequently, GPU execution units spend over 70% to 80% of their execution cycles stalled waiting for memory transfers over the HBM bus.

FlashAttention-1: Online Softmax and Block Tiling

Published in Dao et al. (NeurIPS 2022), FlashAttention-1 proved that exact attention can be computed with an IO complexity of $\Theta(N^2 d^2 / M)$ HBM memory accesses where $M$ is the capacity of fast SRAM, drastically outperforming standard attention's $\Theta(N d + N^2)$ HBM memory accesses.

1. Mathematical Formulation of Online Softmax

The primary theoretical hurdle in avoiding the $N \times N$ global allocation is the normalization term in the softmax denominator. For row vector $x \in \mathbb{R}^N$, softmax is defined as:

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

where $m(x) = \max_j x_j$ is required to maintain numerical stability against floating-point overflow. Standard implementations compute $m(x)$ across the entire row, compute the sum of exponentials $l(x) = \sum_{j} e^{x_j - m(x)}$, and finally normalize each element.

FlashAttention utilizes the online softmax technique (originally introduced by Milakov and Gimelshein, 2018 and extended for attention by Rabe and Staats, 2021). If an input vector $x$ is partitioned into two sub-blocks $x^{(1)}$ and $x^{(2)}$:

1. Let local statistics for block 1 be:

$$m^{(1)} = \max(x^{(1)}), \quad l^{(1)} = \sum e^{x^{(1)} - m^{(1)}}, \quad O^{(1)} = \text{softmax}(x^{(1)}) V^{(1)}$$

1. When processing block 2 with local statistics $m^{(2)} = \max(x^{(2)})$ and $l^{(2)} = \sum e^{x^{(2)} - m^{(2)}}$, the combined running maximum is:

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

1. The combined normalizer updates via rescaling:

$$l^{(new)} = e^{m^{(1)} - m^{(new)}} \cdot l^{(1)} + e^{m^{(2)} - m^{(new)}} \cdot l^{(2)}$$

1. The running output vector updates incrementally without storing intermediate scores:

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

This mathematical identity allows the GPU to compute attention in a single fused pass over localized sub-blocks entirely within on-chip SRAM.

2. Tiling Algorithm and Block Dimensions

FlashAttention-1 partitions the sequence into blocks that fit strictly inside on-chip SRAM cache:

• Let SRAM size be $M$ bytes.
• Choose block sizes $B_c = \lceil M / (4d) \rceil$ and $B_r = \min(\lceil M / (4d) \rceil, d)$.
• Divide $Q$ into $T_r = \lceil N / B_r \rceil$ blocks of shape $B_r \times d$.
• Divide $K$ and $V$ into $T_c = \lceil N / B_c \rceil$ blocks of shape $B_c \times d$.

The algorithm operates via nested loops:

• Outer loop iterates over column blocks $j \in [1, T_c]$ of $K$ and $V$, loading $K_j, V_j$ into SRAM.
• Inner loop iterates over row blocks $i \in [1, T_r]$ of $Q$, loading $Q_i$ and running accumulators $(O_i, l_i, m_i)$ into SRAM, computing local matrix multiplication $S_{ij} = Q_i K_j^T$, updating running statistics via online softmax, and writing back updated accumulators.

# Conceptual FlashAttention-1 forward loop structure
# M: SRAM capacity, N: sequence length, d: head dimension
# Q, K, V stored in global memory (HBM)

def flash_attention_v1_forward(Q, K, V, Br, Bc):
    O = zeros_like(Q)
    l = zeros((N,))
    m = fill((N,), -infinity)
    
    # Outer loop over K, V blocks
    for j in range(0, N, Bc):
        K_j = K[j:j+Bc, :]  # Load Bc x d into SRAM
        V_j = V[j:j+Bc, :]  # Load Bc x d into SRAM
        
        # Inner loop over Q blocks
        for i in range(0, N, Br):
            Q_i = Q[i:i+Br, :]  # Load Br x d into SRAM
            O_i = O[i:i+Br, :]
            l_i = l[i:i+Br]
            m_i = m[i:i+Br]
            
            # Compute block dot products in SRAM
            S_ij = (Q_i @ K_j.T) / sqrt(d)
            
            # Online softmax reduction
            m_ij = rowmax(S_ij)
            m_new = maximum(m_i, m_ij)
            P_tilde = exp(S_ij - m_new[:, None])
            l_tilde = rowsum(P_tilde)
            
            l_new = exp(m_i - m_new) * l_i + l_tilde
            
            # Update output accumulator
            O[i:i+Br, :] = diag(exp(m_i - m_new) * l_i / l_new) @ O_i + \
                           diag(1.0 / l_new) @ (P_tilde @ V_j)
            l[i:i+Br] = l_new
            m[i:i+Br] = m_new
            
    return O

Backward Pass: Selective Recomputation over $O(N^2)$ Storage

In standard transformer training backpropagation, storing the activation matrix $P \in \mathbb{R}^{N \times N}$ for all layers creates an intolerable activation memory footprint during the backward pass. Gradient checkpointing traditionally mitigates this by recomputing the entire layer forward pass, incurring a 33% compute overhead.

FlashAttention introduces selective backward recomputation:

1. During the forward pass, FlashAttention discards $S$ and $P$. It saves only the forward outputs $O \in \mathbb{R}^{N \times d}$ and the log-sum-exp vectors $L \in \mathbb{R}^N$ where $L_i = m_i + \log(l_i)$.
2. In the backward pass, given incoming gradient $dO \in \mathbb{R}^{N \times d}$, FlashAttention loads tiles of $Q, K, V$ and statistics $L$ back into SRAM.
3. It recomputes attention logits $S_{ij} = Q_i K_j^T / \sqrt{d}$ on the fly in SRAM.
4. Softmax probabilities are reconstructed instantly via:

$$P_{ij} = \exp(S_{ij} - L_i)$$

1. Gradients $dQ, dK, dV$ are accumulated in registers and SRAM:

$$dV_j = \sum_i P_{ij}^T dO_i$$ $$dP_{ij} = dO_i V_j^T, \quad dS_{ij} = P_{ij} \circ (dP_{ij} - D_i), \quad \text{where } D_i = \text{rowsum}(dO_i \circ O_i)$$ $$dQ_i = \frac{1}{\sqrt{d}} \sum_j dS_{ij} K_j, \quad dK_j = \frac{1}{\sqrt{d}} \sum_i dS_{ij}^T Q_i$$

Because SRAM compute is orders of magnitude faster than HBM read transfers, recomputing $S_{ij}$ in SRAM during the backward pass is strictly faster than loading the saved $N \times N$ matrix from HBM. Total memory consumption drops from $O(N^2)$ to $O(N)$ linear scaling.

FlashAttention-2: Parallelism, Query Loop Swapping, and Warp Partitioning

While FlashAttention-1 demonstrated a 2x to 4x wall-clock speedup, profiling revealed that it reached only 25% to 40% of theoretical peak FP16 TFLOPS on NVIDIA Ampere GPUs. In Dao (ICLR 2024), FlashAttention-2 resolved three critical architectural inefficiencies:

1. Swapping the Loop Order (Query Outer Loop)

In FlashAttention-1, the outer loop iterated over $K, V$ blocks and the inner loop over $Q$ blocks. This required continuously reading and writing intermediate accumulators $O_i, l_i, m_i$ to and from global memory or shared memory across outer iterations.

FlashAttention-2 inverted the loop order:

• Outer Loop: Iterates over row blocks of $Q$ ($i \in [1, T_r]$).
• Inner Loop: Iterates over column blocks of $K, V$ ($j \in [1, T_c]$).

FlashAttention-1 Loop Hierarchy:
  Outer Loop: K_j, V_j (Column Tiles)
    Inner Loop: Q_i (Row Tiles) -> Repeated read/write of O_i to HBM

FlashAttention-2 Loop Hierarchy:
  Outer Loop: Q_i (Row Tiles) -> Q_i loaded once into SRAM/Registers
    Inner Loop: K_j, V_j (Column Tiles) -> Accumulates locally in SRAM
  End of Inner Loop: Writes final O_i to HBM exactly once

By keeping $Q_i$ and the output accumulator $O_i$ pinned in fast registers throughout the inner loop, FlashAttention-2 eliminates thousands of redundant HBM read/write transactions per head.

2. Reduction of Non-Matmul FLOPs and Deferred Normalization

Tensor Cores perform Matrix Multiply-Accumulate (MMA) instructions at peak throughput, but point-wise arithmetic (scaling, addition, exponential) executes on standard CUDA cores at lower throughput.

FlashAttention-2 eliminates point-wise scaling overhead in the inner loop by maintaining unnormalized output accumulators:

$$O_i^{(new)} = \text{diag}(\exp(m_i^{(old)} - m_i^{(new)})) \cdot O_i^{(old)} + \exp(S_{ij} - m_i^{(new)}) V_j$$

The costly element-wise division by the cumulative sum $l_i$ is deferred until the entire inner loop over keys and values completes. At the very end of the row block processing, a single vector multiplication normalizes the output:

$$O_i = \text{diag}(l_i^{(final)})^{-1} O_i^{(final)}$$

3. Sequence-Length Parallelism and Warp Partitioning

On modern GPUs, high occupancy across all SMs is vital. FlashAttention-1 parallelized exclusively across batch size and attention heads ($B \times H$). When batch size was small (such as single-stream inference or autoregressive prefill with small batch sizes), many SMs remained idle.

FlashAttention-2 parallelizes across three dimensions:

• Batch size ($B$)
• Attention heads ($H$)
• Number of sequence query blocks ($T_r = \lceil N / B_r \rceil$)

Furthermore, within each thread block, FlashAttention-2 restructured warp work allocation. In FlashAttention-1, warps within a block split the $K, V$ column dimensions, requiring frequent inter-warp synchronization via __syncthreads() and shared memory accumulation. FlashAttention-2 assigns different row slices of $Q$ to distinct warps within the block. Warps independently stream through $K$ and $V$ without requiring inter-warp communication, achieving up to 73% peak A100 FLOPS (225 TFLOPS FP16).

FlashAttention-3: Hardware Asynchrony, TMA, Warp Specialization, and FP8

With the release of NVIDIA Hopper architecture (H100/H800), peak FP16 compute jumped to 989 TFLOPS, while memory bandwidth reached 3.35 TB/s. However, running FlashAttention-2 on H100 achieved only ~35% of theoretical maximum compute because the hardware architecture introduced new asynchronous paradigms that software kernels had to manage explicitly.

In Dao et al. (2024), FlashAttention-3 leveraged specific Hopper hardware instructions to hit 740 to 840 TFLOPS in FP16 (75% to 85% utilization) and up to 1.2 PFLOPS in FP8 precision.

+-------------------------------------------------------------------------+
|                  NVIDIA Hopper H100 Execution Pipeline                  |
+-------------------------------------------------------------------------+
|                                                                         |
|  [ Producer Warps ]  ===> Issues Asynchronous TMA Transfers            |
|                             | (HBM -> Shared Memory / SRAM)            |
|                             v                                           |
|  [ Tensor Memory Accelerator (TMA) ] ===> Bypasses Registers/SM ALU     |
|                             |                                           |
|                             v                                           |
|  [ Asynchronous Pipeline Barrier ] ===> Synchronizes Data Readiness     |
|                             |                                           |
|                             v                                           |
|  [ Consumer Warps ]  ===> WGMMA (Warpgroup MMA on Tensor Cores)         |
|                             |                                           |
|                             v                                           |
|  [ Interleaved ALU / Softmax ] ===> Overlaps Exp/Max with Matrix Math   |
|                                                                         |
+-------------------------------------------------------------------------+

1. Hardware Asynchrony via Tensor Memory Accelerator (TMA)

On older architectures, transferring data from HBM to Shared Memory required registers: global memory $\rightarrow$ registers $\rightarrow$ shared memory. This consumed register file space and occupied SM execution pipelines.

Hopper introduced the Tensor Memory Accelerator (TMA), a dedicated hardware unit capable of copying multi-dimensional tensor tiles directly between global memory and shared memory asynchronously:

• TMA instructions are issued in a single clock cycle.
• Execution happens entirely in hardware without occupying SM registers or integer units.
• Hardware asynchronous barriers (cuda::barrier) coordinate memory transfer completion without CPU or SM polling.

2. Warp Specialization and Producer-Consumer Pipelines

FlashAttention-3 partitions the warps within a Hopper thread block into specialized functional roles:

• Producer Warps: Dedicated solely to calculating memory descriptors, configuring barriers, and issuing TMA transfer requests for the next iteration blocks ($Q_{i+1}, K_{j+1}, V_{j+1}$).
• Consumer Warps: Grouped into Warp Groups (128 threads) dedicated exclusively to executing Hopper Warpgroup Matrix Multiply-Accumulate (wgmma) instructions on Tensor Cores.

By decoupling memory orchestration from math execution, producer warps pre-fetch future matrix tiles into shared memory while consumer warps compute current block multiplications, effectively hiding memory transfer latency behind arithmetic compute.

3. Interleaving GEMMs and Softmax Point-Wise Math

In standard kernels, matrix multiplication instructions on Tensor Cores and exponent calculations on SIMT ALU pipelines run sequentially. FlashAttention-3 pipelines them concurrently:

1. Issue asynchronous wgmma for $S_{ij}^{(1)} = Q_i K_j^{(1)T}$.
2. While Tensor Cores calculate $S_{ij}^{(1)}$, execute ALU instructions to compute exponents, row-maxima, and rescale factors for the previous block's output $O_i^{(0)}$.
3. Issue asynchronous wgmma for $P_{ij}^{(0)} V_j^{(0)}$.
4. Complete normalization updates while the next GEMM processes on Tensor Cores.

4. Low-Precision FP8 Attention with Incoherent Processing

FP8 precision (using either E4M3 or E5M2 formats defined in the IEEE 754-2008 standard) doubles tensor core computational throughput. However, applying FP8 directly to attention causes catastrophic quantization error due to outlier activation magnitudes across sequence tokens.

FlashAttention-3 solves this with two structural innovations:

• Block-Wise Quantization Scaling: Computes dynamic scaling factors per local SRAM tile rather than a single static per-tensor scale factor.
• Incoherent Processing (Randomized Hadamard Transformation): Multiplies input matrices by orthonormal Walsh-Hadamard matrices prior to quantization:

$$\tilde{X} = X H$$ Because $H$ is orthogonal ($H^T H = I$), inner products are strictly preserved: $$(\tilde{Q} H) (\tilde{K} H)^T = \tilde{Q} (H H^T) \tilde{K}^T = \tilde{Q} \tilde{K}^T$$ The Hadamard transformation diffuses localized outlier values uniformly across all head dimensions, preventing single-coordinate clipping and reducing FP8 numerical error by $2.6\times$ compared to standard per-tensor quantization.

Architectural Evolution and Comparative Summary

The technical progression from un-fused attention to FlashAttention-3 highlights the transition from algorithm-centric complexity to hardware-aware systems design:

• Standard Multi-Head Attention:
• Arithmetic Complexity: $O(N^2 d)$ FLOPs
• IO Complexity (HBM Access): $\Theta(N d + N^2)$
• Memory Footprint: $O(N^2)$
• GPU Execution Mode: Memory-bound (low arithmetic intensity)
• Peak FLOPS Utilization (A100): 15% to 20%
• FlashAttention-1:
• Arithmetic Complexity: $O(N^2 d)$ FLOPs (plus small recomputation overhead)
• IO Complexity (HBM Access): $\Theta(N^2 d^2 / M)$
• Memory Footprint: $O(N)$ linear
• GPU Execution Mode: Tiled IO-aware kernel fusion with online softmax
• Peak FLOPS Utilization (A100): 25% to 40% (120 TFLOPS)
• FlashAttention-2:
• Arithmetic Complexity: $O(N^2 d)$ FLOPs
• IO Complexity (HBM Access): $\Theta(N^2 d^2 / M)$ (reduced SRAM traffic)
• Memory Footprint: $O(N)$ linear
• GPU Execution Mode: Outer query loop, sequence parallelism, warp non-communication
• Peak FLOPS Utilization (A100): 50% to 73% (225 TFLOPS)
• FlashAttention-3:
• Arithmetic Complexity: $O(N^2 d)$ FLOPs
• IO Complexity (HBM Access): $\Theta(N^2 d^2 / M)$
• Memory Footprint: $O(N)$ linear
• GPU Execution Mode: TMA hardware asynchrony, producer-consumer warp specialization, GEMM-softmax interleaving, FP8 Hadamard transformation
• Peak FLOPS Utilization (H100): 75% to 85% FP16 (740 to 840 TFLOPS), ~1.2 PFLOPS FP8

By systematically addressing hardware bandwidth constraints, cache hierarchies, and asynchronous pipeline mechanics, the FlashAttention paradigm transformed long-context transformer modeling from a theoretical computational barrier into standard production practice across modern open-weights and commercial frontier architectures.

Sources

• Fast and Memory-Efficient Exact Attention with IO-Awareness (FlashAttention-1) - Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré (NeurIPS 2022).
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning - Tri Dao (ICLR 2024).
• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision - Tri Dao, Daniel Haziza, Francisco Massa, Grigory Sizov (2024).
• Online Normalizer Calculation for Softmax - Maxim Milakov, Natalia Gimelshein (NVIDIA, 2018).
• Self-attention Does Not Need $O(n^2)$ Memory - Markus N. Rabe, Charles Staats (2021).
• Roofline: An Insightful Visual Performance Model for Multicore Architectures - Samuel Williams, Andrew Waterman, David Patterson (Communications of the ACM, 2009).
• NVIDIA Hopper Architecture In-Depth - NVIDIA Technical Blog (2022).