mha.py

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

"""
Multi-Head Attention - minimal runnable implementation
======================================================

Educational PyTorch reference for Multi-Head Self-Attention.
Standalone script: runs three sanity checks on CPU in a few seconds.

Pairs with: docs/tutorials/attention_tutorial.md (concept reference).

Key shapes:
    Input  x: [B, N, C]
    qkv:      [B, N, 3*C] -> reshape -> [3, B, H, N, d_k]
    score:    [B, H, N, N]
    Output:   [B, N, C]

Run:
    python mha.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.

    - Single fused Linear for Q/K/V (kernel-fusion friendly).
    - Optional additive attention mask (float, -inf at masked positions).
    - No KV-cache; meant for educational clarity, not inference speed.
    """

    def __init__(self, embed_dim: int, num_heads: int, qkv_bias: bool = False):
        super().__init__()
        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = 1.0 / math.sqrt(self.head_dim)

        # One Linear produces Q, K, V at once.
        self.qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=qkv_bias)
        # Output projection bias follows qkv_bias so the sanity check below can
        # align cleanly with torch.nn.MultiheadAttention (which has no
        # out_proj.bias when bias=False).
        self.proj = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)

    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor | None = None):
        """
        Args:
            x: [B, N, C]
            attn_mask: [N, N] additive mask (float, -inf at masked positions)
        Returns:
            [B, N, C]
        """
        B, N, C = x.shape

        # qkv: [B, N, 3*C] -> [B, N, 3, H, d_k] -> [3, B, H, N, d_k]
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]            # each: [B, H, N, d_k]

        # score: [B, H, N, N]
        score = (q @ k.transpose(-2, -1)) * self.scale
        if attn_mask is not None:
            score = score + attn_mask                # additive mask, broadcasts
        attn = F.softmax(score, dim=-1)

        # out: [B, H, N, d_k] -> [B, N, H, d_k] -> [B, N, C]
        out = (attn @ v).transpose(1, 2).reshape(B, N, C)
        return self.proj(out)


def make_causal_mask(n: int, device: str | torch.device = "cpu") -> torch.Tensor:
    """Causal mask of shape [N, N]: upper triangle filled with -inf."""
    mask = torch.zeros(n, n, device=device)
    mask = mask.masked_fill(
        torch.triu(torch.ones(n, n, device=device), diagonal=1).bool(),
        float("-inf"),
    )
    return mask


def sanity_check():
    """Verify alignment with torch.nn.MultiheadAttention (no mask, no bias)."""
    torch.manual_seed(0)
    B, N, C, H = 2, 17, 64, 4
    x = torch.randn(B, N, C)

    mine = MultiHeadAttention(C, H, qkv_bias=False)
    # bias=False so builtin.in_proj_bias / out_proj.bias are both None.
    builtin = nn.MultiheadAttention(C, H, bias=False, batch_first=True)

    # Copy our qkv + proj weights into the builtin (in_proj_weight = [3*C, C]).
    with torch.no_grad():
        builtin.in_proj_weight.copy_(mine.qkv.weight)
        builtin.out_proj.weight.copy_(mine.proj.weight)
        # Note: bias=False -> builtin.out_proj.bias is None, and mine.proj also
        # has no bias because we passed qkv_bias=False (see __init__).

    y_mine = mine(x)
    y_builtin, _ = builtin(x, x, x, need_weights=False)
    diff = (y_mine - y_builtin).abs().max().item()
    print(f"[sanity] max diff vs nn.MultiheadAttention = {diff:.2e}")
    assert diff < 1e-5, f"mismatch! diff={diff}"
    print("[sanity] PASS")


def shape_check():
    """Common transformer shape: B=2, N=1024 (32x32 tokens), C=768, H=12, d_k=64."""
    B, N, C, H = 2, 1024, 768, 12
    x = torch.randn(B, N, C)
    mha = MultiHeadAttention(C, H)
    y = mha(x)
    print(f"[shape] input  x: {tuple(x.shape)}")
    print(f"[shape] output y: {tuple(y.shape)}")
    assert y.shape == x.shape


def causal_check():
    """Causal mask sanity check for autoregressive transformers."""
    B, N, C, H = 1, 8, 32, 4
    x = torch.randn(B, N, C)
    mha = MultiHeadAttention(C, H)
    mask = make_causal_mask(N)
    y = mha(x, attn_mask=mask)
    print(f"[causal] applied causal mask, output shape {tuple(y.shape)}")


if __name__ == "__main__":
    print("== Multi-Head Attention sanity ==")
    sanity_check()
    shape_check()
    causal_check()
    print("\nAll checks passed.")