Google bringt PyTorch-native Inferenz-Engines für vLLM und SGLang auf die TPU

Warum Google PyTorch-Backends gebaut hat

vLLM und SGLang laufen auf der TPU bereits über ein JAX-basiertes Backend. Google zufolge fragten Entwickler jedoch immer wieder nach einer Möglichkeit, ihren Serving-Stack in PyTorch anzupassen und zu debuggen. Das Unternehmen reagierte mit TorchTPU, das PyTorch nativ auf der TPU ausführt, und baute darauf neue Engines für vLLM und SGLang auf.

Nach Angaben von Google soll die PyTorch-Community damit Modelle auf die TPU bringen und mit Werkzeugen beschleunigen können, die Entwickler bereits kennen. Die Engines kommen mit produktionsreifen Serving-Funktionen und eigenen Pallas-Kernels.

TorchTPU macht die TPU zu einem nativen PyTorch-Gerät. Tensoren sind gewöhnliche PyTorch-Tensoren auf einem tpu-Gerät, ganz ohne Wrapper-Typen. Der Code läuft im Eager-Modus oder wird über torch.compile mit XLA kompiliert, kollektive Operationen laufen über torch.distributed. Beide Engines behalten das Worker-Modell, den Scheduler, den KV-Cache-Manager und den API-Server des jeweiligen Frameworks bei und ergänzen darunter TPU-Implementierungen. Die meisten Modelle nutzen die standardmäßigen PyTorch-Modelldefinitionen aus dem Upstream, TPU-spezifische Arbeit beschränkt sich auf Kernels und Layer-Implementierungen.

Unterstützte Funktionen und KV-Cache-Transfer

Google nennt Unterstützung für mehrere Arten von Workloads:

• Agentische und Multi-Turn-Szenarien: automatisches Prefix-Caching, KV-Cache-Offloading in den Host-Speicher und auf Managed Lustre, Tool Calling und strukturierte Ausgaben.
• Niedrige Latenz: spekulatives Decoding mit MTP, EAGLE-3, DFlash, DSpark und N-Gram-Drafters.
• Langer Kontext: Context Parallelism für Prefill und Decode, FP8-KV-Cache, KV-Offloading in den Host-Speicher.
• Disaggregiertes Serving: Prefill und Decode auf getrennten Hosts.
• Multimodal: Vision-Language-Modelle mit Bild- und Videoeingabe sowie Embedding-Modelle.
• Sehr große Modelle: mehrere Parallelisierungsmodi über Hosts hinweg sowie Checkpoints in FP8, FP4 (MXFP4, NVFP4), INT8 und INT4.

Agenten senden wachsende Konversationen immer wieder neu, deshalb ist es wichtig, ein zwischengespeichertes Präfix zu finden, statt es neu zu berechnen. Beide Engines halten wiederverwendbare Präfixe im HBM und lagern ältere Präfixe bei aktiviertem Host-Offload als zweite Ebene in den Host-Speicher aus.

Disaggregiertes Serving und gestaffeltes Offloading verlagern den KV-Cache jeweils weg von dem Gerät, das ihn berechnet hat. Beide Engines nutzen dafür TPU Sync, eine quelloffene C++-Bibliothek. Laut Google liest sie die Geräte-Puffer hinter den PyTorch-Tensoren direkt aus, sodass auf dem Host keine zusätzliche Zwischenkopie nötig ist, und führt Transfers in eigenen Threads außerhalb von Python aus, während Prefill und Decode weiterlaufen.

Benchmarks und Produktiveinsatz

Google veröffentlicht Benchmarks öffentlich über das Projekt SemiAnalysis InferenceX. Aufbauend auf einer früheren Einreichung von Qwen 3.5 auf der Ironwood TPU (TPU 7x) habe weitere Optimierung insgesamt die niedrigsten Kosten pro Million Token erreicht, nämlich 2,01 Cent in Konfigurationen mit hoher Parallelität, so das Unternehmen. Zudem veröffentlichte es neue Ergebnisse bei AgentX, der Kategorie des Projekts für agentische Workloads mit langem Kontext und mehreren Gesprächsrunden, und erklärt, die Ironwood TPU liege über den größten Teil der Pareto-Front vorn.

Gemeinsam mit Inferact integrierte Google einen Pallas-Megakernel für das Modell Kimi K3 mit 2,8 Billionen Parametern, der den gesamten Decode-Schritt als einen einzigen Kernel ausführt. Google zufolge liefert er auf der Ironwood TPU mehr als 700 Token pro Sekunde für einen einzelnen Nutzer. Inferact führte außerdem eine erste vorläufige AgentX-Einreichung für Kimi K3 auf vLLM-TorchTPU an.

SGLang-TorchTPU, gemeinsam mit RadixArk entwickelt, läuft bei Meta im Produktivbetrieb. Zongwei Zhou, PyTorch-TPU-Leiter bei Meta, sagte, es biete Meta „einen bemerkenswert ausgereiften, PyTorch-nativen Serving-Stack, der sich mühelos in unsere bestehende KI-Infrastruktur einfügt“. Simon Mo, CEO und Mitgründer von Inferact, sagte, Entwickler könnten mit vLLM „sofort Produktions-Workloads auf TPUs ausführen“.

Was das für Entwickler bedeutet

Für Teams, die mit PyTorch arbeiten, eröffnet der Start eine Möglichkeit, offene Modelle auf der TPU bereitzustellen, ohne auf JAX umzusteigen. Google zufolge werden die meisten beliebten offenen Modellfamilien unterstützt, darunter Qwen, DeepSeek, GLM, Kimi, GPT-OSS und Gemma4. Man arbeite an Day-0-Unterstützung für neue offene Modelle.

Außerdem startet vLLM einen TPU-Kapazitätspool für Community-Beiträge, der Entwicklern, die aktiv zu beiden Engines beitragen, kostenlose TPU-Rechenleistung bereitstellt. Google ergänzt, weitere Optimierungen seien in Arbeit, und lädt Entwickler ein, ein GitHub-Issue zu eröffnen, um eine Zusammenarbeit anzufragen.