Flash Attention: The Engine Behind Efficient Transformers

^{SnackOnAI Engineering | Senior AI Systems Researcher | Technical Deep Dive | April 25, 2026}

The transformer community spent years optimizing the wrong thing. Approximate attention methods traded model quality for reduced compute, assuming FLOP count was the bottleneck. FlashAttention (2022) proved the bottleneck was never compute. It was memory traffic. The N×N attention matrix was being written to GPU HBM (High Bandwidth Memory, the GPU's main off-chip memory) and read back three times per forward pass. Every byte of that write-and-read was idle time for the compute cores. The fix was not to do less math. It was to do the same math without ever writing the intermediate matrix to HBM.

This newsletter dissects the full FlashAttention lineage: v1 (IO-aware tiling), v2 (warp-level work partitioning), v3 (Hopper asynchrony and FP8), and the just-published v4 (Blackwell co-design, 1613 TFLOPs/s). It also covers the IO optimality proof (arXiv:2402.07443) that shows FlashAttention is theoretically optimal for its SRAM regime, the block-sparse variant, and the INT8/FP8 quantized variant. This is not a tutorial. It is a systems audit.

Scope: all four FlashAttention versions, IO complexity theory, and the block-sparse and quantized variants. Not covered: linear attention approximations (Linformer, Performer), or the full transformer training stack.

FlashAttention by Tri Dao et al. is an IO-aware exact attention algorithm. "Exact" means no approximation, same mathematical output as standard attention. "IO-aware" means the algorithm is designed around the GPU memory hierarchy, not around FLOP minimization.

The repo has over 15,000 GitHub stars. It is the default attention implementation in PyTorch (torch.nn.functional.scaled_dot_product_attention), Hugging Face Transformers, vLLM, llm.c, MLC-LLM, and virtually every production LLM training and inference stack. Context length in production LLMs grew from 2-4K (GPT-3, OPT) to 128K (GPT-4) to 1M+ (Llama 3). FlashAttention is the enabling mechanism.

Benchmarks across the four versions:

The core insight is that standard attention launches three separate CUDA kernels, each materializing the N×N attention matrix to HBM. FlashAttention fuses all attention operations into one kernel using tiling and online softmax, never writing the N×N matrix.

^{Focus on the memory flow. Standard attention writes the N×N matrix three times. FlashAttention writes it zero times. The speedup is not from doing less math. It is from eliminating the memory bottleneck entirely.}

Three techniques interlock to make this work:

Tiling. Q, K, V are split into blocks. The algorithm loads one block of K,V and loops over all Q blocks in SRAM, computing the partial attention output without ever needing the full N×N matrix resident in memory at once.

Online softmax. Standard softmax over a row requires seeing all N scores before computing the denominator. Online softmax tracks two running statistics per row: m (current row maximum) and l (running sum of exponentials). When a new tile of scores arrives, the output is rescaled using the updated statistics. This allows exact softmax computation one tile at a time with no approximation.

Recomputation. The backward pass needs the attention probability matrix P to compute gradients. Standard implementations save P to HBM (O(N²) memory). FlashAttention saves only the softmax statistics (m, l), which are O(N). In the backward pass, P is recomputed on-chip from Q, K, and the stored statistics. This trades extra FLOPs for eliminated HBM reads, and the trade is favorable because compute is faster than memory on modern GPUs.

v1 achieved IO optimality but left GPU occupancy on the table. The inner loop over Q blocks was handled by a single thread block, limiting parallelism. FlashAttention-2 (2023) parallelized the attention computation across Q blocks as well, distributing work across thread blocks to increase GPU occupancy. It also reduced non-matmul FLOPs in the online softmax update and improved warp-level work partitioning to reduce shared memory communication. Result: 50-73% of theoretical A100 maximum FLOPs/s, 225 TFLOPs/s per A100 for GPT-style training.

H100 (Hopper architecture) introduced WGMMA (Warpgroup Matrix Multiply Accumulate) instructions and TMA (Tensor Memory Accelerator) for asynchronous data movement. FlashAttention-2 achieved only 35% utilization on H100 because it did not use these features. The core bottleneck revealed: the H100 SXM5 has 989 TFLOPs/s of FP16 matmul throughput but only 3.9 TFLOPs/s for exponential (the softmax function). For head dimension 128, exponential can consume 50% of wall-clock time if not overlapped with matmul.

FlashAttention-3 (2024) introduced:

Result: 740 TFLOPs/s FP16 (75% utilization on H100), 1.2 PFLOPs/s FP8. 1.5-2x over FA2.

B200 (Blackwell) doubles tensor core throughput relative to H100 while shared memory bandwidth and exponential unit throughput scale more slowly. This asymmetric scaling breaks the assumptions that FA3's pipeline was designed around. FA4 (March 2026) introduces:

Result: 1613 TFLOPs/s BF16 on B200 (71% utilization), 1.3x over cuDNN 9.13, 2.7x over Triton.

The memory numbers tell the whole story. At 32K context, standard attention needs 2GB just for the attention matrix. FlashAttention needs 16MB. Without this difference, long-context LLMs do not fit in GPU memory.

The IO optimality proof is the part of FlashAttention the community underreacts to. It is not just fast. It is provably the best possible algorithm for this hardware. The evolutionary improvements in v2-v4 are exploiting new hardware features, not finding new algorithmic improvements over v1's IO strategy.

Scenario: Measure the concrete impact of FlashAttention across sequence lengths on an A100 80GB GPU with fp16, batch size 1, 32 heads, head dim 128 (GPT-3-style configuration).

Input: Attention layer, sequence lengths 512 / 4K / 32K / 128K.

Step 1: Memory footprint (the critical metric)

Step 2: Latency (A100 80GB, fp16, measured)

Step 3: Training throughput (GPT-style, A100)

Step 4: IO access counts (theoretical, N=4K, d=128, M=192KB on A100)

The core insight is that modern GPUs are not compute-bound for attention. They are memory-bound. An A100 has 312 TFLOPS of FP16 tensor core throughput but only 2TB/s of HBM bandwidth. Writing a single N×N matrix at N=4096 in fp16 requires ~64MB, which at 2TB/s takes ~32 microseconds just for the write. Standard attention does this multiple times per layer. FlashAttention eliminates all of it.

The algorithm trades FLOPs for memory traffic, and this is the correct tradeoff on every GPU manufactured since 2018. The ratio of compute to memory bandwidth (called arithmetic intensity) has been increasing with every generation. H100 has 989 TFLOPS of FP16 compute and 3.35TB/s HBM bandwidth. B200 doubles the compute again. FlashAttention's trade becomes more favorable with every hardware generation, not less.

What FlashAttention trades away:

Debuggability. When something goes wrong in the attention computation, there is no attention matrix to inspect. It was never materialized. Debugging requires either adding visualization hooks that force materialization (defeating the purpose) or relying on the statistics (m, l) to diagnose numerical issues.

Non-CUDA generalization. The core implementation is tightly coupled to CUDA's memory model, TMA, and WGMMA instructions. Triton implementations exist for AMD and other backends, but they consistently lag NVIDIA's implementations in utilization. FlashAttention v4's 2.7x advantage over Triton on B200 is evidence of this.

Causal masking complexity. For causal (autoregressive) masks, approximately half the SRAM work is on masked-out tiles that contribute zero output. Some tiles can be skipped, but the scheduling adds implementation complexity and reduces the clean IO complexity analysis.

The IO optimality proof (arXiv:2402.07443). FlashAttention's HBM access count is not just low. It is provably optimal. Saha and Ye (2024) prove that for SRAM size M ≥ d², no algorithm can compute exact attention with fewer HBM accesses than FlashAttention, within constant factors. This means the algorithmic improvements in v2, v3, and v4 are not beating v1's IO strategy. They are exploiting new hardware features that the original analysis did not account for (WGMMA, TMA, 2-CTA MMA). The IO strategy itself was optimal in 2022.

Hardware co-design expertise. Each FlashAttention version requires deep knowledge of a specific GPU microarchitecture: A100's shared memory and thread block scheduler for v2, H100's TMA and WGMMA with asynchronous pipelining for v3, B200's 2-CTA MMA mode and asymmetric scaling for v4. FlashAttention v4 is implemented in CuTe-DSL embedded in Python, achieving 20-30x faster compile times than C++ template approaches while maintaining full expressivity. This is not kernel writing. It is compiler-aware microarchitecture programming.

Integration as the default. PyTorch's scaled_dot_product_attention dispatches to FlashAttention when the conditions are met (CUDA, fp16/bf16, headdim ≤ 256). Every framework that calls this gets FlashAttention without knowing it. Displacing this requires not just a faster kernel but a faster kernel that PyTorch's dispatch system prefers over an already-integrated default.

Insight One: FlashAttention did not make attention faster. It revealed that attention was never compute-bound, and the entire "approximate attention" research direction was solving the wrong problem.

Linformer, Performer, Longformer, BigBird, and dozens of approximate attention methods from 2020-2022 accepted approximation error to reduce quadratic FLOP count. They all failed to achieve wall-clock speedup in practice. FlashAttention demonstrated why: the bottleneck was HBM memory bandwidth, not arithmetic operations. All those approximate methods reduced the wrong quantity. They reduced FLOPs while leaving the memory bottleneck intact. FlashAttention increased FLOPs (due to backward recomputation) while eliminating the memory bottleneck. The result is that exact attention is now faster than approximate attention in practice. The entire approximate attention research direction was obsoleted not by a better approximation, but by the insight that approximation was never necessary.

Insight Two: FlashAttention v4's implementation in CuTe-DSL embedded in Python is more significant than its throughput numbers. It signals that the write-attention-kernels-in-CUDA-C++ era is ending.

FlashAttention v3 was implemented in CUDA C++ with complex template metaprogramming. Compile times were on the order of minutes. FlashAttention v4 is written in CuTe-DSL, a Python-embedded DSL for GPU kernel authorship, achieving 20-30x faster compile times while reaching 1613 TFLOPs/s on B200. This is not a productivity improvement. It is a demonstration that the compiler infrastructure for GPU kernel development has matured to the point where Python-level abstractions can match hand-written CUDA for attention workloads. The implication: future attention optimizations will be written in higher-level languages by teams that cannot write CUDA C++. The barrier to contributing production-quality attention kernels is dropping by a generation.

FlashAttention's backward pass does more work than its forward pass, and is still faster than standard attention's backward pass despite the extra FLOPs.

Standard attention's backward pass reads the saved N×N probability matrix P from HBM to compute gradients, then reads Q, K, V again. FlashAttention's backward pass recomputes P on-chip from Q, K, and the stored statistics (m, l), doing a full second forward pass in blocks. This is more FLOPs. It is faster because it eliminates the massive HBM read of the N×N matrix. On A100, reading 64MB (the P matrix at N=4K) takes ~32 microseconds at 2TB/s bandwidth. The arithmetic to recompute P takes ~8 microseconds on tensor cores. Trading 32μs of memory time for 8μs of compute is a 4x favorable exchange. This is the cleanest possible illustration of why FLOP count is the wrong optimization target on modern GPUs.

The history of attention optimization before FlashAttention is a history of solving the wrong problem. FLOP count was the measured quantity because it was easy to measure. HBM bandwidth consumption was the actual bottleneck because it was harder to see. FlashAttention made the bottleneck visible and then eliminated it with a proof of optimality attached. The community has spent the four years since applying the same IO-aware reasoning to other GPU bottlenecks: KV cache management (PagedAttention), weight loading (speculative decoding), and now operator-level pipelining (FA3, FA4). The frame is now standard. It started with attention.