PyTorch手写Scaled Dot-Product Attention需三步：QKᵀ计算、除以√head_dim缩放、softmax后乘V；mask必须在softmax前用masked_fill填-inf，且输入形状为(batch, seq_len, head_dim)。
PyTorch里怎么手写一个标准的Scaled Dot-Product Attention？
PyTorch没有现成的torch.nn.ScaledDotProductAttention（直到2.0才作为实验性模块加入，且默认不启用），所以得自己用torch.bmm或torch.einsum搭。核心就三步：算QKᵀ、缩放、softmax、再乘V。
注意点在于缩放因子不是固定1/√d，必须是1 / math.sqrt(head_dim)，否则梯度会爆炸；另外mask要加在softmax前，且用float('-inf')而不是0，否则softmax会把masked位置变成非零概率。
- 输入Q/K/V形状必须是
(batch_size, seq_len, head_dim)，不能是(batch_size, head_num, seq_len, head_dim)——那是MultiHeadAttention内部拆分后的格式 - 如果要用
torch.bmm，得先transpose(1, 2)把seq_len和head_dim对齐，再bmm(q, k.transpose(-2, -1)) - 避免用
torch.matmul直接连写三层，容易因广播规则出错；推荐显式reshape+bmm或einsum('bqd,bkd->bqk', q, k)
为什么自己实现Attention时mask总不起作用？ 常见现象是训练loss不降、attention权重全均匀、甚至NaN——大概率是mask加错了位置或值设错了。
mask必须是二维的(seq_len, seq_len)（训练时）或(1, seq_len, seq_len)（推理时batch=1），且只在计算Q @ K.T之后、softmax之前应用：
- 错误做法：
attn = F.softmax(Q @ K.T + mask, dim=-1) @ V，其中mask是0/-inf但shape是(batch, seq, seq)，而Q @ K.T是(batch, seq, seq)——看起来对，但若mask是bool类型，PyTorch会自动转成0/1，-inf就丢了 - 正确做法：
attn_weights = Q @ K.T / scale; attn_weights = attn_weights.masked_fill(mask == 0, float('-inf'))，确保mask是byte/bool且fill目标明确 - Decoder自回归mask要用
torch.tril(torch.ones(...))，别手写循环填充，效率低还易索引越界
MultiHeadAttention里为什么要单独投影Q/K/V，不能共享权重？ 因为Q/K/V承担不同角色：Q代表“查询意图”，K是“键匹配依据”，V是“实际携带信息”。共享权重会让三者坍缩成同一语义空间，破坏attention的条件检索本质。
实操中容易踩的坑：
macOS UI 自动化技能包 —— 用 Python 控制你的 Mac
下载
让 macOS 自动化像说话一样简单 —— 用代码控制你的 Mac，无需学习 AppleScript 系统要求 macOS 14+ (Sonoma) Python 3.9+ 需要开启：辅助功能 + 屏幕录制权限
- 用
nn.Linear(d_model, d_model)一次性投影再切分，不如用三个独立nn.Linear——前者参数量一样，但梯度更新耦合，训练不稳定 - 投影后没做
view(..., n_heads, head_dim)和transpose(1, 2)，导致bmm维度对不上，报错mat1 and mat2 shapes cannot be multiplied - 拼接多头输出后忘了过一层
nn.Linear(d_model, d_model)（即“output projection”），结果维度对不上下游层
用torch.einsum写Attention比bmm快吗？
不绝对。小batch、短序列下einsum可读性好、不易写错维度；但大batch或长序列时，bmm底层调用cuBLAS，通常快10%~20%，且内存更可控。
关键差异在编译期优化：
-
einsum('bqd,bkd->bqk', q, k)会被PyTorch JIT尝试融合，但遇到复杂mask或动态shape可能fallback到慢路径 -
bmm要求严格二维，所以必须提前reshape和transpose，看似啰嗦，但每一步都明确可控 - 如果你开了
torch.compile()，两者性能差距会大幅缩小，但einsum的debug信息更友好（报错直接指出哪一维不匹配）
实际部署时，别迷信某一种写法。先用einsum快速验证逻辑，上线前换bmm并profile确认吞吐。最麻烦的从来不是矩阵乘本身，而是Q/K/V的shape变换和mask广播规则——这两处错一点，整个attention就静默失效，还很难定位。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
