PyTorch计算效率：含零块矩阵的多头注意力实现方案速度对比

常规模型场景下，用for循环逐头计算注意力、最后拼接的方案，速度会大幅快于构造含大量零元素的块对角大矩阵的方案，二者的性能差通常能达到数倍甚至一个数量级——你以为能靠去掉循环、省掉拼接提速，实际上零块矩阵带来的冗余计算、冗余显存开销，才是真正的性能大头。

核心开销差异来源 #

• 块对角矩阵的稠密计算冗余
  你构造的带零块的大矩阵本质是稠密张量，PyTorch底层调用的BLAS矩阵乘算子，会严格按照张量的完整维度调度所有乘加计算，完全不会因为矩阵中的元素是0就跳过对应的计算单元。
  你提到的块对角矩阵结构，实际在内存中就是如下的稠密形式：

# 块对角矩阵的实际存储结构
block_diag = [
    [A, 0, 0],
    [0, B, 0],
    [0, 0, C]
]

举个实际量化的例子：假设模型有3个注意力头，每个头的特征维度是128，那么单头的投影计算量是seq_len * 128 * 128，3个头逐头计算的总浮点操作量是3倍的单头量；但如果把3个头的权重堆成3*3的块对角大矩阵，整个权重矩阵的维度会直接变成384 * 384，对应的单次矩阵乘浮点操作量是seq_len * 384 * 384，是逐头计算方案的9倍——这部分完全无效的乘加计算，根本不是省掉拼接操作能抵消的。

• 梯度计算的冗余问题
  你之前判断核心性能差异取决于“PyTorch是否会对乘零的张量部分执行梯度计算”，这个方向完全正确，但实际结果会比预期差很多：PyTorch根本不会跳过乘零部分的梯度计算。它的自动求导引擎会严格按照前向传播构建的计算图形状，为所有参与运算的张量生成形状完全匹配的梯度——哪怕前向计算中某部分的输入是全0，对应的反向计算依然会在全张量范围内执行，既不会自动剪枝无意义的乘0操作，也不会自动压缩梯度的存储形状。也就是说用块对角矩阵的方案，不仅前向多算了数倍的无效乘加，反向传播的时候还要多分配数倍的显存存储无效梯度值，整体显存占用也会远高于逐头实现。
• 待优化操作的实际开销远低于预期
  逐头计算后做的拼接操作，本质是显存上的连续内存拷贝，时间复杂度和最终输出的元素规模成正比，在常见的注意力输出维度下，这点拷贝的耗时通常只占整个注意力计算总耗时的1%都不到；哪怕是Python层面的for循环调度开销，在头数为2~16的常规模型场景下，相比矩阵乘的计算耗时也基本可以忽略。

补充说明 #

PyTorch官方的多头注意力实现之所以速度更快，既不是靠for循环逐头算，也不是靠构造带零块的对角矩阵，而是直接把投影权重reshape成带“头数”维度的形状，通过批量矩阵乘（bmm）把所有头的计算打包成批次维度的并行计算，全程没有任何零元素的冗余计算，也不需要额外的拼接操作，本质是在无无效计算的前提下一次性完成所有头的运算，和上面两种自行实现的方案都有本质区别。

内容的提问来源于stack exchange，提问作者financial_physician