Skip to main content
Back to timeline
vLLM BlogSource publication:

vLLM models now serve on Google TPUs without code changes, as a new backend reuses the existing scheduler and KV cache

Synopsis

Inferact and Google released vLLM TorchTPU, a backend that serves model definitions already in vLLM on Google Cloud TPUs without changing model code. Previously, bringing a model to TPU meant writing separate model code; this backend keeps vLLM's command line, OpenAI-compatible server, scheduler and KV cache manager, with TPU-specific work sitting underneath them. It compiles through torch.compile's TPU backend to XLA and uses custom JAX/Pallas kernels for attention, mixture-of-experts and quantized matmul.

Source-provided article image: Introducing vLLM TorchTPU
vLLM

Interpretation

Model definitions already in vLLM can now be served on Google Cloud TPUs, and most models that work in vLLM today need no changes to model code. Previously, moving a model to TPU required separate model code, while this backend reuses vLLM's existing model code. The backend is built on Google's TorchTPU stack; models run through vLLM's standard model runner and are compiled with torch.compile's tpu backend, which lowers to XLA.

The user-facing experience stays the same: the command line, OpenAI-compatible server, scheduler and KV cache manager match those on other hardware. TPU-specific work sits underneath these components rather than replacing them. Each TPU device gets its own worker process, device communication goes through torch.distributed, and the scheduler, workers, KV cache manager and API server are unmodified upstream vLLM.

Complex operations such as attention, mixture-of-experts and quantized matmul are handled by hand-tuned Pallas kernels. These kernels are registered as custom ops and called from inside the compiled graph. The kernels cover Ragged Paged Attention v3 for MHA/GQA, dense and sparse MLA for DeepSeek, Kimi and GLM, Gated DeltaNet and Kimi Delta Attention for hybrid linear-attention models, fused mixture-of-experts with expert parallelism, and FP8, INT8, MXFP4 and NVFP4 quantized matmuls.

Perspective

The demonstrated scope is that, after installing the plugin on a Cloud TPU VM with Python 3.12, vLLM picks it up automatically through the vllm.platform_plugins entry point, and most models supported by vLLM can be served on TPU. TPU v7x (Ironwood) is the recommended hardware and is tested in CI on every pull request; v6e, v5e and v5p are experimental. First startup compiles TPU graphs for the configured shapes and can take a few minutes, and VLLM_XLA_CACHE_PATH reuses compiled artifacts across restarts. The intended reader is an engineer already serving models with vLLM who wants to place that load on TPU.

Support is defined by the validated-checkpoint list and nightly support matrices rather than all vLLM models, and v6e, v5e and v5p remain labeled experimental. First startup requires minutes of graph compilation, with a cache path for reuse across restarts. Roadmap items such as automated CI onboarding new open-weight models within 24 hours, routing, and diffusion model serving are not listed as delivered capabilities.