1 / 10 Attention Internals

Decode-Step Attention Shapes

A decoder-only model has d_model = 512 and h = 8 heads, so d_k = d_v = 64. It is generating token 1025, with 1024 tokens already in the KV cache. Inside one head, the query for the new token is Q with shape (1, 64), and the cached keys and values are K and V, each (1024, 64).

Work out the shape of Q @ K.T, the axis softmax normalizes over, and the shape of the head's output.

A

Q @ K.T is (1, 512), because the query keeps the full model dimension inside a head and only the output projection splits it into 8 pieces.

B

Q @ K.T is (1, 1024), but softmax normalizes over the 64-wide feature axis of V, so each value dimension sums to 1 before the weighted average.

C

Q @ K.T is (1, 1024): one score per cached token. Softmax normalizes across those 1024 scores, and scores @ V gives (1, 64), which is concatenated across the 8 heads back to 512.

D

Q @ K.T is (1024, 1024), since the attention score matrix is always square in the sequence length; softmax runs over each of its rows and the output is (1024, 64).

E

K.T @ Q is computed instead, giving (64, 1); softmax normalizes over the 64 feature dimensions and the output is (64, 1024).