Attention algorithms and numerical stability
Online Softmax Without the Attention Matrix
Running maxima and rescaled sums recover exact attention arithmetic without retaining every score. A small CPU example makes the invariant explicit.
Does computing exact attention require storing every pairwise score? The answer depends on whether an intermediate value must remain available after it has contributed to the output. A weighted average can often be accumulated without keeping all its weights.
FlashAttention applies this idea with GPU-aware tiling and online softmax statistics, avoiding materialization of the full attention matrix in high-bandwidth memory. Its exact-attention path changes the computation's organization rather than dropping token pairs. The paper's IO analysis and kernel results go beyond the small arithmetic example here. FlashAttention paper
One row is enough to see the mechanism
Consider one query with four already-computed attention scores: 1000 + log(1), 1000 + log(2), 1000 + log(3), and 1000 + log(4). Let the corresponding scalar values be 10, 20, 30, and 40.
Softmax is unchanged by subtracting the same constant from every score. The normalized weights are therefore 1/10, 2/10, 3/10, and 4/10. Their weighted value is 1 + 4 + 9 + 16 = 30.
Computing exp(1000) directly is numerically unsafe. A stable calculation subtracts the maximum score first. Streaming adds one complication: a later tile may contain a larger maximum, changing the reference scale.
Maintain three quantities: the running maximum m, a denominator l containing the sum of scaled exponentials, and a numerator z containing the correspondingly scaled value sum. When a new tile raises the maximum to m_new, multiply the old l and z by exp(m−m_new), then add the new tile's contributions on that same scale. The final answer is z/l.
In the two-tile example, the second tile contains the largest score. Rescaling the first tile's contributions is exactly what prevents the early values from being weighted incorrectly. Simply averaging the two tile outputs would fail because their normalization denominators differ.
Separate storage from arithmetic
The accompanying standard-library code computes both the dense weighted sum and the tiled version, then checks agreement within floating-point tolerance. It operates on one row and scalar values. A full attention implementation additionally computes query-key dot products, handles vector-valued outputs, and applies masks and scaling consistently.
For sequence length 4,096, a materialized score matrix contains 16,777,216 entries per head. A one-query reference implementation processing 64 keys at a time needs only 64 temporary score entries. These counts cover score intermediates alone; they exclude inputs, outputs, and running statistics. They also describe different concurrency: a full score matrix versus one query row processed in tiles. Their ratio is not an application-level memory reduction. Actual GPU kernels tile multiple query rows to obtain reuse and sufficient parallel work.
The dense pairwise arithmetic remains, and this CPU calculation establishes no GPU speedup. Its result is an invariant: after every tile, the accumulated numerator and denominator describe all keys seen so far using one common numerical scale.
Run the example
The Python 3 example uses only the standard library. The notebook, source, figure data and recorded output are included in the complete bundle.
Inspect the code and recorded output
from math import exp, log
scores = [1000 + log(i) for i in range(1, 5)]
values = [10.0, 20.0, 30.0, 40.0]
maximum = max(scores)
weights = [exp(score - maximum) for score in scores]
dense = sum(w * v for w, v in zip(weights, values)) / sum(weights)
def streaming(tile_size):
running_max = float("-inf")
denominator = 0.0
numerator = 0.0
for start in range(0, len(scores), tile_size):
tile_scores = scores[start:start + tile_size]
tile_values = values[start:start + tile_size]
new_max = max(running_max, max(tile_scores))
correction = exp(running_max - new_max)
local_weights = [exp(s - new_max) for s in tile_scores]
denominator = denominator * correction + sum(local_weights)
numerator = numerator * correction + sum(
w * v for w, v in zip(local_weights, tile_values))
running_max = new_max
return numerator / denominator
for tile_size in (1, 2, 4):
result = streaming(tile_size)
assert abs(result - dense) < 1e-10
print(f"tile={tile_size}: output={result:.6f}")
assert abs(dense - 30.0) < 1e-10
length, block = 4096, 64
print(f"score_entries: full_matrix={length * length}, one_row_tile={block}")
Recorded CPU output
tile=1: output=30.000000 tile=2: output=30.000000 tile=4: output=30.000000 score_entries: full_matrix=16777216, one_row_tile=64