03 Chapter 2 · attention
Attention
One query against four keys. The dots are the scores, softmax turns scores into weights, and the output is those weights times V.
scores Q[q] · Kᵀ
weights softmax(scores)
output Σ weightᵢ · Vᵢ
One head, four tokens, d = 4, small integers so every product is readable. This is the dataflow, not a full attention kernel: no scaling by 1/√d, no mask, no extra heads. QKᵀ and the mix with V are the same shape as the GEMM microkernel. Softmax, between them, is not.
one head · QKᵀ → softmax → ×V
Q and K are in place. Play to take the four dots.
Arrow keys step, space plays or pauses.
This browser prefers reduced motion. Play still runs the instruction and jumps straight to the result.
Reading the picture
-
Scores are dots
Each step dots the chosen query with one key and writes that score. Token 2 matches key 2, so that score is the largest. Pick another query and the winning key moves.
-
Softmax is not a GEMM
The weights are the scores after a stable softmax. They sum to 1, and the brightest V row is the key the query matched. The exp and the max subtraction live on their own page. Open the softmax page
-
Two GEMMs, softmax between them
QKᵀ writes the scores and the mix with V writes the output. Both are the same shape as the microkernel, and both want their panels packed first. Open the pack page Open the microkernel
Query t2 = [0, 1, 1, 0]. Keys and V stay fixed.