Multi-Head Self-Attention

#Learning Objectives

• Implement a batched Query/Key/Value projection as a single linear layer split into H heads.
• Compute scaled dot-product attention with the correct normalization and dtype handling.
• Apply a causal mask that prevents a position from attending to future positions.
• Inspect per-head attention weights for a fixed input and reason about what each head looks at.
• Train a small attention block on a toy task and watch the loss fall as the heads specialize.