Dieses Dokument synthetisiert wegweisende Entwicklungen im Bereich der KI-Infrastruktur, hardwarenahen Compiler-Optimierung und systemübergreifenden Skalierung auf Google TPUs und High-End-GPUs. Im Fokus stehen architektonische Paradigmenwechsel: die physische Entkopplung (Disaggregation) von Rechen- und Speicherkomponenten wie Attention- und FFN-Knoten (FastFD), die Einführung mutabler Referenzen im funktionalen JAX-Framework zur Eliminierung von Speicher-Kopier-Overheads sowie die hardwarenahe Register- und Speichersteuerung mittels spezialisierter Compiler (Mosaic, Pallas pl.kernel, Flight DSL). Vergleichende Benchmarks demonstrieren gravierende P99-Latenzeinbrüche bei Multi-GPU-Clustern ohne dedizierte NVLink-Interconnects und unterstreichen den Wandel von theoretischen Peak-FLOPs hin zu praxistauglichem, Service-Level-Objective-konformem „Goodput“.
Wichtigste Punkte
System-Co-Design & TPU-Evolution: Googles Infrastruktur skaliert von historischen In-Memory-Webindizes auf 1.200 Rechnern (1999) über TPUv1-Inferenzbeschleuniger (2015) bis zu modernen Groß-Pods der Generationen v4 bis v7 mit integrierten optischen Schaltern (OCS) und Flüssigkühlung. Durch automatisierte Flottenüberwachung stieg der Goodput von 50–60 % auf 95–98 % bei Systemgrößen von knapp 10.000 Chips.
Latenzkollaps ohne NVLink: Ein realer Serving-Benchmark von Llama 3.1 (8B & 70B in BF16) zeigt, dass Multi-GPU-Workstations mit 8x RTX 6000 Ada über PCIe Gen 5 ohne NVLink beim 70B-Modell extreme Einbrüche in den P99-Latenzen (bis zu 462 ms) erleiden. TPUs (v5e/v6e via MaxText) bieten stabilen Durchsatz unter hoher Last, erfordern jedoch striktes Caching gegen 20–30-minütige OpenXLA-Kaltstarts.
Pallas pl.kernel & Mutable Refs: JAX bricht mit dem rein funktionalen Paradigma durch die Einführung veränderlicher Referenzen (jax.Ref). Die monolithische Schnittstelle pallas_call wird durch die modulare API pl.kernel ersetzt, um speichereffiziente In-Place-Mutationen zu ermöglichen. Das Helper-Tool emit_pipeline erlaubt manuelles Pipelining zwischen HBM, VM und SM sowie die Überlappung von Berechnungen mit asynchroner Kollektivkommunikation (z. B. All-Gather).
TPU Sparse Cores Programmierbarkeit: Zur Beschleunigung unregelmäßiger Vektoroperationen (z. B. Mixture-of-Experts-Routing, Top-K-Sampling) werden TPU Sparse Cores, bestehend aus 1 Skalar-Subcore und 16 Vektor-Subcores (SIMD), über Hardware-Semaphoren koordiniert. Hardwarebeschleunigte Gather/Scatter-Befehle und MPMD-Kernel optimieren hierbei den Datendurchsatz.
Layout-Kompilierung mit Mosaic: Als MLIR-basierte Zwischenschicht bildet Mosaic N-dimensionale Tensoren direkt auf zweidimensionale TPU-Vektorregister (VREGs: 8 Sublanes × 128 Lanes) ab. Der Compiler automatisiert Tiling, Padding und Relayouts und senkt gepackte Sub-32-Bit-Formate (FP8, BF16) hardwarenah ab.
Alternative DSLs & Portierungen: Die Python-basierte Flight DSL umgeht XLA-Einschränkungen auf AMD-Beschleunigern durch feingranulare Layout-Kontrolle auf MFMA- und Thread-Ebene (analog zu CuTe), was bei Mixtral 8x7B (FP8/BF16) Beschleunigungsfaktoren von 2,5x beim Prefill und 3,5x beim Decoding liefert. Das Intel oneAPI-Backend ermöglicht JAX-Workloads auf Arc Pro B70 (Xe2) via oneDNN und MLIR-Absenkung zu SPIR-V und Level-Zero.
TokuMX Kernel-Bibliothek: Die plattformübergreifende Open-Source-Bibliothek TokuMX stellt hochoptimierte Deep-Learning-Kernel (Compressed Attention, Mamba, MoE-Routing, FP8-Quantisierung) für NVIDIA GPUs (H100, B200) und Google TPUs (v5p, TPU 8/8i) bereit, inklusive automatisiertem Autotuning pro Tensorform und nativer PyTorch-Schnittstellen.
XPro Profiling-Metriken: Der OpenXLA Performance Profiler (XPro) erlaubt Zugriff auf bis zu 60.000 TPU-Hardware-Counter (Sampling-Rate bis 833 MHz) und einen 2-GB-Ringspeicher für transiente Fehler. Für NVIDIA Hopper- und Blackwell-GPUs bietet er über CUPTI PM-Sampling (10–100 kHz) und Multi-Pass Range Profiling mit über 6.000 Metriken.
Flax NNXv2 & Grain: Flax 0.13 führt standardmäßig den strikten Tree Mode (graph=False) ein, um dynamischen Laufzeit-Overhead zu eliminieren, unterstützt von Hilfsmitteln wie unpack und jit_partial. Das Python/C++-Framework Grain garantiert deterministisches, fehlertolerantes Laden von Datensätzen und verhindert Trainings-Inferenz-Abweichungen (Skew) durch C++-Backends.
Disaggregiertes Serving mit FastFD: Traditionelle Systeme leiden unter einer sinkenden Hardwareauslastung (MFU von 28 % auf 0,5 %) beim FFN, wenn lange Kontexte den KV-Cache blockieren. FastFD entkoppelt Attention- und FFN/MLP-Knoten physisch. Durch M:N-Dispatch-Muster, One-Step-Ahead-Scheduling und fusionierte DeepSeek-Kernel wird die Transferzeit maskiert, was 1,35x bis 1,45x mehr Durchsatz bei konstanter Latenz bringt.
RPA für optimierte TPU-Attention: Um die physische Ineffizienz von 256×256 Matrix-Units (MXU) bei Head-Dimensionen von $H=128$ zu beheben, fusioniert RPA die KV-Cache-Aktualisierung direkt in den Rechenkernel. Unter Nutzung asynchroner DMAs wird beim Decoding eine Memory Bandwidth Utilization (MBU) von bis zu 86 % in vLLM-TPU und SGLang erzielt.
Modellinnovationen & Validierung:
UniDiffuser: Erzeugt parallel 1024 Bild-Tokens in nur 20 Entrauschungsschritten (Flow-Matching) statt sequenzieller autoregressiver Schritte.
OOD-Erkennung: Untersuchungen zeigen, dass Skalierung komplexe Algorithmen schlägt. Die Berechnung eines OOD-Scores mittels Mahalanobis-Distanz erfolgt in JAX in nur 5 ms (vs. 100 ms Vorwärtslauf).
Klinische LLMs: Ein auf Notaufnahme-Gesundheitsakten (EHR) feinabgestimmtes Qwen 34B verbesserte den F1-Score bei ICD-10-Diagnosen um 16 %, zeigte jedoch partielles katastrophales Vergessen und den Verlust von Chain-of-Thought-Reasoning.
JAX GCM: Portiert Millionen Zeilen Fortran-Klimacode in ein differenzierbares JAX-Modell. Eine 15-tägige atmosphärische Simulation läuft in 26 Sekunden auf einem Standard-Laptop.
Glossar
BF16 (Bfloat16): Ein von Google entwickeltes 16-Bit-Gleitkommaformat, das den gleichen Wertebereich wie das herkömmliche 32-Bit-Format (FP32) besitzt, jedoch nur die Hälfte des Speichers beansprucht. Dies beschleunigt Deep-Learning-Berechnungen drastisch.
CoT (Chain of Thought): Eine Methode, bei der ein Sprachmodell komplexe Aufgaben löst, indem es explizite logische Zwischenschritte formuliert, bevor es die endgültige Antwort ausgibt.
CUPTI (CUDA Profiling Tools Interface): Eine hardwarenahe Programmierschnittstelle von NVIDIA, mit der Entwickler detaillierte Leistungsdaten und Zeitmessungen direkt von der GPU abfragen können.
DMA (Direct Memory Access): Ein Verfahren, bei dem Hardwarekomponenten (wie Grafik- oder TPU-Chips) direkt Daten im Arbeitsspeicher lesen oder schreiben können, ohne den Hauptprozessor (CPU) zu belasten.
DSL (Domain-Specific Language): Eine spezialisierte Programmiersprache, die für ein ganz bestimmtes Problemfeld entwickelt wurde (z. B. mathematische Kernel-Programmierung) und in diesem Bereich effizienter als Universalsprachen ist.
EHR (Electronic Health Record): Die elektronische Patientenakte, welche die Krankengeschichte, Diagnosen und Behandlungsverläufe von Patienten digital speichert.
Flow-Matching: Eine neuartige mathematische Methode zur Definition kontinuierlicher Pfade bei generativen Modellen (wie Bildgeneratoren), die oft schneller und stabiler als klassische Diffusionsprozesse konvergiert.
FP8: Ein extrem kompaktes 8-Bit-Gleitkommaformat für neuronale Netze, das im Vergleich zu 16-Bit-Formaten Speicherplatz spart und Rechenoperationen auf moderner Hardware beschleunigt.
GEMM (General Matrix Multiply): Die grundlegende mathematische Operation der Matrixmultiplikation, die das Kernfundament fast aller modernen Deep-Learning-Berechnungen bildet.
Goodput: Die Menge an nützlichen, fehlerfreien Daten oder Anfragen, die ein System innerhalb eines bestimmten Zeitraums verarbeitet. Im Gegensatz zum reinen Durchsatz (Throughput) fließen Fehler, Verzögerungen und System-Overheads hier negativ ein.
HBM (High Bandwidth Memory): Ein extrem schneller, dreidimensional gestapelter Arbeitsspeicher, der direkt neben oder auf dem Prozessorchip verbaut ist, um die enorme Datenbandbreite für KI-Berechnungen zu liefern.
HLO (High-Level Optimizer): Die High-Level-Zwischenrepräsentation innerhalb des XLA-Compilers, die den mathematischen Rechengraphen eines Modells optimiert, bevor dieser in hardwarenahen Maschinencode übersetzt wird.
ICD-10-Diagnosecodes: Das weltweit anerkannte Klassifikationssystem der Weltgesundheitsorganisation (WHO) zur Verschlüsselung von medizinischen Diagnosen.
JAX: Ein von Google entwickeltes, hochleistungsfähiges Python-Framework für numerische Berechnungen und maschinelles Lernen, das automatische Differenzierung und Just-In-Time-Kompilierung (JIT) via XLA unterstützt.
MBU (Memory Bandwidth Utilization): Ein Prozentsatz, der angibt, wie effizient die theoretisch maximale Speicherbandbreite eines Chips durch eine ausgeführte Rechenoperation tatsächlich ausgenutzt wird.
MFU (Model FLOPs Utilization): Ein Effizienzmaß für KI-Hardware, das angibt, wie viel Prozent der theoretisch maximalen Gleitkomma-Rechenleistung (FLOPs) während der echten Ausführung des Modells effektiv genutzt werden.
MLIR (Multi-Level Intermediate Representation): Eine moderne Compiler-Infrastruktur, die es ermöglicht, Code auf verschiedenen Abstraktionsebenen (von High-Level-Graphen bis zu maschinennahen Registern) einheitlich zu optimieren.
MoE (Mixture of Experts): Eine Netzwerkarchitektur, bei der nicht das gesamte Modell für jede Eingabe aktiv ist. Stattdessen routet ein "Gating-Netzwerk" Daten dynamisch an spezialisierte Teilnetzwerke ("Experten"), was Rechenleistung spart.
MPMD (Multiple Program, Multiple Data): Ein paralleles Ausführungsschema, bei dem verschiedene Prozessoren oder Kerne gleichzeitig unterschiedliche Programme auf unterschiedlichen Datensätzen ausführen.
Mutable Referenzen (jax.Ref): Zeiger auf veränderbare Speicherbereiche innerhalb des sonst strikt funktionalen (unveränderlichen) JAX-Frameworks. Sie verhindern unnötige Kopierprozesse von Daten im Speicher.
MXU (Matrix Multiply Unit): Spezialisierte Rechenwerke auf Google TPUs, die ausschließlich für die extrem schnelle Ausführung von Matrixmultiplikationen ausgelegt sind.
NVLink: Eine proprietäre, extrem schnelle Kommunikationsverbindung von NVIDIA zur Kopplung mehrerer GPUs, die im Vergleich zu Standard-PCIe-Verbindungen eine um Vielfaches höhere Bandbreite bietet.
OCS (Optical Circuit Switch): Ein optischer Schalter, der Datenverbindungen zwischen Servern oder Chips mithilfe von Lichtwellen und Spiegeln anstelle von elektrischen Kabeln herstellt. Dies reduziert Latenz und Energieverbrauch im Rechenzentrum dramatisch.
OOD (Out-of-Distribution): Bezeichnet Eingabedaten, die sich grundlegend von den Trainingsdaten eines Modells unterscheiden. Modelle reagieren auf OOD-Daten oft unvorhersehbar oder mit Halluzinationen.
OpenXLA: Ein Open-Source-Compiler-Konsortium für maschinelles Lernen, das die effiziente Ausführung von Modellen auf verschiedenen Hardwaretypen (TPUs, GPUs, CPUs) vereinheitlicht.
P99-Latenz: Das 99. Perzentil der Antwortzeit. Es gibt an, dass 99 % aller Anfragen schneller als dieser Wert verarbeitet wurden. Dies ist eine kritische Metrik, um die Auswirkungen von seltenen, aber extremen Verzögerungen zu messen.
SIMD (Single Instruction, Multiple Data): Ein Prozessorprinzip, bei dem ein einziger Befehl gleichzeitig auf mehrere Datenelemente angewendet wird (Vektorverarbeitung).
SLO (Service Level Objective): Ein definiertes Leistungsziel für einen Dienst, beispielsweise dass 99 % aller Benutzeranfragen in weniger als 200 Millisekunden beantwortet werden müssen.
SM (Scalar Memory): Ein kleiner, schneller On-Chip-Speicher auf TPUs, der für einfache, nicht-vektorisierte Hilfsdaten und Kontrollstrukturen genutzt wird.
SPIR-V / Level Zero: SPIR-V ist ein standardisiertes Zwischenformat für Computercode im Grafikbereich. Level Zero ist Intels hardwarenahe Schnittstelle zur direkten Steuerung und Programmierung seiner Grafikkarten.
TPU (Tensor Processing Unit): Ein von Google entwickelter anwendungsspezifischer Prozessor (ASIC), der speziell für die Beschleunigung von Berechnungen mit Tensoren (wie im Deep Learning) optimiert ist.
VM (Vector Memory / VMEM): Der schnelle, interne Speicher auf einer TPU, der speziell für die Bereitstellung von Vektordaten an die Vektor-Rechenwerke zuständig ist.
VREGs (Vector Registers): Interne Register im Prozessor, die große Datenvektoren für direkte mathematische Operationen extrem schnell bereithalten.