先记住一句话
编译器不会凭空创造硬件性能:它主要消除 Python/launch overhead、融合 memory-bound 操作并选择更好代码;graph break、动态 shape 和编译时间决定净收益。
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. 评估方法
- 记录 eager、首次 compiled、warm compiled 三种 latency;
- 同时测训练 step 或完整 request,不只 isolated module;
- 统计 graph breaks、compiled graphs、recompiles、cache hit;
- 比较峰值显存与数值;
- 用 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。