主题
第 05 章 多头注意力的实现选择与 buffer 生命周期
同一套注意力公式，可以用逐头 Python 循环实现，也可以写成一次性矩阵乘法、einsum，或交给 SDPA、完整的 MultiheadAttention 类和 FlexAttention。代码看起来差别很大，算的却未必是不同的东西。比较速度之前，先要弄清楚：哪些改动只调整了张量布局，哪些已经改变了参数量或功能。
本章沿着这个问题，阅读上游的高效注意力 notebook 和 buffer notebook：前者帮助比较实现，后者解释计算所需的张量该怎样随模型迁移、保存和恢复。上游的大尺寸计时没有在本书中运行；文中引用的本书实验记录都来自小型 CPU 张量，也不是本次文字修订的新实测。实现集合
5.1 把三次投影合成一次，不改变线性代数
先看最容易对齐的一种改法。若 Q、K、V 都是
Linear(D,3D) 的输出是 [B,T,3D]，再拆成三个 [B,T,D]。有偏置时，偏置也要按同样顺序拼接。原来分开调用的三次投影因此合到了一起，可以减少部分调用和读写开销；但权重里的数并没有变少，主要的矩阵乘法量也还在。
上游 combined-QKV 版本先 view(B,T,3,H,d)，再 permute(2,0,3,1,4)，得到 [3,B,H,T,d]，最后 unbind(0)。本书采用先 chunk 再分头的顺序，得到的布局相同。核对时不能只数元素总量：reshape 的轴顺序一旦弄错，Q 的某些通道就可能被分到 K 中，形状检查却仍然通过。
这里还有一个容易从示例里漏看的条件：上游类最终按输入的 embed_dim 合并输出，而示例恰好满足 d_in=d_out。如果输入宽度与输出宽度不同，合并时应使用 d_out。把这个类用于其他尺寸前，需要先去掉对输入、输出等宽的假定。
5.2 einsum 让索引关系写在代码里
上游 MHAEinsum 用 bnd,do->bno 表示投影：批次 bhnd,bhmd->bhnm，对每头特征 bhnm,bhmd->bhnd，对键位置
这些索引把第 4 章的求和关系直接写进了代码，但写成 einsum 并不保证更快。乘法怎样执行、内存怎样访问，仍取决于张量形状和底层实现。
数值对照还要先处理权重布局：上游用 Parameter 保存 [d_in,d_out] 权重，并自定义初始化；Linear 的布局恰好相反。fan-in 表示每个输出接收的输入数量，它的解释也与轴布局有关。因此，只固定随机 seed 不足以让两版拿到对应的权重，比较前必须显式复制参数。
如果启用 QKV 偏置，einsum 版在投影之后加 [d_out] 向量；禁用时用 register_parameter(name,None)。这保留了统一接口，同时不把不存在的偏置计入参数量。
5.3 SDPA 的 mask 与 dropout 约定
scaled_dot_product_attention 接受已经分头的 Q、K、V，内部完成缩放、掩码、softmax 和加权汇总。使用 is_causal=True 可以表达普通方形因果注意力。本书对照测试的数学式与该接口在 CPU 上一致。
换接口时，尤其要核对布尔 mask 中 True 的含义。手写 masked_fill(mask,-inf) 用 True 表示“禁止访问”；SDPA 的布尔 attn_mask 却用 True 表示“允许访问”，所以上游显式 mask 版本传入的是 ~attn_mask。到了 nn.MultiheadAttention，True 又表示禁止。同一个变量不能不经检查就传给不同接口，否则允许访问的位置可能全部反过来。
dropout 也需要调用者明确控制。函数式 SDPA 的 dropout_p 不会根据外层模块的 eval() 自动归零，所以上游写成 0. if not self.training else self.dropout。漏掉这个判断，评估时仍可能随机丢弃注意力权重；保存恢复或 cache 等价检查也会因此受到干扰。
上游把一个显式 mask 变体命名为 WithoutFlash。阅读时可以用它定位作者比较的代码路径，但不能据此认定任何 PyTorch 版本、任何设备都会禁用某个内核。后端选择还取决于数据类型、形状、mask 支持和运行库。要确认实际执行了哪个内核，需要后端控制或 profiler 记录；本书的 CPU 数值等价结果不能充当 FlashAttention 已运行的证据。
5.4 完整 MultiheadAttention 类不只是公式包装
上游 MHAPyTorchClass 内部使用 nn.MultiheadAttention(batch_first=True)，输入保持 [B,T,D]，把同一 x 作为 query、key、value。设置 need_weights=False 可以避免请求返回注意力权重，也可能开启更合适的执行路径；是否加速要测量。
这个包装类还有两处功能差异，会直接影响比较。它先经过 nn.MultiheadAttention 自带的输出投影，再经过额外的一层 proj，因此参数量未必与只有一层输出投影的手写版本相同。若 qkv_bias=True，代码还会打开 add_bias_kv，引入额外的键值偏置项；这与普通 QKV 线性层的偏置开关并不完全等价。
如果要比较同一计算的速度，就先消除这些差异。若保留它们，也应明确说明比较的是功能不同的模块。
5.5 FlexAttention 表达规则，执行条件另算
上游 causal(b,h,q_idx,kv_idx) 返回 q_idx >= kv_idx。create_block_mask 把这个位置规则组织成 block mask，再交给 flex_attention。与显式存一个完整三角矩阵相比，这种接口更方便表达复杂可见范围；但它的可用设备、版本和编译支持需要独立核验。
读到这里，还需要把接口表达能力与实际运行条件分开。固定版本 notebook 会先去掉 PyTorch 版本号的构建后缀，再比较最低版本，并检查 CUDA 是否可用，以决定是否导入和实例化 FlexAttention。不过，后面某些独立的 %timeit 单元没有同样的保护；在 CPU 环境依次运行全部单元，仍可能引用尚未定义的对象。兼容性注释只能作为线索，运行前还得检查实际安装版本。
这个实现虽然接收 dropout 参数，当前 forward 却没有把它用于 FlexAttention 路径。示例取值为 0，所以这里没有差异；改成非零后，就不能假定它具有与其他版本相同的训练正则化效果。新版分块掩码接口变化后，上游也改为按当前序列长度重新构造 mask，不再切片旧对象。本书没有执行 FlexAttention 或 GPU 编译测试，这些实现细节来自源码阅读。
5.6 buffer 解决的是模块状态管理
前面一直在比较怎样算注意力，接下来要处理一个不在公式里的问题：计算所需的张量由谁管理？因果 mask 不需要优化器更新，却必须与参与计算的张量处在合适的设备上。如果只写 self.mask = tensor，它就是普通属性，model.to(device) 不会自动管理它的迁移，state_dict() 也不会保存它。于是，权重可能已经到了 GPU，mask 却还留在 CPU，直到前向计算时才报错。
register_buffer("mask", tensor) 会把这个张量交给模块管理。默认的持久 buffer 随模型迁移设备，也进入 state_dict，但不会出现在 named_parameters() 中。对于可以重建、不想写入检查点的状态，可以设为 persistent=False；这样仍会迁移设备，只是不保存。要不要保存一个状态，与要不要训练它，是两个独立的决定。
上游 buffer notebook 分别构造无 buffer 和有 buffer 的 CausalAttention，先观察设备，再手动迁移普通属性，最后把 mask 中的 1 改成 2 并保存恢复。由于前向会把 mask 转布尔，1 和 2 都表示禁止，修改是为了检查状态是否被保存，而不是改变因果含义。无 buffer 的新对象会重新得到初始 mask；有 buffer 的对象会加载修改后的值。buffer 实验
本书直接保存布尔 mask，省去每次 .bool() 转换，也不必用浮点张量存储。已有实验记录确认：future 在 state_dict 中、不在参数字典中，dtype 为 bool。
缓存也可以注册成非持久 buffer，但不能照搬静态 mask 的生命周期。mask 描述固定的可见规则，缓存保存的却是某次请求的内容。新请求开始时必须清空缓存，batch、设备或 dtype 改变时也要相应处理。
5.7 怎样测量而不混淆问题
实现和状态管理对齐以后，计时结果才容易解释。上游分别比较前向、前后向，以及 torch.compile 后的前后向。CUDA 计时函数先预热，再记录 CUDA Event，结束后同步，最后统计均值和标准差；绘图函数据此画出带误差条的柱状图。GPU 执行是异步的，这种做法比直接用 Python 时钟包住任务提交，更接近要测的设备执行时间。
同一张计时表里的实验，还应统一是否记录计算图、是否保留输出，以及反向前是否清空参数梯度。上游 forward_backward 会清理输入梯度，却没有在每次重复前清理各模块的参数梯度，因此参数梯度会累加。这能测量重复 backward 的耗时，但没有覆盖完整优化器步骤，不能据此报告端到端训练吞吐。
编译时间也要单独看。torch.compile 的首次编译与后续稳定运行不应混在一起；设置 suppress_errors 后还允许回退，仅仅“不报错”并不能证明整段图都已经编译。
set_float32_matmul_precision("high") 是精度/性能选择，不代表任何 CPU 都拥有或启用了 Tensor Core。应记录设备、dtype、批量、长度、宽度、头数、模式、是否反向、是否同步，以及编译状态。内核优化还可能改变浮点求和顺序，所以性能检查应伴随误差和梯度检查。
5.8 一张比较表，明确各自代价
| 实现 | 主要变化 | 容易遗漏的条件 | 本书验证 |
|---|---|---|---|
| 多头 wrapper | 每头独立模块，拼接 | d_out 是每头宽度 | 解释与形状推导 |
| 拆头 MHA | 一次大投影后 reshape | 总宽度能被头数整除 | 微型 GPT 对照 |
| 合并 QKV | 三投影合成一投影 | QKV 顺序、偏置、输入输出宽度 | 参数映射对照 |
| einsum | 显式写求和索引 | 初始化与权重布局 | 三组缩小尺寸对照 |
| SDPA | 库实现核心注意力 | mask 极性、dropout、后端 | CPU 前向和梯度 |
| MultiheadAttention | 更完整的层接口 | 额外投影、返回权重、bias-kv | 静态解释 |
| FlexAttention | 用规则构造块掩码 | 版本、设备、未应用的 dropout | 未运行硬件路径 |
5.9 练习与答案
先用三个问题检查是否分清了接口、状态和比较条件。第一题：给 masked_fill 的禁止 mask 是严格上三角，改用 SDPA 时应传什么？看清 True 的含义后，答案就是取反后的允许 mask；同时要避免传入与显式 mask 不相容的 causal 参数组合。
第二题：模型 .to(float64) 后，一个普通浮点属性和一个浮点 buffer 会怎样？答案：注册的参数和浮点 buffer 跟随转换，普通属性不会自动受管理；布尔 buffer 仍应保持布尔类型。
第三题：为什么两版 MHA 都输出 [2,5,12]，仍可能不能公平比速度？答案：输出形状没有说明参数量、额外输出投影、dropout、返回权重和反向功能是否相同。先对齐这些条件，再复制权重和比较数值。
上游 tests/test_mha_implementations.py 会显式复制并转置三份权重，在输入输出等宽、变宽、变窄三组尺寸下做对照，还检查显式 SDPA mask 是否阻止未来信息泄漏。本书原稿记载了更小的三组尺寸复查：12→12、8→12、12→8，float64 最大误差均小于
这组历史记录需要保留一个核验边界：当前配套目录能找到生成该结果的检查脚本，却未包含所述 JSON 结果文件。本次文字修订没有重跑该实验，因此这里保留的是原稿的记录，不能视作已由当前交付文件独立复核。
主检查命令是 python examples/core_training/run_checks.py。已保存的 results.json 记录了 SDPA 最大前向误差
参考与覆盖
- 高效 MHA 完整 notebook：九种比较路径、版本检查、计时、绘图、编译包装。
- buffer 生命周期 notebook：普通属性、注册状态、设备迁移与保存恢复。
- 本章没有复用上游图、没有运行其大张量 CUDA 计时，性能断言与数学一致性检查已分开。