Online Softmax Without the Attention Matrix notebook

Python starts when you choose Run.

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.

Original CPU online-softmax illustration for one query row and scalar values; not a FlashAttention GPU kernel.

The code uses Python 3 and the standard library. Saved output is from the checked CPU run.

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}")

tile=1: output=30.000000
tile=2: output=30.000000
tile=4: output=30.000000
score_entries: full_matrix=16777216, one_row_tile=64

Sources

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness: https://arxiv.org/html/2205.14135v2

Python runs locally in this page. Saved results and figures are visible before execution.

Download full experiment notebook