PyTorch性能分析第二篇：从线性层到融合MLP

背景

本系列文章旨在帮助读者逐步掌握 PyTorch Profiler 的阅读技巧，并利用它驱动模型优化。在第一部分中，我们通过 torch.add(torch.matmul(x, w), b) 这一基础操作，学习了如何解读 PyTorch profiler 的追踪信息，涵盖了 CPU 调度链、启动开销、开销受限与计算受限的区别，以及 torch.compile 的一些内部机制。

在第二部分（即本文），我们将爬上一级台阶：把手动实现的 matmul-add 对替换为 nn.Linear（带 bias=True）。这是每个深度学习模型都会使用的基本构建块。接着，我们将三个这样的 nn.Linear 层堆叠在一起，中间插入激活函数，形成一个多层感知机（MLP）模块。通过这个过程，我们将深入理解 nn.Linear 的内部实现、torch.compile 的融合能力以及 MLP 的优化思路。

所有示例脚本可在对应代码仓库中找到，本文使用 NVIDIA A100-SXM4-80GB GPU 运行。在 Hugging Face 基础设施上设置 GPU 非常方便，可以通过 Spaces 的 Dev Mode 或 Hugging Face Jobs 管道来实验。

核心内容

从 matmul-add 到 Linear

nn.Linear 是一个模块封装，内部执行与第一部分相同的矩阵乘法和加法运算。唯一的区别是它拥有自己的权重 w 和偏置 b 作为参数，并公开了 PyTorch 用户熟悉的 forward 方法。

linear_layer = nn.Linear(in_dim, out_dim, bias=True)
y = linear_layer(x)

该操作等价于：y = x @ w.T + b。

运行 02_linear.py 并查看 profiler 追踪（如图 1 所示），可以看到线性层一次前向调用的追踪。使用与之前相同的 schedule 设置（wait=1, warmup=1, active=3），因此 CPU 和 GPU 通道中出现了三个 Profile Step。

转置做了什么？

放大追踪信息（图 2），我们发现在 aten::addmm（乘加操作）之前有一个 aten::t（转置）操作。这说明 nn.Linear 会将权重参数转置后再与输入相乘。值得注意的是，aten::t 并不实际复制或重排数据：它只在 CPU 上重写张量的元数据（形状和步长）以表示转置后的矩阵，不会在 GPU 上启动内核。可以通过检查 GPU 通道或 profiler 表中 aten::t 行的 CUDA 时间来验证这一点。

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

如图 3 所示，线性层的调度链中没有 aten::add（偏置加法）。这是因为偏置加法已被折叠到矩阵乘法内核中，使用了所谓的“epilogue”（后处理）。

Epilogue 是 GEMM 内核在计算结束时执行的一个小计算，就在它将结果写回 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 内核内置了偏置加法变体，这正是 aten::addmm 选用的内核。加法从未作为独立内核出现，因为它已经包含在 matmul 内核的回写过程中——这正是 epilogue 的定义。

这里有一个微妙之处：第一部分中 --compile 版本看到的 addmm 内核，正是 eager 模式下 nn.Linear 已经使用的内核。因此，对于单个带偏置的 GEMM，torch.compile 没有任何可融合的东西。

--compile 能否帮助单个 Linear？

编译前向调用并查看 profiler 追踪（图 4），比较 eager 和编译后的单个 nn.Linear 前向追踪，可以发现：

• GPU 上运行的是相同的 cuBLAS GEMM 内核。
• CPU 上运行的是相同的 aten::addmm 操作。
• 编译版本在 CPU 通道上多了几行独有的记录。

这意味着，当模型感觉慢时，torch.compile 并非万能药。对于单个 GEMM-with-bias，编译几乎没有用武之地。这并非 bug，只是编译需要多个操作才能进行融合。下面通过 MLP 来证明这一点。

转置去哪了？内核布局与预处理操作

仔细对比 eager 和编译版本的两条追踪，会发现 eager 的 CPU 调度链比编译版本更多。eager 的 aten::linear 调度链是 aten::t 后跟 aten::addmm（图 4）。要理解 aten::t 实际做了什么，需要简单了解步长（stride）和视图（view）。

张量将其数据存储为内存中一个连续的平坦数值序列。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)) 形状和步长交换，数据不变
T.flatten()          # 强制物化，数据重新排序为 tensor([0, 2, 4, 1, 3, 5])

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

这正是 aten::t 在线性层内部所做的：它不是分配新张量或复制数据，而是生成一个步长重写的权重视图。如图 5 所示，编译并没有移除 GPU 内核，而是移除了调度该视图的 CPU 开销。Inductor 在编译时追踪视图链，一次性计算出结果步长，并直接发出带有硬编码步长的 aten::addmm 调用。这样消失了几微秒的 CPU 工作，而 GPU 执行完全相同的数学运算。当然，如果输入数据违反了编译器预计算的步长，会抛出错误。

查看两条追踪的 GPU 通道，每个前向调用正好有一个内核，并且两次都是同一个内核：cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8。

从单个 Linear 到 MLP

现在我们将三个 nn.Linear 堆叠成一个简单的 MLP，中间插入 ReLU 激活函数：

class SimpleMLP(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(hidden_dim, out_dim)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu2(x)
        x = self.fc3(x)
        return x

运行脚本 03_simple_mlp.py 并查看 profiler 追踪（图 6）。在 GPU 通道上，可以看到