本文演示了如何使用剖析器跟踪在 PyTorch 中对注意力机制进行剖析，以识别性能瓶颈和优化机会。它将朴素注意力实现与 PyTorch 内置的缩放点积注意力 (SDPA) 进行比较，以说明内核行为。

  • 朴素的因果注意力涉及多个内核，包括 matmul、缩放、掩码、softmax 以及由非原地操作引起的意外内存复制。
  • 用原地 masked_fill_ 替换 masked_fill 可从 GPU 跟踪中移除不必要的内存复制内核。
  • PyTorch 的 SDPA 函数会根据输入类型和硬件自动分派到最快的可用后端（如 Flash Attention 或 cuDNN）。
  • 对各个 SDPA 后端进行剖析会揭示不同的内核签名，使开发人员能够验证哪个实现处于活动状态。

理解这些剖析器跟踪有助于用户通过消除冗余内存操作并利用 PyTorch 高效的注意力后端来优化 Transformer 模型。