FlashAttention: Definition & Examples

FlashAttention computes exact dense attention by processing tiles that fit in fast GPU memory, avoiding storage of the full attention matrix in slower device memory. For a quant model reading long filings or event sequences, this can reduce memory pressure without changing the intended attention rule. “Exact” describes the mathematical operation; it does not promise identical floating-point bits.

Move less data

For one head, scaled dot-product attention is

O=\operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}+M\right)V,

where $Q,K,V$ contain queries, keys and values, and $M$ is an optional additive mask. Softmax acts across keys for each query.

A conventional implementation writes the score and probability matrices to high-bandwidth memory (HBM), then reads them again. FlashAttention forms small tiles in on-chip SRAM, combines their contributions and discards temporary scores. During training, it recomputes needed tiles for the backward pass rather than saving the full matrix. Dense arithmetic still scales quadratically with sequence length; avoiding the quadratic intermediate storage is the distinction. See the original paper, §§2–3.

Why a tile can be discarded

Softmax needs a denominator covering every allowed key, so separately normalising each tile is insufficient. The online normalizer maintains a maximum and a rescalable sum. For one query row, extend that summary to include the value-weighted numerator. For a block

A

, define

m_A=\max_{i\in A}s_i,\qquad \ell_A=\sum_{i\in A}e^{s_i-m_A},\qquad u_A=\sum_{i\in A}e^{s_i-m_A}v_i.

Here $s_i$ already includes scaling and any mask; $v_i$ may be a vector. For two nonempty blocks with finite allowed scores, set $m=\max(m_A,m_B)$ and merge:

\begin{aligned} \ell&=e^{m_A-m}\ell_A+e^{m_B-m}\ell_B,\\ u&=e^{m_A-m}u_A+e^{m_B-m}u_B,\\ o&=u/\ell. \end{aligned}

Each rescaling factor changes the reference maximum, while preserving the underlying exponential sums. The blocks need not retain their individual weights. Fully masked blocks require explicit handling rather than evaluating an indeterminate subtraction of infinities. This is a mathematical summary of the mechanism, not GPU implementation code; FlashAttention-2, §2.3 develops its use in attention.

One row, checked two ways

Take synthetic scores $(0,\ln2,\ln3)$ and scalar values $(1,2,4)$ , with no mask or dropout. Direct softmax gives weights $(1,2,3)/6$ , so

o=\frac{1\cdot1+2\cdot2+3\cdot4}{6}=\frac{17}{6}.

Split after the second score. Block A has

(m_A,\ell_A,u_A)=(\ln2,3/2,5/2)

; block B has

(m_B,\ell_B,u_B)=(\ln3,1,4)

. Raising A’s reference maximum to

\ln3

multiplies both its sums by

2/3

. Consequently,

\ell=\frac23\frac32+1=2,\qquad u=\frac23\frac52+4=\frac{17}{3},\qquad o=\frac{17/3}{2}=\frac{17}{6}.

Synthetic one-query example: block A has scores 0 and ln 2, values 1 and 2, maximum ln 2, normalizer 3/2 and numerator 5/2. Block B has score ln 3, value 4, normalizer 1 and numerator 4. Rescaling block A by 2/3 produces merged normalizer 2 and numerator 17/3, hence output 17/6.

Original synthetic example. Both the normalizer and the value-weighted numerator must be expressed relative to the same maximum before merging.

Adding the unrescaled block sums would instead give $(5/2+4)/(3/2+1)=2.6$ , which is wrong: those sums use different reference maxima.

What the name does and does not guarantee

FlashAttention-2 improves the algorithm’s work partitioning across GPU thread blocks and warps and reduces non-matrix-multiplication work. It retains dense attention rather than replacing it with a sparse or low-rank approximation. These are separate design choices; the 2023 paper, §3 explains the changes.

Floating-point operations depend on evaluation order, so different tiled implementations or precisions can produce small numerical differences even for the same mathematical expression. PyTorch’s numerical-accuracy notes explain this distinction. Actual speed depends on hardware, sequence length, head dimension, batch size, precision and supported masks. A kernel benchmark is not an end-to-end training-speed guarantee.

For the underlying scoring and weighted-context construction, read the SelfMadeQuant query–key–value lesson (public preview; membership required for the full lesson).

Practice: can you average block outputs?

A query has scores $(0,\ln3,\ln2)$ and scalar values $(2,6,10)$ , without masking or dropout. Block A contains the first two positions and block B the third. Their separately normalised outputs are $5$ and $10$ . Is their simple average the full attention output?

Answer. No. Relative to the unshifted exponentials, A has mass

1+3=4

and B has mass

2

. Their contributions must therefore be combined as

(4\cdot5+2\cdot10)/6=20/3

, not

(5+10)/2=7.5

. Equivalently, rescale both block summaries to the common maximum

\ln3

: the merged denominator is

2

, numerator

40/3

, and output

20/3

. A tile’s output alone loses the information needed to weight it against another tile.

Sources

Dao et al. (2022), FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, §§2–3, Algorithm 1 and Theorem 1.
Milakov and Gimelshein (2018), Online normalizer calculation for softmax, §§3–3.1.
Dao (2023), FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, §§2.3 and 3.
PyTorch 2.15 documentation, Numerical accuracy, floating-point ordering and mathematically identical computations.