手写SelfAttention层需四步：线性投影得Q/K/V；缩放点积Q@Kᵀ/√dₖ；mask填-inf后softmax；加权求和再经Wₒ输出，严格控维、控dtype、控mask位置。
PyTorch里怎么手写SelfAttention层（不依赖nn.MultiheadAttention）
直接实现最简自注意力，能帮你真正理解Q、K、V的计算逻辑和缩放因子的作用。别急着套用封装好的模块——很多 bug 出在维度没对齐或mask应用位置错。
关键点：输入x形状必须是(batch_size, seq_len, embed_dim)；W_q、W_k、W_v都是(embed_dim, head_dim)；算Q @ K.T前要除以sqrt(head_dim)，否则 softmax 容易饱和。
- 用
torch.bmm做 batch 内矩阵乘（不能直接@，维度会错） - 如果要做 causal attention，mask 要在 softmax 前加，用
torch.finfo(torch.float32).min填无效位置 - 输出要再过一个线性层
W_o（(head_dim, embed_dim)），不然维度接不上后续层
nn.MultiheadAttention
- 检查输入 shape：若你习惯
(B, S, E)，就显式设batch_first=True，且确保所有前后层统一 - 注意
key_padding_mask形状是(B, S)，而attn_mask（用于 causal 或自定义遮罩）是(S, S)或(B, S, S) - 该模块不自动做 LayerNorm 或残差连接，得自己包一层
TransformerEncoderLayer才完整
常见报错：RuntimeError: mat1 and mat2 shapes cannot be multiplied
90% 出现在Q @ K.T这步。不是模型写错了，而是输入x的embed_dim和线性层的in_features不一致，或者batch_first开关和实际数据 shape 冲突。
- 打印
x.shape和self.W_q.weight.shape，确认第二维匹配 - 如果用了
nn.Embedding，注意它输出是(B, S, E)，但 embedding 层的num_embeddings和embedding_dim别和 attention 的embed_dim搞混 - GPU 上调试时，
torch.cuda.synchronize()后加断点，避免异步报错定位偏移
要不要自己实现，还是直接用nn.MultiheadAttention
除非你在调参、做消融实验，或需要定制梯度（比如让attn_weights参与 loss），否则别手写。PyTorch 版本 ≥1.10 后，nn.MultiheadAttention已支持 FlashAttention 风格优化（通过enable_math=False等 flag），速度比手写快，且自动处理了多头拆分、concat、dropout 等细节。
- 想兼容旧版 PyTorch（
- 训练时发现 attention 权重异常（全 0 或全 1），优先检查
scale因子和mask是否生效，而不是重写整个模块 - 导出 ONNX 时，手写 attention 更容易适配，但
nn.MultiheadAttention从 1.12 起也基本稳定了
维度对齐永远比公式推导更耗时间，尤其是mask的广播行为和bmm的 batch 维位置——这些地方错一点，结果就全崩。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
