S109 · Multi-Head Attention — Parallel Views

# multi_head_attention.py — full MHA with shape wrangling + per-head inspection.
# Run: uv run multi_head_attention.py
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
 
 
class MultiHeadAttention(nn.Module):
    """Standard multi-head self-attention, following the Transformer paper."""
 
    def __init__(self, d_model: int, n_heads: int, causal: bool = False):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.causal = causal
        # Single big projections — will reshape to (N heads, d_head) later.
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)
 
    def forward(
        self, x: torch.Tensor, return_weights: bool = False
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        B, T, D = x.shape
        assert D == self.d_model
 
        # Project + reshape into (B, T, N, d_head), then transpose to (B, N, T, d_head).
        Q = self.W_q(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        K = self.W_k(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        V = self.W_v(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
 
        # Scaled dot-product per head (batch dim = B*N).
        scores = Q @ K.transpose(-2, -1)              # [B, N, T, T]
        scores = scores / math.sqrt(self.d_head)
 
        if self.causal:
            mask = torch.triu(
                torch.ones(T, T, device=x.device), diagonal=1
            ).bool()
            scores = scores.masked_fill(mask, float("-inf"))
 
        weights = F.softmax(scores, dim=-1)           # [B, N, T, T]
        out = weights @ V                              # [B, N, T, d_head]
 
        # Concatenate heads: [B, N, T, d_head] → [B, T, N, d_head] → [B, T, d_model].
        out = out.transpose(1, 2).contiguous().view(B, T, D)
        out = self.W_o(out)                            # final mix
 
        if return_weights:
            return out, weights  # weights: [B, N, T, T]
        return out
 
 
class GroupedQueryAttention(nn.Module):
    """GQA — N Q heads but only n_kv_heads K/V heads. Q groups share K/V."""
 
    def __init__(self, d_model: int, n_heads: int, n_kv_heads: int, causal: bool = False):
        super().__init__()
        assert n_heads % n_kv_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.groups = n_heads // n_kv_heads
        self.d_head = d_model // n_heads
        self.causal = causal
        self.W_q = nn.Linear(d_model, n_heads * self.d_head, bias=False)
        self.W_k = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.W_v = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        Q = self.W_q(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        K = self.W_k(x).view(B, T, self.n_kv_heads, self.d_head).transpose(1, 2)
        V = self.W_v(x).view(B, T, self.n_kv_heads, self.d_head).transpose(1, 2)
 
        # Repeat K, V to match n_heads (each K/V shared across group Q heads).
        K = K.repeat_interleave(self.groups, dim=1)  # [B, N, T, d_head]
        V = V.repeat_interleave(self.groups, dim=1)
 
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_head)
        if self.causal:
            mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool()
            scores = scores.masked_fill(mask, float("-inf"))
        w = F.softmax(scores, dim=-1)
        out = (w @ V).transpose(1, 2).contiguous().view(B, T, D)
        return self.W_o(out)
 
 
def demo() -> None:
    torch.manual_seed(0)
    B, T, D, N = 1, 6, 32, 4
    x = torch.randn(B, T, D)
 
    mha = MultiHeadAttention(D, N)
    out, weights = mha(x, return_weights=True)
    print(f"MHA input:  {x.shape}")
    print(f"MHA output: {out.shape}  (should match input)")
    print(f"weights:    {weights.shape}  (B, N_heads, T, T)")
 
    print("\nAttention pattern per head (row = query token, col = key token):")
    for h in range(N):
        print(f"\n  head {h}:")
        w = weights[0, h].detach().numpy()
        for i in range(T):
            row = "  ".join(f"{v:.2f}" for v in w[i])
            print(f"    t{i}: {row}")
 
    # Verify GQA has fewer params in K/V projections.
    print("\n--- GQA comparison ---")
    for n_kv in [4, 2, 1]:
        gqa = GroupedQueryAttention(D, N, n_kv)
        n_params = sum(p.numel() for p in gqa.parameters())
        print(f"GQA n_heads={N}, n_kv_heads={n_kv}: params={n_params:,}")
 
    # Verify PyTorch's builtin agrees with our MHA.
    print("\n--- Sanity check vs F.scaled_dot_product_attention ---")
    Q = mha.W_q(x).view(B, T, N, D // N).transpose(1, 2)
    K = mha.W_k(x).view(B, T, N, D // N).transpose(1, 2)
    V = mha.W_v(x).view(B, T, N, D // N).transpose(1, 2)
    builtin = F.scaled_dot_product_attention(Q, K, V)
    builtin = builtin.transpose(1, 2).contiguous().view(B, T, D)
    builtin = mha.W_o(builtin)
    diff = (out - builtin).abs().max().item()
    print(f"Max diff (ours vs builtin+W_o): {diff:.2e}")
 
 
if __name__ == "__main__":
    demo()