Cross-checking against PyTorch

The NumPy attention is checked three ways: against explicit loops (one np.dot per entry), against qh @ kh.swapaxes(-1, -2) / sqrt(D_k), and against PyTorch’s F.scaled_dot_product_attention with dropout_p=0.0.

The PyTorch heads are built the same way — reshape(B, S, H, D // H).permute(0, 2, 1, 3) — and match the NumPy qh exactly, even though the permuted tensor is not contiguous.

The PyTorch output matches the NumPy head values to rtol=1e-10, and so does the merged result after .contiguous().view(B, S, D).