One attention head mixes positions by one notion of similarity. Multi-head attention splits the d features of each position into h heads of d / h features, lets each head attend on its own, then joins them and mixes them with one more linear layer.

import torch
import torch.nn.functional as F

B, T, D, H = 2, 5, 8, 2
x = torch.randn(B, T, D)
heads = x.view(B, T, H, D // H)
heads = heads.transpose(1, 2)
print(heads.shape)  # (B, H, T, D // H)
out = F.scaled_dot_product_attention(
    heads, heads, heads)
merged = out.transpose(1, 2)
print(merged.reshape(B, T, D).shape)

The split is a reshape to (B, T, H, D // H) followed by a transpose. Reshaping straight to (B, H, T, D // H) has the same shape but cuts the numbers in the wrong places, mixing features of different positions into one head.

A padding mask stops every query from attending to padded keys: their scores become minus infinity before the softmax. For F.scaled_dot_product_attention, a boolean attn_mask is True where attending is allowed.

Write MultiHead(d, heads), an nn.Module with four nn.Linear(d, d, bias=False) layers q, k, v and o. forward(x, pad=None) takes (batch, positions, d) and an optional boolean pad of shape (batch, positions), True at padded positions no query may attend to, and returns o of the merged heads, shaped like x.

The explanation opens here once every check passes.