Dieser Vortrag präsentiert technische Fortschritte bei der Skalierung großer Sprachmodelle (LLMs) und Mixture-of-Experts-Architekturen (MoE) auf Google Tensor Processing Units (TPUs) unter Verwendung der Open-Source-Inferenz-Engines vLLM und SGLang (JAX-Backend).
Im Mittelpunkt stehen detaillierte System- und Kernel-Optimierungen für XL-MoE-Modelle wie Qwen3-Coder-480B auf der Google Ironwood-Chiparchitektur. Durch horizontale Kommunikationsanpassungen (Ersetzung von all_to_all durch Replikation und all_reduce) sowie vertikale Kernel-Fusionen auf dem TPU-On-Chip-Speicher (VMEM) konnte die Ausführungseffizienz von 9,8 auf 286 Tokens pro Sekunde und Chip gesteigert werden, was 77 % der theoretischen Roofline-Leistung entspricht.
Darüber hinaus wird das „TPU-RL“-Ökosystem vorgestellt, das MaxText, Tunx, TPU-Inference und den Multi-Host-Orchestrierer Pathways kombiniert. Die Infrastruktur ermöglicht effizientes Reinforcement Learning durch entkoppelte oder gemeinsam genutzte Ausführungsmodi, dynamisches KV-Cache-Management für Synchros und Expert Replay. Leistungsdaten demonstrieren bei einem GRPO-Run mit Qwen2.5-35B auf 128 Ironwood-Chips einen Sampler-Durchsatz von 35.000 Tokens/s. Abschließend wird die Integration von TorchTPU angekündigt.
Wichtigste Punkte & Zeitstempel
0:00 Einführung & Hintergrund: Vorstellung der Neuerungen in den JAX-basierten Inferenz-Engines vLLM und SGLang für Google TPUs zur Erzielung wettbewerbsfähiger Kosten-Leistungs-Verhältnisse in Produktionsumgebungen.
0:58 Wertversprechen der TPU-Inferenz: Erweiterung der verfügbaren Rechenkapazitäten und Reduzierung von Entwicklungsaufwänden durch Beibehaltung bestehender Schnittstellen und Kompatibilität mit GKE und Vertex AI.
2:51 Fortschritte der vLLM-TPU-Architektur: Übergang zu einer Hardware-Plugin-Struktur, Erreichen einer 2- bis 5-fachen Leistungssteigerung gegenüber dem Prototyp und Fokus auf die Ironwood-Chipgeneration für XL-MoE- und Argumentationsmodelle.
5:14 Leistungsoptimierung für XL-MoE-Modelle: Optimierungsarbeiten am Modell Qwen3-Coder-480B führen zu einer Roofline-Ausschöpfung von 77 % sowie einer Durchsatzsteigerung von 9,8 auf 286 Tokens/s/Chip (bei Decode-Prozessen von 34 auf 574 Tokens/s/Chip).
10:00 Horizontale Kommunikationstaktiken: Umgehung ineffizienter all_to_all-Operationen auf der Ring-Mesh-Topologie der TPUs durch vollständige Eingabe-Replikation über alle Shards hinweg.
12:44 Lokale Reduktion vor Inter-Node-Kommunikation: Umstrukturierung der Berechnungsreihenfolge zur Durchführung lokaler Token-Reduktionen vor Datenübertragungen; der Wechsel von all_to_all zu all_reduce halbiert das Kommunikationsvolumen und verdoppelt die Leistung.
14:30 Vertikale Kernel-Fusion im On-Chip-Speicher: Ausnutzung des großen TPU-VMEMs zur aggressiven Verschmelzung von GEMM-Operationen, Aktivierungsfunktionen (SiLU/SwiGLU) sowie Quantisierungsstufen, um HBM-Zugriffe zu eliminieren.
16:11 Automatisierte Skalierung auf Qwen 3.5: Nutzung kompilierbasierter Kostenmodelle zur automatischen Kachelgrößen-Auswahl; Wiederverwendung optimierter Module verkürzt die Optimierungsdauer von Qwen 3.5 auf 2 Monate für 50 % Roofline-Effizienz.
19:45 Reinforcement-Learning-Architektur (TPU-RL): Präsentation der RL-Pipeline aus MaxText (Training), Tunx (Post-Training/Orchestrierung), TPU-Inference (vLLM) und Pathways (Multi-Host-Steuerung via Inter-Core Interconnect).
23:11 Ausführungsmodi für RL-Workloads: Gegenüberstellung von Colocated-Ausführung für ressourcenbeschränkte Hardware und Disaggregated-Ausführung zur Entkopplung von Training und Rollouts auf separaten TPU-Knoten.
25:02 vLLM-Funktionen für RL-Anforderungen: Implementierung von Datenparallelität für Rollouts, temporärem Löschen des KV-Caches vor Gewichts-Synchronisierungen zur HBM-Entlastung sowie Expert-Replay für stabiles Routing.
27:51 RL-Benchmark-Leistung auf Ironwood: Ein GRPO-Trainingstestlauf mit Qwen2.5-35B auf 128 Ironwood-Chips erzielt 35.000 Tokens/s Sampler-Durchsatz bei einer globalen Schrittzeit von 6 Minuten.
28:45 Ausblick und Ankündigung von TorchTPU: Ausrichtung auf Multimodalität, lange Kontexte und Rack-Scale-Systeme; Ankündigung des TorchTPU-Starts und neuer vLLM-/SGLang-Backends im Oktober 2026.
Glossar
TPU (Tensor Processing Unit): Von Google entwickelter, anwendungsspezifischer Integrierter Schaltkreis (ASIC) zur Beschleunigung von Arbeitslasten des maschinellen Lernens.
vLLM / SGLang: Hochoptimierte Open-Source-Software-Bibliotheken zur effizienten Bereitstellung und Inferenz großer Sprachmodelle.
MoE (Mixture of Experts): Eine Architektur neuronaler Netze, bei der Token dynamisch an spezialisierte Teilnetzwerke (Experts) geroutet werden, um die Parameterkapazität ohne proportionalen Anstieg der Rechenkosten zu erhöhen.
Ironwood (Trillium): Spezifische Google TPU-Hardware-Generationen, die für hohe Rechenleistung pro Flop sowie optimierte Speicherbandbreiten ausgelegt sind.
EP (Expert Parallelism): Eine Verteilungsstrategie für MoE-Modelle, bei der unterschiedliche Experten auf verschiedenen Beschleunigerchips platziert werden.
TP (Tensor Parallelism): Ein Aufteilungsverfahren, bei dem einzelne Gewichtsmatrizen innerhalb einer Modellschicht über mehrere Beschleuniger hinweg berechnet werden.
Roofline-Modell: Ein visuelles Leistungsmodell zur Bestimmung, ob eine Arbeitslast durch die maximale Rechenleistung (Flops) oder die Speicherbandbreite (Byte/s) der Hardware begrenzt ist.
GEMM (General Matrix Multiply): Standardisierte Operation der linearen Algebra für allgemeine Matrix-Matrix-Multiplikationen, die die Hauptrechenlast in tiefen Netzen bildet.
VMEM (Vector Memory) vs. HBM (High Bandwidth Memory): VMEM bezeichnet den extrem schnellen, lokalen On-Chip-Speicher der TPU-Kerne; HBM ist der große, externe Hauptspeicher mit hoher Bandbreite.
Kernel-Fusion: Ein Optimierungsverfahren von Compilern, bei dem mehrere aufeinanderfolgende mathematische Operationen in einem einzigen Speicher-Ablauf zusammengefasst werden, um Lese- und Schreibzugriffe auf den Hauptspeicher zu reduzieren.
ICI (Inter-Core Interconnect): Googles proprietäre, latenzarme Interconnect-Technologie zur direkten Kopplung mehrerer TPU-Chips innerhalb eines Systems.
Pathways: Googles verteilte Laufzeitumgebung zur transparenten Orchestrierung paralleler Rechenoperationen über Tausende von Beschleunigerchips.
GRPO (Group Relative Policy Optimization): Ein effizienter Algorithmus des Reinforcement Learnings zur Ausrichtung (Alignment) von Sprachmodellen.
TorchTPU: Eine native Integrationsschicht, die es ermöglicht, PyTorch-basierte Modelle direkt auf Google TPUs auszuführen.