19. Multi-Head Attention: Split & Merge Heads

Problem

A single softmax(QKT/√d_k)V computes one weighted average per query — one "view" of how the sequence relates to itself. Multi-head attention runs several of these views in parallel on lower-dimensional slices of the same projected q/k/v, so different heads can specialize (one might track local adjacency, another long-range dependency) without paying for H independent full-width attention computations.

This problem isolates the reshape mechanics that make "run H attentions in parallel" a single batched matmul instead of a Python for-loop over heads: splitting the last dimension d_model into (num_heads, head_dim), doing attention per head using the extra head axis as an additional batch axis, and merging the head outputs back into d_model at the end. It takes already-projected q, k, v as input, isolating just the head split/merge and per-head attention from the surrounding linear projections.

Source: src/19_multi_head_attention.py

def multi_head_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int,
    mask: torch.Tensor | None = None,
) -> torch.Tensor: ...

>>> q = k = v = torch.randn(2, 5, 12)
>>> multi_head_attention(q, k, v, num_heads=3).shape
torch.Size([2, 5, 12])

Step-by-Step Approach

1. Reshape each of q, k, v from (batch, seq, d_model) to (batch, seq, num_heads, head_dim) with view, where head_dim = d_model // num_heads.
2. transpose(1, 2) to move the heads axis before the sequence axis: (batch, num_heads, seq, head_dim) — now heads act like an extra batch dimension for the matmuls that follow.
3. Run ordinary scaled dot-product attention on that 4D tensor: qh @ kh.transpose(-1, -2) / sqrt(head_dim), softmax, then @ vh — PyTorch's batched matmul broadcasts over both the batch and head axes simultaneously.
4. If a mask is given as (batch, seq_q, seq_k), insert a heads axis with unsqueeze(1) so it broadcasts identically to every head before masked_fill.
5. Transpose the head output back: (batch, num_heads, seq, head_dim) -> (batch, seq, num_heads, head_dim), call .contiguous(), then view back to (batch, seq, d_model) to merge the heads.
6. Sanity-check with num_heads=1: it must reduce exactly to ordinary single-head scaled dot-product attention, since one head spanning the full d_model is mathematically identical to no head split at all.

The key insight is that "multiple heads" is not a loop, it's an extra tensor axis: once heads live at dimension 1 (right after batch), every subsequent op is the exact same scaled-dot-product-attention code as the single-head case, just broadcasting over one more dimension.

Reference solution

import math
import torch


def _split_heads(t: torch.Tensor, num_heads: int) -> torch.Tensor:
    # (batch, seq, d_model) -> (batch, seq, num_heads, head_dim) -> (batch, num_heads, seq, head_dim)
    batch, seq_len, d_model = t.shape
    head_dim = d_model // num_heads
    return t.view(batch, seq_len, num_heads, head_dim).transpose(1, 2)


def multi_head_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int,
    mask: torch.Tensor | None = None,
) -> torch.Tensor:
    batch, seq_len, d_model = q.shape
    head_dim = d_model // num_heads

    # fold the head axis into the batch dims so one batched matmul computes
    # every head's attention at once: (batch, heads, seq, head_dim)
    qh = _split_heads(q, num_heads)
    kh = _split_heads(k, num_heads)
    vh = _split_heads(v, num_heads)

    scores = qh @ kh.transpose(-1, -2) / math.sqrt(head_dim)
    if mask is not None:
        # mask is (batch, seq_q, seq_k); unsqueeze a heads axis so it
        # broadcasts identically across every head
        scores = scores.masked_fill(~mask.unsqueeze(1), float("-inf"))
    attn_weights = torch.softmax(scores, dim=-1)
    out_h = attn_weights @ vh  # (batch, heads, seq_q, head_dim)

    # merge heads back; contiguous() is required because transpose() only
    # permutes strides -- view() needs the memory layout to match
    return out_h.transpose(1, 2).contiguous().view(batch, seq_len, d_model)

Key Functions & Tricks

• tensor.view(batch, seq, num_heads, head_dim) — splits the model dimension into heads; requires the tensor to already be contiguous in that layout.
• tensor.transpose(1, 2) — moves heads next to batch so batched matmul treats (batch, heads) as one combined batch dimension.
• tensor.contiguous() before the merge view — transpose only changes strides, not memory layout, and view requires contiguous memory; skipping this raises a RuntimeError.
• mask.unsqueeze(1) — inserts the heads axis into a mask shaped for a single "virtual head" so it broadcasts across all heads without duplicating data.
• Batched @ over 4D tensors — PyTorch broadcasts matmul over every leading dimension beyond the last two, so (B, H, T, Dh) @ (B, H, Dh, T) needs no explicit loop over heads.

How to Recognize This Pattern