Vorstellung von vLLM TorchTPU

Heute veröffentlichen wir vLLM TorchTPU (vllm-torchtpu), ein neues Backend für den Betrieb von vLLM auf Google Cloud TPUs. Es basiert auf Googles TorchTPU Stack und führt die Modelldefinitionen aus, die bereits in vLLM vorhanden sind, sodass die meisten Modelle, die heute in vLLM funktionieren, auf TPU ohne Änderungen am Modellcode bereitgestellt werden können.

Diese Arbeit geht hervor aus dem Ingenieurspartnerschaft zwischen Inferact und Google den wir letzten Monat angekündigt haben. Unsere Ziele für das Projekt sind:

1. Es einfach machen, neue Modelle auf TPU zu bringen, indem vLLMs vorhandener Modellcode wiederverwendet wird.
2. Die TPU-Inferenzleistung in Open Source vorantreiben, mit Pallas-Kernels für Attention, Mixture-of-Experts und quantisiertes Matmul.
3. Das vLLM-Erlebnis beibehalten: dieselbe CLI, derselbe OpenAI-kompatible Server, derselbe Scheduler und derselbe KV-Cache-Manager, die Sie auf anderer Hardware verwenden.

Wie es funktioniert

vllm-torchtpu ist ein vLLM-Plattform-Plugin außerhalb des Quellbaums. Nach der Installation wird es von vLLM automatisch über den vllm.platform_plugins -Einstiegspunkt erkannt.

Modelle laufen über vLLMs Standard-Model Runner und werden mit torch.compile(backend="tpu")kompiliert, was zu XLA herunterbricht. Jedes TPU-Gerät erhält seinen eigenen Worker-Prozess, auf die gleiche Weise, wie vLLM auf GPUs läuft, und die Kommunikation zwischen den Geräten erfolgt über torch.distributed. Das bedeutet, dass Scheduler, Worker, KV-Cache-Manager und API-Server alles unverändertes Upstream-vLLM sind. Die TPU-spezifische Arbeit liegt darunter, hauptsächlich in einer Reihe von JAX/Pallas-Kernels, die als Custom Ops registriert und aus dem kompilierten Graphen heraus aufgerufen werden.

Erste Schritte

1. Installation

   pip install vllm-torchtpu

   Dies erfordert Python 3.12 auf einer Cloud TPU VM und installiert ein passendes vLLM, torch, torch_tpuund libtpu neben dem Plugin. Der Installationsleitfaden behandelt auch uv, Source-Builds und Docker.

2. Ein Modell ausliefern

   vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 \
     --tensor-parallel-size 8 \
     --enable-expert-parallel \
     --kv-cache-dtype fp8 \
     --max-model-len 16384 \
     --async-scheduling

   Der erste Start kompiliert TPU-Graphen für die konfigurierten Shapes, was einige Minuten dauern kann. Setzen Sie VLLM_XLA_CACHE_PATH , um kompilierte Artefakte über Neustarts hinweg wiederzuverwenden (siehe Kompilierung & Warmup-Cache).

Speculative Decoding, Prefix Caching, Multi-LoRA, strukturierte Ausgaben und multimodale Eingaben werden alle mit den üblichen vLLM-Flags aktiviert. Siehe die quickstart für weitere Beispiele, einschließlich Offline-Batch-Inferenz, multimodale Modelle und Multi-LoRA.

Kernels

vllm-torchtpu verwendet benutzerdefinierte Pallas-Kernels für komplexe Operationen, die von manueller Feinabstimmung auf TPU profitieren. Einige Beispiele:

• Ragged Paged Attention v3 für MHA/GQA-Modelle
• Dense und sparse Multi-head Latent Attention (MLA) für DeepSeek, Kimi und GLM
• Gated DeltaNet und Kimi Delta Attention für hybride Linear-Attention-Modelle
• Fused Mixture-of-Experts mit Expertenparallelität
• FP8-, INT8-, MXFP4- und NVFP4-quantisierte Matmuls

Skalierung nach außen

Wir haben uns stark auf die Workloads konzentriert, die heute den größten Druck auf Serving-Systeme ausüben, insbesondere auf Multi-Turn-Agentic- und multimodalen Datenverkehr, bei dem Kontexte ständig wachsen und Tool-Ausgaben lang sein können.

Parallelität

Alle sechs Parallelitätsdimensionen von vLLM werden unterstützt, sowohl auf Single-Host- als auch auf Multi-Host-TPU-Slices, und sie können kombiniert werden:

Für große MoE-Modelle erzielen wir in der Regel den besten Durchsatz mit datenparalleler Attention und Expert Parallelism (zum Beispiel, DP=8, TP=1 mit EP auf einem einzelnen v7x-Host). Dadurch bleibt die Attention lokal auf jedem Gerät, während die Experten über den gesamten Slice verteilt werden. Bei langen Kontexten teilen PCP und DCP Prefill und den KV-Cache auf Geräte auf. Der Parallelism Guide enthält weitere Details und Beispielkonfigurationen.

KV-Cache-Offloading und entkoppeltes Serving

vllm-torchtpu verwendet TPU Sync, eine Open-Source-C++-Bibliothek, die TPU-Gerätepuffer direkt außerhalb der Python-Runtime lesen kann. Über vLLMs KV-Connector-Interface erhalten wir dadurch:

• Zero-Copy-KV-Cache-Offloading in den Host-DRAM, sodass der Präfix-Cache langer Multi-Turn-Sessions nicht durch den HBM begrenzt wird.
• Entkoppeltes Prefill-/Decode-Serving über TPU-Hosts hinweg mit geringem KV-Transfer-Overhead. Prefill und Decode können unterschiedliche Layouts verwenden, zum Beispiel Prefill auf PCP=4, TP=2 und Decode auf DP=8, TP=1, wobei der KV-Cache während der Übertragung resharded wird.

Siehe den Leitfaden zum entkoppelten Serving für die Einrichtung.

Speculative Decoding und Prefix Caching

Speculative Decoding unterstützt MTP, EAGLE-3, DFlash, DSpark, n-gram und Suffix Decoding (Leitfaden). Automatic Prefix Caching funktioniert auch für hybride Modelle wie Qwen3.5 und Kimi K3, bei denen der Linear-Attention-Zustand zusammen mit dem KV-Cache zwischengespeichert wird.

Modelle, Funktionen und der weitere Ausblick

vllm-torchtpu kann die meisten von vLLM unterstützten Modelle bereitstellen, mit TPU-spezifischen Kernel-, Quantisierungs- und Parallelitätsarbeiten für die Familien Qwen, DeepSeek, Kimi, GLM, Llama, GPT-OSS und Gemma. TPU v7x (Ironwood) ist die empfohlene Hardware und wird in der CI bei jedem Pull Request getestet; v6e, v5e und v5p sind experimentell.

Die Seite Supported Models & Features enthält die vollständige Liste der validierten Checkpoints und Funktionen sowie Anleitungen zur Aktivierung. Die nächtlichen Support-Matrizen verfolgen Korrektheits- und Benchmark-Ergebnisse der CI für TPU v7x.

Was kommt als Nächstes?

Wir wollen, dass neue offene Modelle aus den großen Familien (Qwen, DeepSeek, GLM, Kimi, Gemma, GPT-OSS) am Tag ihrer Veröffentlichung auf TPU laufen. Einiges dessen, was kommt, aus der öffentlichen Roadmap:

• Eine automatisierte CI-Pipeline, die neue Open-Weight-Modelle innerhalb von 24 Stunden nach Veröffentlichung onboardet, benchmarkt und Support-Matrizen veröffentlicht.
• Prefix-, Cache- und Session-aware Routing sowie llm-d TPU Memory-Tier-Routing.
• PyTorch-Schnittstellen für Tokamax -Kernels, die mit anderen Projekten im TorchTPU-Ökosystem geteilt werden.
• Serving von Diffusionsmodellen durch vllm-omni.

Beiträge sind willkommen!

TPU-Kapazität für Beitragende

Wir öffnen außerdem einen Community Contribution TPU Capacity Pool. Wenn Sie auf vllm-torchtpubauen, können Sie sich um kostenlose TPU-Zeit bewerben, und die Core-vLLM-Maintainer bei Inferact helfen bei Design und PR-Reviews.

Bewerben Sie sich hier für Contributor-TPU-Zugang.

Probieren Sie es aus

• Schnellstart
• Cloud TPU Einrichtung
• TPU v7x (Ironwood) Serving-Rezepte
• vLLM TorchTPU-Rezepte auf Ironwood
• Kubernetes-Bereitstellung
• Beitragshandbuch

Wenn Sie Fragen haben oder sich beteiligen möchten, finden Sie uns im #sig-tpu auf dem vLLM Developer Slack.