The course

Beyond GEMM · 3/3

1 head · 4 tokens · d = 4

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.

Chapter 2 complete.

Instruction
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.

query

Arrow keys step, space plays or pauses.

Reading the picture

  1. 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.

  2. 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

  3. 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.