torch.compile
也称为 torch.compile、TorchDynamo、AOTAutograd、Inductor
先用大白话
torch.compile 是 PyTorch 的编译入口,把模型捕获成图交给编译器后端,而不是逐个执行每个操作。
技术定义
torch.compile 是 PyTorch 的 API,用 TorchDynamo 把 Python 层的模型代码捕获为 FX 图,用 AOTAutograd 追踪前向和反向,再由 Inductor 或 XLA 等可插拔后端生成设备代码。
工程细节
在 NVIDIA GPU 上,vLLM 使用 torch.compile 配合 Inductor,输出 Triton kernel。在 TPU 上,TorchTPU 保留同一入口但更换后端:FX 图被降级为 StableHLO,由 XLA 编译。Google 明确了这一编译器选择,用 XLA 而不是 Inductor 和 Triton。此前的 TorchAX 路径完全绕过 torch.compile,因为 jax.jit 已经负责图捕获和编译。编译后的图按张量形状特化,因此服务引擎需要对请求形状分桶以限制重编译。
为什么重要
保留 torch.compile 作为入口,意味着 TPU 专用的编译藏在 vLLM 和 SGLang 已经在调用的 API 后面。针对 PyTorch 编译路径编写的引擎功能可以直接沿用,而 XLA 的形状敏感性意味着分桶和 padding 仍然需要针对 TPU 调优。
如何在 InferenceX 中解读
TPU InferenceX 预览中的 Qwen3.5 397B 数据都经过 torch.compile 进入 XLA。低并发调优把请求元数据按活跃请求数而非配置上限分桶,让 GDN 调度开销从 283 微秒降到 97 微秒,8k1k 吞吐量在并发 64 下从每芯片每秒 2,328 token 提高到 2,516。