Scaled scores, normalized over keys

logits  = np.einsum('bhsd,bhtd->bhst', q, k) / np.sqrt(D_k)
weights = softmax(logits, axis=-1)
values  = np.einsum('bhst,bhtd->bhsd', weights, v)

In the einsum, s is the query token, t the key token and d the head feature. Only d disappears in the score contraction.

L holds scores (logits); A holds normalized weights. Dividing by √D_k counteracts the growth in dot-product variance when component variances are roughly one. Softmax runs over the keys, the last axis, and the notebook’s version is unmasked, with no dropout.