PyTorch 中的注意力性能分析

关键信息

文中的注意力流程被拆解为 matmul、缩放、因果掩码、softmax 以及最后一次 matmul，并且文章对由这些原语组成的朴素 `NaiveCausalAttention` 模块进行了分析。trace 中出现了一个意外的内存拷贝内核，文章将其归因于 PyTorch 在 `masked_fill` 周围的非原地行为，说明看似简单的代码也可能引入额外开销。

资讯摘要

《Profiling in PyTorch》系列的目标，是让读者逐步熟悉如何阅读 profiler trace 和 profiler table。第 1 部分分析了加法、乘法等基础算术操作，第 2 部分把这些操作封装进线性层和多层感知机，并顺带分析了融合内核和手工调优内核。第 3 部分把重点转向 Transformer 的下一个核心原语：注意力机制。文章指出，注意力以二次时间复杂度著称，但也有很多技巧可以缓解这一问题并提升速度。作者并不打算逐一展开所有优化方法，而是希望展示不同实现方式在 profiler 中会呈现出怎样的差异。文中给出的示例脚本被拆成多个文件，作者使用 NVIDIA A100-SXM4-80GB GPU 在 Hugging Face 基础设施上运行这些脚本。

朴素的因果注意力模块由 `torch.matmul`、按常数缩放、用于因果掩码的 `masked_fill`、`torch.softmax` 以及最后一次与 values 的 `torch.matmul` 组成。在查看 trace 之前，文章先让读者猜测应该会看到哪些操作，然后确认 profiler 中确实出现了预期的 matmul、乘法、掩码和 softmax。展开 GPU 轨道后，还能看到一次注意力步骤实际启动了哪些内核。这里有一个意外的内核是内存拷贝。文章解释说，这来自 PyTorch 的非原地行为：有些操作不会直接修改原张量，而是先复制张量、在副本上执行操作，再返回副本。

来源与参考

1. 原始链接
2. Profiling in PyTorch (Part 3): Attention is all you profile