Why (B, H) goes first

The final two dimensions are the matrix axes. Putting (B, H) first lets matrix multiplication batch over both axes, so one @ or einsum computes every head of every example at once.

Each head then compares every query token with every key token in its own feature space of D_k features.

The reshape does not learn those spaces: the projection matrices do. Splitting heads only regroups features that W_Q, W_K and W_V already produced.