AOTI

定义

AOTI 指 Ahead-Of-Time Inductor,即 PyTorch 编译栈中的提前编译路径。它由 TorchInductor 派生而来,专门接收经过 Torch.Export 导出的模型,对计算图做优化后产出共享库等工件,这些工件面向非 Python 环境的部署场景,常用于服务器端推理。

这一缩写中的「Ahead-Of-Time」与「Just-in-time」相对:后者在模型首次运行时才编译,前者则把编译提前到构建阶段完成。部分中文技术资料把 AOTI 解释为 Ahead-Of-Time Inference,即围绕推理的提前编译^[2];而 PyTorch 官方文档与推理服务器的说明中,它写作 Ahead-Of-Time Inductor,并以 torch_aoti 作为平台或后端的标识^[3]。

同一缩写在其他语境下另有含义,例如指代医疗器械企业 Advanced Oxygen Therapy 或分析仪器公司 Alpha Omega Technologies,这些用法与人工智能技术无关^[4][5]。

原理

AOTI 的工作流程分三步。第一步用 torch.export.export() 把模型捕获成计算图,该接口对所记录的中问表示给出严格规范,这是后续编译所依赖的前提。第二步调用 torch._inductor.aoti_compile_and_package(),借助 TorchInductor 完成算子融合、代码生成与内核选择,并把结果打包成 PT2 归档格式的文件。第三步在目标环境中加载该包运行,Python 侧使用 torch._inductor.aoti_load_package(),C++ 侧使用 AOTIModelPackageLoader^[6]。

编译能否复用取决于形状信息。若导出时用 Dim 标注了动态维度并给出上下界,编译产物就能在该范围内接受不同批大小的输入;否则推理阶段输入在尺寸、数据类型和步长上须与导出时一致,官方也提供了环境变量用于打开这类运行时检查^[1]。

与运行期即时编译相比,AOTI 把编译器开销移到构建期,运行时只需加载已生成的二进制,因此可以避开即时编译的冷启动延迟;在进程频繁重建的托管环境中,这一差别尤为明显。其整体流程可概括如下:

flowchart LR

A[PyTorch 模型] --> B[torch.export 导出计算图]

B --> C[TorchInductor 编译与优化]

C --> D[打包为 .pt2 / 共享库]

D --> E[Python 加载运行]

D --> F[C++ 加载运行]

发展历程

AOTI 属于 PyTorch 2.x 编译体系的一部分。自 PyTorch 2.0 起,框架形成两条主要编译路径:一条是首次运行时编译并缓存的 torch.compile,另一条是先导出再编译的 torch.export 加 AOTInductor 组合。

早期版本中,编译入口为 torch._export.aot_compile,产物是可在非 Python 环境加载的 .so 共享库;C++ 侧对应 AOTIModelContainerRunnerCuda 与 AOTIModelContainerRunnerCpu 两个运行器^[7]。此后接口调整为 aoti_compile_and_package 与 aoti_load_package,产物格式也改为 .pt2 包^[1]。

2025 年,PyTorch 提出原型接口 torch.library.register_aot_impl,尝试让通过 TVM-FFI 接入的领域专用语言内核也能注册进 AOTI 路径,从而不再局限于即时执行模式^[8]。同一年,社区项目开始在托管 GPU 环境里用 AOTI 编译生成式模型的计算主体并发布为可复用的编译包^[9]。

应用

AOTI 的主要用途是服务端推理部署。官方文档将其定位为面向非 Python 环境的部署手段,典型形态是把模型编译成单个自包含文件后,交给不安装 Python 的 C++ 服务加载执行^[6]。

它与推理服务器和推理优化框架存在集成。英伟达的 Triton 推理服务器可通过 PyTorch 后端加载以 torch_aoti 平台产出的提前编译包来提供模型服务^[3];Torch-TensorRT 则把可转换的子图编成 TensorRT 引擎、其余算子交给 AOTI 生成原生内核,合并为同一个包,推理阶段无需依赖 Torch-TensorRT^[6]。

在面向公众的模型演示方面,AOTI 被用于缩短生成式模型的启动等待。有实践记录显示,把 Flux、Wan、LTX 等模型的 Transformer 部分提前编译后,在托管 GPU 空间上获得约 1.3 至 1.8 倍的提速。相关项目也以 ZeroGPU 与 AOTI 组合发布了多批演示程序^[10]。

局限

AOTI 长期标注为原型状态,官方提示其接口可能发生不向后兼容的改动,使用方需要关注版本迁移成本^[1]。打包能力也存在平台限制,相关文档指出打包目前仅支持 Linux^[6]。

编译与加载必须在同一类设备上完成:在 CUDA 上编译的产物需在 CUDA 上加载,CPU 与 GPU 之间交叉加载不被支持;运行时输入的设备、形状还要与编译时匹配,否则可能出现崩溃、异常或结果错误^[11][12]。

功能覆盖同样有限。以 CUDA 图方式运行模型的接口只存在于旧式运行器,不能经打包模型接口调用,也未在包加载接口中开放^[13]。此外,提前编译只适用于推理场景,无法承担训练期的梯度计算与参数更新。

参见

• PyTorch —— AOTI 所属的深度学习框架。

• TorchInductor —— AOTI 由该编译器派生而来。

• 即时编译 —— 与提前编译相对的另一种编译时机选择。

• TensorRT —— 可与 AOTI 组合打包推理引擎的优化框架。

• 计算图 —— AOTI 导出与优化所作用的对象。