T
traeai
登录
返回首页
Hugging Face Blog

Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP

8.5Score

TL;DR · AI 摘要

PyTorch 中 nn.Linear 的性能分析显示,其内部的矩阵转置操作显著影响计算效率,通过融合 MLP 可以减少开销。

核心要点

  • nn.Linear 的矩阵转置操作会增加计算开销。
  • 使用融合的 MLP 可以减少调度和启动开销。
  • PyTorch 的 profiler 工具能有效识别性能瓶颈。

结构提纲

按章节快速跳转。

  1. 文章回顾了 PyTorch 性能分析的基础知识,并介绍了 nn.Linear 的使用。

  2. nn.Linear 是一个封装了矩阵乘法和加法的模块,其内部使用了权重和偏置参数。

  3. 通过 PyTorch 的 profiler 工具分析了 nn.Linear 的性能瓶颈,发现矩阵转置操作的开销。

  4. 文章讨论了如何通过融合多个 nn.Linear 层来减少调度和启动开销。

  5. 通过融合 MLP,可以显著减少 PyTorch 的调度开销,提高计算效率。

思维导图

用一张图看清主题之间的关系。

查看大纲文本(无障碍 / 无 JS 友好)
  • PyTorch 性能分析
    • nn.Linear 的性能分析
      • 矩阵转置操作的开销
      • PyTorch profiler 工具
    • 融合的 MLP
      • 减少调度和启动开销
      • 提高计算效率

金句 / Highlights

值得收藏与分享的关键句。

#PyTorch#性能优化#深度学习#GPU
打开原文

PyTorch 中的性能分析(第 2 部分):从 nn.Linear 到融合的 MLP

返回文章列表

[-1

]

[0

发布于 2026 年 6 月 11 日

在 GitHub 上更新

点赞

11

[

  • +5

Aritra Roy Gosthipaty

ariG23498

关注

Rémi Ouazan Reboul

ror

Sergio Paniego

sergiopaniego

Pedro Cuenca

pcuenq

Sayak Paul

sayakpaul

在本系列的第一部分“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 非常容易,可以使用 Dev Mode 与 Spaces 进行实验。也可以使用 Hugging Face Jobs 管道来运行这些脚本。

在开始之前,让我们快速回顾两个我们将反复用到的概念:

  • GPU 内核是一种在 GPU 的多个线程上并行运行的程序。
  • CPU 负责调度并启动这些内核。在性能分析追踪中,你看到的大多数 PyTorch 开销就是这个调度工作。

从矩阵乘法加法到 Linear

nn.Linear 是围绕我们已经在第 1 部分中分析过的相同矩阵乘法和加法的模块包装器。唯一的不同是它拥有自己的权重和偏置作为参数,并提供了一个 PyTorch 用户熟悉的 forward 方法。

code
# bias=True 真正地模拟了我们在系列第一部分中看到的乘法和加法操作
linear_layer = nn.Linear(in_dim, out_dim, bias=True)
y = linear_layer(x)

当前的操作可以写成:

code
y = x @ w.T + b

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

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

trace-util 是一个实用工具,它会将你的追踪同步到 Hugging Face 存储桶,并在终端上提供 Preffeto URL。

图 1:nn.Linear 的性能分析追踪

图 1 显示了线性层的前向调用的性能分析追踪。我们使用与之前追踪类似的调度设置来追踪线性层的前向调用,设置为 wait=1、warmup=1 和 active=3。这就是为什么我们在 CPU 和 GPU 轨道上看到了三个 Profile Steps。

转置操作在做什么?

图 2:转置的 CPU 行

如果我们放大性能分析追踪,如图 2 所示,我们注意到在 aten::addmm(乘法和加法)操作之前有一个 aten::t(转置)操作。我们已经可以推断出,nn.Linear 会先转置权重参数,然后将其与输入相乘。这就是为什么我们会看到一个 aten::t 操作。

一个重要的注意事项是,aten::t 并没有真正复制或重新组织数据:它只是在 CPU 上重写张量的元数据(形状和步长),以表示转置后的矩阵。它不会在 GPU 上启动内核。可以通过两种方式验证这一点:查看跟踪中的 GPU 通道,或检查 profiler 表中的 aten::t 行及其在 CUDA 上所花费的时间。

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

图 3:没有

code
aten::add

在线性层的性能分析中

如图 3 所示,线性层的调度链中没有 aten::add(即偏置加法)。这是因为偏置加法已经被合并到矩阵乘法内核中,使用了一种称为 epilogue 的方法。

epilogue 是 GEMM(GEneral Matrix Multiply,通用矩阵乘法)内核在将结果写回 HBM(High Bandwidth Memory,GPU 的主内存)之前执行的一小段计算。添加偏置、应用激活函数或按常数进行缩放都是经典的 epilogue 操作。epilogue 的目的是避免再次加载或写入 HBM,因为内存访问会使操作变得昂贵。

nn.Linear 调用了 torch.nn.functional.linear,而后者又调用了 aten::linear。aten::linear 会检查输入,发现传入了偏置,然后调度 aten::addmm(bias, x, weight),而不是分别执行矩阵乘法和加法。addmm 计算如下:

code
out = x @ weight.T + bias

在 GPU 上运行的 cuBLAS GEMM 内核有一个内置的偏置加法变体,这就是 aten::addmm 所选择的内核。加法操作永远不会作为单独的内核出现,因为它是矩阵乘法内核写回操作的一部分,这正是 epilogue 的定义。

这是需要注意的一个微妙之处。你在第一部分中看到的 --compile 下的内核(addmm)是 eager 模式下的 nn.Linear 已经使用的内核。对于 torch.compile 来说,这里已经没有可以融合的内容,这正是我们接下来要验证的。

--compile 能否帮助单个 Linear?

让我们编译前向调用并查看 profiler 的跟踪信息。(profiler 跟踪信息将在下一节中进行可视化)

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

如果你比较单个 nn.Linear 前向调用的 eager 模式和编译后的跟踪信息,你会发现:

  • GPU 上使用的是相同的 cuBLAS GEMM 内核。
  • CPU 上使用的是相同的 aten::addmm 操作。
  • CPU 通道上有一些额外的行,这些行是编译特有的。

这一点值得深入理解。当模型运行缓慢时,一个常见的反应是使用 torch.compile。但对于单个带有偏置的 GEMM 操作,torch.compile 几乎没有什么可以做。这不是一个 bug,而是因为 torch.compile 需要多个操作才能进行任何融合。让我们通过查看一个 MLP 来验证这一点。

转置去哪了?内核布局和预操作

仔细查看两个跟踪信息(eager 模式与编译模式)的读者会注意到,eager 模式的 CPU 调度链中包含的内容比编译模式的更多。

图 4:eager 模式下的调度链,其中

code
aten::linear

依次经过

code
aten::t

(转置)和

code
aten::addmm

图 5:编译模式下的调度链,其中直接调用,没有转置

在 aten::linear 中的 eager 模式 CPU 调度链是 aten::t 后跟 aten::addmm(图 4)。为了理解 aten::t 实际上做了什么,我们需要简要了解 strides 和 views 的概念。

张量在内存中以一连串连续的数字形式存储其数据。形状和步长是元数据,它们位于这串数据之上,并告诉 PyTorch 如何遍历它:步长为 (s0, s1) 表示“移动一行需要跨过 s0 个元素,移动一列需要跨过 s1 个元素”。更改这些元数据,你将获得对相同原始数据的不同视图,而无需复制:

code
>>>
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
tensor([[
0
,
2
,
4
],
        [
1
,
3
,
5
]])
>>>
T.flatten()
# 强制生成,因此数据被重新排列
tensor([
0
,
2
,
4
,
1
,
3
,
5
])

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

这正是线性层内部的 aten::t 所做的:它不会分配新的张量或复制任何数据,而是生成一个具有重写步长的权重视图。

如图 5 所示,编译器并未移除 GPU 内核:它移除了 CPU 上调度该视图的开销。Inductor 在编译时追踪了视图链,计算了一次最终的步长,并发出一个带有这些硬编码步长的直接 aten::addmm 调用。几微秒的 CPU 工作被消除,而 GPU 执行了相同的数学运算。

正如预期的那样,当输入数据违反了编译器预先计算的步长时,它将抛出错误。

如果你查看两个追踪中的 GPU 轨道,每个前向传递都有一个内核,并且两次都是相同的内核:

code
cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8

如果没有转置内核运行,那么是谁教会了 GEMM 以转置顺序读取权重矩阵?答案在内核名称中。看看后缀:

code
cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8
                                                          ^^

这个 tn 是布局描述符。cuBLAS 和 CUTLASS 会为每种输入布局的组合预编译一个单独的内核二进制文件。

n(非转置)和 t(转置)描述了内核在内循环中如何遍历其输入。调度器的任务是查看输入的步长,决定哪种后缀组合匹配,并选择正确的预编译内核。

性能分析追踪中的内核名称是内核身份的哈希转储。如果两次运行显示相同的内核名称,GPU 执行的是相同的工作。如果它们不同(例如,_tn_ 与 _nn_,bf16 与 fp16,或 s16816gemm 与 s161616gemm),则 GPU 执行的是不同的工作,调度器选择了不同的分支。学习如何阅读这个名称是对比追踪时最有用的习惯之一。

堆叠三个线性层:MLP

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

图 6:论文《GLU Variants Improve Transformer》的结论部分。

code
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。执行方式如下:

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

在我们打开跟踪信息之前,让我们一起思考一下我们应该看到什么。forward 函数执行了相当多的计算,但其中大部分对我们来说已经很熟悉了。

我们预计会看到三个 aten::linear 的调度,每个 nn.Linear 层对应一个。我们还预计会看到两个逐点内核启动,一个用于 GeLU,一个用于乘法。在查看之前形成这种预期是性能分析过程中最有用的习惯:你通过阅读跟踪信息来确认或推翻猜测,而不是从头开始形成一个。

图 7:GeGLU MLP 的性能分析跟踪图

图 8:线性投影 CPU 通道中突出显示的占用率查询

从图 7 可以看出,我们的直觉是正确的。每次 forward 传递(一个 mlp_fwd),GPU 会运行正好 5 个内核。图 8 突出了在线性投影层的 CPU 通道中看到的“占用率查询”。

Op

CPU 操作

GPU 内核

启动

code
gate_proj
code
ampere_bf16_s16816gemm_bf16_128x128_...

占用率查询 + cudaLaunchKernel

code
up_proj
code
gelu
code
aten::gelu
code
vectorized_elementwise_kernel<4, GeluCUDAKernelImpl...>

cudaLaunchKernel

code
h * u
code
aten::mul
code
vectorized_elementwise_kernel<4, ...MulFunctor...>
code
down_proj
code
ampere_bf16_s16816gemm_bf16_128x256_...

每个 GEMM 在启动之前都会进行一次额外的 cudaOccupancyMaxActiveBlocksPerMultiprocessor 调用。我们在第一部分有专门的部分介绍这一点,你可以在这里找到。这是 cuBLAS 在调整网格大小。逐点操作(GeLU 和 mul)直接启动,没有占用率查询。因此,“一个线性”实际上是查询 + 启动,而“一个逐点操作”只是启动。

图 9:表格显示某些操作启动了零个内核

aten::t、aten::transpose、aten::reshape、aten::view、aten::as_strided 和 aten::_unsafe_view 操作启动了零个内核。在表格中(图 9),它们显示 0.000us 的 CUDA 时间,因为它们只在 CPU 上重写张量元数据(形状和步长)。浏览表格的读者可以看到每个线性操作大约有六个操作名称,但其中只有一个(mm)会到达 GPU。

为什么会有两种类型的 GEMM 内核?

MLP 将 [batch, seq, dim] 展平为 [batch * seq, dim] 以进行矩阵乘法。在我们的命令行调用中,我们使用了 64 作为 batch 和 128 作为 seq,因此下面的 8192(batch * seq = 64 * 128)就是由此而来。

从跟踪信息中:

线性

code
aten::mm

输入维度

M·K·N

cuBLAS 内核

平均 CUDA 时间

code
[8192,768] x [768,3072]
code
8192·768·3072
code
…128x128…stages_32x5_tn

0.19ms

code
[8192,3072] x [3072,768]
code
8192·3072·768
code
…128x256…stages_64x3_tn

0.17ms

所有三个 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 方法并可视化之前,让我们再次进行一次思维练习,思考我们期望在追踪中看到什么。这是一次有趣的实验,也是每次你自己进行性能分析时都应该重复的重要步骤。始终基于你的直觉进行分析,当发现某些地方不匹配时,停下来找出原因。

code
uv run 03_simple_mlp.py --batch 64 --
seq
128 --dim 768 --hidden 3072 --compile
uvx trace-util traces -b 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 的调用方式中。我们可以看到三个直接的 aten::mm 外部调用(图 10)。证明它们是同一个 GEMM 的证据是,内核名称与 eager 模式下完全一致:...128x128...stages_32x5_tn 用于 gate 和 up,而 ...128x256...stages_64x3_tn 用于 down。

融合的 Triton 内核

图 11:融合的 Triton 内核

这是整个编译课程的重点。两个 eager 模式的点wise内核(GeLU 和 mul)以及一个 reshape 操作被合并为一个内核,triton_poi_fused__unsafe_view_gelu_mul_0(图 11)。让我们解析这个名字:

  • triton:由 Inductor 的 Triton 后端生成(不是 cuBLAS,也不是 ATen)。
  • poi:点wise(Inductor 用 poi 标记点wise内核,用 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 读取回来。融合操作将其保留在寄存器中(寄存器位于芯片内部,比 HBM 更接近)。Triton 内核只读取 g 和 u 一次,计算 gelu(g) * u,并将结果写入一次。中间变量通过全局内存的一次往返被省略了。

让我们使用手工调优的内核

到目前为止,我们让 PyTorch(eager 模式)和编译器(torch.compile)为我们选择内核。现在,我们引入一个人工专家手工编写和调优的内核。我们使用 LigerGEGLUMLP 层,可以轻松地从 Hugging Face Hub 获取,使用 kernels 库。

code
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 。

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

图 12:LigerGEGLUMLP 层的性能分析追踪

图 12 显示了使用来自 Hub 的 Liger 内核对 LigerGEGLUMLP 层进行的性能分析。

为什么使用内核库

在 Triton 或 CUDA 中编写内核是一个问题,而将它们部署出去则是另一个问题。内核必须为你的 GPU 架构、CUDA 版本和 PyTorch 版本的精确组合进行编译。这通常是导致问题的步骤(“在我的机器上可以运行”,缺少 nvcc,错误的 Triton 版本)。

内核库将这个构建步骤从你的机器上移除。get_kernel("kernels-community/liger-kernels", version=1) 从 Hugging Face Hub 下载一个预构建、版本固定的内核包,并将其缓存到本地(此处为 ~/.cache/...kernels-community--liger-kernels)。其优势包括:

  • 内核在 CI 中一次性编译,适用于许多架构和版本组合。你下载的是正确的二进制文件,而不是自己编译。
  • version=1 固定构建的精确版本,因此运行你脚本的每个人都能获得相同的内核。不会有“在我更新包之后变慢”的情况。
  • 该包暴露了一个 .layers 属性,其中包含可以直接替换的 nn.Module(如 LigerGEGLUMLP)。你只需将模块替换为它们的,模型中的其他部分不会发生变化。

为什么调优后的内核更好

当我们说“调优”时,我们指的是两个具体的事情,这两点在追踪中都可以看到。

图 13:编译后的运行在任何 GEMM 运行之前支付预操作(Dynamo、guards、prologue)的开销

图 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 guards,没有编译延迟,没有重新编译的风险)。
  • 启动参数是为硬件选择的。内核不会随机猜测其块大小。Liger 的 calculate_settings 从列数中选择它们。

在这里诚实地面对权衡是值得的,因为原始数据可能具有误导性。Liger 内核运行时间为 92.8 微秒,而 Inductor 编译运行的融合内核为 89.4 微秒。乍一看,手写内核似乎稍慢,但这个比较隐藏了使它值得的代价。

torch.compile 是为静态形状专门化的。Inductor 的 89.4 微秒内核之所以快,是因为它为这个特定的 [8192, 3072] 问题生成。如果你更改批处理大小、序列长度或隐藏维度,Dynamo 会重新追踪,并且你将再次支付编译成本以获得新的专用内核。

因此,真正的选择不是“慢的人工内核 vs 快的编译内核”。而是快速的通用内核 vs 为特定输入形状专门化的内核。Liger 内核使用一组启动参数,无需重新编译即可运行任何形状。它放弃了每个形状专门化可能带来的最后几微秒,以换取对形状变化的鲁棒性。

结论

下表总结了每个步骤在 GPU 上的更改和未更改的内容。

设置

更改的内容

保持不变的内容

Eager

基准:偏差加法已经折叠到 GEMM 的 epilogue 中(

code
addmm

),所以它是

一个

cuBLAS 内核,而不是矩阵乘法加上一个加法

已编译

一些 CPU 分发操作(视图簿记)消失

与之前完全相同的一个 cuBLAS GEMM 内核,字节对字节。编译过程中没有进行任何融合

急切的 MLP

5 个 GPU 内核:3 个 GEMM + 一个 GeLU + 一个乘法。中间结果

code
[8192, 3072]

在 HBM 中完成了一次完整的往返

每个 GEMM 仍然是与独立线性层相同的无偏 cuBLAS 内核

编译后的 MLP

GeLU + 乘法 + reshape 被合并为一个融合的 Triton 内核;中间结果保留在寄存器中。支付编译预操作(Dynamo,guards)

3 个 GEMM 保持不变,使用相同的 cuBLAS 内核名称

Liger MLP

相同的融合,但被嵌入到一个手工编写的 Triton 内核中,使用了针对硬件优化的启动参数

没有

Dynamo、guards 或编译延迟

3 个 GEMM 仍然是相同的 cuBLAS 内核

如果要延续一个习惯,那就是我们在每次跟踪之前都实践过的:先猜测,然后查看。说出你期望跟踪中包含的内容,打开它,把任何不匹配视为屏幕上最有趣的事情。

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

感谢 Noe Flandre 和 Pedro Gabriel Gengo Lourenço 对文章初稿的审阅!

本文中提到的数据集 1

更多来自我们博客的文章

torch

profile

CUDA

在 PyTorch 中进行性能分析(第 1 部分):torch.profiler 入门指南

  • +1

94

2026 年 5 月 29 日

open-source

kernels

从 Codex 和 Claude 为所有人定制内核

80

2026 年 2 月 13 日

社区

编辑

预览

通过拖拽到文本输入框、粘贴或

点击此处

上传图片、音频和视频。

点击此处上传图片

评论

· 注册或登录以发表评论

AI 可能会生成不准确的信息,请核实重要内容

Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP | Hugging Face Blog | traeai