PyTorch性能分析第三部分：注意力即剖析
速览
PyTorch官方推出性能分析系列第三篇，聚焦注意力机制的性能剖析方法。文章以Transformer模型为例，演示如何使用PyTorch Profiler工具定位注意力计算瓶颈，并给出优化建议。这对于提升大模型训练效率具有实际指导意义。
AI 深度解读
背景
“Profiling in PyTorch” 系列旨在帮助开发者掌握读取性能分析器（profiler）跟踪轨迹（traces）和表格的能力，并以此驱动优化。在系列的第一部分中，我们分析了加法、乘法等基本数学运算的 profiling，展示了 profiler 表格如何揭示热点，以及 trace 如何展示算法随时间的执行顺序。第二部分则将加法和乘法封装进 torch.nn.Linear 层，接着堆叠多个线性层组成多层感知机（MLP）并进行 profiling，同时分析了融合内核（fused kernel）和手工调优内核。
从 Transformer 架构的角度来看，下一个合乎逻辑的 profiling 对象是另一个基础算法——注意力机制（Attention）。尽管注意力因其二次时间复杂度而“臭名昭著”，但存在许多巧妙技巧来缓解这一问题并使其变快。本文的目标不是详述每个技巧，而是观察每种技巧在 profiler 下呈现的不同形态。
所有脚本托管在相应仓库中，包含四个文件：04_a_naive_attention.py、04_b_inplace_ops_attention.py、04_c_sdpa_attention.py 和 04_d_kernels_attention.py。建议在阅读时打开脚本并对照代码。实验使用 NVIDIA A100-SXM4-80GB GPU，可以在 Hugging Face 基础设施上方便地设置 GPU，并通过 Spaces 的 Dev Mode 运行脚本，或使用 Hugging Face Jobs 管道。
核心内容
朴素注意力（Naive Attention）
注意力机制涉及查询（Query, q）、键（Key, k）和值（Value, v）。它们之间的交互可写为一组简短的步骤：
- 构建注意力得分
scores：matmul(q, k.T) - 缩放得分：
scores * scale - 应用因果掩码：
scores.masked_fill(mask, "-inf") - 用 softmax 归一化得到注意力权重
attn：softmax(scores) - 用权重重新加权值：
matmul(attn, v)
因此，注意力实际上是若干基本操作的集合，其中矩阵乘法我们已经熟悉，其余操作也容易辨认。下面用 PyTorch 编写一个朴素因果注意力模块并对其 profiling：
class NaiveCausalAttention(nn.Module):
def __init__(self, head_dim):
super().__init__()
self.scale = 1.0 / math.sqrt(head_dim)
def forward(self, q, k, v, mask):
scores = torch.matmul(q, k.transpose(-2, -1))
scores = scores * self.scale
scores = scores.masked_fill(mask, float("-inf"))
attn = torch.softmax(scores, dim=-1)
out = torch.matmul(attn, v)
return out
在打开 trace 之前，我们做常规的预测练习。跟踪 forward 过程，预期看到：
- 一个 matmul 内核（
q · k.T） - 一个 mul 内核（缩放）
- 一个掩码操作
- 一个 softmax 内核
- 一个 matmul 内核（
attn · v）
运行脚本并可视化 trace（命令：uv run 04_a_naive_attention.py 然后 uvx trace-util -f traces/ -b <hf_uname>/traces）。图 1 展示了 CPU 通道（GPU 通道被折叠）。在 attn_fwd（我们标注的前向调用）内部，正好看到我们猜测的那些操作：matmul、mul、masked_fill、softmax。展开 GPU 通道后，可以看到实际启动的内核序列（图 2）。放大单个 attn_fwd 块（图 3），逐个识别内核：
- matmul（query 和 key）
- mul（缩放）
- 内存拷贝（Memcpy）??
- 因果掩码
- softmax（产生注意力权重）
- matmul（注意力权重和值）
前五个都在预期之内，但内存拷贝是意料之外的。它从何而来？线索是 PyTorch 本身有 in-place 操作。当你以普通（out-of-place）方式操作张量时，PyTorch 经常制作一份副本，对副本应用操作，然后返回副本。根据操作顺序，罪魁祸首是 masked_fill。如果将其替换为 in-place 操作呢？
带 In-place 掩码的朴素注意力
只需将 masked_fill 改为 masked_fill_（注意下划线，这是 PyTorch 中 in-place 操作的约定），并运行相同脚本。
def forward(self, q, k, v, mask):
scores = torch.matmul(q, k.transpose(-2, -1))
scores = torch.mul(scores, self.scale)
scores.masked_fill_(mask, float("-inf")) # 改为 in-place
attn = torch.softmax(scores, dim=-1)
out = torch.matmul(attn, v)
return out
查看 trace（图 4 与图 5对比）：in-place 版本在掩码步骤内部包裹的 CPU 操作远少于 out-of-place 版本，这是一个令人鼓舞的信号。展开 GPU 通道（图 6 和图 7），可以看到 Memcpy 内核彻底消失了。仅凭一行改动，我们在每次前向传递中砍掉了一个完整的内核。单独看可能不算多，但记住这是一个注意力操作。在基于 Transformer 的大型模型（LLM、扩散模型等）中，每层都会重复一次，而层数众多，节省量会迅速累积。
为何 out-of-place 是 PyTorch 的默认选择？为了计算梯度，autograd 必须记住前向传递中看到的张量值，因为许多反向公式会重复使用它们。In-place 操作会覆盖那些内存中的值，导致反向传递读到错误的数据。由于我们在 torch.no_grad 下运行前向，in-place 操作是安全的，没有反向传递，也不会破坏任何东西。值得一提的是，in-place 操作不仅节省时间（如同我们看到的），还节省内存（因为没有额外副本），这对 logits 这类大张量非常有利。
缩放点积注意力（Scaled Dot Product Attention, SDPA）
我们已经从基本操作构建了注意力，甚至砍掉了 Memcpy。好消息是 PyTorch 团队已经为我们完成了所有这些工作，并将整个流水线打包成一个函数：
from torch.nn import functional as F
F.scaled_dot_product_attention(q, k, v, is_causal=True)
这一行代码替代了我们手工编写的模块，is_causal=True 甚至免去了手动构建掩码。值得停下来体会这一调用隐藏了多少细节。它隐藏的不仅仅是代码行。SDPA 并非单一实现，它在底层会分派到多个后端之一，并选择支持我们输入（dtype、head 维度、掩码、硬件等）的最快后端。官方 SDPA 教程介绍了这一选择过程，后端列表定义在 torch.nn.attention.SDPBackend 枚举中：
from torch.nn.attention import SDPBackend
BACKENDS = {
"math": SDPBackend.MATH,
"flash": SDPBackend.FLASH_ATTENTION,
"efficient": SDPBackend.EFFICIENT_ATTENTION,
"cudnn": SDPBackend.CUDNN_ATTENTION,
}
通常 SDPA 会自动选择，但我们可以通过 torch.nn.attention.sdpa_kernel 上下文管理器锁定特定后端。脚本中正是这样做的，从而可以分别分析每个后端的 trace，观察它们在 trace 中的不同表现。
Math 后端
运行命令：uv run 04_c_sdpa_attention.py --backend math 然后可视化。在打开 trace 前，先猜测：我们已经将手工注意力（matmul, mul, mask, softmax, matmul）替换为一行代码，期望 trace 变得更简单、更快——更少的内核、更少的 CPU 分发，甚至可能有一个融合内核。
（原文至此中断，但