| Documentation | Blog | User Forum | Developer Slack (#sig-tpu) |
vLLM TorchTPU is the PyTorch-native TPU backend for vLLM. It is an out-of-tree vLLM platform plugin, packaged as vllm-torchtpu, that runs vLLM's standard PyTorch model definitions on Google Cloud TPUs through TorchTPU (torch_tpu, via torch.compile(backend="tpu")) and a library of custom JAX/Pallas kernels for attention, Mixture-of-Experts, and quantized matmul.
The goal is simple: the same vLLM you already use — same CLI, same OpenAI-compatible server, same model definitions — running at TPU speed with no model code changes.
[!NOTE] Looking for the JAX backend? The original TPU backend for vLLM,
tpu-inference, has been rebranded as vLLM JAX. It remains available for serving pure JAX model implementations on TPU.vllm-torchtpuis the recommended path for serving PyTorch model definitions on TPU.
| TPU generation | Status |
|---|---|
| TPU v7x (Ironwood) | Recommended — validated in CI on every pull request and nightly |
| TPU v6e / v5e / v5p | Experimental — not covered by CI |
See Cloud TPU Setup for project setup, quota, and provisioning commands.
vllm-torchtpu requires Python 3.12 on a Cloud TPU VM. Install the plugin and a matching vLLM with a single command:
pip install \
--extra-index-url https://download.pytorch.org/whl/cpu \
vllm-torchtpu
For uv, editable source installs, cross-repo development against a local vllm checkout, and Docker images, see the full Installation guide and Docker guide.
Start an OpenAI-compatible server exactly as you would on any other vLLM platform:
vllm serve meta-llama/Llama-3.2-1B-Instruct \
--tensor-parallel-size 1 \
--max-model-len 4096
Then send a request:
A stage may instead span a prefill context-parallel group
(--prefill-context-parallel-size), which splits each chunk's attention over
the stage's cores; MoE models then need --enable-expert-parallel. This needs
a torch_tpu that compiles a stage's kernels and collectives for the stage's
devices only (the 2026-09-28 build of google-pytorch/torch_tpu#4387); on a
torch_tpu that builds world-size meshes only, the platform refuses such stages.
A two-core stage must not be the two cores of one chip: set
VLLM_TPU_DEVICE_ORDER to the device each local rank binds, e.g.
0,2,1,3,4,6,5,7 for four stages of two cores on one host.
curl http://localhost:8000/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "meta-llama/Llama-3.2-1B-Instruct",
"messages": [{"role": "user", "content": "Explain how TPUs accelerate matrix multiplication."}],
"max_tokens": 128
}'
[!TIP] The first startup compiles TPU graphs for the configured shapes and can take several minutes. Set
VLLM_XLA_CACHE_PATHto persist compiled artifacts across restarts — see Compilation & Warmup Cache.
The Quickstart covers offline batch inference with vllm.LLM, multi-device serving with TP/EP, and multimodal examples.
vllm-torchtpu can serve most PyTorch models supported by vLLM, with first-class TPU kernel, quantization, and parallelism support for Qwen3 / Qwen3.5, DeepSeek-V3 / R1 / V4, Kimi-K2.5 / K3, GLM-5, Llama 3 / 4, GPT-OSS, Gemma 3 / 4, Mixtral, Qwen-VL and Gemma VL families, and diffusion models via vllm-omni.
We welcome contributions of all kinds — bug reports, model and feature requests, documentation fixes, and code.
ready-label CI gate, and the PR lifecycle. Good first issues are a great place to begin.Developer references: Testing & CI · Profiling · Benchmarking · Coding guidance · Optimization best practices
#sig-tpu).
| Documentation | Blog | User Forum | Developer Slack (#sig-tpu) |
vLLM TorchTPU is the PyTorch-native TPU backend for vLLM. It is an out-of-tree vLLM platform plugin, packaged as vllm-torchtpu, that runs vLLM's standard PyTorch model definitions on Google Cloud TPUs through TorchTPU (torch_tpu, via torch.compile(backend="tpu")) and a library of custom JAX/Pallas kernels for attention, Mixture-of-Experts, and quantized matmul.
The goal is simple: the same vLLM you already use — same CLI, same OpenAI-compatible server, same model definitions — running at TPU speed with no model code changes.
[!NOTE] Looking for the JAX backend? The original TPU backend for vLLM,
tpu-inference, has been rebranded as vLLM JAX. It remains available for serving pure JAX model implementations on TPU.vllm-torchtpuis the recommended path for serving PyTorch model definitions on TPU.
| TPU generation | Status |
|---|---|
| TPU v7x (Ironwood) | Recommended — validated in CI on every pull request and nightly |
| TPU v6e / v5e / v5p | Experimental — not covered by CI |
See Cloud TPU Setup for project setup, quota, and provisioning commands.
vllm-torchtpu requires Python 3.12 on a Cloud TPU VM. Install the plugin and a matching vLLM with a single command:
pip install \
--extra-index-url https://download.pytorch.org/whl/cpu \
vllm-torchtpu
For uv, editable source installs, cross-repo development against a local vllm checkout, and Docker images, see the full Installation guide and Docker guide.
Start an OpenAI-compatible server exactly as you would on any other vLLM platform:
vllm serve meta-llama/Llama-3.2-1B-Instruct \
--tensor-parallel-size 1 \
--max-model-len 4096
Then send a request:
A stage may instead span a prefill context-parallel group
(--prefill-context-parallel-size), which splits each chunk's attention over
the stage's cores; MoE models then need --enable-expert-parallel. This needs
a torch_tpu that compiles a stage's kernels and collectives for the stage's
devices only (the 2026-09-28 build of google-pytorch/torch_tpu#4387); on a
torch_tpu that builds world-size meshes only, the platform refuses such stages.
A two-core stage must not be the two cores of one chip: set
VLLM_TPU_DEVICE_ORDER to the device each local rank binds, e.g.
0,2,1,3,4,6,5,7 for four stages of two cores on one host.
curl http://localhost:8000/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "meta-llama/Llama-3.2-1B-Instruct",
"messages": [{"role": "user", "content": "Explain how TPUs accelerate matrix multiplication."}],
"max_tokens": 128
}'
[!TIP] The first startup compiles TPU graphs for the configured shapes and can take several minutes. Set
VLLM_XLA_CACHE_PATHto persist compiled artifacts across restarts — see Compilation & Warmup Cache.
The Quickstart covers offline batch inference with vllm.LLM, multi-device serving with TP/EP, and multimodal examples.
vllm-torchtpu can serve most PyTorch models supported by vLLM, with first-class TPU kernel, quantization, and parallelism support for Qwen3 / Qwen3.5, DeepSeek-V3 / R1 / V4, Kimi-K2.5 / K3, GLM-5, Llama 3 / 4, GPT-OSS, Gemma 3 / 4, Mixtral, Qwen-VL and Gemma VL families, and diffusion models via vllm-omni.
We welcome contributions of all kinds — bug reports, model and feature requests, documentation fixes, and code.
ready-label CI gate, and the PR lifecycle. Good first issues are a great place to begin.Developer references: Testing & CI · Profiling · Benchmarking · Coding guidance · Optimization best practices
#sig-tpu).