FlashAttention

FlashAttention is an implementation strategy for exact softmax attention. It reduces data movement between GPU high-bandwidth memory (HBM) and on-chip memory by tiling the computation. It does not replace softmax with a linear attention rule.

That distinction matters when reading model papers: an attention architecture and the kernel that executes it are separate choices. A dense attention layer can use an efficient kernel while still doing quadratic arithmetic in sequence length.

A conventional implementation can materialize scores and probabilities for every query–key pair. FlashAttention processes blocks and keeps running row statistics, avoiding those full intermediate matrices in HBM. The visual compares an illustrative score tile with the full matrix; it is not a GPU memory profiler.

The hard part is not splitting a matrix into tiles. It is combining softmax contributions when a later tile contains a larger score. Earlier contributions must be rescaled into the new exponent range before adding the next block.

For one query row, keep the maximum $m$, denominator $\ell$, and unnormalized weighted-value numerator $u$. For a new score block $s$ with values $V$:

$m' = \max(m, \max s), \qquad a = e^{m-m'}$ $p = e^{s-m'}, \qquad \ell' = a\ell + \sum p, \qquad u' = au + pV$

After the final block, output $u/\ell$. If you store a normalized output instead of $u$, its update must also include the old denominator. Mixing those two conventions is a common implementation bug.

For fixed head width, dense attention still forms a quadratic number of pairwise interactions. Avoiding a stored $T\times T$ matrix changes intermediate storage and memory traffic, not that arithmetic count.

This executable scalar-value example checks the online invariant. It is a correctness illustration, not a GPU kernel.

1from math import exp
2scores, values = [1., 2., 1001., 1002.], [2., -1., 4., 8.]
3m, denominator, numerator = float('-inf'), 0., 0.
4for start in range(0, len(scores), 2):
5    block = scores[start:start + 2]
6    new_m = max(m, max(block))
7    scale = exp(m - new_m)
8    denominator *= scale
9    numerator *= scale
10    for offset, score in enumerate(block):
11        weight = exp(score - new_m)
12        denominator += weight
13        numerator += weight * values[start + offset]
14    m = new_m
15online = numerator / denominator
16weights = [exp(s - max(scores)) for s in scores]
17reference = sum(w*v for w,v in zip(weights, values)) / sum(weights)
18assert abs(online - reference) < 1e-12
19print(online)

Why use it? Memory traffic and temporary storage can limit practical attention performance. Tiling, fusion, and recomputation target those implementation costs. Actual speed depends on the GPU, shapes, dtype, masking, and kernel version.

Is it sparse attention? Standard FlashAttention computes dense attention. Sparsity specifies which interactions are retained; it is a separate design dimension.

Is DeepSeek V4.1-Flash a FlashAttention paper? No. The model name is not an attention taxonomy. Its documented attention system includes compression and sparse indexing; read the model breakdown separately.

A new tile raises the running maximum from 5 to 10. What must happen to the previous numerator and denominator? Both are multiplied by $e^{-5}$ before adding the new tile's contributions. Rescaling only one changes the result.

Sources: FlashAttention, Dao et al., official implementation and supported kernels. Compare the architectural alternative in linear attention.