Online softmax

A softmax over a row needs the row's maximum before it can exponentiate. The usual recipe makes three passes over the row: find the maximum, sum the exponentials of the shifted values, then divide. A tiled engine sees the row a few keys at a time, so it cannot do that without storing the whole row.

The online softmax keeps two running values for each row: the maximum m so far, and the sum l of exponentials relative to it. When a new tile brings keys with a larger maximum, the old sum was relative to a smaller maximum, so it must shrink. It is multiplied by

α = exp(m_old − m_new)

and the same factor rescales the output accumulator O, which is a sum of exponentials times values. After the last tile, O / l is the exact attention output for that row. Nothing about the row has to be kept except m, l and O.

In the engine

The engine works in log base 2, so it uses 2^z in place of exp. The scores are scaled to z = s · log₂e / √d and the rescale is α = 2^(m_old − m_new). The power of two is computed by the exp2 unit (see Fixed point), and α is stored as an unsigned 16-bit fraction.

Watch one row

The demo runs the engine on the toy sentence, records the row state after each K/V tile, and shows it next to the float softmax of the same inputs. The difference is fixed-point rounding only.