torch.compile
Also known as torch.compile, TorchDynamo, AOTAutograd, Inductor
In plain English
torch.compile is the PyTorch entry point that captures a model into a graph and hands it to a compiler backend instead of running each operation one at a time.
Technical definition
torch.compile is the PyTorch API that uses TorchDynamo to capture Python-level model code into an FX graph, AOTAutograd to trace forward and backward, and a pluggable backend such as Inductor or XLA to generate device code.
Engineering details
On NVIDIA GPUs, vLLM uses torch.compile with Inductor, which emits Triton kernels. On TPU, TorchTPU keeps the same entry point but swaps the backend: the FX graph is lowered to StableHLO and compiled by XLA. Google made that compiler choice explicit, XLA rather than Inductor and Triton. The earlier TorchAX path bypassed torch.compile entirely because jax.jit already performed graph capture and compilation. Compiled graphs are specialized on tensor shapes, so serving engines bucket request shapes to limit recompilation.
Why it matters
Keeping torch.compile as the entry point means TPU-specific compilation hides behind an API vLLM and SGLang already call. Engine features written against the PyTorch compile path carry over, while shape sensitivity from XLA means bucketing and padding still need TPU-specific tuning.
How to read it in InferenceX
The Qwen3.5 397B numbers in the TPU InferenceX preview run through torch.compile into XLA. Low-concurrency tuning that bucketed request metadata by active requests rather than the configured maximum cut GDN scheduling overhead from 283 to 97 microseconds and raised 8k1k throughput from 2,328 to 2,516 tokens per chip per second at concurrency 64.