前面几篇把武器备齐了：链式法则的矩阵形态、微分法、VJP、softmax 雅可比。这篇打第一场硬仗——注意力机制的反向传播。注意力是三段复合：打分 $S=QK^\top/\sqrt d$、归一化 $P=\mathrm{softmax}(S)$、加权求和 $O=PV$。每一段单独看都不难，难在 softmax 夹在两个矩阵乘中间，梯度要穿它两次。推完这条链，Transformer 里所有注意力的变体（带 mask、带相对位置偏置、GQA）的反向都只剩形状调整。
前向拆解与记号
单头注意力，忽略 batch 与头维度（它们只是外层的独立循环，梯度互不耦合）。设
$$ Q\in\mathbb{R}^{n\times d},\quad K,V\in\mathbb{R}^{m\times d},\quad S=\frac{QK^\top}{\sqrt d}\in\mathbb{R}^{n\times m},\quad P=\mathrm{softmax}_{\text{row}}(S),\quad O=PV\in\mathbb{R}^{n\times d}. $$
$n$ 是 query 数、$m$ 是 key/value 数（自注意力里 $n=m$）。softmax 按行作用：$P$ 的每一行是一个概率分布。上游传来 $\bar O=\partial\ell/\partial O$（形状与 $O$ 相同），目标是 $\bar Q,\bar K,\bar V$。
第一段：O = PV 的反向
矩阵乘的梯度用微分法两秒出：$d\ell=\mathrm{tr}(\bar O^\top dO)=\mathrm{tr}(\bar O^\top(dP\cdot V+P\,dV))$，整理得
$$ \bar P=\bar O V^\top,\qquad \bar V=P^\top\bar O. $$
形状自检：$\bar O V^\top$ 是 $(n\times d)(d\times m)=n\times m$，与 $P$ 一致；$P^\top\bar O$ 是 $m\times d$，与 $V$ 一致。"梯度 = 一侧的转置乘另一侧"，这是所有矩阵乘梯度的统一形态，后面还会用两次。
$\bar V=P^\top\bar O$ 值得看成分量：$\bar V_j=\sum_i P_{ij}\bar O_i$——第 $j$ 个 value 的梯度是所有 query 对它输出的梯度按注意力权重加权的和。前向是"按权重混合 value"，反向是"按同一组权重汇集梯度"，结构对偶。
第二段：穿过 softmax
$\bar P$ 拿到后要穿过按行的 softmax 得到 $\bar S$。softmax 逐行独立，所以对每行用上一篇的结论：行向量 $p_i$、上游行向量 $\bar p_i$，VJP 为
$$ \bar s_i=J^\top\bar p_i=\Big(\mathrm{diag}(p_i)-p_ip_i^\top\Big)\bar p_i=p_i\odot\bar p_i-p_i\,\langle p_i,\bar p_i\rangle. $$
合成矩阵形式：
$$ \boxed{\;\bar S=P\odot\bar P-P\odot\big(\mathrm{rowsum}(P\odot\bar P)\,\mathbf{1}^\top\big)\;} $$
其中 $\mathrm{rowsum}$ 把 $P\odot\bar P$ 按行求和成 $n$ 维列向量，再广播回 $n\times m$。两行代码就能写完，而且注意全程没有实例化任何 $m\times m$ 的行雅可比——$O(nm)$ 收场。这里也有一个与交叉熵情形平行的观察：softmax 行和为 1，所以 $\mathrm{rowsum}(\bar S)=\mathrm{rowsum}(P\odot\bar P)-\mathrm{rowsum}(P\odot\bar P)\cdot 1=0$，$\bar S$ 每行和恒为零。这对应打分矩阵的平移自由度：给某一行 logits 同加常数不改变输出，梯度在该方向自然为零。做数值梯度检验时这是个免费的 sanity check。
第三段：S = QKᵀ/√d 的反向
$S=QK^\top/\sqrt d$ 是转置版的矩阵乘。微分：$dS=(dQ\cdot K^\top+Q\,dK^\top)/\sqrt d$，代入 $d\ell=\mathrm{tr}(\bar S^\top dS)$：
$$ \mathrm{tr}\Big(\bar S^\top\frac{dQ\,K^\top}{\sqrt d}\Big)=\mathrm{tr}\Big(\frac{K\bar S^\top}{\sqrt d}dQ\Big)\cdot(\text{取转置})\ \Rightarrow\ \bar Q=\frac{\bar S K}{\sqrt d}, $$
同理
$$ \bar K=\frac{\bar S^\top Q}{\sqrt d}. $$
（$\bar S^\top Q$ 的形状是 $(m\times n)(n\times d)=m\times d$，与 $K$ 一致。）把三段拼起来，完整反向就是四步矩阵运算：
# 前向缓存: P (n×m), Q, K, V；上游: dO (n×d)
dP = dO @ V.T # (n×m)
dV = P.T @ dO # (m×d)
dS = P * (dP - (P * dP).sum(-1, keepdim=True)) # softmax 行 VJP
dQ = dS @ K / sqrt(d) # (n×d)
dK = dS.T @ Q / sqrt(d) # (m×d)
五个矩阵乘加一次逐元素行操作，就是注意力反向的全部。写完立刻用第一章的数值梯度验收标准动作过一遍——float64、与 autograd 逐项对照：
import torch
torch.manual_seed(0)
n, m, d = 5, 7, 4 # 非方阵情形更有检验价值
Q = torch.randn(n, d, dtype=torch.float64, requires_grad=True)
K = torch.randn(m, d, dtype=torch.float64, requires_grad=True)
V = torch.randn(m, d, dtype=torch.float64, requires_grad=True)
dO = torch.randn(n, d, dtype=torch.float64)
P = torch.softmax(Q @ K.T / d**0.5, dim=-1)
O = P @ V
O.backward(dO)
# 手写反向
dP = dO @ V.detach().T
dV = P.detach().T @ dO
dS = P.detach() * (dP - (P.detach() * dP).sum(-1, keepdim=True))
dQ = dS @ K.detach() / d**0.5
dK = dS.T @ Q.detach() / d**0.5
assert torch.allclose(Q.grad, dQ, atol=1e-12)
assert torch.allclose(K.grad, dK, atol=1e-12)
assert torch.allclose(V.grad, dV, atol=1e-12)
# 免费自查：S̄ 每行和为零（logits 行平移自由度）
assert torch.allclose(dS.sum(-1), torch.zeros(n, dtype=torch.float64), atol=1e-12)
全部通过。注意这里故意用 $n\neq m$ 的非方阵——方阵下转置错位的梯度形状碰巧也对，只有数值检验能抓；非方阵时形状自检就直接报警了，两类防线各演示一次。
causal mask 在反向里做什么
decoder 的因果掩码在前向把 $S_{ij}$（$j>i$）置为 $-\infty$，softmax 后对应 $P_{ij}=0$。反向不需要单独处理：$\bar S=P\odot(\cdot)$ 里的因子 $P$ 已经把被掩位置的梯度自动清零——未来位置收不到任何梯度，信息泄漏在前向和反向两个方向都被堵死。实现上只需保证前向掩码正确，反向跟着 $P$ 走即可。这也是"缓存 $P$"成为注意力反向必要开销的原因之一：掩码信息、归一化信息全在它里面。
FlashAttention 的反向：重算而不是缓存
标准实现要缓存 $P\in\mathbb{R}^{n\times m}$，对长序列这是显存灾难（$n=m=128\text{k}$ 时一个头一层就是 64GB 量级）。FlashAttention 的做法是前向不存 $P$，只存输出 $O$ 和每行的 logsumexp 统计量 $\ell_i=\log\sum_j e^{S_{ij}}$（在线 softmax 的副产品）；反向时按块重算 $S_{ij}=Q_iK_j^\top/\sqrt d$，用缓存的 $\ell_i$ 一行还原 $P_{ij}=e^{S_{ij}-\ell_i}$，然后套用上面的公式在块内累加 $\bar Q,\bar K,\bar V$。
代价账很清楚：反向多算一次 $QK^\top$（重算 $S$），总 FLOPs 从 4 GEMM 涨到约 5 个，但 $P$ 的 $O(nm)$ 显存读写全省了。注意力是访存受限算子，用 25% 的计算换数量级的显存访问下降，这笔交易在 HBM 带宽是瓶颈的硬件上稳赚。反向公式本身一行没变——FlashAttention 改的是存储与重算的调度，不是数学。
多头与 GQA 的形状注记
多头情形下 $Q,K,V$ 按头切开独立算，梯度按头拼接，没有交叉项。GQA/MQA 里多个 query 头共享一组 $K,V$，反向时 $\bar K,\bar V$ 要把共享组内所有 query 头的贡献求和——又是一次广播梯度归约（共享即广播，广播反向即求和，同一原则第三次出现了）。PyTorch 里 enable_gqa=True 的 SDPA 内部就是先 repeat_interleave 展开再按组求和梯度，数学上没有任何新东西。
小结
注意力反向 = 三段 VJP 串接：加权求和段给出 $\bar P=\bar OV^\top$、$\bar V=P^\top\bar O$；softmax 段做行内 VJP $\bar S=P\odot\bar P-P\,\mathrm{rowsum}(P\odot\bar P)^\top$ 广播版；打分段给出 $\bar Q=\bar SK/\sqrt d$、$\bar K=\bar S^\top Q/\sqrt d$。免费自查有两个：$\bar S$ 行和为零，$\bar Q/\bar K/\bar V$ 形状各自等于参数形状。推导用到的每一件工具都来自前几篇，这正是把矩阵求导基础打牢的回报——第四章讲 FlashAttention、MLA 时，反向的数学部分已经全部提前结清了。下一篇换个节奏，把常用梯度公式汇总成一张速查表，每条配一行微分推导。