PyTorch性能分析（第三部分）：注意力即你所分析 来源：https://huggingface.co/blog/torch-attention-profile 返回文章列表 (https://huggingface.co/blog)

• 朴素注意力 (https://huggingface.co/blog/torch-attention-profile#naive-attention)
• 使用原地因果掩码的朴素注意力 (https://huggingface.co/blog/torch-attention-profile#naive-attention-with-inplace-causal-masking)
• 缩放点积注意力 (https://huggingface.co/blog/torch-attention-profile#scaled-dot-product-attention)
  • Math 后端 (https://huggingface.co/blog/torch-attention-profile#math-backend)
  • Efficient 后端 (https://huggingface.co/blog/torch-attention-profile#efficient-backend)
  • Flash 后端 (https://huggingface.co/blog/torch-attention-profile#flash-backend)
  • cuDNN 后端 (https://huggingface.co/blog/torch-attention-profile#cudnn-backend)
• 一览全貌 (https://huggingface.co/blog/torch-attention-profile#everything-we-covered-at-a-glance)
• 系列总结 (https://huggingface.co/blog/torch-attention-profile#concluding-the-series)

博客文章缩略图 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-attention-profile/profile-3-thumbnail.png)

本系列“PyTorch性能分析”旨在让您轻松阅读性能分析器追踪和表格。在
第一部分 (https://huggingface.co/blog/torch-profiler) 中，我们对加法和乘法等基本数学运算进行了性能分析。我们看到了性能分析表格如何揭示热点，以及性能分析追踪如何展示算法随时间运行的顺序。在
第二部分 (https://huggingface.co/blog/torch-mlp-fusion) 中，我们将这些加法和乘法封装到一个 torch 线性层中。然后我们将多个线性层堆叠在一起（一个多层感知机）并对其进行了分析。在此过程中，我们还分析了融合和手工调优的内核。从 Transformer 架构的角度来看，下一个合乎逻辑的分析步骤是另一个基础算法：注意力。尽管以其二次时间复杂度而臭名昭著，但存在许多巧妙的技巧来缓解该问题并使其更快。我们的目标不是详细涵盖每个技巧。相反，我们想了解每个技巧在分析器下有何不同表现。

此博客文章的脚本位于：
04_a_naive_attention.py (https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/04_a_naive_attention.py),
04_b_inplace_ops_attention.py (https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/04_b_inplace_ops_attention.py),
04_c_sdpa_attention.py (https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/04_c_sdpa_attention.py),
04_d_kernels_attention.py (https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/04_d_kernels_attention.py).
跟之前一样，建议在新标签页中打开它们，并边读边浏览代码。我们使用
NVIDIA A100-SXM4-80GB GPU 来运行脚本。在 Hugging Face 基础设施上设置 GPU 并使用 Spaces 的 Dev Mode (https://huggingface.co/docs/hub/spaces-dev-mode) 非常容易。也可以使用 Hugging Face Jobs 管道 (https://huggingface.co/docs/huggingface_hub/en/guides/jobs) 运行脚本。

https://huggingface.co/blog/torch-attention-profile#naive-attention 朴素注意力

注意力机制使用查询（q）、键（k）和值（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)

所以注意力本质上是一组基本操作的集合。其中一些我们已经知道（矩阵乘法），其余的很容易辨认。让我们在 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，我们预期：

• 一个矩阵乘法内核（q . k.T）
• 一个乘法内核（缩放）
• 一个用于掩码的操作
• 一个 softmax 内核
• 一个矩阵乘法内核（atten . v）

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

朴素注意力分析器追踪的 CPU 线程，其中 attn_fwd 块展开显示其 matmul、mul、masked_fill 和 softmax 操作
(https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-attention-profile/cpu-profile-naive.png)
图 1：朴素注意力的分析器追踪 CPU 线程，突出显示各个离散操作

图 1 显示了分析的 CPU 线程（GPU 线程已折叠以免显得杂乱）。在 attn_fwd（我们注释的前向调用）内部，我们可以确切地看到我们猜测的操作。矩阵乘法已经是老朋友了，新操作也容易辨认：

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

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

朴素注意力的分析器追踪，显示 CPU 线程在上方，GPU 线程在下方，每个 attn_fwd 步骤映射到一组 GPU 内核
(https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-attention-profile/gpu-profile-naive.png)
图 2：朴素注意力的分析器追踪的 GPU 和 CPU 线程，突出显示对应于一个分析器步骤的一组内核。

图 2 显示了 GPU 线程与 CPU 线程相邻。让我们放大 GPU 线程上的单个 attn_fwd 块，逐一查看内核。

朴素注意力 GPU 线程的放大视图，显示一个步骤的各个内核：两个 matmul、一个 mul、一个内存复制、一个掩码内核和一个 softmax
(https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-attention-profile/each-kernels-naive.png)
图 3：朴素注意力实现的 GPU 线程分析器追踪放大视图。

图 3 让我们可以读出一个分析器步骤的各个内核：

1. matmul（查询和键）
2. mul（缩放）
3. 内存复制 🤔
4. 因果掩码
5. softmax（生成注意力权重）
6. matmul（注意力权重和值）

其中五个是预期的。内存复制是意料之外的，那么它来自哪里？线索在于 PyTorch 有原地操作。当你以普通（非原地）方式操作张量时，PyTorch 通常会创建一个副本，对副本应用请求的操作，然后返回副本。按照操作顺序，罪魁祸首是我们的 masked_fill (https://docs.pytorch.org/docs/2.13/generated/torch.Tensor.masked_fill.html)。如果我们将其替换为原地操作会怎样？

https://huggingface.co/blog/torch-attention-profile#naive-attention-with-inplace-causal-masking 使用原地因果掩码的朴素注意力

我们只将 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 /traces

原地版本（图 5）在掩码步骤内包装的 CPU 操作远少于非原地版本（图 4）。这是一个令人鼓舞的信号。让我们展开 GPU 线程来确认发生了什么。

在 GPU 线程上，Memcpy 内核彻底消失了（图 6 和图 7）。通过一行更改，我们在每次前向传播中减少了一个完整的内核。单独来看这可能不算什么，但请记住这只是单个注意力操作。在以 Transformer 为基础的大型模型（LLM、扩散模型等）中，每层都会重复一次，而且有很多层，因此节省的开销会迅速累积（如果这给您带来了加薪，分享至少 10% 给我们感觉才公平）。

非原地操作是 PyTorch 的默认行为，这是有原因的。为了计算梯度，autograd 必须记住在前向传播中看到的张量值，因为许多反向公式会重复使用它们。原地操作会覆盖内存中的这些值，因此反向传播会读取错误的数字。由于我们在 torch.no_grad 下运行 forward，原地操作对我们来说是安全的，没有反向传播，也就没有损坏的风险。另外值得注意，原地操作不仅节省时间（就像我们在例子中看到的），还节省内存（由于没有额外副本），对于像 logits 这样的大张量来说非常棒！

https://huggingface.co/blog/torch-attention-profile#scaled-dot-product-attention 缩放点积注意力

我们刚刚从基本操作构建了注意力，甚至缩减了一个 Memcpy。好消息是 PyTorch 团队已经为我们完成了所有这些工作，并将整个流程打包成了一个函数：

from torch.nn import functional as F
F.scaled_dot_product_attention(q, k, v, is_causal=True)

这一行代替了我们手写的模块，而 is_causal=True 甚至省去了我们手动构建掩码的麻烦。值得花点时间体会一下这一个调用隐藏了多少东西。而且它隐藏的不仅仅是代码行。缩放点积注意力（SDPA）并没有一个单一的实现。在底层，它会分派到几个后端之一，并选择支持我们的输入（数据类型、头维度、掩码、硬件等）中最快的一个。
官方 SDPA 教程 (https://docs.pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial.html) 介绍了这个选择过程，而后端本身在 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 上下文管理器来固定特定的后端。这就是我们在脚本中所做的。这让我们可以分别分析每个后端，并了解它们在追踪中的不同表现。让我们逐一进行。

https://huggingface.co/blog/torch-attention-profile#math-backend Math 后端

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

在打开任何东西之前，让我们先猜测一下。我们已经将手写的注意力（matmul、mul、mask、softmax、matmul）替换为单行代码，因此我们预期追踪会变得更简单且更快。更少的内核，更少的 CPU 分派，甚至可能有一个融合内核。让我们先查看性能分析表格。

这是我们的第一个惊喜：单行代码慢了 3.7x。打开追踪（图 9）显示了警报响起的原因：math 后端每次前向启动 20 个 GPU 内核，而不是我们朴素注意力实现启动的 5 个（图 8）。这与我们的猜测相反。让我们弄清楚为什么会这样。

https://huggingface.co/blog/torch-attention-profile#tensor-cores-left-vacant 未使用的张量核心

在
第二部分 (https://huggingface.co/blog/torch-mlp-fusion#where-did-the-transpose-go-kernel-layouts-and-pre-ops) 中，我们学会了像读取指纹一样读取内核名称。让我们在这里沿用这个习惯：我们用来捕获这些追踪的 A100 搭载了
Tensor Cores (https://www.nvidia.com/en-us/data-center/tensor-cores/)，这是用于加速矩阵乘法的专用硬件，众所周知比普通的 CUDA 核心快得多。要理解为什么这很重要，了解 GPU 内部结构会有所帮助。流多处理器（SM）是 GPU 的计算单元，每个 SM 有两种算术单元：CUDA 核心和张量核心。CUDA 核心是通用型的，一次处理少量元素，而张量核心则在单条指令中乘法和累加整个小矩阵块。那么问题很简单：“每个后端是否真正使用了快速路径？”内核名称给出了答案。

朴素内核（图 10）中的 s16816 是 bfloat16 张量核心矩阵乘法的签名（16x8x16 张量核心指令），因此朴素版本走的是快速路径。sgemm（图 11）是在普通 CUDA 核心上运行的经典单精度（FP32）矩阵乘法。换句话说，math 后端从不触碰张量核心：为了用速度换取数值精度，它将张量向上转换为 FP32（数据移动量翻倍，即使输入是 bf16），并回退到较慢的 CUDA 核心。

https://huggingface.co/blog/torch-attention-profile#causal-masks-built 构建因果掩码

在朴素版本中，我们构建了一次因果掩码并重复使用。在这里，我们传递了 is_causal=True，而 math 后端为我们物化了一个掩码，在每次调用时都这样做。您可以在 CPU 线程上看到这个过程：

SDPA math 后端的 CPU 线程，显示重建因果掩码的操作：aten::ones、aten::tril、aten::scalar_tensor、aten::fill_ 和 aten::where
(https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-attention-profile/mask-math.png)
图 12：显示掩码操作的 CPU 线程

我们在图 12 中看到的是：

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

在 GPU 上，这表现为一个 triu_tril_kernel、几个 where 内核和一个 add_。让我们不再操心掩码的便利标志并没有消除工作量，只是将其下移了一层，掩码在每次前向传播中都被重新从头构建。

https://huggingface.co/blog/torch-attention-profile#the-safe-softmax 安全 softmax

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

SDPA math 后端的 GPU 线程，显示 aten::_safe_softmax 相比普通 softmax 启动的额外内核
(https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-attention-profile/safe-softmax-extra-kernels.png)
图 13：安全 softmax 突出显示与普通 softmax 相比的额外内核

一行完全被掩码（每个条目都是 -inf）的行会使普通 softmax 计算 exp(-inf)/sum(exp(-inf)) = 0/0 = NaN。_safe_softmax 正好防止了这种情况。我们的朴素内核从未费心，在这种边界情况下会悄无声息地产生 NaN。

https://huggingface.co/blog/torch-attention-profile#so-what-is-the-math-backend-for 那么 math 后端有什么用？

综合来看，math 后端是参考实现。它是注意力操作直接了当、数据类型安全、NaN 安全的分解，分解为基本的 ATen 操作。它本质上就是我们手写的朴素注意力，但更谨慎。那种谨慎正是