Online Softmax, One Tile at a Time

Chapter 10 · Mathematics for AI and Machine Learning · Advanced · 1:30

An interactive live film: FlashAttention reads a row of attention scores tile by tile, keeping a running maximum, sum and output that are rescaled by e^(m_old - m_new) whenever a larger score appears; drag the scores or change the tile size and the result still equals the full softmax exactly.

Narration

Here is one row of attention scores s, the eight blue bars, with a value v under each one. The safe softmax needs the maximum and the sum over the whole row.

FlashAttention reads the row in tiles, two scores at a time, in the orange box. It keeps only three running numbers: the maximum m, drawn as the red line, the sum l, and the output o.

The rest of the narration is in the book's Companion.