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

](https://huggingface.co/ariG23498)

在本系列的第一部分"PyTorch 中的性能分析"中，我们使用 torch.add(torch.matmul(x, w), b) 来学习如何阅读 PyTorch profiler 的 trace。我们还讨论了其他几个相关主题——CPU 调度链、启动开销、开销受限与计算受限场景的区别，以及 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 并使用 Spaces 的 Dev Mode 来实验脚本非常容易。也可以使用 Hugging Face Jobs pipeline 来运行脚本。

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

1. GPU kernel 是一个在 GPU 的多个线程上并行运行的程序。
2. CPU 调度并启动这些 kernel。你在 profiler trace 中看到的大部分 PyTorch 开销就是这种调度工作。

从 matmul-add 到 Linear

nn.Linear 是一个模块封装器，封装了我们已经在第一部分中分析过的矩阵乘法和加法。唯一的区别是它拥有自己的权重和偏置作为参数，并暴露了一个 PyTorch 用户已经熟悉的 forward 方法。

# bias=True 将真正模拟我们在本系列第一部分中看到的乘法和加法操作
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 并检查 profile。

uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64
uvx trace-util traces -b traces

trace-util 是一个工具，它会将你的 trace 同步到一个 Hugging Face bucket，然后在你的终端上提供 Preffeto URLs。

图 1：nn.Linear 的 Profiler trace

图 1 展示了线性层一次前向调用的 profiler trace。我们使用与之前 trace 类似的 schedule 设置来跟踪线性层的 forward 调用，其中 wait=1、warmup=1 和 active=3。这就是为什么我们在 CPU 和 GPU 通道上看到三个 Profile Step。

转置在做什么？

图 2：转置的 CPU 行

如果我们放大 profiler trace，如图 2 所示，我们会注意到在 aten::addmm（乘法和加法）操作之前有一个 aten::t（转置）操作。我们已经可以推断出 nn.Linear 会转置权重参数，然后将其与输入相乘。这就是我们看到 aten::t 操作的原因。

需要注意的一个重要点是，aten::t 实际上并不复制或重新组织数据：它只在 CPU 上重写张量元数据（形状和 stride）来表示转置后的矩阵。它不会在 GPU 上启动 kernel。可以通过两种方式验证这一点：查看 trace 中的 GPU 通道，或者检查 profiler 表中的 aten::t 行及其在 CUDA 上花费的时间。

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

图 3：线性层的 profile 中没有 aten::add

如图 3 所示，线性层的调度链中没有 aten::add（偏置加法）。这是因为偏置加法已经被折叠进了矩阵乘法 kernel，使用了所谓的 epilogue。

epilogue 是 GEMM（通用矩阵乘法）kernel 在最后、即将将结果写回 HBM（高带宽内存，GPU 的主内存）之前执行的一个小型计算。添加偏置、应用激活函数或乘以常数都是经典的 epilogue。使用 epilogue 的目的是避免第二次加载或写入 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 的定义。

现在是时候注意一个微妙之处了。你在第一部分的 --compile 下看到的 kernel（addmm）正是 eager 模式下的 nn.Linear 已经使用的 kernel。这里没有留给 torch.compile 去融合的东西，这是我们接下来要验证的。

--compile 能帮助单个 Linear 吗？

让我们编译前向调用并查看 profiler trace。（profiler trace 在下一节中可视化）

uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64 --compile
uvx trace-util traces -b traces

如果你比较单个 nn.Linear 的 forward 的 eager 和编译后的 trace，你会发现：

• GPU 上是相同的 cuBLAS GEMM kernel。
• CPU 上是相同的 aten::addmm 操作。
• CPU 通道上多了几行编译特有的内容。

这一点值得内化。一个常见的反应是，每当模型感觉慢时就使用 torch.compile。对于单个带偏置的 GEMM，compile 几乎无事可做。这不是 bug，只是因为 compile 需要多于一个操作才能进行可能的融合。让我们通过查看 MLP 来证明这一点。

转置去哪了？Kernel 布局和预操作

仔细阅读两个 trace（eager vs compile）的读者会注意到，eager 的 CPU 调度链比编译后的包含更多内容。

图 4：Eager 调度链，其中 aten::linear 经过 aten::t（转置）然后 aten::addmm

图 5：编译后的调度链，其中直接调用 aten::addmm，没有转置

eager 模式下 aten::linear 内部的 CPU 调度链是 aten::t 后跟 aten::addmm（图 4）。要理解 aten::t 实际做了什么，我们需要快速绕道了解 strides 和 views。

张量将其数据存储为内存中一个扁平的、连续的数值序列。shape 和 stride 是位于该序列之上的元数据，告诉 PyTorch 如何遍历它：stride (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))   # shape 和 stride 交换，数据未动
>>> T
tensor([[0, 2, 4],
        [1, 3, 5]])
>>> T.flatten()                # 强制物化，因此数据被重新排序
tensor([0, 2, 4, 1, 3, 5])

M.t() 没有移动任何一个数字。它返回了一个新的视图，其 strides 被交换了，因此逐行读取现在会以转置顺序遍历原始缓冲区 0, 1, 2, 3, 4, 5。底层数据是相同的；只有元数据不同。

这正是 aten::t 在线性层内部所做的：它不分配新张量或复制任何数据，它产生一个具有重写 strides 的权重_视图_。

正如我们在图 5 中看到的，compile 并没有移除 GPU kernel：它移除了调度该视图的 CPU 开销。Inductor 在编译时跟踪了视图链，一次性计算出最终的 strides，并发出一个直接调用 aten::addmm 的指令，其中硬编码了这些 strides。几微秒的 CPU 工作消失了，而 GPU 执行完全相同的数学运算。

正如预期的那样，当输入数据违反了编译器预计算的 strides 时，它会抛出一个错误。

如果你查看两个 trace 中的 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 在其内部循环中如何遍历其输入。调度器的工作是查看输入 strides，决定哪个后缀组合匹配，并选择正确的预编译 kernel。

profiler trace 中的 kernel 名称是 kernel 身份的哈希转储。如果两次运行显示相同的 kernel 名称，则 GPU 在做相同的工作。如果它们不同（例如，_tn_ vs _nn_、bf16 vs fp16，或 s16816gemm vs s161616gemm），则 GPU 在做不同的工作，调度器走了一条不同的分支。学会读取这个名称是比较 trace 时最有用的习惯之一。

堆叠三个 Linear：MLP

在本节中，我们将分析一个多层感知机（MLP）。为了更有趣，我们将分析一个使用 GeGLU 激活变体的前馈网络（这在实践中使用非常广泛）。这也是我们向深度学习研究史上最伟大的结语之一致敬的方式（图 6）。

图 6：GLU Variants Improve Transformer 论文的结论部分。

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 traces -b traces

在打开 trace 之前，让我们一起思考一下我们期望看到什么。forward 函数做了相当多的计算，但其中大部分我们已经熟悉了。

我们应该期望看到三个 aten::linear 调度，每个 nn.Linear 层一个。我们还应该期望看到两次 pointwise kernel 启动，一次用于 GeLU，一次用于乘法。在查看之前形成这种期望是性能分析过程中最有用的习惯：你阅读 trace 是为了_确认或推翻_一个猜测，而不是从头开始形成一个猜测。

图 7：GeGLU MLP 的 Profiler trace

图 8：线性投影 CPU 通道中高亮的 occupancy queries

从图 7 我们可以自我肯定，因为我们的直觉是正确的。每次前向传播（一个 mlp_fwd），GPU 恰好运行 5 个 kernel。图 8 高亮了线性投影层 CPU 通道中看到的 "occupancy query"。

操作	CPU 操作	GPU kernel	启动次数
gate_proj	aten::linear	ampere_bf16_s16816gemm_bf16_128x128_...	occupancy query + cudaLaunchKernel
up_proj	aten::linear	ampere_bf16_s16816gemm_bf16_128x128_...	occupancy query + cudaLaunchKernel
gelu	aten::gelu	vectorized_elementwise_kernel<4, GeluCUDAKernelImpl...>	cudaLaunchKernel
h * u	aten::mul	vectorized_elementwise_kernel<4, ...MulFunctor...>	cudaLaunchKernel
down_proj	aten::linear	ampere_bf16_s16816gemm_bf16_128x256_...	occupancy query + cudaLaunchKernel

三个 GEMM 在启动前各自执行了一个额外的 cudaOccupancyMaxActiveBlocksPerMultiprocessor 调用。我们在第一部分有一个专门的章节讨论这个，你可以在这里找到。那是 cuBLAS 在确定 grid 大小。Pointwise 操作（GeLU 和 mul）直接启动，没有 occupancy query。所以 "一个 linear" 实际上是 query + launch，而 "一个 pointwise op" 只是 launch。

图 9：该表显示某些操作启动了零个 kernel

aten::t、aten::transpose、aten::reshape、aten::view、aten::as_strided 和 aten::_unsafe_view 这些操作启动了零个 kernel。它们在表中显示 0.000us 的 CUDA 时间（图 9），因为它们只在 CPU 上重写张量元数据（形状和 stride）。扫描该表的读者会看到每个 linear 大约有六个操作名称，但其中只有一个（mm）会到达 GPU。

为什么有两种类型的 GEMM kernel？

MLP 将 [batch, seq, dim] 展平为 [batch * seq, dim] 以进行 matmul。在我们的命令行调用中，我们对 batch 使用了 64，对 seq 使用了 128，所以这就是下面 8192（batch * seq = 64 * 128）的来源。

从 trace 中：

Linear	aten::mm 输入维度	M·K·N	cuBLAS kernel	平均 CUDA 时间
gate_proj	[8192,768] x [768,3072]	8192·768·3072	…128x128…stages_32x5_tn	0.19ms
up_proj	[8192,768] x [768,3072]	8192·768·3072	…128x128…stages_32x5_tn	0.19ms
down_proj	[8192,3072] x [3072,768]	8192·3072·768	…128x256…stages_64x3_tn	0.17ms

所有三个 GEMM 的 FLOP 计数相同，每个都是 2·8192·768·3072 ≈ 38.7 GFLOP，但 down_proj 快了大约 10%。相同的工作量，不同的形状（N=768 而不是 3072），因此 cuBLAS 选择了一个不同的 tile（128×256，具有更深的 stages_64x3 流水线），该 tile 对该形状有更好的重用。

如果你想更深入地了解 tiling，这里有一个很好的入门资源。

这正是表中有两行 GEMM 的原因（图 9）：128x128 行对应 gate+up，128x256 行对应 down。

torch.compile 做了什么？

在编译 forward 方法并可视化它之前，让我们再次进行脑力练习，问自己期望在 trace 中看到什么。这是一个有趣的实验，也是每次你自己做性能分析时都应该重复的重要步骤。始终建立你的直觉，一旦出现不匹配，就停下来弄清楚原因。

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

图 10：编译后的 GeGLU MLP 的 Profiler trace

在 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）。证明它是同一个 GEMM 的证据是 kernel 名称与 eager 模式逐字节相同：gate 和 up 是 ...128x128...stages_32x5_tn，down 是 ...128x256...stages_64x3_tn。

融合的 Triton kernel

图 11：融合的 Triton kernel

这是整个 compile 课程的重点。两个 eager pointwise kernel（GeLU 和 mul）加上一个 reshape 合并成了一个 kernel，triton_poi_fused__unsafe_view_gelu_mul_0（图 11）。让我们解码这个名称：

• triton：由 Inductor 的 Triton 后端生成（不是 cuBLAS，不是 ATen）。
• poi：pointwise（Inductor 将 pointwise kernel 标记为 poi，reduction 标记为 red，persistent reduction 标记为 per）。
• fused__unsafe_view_gelu_mul：它合并的操作：_unsafe_view（reshape）、GeLU 和 mul。
• 0：图内的唯一 id。

为什么这是一个胜利？在 eager 模式下，中间结果 h = gelu(g) 是一个完整的 [8192, 3072] bf16 张量（约 50 MB），GeLU kernel 将其写入 HBM，mul kernel 立即将其读回。融合将其保留在寄存器（位于芯片内部且比 HBM 更近的内存）中。Triton kernel 读取 g 和 u 一次，计算 gelu(g) * u，并写入结果一次。中间结果通过全局内存的一次完整往返被消除了。

让我们使用手工调优的 kernel

到目前为止，我们让 PyTorch（eager）和编译器（torch.compile）来选择我们的 kernel。现在我们插入一个由人类专家编写并手工调优的 kernel。我们使用 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 traces -b traces

图 12：LigerGEGLUMLP 层的 Profiler trace

图 12 展示了使用来自 Hub 的 Liger kernel 的 LigerGEGLUMLP 层的 profile。

为什么使用 kernels 库

用 Triton 或 CUDA 编写 kernel 是一个问题，而_分发_它们是另一个问题。kernel 必须针对你的 GPU 架构、CUDA 版本和 PyTorch 版本的确切组合进行编译。这一步通常会出问题（"在我机器上能跑"、缺少 nvcc、错误的 Triton 版本）。

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

• Kernel 在 CI 中为多种架构和版本组合编译一次。你下载正确的二进制文件，而不是自己编译。
• version=1 固定了确切的构建，因此运行你脚本的每个人都会得到相同的 kernel。不存在"我更新了一个包后它变慢了"的情况。
• 该包暴露了一个 .layers 属性，其中包含即插即用的 nn.Module（如 LigerGEGLUMLP）。你将你的模块换成它们的，模型中的其他部分无需更改。

为什么调优的 kernel 更好

当我们说"调优"时，我们指的是两个具体的事情，两者都在 trace 中可见。

图 13：编译后的运行在任何 GEMM 运行之前需要支付预操作（Dynamo、guards、prologue）的开销

图 14：Liger kernel 没有预操作——它们本应出现的框是空的

1. 融合是内置的。 LigerGEGLUMLP 的 forward 是 down_proj(LigerGELUMulFunction.apply(gate_proj(x), up_proj(x)))。LigerGELUMulFunction 运行一个单一的 Triton kernel _geglu_tanh_forward_kernel，它一次性计算 gelu(gate) * up。这正是我们从 torch.compile 中看到的，中间结果不会通过 HBM 往返。我们在这里无需编译器就得到了它，如图 13 和 14 所示（没有 Dynamo guards，没有编译延迟，没有重新编译风险）。

2. 启动参数是为硬件选择的。 Kernel 不会随机猜测其 block 大小。Liger 的 calculate_settings 根据列数来选择它们。

诚实地看待这里的权衡是值得的，因为原始数字可能会产生误导。Liger kernel 运行时间为 92.8 µs，而编译运行中 Inductor 的融合 kernel 是 89.4 µs。乍一看，手写 kernel 似乎稍慢，但这种比较隐藏了使其有价值的成本。

torch.compile 针对静态形状进行特化。Inductor 的 89.4 µs kernel 之所以快，正是因为它针对_这个确切_的 [8192, 3072] 问题生成的。更改 batch size、sequence length 或 hidden dimension，Dynamo 会重新跟踪，你需要再次支付编译成本才能获得一个新的特化 kernel。

所以真正的选择不是"慢的人类 kernel vs 快的编译 kernel"。而是一个快速的通用 kernel vs 一个针对特定输入形状特化的 kernel。Liger kernel 采用一组启动参数，并针对_任何_形状运行它们，无需重新编译。它放弃了形状特化所能带来的最后几微秒，以换取对形状变化的鲁棒性。

结论

下表总结了每一步在 GPU 上改变了什么，以及什么保持不变。

设置	改变的内容	保持不变的内容
Eager nn.Linear	基线：偏置加法已折叠进 GEMM epilogue（addmm），因此是_一个_ cuBLAS kernel，而不是 matmul 加 add	—
编译后的 nn.Linear	一些 CPU 调度操作（aten::t 视图簿记）消失	相同的单个 cuBLAS GEMM kernel，逐字节相同。Compile 没有东西可融合
Eager MLP	5 个 GPU kernel：3 个 GEMM + 一个 GeLU + 一个 mul。[8192, 3072] 中间结果通过 HBM 进行一次完整往返	每个 GEMM 仍然是相同的无偏置 cuBLAS kernel，如同独立的 linear
编译后的 MLP	GeLU + mul + reshape 合并为一个融合的 Triton kernel；中间结果保留在寄存器中。支付编译预操作（Dynamo、guards）	3 个 GEMM 保持不变，具有相同的 cuBLAS kernel 名称
Liger MLP	相同的融合，但内置于手写的 Triton kernel 中，具有硬件调优的启动参数，没有 Dynamo、guards 或编译延迟	3 个 GEMM 仍然是相同的 cuBLAS kernel

如果有一个习惯需要坚持，那就是我们在每次 trace 之前练习的：先猜测，再看。 陈述你期望 trace 包含的内容，打开它，并将任何不匹配视为屏幕上最有趣的事情。

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

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