Hugging Face 6 月 11 日发布 PyTorch Profiling 系列第二篇,测试环境里用到了 NVIDIA A100-SXM4-80GB。文章看的不是新模型,也不是新框架,而是一个更容易误读的问题:torch.compile 到底改了什么。

最反常的地方在这里。CPU trace 里一串 aten::ttransposeviewreshape,看起来很忙;但 GPU 上真正跑的 kernel,可能几乎没变。

这对做 PyTorch 推理、训练优化的人很要紧。因为很多性能报告会把 CPU trace 里的 op 数量,直接当成 GPU kernel 数量。这个错一旦带进优化方案,后面很容易走偏。

单个 Linear:bias 已经折进 GEMM,compile 空间很小

nn.Linear 本质上就是:x @ w.T + b

但 PyTorch 不会真的先在 GPU 上复制一份转置后的权重,再单独做一次 bias add。Hugging Face 的 profiler trace 指向的是 aten::addmm。bias 会被放进 GEMM 的 epilogue,在结果写回显存前完成加法。

所以,单个带 bias 的 Linear,本来就已经比较紧。

CPU trace 里出现的 aten::taten::transposeviewreshapeas_strided,大多是元数据操作。它们改的是 shape 和 stride,不等于 GPU 上真的搬了一遍矩阵。

这也是很多人读 profiler 时最容易踩的坑:CPU 侧看到多个 aten op,不等于 GPU 侧 launch 了多个 kernel。

放到 torch.compile 上,结论就很清楚。对单个 nn.Linear,compile 前后 GPU 上的 GEMM kernel 基本相同。它能省的主要是部分 CPU dispatch、view 链路,不是把矩阵乘法本身变成了另一种神仙算法。

对工程团队来说,这会影响动作选择。

如果你只是在一个大 Linear 上套 torch.compile,不要急着把它写进“核心加速手段”。更现实的做法是先确认:耗时到底在 GEMM,还是在 Python 调度和小算子链。

GeGLU MLP:真正能省的是 launch 和 pointwise

GeGLU MLP 的情况不一样。它不是一个孤立 Linear,而是多个 Linear 加上激活和乘法。

在 eager 模式下,原文里的预期路径是:3 个线性层对应 3 个 GEMM,GeLU 是 1 个 pointwise kernel,mul 也是 1 个 pointwise kernel。合起来是 5 个 GPU kernel。

这时 compile 才有更明确的空间。

它不需要改写 3 个 GEMM 的本质。它更可能做的是减少 CPU dispatch,并把 GeLU、mul 这类 pointwise 操作合并成更少的 GPU kernel。

场景CPU trace 里容易看到什么GPU 侧更该看什么工程判断
单个 nn.Linearaten::taten::addmm、view 类操作仍是 GEMM-with-biascompile 收益有限,主要省调度链
GeGLU MLP eager3 个 Linear、GeLU、mul3 个 GEMM + 2 个 pointwise kernel预期约 5 个 GPU kernel
GeGLU MLP compile3 个 aten::mm 加 fused pointwiseGEMM 名称大体相近,pointwise 可能合并收益来自少 launch、少 dispatch

这里的主线不是“compile 有没有用”。

它当然有用,但用处有边界。单个大 GEMM 已经由 cuBLAS、CUTLASS 这类库打磨多年,能挤出的空间有限。多个小算子夹在 GEMM 中间时,融合才更容易产生体感收益。

这也解释了为什么同样是 torch.compile,有人测得明显变快,有人几乎没感觉。shape、dtype、算子组合、后端 kernel 选择,都会改变结果。没有这些条件,单说快或慢,都太粗。

性能报告该怎么改:少数 op,不如看 kernel 名称

读这类 profiler trace,我更在意 kernel 名称,而不是 CPU op 数量。

比如 CUTLASS kernel 名称里常能看到 bf16、tile 形状、tn 等信息。tn 说明 GEMM 读取矩阵布局的方式,kernel 可以按转置视图去读权重,并不需要先生成一个真实转置矩阵。

如果 eager 和 compile 下 kernel 名称逐字一致,至少说明 GPU 核心计算路径没有明显变化。此时把加速归因到“GEMM 被优化了”,就站不住。

更稳的性能报告,应该把几件事拆开写:

要检查的问题看哪里能避免什么误判
CPU 调度是否减少CPU trace、op 调用链把 Python/dispatch 开销误认为 GPU 计算变化
GPU kernel 是否变少CUDA timeline、kernel launch 数把多个 aten op 误当多个 GPU kernel
GEMM 是否真的换了kernel 名称、tile、dtype、tn 等标记把同一个 GEMM 包装成“底层计算升级”
收益是否可迁移不同 shape、batch、dtype 下复测把单点实验写成通用结论

最相关的两类人,动作也不一样。

做模型推理优化的工程师,不该只给 forward 套一层 torch.compile 就交差。更应该先把 trace 分成三层:Python/CPU dispatch、CUDA kernel launch、GEMM 本体。瓶颈在哪一层,决定该调 batch、换 dtype、改图,还是继续看 Inductor 融合。

写性能报告或做迁移评估的人,也要收一下结论强度。若只在单个 Linear 上看到很小变化,就不适合推动大规模迁移。若 MLP、激活、残差、小算子链里 kernel 数明显减少,再谈 compile 收益会更稳。

接下来最该看的不是一个固定加速倍数,而是三个变量:不同 shape 下 cuBLAS 或 CUTLASS 选到的 tile 是否变化;bf16、fp16 等 dtype 是否触发不同 kernel;Inductor 对激活函数、归一化、残差结构能融合到什么程度。

这几个变量不清楚,采购新 GPU、调整 batch size、迁移到 compile,都应该慢半拍。磨刀不误砍柴工,前提是先确认刀钝在哪里。