All four scores at once

Softmax turns the scores into weights. Each score is exponentiated, and each result is divided by the sum of all of them. The attention output is then the weighted sum of the values. In practice the largest score is subtracted from every score before exponentiating, because exp of a large score overflows a floating-point number. Subtracting the max fixes that without changing the answer, since the same factor exp(-max) ends up in every numerator and in the denominator, where it cancels.

Here the max is 5, so the four weights before dividing are exp(1-5), exp(3-5), exp(5-5) and exp(2-5):

0.0183   0.1353   1.0000   0.0498      sum ≈ 1.2034

Multiplying each by its value and adding gives 10 × 0.0183 + 20 × 0.1353 + 30 × 1 + 40 × 0.0498 ≈ 34.8813, and the output is 34.8813 / 1.2034 ≈ 28.9847. (Each figure in this post is worked out from the unrounded numbers and then rounded to four places, so redoing a sum from the printed figures can come out a digit off in the last place. The code further down does the whole calculation unrounded.)

The calculation needed the max of all four scores before any weight could be computed, and the sum of all four weights before any of them could be used. For a long sequence that is a whole row of the score matrix, one entry per key. The main goal in the FlashAttention paper is “to avoid reading and writing the attention matrix to and from HBM”, the GPU’s large but comparatively slow memory, and that requires “computing the softmax reduction without access to the whole input”. FlashAttention loads the inputs in blocks into the GPU’s small on-chip memory instead, and works on one block at a time.

Two blocks

So split the four keys into two blocks of two, and pretend the second block hasn’t arrived yet.

The first block holds the scores 1 and 3, with values 10 and 20. Its largest score, 3, is the max so far, and the weights are computed against it: exp(1-3) ≈ 0.1353 and exp(3-3) = 1. They add up to about 1.1353, and the weighted sum of the values is 10 × 0.1353 + 20 × 1 ≈ 21.3534. If the keys ended here, the output would be 21.3534 / 1.1353, about 18.81, which is the right answer for a query with only these two keys. But they don’t end here, so nothing is divided yet. What gets carried forward is three numbers: the max so far (3), the sum so far (1.1353) and the weighted sum so far (21.3534).

The second block holds the scores 5 and 2. The 5 is larger than the 3 that every weight so far was computed against, so each of those weights is too large by a factor of exp(5-3). To undo that, both carried numbers get multiplied by exp(3-5) ≈ 0.1353. The sum becomes 1.1353 × 0.1353, about 0.1537, and the weighted sum becomes 21.3534 × 0.1353, about 2.8899.

The multiplication also changed each weight inside the sum. The 1 that belonged to the score 3 became 0.1353, which is exp(3-5). The 0.1353 that belonged to the score 1 became 0.0183, which is exp(1-5). Those are the first two numbers of the all-at-once list above. Multiplying by exp(3-5) is the same as subtracting 5 instead of 3 inside each exponent, because exp(s-3) × exp(3-5) = exp(s-5) for any score s, so after the correction the first block looks exactly as it would have if 5 had been the max from the beginning. The weighted sum is made of those same weights times the values, so multiplying it by 0.1353 corrects it too.

The second block itself is computed against the new max directly: exp(5-5) = 1 and exp(2-5) ≈ 0.0498, which add up to about 1.0498, with a weighted sum of 30 × 1 + 40 × 0.0498 ≈ 31.9915. Then adding the two blocks together gives a sum of 0.1537 + 1.0498 ≈ 1.2034 and a weighted sum of 2.8899 + 31.9915 ≈ 34.8813. Those are the two totals from the all-at-once calculation, and 34.8813 / 1.2034 is again about 28.9847.

More than two blocks

Nothing in the second step depended on there being only two blocks. Every new block does the same thing: compare its max with the max so far, multiply the carried sum and weighted sum by exp(old max - new max), and add the new block’s own sum and weighted sum, computed against the new max. If the new block’s max is not larger, the factor is exp(0) = 1 and the carried numbers stay as they are. The update itself is older than FlashAttention: both FlashAttention papers cite Milakov and Gimelshein, who in 2018 used it to get the max and the sum of a plain softmax “in a single pass over input vector”. FlashAttention carries the weighted sum of the values along with those two, and FlashAttention-2 leaves it undivided, “un-scaled” in the paper’s word, until after the last block, as the example above did.

Here it is in NumPy, for any block size, next to the all-at-once version:

import numpy as np

def softmax_weighted_sum(s, v):
    w = np.exp(s - s.max())
    return (w[:, None] * v).sum(axis=0) / w.sum()

def online(s, v, block):
    m, l, acc = -np.inf, 0.0, np.zeros(v.shape[1])
    for i in range(0, len(s), block):
        sb, vb = s[i:i + block], v[i:i + block]
        m_new = max(m, sb.max())
        c = np.exp(m - m_new)          # rescales what is already there
        p = np.exp(sb - m_new)
        l = l * c + p.sum()
        acc = acc * c + (p[:, None] * vb).sum(axis=0)
        m = m_new
    return acc / l

s = np.array([1.0, 3.0, 5.0, 2.0])
v = np.array([[10.0], [20.0], [30.0], [40.0]])
print(online(s, v, block=2), softmax_weighted_sum(s, v))

rng = np.random.default_rng(0)
s = rng.normal(0, 4, size=1000)
v = rng.normal(size=(1000, 8))
diff = np.abs(online(s, v, block=64) - softmax_weighted_sum(s, v)).max()
print(diff)

[28.98474626] [28.98474626]
1.3322676295501878e-15

On the worked example both versions print 28.98474626. On 1,000 random scores in blocks of 64, they differ by at most about 1e-15, at the level of floating-point rounding. That is what “exact attention algorithm” means in the first FlashAttention paper: it computes the same attention, not an approximation of it. The loop starts with the max at negative infinity and the sum at zero, so on the first block the correction factor is exp(-inf) = 0 and there is nothing yet to correct.

On a GPU the loop runs over blocks of queries as well as blocks of keys, with the block sizes chosen to fit the on-chip memory, and each query in a block has its own max, sum and partial output. FlashAttention-2 puts the blocks of queries in the outer loop and walks the blocks of keys inside it, so one block of queries goes through every block of keys before its output is finished; the original FlashAttention had the two loops the other way round.

The worked figures were checked against the NumPy code above, and the update is the same one the FlashAttention papers write out in matrix form.