PyTorch Attention 层性能分析与优化

PyTorch Attention 层性能分析与优化 | 前端进阶之旅

Full translation

完整中文译文

这是 PyTorch 性能分析系列的第三篇文章。在这个系列中，我们逐步培养阅读性能分析器跟踪记录的能力，并用它来指导优化：

PyTorch 性能分析（第 1 部分）：torch.profiler 初学者指南

PyTorch 性能分析（第 2 部分）：从 nn.Linear 到融合 MLP

PyTorch 性能分析（第 3 部分）：Attention is all you profile（本文）

“PyTorch 性能分析”系列旨在帮助你熟练阅读性能分析器的跟踪记录和表格。在第 1 部分中，我们分析了加法和乘法等基础数学运算。我们了解了如何通过性能分析器表格发现热点，以及如何通过性能分析器跟踪记录观察算法随时间推移的执行顺序。

在第 2 部分中，我们将这些加法和乘法封装进了一个 torch 线性层。随后，我们将多个线性层堆叠在一起，构成一个多层感知机，并对其进行了性能分析。在此过程中，我们还分析了融合内核和手工调优的内核。

从 Transformer 架构的角度来看，接下来合乎逻辑的分析对象是另一个基础算法：注意力。尽管注意力机制因其平方时间复杂度而广为人知，但已经有许多巧妙的技巧可以缓解这一问题，使其高效运行。我们的目标并不是详细介绍每一种技巧，而是观察它们在性能分析器中分别呈现出怎样不同的特征。

本文使用的脚本位于这里：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 流水线运行它们。

注意力机制使用查询（Queries，q）、键（Keys，k）和值（Values，v）。它们之间的交互可以表示为一组简短的步骤：

构建注意力分数 scores：matmul(q, k.T)

缩放分数：scores * scale

对分数应用因果掩码：scores.masked_fill(mask, "-inf")

使用 softmax 对分数进行归一化，得到注意力权重 attn：softmax(scores)

使用这些权重重新加权值：matmul(attn, v)

因此，注意力实际上是一组原始操作的集合。其中一些我们已经熟悉了（矩阵乘法），其余操作也很容易识别。接下来，让我们用 PyTorch 编写一个朴素的注意力模块，并对它进行性能分析。

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

在打开跟踪记录之前，我们先像往常一样练习一下，猜测其中应该出现什么。跟踪这个模块的前向传播时，我们预计会看到：

一个 matmul 内核（q . k.T）

一个 mul 内核（缩放操作）

一个执行掩码操作的算子

一个 matmul 内核（atten . v）

uv run 04_a_naive_attention.py
uvx trace-util -f traces/ -b <hf_uname>/traces

图 1 展示了性能分析结果的 CPU 轨道（GPU 轨道已折叠，以免信息过多）。在 attn_fwd（我们标注的前向调用）内部，可以准确看到我们刚才猜测的那些操作。matmul 到现在已经是我们的老朋友了，新出现的操作也很容易识别：

masked_fill：因果掩码

softmax：softmax 内核

现在，让我们展开 GPU 轨道，看看实际启动了哪些内核。

图 2 展示了紧邻 CPU 轨道的 GPU 轨道。让我们放大 GPU 轨道上的一个 attn_fwd 块，逐一查看其中的内核。

通过图 3，我们可以读出单个性能分析步骤中的各个内核：

matmul（查询和键）

softmax（生成注意力权重）

matmul（注意力权重和值）

其中五个内核都符合预期。内存复制是唯一一个意料之外的内核，那么它来自哪里？线索在于 PyTorch 支持原地操作。当你以普通方式（非原地方式）操作张量时，PyTorch 通常会创建一个副本，对该副本执行所请求的操作，然后返回这个副本。沿着操作序列排查，可以发现这里的罪魁祸首是 masked_fill。

如果把它替换为原地操作，会发生什么？

使用原地因果掩码的朴素注意力

我们只需要将 masked_fill 改为 masked_fill_（注意末尾的下划线，这是 PyTorch 对原地操作的命名约定），然后运行同一个脚本。

def forward(self, q, k, v, mask):
    # q, k, v: [batch, heads, seq, head_dim]
    scores = torch.matmul(q, k.transpose(-2, -1))  # [batch, heads, seq, seq]
    scores = torch.mul(scores, self.scale)
-    scores = scores.masked_fill(mask, float("-inf"))
+    scores.masked_fill_(mask, float("-inf"))
    attn = torch.softmax(scores, dim=-1)
    out = torch.matmul(attn, v)  # [batch, heads, seq, head_dim]
    return out

让我们查看跟踪记录，看看是否发生了变化。

uv run 04_b_inplace_ops_attention.py
uvx trace-util -f traces/ -b <hf_uname>/traces

与非原地版本（图 4）相比，原地版本（图 5）在掩码步骤中包含的 CPU 操作少得多。这是一个令人鼓舞的信号。接下来展开 GPU 轨道，确认那里发生了什么。

在 GPU 轨道上，Memcpy 内核彻底消失了（图 6 和图 7）。仅仅修改一行代码，我们就在每次前向传播中减少了一个完整的内核。单独来看，这似乎算不了什么，但请记住，这只是一次注意力操作。在基于 Transformer 的大模型（LLM、扩散模型等）中，每一层都会执行一次这样的操作，而模型通常有很多层，因此节省下来的开销会迅速累积起来（如果这项优化让你加薪了，至少分给我们 10% 才算公平）。

PyTorch 默认使用非原地操作是有原因的。为了计算梯度，autograd 必须记住它在前向传播中看到的张量值，因为许多反向传播公式会再次使用这些值。原地操作会覆盖内存中的这些值，导致反向传播读取到错误的数据。由于我们在 torch.no_grad 下执行前向传播，不存在反向传播，也就没有任何会被破坏的数据，因此使用原地操作对我们来说是安全的。还值得注意的是，原地操作不仅能节省时间（正如本例所展示的那样），还能节省内存（因为不需要额外的副本），这对于 logits 之类的大型张量尤为有益！

缩放点积注意力

我们刚刚使用原始操作构建了注意力，甚至还消除了一个 Memcpy。好消息是，PyTorch 团队已经替我们完成了所有这些工作，并将整套流程封装进了一个函数：

from torch.nn import functional as F

F.scaled_dot_product_attention(q, k, v, is_causal=True)

这一行代码就可以替代我们手写的模块，而且 is_causal=True 甚至免去了手动构建掩码的工作。值得停下来体会一下，这一个函数调用隐藏了多少内容。它隐藏的不仅仅是代码行数。缩放点积注意力（Scaled Dot Product Attention，SDPA）并不只有一种实现。在底层，它会分派到多个后端之一，并选择支持当前输入条件（dtype、头维度、掩码、硬件等）的最快后端。

官方 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 上下文管理器固定使用某个特定后端。我们的脚本采用的正是这种方式。这样一来，我们就能分别分析每个后端，并观察它们在跟踪记录中呈现出的不同特征。下面逐个来看。

uv run 04_c_sdpa_attention.py --backend math
uvx trace-util -f traces/ -b <hf_uname>/traces

在打开任何内容之前，我们先来猜一下。我们已经用一行代码替换了手写的注意力实现（matmul、mul、mask、softmax、matmul），所以跟踪记录理应变得更简单，执行速度也应该更快。内核数量更少、CPU 分派更少，甚至可能出现融合内核。让我们先查看性能分析器表格。

这是第一个出乎意料的结果：这一行代码反而慢了 3.7 倍。

打开跟踪记录（图 9）后，就能看出为什么警报会响起：math 后端在每次前向传播中启动了 20 个 GPU 内核，而我们的朴素注意力实现（图 8）只启动了 5 个。这与我们的猜测完全相反。让我们分析一下为什么会出现这种情况。

在第 2 部分中，我们学会了像读取指纹一样读取内核名称。现在就在这里运用这个习惯：

我们用来捕获这些迹线的 A100 GPU 配备了张量核心，这是用于加速矩阵乘法的专门硬件，众所周知它比普通 CUDA 核心快得多。要理解为什么这很重要，需要了解 GPU 内部的构成。流式多处理器（SM）是 GPU 的计算单元，每个 SM 有两种算术单元：CUDA 核心和张量核心。CUDA 核心是通用的，每次处理少量元素，而张量核心在单条指令中就能完成整个小矩阵块的乘法累加。所以问题很简单："每个后端是否真的在使用高效路径？"

内核名称可以回答这个问题。朴素内核中的 s16816（图 10）是 bfloat16 张量核心矩阵乘法的特征（16x8x16 张量核心指令），所以朴素版本在高效路径上。sgemm（图 11）是在普通 CUDA 核心上运行的经典单精度（FP32）矩阵乘法。换句话说，数学后端根本不使用张量核心：为了在速度和数值精度之间做出权衡，它将张量转换为 FP32（即使输入是 bf16 也会加倍数据移动），并回退到较慢的 CUDA 核心。

在我们构建的朴素版本中，我们构建一次因果掩码并重复使用。这里我们传入 is_causal=True，数学后端为我们生成了一个，在每一次调用时都要这样做。你可以在 CPU 轨道上看到它发生：

这是图 12 中我们看到的

aten::ones -> aten::tril            构建 [seq, seq] 下三角矩阵
aten::scalar_tensor -> aten::fill_  构造 -inf 填充值
aten::where                         将其转换为加法偏置（0 或 -inf）

在 GPU 上，这表现为 triu_tril_kernel、多个 where 内核和一个 add_。便利的标志让我们不用思考掩码，但并没有消除工作，只是将其向下移动一层，掩码在每次前向传播时都重新构建。

我们手工编写的版本调用了普通的 aten::softmax。数学后端调用 aten::_safe_softmax，区别再次表现为额外的内核（图 13）：

当一整行被完全掩码（每个条目都是 -inf）时，普通 softmax 会计算 exp(-inf)/sum(exp(-inf)) = 0/0 = NaN。_safe_softmax 正是防范这种情况。我们的朴素内核从未考虑过，会在这个极端情况下悄悄产生 NaN。

综合来看，数学后端是参考实现。它是对注意力的直接、类型安全、NaN 安全的原始 ATen 操作分解。它本质上就是我们手工编写的朴素注意力，但更加细致谨慎。正是这种细致谨慎使得它极其缓慢。

它的目的不是速度，而是正确性。这使它成为完美的基线。我们接下来分析的每个后端（flash、efficient、cudnn）都试图将 20 个 GPU 内核合并为一个融合内核，保持 bf16，永远不具体化中间矩阵。

uv run 04_c_sdpa_attention.py --backend efficient
uvx trace-util -f traces -b <hf_uname>/traces

其中数学后端在一个分析器步骤中启动 20 个内核，efficient 后端仅启动一个 fmha_cutlassF_bf16_aligned_64x64_rf_sm80（如图 14 所示）。

让我们解读这个内核的名称：

fmha（融合多头注意力）：注意力中的所有原始操作现在在一个操作中"融合"。

cutlassF：基于 CUTLASS（NVIDIA 的张量核心 GEMM 开源模板），F 代表前向。

bf16_aligned：以 bfloat16 运行（没有 FP32 转换，不像数学后端）。

64x64：块大小。

rf（寄存器文件）：工作集保存在寄存器中，芯片上最快的内存。

sm80：为 Ampere（A100 的计算能力 8.0）编译。

这是从 Meta 的 xformers 库演变而来、后来被上游提交到 PyTorch 的内存高效注意力内核。当人们说"xformers 后端"时，他们指的就是这个 fmha_cutlassF 内核。

uv run 04_c_sdpa_attention.py --backend flash
uvx trace-util -f traces -b <hf_uname>/traces

void pytorch_flash 内核（图 15）是 FlashAttention-2（Tri Dao 的实现），供应商预装到 PyTorch 中。

在我们进一步阅读迹线之前，值得回答你现在应该问的问题：为什么有整个一个名为"flash"的后端，为什么它这么重要？

让我们回到数学后端一会。它真正的问题不是 20 个内核的数量，而是这些内核相互传递的东西。

第 1 步构建完整的分数矩阵 attn = q . k.T，即每个头的 [seq, seq]。对于序列长度 4096，单个头的矩阵是 4096 x 4096 ≈ 1600 万个数字。如果有足够的空间，该矩阵被写入到 HBM（GPU 的主内存）。然后，它被读回以进行缩放、再次写入掩码、再次读取进行 softmax，等等。注意力的成本由这种到 HBM 的往返流量主导，而不是由矩阵乘法本身。

FlashAttention 正是攻击这一点。它不是计算整个 s 矩阵然后才缩减它，而是按块遍历 k 和 v，边走边维护一个运行中的 softmax（"在线 softmax"技巧），并一次一块地累积输出。完整的 [seq, seq] 分数矩阵永远不会写入 HBM，它仅活在芯片上。这是让整个注意力管道合并为一个融合内核的单一想法，该内核在张量核心上保持 bf16。

这是 flash 让阅读分析器足迹的人惊讶的地方。它是最快的后端，但分析器用非常低的占有率报告它（如图 16 所示）。要理解为什么这没问题，我们需要三个快速定义。

GPU 内核本质上是由许多小执行单元执行的一系列指令。这些个别的执行单元（线程）负责加载变量、将它们相加、存储回去等。对于每个内核，我们启动许多线程，为了跟踪它们，我们按块对它们进行分组。

块被调度到流式多处理器（SM），GPU 的主要计算单元。一个块完全驻留在一个 SM 上，如果 SM 有足够的资源，可以同时托管多个块。这些资源包括寄存器、共享内存、最大常驻线程数和最大常驻束数。所以当我们说一个内核占有率低时，我们的意思是每个 SM 拥有的常驻束少于它理论上可以支持的数量。

如果你想了解更多关于线程、块、网格等的信息，这里有一个很好的资源。

如果你点击迹线中的 flash 内核，其足迹讲述了故事（图 17）。

Flash 使用大量的每线程寄存器和每块的大量共享内存。例如，如果一个块有 128 个线程，每个线程使用 255 个寄存器，那么该块需要 128 × 255 = 32,640 个寄存器。在拥有 65,536 个寄存器的 Ampere SM 上，一次最多装下两个这样的块。每个 128 线程块有 128 / 32 = 4 个束，所以两个块仅给出 8 个常驻束。对比最多 64 个常驻束，那大约是 13% 的占有率。Flash 占有率低不是因为它优化不好，而是因为每个块故意在芯片上资源使用上很"重"。

这正是整个要点。高占有率通过保持许多束准备好运行来帮助隐藏延迟，但它不能让工作本身高效。Flash 故意花费那些寄存器和共享内存，以保持注意力块在芯片上、积极地重用数据，并避免在全局内存中具体化完整的注意力矩阵。

uv run 04_c_sdpa_attention.py --backend cudnn
uvx trace-util -f traces -b <hf_uname>/traces

现在模式已经很熟悉了。像 flash 和 efficient 一样，cuDNN 为我们提供了每个前向的一个融合、flash 风格的内核（图 18）。所以自然的问题是：如果 flash 已经融合了注意力，为什么 PyTorch 还要发布另一个 flash 后端？答案是谁写了内核以及它是如何构建的，而这种差异就是使迹线看起来不同的原因。

Flash 和 efficient 是固定的、预编译的内核，供应商预装到 PyTorch 中。你每次都得到相同的二进制。cuDNN 是 NVIDIA 自己的深度学习库，其注意力内核是为手头的特定问题生成和调优的。它在精神上更接近 torch.compile 的代码生成，而不是固定的 cuBLAS 二进制。你可以直接从（非常长的）内核名称中读出来：

cudnn_generated_fort_native_sdpa_sm80_flash_fprop_wmma_f16_knob_6_128x64x64_4x1x1_cga1x1x1_kernel0_0

cudnn_generated：不是预先发布的二进制，是由 cuDNN 生成的。

flash_fprop：一个 flash 注意力风格的前向传播。所以算法与 flash 后端属于同一族。

wmma_f16：它使用 warp 级矩阵乘法累加（WMMA）API，16 位浮点管道上的张量核心路径。

knob_6：cuDNN 从一组预调优的配置（"knob"）中选择。不同的形状选择不同的 knob，很像 cuBLAS 选择一个块变体。

128x64x64：它选择的块维度。

这一个事实，每个问题生成一次，解释了迹线中其他看起来不寻常的所有内容。

无转置：CPU 泳道从 _cudnn_attention_forward 直接进入几次 aten::empty 内存分配，随后执行 kernel，期间没有任何 aten::transpose（图 19、20 和 21）。Flash 和 efficient 各自插入了四次（元数据）转置来重塑张量，而 cuDNN 会直接处理原生的 [B, H, S, D] 布局，因为它的生成器会针对该布局生成 kernel。变体轨迹图 19：Flash 图 20：Efficient 图 21：cuDNN

无转置：CPU 泳道从 _cudnn_attention_forward 直接进入几次 aten::empty 内存分配，随后执行 kernel，期间没有任何 aten::transpose（图 19、20 和 21）。Flash 和 efficient 各自插入了四次（元数据）转置来重塑张量，而 cuDNN 会直接处理原生的 [B, H, S, D] 布局，因为它的生成器会针对该布局生成 kernel。

它通过 cuLaunchKernelEx 而不是 cudaLaunchKernel 启动：本系列中其他所有 kernel 都通过运行时 API cudaLaunchKernel 启动。cuDNN 使用驱动层的扩展启动接口，该接口会携带启动属性（图 22）。图 22：cuDNN 后端的 CPU 泳道，展示了驱动层的 cuLaunchKernelEx 启动，而非 cudaLaunchKernel

它通过 cuLaunchKernelEx 而不是 cudaLaunchKernel 启动：本系列中其他所有 kernel 都通过运行时 API cudaLaunchKernel 启动。cuDNN 使用驱动层的扩展启动接口，该接口会携带启动属性（图 22）。

分析器报告实现占用率为 0%：不要只看这个数值就下结论——这是测量缺失，并非 GPU 停滞。CUPTI（分析后端）无法像处理 cudaLaunchKernel 那样，将占用率归因于通过驱动 API（cuLaunchKernelEx）发起的启动，因此该字段显示为 0。资源占用情况揭示了真相（图 23）：240 个寄存器 × 256 个线程 = 每个线程块使用 61,440 个寄存器，而 SM 只有 65,536 个寄存器，因此每个 SM 只能容纳一个线程块（8 个 warp，约为 12.5%），这与 Flash 的情况非常吻合。图 23：cuDNN kernel 报告实现占用率为 0%，每个线程使用 240 个寄存器，每个线程块包含 256 个线程

分析器报告实现占用率为 0%：不要只看这个数值就下结论——这是测量缺失，并非 GPU 停滞。CUPTI（分析后端）无法像处理 cudaLaunchKernel 那样，将占用率归因于通过驱动 API（cuLaunchKernelEx）发起的启动，因此该字段显示为 0。资源占用情况揭示了真相（图 23）：240 个寄存器 × 256 个线程 = 每个线程块使用 61,440 个寄存器，而 SM 只有 65,536 个寄存器，因此每个 SM 只能容纳一个线程块（8 个 warp，约为 12.5%），这与 Flash 的情况非常吻合。

“没有转置”很容易让我们以为 cuDNN 会是 CPU 端最轻量的后端。然而事实恰恰相反。

即使完全没有转置操作，cuDNN 每次前向传播仍会在 CPU 上花费约 214 µs，比 Flash（138 µs）或 efficient（117 µs）都要多。几乎所有时间都落在 aten::scaled_dot_product_attention 的自身耗时（占整个运行过程的 26%）以及 _cudnn_attention_forward 中。这是因为 cuDNN 的运行时引擎会在每次调用时选择并准备执行计划，也就是执行“旋钮”（knob）搜索。

可见的 ATen 操作更少，并不意味着 CPU 工作量更小；这些工作只是被移入了库中，分析器只能将其显示为一条粗大的不透明条带。当轨迹突然变得更简洁时，工作并不一定已经消失，有时它只是转移到了分析器无法进一步拆解的位置。

在 GPU 上，cuDNN（186.3 µs）的表现介于 efficient 和 Flash 之间。对于这个非常适合 Flash 的张量形状，手写的 FlashAttention-2 略胜一筹。cuDNN 在其他形状下往往能够胜出，例如更大的头维度或不同的序列长度，原因恰恰在于它的生成器会针对每个具体问题重新调优；但这种重新调优，也正是你刚才在 CPU 端付出的代价。

我们所讲内容一览

在结束之前，下面用一张表回顾我们分析过的每种注意力变体，以及每条轨迹教给我们的一个关键经验。

系列总结

如果你只能从整个系列中记住一件事，那就记住我们在查看每一条轨迹前都反复实践的习惯：先猜测，再查看。

大声说出你预期轨迹中会包含什么，然后打开轨迹，把任何不符合预期之处视为屏幕上最有趣的东西。这三篇文章中的每个真知灼见——隐藏的 Memcpy、addmm 的尾声阶段、包含 20 个 kernel 的数学后端、Flash 那个“看起来不对”的占用率，以及 cuDNN 那条粗大的 CPU 条带——都源于某个与实际轨迹不符的猜测。

性能分析并不是一项独立且令人生畏、只为 GPU 专家保留的技能。它只是一种严谨的习惯：仔细观察，并不断追问“等等，为什么会这样？”，直到你彻底理解答案。现在，你已经掌握了在自己的模型上这样做所需的术语和直觉。打开一条轨迹，提出一个猜测，然后去寻找其中不符合预期的地方。

感谢阅读《PyTorch 性能分析》系列。现在，去分析点什么吧。🤗

感谢 Noe Flandre 对本文早期草稿的审阅！

本篇博客文章使用了 LLM 进行润色。这绝不意味着我们让某个智能体在后台运行并生成了整篇博客。团队中的一些成员并非英语母语者，我们认为 LLM（其训练语料主要为英语）可以纠正一些低级语法错误，或者重新表述某些句子，使其读起来不那么令人生畏，也更加简洁。希望这能帮助回应“如果这是由 LLM 生成的，我为什么还要阅读”这一疑问。🤗

本文提及的数据集 1

更多博客文章

PyTorch 性能分析（第 2 部分）：从 nn.Linear 到融合 MLP

PyTorch 性能分析（第 1 部分）：torch.profiler 新手指南

感谢这个系列，读起来真的很愉快！

期待中文版

你好，Aritra（@ariG23498），非常感谢你带来这个系列！我有一个疑问：你是如何确定 GPU 中的 Memcpy（内存复制）操作确实是由非原地的 masked_fill 引起的？

我能想到的一种判断方式是，它恰好出现在 masked_fill kernel 操作之前？但我不太明白你实际上是如何验证这一点的。

编辑：我在自己动手研究时弄明白了，不过还是把这条评论留在这里，以防其他人也有同样的疑问。Perfetto 中的箭头映射准确展示了 CPU 上的哪个操作调用了 GPU 上的对应 kernel。

· 注册或登录后发表评论

本文提及的数据集 1