FlashAttention

FlashAttention computes attention in tiles. A block of query rows stays on chip. The engine streams key and value tiles past it, one 4-row tile at a time, folds each tile into the running state of online softmax, and writes the output once at the end. The N × N score matrix never exists in memory.

The schedule

• Q rows are split into blocks. A block holds cfg_qrows rows (4, 8, 16 or 32 here).
• For each block, every key and value tile is loaded once and used by all the row groups in the block.
• With the causal mask, tiles above the diagonal are skipped, and the K/V loads stop at the block's last needed tile.
• The output of a block is normalized and written once.

The traffic

Each K or V tile is 16 bytes (4 rows of 4 int8 values), each Q tile is 16 bytes, and each O tile is 32 bytes (4 rows of 4 int16 values). With nb = N / 4 row groups:

• FLASH: reads 16·nb bytes of Q, plus 32 bytes per K/V tile load. Writes 32·nb bytes of O. A larger Q block means fewer K/V loads.
• NAIVE: writes and reads the score tiles Z (32 bytes each) and the probability tiles P (32 bytes each) through HBM, once per pass.

The NAIVE schedule is the standard three-pass softmax: pass 1 writes the scores, pass 2 finds the row maxima, pass 3 writes the probabilities, and pass 4 multiplies them by V.

Compare the bytes

The tables come from the golden traffic model, which the cocotb end-to-end test checks against the RTL byte counters.