PyTorch Profiling 第三篇：如何从轨迹看懂注意力实现

系列《Profiling in 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，并使用 Dev Mode with Spaces 来试验这些脚本非常容易。也可以用 Hugging Face Jobs pipeline 来运行这些脚本。

朴素注意力

注意力使用查询（Queries，q）、键（Keys，k）和值（Values，v）。它们之间的交互可以写成一短串步骤：

1. 构建注意力分数 scores：matmul(q, k.T)
2. 缩放这些分数：scores * scale
3. 对分数应用因果掩码：scores.masked_fill(mask, "-inf")
4. 用 softmax 对分数进行归一化，得到注意力权重 attn：softmax(scores)
5. 用这些权重对值重新加权：matmul(attn, v)

因此，注意力实际上是一组基本运算的集合。其中一些我们已经了解（矩阵乘法 matmul），其余的也很容易识别。让我们在 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

在打开追踪之前，让我们做一贯的练习，猜猜应该看到什么。追踪该模块的 forward，我们预期会看到：

• 一个 matmul 内核（q . k.T）
• 一个 mul 内核（缩放）
• 一个用于掩码的运算
• 一个 softmax 内核
• 一个 matmul 内核（atten . v）

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

图 1：朴素注意力性能分析追踪的 CPU 通道，突出显示各个离散运算

图 1 展示了性能分析的 CPU 通道（GPU 通道已折叠，以免信息过多）。在 attn_fwd（我们标注的前向调用）内部，可以看到恰好是我们猜测的那些运算。matmul 如今已是老朋友，新的运算也很容易辨认：

• mul：缩放
• masked_fill：因果掩码
• softmax：softmax 内核

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

图 2：朴素注意力性能分析追踪的 GPU 与 CPU 通道，突出显示对应于一个分析器步骤的一组内核。

图 2 在 CPU 通道旁展示了 GPU 通道。让我们放大 GPU 通道上的单个 attn_fwd 块，逐个查看这些内核。

图 3：朴素注意力实现的性能分析器追踪中放大后的 GPU 通道。

图 3 让我们读出单个分析器步骤中的各个内核：

1. matmul（查询与键）
2. mul（缩放）
3. 内存拷贝 🤔
4. 因果掩码
5. softmax（生成注意力权重）
6. 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

类型	CPU 流
图 4：朴素掩码
图 5：原地掩码

原地版本（图 5）在掩码步骤中包裹的 CPU 操作远少于非原地版本（图 4）。这是一个令人鼓舞的信号。让我们展开 GPU 通道，确认那里发生了什么。

类型	GPU 流
图 6：朴素掩码
图 7：原地掩码

在 GPU 通道上，Memcpy 内核彻底消失了（图 6 和图 7）。只改一行，我们就在每次前向传播中省掉了整整一个内核。单看这一点或许不算什么，但请记住这只是一次注意力运算。在基于 transformer 的大模型（LLM、扩散模型等）中，它会在每一层重复一次，而层数很多，因此节省会迅速累积（如果这能为你赢得加薪，至少分给我们 10% 才算公平）。

非原地之所以是 PyTorch 的默认方式，是有原因的。为了计算梯度，autograd 必须记住前向传播中看到的张量值，因为许多反向公式会复用它们。原地操作会在内存中覆盖这些值，于是反向传播就会读到错误的数字。由于我们是在 torch.no_grad 下运行 forward，原地操作对我们是安全的：没有反向传播，也就没有可被破坏的东西。同样值得注意的是，原地操作不仅节省时间（正如我们在本例中所见），还节省内存（因为没有额外拷贝），这对 logits 这类大型张量非常有利！

缩放点积注意力

我们刚刚用基础算子搭建了注意力，甚至还省掉了一个 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、头维度、掩码、硬件等）的最快后端。

官方 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 上下文管理器固定某个特定后端。我们的脚本正是这样做的。这样我们就能单独剖析每个后端，并看出它们在跟踪中的表现有何不同。让我们逐个来看。

Math 后端

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.7x。

	性能分析器追踪
图 8：朴素原地注意力的性能分析器追踪，显示一次前向会启动五个 GPU 内核
图 9：SDPA math 后端的性能分析器追踪，显示单次注意力前向会启动 20 个 GPU 内核

打开追踪（图 9）就能看出警报为何响起：math 后端每次前向启动 20 个 GPU 内核，而不是我们的朴素注意力实现所启动的 5 个（图 8）。这与我们的猜测正好相反。让我们弄清楚为什么会这样。

Tensor Core 被闲置

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

运行	matmul 内核
图 10：朴素注意力
图 11：使用 math 后端的 SDPA

我们用来采集这些追踪的 A100 配备了 Tensor Cores，这是用于加速矩阵乘法的专用硬件，众所周知比普通 CUDA core 快得多。要明白这在这里为何重要，先了解 GPU 内部有什么会有帮助。流式多处理器（SM）是 GPU 的计算单元，每个 SM 有两类算术单元，即 CUDA core 和 Tensor Core。CUDA core 是通用的，一次处理少量元素，而 Tensor Core 能在一条指令中对一整块小矩阵做乘加。所以问题很简单：“每个后端真的在走快速路径吗？”

内核名称给出了答案。朴素内核中的 s16816（图 10）是 bfloat16 Tensor Core 矩阵乘法的标志（即 16x8x16 Tensor Core 指令），因此朴素版本走的是快速路径。sgemm（图 11）是在普通 CUDA core 上运行的经典单精度（FP32）矩阵乘法。换句话说，math 后端完全不接触 Tensor Core：为了用速度换取数值精度，它把张量升精度到 FP32（搬运的数据量翻倍，即便输入是 bf16），并回退到更慢的 CUDA core。

因果掩码已构建

在朴素版本中，我们只构建一次因果掩码并加以复用。这里我们传入了 is_causal=True，math 后端为我们物化了一个，而且是在每一次调用时。你可以在 CPU 通道上看到这一过程：

图 12：显示掩码相关算子的 CPU 通道

下面是我们在图 12 中看到的

aten::ones -> aten::tril            build a [seq, seq] lower-triangular matrix
aten::scalar_tensor -> aten::fill_  make the -inf fill value
aten::where                         turn it into an additive bias (0 or -inf)

在 GPU 上，这表现为一个 triu_tril_kernel、若干个 where 内核，以及一个 add_。那个让我们不必再考虑掩码的便利标志并没有去掉这些工作，只是把它下移了一层，掩码在每次前向时都从头重建。

安全 softmax

我们手写的版本调用的是普通的 aten::softmax。math 后端调用的是 aten::_safe_softmax，差异再次体现为额外的内核（图 13）：

图 13：安全 softmax，突出显示了与通用 softmax 相比多出的内核

一整行都被掩码（每个元素都是 -inf）时，普通 softmax 会算出 exp(-inf)/sum(exp(-inf)) = 0/0 = NaN。_safe_softmax 防范的正是这种情况。我们的朴素内核从不管这一点，在这种极端情况下本来会悄悄产生 NaN。

那么 math 后端是用来做什么的？

总的来说，math 后端是参考实现。它把注意力以直接、对 dtype 安全、对 NaN 安全的方式分解为基本的 ATen 算子。它本质上就是我们手写的朴素注意力，但更加谨慎。正是这种谨慎使它极其缓慢。

它的任务不是追求快，而是始终能工作。这使它成为完美的基线。接下来我们要分析的每个后端（flash、efficient、cudnn）都在试图把这 20 个 GPU kernel 折叠成基本上一个融合 kernel，它保持在 bf16，并且完全不物化中间矩阵。

Efficient 后端

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

图 14：使用 efficient 后端的 sdpa 的性能分析器轨迹

math 后端在一个分析器步骤中启动了 20 个 kernel，而 efficient 后端只启动一个 fmha_cutlassF_bf16_aligned_64x64_rf_sm80（如图 14 所示）。

让我们来解读这个 kernel 的名称：

• fmha（融合多头注意力）：注意力中的所有基本运算现在都被“融合”到一个运算中。
• cutlassF：基于 CUTLASS（NVIDIA 用于 Tensor Core GEMM 的开源模板）构建，F 表示前向。
• bf16_aligned：以 bfloat16 运行（与 math 不同，不会上转为 FP32）。
• 64x64：tile 大小。
• rf（寄存器文件）：工作集保存在寄存器中，这是芯片上最快的存储器。
• sm80：为 Ampere 编译（A100 的计算能力 8.0）。

这就是从 Meta 的 xformers 库发展而来、并被上游合入 PyTorch 的内存高效注意力 kernel。当人们说“xformers 后端”时，指的就是这个 fmha_cutlassF kernel。

Flash 后端

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

图 15：flash 后端的轨迹，每次前向一个融合的 pytorch_flash kernel

void pytorch_flash kernel（图 15）是 FlashAttention-2（Tri Dao 的实现），以内嵌源码的方式纳入 PyTorch。

在我们进一步阅读轨迹之前，值得回答一个你现在应该会问的问题：为什么会有一个名为“flash”的完整后端，以及它为什么如此重要？

为什么存在 flash attention？

让我们暂时回到 math 后端。它真正的问题不是 20 个 kernel 的数量，而是这些 kernel 彼此传递的内容。

第 1 步构建完整的分数矩阵 attn = q . k.T，其大小为 [seq, seq] 每个头。对于序列长度 4096，单个头就是 4096 x 4096 ≈ 16 million 个数。该矩阵会被写出到 HBM（GPU 的主存），前提是甚至有足够的空间来这样做。然后，它被读回以进行缩放，为掩码再次写出，为 softmax 再次读入，如此反复。注意力的开销主要由这种与 HBM 之间的来回传输主导，而不是由矩阵乘法本身主导。

FlashAttention 正是针对这一点。它不是先计算整个 s 矩阵然后再归约，而是以分块（tile）的方式遍历 k 和 v，在过程中维护一个运行中的 softmax（即“在线 softmax”技巧），并一次一块地累加输出。完整的 [seq, seq] 分数矩阵从不写入 HBM，它只存在于片上。正是这一个想法，让整个注意力流水线折叠成一个融合 kernel，并在 Tensor Core 上保持 bf16。

为什么 flash 在性能分析器下看起来“不对”

图 16：可以看到 flash kernel 的估计占用率为 13%

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

GPU kernel 本质上是由许多小型执行单元执行的一系列指令。这些单独的执行单元（线程）负责加载变量、将它们相加、存回等等。对于每个 kernel，我们会启动非常多的线程，而为了跟踪它们，我们按 block 对它们分组。

Block 被调度到流式多处理器（Streaming Multiprocessors，SM）上，SM 是 GPU 的主要计算单元。一个 block 完全驻留在一个 SM 上，并且一个 SM 可以同时承载多个 block，前提是它有足够的资源。这些资源包括寄存器、共享内存、最大驻留线程数以及最大驻留 warp 数。因此，当我们说某个 kernel 的 occupancy 较低时，指的是每个 SM 上的驻留 warp 数少于其理论上能够支持的数量。

如果你想进一步了解 thread、block、grid 等概念，这里有一份很好的资料。

如果你在 trace 中点击 flash kernel，它的资源占用情况就能说明问题（图 17）。

图 17：flash kernel 的资源占用，每个 block 大量使用寄存器和共享内存。

Flash 使用大量的每线程寄存器，以及每个 block 大量的共享内存。例如，如果一个 block 有 128 个线程，每个线程使用 255 个寄存器，那么该 block 需要 128 × 255 = 32,640 个寄存器。在拥有 65,536 个寄存器的 Ampere SM 上，一次只能容纳两个这样的 block。每个 128 线程的 block 有 128 / 32 = 4 个 warp，因此两个 block 只有 8 个驻留 warp。相对于最多 64 个驻留 warp，这大约是 13% 的 occupancy。Flash 的 occupancy 低，并不是因为它优化得不好，而是因为每个 block 在片上资源使用上被故意设计得非常“重”。

而这正是关键所在。高 occupancy 通过让许多 warp 随时准备运行来帮助隐藏延迟，但它并不会让工作本身更高效。Flash 有意消耗这些寄存器和共享内存，以便将 attention tile 保留在片上、积极复用数据，并避免在全局内存中实体化完整的 attention 矩阵。

cuDNN 后端

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

图 18：cuDNN 后端的 trace，每次前向只有一个生成的 attention kernel。

到这里，这个模式已经很熟悉了。和 flash 以及 efficient 一样，cuDNN 在每次前向中给我们一个融合的、flash 风格的 kernel（图 18）。于是自然的问题是：如果 flash 已经融合了 attention，为什么 PyTorch 还要再提供另一个 flash 后端？答案在于kernel 由谁编写以及它是如何构建的，正是这一差异让 trace 看起来不同。

cuDNN kernel 有何不同

Flash 和 efficient 是被内置进 PyTorch 的固定的、预编译的 kernel。你每次得到的都是同一个二进制文件。cuDNN 是 NVIDIA 自己的深度学习库，它的 attention kernel 是针对具体问题生成并调优的。它在理念上更接近 torch.compile 的代码生成，而不是固定的 cuBLAS 二进制文件。你可以直接从（非常长的）kernel 名称中读出这一点：

cudnn_generated_fort_native_sdpa_sm80_flash_fprop_wmma_f16_knob_6_128x64x64_4x1x1_cga1x1x1_kernel0_0

• cudnn_generated：不是预先随附的二进制文件，它是由 cuDNN 生成的。
• flash_fprop：一种 flash attention 风格的前向计算。因此该算法与 flash 后端属于同一家族。
• wmma_f16：它使用 warp 级矩阵乘加（WMMA）API，即 16 位浮点流水线上的 Tensor Core 路径。
• knob_6：cuDNN 从一组预先调优的配置（“knobs”）中进行选择。不同的 shape 会选择不同的 knob，很像 cuBLAS 选择一种 tile 变体。
• 128x64x64：它所选择的 tile 维度。

这一事实——针对每个问题生成——解释了 trace 中其他所有看起来不寻常的地方。

1. 没有转置：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

2. 它通过 cuLaunchKernelEx 启动，而不是 cudaLaunchKernel：本系列中的其他所有内核都走运行时 API cudaLaunchKernel。cuDNN 使用驱动级的扩展启动，其中携带启动属性（图 22）。

   图 22：cuDNN 后端的 CPU 泳道，显示的是驱动级启动 cuLaunchKernelEx，而不是 cudaLaunchKernel

3. 性能分析器报告达到的占用率为 0%：不要按表面数值理解，这是测量缺口，而不是 GPU 停顿。CUPTI（性能分析后端）无法像对待 cudaLaunchKernel 那样，把占用率归属到一次驱动 API（cuLaunchKernelEx）启动上，因此该字段读数为 0。资源占用补上了真相（图 23）：每个 block 有 240 registers × 256 threads = 61,440 个寄存器，而 SM 共有 65,536 个，因此每个 SM 只能容纳一个 block（8 个 warp ≈ 12.5%），与 flash 一致。

   图 23：cuDNN 内核报告达到的占用率为 0%，每线程 240 个寄存器，每 block 256 个线程

成本转移到了 CPU

“没有转置”的说法容易让人以为 cuDNN 会是 CPU 上最精简的后端。事实恰恰相反。

即便转置操作为零，cuDNN 大约花费在 CPU 上每次前向 214 µs，多于 flash（138）或 efficient（117）。几乎全部都落在 aten::scaled_dot_product_attention 的自身时间（占整个运行的 26%）和 _cudnn_attention_forward 上。那是 cuDNN 的运行时引擎在每次调用时选择并准备计划（“旋钮”搜索）。

可见的 ATen 操作变少并不意味着 CPU 工作变少，它把工作移进了库里，性能分析器只能把它显示成一根又粗又不透明的条。当一条跟踪突然变得更干净时，工作并不总是消失了，有时只是挪到了性能分析器无法拆解的地方。

在 GPU 上，cuDNN（186.3 µs）落在 efficient 和 flash 之间。在这个非常适合 flash 的形状上，手写的 FlashAttention-2 略胜一筹。cuDNN 常常在其他形状上胜出（更大的 head 维度、不同的序列长度），正是因为它的生成器会按问题重新调优，但这次重新调优也正是你刚刚在 CPU 上付出的代价。

我们所讲的一切，一览

在收尾之前，这里用一张表回顾我们分析过的每一种注意力变体，以及每条跟踪教会我们的一个教训。

系列结语

如果整个系列你只记住一件事，那就记住我们在每一次跟踪之前反复做的习惯：先猜测，再查看。

大声说出你预期这份 trace 会包含什么，打开它，并把任何不匹配都当作屏幕上最有意思的事。这三篇文章里的每一个真正洞见——隐藏的 Memcpy、addmm 后记、20 kernel math backend、flash 那“看起来不对”的 occupancy、cuDNN 那条粗壮的 CPU 柱——都来自一个与 trace 不符的猜测。

性能分析并不是一项专留给 GPU 专家的、独立且令人望而生畏的技能。它不过是仔细观察、并追问“等等，为什么那件事会发生？”直到答案豁然开朗的习惯。你现在已经具备在自己的模型上这样做的词汇和反应。打开一份 trace，形成一个猜测，然后去找出不匹配之处。

感谢阅读 Profiling in PyTorch 系列。现在去分析点什么吧。🤗

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

这篇博文使用 LLM 进行了润色。这绝不意味着我们让一个 agent 在后台运行并让它生成了这篇博客。团队中有些人并非英语母语者，他们认为 LLM（大多以英语训练）可以纠正一些低级语法错误，或把句子改写得不那么吓人、更干净。希望这能回应“如果这是 LLM 生成的，我为什么还要读”这个想法。🤗