Dieses Transkript dokumentiert eine Präsentation des Google-Teams (Britney Rockwell, Kuan Kim und Winshin Dong) auf einem DevLab-Event, in der die Skalierung von Large Language Models (LLMs) und Mixture-of-Experts-Modellen (MoE) für die Inferenz auf Google-TPUs (Tensor Processing Units) behandelt wird. Im Fokus stehen die Integration von TPUs in Open-Source-Inferenz-Engines wie vLLM und SGLang (mittels JAX), horizontale und vertikale Optimierungstechniken (wie Instruction Reordering und Kernel-Fusion), die Unterstützung von Reinforcement Learning (RL) über Pathways sowie die kommende Einführung des PyTorch-TPU-Backends im Oktober.
Wichtigste Punkte & Zeitstempel
0:30 vLLM- und SGLang-JAX-Backends: Integration von TPUs direkt in Open-Source-Inferenz-Engines zur Steigerung des Preis-Leistungs-Verhältnisses (Perf-per-Dollar), Erhöhung der Hardware-Verfügbarkeit und Minimierung von Entwickler-Reibungsverlusten.
3:50 Fokus auf Ironwood und Trillium: Anpassung der Engines an den neuen Ironwood-Chip (XLA-Optimierungen), um speicher- und rechenintensive Reasoning-Modelle im Vergleich zum kleineren Trillium-Chip effizienter auszuführen.
5:14 Mixture-of-Experts (MoE) Herausforderungen: Bewältigung von Token-Ungleichgewichten, Expert Parallelism (EP) und Routing-Strategien bei extrem skalierten MoE-Modellen (XL-MoE).
7:38 Leistungssteigerung bei Qwen-Modellen: Steigerung der theoretischen Roofline-Effizienz von Qwen-Modellen von 2 % auf bis zu 77 % (ca. 95 % bei realistischen Annahmen) durch neue GMM-Kernel (General Matrix Multiply), Ragged Gather/Scale und Batched-RPA-Kernel.
11:02 Horizontale Optimierungen (Token Routing): Ersetzen des ineffizienten All-to-All-Netzwerk-Overheads durch vollständige Token-Replikation und Vorziehen von Reduktionsoperationen, was zu einer Verdopplung der Performance führt.
14:31 Vertikale Optimierungen (Kernel-Fusion): Nutzung des großen, exklusiven Vector Memory (VM) pro TPU-Kern für aggressive Fusionsoperationen (Kombination von GMM, Aktivierungsfunktionen wie SiLU sowie In- und Dequantisierung).
16:58 Modellübergreifende Code-Wiederverwendung: Einsatz von Compiler-basierten Kostenmodellen und automatisierter Kachelgrößen-Optimierung (Tile Sizing), wodurch die Portierungszeit für Qwen 3.5 auf zwei Monate verkürzt wurde.
19:45 Reinforcement Learning (RL) auf TPUs: Winshin Dong erläutert die Orchestrierung von MaxText, Tunx, TPU-Inference (vLLM) und Pathways für skalierbare RL-Infrastrukturen in kolokierten oder disaggregierten Modi.
25:12 Native vLLM-Erweiterungen für RL: Implementierung von Data Parallelism, effizientem Gewichts-Sync (Weight Sync) über Cache-Neuzuweisung, Pathways-Integration und Expert Replay zur Beschleunigung der Konvergenz.
27:52 GRPO-Skalierungsergebnisse: Erzielung einer Sampler-Throughput-Leistung von 35k Token/Sekunde und einer globalen Schrittzeit von 6 Minuten bei einem GRPO-Lauf (Group Relative Policy Optimization) mit Qwen3-35B auf 128 Ironwood-Chips.
28:50 Zukünftige Roadmap & PyTorch-TPU-Launch: Die kommenden sechs Monate konzentrieren sich auf Multimodalität, langen Kontext (Agentic Workloads), Rack-Scale-Multi-Host-Support und den Start des PyTorch-TPU-Backends im Oktober für vLLM und SGLang.
Glossar
Large Language Model (LLM): Ein fortgeschrittenes KI-Modell, das auf massiven Textdaten trainiert wurde, um menschliche Sprache zu verstehen und zu generieren.
Tensor Processing Unit (TPU): Ein von Google entwickelter, spezialisierter Mikrochip, der gezielt für die Beschleunigung von rechenintensiven Workloads im Bereich des maschinellen Lernens optimiert ist.
Inferenz (Inference): Die Betriebsphase eines KI-Modells, in der dieses auf Basis bereits gelerntem Wissens Vorhersagen trifft oder Antworten auf neue Anfragen generiert.
Mixture-of-Experts (MoE): Eine neuronale Netzwerkarchitektur, bei der ein Modell in mehrere spezialisierte Subnetze ("Experten") unterteilt wird, von denen bei jeder Anfrage nur ein Bruchteil aktiv geschaltet wird, um Rechenleistung zu sparen.
vLLM & SGLang: Weit verbreitete Open-Source-Inferenz-Engines, die für den Hochdurchsatz und die speichereffiziente Ausführung von Sprachmodellen optimiert sind.
JAX: Ein von Google entwickeltes Programmierframework für Hochleistungs-Rechenoperationen und automatische Differenzierung, das häufig im modernen KI-Training und -Inferenz-Stack eingesetzt wird.
XLA (Accelerated Linear Algebra): Ein domainspezifischer Compiler von Google, der Berechnungen linearer Algebren optimiert, um die Ausführungsgeschwindigkeit und Speichereffizienz auf Hardware wie TPUs zu maximieren.
GMM (General Matrix Multiply): Eine grundlegende mathematische Rechenoperation (Matrixmultiplikation), die den überwiegenden Teil der Rechenlast in neuronalen Netzen darstellt.
Vector Memory (VM): Ein extrem schneller, lokaler Hochgeschwindigkeitsspeicher, der direkt in einen TPU-Kern integriert ist.
Pathways: Ein von Google entwickeltes, verteiltes Multi-Host-Orchestrierungssystem, das die Ausführung komplexer, paralleler KI-Workloads über Tausende von TPU-Chips hinweg koordiniert.
Reinforcement Learning (RL): Ein maschinelles Lernverfahren, bei dem ein Modell durch Versuch und Irrtum (Belohnungen und Bestrafungen) lernt, optimale Handlungsstrategien zu entwickeln.
GRPO (Group Relative Policy Optimization): Ein effizientes Reinforcement-Learning-Verfahren zur Optimierung von Sprachmodellen, das im Gegensatz zu älteren Methoden ohne ein separates Bewertungsnetzwerk (Critic) auskommt.