PyTorch 性能分析（第二部分）：从 nn.Linear 到融合 MLP-未完纪

这是“PyTorch 性能分析”系列的第二篇文章。本系列逐步培养阅读分析器跟踪的能力，并用它指导优化：

在本系列第一部分中，我们用 torch.add(torch.matmul(x, w), b) 学习如何阅读 PyTorch 分析器跟踪。途中还讨论了 CPU 分派链、启动开销、开销受限与计算受限场景的区别，以及 torch.compile 的一些内部机制。

在本文中，我们再上一层：用 nn.Linear（bias=True）替换手写的矩阵乘法加加法。这是各类深度学习模型使用的基本构件。然后，针对本文示例，将三个 Linear 堆叠起来，中间加入激活，形成多层感知机（MLP）块。

本文脚本在这里：02_linear.py、03_simple_mlp.py 和 03_kernels_mlp.py。与前文一样，建议单独打开它们，阅读时对照代码。原作者使用 NVIDIA A100-SXM4-80GB GPU 运行脚本。在 Hugging Face 基础设施上配置 GPU，并通过 Spaces 开发模式实验很容易；也可以使用 Hugging Face Jobs 流水线运行脚本。

开始之前，先回顾两个后文会反复依赖的概念：

GPU 内核是在 GPU 的许多线程上并行运行的程序。
CPU 调度并启动这些内核。分析器跟踪中看到的大部分 PyTorch 开销，就是这种调度工作。

从 matmul-add 到 Linear

nn.Linear 是模块封装，执行的仍是第一部分分析过的矩阵乘法和加法。区别只是它把权重和偏置作为自己的参数，并提供 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：nn.Linear 的分析器跟踪

图 1 展示了线性层一次前向调用的分析器跟踪。我们采用与此前类似的 schedule 设置来跟踪 forward：wait=1、warmup=1、active=3。因此，CPU 和 GPU 轨道中能看到三个 Profile Step。

转置在做什么？


图 2：转置的 CPU 行

像图 2 一样放大跟踪，会看到 aten::addmm（乘法和加法）之前有一个 aten::t（转置）操作。由此已经可以判断，nn.Linear 先转置权重参数，再与输入相乘；这就是出现 aten::t 的原因。

aten::t 实际上不会复制或重排数据。它只在 CPU 上重写张量的形状和步幅元数据，用它们表示转置后的矩阵，不会在 GPU 上启动内核。可以通过两种方式核验：查看跟踪中的 GPU 轨道，或检查分析器表中的 aten::t 行及其 CUDA 耗时。

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


图 3：线性层分析结果中没有 aten::add

如图 3，线性层分派链中没有负责偏置加法的 aten::add。这是因为偏置加法已经通过称为 epilogue（尾声）的机制合并进矩阵乘法内核。

epilogue 是 GEMM（通用矩阵乘法）内核在最末尾、把结果写回 HBM（高带宽内存，即 GPU 主存）前进行的一小段计算。添加偏置、应用激活或乘以常数，都是典型的 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 内核内置偏置加法变体，aten::addmm 选择的正是它。加法不会作为独立内核出现，因为它是 matmul 内核写回过程的一部分，这也正是 epilogue 的含义。

这里有一个微妙之处：第一部分中使用 --compile 时看到的 addmm 内核，正是 eager 模式的 nn.Linear 已经使用的内核。这里已经没有什么可供 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 的 forward 在 eager 与编译模式下的跟踪，会发现：

GPU 上是相同的 cuBLAS GEMM 内核。
CPU 上是相同的 aten::addmm 操作。
CPU 轨道额外出现几行编译模式特有的内容。

理解这一点很有价值。模型变慢时，人们常会下意识使用 torch.compile。对于单个带偏置 GEMM，compile 能做的事情很少。这不是缺陷，而是融合至少需要多个操作才可能发生。下面通过分析 MLP来证明。

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

仔细对照 eager 与 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 的实际工作，需要先了解步幅与视图。

张量在内存中把数据存为一段平坦、连续的数字。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))   # 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，compile 没有移除 GPU 内核，而是移除了分派该视图的 CPU 开销。Inductor 在编译时跟踪视图链，一次性计算最终步幅，生成将这些步幅硬编码进去的直接 aten::addmm 调用。几个微秒的 CPU 工作消失了，而 GPU 执行完全相同的计算。

正如预期，当输入数据违反编译器预先计算的步幅时，会抛出错误。

查看两份跟踪的 GPU 轨道，每次前向都恰好只有一个内核，而且两次都是相同的内核：

cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8

既然没有运行转置内核，是什么让 GEMM 按转置顺序读取权重矩阵？答案在内核名称中，请看后缀：

cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8
                                                          ^^

tn 是布局描述符。cuBLAS 与 CUTLASS 会针对输入布局的每一种组合，预编译单独的内核二进制。

n（未转置）与 t（已转置）描述内核在内循环中如何遍历输入。分派器检查输入步幅，确定匹配的后缀组合，再选择正确的预编译内核。

分析器跟踪中的内核名称，可以看作内核身份的散列式信息集合。如果两次运行显示同一个内核名，GPU 就在执行相同工作。如果名字不同，例如 _tn_ 与 _nn_、bf16 与 fp16，或 s16816gemm 与 s161616gemm，GPU 做的工作就不同，分派器选择了不同分支。比较跟踪时，学会阅读这个名字是最有用的习惯之一。

堆叠三个 Linear：MLP

本节分析多层感知机（MLP）。为了更有趣，使用带 GeGLU 激活变体的前馈网络，这种变体在实践中很常见。这也是向深度学习研究历史中一段极精彩的话致意，见图 6。


图 6：《GLU Variants Improve Transformer》论文的结论部分。（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 -f traces -b <hf_uname>/traces

打开跟踪前，先思考预计会看到什么。forward 进行了不少计算，但其中多数已经很熟悉。

预计有三次 aten::linear 分派，分别对应三个 nn.Linear 层，还应有两次逐元素内核启动，一次 GeLU、一次乘法。在查看前形成预期，是性能分析过程中最有用的习惯：阅读跟踪是为了确认或推翻猜测，而不是从零开始猜。


图 7：GeGLU MLP 的分析器跟踪


图 8：线性投影 CPU 轨道中高亮显示的占用率查询

图 7 证明了我们的直觉是正确的。每次前向传播，即一次 mlp_fwd，GPU 恰好运行 5 个内核。图 8 高亮了线性投影层 CPU 轨道中的“占用率查询”。

操作	CPU 操作	GPU 内核	启动操作
gate_proj	aten::linear	ampere_bf16_s16816gemm_bf16_128x128_…	占用率查询 + cudaLaunchKernel
up_proj	aten::linear	ampere_bf16_s16816gemm_bf16_128x128_…	占用率查询 + 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_…	占用率查询 + cudaLaunchKernel

三个 GEMM 在启动前各执行一次额外的 cudaOccupancyMaxActiveBlocksPerMultiprocessor 调用。第一部分有单独一节介绍它，见这里。这是 cuBLAS 在确定网格大小。逐元素操作 GeLU 与 mul 则直接启动，没有占用率查询。所以“一个线性层”实际上是查询加启动，而“一个逐元素操作”只有启动。


图 9：表格显示有些操作没有启动任何内核

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。

从跟踪中可以看到：

Linear	aten::mm 输入维度	M·K·N	cuBLAS 内核	平均 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 的浮点运算数相同，每个都是 2·8192·768·3072 ≈ 38.7 GFLOP，但 down_proj 约快 10%。工作量相同、形状不同——N 为 768 而不是 3072——因此 cuBLAS 选择不同分块，即 128×256、具有更深的 stages_64x3 流水线；它对该形状有更好的复用。

若想深入了解 tiling（分块），这个资源是不错的起点。

这正是图 9 表格有两个 GEMM 行的原因：128×128 对应 gate 和 up，128×256 对应 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


图 10：编译后 GeGLU MLP 的分析器跟踪

在 eager 模式，每个 nn.Linear 展开为一串分派器操作：aten::linear → aten::t → aten::transpose → aten::matmul → aten::reshape → aten::mm。这些是 ATen 到达真正 GEMM 前经过的高层封装。torch.compile 移除了这条链。

编译图运行时，已没有 linear、matmul、transpose 或 reshape，这些元数据操作被并入 mm 的调用方式。图 10 中能看到三个直接的 aten::mm 外部调用。证明 GEMM 相同的依据是：内核名字与 eager 模式逐字节一致，gate 和 up 是 ...128x128...stages_32x5_tn，down 是 ...128x256...stages_64x3_tn。

融合的 Triton 内核


图 11：融合的 Triton 内核

这是整段 compile 学习中最重要的结果：两个 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 获取。（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：LigerGEGLUMLP 层的分析器跟踪

图 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。好处有：（kernels）

内核在 CI 中针对多种架构与版本组合编译一次。你下载正确的二进制，而无须自行编译。
version=1 固定确切构建，因此运行脚本的人得到同一内核，不会出现更新软件包后变慢的问题。
软件包通过 .layers 属性暴露可直接替换的 nn.Module，如 LigerGEGLUMLP。用它替换自己的模块，模型其他部分无须改变。

为什么调优过的内核更好

这里的“调优”具体指两件事，跟踪中都能看到。


图 13：编译运行在任何 GEMM 开始前都要付出前置操作成本，包括 Dynamo、守卫检查和序言


图 14：Liger 内核没有这些前置操作，对应位置为空

融合已经内置。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；LigerGELUMulFunction；_geglu_tanh_forward_kernel）
启动参数针对硬件选择。内核不会随机猜测块大小，Liger 的 calculate_settings 根据列数选择它们。（calculate_settings）

必须如实说明这里的权衡，因为原始数字可能产生误导。Liger 内核运行 92.8 μs，而 compile 运行中 Inductor 的融合内核是 89.4 μs。乍看手写内核稍慢，但这种比较隐藏了它值得采用的成本因素。

torch.compile 针对静态形状特化。Inductor 的 89.4 μs 内核之所以快，正是因为针对确切的 [8192, 3072] 问题生成。改变 batch 大小、序列长度或隐藏维度，Dynamo 就重新跟踪，需要再次支付编译成本，才能得到新的特化内核。

真正的选择是快速的通用内核与针对一种特定输入形状特化的内核之间的权衡。Liger 内核使用一组启动参数处理任意形状，无须重新编译。它牺牲逐形状特化带来的最后几个微秒，换取对形状变化的稳健性。

结语

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

配置	发生的变化	保持不变的部分
eager nn.Linear	基线：偏置加法已经合并进 GEMM epilogue（addmm），因此只有一个 cuBLAS 内核，而不是 matmul 加独立 add。	—
编译后的 nn.Linear	少数 CPU 分派操作消失，例如 aten::t 的视图记账工作。	单个 cuBLAS GEMM 内核逐字节相同；compile 没有什么可融合。
eager MLP	5 个 GPU 内核：3 个 GEMM、1 个 GeLU、1 个 mul。[8192, 3072] 中间结果在 HBM 中完整往返一次。	每个 GEMM 仍然是与独立 linear 相同的不带偏置 cuBLAS 内核。
编译后的 MLP	GeLU、mul、reshape 合为一个融合 Triton 内核，中间结果保留在寄存器中；需支付编译前置操作成本，包括 Dynamo 和守卫检查。	3 个 GEMM 未被改变，cuBLAS 内核名相同。
Liger MLP	相同的融合，但直接内置于手写 Triton 内核，并使用针对硬件调优的启动参数；没有 Dynamo、守卫或编译延迟。	3 个 GEMM 仍是相同的 cuBLAS 内核。