Where attention’s square comes from
One query compares with every key to make one row. Repeating that for every query fills an N × N attention matrix.
One row, then N rows
N queries × N keys = N² pairwise scores
Follow the blue query as it meets every yellow key. Each magenta score completes one cell in its row; row-wise softmax turns that row into green probabilities. Then the next query repeats the same work.
512² cells · 75.5 MB (72 MiB)
Naively materialized fp16 score or probability tensor, batch size 1, across 12 layers × 12 heads. Excludes gradients and all other activations. Standard materialized attention stores N²; FlashAttention is exact and still quadratic in arithmetic, but avoids materializing the full N × N matrix in HBM.
Query 1 compares with key 1 → score cell (1, 1)
query
key
score
probability
ready · 0 / 64 cells