PyTorch 性能剖析（第 1 部分）：torch.profiler 入门指南

无法进行性能剖析( profiling )，就无法优化。

无论目标是让大语言模型（LLM）每秒生成更多词元（token）、把推理耗时缩短几毫秒，还是弄清楚训练循环为什么比规格表承诺的速度慢，最终都绕不开性能剖析。

难点在于，性能剖析的入门门槛很高。性能轨迹像一面密密麻麻的彩色矩形墙，各类事件名称也常常让人望而生畏。大多数教程默认读者已经知道如何阅读轨迹。因此，即便已经知道应该做性能剖析，打开一份轨迹也常常像是一件可以留到以后、或交给别人处理的麻烦事。本文以及这个系列，正是为了降低这道门槛。

本文是《PyTorch 性能剖析》系列的开篇文章。本系列会循序渐进地介绍如何阅读性能轨迹，并把这些观察用于指导优化。计划如下：

1. 第 1 部分（本文）： 从最简单的操作开始：一次矩阵乘法后接偏置加法，学习如何阅读 torch.profiler 返回的结果。
2. 第 2 部分： 扩展到 nn.Linear 和一个小型多层感知机（MLP），借助轨迹数据找准优化方向，并初步查看底层 CUDA 内核。
3. 第 3 部分： 结合 transformers 库，在大语言模型上综合运用前述全部知识。

本文从初学者视角记录这段过程。除基础 PyTorch 知识外，不需要其他前置知识；可以把它当作一篇轻松阅读的教程，过程中会穿插一些「原来如此」的瞬间。文章采用问题导向结构：打开一份轨迹，提出「等等，为什么会这样？」这样的问题，再顺着线索追到答案。读完本文后，应能理解以下内容：

• 如何设置 torch.profiler，以及该工具的输出数据含义；
• 读懂性能统计表与性能轨迹（CPU 时间线、GPU 时间线，以及两者间反常的空闲间隔）；
• 从 Python 接口调用，逐层下沉至 CUDA 内核执行的完整事件链路；
• 启用 torch.compile 后，哪些性能指标会变化，更关键的是：哪些指标保持不变。

开始之前，先明确两个关键定义，方便后续阅读：

1. GPU 内核（kernel）：能够依托 GPU 海量线程并行运行的程序单元。
2. CPU 负责调度并启动各类 GPU 内核。

日常开发一般无需手动编写 GPU 内核：调用 PyTorch 算子时，框架会自动将运算转化为一个或多个 GPU 内核，在显卡上完成计算。

掌握以上两个概念，我们就开始逐一探究问题。

本文使用的完整脚本在这里：01_matmul_add.py。建议在单独的标签页中打开脚本，并逐步阅读代码。本文使用 NVIDIA A100-SXM4-80GB GPU 运行这些脚本。

矩阵乘法与加法运算

正如 Sara Hooker 博士的准确调侃：人主要由水构成，而深度神经网络主要由矩阵乘法构成。矩阵乘法是深度学习的基石，用它作为性能剖析之旅的起点再合适不过。

def fn(x, w, b):
  return torch.add(torch.matmul(x, w), b)

矩阵乘法搭配偏置加法，模拟了神经元中权重与偏置的交互方式。顺带一提，这里的加法也会帮助我们理解 后文 中的编译机制。

我们将借助 torch.profiler 模块完成性能剖析，整体分为四步：

1. 准备好 需要剖析的代码（这里是 def fn，它封装了矩阵乘法和偏置加法）。
2. 标注算法。这个步骤完全可选，但建议执行。record_function 会把函数标注为 matmul_add，后续在性能轨迹中可以快速定位对应模块。

def step():
  with torch.profiler.record_function("matmul_add"):
    return fn(x, w, b)

3. 用 torch.profiler.profile 上下文管理器 包裹代码。

  with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,  # the cpu activities
        torch.profiler.ProfilerActivity.CUDA, # the gpu activities
    ],
  ) as prof:
    # it is recommended to run events multiple times to warm up the GPUs
    for _ in range(5):
      step()
      prof.step()

4. 导出 剖析结果。

# the profiler table
prof.key_averages().table(sort_by="cuda_time_total", row_limit=15)

# the profiler trace
prof.export_chrome_trace(trace_path)

剖析器会导出两类不同产物：

1. 性能统计表：提供算法的统计摘要，用来回答「什么耗时最多」。这对定位性能热点很有帮助。热点可能是最耗时的事件，也可能是流水线瓶颈，或是触发次数特别多的事件。
2. 性能轨迹：提供时间维度上的执行视图，用来回答「某个操作何时发生、为什么发生」，并展示 CPU 与 GPU 上发生的活动。当需要调查启动了哪些内核、内核启动是否延迟、CPU 与 GPU 活动是否重叠时，它很有用。

现在通过第一次执行看看这两者如何工作。（完整源码 01_matmul_add.py 在这里）

建议在搭载 GPU 的设备上运行该脚本。

uv run 01_matmul_add.py --size 64

在搭载 GPU 的机器上运行上述脚本后，会生成 traces/01_matmul_add 文件夹，里面包含两类分析结果文件：

64_bf16_cold_eager.json
64_bf16_cold_eager.txt

.txt 文件存放性能剖析汇总表。打开文件后（如图 1），可以看到一张大型表格：首列是剖析区间内执行的各类事件名称。

其余各列记录对应事件在 CPU、GPU，或 torch.profiler.profile 的 activities 参数指定设备上的耗时。重点观察耗时占比最高的事件，并结合业务逻辑判断该耗时是否合理；同时也要关注 # of Calls 这一列，它表示对应事件的触发频次。

顺便说明一下自身耗时「Self CPU/CUDA」与总耗时「CPU/CUDA total」的区别。「Self」列只统计事件自身花费的时间，不含内部子事件耗时。「total」列则包含事件自身及其所有嵌套子事件的耗时总和。因此，查看 matmul_add 的「CPU total」时，它包含自身耗时，以及它触发的所有内部子算子的 CPU 耗时。这是一个需要注意的重要细节。

如果查看表格最后两行，会发现剖析器给出：

Self CPU time total: 2.314ms
Self CUDA time total: 23.104us

CPU 耗时单位是 ms，而 GPU 耗时单位是 us。为了建立直观感受：GPU 上花费的时间（内核 ampere_bf16_s16816gemm...）不到 CPU 上花费时间（matmul_add 操作）的 1%。GPU 大部分时间处于空闲状态，这是一个明显的警示信号。出现这种情况的原因是，小规模矩阵乘法在 GPU 上运算极快，程序绝大部分开销都消耗在内核准备、GPU 任务启动、发送待乘数据和收集结果等环节。这类受调度开销制约的算法被称作 开销受限型（overhead-bound） 算法。

脱离这种状态最简单的方法，是改用更大尺寸的矩阵执行乘法运算。

uv run 01_matmul_add.py --size 4096

图 2 最后两行是：

Self CPU time total: 4.908ms
Self CUDA time total: 4.495ms

两项耗时单位均为毫秒，可见仅通过增大矩阵尺寸，就能让 GPU 承担更多实际计算。对照图 2 能发现：如今绝大部分 CUDA 耗时消耗在 GPU 内核（ampere_bf16_s16816gemm_..）上，而非发起内核调用的 CPU 操作（matmul_add）。这说明程序成功从开销受限型（overhead-bound）转为计算受限型（compute-bound）。

接下来通过 .json 结果文件可视化任务下发调用链。可以将 JSON 文件上传至 Perfetto UI 查看性能轨迹，也可以执行命令 uvx trace-util traces -b traces 直接生成 Perfetto 访问链接。

64×64 性能轨迹

在图 3 中，我们能看到矩阵乘加运算的性能剖析时序轨迹。矩形条的宽度代表事件持续时长，纵向嵌套结构代表调用层级；CPU 时间线展示 CPU 侧发生的各类事件，GPU 时间线则呈现 GPU 内核的实际执行过程。图中空白区间代表等待或硬件空闲时段。

脚本采用如下默认配置运行：

• 矩阵尺寸 64：输入、权重与偏置张量形状均为 (64, 64)
• 数据类型 bf16：采用 bfloat16（16 位脑浮点）精度
• 未开启编译：没有编译 torch 算子
• 无预热：性能采集前未做 GPU 预热

使用 Perfetto 工具时，推荐快捷键操作来快速浏览时序图：W/A/S/D 按键控制视图移动缩放。

图 4 包含两条时序泳道，分别对应 CPU 行为与 GPU 行为。CPU 时间线里能看到三段剖析采集周期（从 ProfilerStep#2 开始），该配置由采样调度器决定：

schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)

wait 参数：跳过干扰数据较多的初始化阶段（对应 ProfilerStep#0），

warmup 参数：程序预热运行，但剖析器不采集日志（对应 ProfilerStep#1），

active 参数：正式采集、最终体现在时序图中的运行轮次。调度配置可以在 这段脚本 中找到。

接下来围绕时序图展开排查。

为什么 ProfilerStep#2 耗时这么长？

图 5 中可以看到，ProfileStep#2 耗时远超后续轮次。仔细观察后，会发现自定义标记 matmul_add 也呈现相同规律。问题根源藏在标记内部的子事件中，而非标记本身。

图 6 中标注的约 228 微秒，是从进入 record_function("matmul_add") 代码块，到 PyTorch 真正下发 aten::matmul 算子之间的空窗期。造成这段耗时的诱因有很多：临时显存工作区申请、cuBLAS（NVIDIA 专有的 GPU 加速基础线性代数库）内部启发式策略开销，或是动态懒加载模块。对此可以先暂时跳过，也可以遵循通用做法：在性能采集前补充多轮预热运行。

在性能剖析场景里，预热指正式采样前预先反复执行目标算子。上文提到的各类 GPU 初始化工作仅会执行一次，不应把这类一次性开销纳入剖析统计。本示例包含两级预热：一是开启剖析器之前手动循环运行函数；二是依托剖析器配置里的 warmup 参数，在剖析周期内部完成预热。本段代码已启用循环运行和采样调度配置。

uv run 01_matmul_add.py --warmup

经过预热、64×64 矩阵的 Perfetto 性能轨迹

从图 7 能够看出，各轮剖析任务耗时基本持平，但这不代表一次性初始化开销已经被优化掉；预热只是让这类开销避开了性能采样。关于如何进一步缩减内核启动开销，可以继续阅读这篇 文章。

为什么 CPU 时序泳道和 GPU 时序泳道存在约 2.5 毫秒的时间差？

图 8 显示 CPU 和 GPU 时间线之间约有 2.5 ms 的偏移：CPU 提交 CUDA 内核之后，GPU 真正开始执行之前会有一段延迟。直觉上，预热阶段加上调度器的 wait 和 warmup 应该能让 GPU 保持忙碌，并缩小这个偏移。

为探明真实原因，我们小幅修改采样调度参数：

- schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
+ schedule = torch.profiler.schedule(wait=0, warmup=0, active=3, repeat=1)

图 9 显示，所有 GPU 运算开始前，GPU 时间线里出现了一项 Activity Buffer Request（活动缓冲区申请）事件，接下来我们放大视图细看细节。

放大 GPU 时序轨迹后能够发现：ProfileStep#0 对应的矩阵乘法和加法内核是连续串行执行的（该轮 CPU 时序未在图中展示），而 ProfileStep#1 的两个内核中间多出一段空窗。最合理的解释是缓冲区溢出，在内核运行途中触发了一次新的缓冲区申请（向 GPU 显存申请分配内存）。

想要排除其他诱因，最稳妥的办法是增加采样迭代次数进行观察：修改参数 active=20，重新采集性能数据。

如图 11 所示，仅在 ProfileStep#1 出现该空隙，和之前推断吻合，因此可以确定空档源自缓冲区内存申请。

事件链

图 12 展示了层层嵌套的 CPU 调用栈，是理解算子下发全链路的关键视图。

最外层的 ProfileStep#编号 包裹整轮性能采样逻辑；因为代码里手动添加了标注，视图里出现 matmul_add 节点，其内部包含两次 ATen 算子调用：矩阵乘法和矩阵加法。

面向用户的 PyTorch 矩阵乘法调用最终会落到 ATen 层 的 aten::matmul 调度接口。aten::mm 是专用于二维矩阵相乘的底层后端算子。

补充一个有趣细节：若给张量新增批量维度（batch 轴），PyTorch 会自动选用批量矩阵乘算子 aten::bmm。下面顺带切换样例，观察 aten::bmm 的实际调用表现。

- x = torch.randn(args.size, args.size, device=device, dtype=dtype)
- w = torch.randn( args.size, args.size, device=device, dtype=dtype)
- b = torch.randn(args.size, args.size, device=device, dtype=dtype)

+ # adding a batch size of 8
+ x = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
+ w = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
+ b = torch.randn(8, args.size, args.size, device=device, dtype=dtype)

在图 13 中，给输入张量新增批量维度后，aten::matmul 内部不再调用 aten::mm，转而使用 aten::bmm，同时还附带一系列前置 CUDA 运行时调用。这也侧面体现出 cuBLAS 会依靠内部启发式逻辑，为当前运算择优调度最合适的 GPU 内核。

后文若无特殊说明，我们均采用普通二维矩阵开展实验。

为什么矩阵乘法会多出一项 CUDA 运行时调用？

可以观察到：aten::mm 对应两次 CUDA Runtime 调用，分别是图中标注的 cudaOccupancyMaxActiveBlocksPerMultiprocessor 与 cudaLaunchKernel；而 aten::add 仅存在 cudaLaunchKernel 一次调用。

cudaOccupancyMaxActiveBlocksPerMultiprocessor 是 CPU 侧的资源预规划接口，作用为：给定内核函数、线程块尺寸与动态共享内存大小，查询单个流式多处理器（SM）最多能同时容纳多少个该内核的线程块。

由此产生疑问：为何矩阵乘法需要资源预规划，加法运算却不需要？

答案取决于内核的硬件资源占用。点击时序图里的 GPU 内核条目，即可查看对应内核的资源开销详情。

从图 15 可见，矩阵乘法的单线程寄存器占用、共享内存占用会随矩阵尺寸动态变化。cuBLAS 内置上百种不同版本的矩阵乘内核，每种内核都依托启发式策略完成启动，而该策略需要实时获取硬件算力信息，占用率查询就是启发逻辑的一环。从原理上理解，GPU 矩阵乘法基于分块（tile）计算：分块数量、单块大小由输入矩阵规格与硬件参数共同决定；现代优化算法细节虽更加复杂，但这套分块思路仍是很好的理解框架。

反观图 16，加法内核固定占用 32 个寄存器、不使用共享内存，资源需求极小，硬件资源不会成为运行瓶颈，因此无需提前查询占用率。这类算子天生轻量化。

该规律可作为性能排查快捷手段：浏览 CPU 时间线，只要出现 cudaOccupancyMaxActiveBlocksPerMultiprocessor 调用，就代表后方是资源开销大、自适应调度启动的内核，常见于通用矩阵乘 GEMM、卷积等算子；前面无占用率查询的内核，大多是逐元素运算、规约类算子，由 PyTorch 按固定逻辑直接发起调度。

cudaDeviceSynchronize 耗时为何高达约 1.78 毫秒？

cudaDeviceSynchronize 会阻塞 CPU 线程，直至当前 GPU 上所有计算任务全部执行完毕。剖析器会在有效采样窗口末尾执行该同步操作，目的是刷新缓存的性能事件，缺少这一步就会丢失内核耗时统计数据。

一次仅 26 微秒 GPU 有效运算，却伴随 1.78 毫秒同步耗时，说明程序空闲占比高达 98%，是典型的开销受限型性能特征。

4096×4096 性能轨迹

前文通过性能统计表分析已知：增大矩阵尺寸后，算法会从开销受限转为计算受限。

执行下述命令，深入分析性能轨迹：

uv run 01_matmul_add.py --size 4096 --warmup

相同内核，单次运行耗时为何参差不齐？

从图 17 可见，ProfileStep#3 对应的矩阵乘 GPU 内核运行耗时显著高于其他轮次。值得留意的是，各轮次输入完全一致，cuBLAS 启发式适配逻辑不会带来差异；CPU 任务下发时序无空隙，也并非剖析工具本身造成的数据失真。

该轨迹点明一个理想化示例里容易忽略的关键点：即便硬件环境、代码、输入数据完全相同，GPU 内核的实际运行耗时也并非固定值。

微调脚本参数，循环运行 20 次并采集每个步骤：

- schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
+ schedule = torch.profiler.schedule(wait=0, warmup=0, active=20, repeat=1)

- for _ in range(5):
+ for _ in range(20):

图 18 同样验证了上述现象：内核代码完全一致，但运算耗时各不相同。造成耗时波动的常见诱因：

• GPU 空闲或加速状态下的时钟频率变化
• 显卡温升带来的性能波动
• 显卡功耗调度策略
• 显卡驱动后台例行维护任务

只看平均耗时的开发者会得出结论：矩阵乘法平均耗时约 1 毫秒（5 次均值为 1084 微秒）；而查看时序轨迹后能发现：除个别硬件扰动时刻外，矩阵乘法大多只需约 580 微秒。两种结论对应的优化思路天差地别，后者才符合真实运行情况。

看看 torch.compile 的实际表现

使用 torch.compile 一直让我们感到惊喜。开发者编写普通 eager PyTorch 代码，PyTorch 会尝试捕获张量密集区域，将其转为计算图、优化，并运行生成代码。默认后端通常是 TorchInductor，整体编译链路如下：

1. TorchDynamo：抓取 Python 执行过程，转化为 FX 计算图
2. AOTAutograd：涉及梯度计算时，构建前向、反向传播计算图
3. Inductor：将上层计算图下沉编译为优化后的 CPU/GPU 代码

本节结合性能剖析轨迹，讲解编译优化细节：

uv run 01_matmul_add.py --size 4096 --warmup --compile

开启编译参数 args.compile 后，执行逻辑变更如下：

def fn(x, w, b):
  return torch.add(torch.matmul(x, w), b)

fn = torch.compile(fn) if args.compile else fn

图 19 中可以看到名为 Torch-Compiled Region: 0/0 的条目，代表当前正在使用编译后的函数。

矩阵乘与加法算子是否融合为单个 GPU 内核？

从图 20 可以提出一个问题：乘法与加法是否真正融合为单次算子运算？

这属于计算图层面的算子融合：Inductor 将 torch.add(torch.matmul(x, w), b) 改写为单次 aten::addmm(b, x, w) 调用。关键细节是：该优化并未生成全新的融合 CUDA 内核。GPU 底层仍沿用动态执行模式使用的同一个 cuBLAS 内核 ampere_bf16_s16816gemm_bf16_128x256_ldg8_f2f_stages_64x3_nn。换言之，本次融合仅发生在算子调度层，而非 GPU 内核层。

PyTorch 原生接口 torch.addmm 本身就整合了矩阵相乘和偏置相加两步运算。可以采集这个函数的性能轨迹，并在评论区分享观察结果。

torch.compile 的运行时架构

理论上理解函数编译会发生什么还不够，同样重要的是在轨迹里看到它实际如何运行。下面看一下反映 torch.compile 运行时架构的 CPU 侧层级：

TorchDynamo Cache Lookup：Dynamo 在这里校验当前调用是否仍匹配已编译版本，包括输入形状、数据类型、设备和张量元数据。如果存在不匹配，Dynamo 会重新编译；这项查询开销每次调用都会产生，即使编译完成后也无法省略。

Torch-Compiled Region：进入编译后版本的包装入口。

AOTDispatcher Runtime Wrapper Prologue：AOTAutograd 的运行时包装器。即便本示例不需要梯度，AOTDispatcher 也始终位于调用栈中，负责张量元数据管理和视图追踪；如果 requires_grad 为真，它还会设置反向传播过程。

**## Call CompiledFxGraph **：实际生成的代码在这里运行。「CompiledFxGraph」后面的字符串是 FX 计算图的内容哈希。三轮有效采样中的哈希一致，说明命中了编译缓存。

可通过哈希值在 /tmp/torchinductor_<user>/fxgraph 路径下找到 Inductor 实际生成的 Triton/C++ 源码，便于进一步调试内核。

CUDA 内核启动次数是否减半？

初看图 21 的 CPU 轨迹时，我们一度高兴地以为每轮只有一次 cudaLaunchKernel 调用，但这与 GPU 轨迹直接矛盾：每轮依然启动了两个内核，分别是 设备到设备内存拷贝（Memcpy DtoD）和 GEMM。回看 CPU 调用栈后可以发现，此前完全遗漏了 cudaMemcpyAsync 下发。

addmm 计算的是 out = α·A·B + β·C。cuBLAS 带偏置加法收尾阶段的 GEMM 会写入目标缓冲区，而这个目标缓冲区需要预先包含偏置。运算收尾（epilogue）可以理解为 GEMM 主计算之后发生的所有附加操作。在深度学习中，激活函数、偏置相加、归一化等都是常见的 GEMM 收尾逻辑，因此 cuBLAS 提供了大量带特定收尾逻辑的 GEMM 内核变体。

使用 torch.compile 的不同 mode 参数会触发不同内核变体。可以自行测试，并在评论区分享观察结果。

因此，Inductor 生成的代码实际执行两步：

• out = copy(C) ← 这是 DtoD 内存拷贝（32 MB，耗时约 33 µs）
• out = α·(A·B) + β·out ← 使用 α=β=1 的 GEMM，将偏置加法融合到写回阶段

数学结果等价，但优化存在隐性开销：额外一次显存拷贝 + 复杂度略高的 GEMM 收尾逻辑。

原本期待的融合是让 x·w + b（这里是 out = α·A·B + β·C）折叠成一个没有额外内存访问的单一内核，但实际并非如此。Inductor 保留了两次访存操作，只是把偏置拷贝表示为内存拷贝，把加法表示为 GEMM 收尾阶段。

真正的融合实现会跳过这次内存拷贝。FlashAttention 这类手写内核会这么做，Inductor 借助 Triton 代码生成也能做到；但针对 4096×4096 bf16 matmul，Inductor 显然判断「使用 cuBLAS，并通过收尾阶段处理偏置」是更合适的路径。

CPU 开销不降反升

对比动态执行与编译运行极易忽略的细节：

编译版本单轮 CPU 耗时约为原生的 2 倍。原因是每次调用都完整经过 Dynamo→AOTAutograd→Inductor 整套调度链路，同时保留原有 aten::addmm 调度开销。编译框架面向多算子大型模型设计，海量算子可平摊单次调用的编译开销；单算子场景下，编译调度反而带来额外性能损耗。

作为课后练习，可以查阅 torch.compile 文档，尝试找到能够降低 CPU 调用开销的 mode 参数。🤗

轨迹阅读速查表

本小节汇总前文出现过的模式：在性能轨迹中看到相应现象时，可以快速对照判断含义。

性能统计表

CPU 时序泳道（CPU lane）

GPU 时序泳道（GPU lane）

任务下发调用链

torch.compile 编译相关

结论

本文从极简的 matmul + add 入手，借它学习如何阅读 PyTorch 剖析器输出。过程中得到的一些分析思路，也适用于更大的工作负载。

本文是《PyTorch 性能剖析》系列的第一站。后续文章会逐渐离开这个双算子示例，走向更大的构建块，并最终分析真实模型。

感谢 Noe Flandre、Suvaditya Mukherjee 和 Vidit Ostwal 审阅本文初稿。