torch.compile

定义

torch.compile 是 PyTorch 2.0(2023 年发布)引入的一行式编译加速接口,把动态图 Python 代码即时编译为优化后的内核,兼顾易用性与性能。

原理

其技术栈分三层:TorchDynamo 在 Python 字节码层安全地捕获计算图,遇到无法处理的代码自动断图回退;AOTAutograd 同时追踪前向与反向图;默认后端 TorchInductor 做算子融合等优化,并用 OpenAI Triton 生成 GPU 内核、用 C++/OpenMP 生成 CPU 代码。

应用

训练和推理均可使用,对小算子多、Python 开销大的模型提速明显;推理部署中常与 CUDA Graph、量化配合,vLLM、SGLang 等框架也逐步集成 torch.compile 优化路径。

优缺点

优点是几乎不改代码、保留动态图调试体验;缺点是首次编译耗时、频繁断图或动态形状会削弱收益。