Implement multi-head self-attention from scratch

Write masked multi-head self-attention from scratch, the layer at the core of a Transformer. Given a batch of token embeddings, project them to queries, keys and values, split them into heads, run scaled dot-product attention in each head, merge the heads and apply the output projection. Support a padding mask and a causal mask.

Requirements

• Functional: a module MultiHeadSelfAttention(d_model, num_heads) with forward(x, padding_mask=None, causal=False). x has shape (batch, seq_len, d_model) and the output has the same shape.
• Separate learned projections for queries, keys and values, and one output projection. d_model must divide evenly by num_heads.
• Scores are Q·Kᵀ / sqrt(d_head), softmax over the key positions, then a weighted sum of the values.
• padding_mask has shape (batch, seq_len) and marks padded keys that no query may attend to. With causal=True, position i may attend only to positions <= i.
• You may use tensor operations (matmul, softmax), but not a built-in attention layer.
• Non-functional: the cost in time and memory as seq_len grows, and numerical stability.

Clarifying questions

• Which framework, and may I use its built-in softmax and linear layers?
• Should a padding mask hide padded keys only, or padded queries as well?
• What should a query row return when every key is masked?
• Do you want dropout on the attention weights?

What a strong answer covers

• The shapes at every step: (B, T, D) → (B, H, T, D/H) for the heads, (B, H, T, T) for the scores, then back to (B, T, D), and where the transposes and reshapes go.
• Why the scores are divided by sqrt(d_head): without it large dot products push softmax into regions with tiny gradients.
• Masks applied before the softmax by setting blocked scores to a large negative value, and the fully masked row handled without producing NaNs.
• Cost: O(T²·D) time and O(T²) memory per head for the score matrix.
• A test: with one head and the projections set to identity, compare against a hand-computed softmax, and check that a causal output at position i doesn't change when later tokens change.

Follow-ups

• Add a key-value cache for autoregressive decoding and say what it saves.
• Multi-query or grouped-query attention: what changes and why it helps inference.
• How do FlashAttention-style kernels avoid materializing the T x T matrix?
• Where do positional encodings such as RoPE enter this layer?