Is a tiled attention kernel an approximation

Why standard attention is slow

S = Q @ K^T       read Q and K, write an N x N matrix
P = softmax(S)    read N x N, write N x N
O = P @ V         read N x N and V, write N x d

Three round trips through high-bandwidth memory. On the reference accelerator that memory moves about 3 TB/s while the on-chip memory moves about 30 TB/s, so every round trip is a factor-of-ten slowdown against keeping the data on chip. From about 2K tokens, attention becomes memory-bound before it becomes compute-bound, and classic kernels underuse the hardware by 4 to 10 times.

The tiling

for each block of Q (about 128 x 128):
    load the Q tile on chip
    for each block of K, V:
        load the K and V tiles on chip
        scores for the tile
        running softmax aggregation
        accumulate into the output tile
    write the output tile once

The score matrix exists only one tile at a time. The backward pass recomputes some values instead of storing them, which is a second memory win.

Why the result is exact

The softmax needs the global maximum and the global sum over a row, and the kernel only sees one tile at a time. It keeps a running pair, maximum and sum, and whenever a new tile raises the maximum it rescales the earlier sum and the earlier partial output by the exponential of the old maximum minus the new one. At the end every term is expressed against the same global maximum and divided by the same global sum. The lesson's description is "bit-identical output to standard attention", with the caveat that 16-bit arithmetic is not associative, so reordering can move the last bits.

Versions

What to tell the colleague

Run the two kernels on the same inputs and compare the outputs. The difference is at the level of 16-bit rounding, the same noise you get from a different batch size, not a change to the model.