VLLM TPU: A New Unified Back End Supporting PyTorch and Jax on TPU (blog.vllm.ai)

vLLM TPU now ships a unified backend powered by tpu-inference that lowers all models through JAX (even PyTorch model definitions via Torchax), giving a single, high-performance compilation path into XLA. The release replaces the older PTXLA flow with a JAX-first lowering, yielding ~20% higher throughput out of the box and dramatic end-to-end gains (e.g., 3.6x for Llama 3.1-8B on v6e-1 and 2.1x for Llama 3.1-70B on v6e-8), and overall 2–5x improvement vs the Feb 2025 TPU prototype. vllm-tpu is a single install path and will prefer TPU-optimized models in tpu-inference, falling back to vLLM upstream models lowered via Torchax when needed. Key technical advances include SPMD as the default programming model (enabling compiler-driven partitioning and overlapping comm/compute), and Ragged Paged Attention v3 (RPA v3): a flexible, fused attention kernel that supports arbitrary head dims, quantization dtypes, and TP configurations, compiles into prefill/decode/mixed sub-kernels, hides KV scatter latency, and adds ~10% throughput over RPA v2 on Trillium (v6e). The stack also supports prefix caching, chunked prefill, quantized matmuls/attention/KV cache, multimodal inputs, and various TPU-friendly kernels. Implications: simpler cross-framework deployment, more time for kernel and compiler optimization, and an OSS foundation for future TPU-optimized work (XL MoE, MLA, sparse offload, advanced speculative decoding).