FlashAttention

Standard attention is not slow because of FLOPs. It is slow because it materialises the n x n attention matrix in HBM, then reads it back to multiply by V. On an A100, HBM is roughly 1.5 TB/s and on-chip SRAM is ~19 TB/s; you spend most of the wall-clock waiting on HBM traffic, not on tensor cores. Tri Dao's FlashAttention (2022) reframes attention as an IO problem and solves it with classic numerical-analysis tools: tiling, online softmax, and recomputation in the backward pass.

The bandwidth bottleneck

A textbook attention forward pass does roughly:

S = Q K^T          # write n*n matrix to HBM
P = softmax(S)     # read S, write P
O = P V            # read P, write O

For n = 8192, fp16, that is 256 MiB of attention matrix written and read twice. The matmuls themselves are cheap relative to those HBM round-trips. The standard implementation is HBM-bandwidth-bound at long context lengths and burns wall-clock waiting on memory.

The fix: tiling + online softmax

FlashAttention loops over blocks of Q (outer) and K/V (inner), holding small tiles in SRAM and computing softmax incrementally using an online algorithm (Milakov & Gimelshein 2018) that keeps running max and sum-of-exp values. The full n x n matrix never exists. The forward pass writes only O and a tiny log-sum-exp vector.

for Qi in blocks(Q):
    mi, li = -inf, 0
    Oi = 0
    for Kj, Vj in blocks(K, V):
        Sij = Qi @ Kj.T / sqrt(d)
        mij = max(mi, rowmax(Sij))
        Pij = exp(Sij - mij)
        li  = exp(mi - mij) * li + rowsum(Pij)
        Oi  = exp(mi - mij) * Oi + Pij @ Vj
        mi  = mij
    write Oi, (mi, li)

For the backward, you do not store the full attention probabilities (the textbook approach) - you recompute them block by block from the saved (m, l) statistics. Recomputation is cheap because the FLOPs were never the bottleneck.

Why it is faster and uses less memory

The counterintuitive bit: FlashAttention does more FLOPs (because of the recomputation in backward) but is 2-4x faster end-to-end and uses O(n) memory instead of O(n^2). The win comes from doing those extra FLOPs on data that lives in SRAM, where bandwidth is an order of magnitude higher than HBM. You trade abundant compute for scarce bandwidth, which is exactly the trade modern GPUs reward.

The 1 -> 2 -> 3 progression

FlashAttention-2 was largely a kernel rewrite: the v1 backward had too much shared-memory traffic between warps, and the parallel decomposition kept some SMs idle on long sequences. v2 fixed both.

What Hopper bought FlashAttention-3

Hopper's selling features matter here because they are asynchronous:

• WGMMA (Warpgroup Matrix Multiply Accumulate) runs matmuls in the background while the warp issues more loads. v3 overlaps the softmax of block j with the matmul of block j+1.
• TMA (Tensor Memory Accelerator) handles HBM-to-SMEM copies as a separate hardware unit, freeing the warp from address calculation.
• FP8 tensor cores at 2x the throughput of FP16. v3 has an FP8 path that hits ~1.2 PFLOPS on H100, with accuracy preserved by per-block scaling.

The combined effect is a kernel that is asynchronous end-to-end - copies, matmuls, and softmax run in parallel rather than serially. On Blackwell (B100/B200) the same ideas extend to FP4 and a wider TMA.

When it falls down

• Short sequences (n < 512). The setup overhead beats the bandwidth savings; PyTorch SDPA is fine.
• Non-standard masks. Custom block-sparse, sliding-window, or document masks need a kernel variant. FlexAttention (PyTorch 2.5+) generates these from a mask description; FA's own block-sparse path covers the common cases.
• Training stability at FP8. The v3 FP8 path needs scaling discipline; a careless mix of FP8 attention with FP32 norms still works, but a careless mix everywhere does not.

Further reading

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness - the original paper with the IO model and roofline analysis.
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning - the v2 kernel rewrite.
• FlashAttention-3 blog by Tri Dao - readable walk-through of the Hopper-specific changes with benchmark plots.
• Dao-AILab/flash-attention GitHub - reference implementation and the practical install matrix per GPU.