前面多篇把注意力的每个组件拆开了讲，这一篇把它们装回去——不用任何深度学习框架，纯 NumPy 写一个完整的因果多头自注意力，再和 PyTorch 的 F.scaled_dot_product_attention 对拍误差。要求的先验只有：会矩阵乘法、知道 softmax 和 mask 是什么。
第一步：缩放点积注意力
先写核心公式 $\mathrm{softmax}(QK^\top/\sqrt{d_k})V$。三个容易写错的点先立住：缩放因子除的是每个头的维度 $d_k$ 的平方根（不是模型维度）；softmax 要做数值稳定化（先减行最大值，否则大 logits 下 exp 溢出）；mask 要在 softmax 之前把非法位置置为 $-\infty$（softmax 之后置零是错的，归一化分母已经把它们算进去了）。
import numpy as np
def softmax(x, axis=-1):
x = x - x.max(axis=axis, keepdims=True) # 数值稳定
e = np.exp(x)
return e / e.sum(axis=axis, keepdims=True)
def sdpa(Q, K, V, mask=None):
# Q: (n, d_k) K: (m, d_k) V: (m, d_v)
S = Q @ K.T / np.sqrt(Q.shape[-1]) # (n, m) 打分
if mask is not None: # mask: (n, m)，True = 允许
S = np.where(mask, S, -np.inf)
P = softmax(S) # (n, m) 行归一
return P @ V # (n, d_v)
五行核心代码，前面讲的所有理论都能在这里对上号：S 就是会爆显存的 $n\times n$ 矩阵；-np.inf 那些行在 softmax 里指数为零——这就是"掩码位置拿不到梯度"的代码形态。
第二步：QKV 投影与因果掩码
自注意力的 Q/K/V 不是输入本身，是三个投影。设输入 $X\in\mathbb{R}^{n\times d}$（这里跳过训练，投影矩阵用固定随机值代替），$W^Q, W^K\in\mathbb{R}^{d\times d_k}$，$W^V\in\mathbb{R}^{d\times d_v}$。
rng = np.random.default_rng(0)
n, d, d_k, d_v = 8, 16, 4, 6
X = rng.normal(size=(n, d))
Wq = rng.normal(size=(d, d_k))
Wk = rng.normal(size=(d, d_k))
Wv = rng.normal(size=(d, d_v))
Q, K, V = X @ Wq, X @ Wk, X @ Wv
causal = np.tril(np.ones((n, n), dtype=bool)) # 下三角：位置 i 只能看 ≤ i
O = sdpa(Q, K, V, mask=causal)
因果掩码就是单位下三角矩阵的逐元素版本——np.tril 生成它的时候，"信息不能从未来流回过去"这条约束就从一个抽象原则变成了一张布尔表。验证掩码正确性的最快办法：把最后一行的 attention 分布打出来，最后一行是全可见的，它的权重应该和"不带 mask 直接 softmax"一致：
assert np.allclose(O[-1], sdpa(Q, K, V)[-1]) # 最后一行不受因果掩码影响
第三步：多头重组
多头不是循环，是一次 reshape。$n_h$ 个头各自在 $d_k$ 维子空间做注意力，实现上把 $d = n_h\cdot d_k$ 的通道维折叠：
n_h = 4
Wh = rng.normal(size=(d, d)) # 简化为 n_h*d_k = d
Wq_h = rng.normal(size=(d, n_h * d_k)) # 每头独立：输出 n_h*d_k 维
Wk_h = rng.normal(size=(d, n_h * d_k))
Xh = X @ Wh # (n, d)
Qh = (Xh @ Wq_h).reshape(n, n_h, d_k).transpose(1, 0, 2) # (n_h, n, d_k)
Kh = (Xh @ Wk_h).reshape(n, n_h, d_k).transpose(1, 0, 2)
Wv_h = rng.normal(size=(d, n_h * d_v)) # V 要输出 n_h*d_v 维再按头切
Vh = (Xh @ Wv_h).reshape(n, n_h, d_v).transpose(1, 0, 2)
O_h = np.stack([sdpa(Qh[h], Kh[h], Vh[h], causal) for h in range(n_h)]) # (n_h, n, d_v)
O = O_h.transpose(1, 0, 2).reshape(n, n_h * d_v) # 拼回头维
两个 reshape 是整个多头机制的全部秘密：第一个把"按头分段"的通道维挪到最前面变成 batch 维，逐头算注意力就是一次循环（或 einsum）；第二个把结果拼回 $n\times(n_h d_v)$。概念上"每个头有独立的 K/V 空间"，代码上只是 reshape 前后谁挨着谁。真正的独立投影（每头一份 $W^Q_h$）写成代码也只是把 Wq 从 $(d, d_k)$ 换成 $(d, n_h d_k)$。
对拍：和手算、和 PyTorch
三层验证，每层抓不同的 bug：
# 1) 手写 softmax 对拍 scipy 无关实现：逐行和为 1
P = softmax(Q @ K.T / np.sqrt(d_k))
assert np.allclose(P.sum(-1), 1.0)
# 2) 数值稳定：把 logits 放大 1e4 倍，naive 版本溢出，稳定版仍有限且归一
x = Q @ K.T / np.sqrt(d_k) * 1e4
naive = np.exp(x); naive = naive / naive.sum(-1, keepdims=True)
stable = softmax(x)
assert np.isnan(naive).any() and np.isfinite(stable).all()
assert np.allclose(stable.sum(-1), 1.0)
# 3) 与 PyTorch 的 F.scaled_dot_product_attention 对拍（float64 下应到 1e-12）
import torch
ref = torch.nn.functional.scaled_dot_product_attention(
torch.tensor(Qh), torch.tensor(Kh), torch.tensor(Vh), is_causal=True)
assert np.allclose(O_h, ref.numpy(), atol=1e-12)
第 2 条值得展开：$10^4$ 倍放大后 logits 量级在数万，naive softmax 的 exp 直接溢出成 inf 再算出 nan；减行 max 的版本全程有限。这也是混合精度训练里一切 loss scale 机制要解决的问题的迷你版。第 3 条用 is_causal=True 顺带验证了我们的下三角掩码语义和 PyTorch 一致（PyTorch 的 is_causal 等价于下三角为 1 的掩码）。
小结
60 行 NumPy 复现自注意力的意义不在代码本身，而在三个"概念到形状"的对照被亲手验证过：掩码是 softmax 前的布尔表、多头是两次 reshape、数值稳定是减一行 max。写完后回头看框架里的 nn.MultiheadAttention，它不再是一个黑盒，而是这 60 行经过向量化、融合、 fused kernel 之后的工业形态——下一篇就拆那个工业形态。