TorchAX
Also known as torchax, tpu-inference backend, PyTorch to JAX translation
In plain English
TorchAX let PyTorch model code run on TPUs by secretly translating each operation into JAX, an approach TorchTPU is now replacing.
Technical definition
TorchAX is a translation layer that intercepts PyTorch ATen operations through the __torch_dispatch__ hook and executes them as JAX operations, backing each torchax.tensor.Tensor with a jax.Array.
Engineering details
vLLM TPU support went through three stages. The first prototype used PyTorch/XLA lazy execution to collect operations into graphs for XLA. The current public tpu-inference backend uses a TPU-optimized JAX model when one exists and otherwise runs the PyTorch model through TorchAX. State such as weights and the KV cache is exposed to JAX explicitly through functionalization so jax.jit can capture each step, and vLLM bypasses its usual torch.compile path because the JAX pipeline already handles graph capture. SGLang-JAX is a separate JAX-native engine that made a similar trade. Google chose JAX because it had more mature TPU primitives and parallelism support at the time.
Why it matters
Translating across two frameworks caused recurring issues with low-level optimization, paged attention, and fitting vLLM’s worker model to TPU execution, and it forced engine features to be reimplemented on the JAX side. Those costs prompted Google, PyTorch, vLLM, and SGLang to build TorchTPU instead.
How to read it in InferenceX
The Pallas kernels developed under TorchAX, including the MoE, GDN, and paged attention work described in the InferenceX Official Preview, carry over to TorchTPU because both stacks can invoke Pallas and JAX kernels. TorchAX itself is scheduled for deprecation once TorchTPU ships.