PyTorch 性能优化：nn.Linear 到融合 MLP 的实战指南

这是“PyTorch 性能分析”系列的第二篇文章。在这个系列中，我们将逐步培养阅读性能分析器跟踪记录的能力，并利用它推动性能优化：

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

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

PyTorch 性能分析（第 3 部分）：你需要分析的只有注意力

在本系列的第一部分“PyTorch 性能分析”中，我们通过 torch.add(torch.matmul(x, w), b) 学习了如何阅读 PyTorch 性能分析器的跟踪记录。我们还讨论了过程中涉及的其他几个主题：CPU 分发链、启动开销、受开销限制与受计算限制两种状态之间的区别，以及 torch.compile 的一些内部机制。

在第二轮（也就是本文）中，我们沿着阶梯再向上迈一级。我们用一个 nn.Linear（设置 bias=True）替换手写的矩阵乘法与加法组合。这是每个深度学习模型都会使用的基础构建块。接着，我们将三个这样的层堆叠起来（这是本示例采用的具体数量），并在它们之间加入激活函数，从而组成一个多层感知机（MLP）块。

本文使用的脚本包括：02_linear.py、03_simple_mlp.py 和 03_kernels_mlp.py。与之前一样，建议在单独的标签页中打开这些脚本，边阅读文章边浏览代码。我们使用 NVIDIA A100-SXM4-80GB GPU 运行这些脚本。在 Hugging Face 基础设施上配置一块 GPU 非常容易，你可以通过 Spaces 的 Dev Mode 运行并试验这些脚本，也可以使用 Hugging Face Jobs 流水线来运行它们。

在开始之前，先快速回顾两个我们会反复用到的概念：

GPU kernel 是一个在 GPU 的大量线程上并行运行的程序。

CPU 负责调度和启动这些 kernel。你在性能分析器跟踪记录中看到的大部分 PyTorch 开销，都来自这些调度工作。

从 matmul-add 到 Linear

nn.Linear 是对我们在第 1 部分中已经分析过的同一种矩阵乘法和加法操作的模块化封装。唯一的区别是，它将权重和偏置作为自己的参数持有，并提供了 PyTorch 用户已经非常熟悉的 forward 方法。

# bias=True would truly emulate the multiplication and addition
# operations we have seen in part 1 of the series
linear_layer = nn.Linear(in_dim, out_dim, bias=True)
y = linear_layer(x)

这里的运算可以写成：

y = x @ w.T + b

其中，x 是输入，w 是权重，b 是偏置。让我们运行 02_linear.py 并查看性能分析结果。

uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64
uvx trace-util -f traces -b <hf_uname>/traces

trace-util 是一个实用工具，它会将你的跟踪记录同步到 Hugging Face bucket，然后在终端中提供 Perfetto URL。

图 1 展示了线性层一次前向调用的性能分析器跟踪记录。我们使用与之前类似的调度配置来跟踪线性层的前向调用，其中 wait=1、warmup=1、active=3。因此，在 CPU 和 GPU 轨道中都能看到三个 Profile Step。

转置操作在做什么？

如果像图 2 那样放大查看性能分析器跟踪记录，就会注意到，在 aten::addmm（乘法与加法）操作之前有一个 aten::t（转置）操作。由此我们已经可以判断，nn.Linear 会先转置权重参数，然后再将其与输入相乘。这就是我们看到 aten::t 操作的原因。

需要特别注意的是，aten::t 并不会真正复制或重新组织数据：它只在 CPU 上重写张量的元数据（形状和步幅），以表示转置后的矩阵。它不会在 GPU 上启动 kernel。可以通过两种方式验证这一点：查看跟踪记录中的 GPU 轨道，或者检查性能分析器表格中的 aten::t 行以及它在 CUDA 上耗费的时间。

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

如图 3 所示，线性层的分发链中没有 aten::add（偏置加法）。这是因为偏置加法已经通过所谓的 epilogue（尾声操作）被折叠进矩阵乘法 kernel 中。

epilogue 是 GEMM（通用矩阵乘法）kernel 在执行的最后阶段、即将结果写回 HBM（高带宽内存，GPU 的主内存）之前完成的一小段计算。添加偏置、应用激活函数或乘以常数，都是典型的 epilogue。epilogue 的意义在于避免第二次从 HBM 加载数据或向 HBM 写入数据，因为内存流量会让运算变得昂贵。

nn.Linear 会调用 torch.nn.functional.linear，后者再调用 aten::linear。aten::linear 检查输入时会注意到传入了偏置，于是分发 aten::addmm(bias, x, weight)，而不是分别执行 matmul 和 add。addmm 计算的是：

out = x @ weight.T + bias

在 GPU 上运行的 cuBLAS GEMM kernel 内置了添加偏置的变体，aten::addmm 选择的正是这个 kernel。add 不会以单独 kernel 的形式出现，因为它属于 matmul kernel 写回结果过程的一部分，而这正是 epilogue 的含义。

此时需要留意一个不太明显的事实。你在第 1 部分使用 --compile 时看到的 kernel（addmm），正是 eager 模式下的 nn.Linear 已经在使用的 kernel。因此，这里已经没有什么可以留给 torch.compile 继续融合了，接下来我们将验证这一点。

--compile 能优化单个 Linear 吗？

让我们编译前向调用，并查看性能分析器跟踪记录。（跟踪记录的可视化将在下一节展示。）

uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64 --compile
uvx trace-util -f traces -b <hf_uname>/traces

如果比较单个 nn.Linear 前向传播在 eager 和编译模式下的跟踪记录，你会发现：

GPU 上运行的是同一个 cuBLAS GEMM kernel。

CPU 上执行的是同一个 aten::addmm 操作。

CPU 轨道上多出了几行编译模式独有的记录。

这一点值得牢记。当模型运行缓慢时，人们常常会下意识地使用 torch.compile。但对于单个带偏置的 GEMM，编译器几乎没有什么可做的。这并不是 bug，只是因为编译器至少需要面对多个操作，才有可能执行融合。接下来，让我们通过观察一个 MLP 来证明这一点。

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

仔细观察两份跟踪记录（eager 与 compile）的读者会发现，eager 模式的 CPU 分发链中包含的内容比编译模式更多。

在 eager 模式下，aten::linear 内部的 CPU 分发链是先执行 aten::t，再执行 aten::addmm（图 4）。要理解 aten::t 实际做了什么，我们需要先简单了解一下步幅和视图。

张量会将数据存储为内存中一段扁平、连续的数值序列。形状和步幅是位于这段数据之上的元数据，用来告诉 PyTorch 应该如何遍历它：步幅 (s0, s1) 表示“移动一行需要前进 s0 个元素，移动一列需要前进 s1 个元素”。改变元数据，就能在完全不复制数据的情况下得到同一份原始数据的不同视图：

>>> M = torch.tensor([[0, 1],
...                   [2, 3],
...                   [4, 5]])
>>> M.shape, M.stride()
(torch.Size([3, 2]), (2, 1))   # two steps per row, one step per column

>>> T = M.t()                  # transpose
>>> T.shape, T.stride()
(torch.Size([2, 3]), (1, 2))   # shape and stride swapped, data untouched
>>> T
tensor([[0, 2, 4],
        [1, 3, 5]])
>>> T.flatten()                # forced to materialize, so the data is reordered
tensor([0, 2, 4, 1, 3, 5])

M.t() 没有移动任何一个数值。它返回了一个步幅被交换的新视图，因此现在逐行读取它时，会以转置后的顺序遍历原始缓冲区中的 0, 1, 2, 3, 4, 5。底层数据完全相同，变化的只有元数据。

这正是线性层内部的 aten::t 所做的事情：它不会分配新张量，也不会复制任何数据，只会生成一个重写了步幅的权重视图。

如图 5 所示，编译并没有消除某个 GPU kernel，而是消除了分发该视图所需的 CPU 开销。Inductor 在编译期追踪了整个视图链，预先计算出最终的步幅，然后生成一个直接调用 aten::addmm 的操作，并将这些步幅硬编码进去。这样，CPU 上几微秒的工作便消失了，而 GPU 执行的数学运算完全相同。

正如预期的那样，如果输入数据不符合编译器预先计算的步幅要求，就会抛出错误。

如果查看两份跟踪记录中的 GPU 轨道，会发现每次前向传播都恰好只有一个 kernel，并且两次使用的是同一个 kernel：

cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8

既然没有运行转置 kernel，那么是谁告诉 GEMM 要以转置后的顺序读取权重矩阵？答案就在 kernel 的名称中。请看它的后缀：

cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8
                                                          ^^

其中的 tn 是布局描述符。cuBLAS 和 CUTLASS 会针对输入布局的每一种组合，分别预编译对应的 kernel 二进制文件。

n（非转置）和 t（转置）描述了 kernel 在内部循环中遍历输入的方式。分发器的工作是检查输入的步幅，判断哪种后缀组合与之匹配，然后选择正确的预编译 kernel。

性能分析器轨迹中的内核名称，是内核身份信息的哈希式转储。如果两次运行显示相同的内核名称，说明 GPU 正在执行相同的工作。如果名称不同（例如 _tn_ 与 _nn_、bf16 与 fp16，或者 s16816gemm 与 s161616gemm），那么 GPU 执行的工作就不同，分派器也走了不同的分支。在比较轨迹时，学会读懂这个名称是最有用的习惯之一。

堆叠三个 Linear：MLP

本节将对多层感知机（MLP）进行性能分析。为了让内容更有意思，我们将分析一个采用 GeGLU 激活变体的前馈网络（它在实践中被广泛使用）。这也是我们向深度学习研究史上最伟大的代码行之一致敬的方式（图 6）。

class SimpleGeGLUMLP(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.gate_proj = nn.Linear(dim, hidden, bias=False)
        self.up_proj = nn.Linear(dim, hidden, bias=False)
        self.down_proj = nn.Linear(hidden, dim, bias=False)

def forward(self, x):
        g = self.gate_proj(x)
        u = self.up_proj(x)
        h = F.gelu(g, approximate="tanh")
        m = h * u
        y = self.down_proj(m)
        return y

完整脚本位于：03_simple_mlp.py。执行方式如下：

uv run 03_simple_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072
uvx trace-util -f traces -b <hf_uname>/traces

在打开轨迹之前，先一起思考一下我们应该看到什么。forward 函数执行了相当多的计算，但其中大部分内容我们已经很熟悉了。

我们应该会看到三次 aten::linear 分派，分别对应三个 nn.Linear 层。还应该会看到两个逐元素内核启动，一个用于 GeLU，另一个用于乘法。在查看轨迹之前先形成这种预期，是整个性能分析过程中最有用的习惯：阅读轨迹是为了验证或推翻猜想，而不是从零开始形成猜想。

从图 7 可以看出，我们的直觉是正确的，值得给自己一点鼓励。每次前向传播（一个 mlp_fwd）中，GPU 恰好运行 5 个内核。图 8 突出了线性投影层在 CPU 轨道中显示的“占用率查询”。

三个 GEMM 在启动前都会额外调用一次 cudaOccupancyMaxActiveBlocksPerMultiprocessor。我们在第 1 部分专门介绍过这一点。这里是 cuBLAS 在确定网格大小。逐元素操作（GeLU 和 mul）则直接启动，不进行占用率查询。因此，“一个 linear”实际上是“查询 + 启动”，而“一个逐元素操作”只是“启动”。

aten::t、aten::transpose、aten::reshape、aten::view、aten::as_strided 和 aten::_unsafe_view 操作不会启动任何内核。它们在表格中（图 9）的 CUDA 时间显示为 0.000us，因为它们只在 CPU 上改写张量元数据（形状和步幅）。浏览表格时，每个 linear 周围大约能看到六个操作名称，但其中只有一个（mm）真正到达 GPU。

为什么有两种 GEMM 内核？

MLP 会将 [batch, seq, dim] 展平为 [batch * seq, dim]，以便执行矩阵乘法。在我们的命令行调用中，batch 使用 64，seq 使用 128，因此下面的 8192 就来自这里（batch * seq = 64 * 128）。

三个 GEMM 的 FLOP 数量都相同，每个都是 2·8192·768·3072 ≈ 38.7 GFLOP，但 down_proj 的速度却快了约 10%。工作量相同，但形状不同（N=768，而不是 3072），因此 cuBLAS 选择了不同的分块（128×256，并采用层数更深的 stages_64x3 流水线），从而针对这一形状获得了更好的数据复用效果。

如果你想深入了解分块，这里有一份很适合作为起点的优质资料。

这正是表格中出现两行 GEMM（图 9）的原因：128x128 这一行对应 gate 和 up，128x256 这一行对应 down。

torch.compile 做了什么？

在编译并可视化 forward 方法之前，我们再次进行一次思维练习，思考预期会在轨迹中看到什么。这是一个很有趣的实验，也是你每次自行开展性能分析时都应该重复的重要步骤。始终以自己的直觉为基础，一旦发现有任何不一致，就停下来弄清楚原因。

uv run 03_simple_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072 --compile
uvx trace-util -f traces -b <hf_uname>/traces

在 eager 模式下，每个 nn.Linear 都会展开为一系列分派器操作（aten::linear → aten::t → aten::transpose → aten::matmul → aten::reshape → aten::mm）。这些是 ATen 在到达真正的 GEMM 之前依次经过的高层封装。torch.compile 会移除这条调用链。

当编译后的计算图开始运行时，已经不存在 linear、matmul、transpose 或 reshape，这些元数据操作都被折叠进了 mm 的调用方式中。我们可以看到三个单独的外部 aten::mm 调用（图 10）。能够证明它们与 eager 模式使用相同 GEMM 的依据是，内核名称逐字节完全一致：gate 和 up 使用 ...128x128...stages_32x5_tn，down 使用 ...128x256...stages_64x3_tn。

融合的 Triton 内核

这是整个编译课程的核心内容。eager 模式下的两个逐元素内核（GeLU 和 mul）以及一个 reshape，被合并成了一个内核：triton_poi_fused__unsafe_view_gelu_mul_0（图 11）。下面来解读这个名称：

triton：由 Inductor 的 Triton 后端生成（不是 cuBLAS，也不是 ATen）。

poi：逐元素操作（Inductor 使用 poi 标记逐元素内核，使用 red 标记归约内核，使用 per 标记持久化归约内核）。

fused__unsafe_view_gelu_mul：它合并的操作，包括 _unsafe_view（reshape）、GeLU 和 mul。

0：该计算图中的唯一 ID。

为什么这是一种优化？在 eager 模式下，中间结果 h = gelu(g) 是一个完整的 [8192, 3072] bf16 张量（约 50 MB），GeLU 内核会将其写入 HBM，随后 mul 内核又会立即将其读回。融合会把它保留在寄存器中（寄存器是位于芯片内部、比 HBM 更近的存储空间）。Triton 内核只需读取一次 g 和 u，计算 gelu(g) * u，再将结果写入一次。这样就省去了中间结果通过全局内存的一整次往返。

使用手工调优的内核

到目前为止，我们一直让 PyTorch（eager）和编译器（torch.compile）来选择内核。现在，我们接入一个由人类专家编写并手工调优的内核。这里使用 LigerGEGLUMLP 层，可以通过 kernels 库轻松地从 Hugging Face Hub 获取。

from kernels import get_kernel

kernels_layers = get_kernel("kernels-community/liger-kernels", version=1).layers
kernels_geglu_mlp = kernels_layers.LigerGEGLUMLP(Config()).to(device, dtype=torch.bfloat16).eval()

完整脚本位于：03_kernels_mlp.py。

uv run 03_kernels_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072
uvx trace-util -f traces -b <hf_uname>/traces

图 12 展示了使用 Hub 上 Liger 内核的 LigerGEGLUMLP 层性能分析结果。

为什么使用 kernels 库

使用 Triton 或 CUDA 编写内核是一个问题，交付这些内核则是另一个问题。内核必须针对你所使用的 GPU 架构、CUDA 版本和 PyTorch 版本的确切组合进行编译。通常出问题的正是这一步（“在我的机器上能运行”、缺少 nvcc、Triton 版本错误）。

kernels 库将这一构建步骤从你的机器上移走。get_kernel("kernels-community/liger-kernels", version=1) 会从 Hugging Face Hub 下载预构建且固定版本的内核包，并将其缓存在本地（这里位于 ~/.cache/...kernels-community--liger-kernels）。这样做的好处包括：

内核只需在 CI 中编译一次，并面向多种架构和版本组合构建。你只需下载正确的二进制文件，而不必自行编译。

version=1 固定了确切的构建版本，因此所有运行你脚本的人都会得到相同的内核。不会出现“更新软件包后速度变慢了”的情况。

该软件包通过 .layers 属性提供可直接替换的 nn.Module（例如 LigerGEGLUMLP）。你只需将自己的模块替换为它们提供的模块，模型中的其他部分都无需改变。

为什么调优后的内核更好

当我们说“调优”时，具体指的是两件事，而且两者都可以在轨迹中看到。

融合已经内置其中。LigerGEGLUMLP 的前向传播是 down_proj(LigerGELUMulFunction.apply(gate_proj(x), up_proj(x)))。LigerGELUMulFunction 会运行单个 Triton 内核 _geglu_tanh_forward_kernel，在一次遍历中计算 gelu(gate) * up。这与我们在 torch.compile 中看到的情况完全相同：中间结果不再通过 HBM 往返。正如图 13 和图 14 所示，我们在这里无需使用编译器就能获得这种效果（没有 Dynamo 护卫条件，没有编译延迟，也没有重新编译的风险）。

融合已经内置其中。LigerGEGLUMLP 的前向传播是 down_proj(LigerGELUMulFunction.apply(gate_proj(x), up_proj(x)))。LigerGELUMulFunction 会运行单个 Triton 内核 _geglu_tanh_forward_kernel，在一次遍历中计算 gelu(gate) * up。这与我们在 torch.compile 中看到的情况完全相同：中间结果不再通过 HBM 往返。正如图 13 和图 14 所示，我们在这里无需使用编译器就能获得这种效果（没有 Dynamo 护卫条件，没有编译延迟，也没有重新编译的风险）。

启动参数是针对硬件选择的。内核并不会随机猜测其块大小。Liger 的 calculate_settings 会根据列数选择这些参数。

启动参数是针对硬件选择的。内核并不会随机猜测其块大小。Liger 的 calculate_settings 会根据列数选择这些参数。

这里有必要坦诚说明其中的权衡，因为原始数字可能会造成误导。Liger 内核的运行时间为 92.8 µs，而编译运行中 Inductor 融合内核的运行时间为 89.4 µs。乍看之下，手写内核似乎稍慢一些，但这样的比较掩盖了让它值得使用的成本因素。

torch.compile 会针对静态形状进行特化。Inductor 的 89.4 µs 内核之所以快，正是因为它是专门为这个确切的 [8192, 3072] 问题生成的。只要改变批次大小、序列长度或隐藏维度，Dynamo 就会重新追踪，而你必须再次支付编译成本，才能得到一个新的特化内核。

因此，真正的选择并不是“慢速的人工编写内核与快速的编译内核”。而是快速的通用内核与针对某一种特定输入形状特化的内核之间的选择。Liger 内核采用一组启动参数，无需重新编译即可在任意形状上运行。它放弃了按形状特化所能节省的最后几微秒，换来对形状变化的稳健适应能力。

下表汇总了每一步在 GPU 上改变了什么，以及哪些部分保持不变。

如果说有一个习惯值得延续下去，那就是我们在查看每个跟踪结果之前所实践的习惯：先猜，再看。先说明你预期跟踪结果中会包含什么，然后打开它，并将任何与预期不符之处视为屏幕上最值得关注的内容。

这是《PyTorch 性能分析》系列的第二站。在下一篇文章中，我们将继续向上攀登，从这个 MLP 块转向注意力块，并最终分析一个完整模型。

感谢 Noe Flandre 和 Pedro Gabriel Gengo Lourenço 对本文早期草稿的审阅！

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

本文提及的数据集 1

我们博客中的更多文章

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

PyTorch 性能分析（第 3 部分）：你所分析的全都是注意力

谢谢兄弟，真的很有帮助，期待你的下一篇博客！

感谢你的友善留言！🤗

非常感谢这个系列的博文！一个好习惯：先推理，再查看 👍。

有几个部分我一直在思考：

当输入数据不符合编译器预先计算的步幅时，它会抛出错误还是重新编译？

文中说“每个 GEMM 仍然是与独立线性层相同的无偏置 cuBLAS 内核”，但当前这个线性层原本就从未使用过偏置。之前的是无偏置的，使用 aten::addmm；当前这个同样无偏置，但使用的是 aten::mm。

你怎么看？我很期待第三篇博文！

一如既往，都是很棒的问题。

我认为它会重新编译，torch.compile 一直以来都是这样处理的。PyTorch 团队有一篇关于如何处理重新编译的官方文章。

抓得很准！确实如此。

感谢这篇以直觉为引导的博客。真的非常感谢！:)

很高兴你喜欢，@deependujha 🔥

我无比期待第三篇博文。

· 注册或登录后发表评论

本文提及的数据集 1