How can attention become faster without ignoring token pairs? Look at the large intermediate matrices a straightforward implementation writes to memory. FlashAttention changes how the calculation moves and combines those numbers.
Before you begin: You understand attention scores, softmax, and the difference between memory and arithmetic cost.
FlashAttention applies this kind of thinking to GPU attention. It is an exact attention algorithm up to normal floating-point differences. It does not achieve its main saving by simply ignoring most token pairs.
Find the large intermediate
With n query positions and n key positions, ordinary full attention has n * n pairwise scores per head. A straightforward implementation writes the score matrix and attention probabilities to high-bandwidth GPU memory.
For 8,192 positions, one score matrix has over 67 million entries. At two bytes each, that is 128 MiB before considering additional heads, batches, or other tensors. Doubling sequence length quadruples this matrix size.
FlashAttention processes blocks of queries, keys, and values using fast on-chip storage. It avoids materializing the full score or probability matrix in high-bandwidth memory. Dense attention still has quadratic pairwise arithmetic; reduced memory traffic is a different improvement from linear-time attention.
Why softmax seems to get in the way
Softmax divides each exponential score by the sum over all available keys. If keys arrive in blocks, how can an early block know the final denominator?
Maintain a running maximum, a scaled exponential sum, and a scaled weighted-value sum. When a later block has a larger maximum, rescale the previous accumulators before adding the new block. The common rescaling leaves the final ratio unchanged.
This scalar version shows the principle. A GPU kernel works with many rows and vector values at once.
# Runnable: Python 3, standard library.
import math
scores, values = [2.0, -1.0, 3.0], [10.0, 20.0, 30.0]
peak, denominator, numerator = -math.inf, 0.0, 0.0
for score, value in zip(scores, values):
new_peak = max(peak, score)
rescale = math.exp(peak - new_peak)
weight = math.exp(score - new_peak)
denominator = denominator * rescale + weight
numerator = numerator * rescale + weight * value
peak = new_peak
online = numerator / denominator
raw = [math.exp(s - max(scores)) for s in scores]
reference = sum(w * v for w, v in zip(raw, values)) / sum(raw)
assert math.isclose(online, reference)
print(round(online, 3))
The assertion checks numerical agreement for a teaching example. It is not a GPU performance benchmark or a replacement for a tested attention kernel.
Make a later score much larger
The online example sees score 3 after score 2. Change the last score to 1000. Predict which calculation becomes numerically dangerous: directly evaluating exp(1000), or subtracting a running maximum before exponentiation.
Follow the rescaling
The direct exponential can overflow. With the running maximum, the new score has relative exponential one and the old accumulators are scaled by a very small factor. The final weighted average approaches the last value, 30, while remaining representable in this example. This is why the running maximum is part of the algorithm, not merely a performance trick.
Run the reference calculation with the same max-subtraction rule and compare within a tolerance. Then reverse the input order. Small rounding differences are possible, but the mathematical result should agree. Neither test measures GPU speed; both test the numerical idea that makes blockwise computation possible.
Understand what the saving does not promise
Tiling is not the same as block-sparse attention. A tiled algorithm can still cover every allowed pair. A causal mask still excludes future positions, and its handling must remain correct at block boundaries.
The backward pass can recompute intermediate values instead of storing the full attention probabilities. This trades additional arithmetic for less memory movement. Actual speed depends on sequence length, head dimensions, dtype, hardware, and the selected backend. Avoid a universal “three times faster” claim without a matching measurement.
Use a framework's supported optimized attention path for real work. Benchmark against its actual baseline, synchronize GPU timing, warm up kernels, and compare outputs and gradients within appropriate tolerances. A CPU loop like the one above teaches the mathematics but will not reproduce GPU speedups.
Exercise: sequence length doubles. Does FlashAttention make dense pairwise attention computation grow only twofold?
Compare your reasoning
No. Pairwise arithmetic still grows roughly fourfold. The algorithm avoids storing the quadratic attention matrix and can use memory bandwidth much better, but those are separate claims.
Next, examine how positional rotations change the scores being computed.
Sources
The central algorithm and memory analysis are in FlashAttention. Later improvements are described in FlashAttention-2. Consult your framework's documentation to determine which backend actually runs on your hardware.
Continue: Rotary Position Embeddings (RoPE).