Profiling in PyTorch (Part 2): 从 nn.Linear 到 Fused MLP 深度解读

背景

在 PyTorch 性能剖析系列文章的第一部分中，我们使用 torch.add(torch.matmul(x, w), b) 这一基础操作作为切入点，学习了如何阅读 PyTorch Profiler 的追踪记录（traces）。当时我们探讨了 CPU 分发链（dispatch chain）、启动开销（launch overhead）、开销主导型（overhead-bound）与计算主导型（compute-bound）模式的区别，以及 torch.compile 的一些内部机制。

在本系列文章的第二部分（即本文），我们将分析层级再提升一步。我们将之前手写的 matmul-add 操作对替换为 nn.Linear（带偏置 bias=True）。这是每个深度学习模型都使用的基础构建模块。随后，我们将三个这样的线性层堆叠起来（针对本例的具体情况），并在它们之间加入激活函数，从而形成一个多层感知机（MLP）模块。

为了配合阅读，本文涉及的脚本位于：02_linear.py、03_simple_mlp.py 和 03_kernels_mlp.py。建议打开这些脚本并在阅读时对照代码。所有实验均在 NVIDIA A100-SXM4-80GB GPU 上运行。在 Hugging Face 基础设施上设置 GPU 并实验非常容易，可以使用 Dev Mode with Spaces 或 Hugging Face Jobs pipeline 来运行这些脚本。

在开始之前，回顾两个我们将反复依赖的核心概念：

1. GPU Kernel：是在 GPU 众多线程上并行运行的程序。
2. CPU 调度：CPU 负责调度和启动这些内核。你在 Profiler 追踪中看到的绝大多数 PyTorch 开销都源于这种调度工作。

核心内容

从 matmul-add 到 Linear

nn.Linear 本质上是我们已在第一部分剖析过的矩阵乘法和加法操作的模块封装。唯一的区别在于，它拥有自己的权重（weight）和偏置（bias）作为参数，并暴露了 PyTorch 用户熟悉的 forward 方法。

其操作可以表示为： $$y = x @ w.T + b$$ 其中 $x$ 是输入，$w$ 是权重，$b$ 是偏置。

当我们运行 02_linear.py 并检查配置文件时，我们会看到 nn.Linear 的前向传播追踪记录。值得注意的是，在 aten::addmm（乘法与加法）操作之前，出现了一个 aten::t（转置）操作。这表明 nn.Linear 在内部对权重参数进行了转置，然后将其与输入相乘。

关键点：aten::t 并没有真正复制或重组数据。 它仅在 CPU 上重写张量的元数据（形状和步幅），以表示转置后的矩阵。它不会在 GPU 上启动任何内核。我们可以通过查看追踪记录中的 GPU 通道，或者检查 Profiler 表中 aten::t 行在 CUDA 上花费的时间来验证这一点。

为什么没有单独的 mul 和 add 内核？

在 nn.Linear 的分发链中，我们没有看到单独的 aten::add（偏置加法）。这是因为偏置加法已经被折叠进了矩阵乘法内核中，这种技术被称为 Epilogue（尾随计算/后处理）。

Epilogue 是 GEMM（GEneral Matrix Multiply，通用矩阵乘法）内核在将结果写回 HBM（High Bandwidth Memory，GPU 的主内存）之前执行的少量计算。添加偏置、应用激活函数或乘以常数都是经典的 Epilogue。其核心目的是避免第二次加载或写入 HBM，因为内存流量会使操作变得昂贵。

具体流程如下：

1. nn.Linear 调用 torch.nn.functional.linear。
2. 后者调用 aten::linear。
3. aten::linear 检查输入，发现传入了偏置，于是分发 aten::addmm(bias, x, weight)，而不是分别执行 matmul 和 add。
4. addmm 计算：$out = x @ weight.T + bias$。
5. GPU 上运行的 cuBLAS GEMM 内核内置了偏置加法变体，aten::addmm 正是选择了这个内核。

因此，加法从未作为单独的内核出现，因为它是矩阵乘法内核写回过程的一部分，这正是 Epilogue 的定义。这也意味着，你在第一部分 --compile 模式下看到的 addmm 内核，正是 eager 模式下的 nn.Linear 已经使用的内核。对于单个 GEMM 带偏置的情况，torch.compile 没有更多的融合空间。

编译能否帮助单个 Linear？

让我们编译前向传播调用并查看 Profiler 追踪记录。对比 eager 模式和编译模式下的单个 nn.Linear 前向传播追踪记录，我们会发现：

• GPU 上使用的是相同的 cuBLAS GEMM 内核。
• CPU 上使用的是相同的 aten::addmm 操作。
• CPU 通道上仅有几行属于编译模式的额外开销。

这值得深入理解。常见的直觉是当模型变慢时就使用 torch.compile。但对于单个带偏置的 GEMM，compile 的作用微乎其微。这不是 bug，而是因为 compile 需要多个操作才可能进行融合。

转置去哪了？内核布局与前操作

仔细对比两份追踪记录（eager vs compile）的读者会发现，eager 模式的 CPU 分发链比编译模式的更复杂。

在 eager 模式下，aten::linear 内部的分发链是 aten::t 后跟 aten::addmm。要理解 aten::t 的实际作用，我们需要快速回顾一下 Strides（步幅） 和 Views（视图）。

张量在内存中存储为一段连续的扁平数字序列。shape 和 stride 是位于该序列之上的元数据，告诉 PyTorch 如何遍历它：步幅 $(s0, s1)$ 意味着“移动一行需跨越 $s0$ 个元素，移动一列需跨越 $s1$ 个元素”。更改元数据即可在不复制数据的情况下获得同一原始数据的不同视图。

例如：

>>> M = torch.tensor([[0, 1],
... [2, 3],
... [4, 5]])
>>> M.shape, M.stride()
(torch.Size([3, 2]), (2, 1)) # 每行两步，每列一步
>>> T = M.t() # 转置
>>> T.shape, T.stride()
(torch.Size([2, 3]), (1, 2)) # 形状和步幅交换，数据未动

M.t() 没有移动任何一个数字，它返回了一个新视图，其步幅已交换。读取它时，现在按转置顺序遍历原始缓冲区。底层数据完全相同，仅元数据不同。

这正是 nn.Linear 内部 aten::t 所做的：它不分配新张量或复制数据，而是生成一个具有重写步幅的权重视图。

在编译模式下，Inductor 在编译时追踪了视图链，计算了最终的步幅，并生成了带有硬编码步幅的直接 aten::addmm 调用。这消除了几微秒的 CPU 分发工作，而 GPU 执行相同的数学运算。因此，编译并没有移除 GPU 内核，而是移除了分发该视图的 CPU 开销。

如果输入数据违反了编译器预计算的步幅，就会抛出错误。

谁教会了 GEMM 以转置顺序读取权重？

如果在两份追踪记录的 GPU 通道中，每次前向传播都只有一个内核，且内核相同，那么如果没有运行转置内核，是谁教会了 GEMM 以转置顺序读取权重矩阵？答案就在内核名称中。

观察内核名称的后缀： cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8

其中的 tn 后缀表明，该内核被配置为以转置（transposed）形式读取权重（第二个操作数），