7.4 Flash Attention 1, 2, 3

In Chapter 3, we analyzed the complexity of self-attention and noted that its time and memory complexity scale quadratically with the sequence length ($O(L^2)$). As we push foundation models to handle massive contexts (e.g., 128k to 1M+ tokens), this quadratic scaling becomes a fatal bottleneck.

However, the primary bottleneck in scaling attention is often not compute (FLOPs), but memory bandwidth (I/O). The time spent moving the large attention matrix between slow GPU High Bandwidth Memory (HBM) and fast on-chip SRAM dominates the execution time.

FlashAttention revolutionized transformer training by making attention IO-Aware. It reduces the number of memory reads/writes between HBM and SRAM, achieving massive speedups while yielding mathematically exact results.

In this section, we trace the evolution of FlashAttention from its original tiling concept to the hardware-specific optimizations of FlashAttention-3.

1. The Memory Wall in Attention

Standard attention computes the matrix $\mathbf{S} = \mathbf{Q}\mathbf{K}^T$, writes it to HBM, reads it back to compute softmax, writes $\mathbf{P} = \text{softmax}(\mathbf{S})$ to HBM, and finally reads it back to compute $\mathbf{O} = \mathbf{P}\mathbf{V}$.

For a sequence length of $L=4096$, the attention matrix alone consumes $4096 \times 4096 \times 4 \text{ bytes} \approx 64 \text{ MB}$ per head. Moving this matrix back and forth between HBM and SRAM wastes immense amounts of time.

2. FlashAttention-1: Tiling and Recomputation

Introduced by Dao et al. (2022) [1], FlashAttention-1 addresses the memory wall by avoiding the materialization of the full $O(L^2)$ attention matrix in slow HBM. It relies on two core ideas:

2.1 Tiling (Forward Pass)

FlashAttention loads blocks of $\mathbf{Q}, \mathbf{K}, \mathbf{V}$ from HBM to the fast, small on-chip SRAM. It computes attention for these blocks and writes the output back to HBM. To do this without seeing the full row of the attention matrix (which is needed for the Softmax denominator), it uses Online Softmax (based on work by Milakov and Gimelshein). This algorithm tracks running maximums and sums of exponentials to rescale the output incrementally, yielding mathematically exact results without materializing the full matrix.

2.2 Recomputation (Backward Pass)

To compute gradients during the backward pass, standard attention requires the stored $L \times L$ attention matrix. FlashAttention avoids storing this by recomputing it on the fly in SRAM during the backward pass using the stored blocks of $\mathbf{Q}, \mathbf{K}, \mathbf{V}$. While this adds some FLOPs (recomputation), it drastically reduces HBM reads/writes, resulting in a net speedup of 2-4x.

3. FlashAttention-2: Better Parallelism and Work Partitioning

FlashAttention-2 (2023) [2] recognized that while FA-1 reduced I/O, it left some GPU compute resources underutilized. It introduced several algorithmic refinements:

• Parallelism over Sequence Length: FA-1 parallelized over batch size and number of heads. For small batch sizes or long sequences, this left many GPU streaming multiprocessors (SMs) idle. FA-2 adds parallelism over the sequence length dimension (blocks of $\mathbf{Q}$), significantly improving utilization.
• Refactored Online Softmax: FA-2 refactored the online softmax to reduce the number of non-matrix-multiplication operations (like exponentials), which are slow on GPUs compared to the heavily optimized Tensor Cores.
• Support for Head Dimensions up to 256: Expanded support for larger head dimensions used in some modern architectures.

These changes resulted in a 2x speedup over FA-1, achieving up to 70% of the theoretical peak FLOPs on A100 GPUs.

4. FlashAttention-3: Asynchrony and Low-Precision for Hopper

FlashAttention-3 (2024) [3] targets the specific architectural features of NVIDIA’s Hopper architecture (H100) to push performance even further. The Hopper architecture introduced features that FlashAttention-3 exploits:

• Asynchronous Execution (Overlapping Compute and I/O): Hopper introduced TMA (Tensor Memory Accelerator), which can move data between HBM and SRAM asynchronously, independent of the Tensor Cores. FA-3 overlaps the loading of the next block of $\mathbf{K}$ and $\mathbf{V}$ with the computation of the current block, hiding I/O latency completely.
• WGMMA (Warpgroup Matrix Multiply-Accumulate): FA-3 utilizes H100’s new WGMMA instructions, which are designed for larger matrix operations and offer higher throughput than legacy instructions.
• Natively Support FP8: FP8 reduces memory bandwidth requirements by half compared to FP16. FA-3 handles the scaling factors required for low-precision matrix multiplication accurately, maintaining model quality while doubling throughput.

4.5 Comparison Table: Evolution of FlashAttention

5. PyTorch SDPA: Verify That Flash Is Actually Selected

PyTorch’s scaled-dot-product attention (SDPA) dispatch depends on the installed PyTorch/CUDA version, GPU capability, dtype, shapes, masks, and dropout. The example below is intentionally CUDA-only and uses current sdpa_kernel backend selection. It demonstrates causal language-model attention; it is not a variable-length training loop.

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

if not torch.cuda.is_available():
    raise RuntimeError("Flash SDPA requires a supported CUDA GPU")

device = torch.device("cuda")
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
batch, heads, sequence, head_dim = 2, 8, 2048, 64

q = torch.randn(batch, heads, sequence, head_dim, device=device, dtype=dtype)
k = torch.randn_like(q)
v = torch.randn_like(q)

with torch.inference_mode(), sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
    output = F.scaled_dot_product_attention(
        q, k, v,
        dropout_p=0.0,
        is_causal=True,
    )

assert output.shape == q.shape
assert output.dtype == dtype and torch.isfinite(output).all()
print(torch.backends.cuda.flash_sdp_enabled(), output.shape, output.dtype)

For training, set dropout_p explicitly because SDPA applies the supplied probability independently of module training mode. is_causal=True does not remove padding from variable-length batches. Use the framework’s supported combined causal/padding or variable-length representation, then confirm that it remains eligible for the intended backend. Compare output and gradients with the math backend on small tensors within a dtype-appropriate tolerance, and use profiler/kernel logs to prove which backend ran [4]. If forced Flash dispatch rejects the inputs, change the configuration or allow a documented fallback; do not report FlashAttention performance from the math path.

Quizzes

References

1. Dao, T., Fu, D., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and memory-efficient exact attention with IO-awareness. arXiv:2205.14135.
2. Dao, T. (2023). FlashAttention-2: Faster attention with better parallelism and work partitioning. arXiv:2307.08691.
3. Dao, T., & Haziza, N. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608.
4. PyTorch. Scaled Dot Product Attention. PyTorch documentation.