Dieser Bericht fasst die Kerninhalte einer Fachtagung zu moderner KI-Infrastruktur, Compiler-Technologien und LLM-Serving-Architekturen zusammen. Im Fokus stehen die Skalierung von Google TPU-Pods (v4 bis v7), hardwarenahe Optimierungen durch Low-Level-APIs wie Pallas und Mosaic, sowie systemarchitektonische Paradigmenwechsel wie die Entkopplung von Attention- und FFN-Modulen (FastFD) zur Maximierung des effektiven Goodputs.
Wichtigste Punkte
Google TPU Pods & Goodput-Optimierung: Skalierung von Groß-Pods auf knapp 10.000 Chips mit Flüssigkühlung und optischen Schaltern (OCS). Automatische Flottenüberwachung steigert den Goodput (Nutzleistung unter strikten Latenz-SLOs) von 50–60 % auf 95–98 %.
LLM-Serving-Vergleich (GPU vs. TPU): Benchmarks von Llama 3.1 zeigen, dass Multi-GPU-Setups ohne NVLink bei 70B-Modellen in P99-Latenzen (bis 462 ms) einbrechen. TPUs bewältigen hohe Lasten effizient, erfordern jedoch ein striktes Caching zur Vermeidung von OpenXLA-Kaltstarts (20–30 Minuten).
Pallas-API & Mutable Refs (jax.Ref): Abkehr vom rein funktionalen Paradigma hin zu mutablen Referenzen. Ermöglicht speichereffiziente In-Place-Mutationen und direkte manuelle Kontrolle über das Pipelining zwischen HBM, VM und SM (emit_pipeline) ohne Kopier-Overhead.
Mosaic Vektor-Layout-Compiler: MLIR-basierte Compiler-Schicht zwischen Pallas und TPU-Backends. Zuständig für die Abbildung N-dimensionaler Tensoren auf zweidimensionale TPU-Vektorregister (VREGs: 8 Sublanes × 128 Lanes) inklusive automatischem Tiling und FP8/BF16-Packaging.
FastFD (Attention-FFN-Disaggregation): Physische Entkopplung von zustandsbehafteter Attention und zustandslosem FFN/MLP in separate Knotenpools. Ein asymmetrisches M:N-Dispatch-Muster und Microbatch-Overlapping eliminieren den Einbruch der Hardwareauslastung bei langen Kontexten und steigern den Durchsatz um das 1,35- bis 1,45-fache.
RPA (Recurrent / Rapid-Paged Attention): Fusionierung der KV-Cache-Aktualisierung direkt in den Attention-Kernel unter Nutzung asynchroner DMAs. Erreicht eine Memory Bandwidth Utilization (MBU) von bis zu 86 % beim reinen Decoding auf TPU-MXUs.
Grain Deterministische Daten-Pipelines: C++/Python-Framework zur Behebung des „Last-Mile-Problems“ von Datenzuführungen. Garantiert fehlerfreies, deterministisches Checkpointing und eliminiert Trainings-Inferenz-Abweichungen (Skew) über Formate wie ArrayRecord und Zarr.
TokuMX & Kuratierte Hochleistungskernel: Plattformübergreifende Open-Source-Bibliothek für NVIDIA GPUs (H100, B200) und Google TPUs. Bietet optimierte Kernel für Compressed Attention, Mamba, MoE-Routing und fusionierte FP8-Quantisierung mit automatisiertem Autotuning.
Glossar
Goodput: Die reale, fehlerfreie und SLO-konforme Nutzleistung (z. B. erfolgreiche Anfragen pro Sekunde unter Einhaltung strenger Latenzgrenzen) im Gegensatz zum theoretischen Maximaldurchsatz (Peak-Throughput).
TPU (Tensor Processing Unit): Von Google entwickelte anwendungsspezifische integrierte Schaltungen (ASICs), die massiv parallelisiert für Deep-Learning-Workloads optimiert sind.
Pallas: Ein Low-Level-Framework in JAX, das Entwicklern direkte Kontrolle über On-Chip-Speicher und asynchrone Speicheroperationen auf TPU-Architekturen gewährt.
XLA (Accelerated Linear Algebra): Ein Domain-Specific Compiler von Google zur Optimierung linearer Algebra-Graphen in ML-Frameworks.
MLIR (Multi-Level Intermediate Representation): Ein modulares Compiler-Infrastruktur-Framework zur Erstellung optimierter Software-Repräsentationen auf verschiedenen Abstraktionsebenen.
FastFD: Systemarchitektur, die Attention-Mechanismen und Feed-Forward-Netzwerke (FFN) physisch trennt, um Hardware-Ressourcen bei langen Textkontexten optimal auszulasten.
KV-Cache (Key-Value Cache): Ein Zwischenspeicher für Schlüssel- und Wertvektoren in Attention-Schichten, der die Neuberechnung bei der sequenziellen Textgenerierung (Autoregressive Generation) vermeidet.
MBU (Memory Bandwidth Utilization): Die prozentuale Auslastung der maximal physikalisch verfügbaren Speicherbandbreite eines Prozessors.
MoE (Mixture of Experts): Eine Modellarchitektur, bei der Eingaben dynamisch an bestimmte Teilnetzwerke („Experten“) geroutet werden, um Parameterkapazitäten bei geringem Rechenaufwand zu skalieren.
OOD (Out-of-Distribution): Eingabedaten oder Anfragen, die außerhalb der Verteilung der ursprünglichen Trainingsdaten liegen und oft zu Fehlverhalten oder Halluzinationen führen.