FlashAttention

Summary

Attention needs a normalized weighted sum, not a permanently stored score matrix. FlashAttention tiles that computation to reduce global-memory storage and data movement while preserving the mathematical attention operator. Floating-point reduction order can change bits; validate numerical tolerance rather than claiming bitwise identity.

Task: one query sees scores $[0,\log 3]$ with scalar values $[2,6]$. Process one score/value pair per tile without retaining prior scores. Recover the same output as dense softmax, then change masking and the performance regime.

Use finite allowed scores and a float64 reference tolerance of $10^{-12}$ for these checks. Production tolerances depend on dtype, shape, and kernel.

The mechanism

A materialized baseline writes score and probability matrices of size $n\times n$ to high-bandwidth memory, HBM. The tiled algorithm keeps a score tile and row summaries in on-chip working storage. It still computes every allowed dense pair; memory traffic, auxiliary storage, arithmetic, and latency are different quantities.

For processed scores $s_j$ and value vectors $v_j$, retain

Here $u$ is an unnormalized numerator; output is $O=u/\ell$. For a new valid block, set

Changing reference from $m$ to $m^{\prime }$ multiplies every earlier exponential by $a$. Rescaling both totals preserves their ratio and admits the new block under a common reference. Initialize $m=-\infty$, $\ell =0$, and $u=0$; the first valid block has old-state scale zero. Skip entirely forbidden blocks without evaluating an undefined $-\infty -(-\infty )$.

The second maximum rescales old totals by $1/3$. Dense evaluation agrees: $(1\cdot 2+3\cdot 6)/(1+3)=5$. Forgetting rescaling gives $(2+6)/(1+1)=4$, a useful failing check.

Change masking

New task: the second score/value pair is future-only for this causal query. Exclude it before local reductions and skip that empty tile. Check: retained state stays $(m,\ell ,u)=(0,1,2)$, giving output $2$, not $5$. Use when: visibility, padding, or a window restricts the allowed pairs. This reference rejects a fully masked row; match the chosen library’s declared empty-row behavior rather than dividing by zero.

Change the performance regime

New task: compare a short sequence on a different GPU. Keep output tolerance checks and measure kernel time, launch overhead, utilization, and peak memory separately. Check: illustrative baseline $7$ microseconds versus tiled $8$ gives speedup $7/8=0.875$, despite unchanged mathematical output. Those timings are invented inputs to the comparison, not measurements. Use when: shape or hardware changes the bottleneck; asymptotics cannot supply a universal crossover length.

Backward pass: trade memory for compute

Backward can recompute score/probability tiles from Q/K and saved row normalization information. It need not retain the entire probability matrix. Saved outputs, statistics, and RNG requirements depend on the algorithm and dropout convention. This is a kernel-level recompute-versus-store tradeoff, not cost-free memory removal.

Dense arithmetic remains $O(n^{2}d)$. Auxiliary attention storage can be $O(nd)$, linear in sequence length at fixed width. This does not imply linear HBM traffic or latency; tile sizes, on-chip capacity, and scheduling matter.

Reconstruct: does reducing quadratic storage make dense attention subquadratic in sequence length?

m_new = max(m_i, max(S_ij))
ell_new = exp(m_i - m_new) * ell_i + sum(exp(S_ij - m_new))
O_new   = (exp(m_i - m_new) * ell_i * O_i + exp(S_ij - m_new) @ V_j) / ell_new