FlashAttention-2 from Scratch: Tiled Online Softmax

The Problem

A transformer attention head computes:

The naive implementation is simple, correct, and slow for long sequences for one structural reason: softmax must see the whole row before it can normalize it. That means the kernel has to materialize the entire N×N score matrix, run a row-wise max and exponential sum, materialize the N×N probability matrix, and only then multiply by V.

On an H100, HBM bandwidth is about 3.35 TB/s while SRAM bandwidth is an order of magnitude higher. Every round trip to HBM is expensive. For N=4K and d=128, the score matrix alone is 16M entries; in FP16 that is 32 MB read and 32 MB written before the next kernel can start. The work is not FLOP-bound — it is memory-bandwidth-bound, and the wall-clock time is dominated by moving bytes, not by arithmetic.

FlashAttention-1 (Dao et al. 2022) solved this by tiling the sequence dimension and maintaining running softmax statistics inside SRAM. FlashAttention-2 (Dao 2023) kept the same IO-aware algorithm but restructured the loop order and work partitioning so that:

• the forward pass iterates over Q tiles on the outside and K/V tiles on the inside, reducing non-matmul FLOPs;
• parallelism is exploited along the sequence dimension, not just batch and heads;
• warps are assigned consistent sub-tiles of Q, K, and V to reduce synchronization.

The result is the same numerical answer as standard attention, but with HBM traffic that scales linearly in N instead of quadratically.

The Concept

Why softmax forces a full-row buffer

The softmax of a vector x is:

To evaluate the denominator, every element of the row must be exponentiated and summed. A naive kernel therefore computes the whole S=QK^T matrix, stores it, reads it back for softmax, stores P=softmax(S), and reads P back for PV. That is three HBM round trips for an N×N tensor. On modern accelerators the HBM trips cost more than the matmuls.

Online softmax: the rescaling identity

FlashAttention replaces the global softmax with an online reduction over tiles. For a single attention row, let:

• m_{\text{old}} be the running maximum over all key tiles seen so far;
• \ell_{\text{old}} be the running sum of exponentials relative to that old maximum;
• O_{\text{old}} be the running weighted-value sum, also relative to the old maximum.

When a new key tile arrives, it has its own tile maximum m_{\text{tile}} and tile sum \ell_{\text{tile}}. The new row statistics are:

The output accumulator is rescaled by the same factor:

Why is this equivalent to the one-shot formula? Because shifting every exponent by a constant c multiplies numerator and denominator by the same constant:

So rescaling every previously accumulated term by e^{m_{\text{old}} - m_{\text{new}}} makes the entire partial result consistent with the new maximum. After the last key tile, the final output is O = O_{\text{final}} / \ell_{\text{final}}. The computation is exact up to floating-point reassociation.

This is also the source of numerical stability. The exponent is always taken relative to the largest value seen so far, never relative to an arbitrary global constant. Even when raw logits are shifted by hundreds of units, the running-max subtraction keeps the exponentials in a safe range.

From FlashAttention-1 to FlashAttention-2

Both versions avoid materializing the full N×N matrix. The difference is in scheduling and parallelism:

The Q-outer ordering matters because a Q tile stays resident in SRAM while the kernel streams over all K and V tiles. In the K-outer ordering, Q must be reloaded for each K block, which wastes memory bandwidth and increases non-matmul work. FlashAttention-2 also removes unnecessary synchronizations between warps working on the same Q tile.

Backward pass: recompute, do not store

Training requires gradients through attention. The gradient with respect to Q involves terms like dO \cdot P^T, where P is the N×N softmax matrix. Storing P for the backward pass would defeat the whole memory-saving point of FlashAttention. Instead, FlashAttention recomputes the score tiles during the backward pass using the saved Q, K, V and the forward-pass output statistics (the per-row max and sum). The recomputation is cheap because it happens tile-by-tile in SRAM, and the HBM traffic remains O(N).

For causal attention, the backward pass is especially fast in FlashAttention-2 because the triangular mask is applied inside the same tile loop. No separate sparse mask tensor is needed, and only the lower-triangular tiles are computed.

Build It

The lesson ships a pure-NumPy simulation of FlashAttention-2. It is not a GPU kernel, but it implements the exact tile-loop algebra so you can see the math work without CUDA scaffolding.

Run it:

consoleCopypython3 main.py

The file contains:

• naive_attention(Q, K, V, causal=False) — a stable reference that subtracts the row max before exp.
• flash_attention_forward(Q, K, V, Br, Bc, causal=False) — the tiled online-softmax forward pass with Q-outer loops and causal masking inside tiles.
• memory_work(N, d) — an analytical HBM traffic model that counts full N×N matrix touches and total bytes moved.
• print_traffic_table(seq_lengths) — prints the O(N²) vs O(N) comparison from N=128 to 4096.

The demo exercises four things:

1. Numerical agreement between flash_attention_forward and naive_attention across full and causal masks, with max absolute differences at float64 machine epsilon.
2. An HBM traffic table showing naive attention's N×N touches growing quadratically while FlashAttention's stay at zero.
3. An extreme-value stability check where logits are shifted by +500. A naive exp/sum/divide in float32 overflows to NaN; the online softmax remains finite and matches a stable reference.
4. A causal attention invariant: row i only attends to keys 0..i.

The shipped artifact is flash_attention_forward: you can import it into any NumPy project and use it as a drop-in, numerically exact attention implementation for experiments and teaching.

Verify It

code/tests/test_main.py locks down the implementation:

• bit-identical agreement between tiled and naive attention for multiple non-causal shapes;
• bit-identical agreement for causal masking;
• tile-size independence: changing Br and Bc does not change the result;
• causal dependency: row i depends only on keys 0..i;
• extreme-value stability under a +500 logit offset;
• unsafe softmax without max-shift overflows in float32;
• memory-work model: naive N×N touches grow quadratically while FlashAttention HBM bytes grow linearly;
• flash_attention_forward is importable and main() exits cleanly.

Why This Matters for Engineers

• Attention backend selection is no longer black magic. When vLLM or SGLang offers flash_attention_2, flash_attention_3, or flashinfer, you know the trade-off: tiled online softmax avoids the N×N HBM write, and later versions add warp-level parallelism and Hopper-specific asynchrony on top of the same idea.
• Causal attention is where FlashAttention-2 wins most. Decoder prefill with causal masking used to spend time materializing and zeroing a triangular matrix. Fusing the mask inside the tile loop removes that work entirely.
• The backward pass is the memory story. Recomputing score tiles from Q/K/V keeps activation memory at O(N), which is why long-context training with FlashAttention is feasible.
• Know when to stop hand-rolling. PyTorch scaled_dot_product_attention dispatches to FlashAttention, memory-efficient attention, or a math fallback depending on hardware and mask. For production serving, use the framework dispatch; for custom masks, sparse patterns, or research kernels, the loop algebra you just wrote is the starting point.
• Numerical stability is not an afterthought. The running-max subtraction is the same reason logsumexp is safe and why transformers do not need per-head logit clamping.

Hints

• If the tiled output differs from the naive reference by more than 1e-6, check that both are using the same scale factor and dtype; float64 is recommended for this lesson's verification.
• Tile sizes Br and Bc change the SRAM working set but not the mathematical result. Pick values that divide the sequence length evenly to avoid boundary-condition bugs in real kernels; here the code handles ragged tiles explicitly.
• Causal masking inside the tile loop means you only mask entries where the key position is greater than the query position. Do not build a separate N×N mask.
• The rescaling factor e^{m_{\text{old}} - m_{\text{new}}} is exactly 1 when the running maximum does not change, so the common case is cheap.
• Online softmax stays stable because every exponent is computed relative to the largest value seen so far, not a fixed global constant.

Key Terms