先记住一句话

编译器不会凭空创造硬件性能:它主要消除 Python/launch overhead、融合 memory-bound 操作并选择更好代码;graph break、动态 shape 和编译时间决定净收益。

capture Python graphguards/specializelower IRfuse/codegencompile/cacherun and inspect graph breaks
Pythontorch.compile
Dynamo
FX graphguards
AOTAutograd
forward/backward graphsfunctional IR
Inductor
kernelsTriton/C++
Graph break 把 pipeline 切成多个 compiled regions,并重新暴露 Python 和 launch overhead。

1. PyTorch 编译路径

Python eager program
  ↓ TorchDynamo 捕获 bytecode 执行
FX graph + guards
  ↓ AOTAutograd 拆 forward/backward
ATen graph
  ↓ Inductor 优化、融合、调度
Triton / C++ / vendor-library kernels

这是理解模型,不是保证每个 graph 都严格走同一形式。library call 可能保留为外部高性能 kernel,小 pointwise/reduction 则适合生成 fused code。

2. Graph capture 与 guards

Dynamo 观察 Python 执行并为控制流、类型、shape 等假设生成 guards。下一次输入满足 guard 就复用 compiled artifact;不满足可能重新捕获/编译。shape 每步乱变会制造 variant storm,因此要看真实 recompilation 日志,而不是只看第二次固定输入。

3. Fusion 为什么有效

read x → op A → write t → read t → op B → write y
可融合为:read x → A → B → write y

省掉中间 tensor 与 launch,尤其适合 pointwise chains、normalization 局部 reduction、activation/bias。融合过大也可能增加 registers、降低 occupancy 或阻碍 vendor GEMM,所以不是越多越好。

4. Graph break

不支持的 Python/C extension、依赖 tensor data 的 Python branch、I/O 或 side effect 可能中断 graph。break 两侧仍可分别编译,但小碎图会损失融合并增加 overhead。先用 explain/log 找 break 原因,再判断改代码、注册 custom op,还是保留 eager island。

5. 动态 shape

完全静态可生成专用代码;symbolic dynamic shape 能复用范围内的实现,但 codegen/tuning 更保守。常用策略是 bucket sequence/image sizes、对少数热点 shape 专门化,其余走 general fallback。padding 的额外计算要与减少重编译和更规则 kernel 的收益一起算。

6. Custom op 是优化边界

没有 meta/FakeTensor 信息的 custom op 对编译器近似黑盒;它仍可被调用,但编译器难以推断 shape 或穿过它融合。可注册 decomposition,把复合 op 展开为编译器理解的 primitives;也可保留高性能 opaque kernel,只提供准确 metadata。两种选择取决于谁能生成更好的实现。

7. 编译与 CUDA Graphs 不同

图编译CUDA Graph replay
主要目标改写/融合/生成 kernel压低重复 launch 的 CPU overhead
要求程序可捕获并满足 guards地址/控制流等 replay 条件稳定
改变 kernel可以本身不改变

二者可叠加:先编译出 kernels,再 capture 稳定执行段。

8. 冷启动与缓存

第一次调用可能包含 capture、codegen、NVRTC/assembler 和 autotune。离线服务可 warm up 热门 buckets;短命进程/交互 workload 需显式计算 compile amortization。cache key 与 framework、GPU architecture、driver、flags 和 shapes 有关,部署不能假定任意环境通用。

9. 评估方法

  1. 记录 eager、首次 compiled、warm compiled 三种 latency;
  2. 同时测训练 step 或完整 request,不只 isolated module;
  3. 统计 graph breaks、compiled graphs、recompiles、cache hit;
  4. 比较峰值显存与数值;
  5. 用 profiler 确认 kernel 数、HBM traffic 和 CPU gaps 真减少。

10. 四个 compiler 收益手算

例 1:launch fusion

10 个 kernels 各 compute 8 μs、launch 5 μs,总 130 μs。融合成 2 个 kernels、总 compute 70 μs,launch 10 μs,总 80 μs,speedup=1.625×

例 2:cold-start 摊销

Compile 用 12 s,eager step=100 ms、compiled=70 ms,每步省 30 ms;break-even steps=12/0.03=400。只跑 100 steps 时净收益为负。

例 3:shape recompilation

输入长度出现 128/256/512 三种,每种 compile 4 s,总 cold cost 12 s。若通过 dynamic guard 只编一次 6 s,省 6 s,但 generated kernel 可能不如专用 shape 快。

例 4:graph-break overhead

原图被 4 个 breaks 切成 5 regions,每边回 Python/dispatch 20 μs,额外约 80 μs/iteration。1M iterations 就是 80 s,值得先修高频 break。

常见误解:

torch.compile 不是一个固定“加速百分比”开关。模型结构、shape 稳定性、已有 fused libraries、Python control flow 和运行时长都会改变收益。

自测

1. guard 的作用是什么?

记录 compiled graph 有效所依赖的输入/程序假设;满足时安全复用,不满足时回退或重编译。

2. 为什么过度 fusion 可能更慢?

更大 kernel 可能占更多寄存器/shared memory、降低 occupancy,或破坏已高度优化的库调用边界。

3. 为什么只报告第二次调用不完整?

它忽略首次编译和实际 shape 变化导致的重编译;只有长时间重复运行才可能完全摊薄。

官方资料

PyTorch torch.compile tutorial介绍编译入口、Dynamo/Inductor 与 graph breaks;Custom Operators tutorial说明 custom op 如何接入 torch.compile