PyTorch 性能分析（上）：torch.profiler 入门指南

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

无法 profiling 的东西，就无法优化。

无论你是想从大语言模型（LLM）中榨取更多每秒 token 数，还是想削减推理的毫秒级延迟，抑或只是想弄明白为什么你的训练循环跑得比规格表承诺的慢，最终都绕不开 profiling。

问题在于，profiling 的入门门槛很高。trace 是密密麻麻的彩色矩形墙。事件名称令人望而生畏。大多数教程都假设你已经能读懂它们。所以，即使我们知道应该做 profiling，打开一个 trace 也常常感觉像是一件最好留到以后（或者留给别人）做的苦差事。这篇博文，以及它开启的系列文章，就是我们降低这个门槛的尝试。

这是 PyTorch 中的 Profiling 系列的开篇，我们将逐步培养阅读 profiler trace 的技能，并用它来驱动优化。计划如下：

1. 第 1 部分（本文）： 从最简单的操作开始——矩阵乘法后接偏置加法，学习如何阅读 profiler 返回的信息。
2. 第 2 部分： 扩展到 nn.Linear 和一个小的 MLP，利用 trace 来推动优化，并窥探底层的 kernels。
3. 第 3 部分： 结合 transformers，在大语言模型上进行综合实践。

我们从初学者的角度记录这段旅程。除了基本的 PyTorch 知识外，无需任何先决条件。请把它当作一次轻松的阅读，其中穿插着一些"啊哈！"时刻。文章结构有意采用问题驱动：我们打开一个 trace，问"等等，为什么会这样？"，然后追寻答案，直到豁然开朗。读完本文，你应该能了解：

• 如何设置 torch.profiler 以及它实际返回什么，
• 如何阅读 profiler 表格和 trace（CPU 通道、GPU 通道以及两者之间可疑的间隙），
• 从 Python 调用一路到 CUDA kernel 的事件链，
• 当你加上 torch.compile 时，什么发生了变化（以及更有趣的是，什么没有变化）。

在开始之前，有两个定义能让下面的内容读起来更顺畅：

1. GPU kernel 是一个在 GPU 的多个线程上并行运行的程序。
2. CPU 调度并启动这些 kernels。

你通常不需要自己编写 GPU kernel；当你使用 PyTorch 操作时，它会自动被翻译成一个或多个在 GPU 上完成工作的 kernel。

记住这两个概念，让我们开始提问吧。

以下是本文使用的完整脚本：01_matmul_add.py。建议在新标签页中打开此脚本，并逐步阅读代码。我们使用 NVIDIA A100-SXM4-80GB GPU 来运行脚本。

矩阵乘法与加法操作

正如 Dr. Sara Hooker 精辟地指出，就像我们主要由水构成一样，深度神经网络主要由矩阵乘法构成。既然它们如此基础，用其他任何东西来开启我们的 profiling 之旅都将是一种遗憾。

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

矩阵加法与矩阵乘法一起，模拟了权重和偏置在神经元中的交互方式。这个加法（双关语）将帮助我们理解它如何为本文后面的编译铺平道路。

为了进行 profiling，我们将使用 torch.profiler 模块。涉及的步骤如下：

1. 准备好要 profiling 的代码（这里是 def fn，它封装了矩阵乘法和矩阵加法）
2. 注释算法。虽然这完全是可选的，但我们建议这样做。record_function 将我们的函数注释为 matmul_add，这将便于在 trace 中导航（我们稍后会提到）

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

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

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,  # CPU 活动
        torch.profiler.ProfilerActivity.CUDA, # GPU 活动
    ],
  ) as prof:
    # 建议多次运行事件以预热 GPU
    for _ in range(5):
      step()
      prof.step()

1. 导出 profile

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

# profiler trace
prof.export_chrome_trace(trace_path)

profiler 导出两种不同的产物：

1. profiler 表格： 提供算法的统计摘要。它回答"什么花费了最多时间"。这对于找出热点非常有帮助。热点是指花费时间最多的事件，可能是管道的瓶颈，或者被触发了很多次的事件。
2. profiler trace： 提供时间上的执行视图。回答"操作何时以及为何发生"，描绘了 CPU 和 GPU 上发生的活动。当我们想要调查启动的 kernel、启动它们的任何延迟、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

图 1：64 大小矩阵的 matmul add 的 profiler 表格

.txt 文件保存了 profiler 表格。打开文件后，如图 1 所示，你会看到一个很大的表格，第一列由在 profile 作用域内触发的事件组成。

其他列与事件在 CPU、GPU 或 torch.profiler.profile 的 activities 中指定的任何其他设备上花费的时间有关。查看哪些事件花费了最多时间，并尝试直观地理解该事件是否确实应该花费那么多时间。查看"# of Calls"列也很重要，它指示事件被触发了多少次。

趁此机会，我们也谈谈"Self CPU/CUDA"与"CPU/CUDA total"。"Self"列仅测量事件本身内部花费的时间，不包括其子事件。"total"列包括事件及其所有子事件的总和。所以，如果你查看 matmul_add 的"CPU total"，它包含了自身花费的时间加上它触发的子事件的时间。这是一个需要注意的重要细微差别。

如果你查看表格的最后两行，你会注意到 profiler 告诉我们：

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

CPU 时间以 ms 为单位，而 GPU 时间以 us 为单位。换个角度看，GPU 上花费的时间（kernel ampere_bf16_s16816gemm...）不到 CPU 上花费的时间（matmul_add 操作）的 1%。GPU 大部分时间处于空闲状态，这是一个直接的警示信号。发生这种情况的原因是 GPU 可以非常快速地计算一个小型 matmul，因此我们的代码大部分时间都花在准备 kernels、将它们启动到 GPU、发送要相乘的数据以及收集结果上。这个概念被称为开销受限算法。

摆脱这种状态的最简单方法是使用更大的矩阵乘法。

uv run 01_matmul_add.py --size 4096

图 2：4096 大小矩阵的 matmul add 的 profiler 表格

图 2 中的最后两行是：

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

两个时间都以 ms 为单位，这意味着我们仅仅通过增加矩阵乘法的大小就实现了更多的 GPU 时间。如果你查看图 2，你还会注意到，现在大部分 CUDA 时间被 GPU kernel（ampere_bf16_s16816gemm_..）占用，而不是启动它的 CPU 操作（matmul_add）。这意味着我们确实能够从开销受限转向计算受限。

现在我们进入可视化调度链的部分，它存在于 .json 产物中。你可以将它们上传到 Perfetto UI 查看 trace，或者使用 uvx trace-util traces -b traces 直接生成 Perfetto 链接。

64x64 的 traces

图 3：64 大小矩阵的 matmul 和 add 的 profiler trace

在图 3 中，我们看到了矩阵乘法和加法的 profiler trace。这里，条形宽度表示事件的持续时间，垂直嵌套是调用层次结构，CPU 通道表示 CPU 上发生的事件，而 GPU 通道显示实际的 kernel 执行。你可能还会注意到空白区域，它们是等待或空闲时间。

该脚本使用默认配置运行，这些配置是：

• size 64：输入、权重和偏置的大小为 (64, 64)
• dtype bf16：数据类型是 bfloat16
• no compile：我们尚未编译 torch 操作
• no warmup：我们在 profiling 之前没有预热 GPU

对于 Perfetto，我们建议使用键盘以便更快地浏览 trace。你可以使用"W A S D"键来导航 trace。

图 4：PyTorch profiler trace 的 CPU 和 GPU 通道

图 4 中有两个通道，一个用于 CPU 活动，一个用于 GPU 活动。在 CPU 通道中，你会注意到三个 profile 步骤（从 ProfilerStep#2 开始）。这来自于 schedule。

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

wait 跳过嘈杂的初始化（ProfilerStep#0），warmup 在不记录的情况下运行 profiler（ProfilerStep#1），而 active 是显示在 trace 中的内容。你可以在这里的脚本中找到正在使用的 schedule。

让我们戴上侦探帽，调查 trace 并提出一些问题。

为什么 ProfilerStep#2 花费这么长时间？

图 5：ProfileStep#2 明显比后续步骤更宽

在图 5 中，我们注意到 ProfileStep#2 比其他步骤花费更多时间，仔细观察，你会在 matmul_add 注释中看到类似的模式。问题出在注释内部，而不是注释本身：

步骤	matmul_add 开始	aten::matmul 开始	间隙
#2	138.736	366.493	227.757 µs
#3	517.926	523.447	5.521 µs
#4	610.039	614.527	4.488 µs

图 6：record_function("matmul_add") 和 aten::matmul 之间约 228 µs 的死窗口

图 6 中显示的约 228 µs 是进入 record_function("matmul_add") 和 PyTorch 实际调度 aten::matmul 之间的"死窗口"。这可能是由多种原因造成的，包括工作空间分配、cuBLAS（NVIDIA 专有的、GPU 加速的基本线性代数运算库）启发式算法或惰性模块加载。我们可以选择忽略它，或者在 profiling 之前运行更多的预热步骤（这是标准做法）。

在 profiling 中，预热是指在实际 profiling 之前运行事件几次。GPU 完成的准备工作（包括上述几点）是一次性工作，我们不想对其进行 profiling。在我们的示例中，我们有两个预热阶段：一个是在进入 profiler 之前实际循环函数，另一个是在 profiler 内部，通过 warmup 参数实现。在本节中，我们启用了实际迭代以及 schedule。

uv run 01_matmul_add.py --warmup

64x64 带 Warmup 的 Perfetto Trace

图 7：预热后，每个 profile 步骤花费的时间相似

在图 7 中，我们看到每个 profile 步骤花费的时间相似，但这并不意味着我们能够优化一次性开销。我们预热了运行，以便这些开销不被 profiling。我们认为，在不提供解决此问题的提示的情况下突然结束本节对读者是不公平的，所以这里有一个链接，可以阅读有关进一步优化启动开销的内容。

为什么 CPU 和 GPU 通道之间存在约 2.5 ms 的偏移？

图 8：CPU 和 GPU 通道之间约 2.5 ms 的偏移

在图 8 中，我们看到 CPU 和 GPU 通道之间存在大约 2.5 ms 的偏移：这是 CPU 提交 CUDA kernels 到它们实际开始执行之间的延迟。人们可能会认为，预热阶段加上 schedule 的 wait 和 warmup 应该能让 GPU 保持忙碌，并减少这种偏移。

为了揭示真正发生了什么，让我们稍微改变一下 schedule：

- 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：使用 wait=0 和 warmup=0，trace 显示在任何操作之前有一个 Activity Buffer Request

图 9 向我们展示了在任何操作之前，GPU 通道中有一个 Activity Buffer Request。让我们再放大一点。

图 10：在 profile 步骤 1 中，matmul 和 add kernel 之间出现了一个间隙

放大 GPU trace 后，我们注意到 ProfileStep#0（其 CPU trace 在图 10 中不可见）的 matmul 和 add kernel 一个接一个地发生，而 ProfileStep#1 的 kernel 之间有一个窗口。对此最好的解释是缓冲区溢出，并且在 kernel 执行期间发出了另一个缓冲区请求（请求在 GPU VRAM 上分配一些内存）。

排除其他可能性的最佳方法是 profiling 更多迭代，并查看 trace 的其他部分是否出现类似的窗口。为此，我们使用 active=20 运行。

图 11：使用 20 个活跃步骤，间隙只出现一次，确认它是缓冲区请求

如图 11 所示，我们在 ProfileStep#1 中看到了类似的趋势。这与我们之前的发现一致，我们可以安全地得出结论，这确实是另一个缓冲区请求。

事件链

图 12：调度链

在图 12 中，我们看到了嵌套的 CPU 调用。这是一个重要的可视化，你可以了解调度链的真正样子。

我们从 ProfileStep#<id> 开始，它封装了 profiling 步骤。由于我们注释了步骤，我们看到了 matmul_add 行。matmul_add 包含两个 aten 调用，一个用于矩阵乘法，一个用于矩阵加法。

aten::matmul 是 ATen 级别的调度，用户层面的 PyTorch matmul 调用会落到这里。aten::mm 是 2D 矩阵-矩阵乘法后端。

非常有趣的是，如果我们为矩阵添加批次轴，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)

+ # 添加批次大小 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：批量矩阵乘法

在图 13 中，为输入添加批次轴后，aten::matmul 现在封装了一堆其他必需的 CUDA 运行时调用以及 aten::bmm（而不是 aten::mm）。这也暗示了 cuBLAS 为了调度适合程序的最合适的 kernel 而需要执行的启发式算法。

在本文的其余部分，除非另有说明，我们将使用简单的 2D 矩阵。

为什么 matmul 有一个额外的 CUDA 运行时调用？

图 14：在 matmul kernel 启动之前触发了一个 CUDA occupancy 查询

我们注意到，对于 aten::mm，有两个 CUDA 运行时调用，即 cudaOccupancyMaxActiveBlocksPerMultiprocessor（在图 14 中框出）和 cudaLaunchKernel，而对于 aten::add，只有 cudaLaunchKernel。

cudaOccupancyMaxActiveBlocksPerMultiprocessor 是一个规划调用，纯粹在 CPU 端进行。它询问："给定一个 kernel 函数、一个选定的块大小和一个选定的动态共享内存大小，这个 kernel 的多少个块可以同时驻留在一个 SM（流式多处理器）上？"

这就引出了一个问题，为什么我们需要为 matmul 做规划，而不需要为 add 做规划？

要理解这一点，我们必须查看 kernel 的资源占用。如果你点击 GPU kernels，你将能够检查相应 kernel 的资源占用。

图 15：Matmul 占用	图 16：Add 占用

在图 15 中，我们注意到对于矩阵乘法，registers per thread 和 shared memory 是动态的（基于矩阵的大小）。cuBLAS 提供了数百种 kernel 变体，每种都有一个由启发式驱动的启动路径，需要关于硬件容量的运行时信息。occupancy 查询是该启发式算法的一部分。从概念上讲，我们可以将 GPU 加速的 matmul 视为在独立的 tile 上工作：我们使用多少个 tile 以及每个 tile 需要多大取决于矩阵和硬件。现代算法比这复杂得多，但这仍然是一个很好的参考框架。

从图 16 中，我们看到加法的占用显示为 32 个寄存器和零共享内存。这很容易满足。没有什么可查询的，因为没有硬件资源会限制 occupancy。该 kernel 在设计上是资源轻量级的。

你可以在阅读任何 trace 时将其用作快速诊断工具。扫描 CPU 通道中的 cudaOccupancyMaxActiveBlocksPerMultiprocessor。每次出现都标记了一个"重量级、自适应启动的"kernel，通常是 GEMM（通用矩阵乘法）、conv 或类似操作。没有前置 occupancy 查询的 kernel 是 PyTorch 机械式启动的逐元素/规约类 kernel。

为什么 cudaDeviceSynchronize 这么大（约 1.78 ms）？

cudaDeviceSynchronize 会阻塞 CPU，直到此设备上的所有 GPU 工作完成。profiler 在活跃窗口结束时发出此同步以刷新事件。没有它，kernel 计时将会缺失。

一个 1.78 ms 的同步覆盖了 26 µs 的实际 GPU 工作，这告诉你这次运行有 98% 的时间是空闲的。这是典型开销受限的症状。

4096x4096 的 traces

从上面的 profiler 表格分析中我们已经知道，为算法提供更大的矩阵可以使其从开销受限区域转移到计算受限区域。

让我们运行命令并深入研究 trace。

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

为什么同一个 kernel 比其他 kernel 花费更多时间？

图 17：尽管输入相同，一个 matmul kernel 的运行时间比其他步骤长

在图 17 中，我们注意到 ProfileStep#3 的 matmul kernel 在 GPU 上花费的时间比其他步骤长。这一点特别有趣，因为启动的其他 kernel 完全相同，这意味着没有涉及 cuBLAS 启发式算法。没有调度间隙，CPU 启动正常，并且这不是 profiler 的产物。

图 17 中的这个 trace 提出了一个在理想化示例中很容易被忽略的有用观点：kernel 运行时不是常数，即使在相同的硬件环境、运行相同的代码和相同的数据上也是如此。

让我们通过稍微修改脚本来使这一点更具体。我们运行迭代 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：在 20 次迭代中，相同的 matmul kernel 以不同的速度运行

图 18 揭示了类似的发现。虽然每个 kernel 完全相同，但它们的计时不同。不同的计算时间可以归咎于一系列原因：

• GPU 时钟在空闲和升频之间变化
• GPU 发热
• GPU 电源管理
• 驱动端的维护工作

只看平均值的读者会得出结论，matmul 花费了大约 1 ms（5 次的平均值 = 1084 µs）；而查看 trace 的读者会看到，除了 GPU 偶尔"闹情绪"时，matmul 花费大约 580 µs。这是两种截然不同的心智模型，只有一种是正确的。

让我们看看 torch compile 的实际效果

使用 torch.compile 总是让我们感到惊奇。你编写普通的 eager PyTorch 代码，但 PyTorch 会尝试捕获张量密集的区域，将它们转换为图，进行优化，并运行生成的代码。默认后端通常是 TorchInductor，其大致流程是：

1. TorchDynamo 将 Python 执行捕获为 FX 图
2. AOTAutograd 在涉及梯度时准备前向/反向图
3. Inductor 将图降级为优化的 CPU 或 GPU 代码。

在本节中，我们将讨论编译并查看 profiler trace。

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：编译后的区域在 trace 中显示为 TorchDynamo 和 Inductor 帧

在图 19 中，我们看到新的 CPU 行名为 Torch-Compiled Region: 0/0，这指向正在使用的编译函数。

我们是否将 matmul 和 add kernel 融合成了一个？

图 20：编译后的运行调度了一个单一的 aten::addmm

查看图 20，我们问自己，我们是否真的将乘法和加法操作融合在了一起？

这是图级别的算子融合。Inductor 将我们的 torch.add(torch.matmul(x, w), b) 重写为单个 aten::addmm(b, x, w) 调用。这里需要注意的重要事情是，它没有产生一个新的融合 CUDA kernel。实际的 GPU 工作仍然是 ampere_bf16_s16816gemm_bf16_128x256_ldg8_f2f_stages_64x3_nn，与 eager 模式使用的 cuBLAS kernel 相同。所以这里的"融合"是在调度器级别，而不是在 kernel 级别。

PyTorch 提供了 torch.addmm 函数，它分两步完成我们所做的事情，即乘法和加法。我们鼓励读者查看此函数的 trace，并在下面的评论中分享你的观察！

torch.compile 的运行时架构

虽然我们在理论上知道编译函数时会发生什么，但同样重要的是看到它的实际运行。让我们看看反映 torch.compile 运行时架构的 CPU 端层次结构。

TorchDynamo Cache Lookup 是 Dynamo 检查当前调用是否仍然与使用相同输入形状、dtypes、设备和张量元数据编译的内容匹配的地方。如果任何内容不匹配，Dynamo 将重新编译。即使在编译之后，每次调用都会付出这个代价。

Torch-Compiled Region 是"进入"编译版本的包装器。AOTDispatcher Runtime Wrapper Prologue 是 AOT Autograd 的运行时包装器。即使我们这里不需要梯度，AOTDispatcher 也始终在栈中处理张量元数据、视图跟踪，并且如果 requires_grad 为 true，则会设置反向传播。

## Call CompiledFxGraph 是实际生成的代码运行的地方。"CompiledFxGraph"后面的字符串是 FX 图的内容哈希。它在所有三个活跃步骤中都是相同的，确认了缓存命中。

你可以在磁盘上的 /tmp/torchinductor_<user>/fxgraph 下找到生成的代码，以该哈希为键，当你想要阅读 Inductor 实际生成的 Triton/C++ 代码时，这很有用。

CUDA 启动次数是否减少了一半？

图 21：每个编译后的步骤仍然启动两个 GPU kernel，一个 Device-to-Device memcpy 和 GEMM

查看图 21 中的 trace，我们很高兴地注意到每个步骤只有一个 cudaLaunchKernel。这个观察结果与我们之前在 GPU trace 中看到的情况直接矛盾。每个步骤仍然启动了两个 kernel，即 Memcpy DtoD (Device -> Device) 和 GEMM。回到 CPU trace，我们注意到我们完全错过了 cudaMemcpyAsync 调度。

addmm 计算 out = α·A·B + β·C，而 cuBLAS 的带偏置加法 epilogue 的 GEMM 写入一个目标缓冲区，该缓冲区需要已经包含偏置。Epilogue 可以被认为是 GEMM 之后发生的所有操作。在深度学习的世界中，我们不断遇到 GEMM-Epilogue，如激活函数、偏置加法、归一化等等。这就是为什么存在 cuBLAS GEMM-with- kernel 变体的原因。

如果你为 torch.compile 使用不同的 mode，你会注意到启动了不同的 kernel 变体。你可以自己尝试一下，并在下面的评论中分享你的观察！

所以 Inductor 生成的代码执行以下操作：

• out = copy(C) ← 这是 DtoD memcpy（32 MB，花费约 33 µs）
• out = α·(A·B) + β·out ← 使用 α=β=1 的 GEMM，将偏置加法融合到回写中

结果在数学上仍然是相同的。偏置加法不是免费的，我们预先支付了一个 memcpy，加上一个稍微昂贵一点的 GEMM epilogue。

人们可能希望的那种融合，即 x·w + b（这里 out = α·A·B + β·C）坍缩成一个没有额外内存流量的单一 kernel，并没有发生。Inductor 保留了这两个内存访问操作，它只是将偏置复制重新标记为 memcpy，并将加法重新标记为 GEMM epilogue。

一个真正融合的实现会跳过 memcpy。这就是 FlashAttention 风格的手写 kernel 所做的，也是 Inductor 可以通过 Triton 代码生成做到的，但对于一个 4096×4096 bf16 matmul，Inductor 显然认为"使用 cuBLAS，通过 epilogue 设置来做偏置"是最佳路径。

CPU 开销增加了，而不是减少了

这是比较 eager 运行和编译运行时最容易忽略的事情：

步骤	eager 持续时间 (ms)	compile 持续时间 (ms)
#2	0.1	0.2
#3	0.07	0.1
#4	0.07	0.1

编译后每个步骤的 CPU 开销大约是 eager 模式的 2 倍。这是因为每次调用都会遍历完整的 Dynamo > AOTAutograd > Inductor 栈，再加上我们无论如何都会有的 aten::addmm 调度。编译管道是为包含数十个操作的 ML 模型构建的，在这些模型中，每次调用的开销可以被摊销（对于单个操作来说，这是一种税）。

torch.compile 有一个 mode 参数。留给读者作为作业，去阅读文档并提出一个可以降低 CPU 开销的 mode。🤗

Trace 阅读速查表

这是我们讨论过的模式的快速参考。思路是：如果你在 trace 中看到这个，它通常意味着什么。

Profiler 表格

你看到的	通常含义
Self CPU time total ≫ Self CUDA time total（CPU 为 ms，GPU 为 µs）	开销受限。CPU 花在调度上的时间比 GPU 花在计算上的时间多。增大工作负载（更大的矩阵、批量操作）或融合调用。
Self CPU time total ≈ Self CUDA time total，两者均为 ms	计算受限。GPU 是瓶颈，这通常是你想要的。
一个事件主导了 CUDA total	那是你的热点。从那里开始优化。
一个事件有巨大的 # of Calls	即使每次调用很便宜，也可能是潜在瓶颈。检查是否可以融合或批处理。
某行的 CPU total ≫ Self CPU	大部分成本存在于子事件中。深入嵌套事件，而不是父事件。

CPU 通道

你看到的	通常含义
第一个 ProfileStep 比其他宽得多	冷启动开销：工作空间分配、cuBLAS 启发式算法、惰性模块加载。添加预热迭代和/或 schedule 的 warmup 参数。
record_function("...") 开始和其内部的第一个 aten::* 之间有大的间隙	相同的冷启动开销，只是放大了。注释已进入，但调度尚未发生。
cudaLaunchKernel 之前的 cudaOccupancyMaxActiveBlocksPerMultiprocessor	一个重量级、自适应启动的 kernel（GEMM、conv 等）。cuBLAS 正在询问驱动程序一个 SM 上能容纳多少个块，以便选择 kernel 变体。
没有前置 occupancy 查询的 cudaLaunchKernel	一个具有固定、资源轻量级占用的逐元素或规约 kernel。无需规划。
活跃窗口末尾的一个长 cudaDeviceSynchronize	Profiler 刷新事件。其持续时间主要是 GPU 完成待处理工作，而不是真正的 CPU 成本。覆盖微小 GPU 工作的同步是经典的开销受限症状。
你没有编写的 cudaMemcpyAsync	通常是一个隐藏的 Device-to-Device 复制。当 addmm 在 GEMM epilogue 之前用偏置填充其目标缓冲区时很常见。

GPU 通道

你看到的	通常含义
GPU 通道上的 Activity Buffer Request	Profiler 正在分配/重新填充自己的事件缓冲区。第一个通常解释了初始的 CPU↔GPU 通道偏移。
单个步骤中两个 kernel 之间的间隙	可能是执行过程中的另一个缓冲区请求。通过运行更多迭代来确认：如果它只出现一次，那是 profiler 的问题，不是你的代码。
同一个 kernel 在不同步骤中计时不同	GPU 时钟、热管理、电源管理、驱动维护。阅读 trace，而不仅仅是平均值。
名为 ampere_bf16_s16816gemm_... 的 kernel	matmul 的实际 cuBLAS GPU 工作。对于相同的形状/dtypes，kernel 名称在 eager 模式和编译模式下通常是相同的。
GEMM 之前的 Memcpy DtoD	addmm epilogue 的偏置复制。"融合"是在调度器级别，而不是在 kernel 中。

调度链

你看到的	通常含义
ProfileStep#N → <record_function name> → aten::* → aten::mm / aten::bmm / aten::add	标准的嵌套调用层次结构。Self time 排除子事件；Total time 包括它们。
aten::matmul 解析为 aten::mm	2D × 2D 矩阵乘法。
aten::matmul 解析为 aten::bmm（带有额外的 CUDA 运行时调用）	3D+ 张量上的批量 matmul。cuBLAS 做更多启发式工作来选择变体。
aten::addmm(b, x, w) 代替单独的 aten::add + aten::mm 对	调度器级别的算子融合。GPU kernel 仍然是相同的 GEMM，偏置加法被折叠到 epilogue 中。

torch.compile

你看到的	通常含义
CPU 通道中的 Torch-Compiled Region: K/M 行	你在一个编译函数内部。
每个步骤都有 TorchDynamo Cache Lookup	Dynamo 正在验证形状/dtypes/设备与缓存的编译匹配。即使在编译之后，每次调用都会付出这个代价。
即使没有梯度也有 AOTDispatcher Runtime Wrapper Prologue	AOTAutograd 的运行时包装器始终在栈中，处理张量元数据和视图跟踪。
跨步骤具有相同哈希的 ## Call CompiledFxGraph <hash>	生成代码的缓存命中。生成的源代码位于 /tmp/torchinductor_<user>/fxgraph/<hash>。
对于小操作，torch.compile 下每个步骤的 CPU 时间高于 eager 模式	预期行为。Dynamo → AOTAutograd → Inductor 栈是一种税，只有在许多操作上才能摊销。

结论

我们从一个小小的 matmul + add 开始，并以此为借口学习如何阅读 PyTorch profiler。在此过程中，我们掌握了一些可以很好地迁移到更大工作负载的心智模型。这是 Profiling PyTorch 系列的第一站。在接下来的文章中，我们将逐渐离开这个两操作玩具，沿着复杂度的阶梯向上走，查看更大的构建块，并最终查看真实的模型。

感谢 Noe Flandre、Suvaditya Mukherjee 和 Vidit Ostwal 对本文早期草稿的审阅！