{"slug": "pytorch-native-tpu-back-end-for-vllm", "title": "PyTorch-native TPU back end for vLLM", "summary": "VLLM released vllm-torchtpu, a PyTorch-native TPU backend packaged as an out-of-tree vLLM platform plugin that runs vLLM's standard PyTorch model definitions on Google Cloud TPUs via TorchTPU (torch_tpu, through torch.compile(backend=\"tpu\")) plus custom JAX/Pallas kernels for attention, Mixture-of-Experts and quantized matmul. The plugin requires Python 3.12 on a Cloud TPU VM, is recommended and CI-validated on TPU v7x (Ironwood) while TPU v6e/v5e/v5p remain experimental and outside CI, and keeps the same CLI, OpenAI-compatible server and model definitions with no model code changes. The release also rebrands the original TPU backend tpu-inference as vLLM JAX, which remains available for pure JAX model implementations, with vllm-torchtpu now the recommended path for PyTorch models on TPU.", "body_md": "| [**Documentation**](https://vllm-torchtpu.readthedocs.io/en/latest/) | [** Blog**](https://blog.vllm.ai/) | [** User Forum**](https://discuss.vllm.ai/c/hardware-support/google-tpu-support/27) | [** Developer Slack**](https://slack.vllm.ai) (#sig-tpu) |\n\n**vLLM TorchTPU** is the PyTorch-native TPU backend for [vLLM](https://github.com/vllm-project/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.\n\nThe 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.**\n\n- **PyTorch first.** Models run through vLLM's native PyTorch runner; no re-implementation in another framework is required to serve on TPU.\n- **TPU-tuned kernels.** Ragged Paged Attention, dense and sparse MLA, Gated DeltaNet, fused MoE (including SparseCore offload), and FP8/INT8/MXFP4 quantized matmuls written in Pallas.\n- **Scale-out built in.** TP, DP, PP, EP, and prefill/decode context parallelism compose across single-host and multi-host TPU slices, with disaggregated prefill/decode serving.\n- **Standard vLLM experience.** Speculative decoding, prefix caching, multi-LoRA, structured output, and multimodal inputs work through the normal vLLM flags.\n\nNote\n\n**Looking for the JAX backend?** The original TPU backend for vLLM, [`tpu-inference`](https://github.com/vllm-project/tpu-inference), has been rebranded as **vLLM JAX**. It remains available for serving pure JAX model implementations on TPU. `vllm-torchtpu` is the recommended path for serving PyTorch model definitions on TPU.\n\n| TPU generation | Status | \n|---|---|\n| **TPU v7x (Ironwood)** | Recommended — validated in CI on every pull request and nightly | \n| TPU v6e / v5e / v5p | Experimental — not covered by CI | \n\nSee [Cloud TPU Setup](https://vllm-torchtpu.readthedocs.io/en/latest/getting_started/tpu_setup/) for project setup, quota, and provisioning commands.\n\n`vllm-torchtpu` requires **Python 3.12** on a Cloud TPU VM. Install the plugin and a matching vLLM with a single command:\n\n```\npip install \\\n  --extra-index-url https://download.pytorch.org/whl/cpu \\\n  vllm-torchtpu\n```\n\nFor `uv`, editable source installs, cross-repo development against a local `vllm` checkout, and Docker images, see the full [Installation guide](https://vllm-torchtpu.readthedocs.io/en/latest/getting_started/installation/) and [Docker guide](https://vllm-torchtpu.readthedocs.io/en/latest/getting_started/docker/).\n\nStart an OpenAI-compatible server exactly as you would on any other vLLM platform:\n\n```\nvllm serve meta-llama/Llama-3.2-1B-Instruct \\\n  --tensor-parallel-size 1 \\\n  --max-model-len 4096\n```\n\nThen send a request:\n\nA stage may instead span a prefill context-parallel group\n(`--prefill-context-parallel-size`), which splits each chunk's attention over\nthe stage's cores; MoE models then need `--enable-expert-parallel`. This needs\na torch_tpu that compiles a stage's kernels and collectives for the stage's\ndevices only (the 2026-09-28 build of [google-pytorch/torch_tpu#4387](https://github.com/google-pytorch/torch_tpu/pull/4387)); on a\ntorch_tpu that builds world-size meshes only, the platform refuses such stages.\nA two-core stage must not be the two cores of one chip: set\n`VLLM_TPU_DEVICE_ORDER` to the device each local rank binds, e.g.\n`0,2,1,3,4,6,5,7` for four stages of two cores on one host.\n\n```\ncurl http://localhost:8000/v1/chat/completions \\\n  -H \"Content-Type: application/json\" \\\n  -d '{\n    \"model\": \"meta-llama/Llama-3.2-1B-Instruct\",\n    \"messages\": [{\"role\": \"user\", \"content\": \"Explain how TPUs accelerate matrix multiplication.\"}],\n    \"max_tokens\": 128\n  }'\n```\n\nTip\n\nThe first startup compiles TPU graphs for the configured shapes and can take several minutes. Set `VLLM_XLA_CACHE_PATH` to persist compiled artifacts across restarts — see [Compilation & Warmup Cache](https://vllm-torchtpu.readthedocs.io/en/latest/user_guide/compilation_cache/).\n\nThe [Quickstart](https://vllm-torchtpu.readthedocs.io/en/latest/getting_started/quickstart/) covers offline batch inference with `vllm.LLM`, multi-device serving with TP/EP, and multimodal examples.\n\n`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`.\n\n- **[Supported Models & Features](https://vllm-torchtpu.readthedocs.io/en/latest/user_guide/supported_models/)** — validated model families, feature status, and quantization formats.\n- **[Nightly support matrices](https://github.com/vllm-project/vllm-torchtpu/blob/main/support_matrices/nightly)** — TPU v7x unit, correctness, and benchmark status, updated automatically from nightly CI.\n\nWe welcome contributions of all kinds — bug reports, model and feature requests, documentation fixes, and code.\n\n- **🐞 Report bugs or request features:** open a[new issue](https://github.com/vllm-project/vllm-torchtpu/issues/new/choose) .\n- **👀 Review pull requests:** help us keep quality high by reviewing[open PRs](https://github.com/vllm-project/vllm-torchtpu/pulls) .\n- **💻 Write code:** start with the**[Contributing Guide](https://github.com/vllm-project/vllm-torchtpu/blob/main/CONTRIBUTING.md)** , which covers environment setup, DCO sign-off, the`ready` -label CI gate, and the PR lifecycle.[Good first issues](https://github.com/vllm-project/vllm-torchtpu/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22) are a great place to begin.\n\nDeveloper references: [Testing & CI](https://vllm-torchtpu.readthedocs.io/en/latest/developers_guide/testing_and_ci/) · [Profiling](https://vllm-torchtpu.readthedocs.io/en/latest/developers_guide/profiling/) · [Benchmarking](https://vllm-torchtpu.readthedocs.io/en/latest/developers_guide/benchmarking/) · [Coding guidance](https://vllm-torchtpu.readthedocs.io/en/latest/developers_guide/coding_guidance/) · [Optimization best practices](https://vllm-torchtpu.readthedocs.io/en/latest/developers_guide/optimization_best_practices/)\n\n- For technical questions and feature requests, open a [GitHub issue](https://github.com/vllm-project/vllm-torchtpu/issues) .\n- For discussion with fellow users, use the [TPU support topic on the vLLM Forum](https://discuss.vllm.ai/c/hardware-support/google-tpu-support/27) .\n- For coordinating contributions and development, join the [vLLM Developer Slack](https://slack.vllm.ai) (`#sig-tpu` ).", "url": "https://wpnews.pro/news/pytorch-native-tpu-back-end-for-vllm", "canonical_source": "https://github.com/vllm-project/vllm-torchtpu", "published_at": "2026-10-19 22:34:38+00:00", "updated_at": "2026-10-19 22:45:54.803310+00:00", "lang": "en", "topics": ["ai-infrastructure", "ai-tools", "large-language-models", "developer-tools", "mlops"], "entities": ["vLLM", "vllm-torchtpu", "TorchTPU", "Google Cloud TPUs", "TPU v7x (Ironwood)", "vLLM JAX", "tpu-inference", "PyTorch"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/pytorch-native-tpu-back-end-for-vllm", "markdown": "https://wpnews.pro/news/pytorch-native-tpu-back-end-for-vllm.md", "text": "https://wpnews.pro/news/pytorch-native-tpu-back-end-for-vllm.txt", "jsonld": "https://wpnews.pro/news/pytorch-native-tpu-back-end-for-vllm.jsonld"}}