Mistake: all the attention weights sum to 1

“Softmax normalizes, so the whole weight matrix adds up to 1.”

scores = np.array([[2., 1., 0.], [0., 2., 1.]])
rowwise = np.exp(scores) / np.exp(scores).sum(axis=-1, keepdims=True)
assert np.allclose(rowwise.sum(axis=-1), 1.0)
assert np.isclose(rowwise.sum(), 2.0)

Each row sums to 1 and the matrix sums to 2 — one per query. Softmax runs along the last axis, so every row is its own distribution over the keys.

Normalizing the whole matrix also totals 1 and is a different thing entirely, which is why the grand total is a bad check.

For a (batch, heads, queries, keys) score tensor, softmax takes the last axis, the keys.