Browse Summaries

Zusammenfassung

Der Vortrag stellt „FastFD“ (Fast Feed-Forward-Attention Disaggregation) vor, ein Open-Source-System zur Entkopplung von Attention- und Feed-Forward-Netzwerken (FFN/MLP) zur Optimierung großer Large Language Model (LLM)-Serving-Infrastrukturen. Vorgestellt von Yujia Zhang (UC San Diego und Google-Intern). Das Kernproblem in traditionellen, ko-lozierten Systemen (wie vLLM) besteht darin, dass bei langen Dekodierungskontexten die Attention-Batch-Größe schrumpft, wodurch die Rechenauslastung (MFU) der FFN-Knoten drastisch von 28 % auf 0,5 % einbricht. FastFD entkoppelt Attention-Server (mit KV-Cache-Verwaltung) und MLP-Server (für Expertenberechnungen) räumlich und nutzt modifizierte M2N-Dispatch- sowie N2M-Kombinationsmuster. Um die durch die Entkopplung drohende Latenzverdopplung zu verhindern, setzt das System auf Microbatch-Overlapping über Schichten hinweg, One-Step-Ahead-Scheduling auf der CPU sowie angepasste, fusionierte M2N/M2M-Mega-Kerne. In Benchmarks mit Qwen (35B) und MiniMax (229B) in FP8 auf H100/H200-Clustern erzielt das System einen Durchsatz-Speedup von 1,35x bis 1,45x gegenüber ko-lozierten Systemen bei konstanter Latenz.

Wichtigste Punkte & Zeitstempel

• 0:00 Architektur von Transformern: Transformer-Schichten bestehen aus stateful Attention (mit einem mit der Kontextlänge wachsenden KV-Cache) und stateless Feed-Forward-Netzwerken (FFN/MLP) (abhängig von Batch-Größe und Hidden-Dimension).
• 0:37 Underutilization bei langen Kontexten: In ko-lozierten Systemen führt eine lange Dekodierung zu schrumpfenden Attention-Batches, wodurch die Modell-Rechenleistungs-Auslastung (MFU) der FFNs von 28 % auf 0,5 % abfällt.
• 0:52 Konzept der Disaggregation: Räumliche Trennung in Attention-Server (für KV-Cache und Routing) und MLP-Server (für Expertenberechnungen) unter Verwendung angepasster M2N-Dispatch- und N2M-Kombinationsmuster zur Aggregation größerer Batches.
• 4:17 Herausforderung des Space-Sharing: Die reine Entkopplung zwingt zu sequenzieller Ausführung und verdoppelt im Worst-Case die Gesamtlatenz im Vergleich zur herkömmlichen Time-Sharing-Ko-lokation.
• 5:32 Microbatch-Overlapping: Aufteilung des Dekodierungsschritts in Microbatches, sodass Attention- und FFN-Server auf unterschiedlichen Microbatches gleichzeitig arbeiten und Rechenzeiten überlappen.
• 6:35 Erweiterte Scheduling- und Kernel-Optimierungen: Einsatz von One-Step-Ahead-Scheduling durch die CPU während der GPU-Ausführung sowie angepasster M2N/M2M-Mega-Kernels (abgeleitet von DeepSeek-Architekturen) zur Beschleunigung geräteübergreifender Kommunikation.
• 8:09 Leistungssteigerung und Benchmarks: Tests mit dem Qwen-35B- und MiniMax-229B-Modell (FP8) demonstrieren einen Durchsatz-Speedup von 1,35x bis 1,45x gegenüber Standard-vLLM-Konfigurationen.
• 9:55 Zukunftsaussichten für Hardware: Zukünftige Hardware-Trends mit verdreifachtem SRAM in modernen Beschleunigern versprechen eine noch stärkere Beschleunigung von FFN-Berechnungen und grössere Attention-zu-FFN-Ausnutzungsverhältnisse.

Glossar

• Attention (Aufmerksamkeitsmechanismus): Ein rechenintensiver Kernbestandteil von Transformern, der berechnet, welche Teile eines Textes beim Generieren des nächsten Wortes relevant sind. Speichert den sogenannten KV-Cache.
• KV-Cache (Key-Value Cache): Ein Speicherbereich in LLM-Infrastrukturen, der die berechneten Schlüssel- und Wert-Vektoren vergangener Token zwischenspeichert, um die Generierung neuer Token zu beschleunigen.
• Feed-Forward Network (FFN) / MLP (Multi-Layer Perceptron): Ein standardmäßiger neuronaler Netzwerkblock innerhalb von Transformer-Schichten, der die Daten nach dem Attention-Schritt verarbeitet; arbeitet zustandslos (stateless).
• Disaggregation (Entkopplung): Die räumliche Trennung verschiedener Rechenaufgaben (hier: Attention und FFN) auf unterschiedliche physische Server oder Hardware-Knoten, anstatt sie auf denselben GPUs auszuführen.
• MFU (Model FLOPs Utilization / Modell-Rechenleistungs-Auslastung): Eine Metrik, die misst, wie nah eine KI-Hardware (wie eine GPU) an ihrer maximal theoretischen Rechenleistung (FLOPs) betrieben wird.
• Microbatch (Mikro-Charge): Die Aufteilung einer großen Datenmenge (Batch) in kleinere Teilmengen, um Berechnungen parallel oder in Pipelines verschachtelt auszuführen.
• FP8 (8-Bit Floating Point): Ein datenarmes, komprimiertes Gleitkommaformat mit 8 Bit Präzision, das in modernen KI-Prozessoren eingesetzt wird, um den Speicherbedarf und die Übertragungszeiten zu halbieren.
• Mega-Kernel (Fusionierter Kernel): Ein speziell optimierter Software-Baustein für GPUs, der mehrere aufeinanderfolgende Rechen- und Kommunikationsschritte zu einer einzigen Operation zusammenfasst, um Overhead zu reduzieren.
• NVLink: Eine proprietäre Hochgeschwindigkeits-Verbindungstechnologie von NVIDIA, um mehrere GPUs innerhalb eines Knotens direkt und mit maximaler Bandbreite zu vernetzen.

Zusammenfassung

Dieses Dokument fasst einen Google-Technologievortrag über die Optimierung von Large Language Models (LLMs) auf Tensor Processing Units (TPUs) zusammen, insbesondere im Hinblick auf den im April auf arXiv veröffentlichten „RPA“-Paper. Es beleuchtet TPU-Interna wie die Speicherhierarchie, das Tiling und die Matrix Multiply Units (MXU, 256x256 ab V5). Zu den Kernherausforderungen zählen unregelmäßige ("ragged") Eingaben aus Mixed-Batching und ineffiziente Scatter/Gather-Operationen bei PagedAttention. Die vorgestellte RPA-Architektur (Recurrent/Rapid-Paged Attention) löst dies durch die Fusion von KV-Cache-Aktualisierungen und asynchronen DMA-Pipelines, wodurch Bandbreitenauslastungen (MBU) von bis zu 86 % bei Decoder-Workloads und 73 % bei Prefill-Workloads erreicht werden. Abschließend werden zukünftige Optimierungen wie Sparse-Core-Auslagerung und gleichzeitige Mehrfachabfragen diskutiert.

Wichtigste Punkte & Zeitstempel

• 0:00 RPA-Paper auf TPUs: Vorstellung der ersten systematischen Analyse von TPU-Workloads für LLMs, publiziert im April auf arXiv durch Google-Softwareingenieur Jefen.
• 1:30 Herausforderungen bei PageAttention & Mixed Batching: PagedAttention erfordert dynamische Scatter/Gather-Operationen, die von TPU-Tensor-Cores unzureichend unterstützt werden; Mixed Batching erzeugt unregelmäßige ("ragged") Eingaben, was die statische Optimierung des XLA-Compilers erschwert.
• 3:08 TPU-Speicherhierarchie & Layout: Ein logischer Datenchunk entspricht 128 Werten zu 32 Bit (512 Bytes, die minimale DMA-Granularität eines Vektorworts). Physische Speicherlayouts müssen das hardwarenahe Tiling berücksichtigen.
• 4:42 MXU-Auslastungseinschränkungen: TPUs ab Version 5 verwenden eine 256x256-MXU (Matrix Multiply Unit). Bei einer Head-Dimension von $H_{dim} = 128$ wird die MXU-Auslastung physikalisch auf maximal 50 % limitiert.
• 5:56 Single-Threaded-Modell & Asynchrones DMA: TPU-Tensor-Cores arbeiten single-threaded und erzielen Parallelisierung durch asynchrone DMA-Befehle zur Überlappung von Prefetching und Berechnungen.
• 6:36 Fusionsdesign der RPA-Architektur: Integration der KV-Cache-Aktualisierung direkt in den Attention-Kernel, um die Latenz von Scatter/Gather-Operationen durch parallele DMA-Transfers vollständig zu kaschieren.
• 8:30 Decoder-Performance-Analyse: Decoder-Workloads erreichen bis zu 86 % Memory Bandwidth Utilization (MBU), erfordern jedoch Kontextlängen von 80K ($H_{dim}=128$) bzw. 4K ($H_{dim}=256$), um die Speicherbandbreite vollständig zu saturieren.
• 11:09 Prefill-Performance-Analyse: Prefill-Workloads erzielen 73 % MBU ohne Kausalmaske und 63 % mit Kausalmaske, wobei Vorverarbeitungs-Reshapes einen unvermeidbaren Overhead von 2 % bis 8 % verursachen.
• 16:06 Instruktionsbündelung & IDLE-Zustände: Die MXU ist durchgehend ausgelastet, während die Skalar-ALU (SAOU) bei Metadatenberechnungen untätig bleibt, was auf Einsparpotenziale durch Sparse-Core-Auslagerung hinweist.
• 17:34 Framework-Adoption: RPA dient als offizielles TPU-Backend für vLLM-TPU (seit Februar des Vorjahres), SGGLM (seit Oktober) und DeepMind-RL-Forschungsrepos.
• 18:19 Zukünftige Optimierungsansätze: Geplante Weiterentwicklungen umfassen die XLA-Unterstützung für unregelmäßige Formen, parallele Verarbeitung mehrerer Anfragen zur Bandbreitenauslastung bei kleineren Kontextlängen sowie Sparse-Core-Offloading.

Glossar

• TPU (Tensor Processing Unit): Von Google entwickelte, anwendungsspezifische integrierte Schaltung (ASIC), die speziell für die Beschleunigung von Deep-Learning- und Machine-Learning-Workloads optimiert ist.
• LLM (Large Language Model): Ein auf massiven Textdaten trainiertes Sprachmodell auf Basis von Transformern zur automatisierten Textgenerierung und -verarbeitung.
• KV-Cache (Key-Value Cache): Ein Zwischenspeicher in Transformermodellen, der berechnete Schlüssel- (Key) und Wertvektoren (Value) früherer Tokens sichert, um redundante Berechnungen bei der Inferenz zu verhindern.
• PagedAttention: Ein Speicherverwaltungssystem für den KV-Cache, das den Speicher in diskrete Blöcke (Pages) unterteilt, um Fragmentierung zu minimieren und dynamisches Caching zu ermöglichen.
• XLA (Accelerated Linear Algebra): Ein domänenspezifischer Compiler von Google, der Rechengraphen linearer Algebra optimiert, um die Ausführungsgeschwindigkeit auf CPUs, GPUs und TPUs zu maximieren.
• DMA (Direct Memory Access): Ein Hardware-Verfahren, bei dem Peripheriekomponenten oder Speicherbereiche Daten direkt unter Umgehung der Haupt-CPU austauschen, um den Datentransfer zu beschleunigen.
• MXU (Matrix Multiply Unit): Die zentrale Recheneinheit eines TPU-Tensor-Cores, die für massiv parallele Matrix-Multiplikationen optimiert ist.
• MBU (Memory Bandwidth Utilization): Eine prozentuale Kennzahl, die angibt, wie viel der maximal physikalisch verfügbaren Speicherbandbreite eines Prozessors während der Programmausführung genutzt wird.
• HFU (Hardware Flops Utilization): Das Verhältnis der tatsächlich erreichten Gleitkommaoperationen pro Sekunde (FLOPS) zur theoretischen Höchstleistung der Hardware.
• Scatter/Gather: Speicheroperationen, bei denen nicht zusammenhängende Daten aus verschiedenen Speicherbereichen gesammelt (Gather) oder verteilt (Scatter) werden, was auf bestimmten Hardware-Architekturen rechenintensiv ist.
• Ragged Inputs: Unregelmäßige Datenstrukturen in Batches, bei denen Sequenzen unterschiedliche Längen aufweisen, was die hardwarenahe Vektorisierung und Kompilierung erschwert.
• Causal Mask (Kausalmaske): Eine binäre Maske in Attention-Mechanismen, die verhindert, dass das Modell während der Inferenz oder des Trainings auf zukünftige Tokens zugreift.
• Sparse Core: Ein spezialisierter Rechenkern auf TPU-Architekturen, der für spärliche (sparse) Berechnungen, Indizierungen und unregelmäßige Speicherzugriffe (wie Scatter/Gather) optimiert ist.
• SAOU (Scalar ALU): Skalare Arithmetisch-Logische Einheit eines Prozessors, die im Gegensatz zu Vektoreinheiten einzelne, nicht-vektorisierte Werte verarbeitet.

Zusammenfassung

Diese Präsentation stellt NNXv2 vor, eine grundlegende Überarbeitung des NNX-Frameworks für JAX, die in Kürze als Teil von Flax 0.13 veröffentlicht wird. Der Schwerpunkt liegt auf der Vereinfachung von APIs, der Durchsetzung reiner Baumstrukturen (PyTrees) anstelle von Graphen, der Entfernung von Graph-Updates als Standardverhalten und der engeren Angleichung an native JAX-Transformationen. Neue APIs wie unpack, jit_partial sowie nx.capture und so reduzieren den Boilerplate-Code und minimieren den Python-Overhead bei komplexen neuronalen Netzen.

Wichtigste Punkte & Zeitstempel

• 0:00 NNXv2 Ankündigung: Vorstellung von NNXv2 als neuer Version zur Vereinfachung von APIs und zur Umsetzung von Lektionen aus den Vorjahren.
• 1:15 Funktionsweise von NNX: NNX nutzt Python-Code zur Erstellung von PyTrees und bildet veränderliche Zustände (Mutabilität) wie Modellparameter oder Zähler ab.
• 4:18 Problem von Graph-Updates: Historisch unterstützte Graph-Updates erlaubten das dynamische Hinzufügen von Attributen im laufenden Betrieb, verursachten jedoch erhebliche interne Komplexität.
• 5:27 Übergang zum Tree-Modus: NNXv2 setzt standardmäßig strikte Baumstrukturen voraus (graph=false), wodurch automatische Transformationen drastisch vereinfacht und beschleunigt werden.
• 8:18 Variablen-Schließungen und Erfassung: Einführung nativer Unterstützung für Variable-Closures (Variablen-Schließungen) und neue Mechanismen zur Erfassung von Zwischenwerten.
• 9:33 Neue APIs unpack und jit_partial: Nutzung von unpack zur manuellen Trennung trainierbarer Parameter von Nicht-Parametern, sowie jit_partial zur Vorab-Flattung von verschachtelten PyTrees zur Minimierung des Python-Overheads.
• 11:55 Metrik-Erfassung mit so und nx.capture: Ersatz für Graph-Updates durch die so- und nx.capture-APIs zur transparenten und kontrollierten Weiterleitung von Zwischenwerten und Metriken aus verschachtelten Schichten.
• 13:40 Zukunftsausblick und Release: NNXv2 ist stark auf das zukünftige hijacks-Konzept ausgerichtet und erscheint als Breaking Change in Flax 0.13.

Glossar

• NNX: Eine Erweiterungsbibliothek für das JAX-Ökosystem, die objektorientierte Programmierkonzepte für den Bau neuronaler Netze bereitstellt.
• JAX: Ein von Google entwickeltes Python-Framework für hochleistungsfähiges numerisches Rechnen und automatisches Differenzieren, optimiert für GPUs und TPUs.
• PyTree (Py-Baum): Eine baumartige Datenstruktur aus Python-Standardtypen (wie Listen, Tupeln und Dictionaries), die von JAX nativ verarbeitet und transformiert werden kann.
• Mutabilität (Veränderbarkeit): Die Eigenschaft von Datenstrukturen oder Variablen, ihren Zustand nach der Erstellung während der Programmausführung zu verändern (z. B. aktualisierte Modellgewichte).
• JIT (Just-In-Time-Kompilierung): Ein Verfahren, bei dem Programmcode zur Laufzeit in maschinennahen Code kompiliert wird, um die Ausführungsgeschwindigkeit drastisch zu steigern.
• API (Programmierschnittstelle): Eine definierte Schnittstelle, die es verschiedenen Softwarekomponenten ermöglicht, miteinander zu kommunizieren.
• Closure (Funktionsschließung): Ein Programmierkonzept, bei dem eine Funktion auf Variablen außerhalb ihres eigenen lokalen Gültigkeitsbereichs zugreift.
• Flattung / Unflattung: Das Umwandeln verschachtelter Datenstrukturen (wie PyTrees) in eine lineare, flache Liste (und umgekehrt), was für bestimmte Optimierungsschritte in JAX erforderlich ist.
• Flax: Ein populäres Open-Source-Framework für neuronale Netze, das auf JAX aufbaut.
• Loss (Verlustfunktion): Eine mathematische Funktion, die misst, wie fehlerhaft die Vorhersagen eines Modells im Vergleich zu den tatsächlichen Zielwerten sind.

Zusammenfassung

Diese Präsentation von Arty (AMD Jax Team) beschreibt die Integration von Flight-DSL-Kerneln in JAX für AMD-GPUs, mit einem Schwerpunkt auf leistungskritischen Workloads wie Mixture-of-Experts (MoE) GEMMs. Da Standard-XLA-Kompilierungen bei bestimmten Workloads wie Fused Attention oder MoE an Grenzen stoßen, bietet Flight DSL eine Python-basierte Layout-Kontrolle bis auf Thread- und MFMA-Ebene. Durch eine neu entwickelte JAX-Bridge, die Tensoren über DLPack in von Flight DSL verarbeitbare Memrefs konvertiert, lassen sich Group-GEMMs und Kernel-Fusions realisieren. Evaluierungen am Mixtral-8x7B-Modell (unter Verwendung von FP8 und BF16) zeigen massive Leistungssteigerungen (bis zu 3,5-facher Speedup beim Decoding) im Vergleich zu traditionellen Sparse-Matmul-Pfaden, wobei zukünftige Entwicklungen auf Trainingsunterstützung und dynamische Skalare ausgerichtet sind.

Wichtigste Punkte & Zeitstempel

• 0:00 Einführung in Flight DSL: Vorstellung der Integration von Flight-DSL-Kerneln in JAX für AMD-GPUs durch Arty vom AMD-Jax-Team.
• 0:27 Grenzen von XLA: XLA generiert zwar effiziente fusionierte Kernel, aber leistungskritische Workloads wie MoE-GEMMs erfordern eine präzise Hardware-Kontrolle über Tiling, Layouts und MFMA-Kerne.
• 1:38 Architektur von Flight DSL: Layout-fokussiertes Design (inspiriert von Cute DSL) mit voller Kontrolle von der Block- bis zur Thread- und Atom-Ebene sowie Unterstützung für manuelle Kernel-Fusions.
• 2:41 Erstellung und JAX-Brücke: Verwendung von Python-Dekoratoren (@flyt.kernel, @flyt.jit) und einer neuartigen JAX-zu-Flight-DSL-Brücke, die JAX-Tensoren via DLPack in Memrefs übersetzt.
• 5:08 Kompilierung und Dispatch-Pipeline: Compile-Time-Kachelung, Absenkung zu FLIR/HCCS und C-Trampoline-Schichten zur Puffer-Neuverpackung für den mandantenbasierten XLA-Dispatch pro Shard.
• 5:58 Optimierung von MoE-Workloads: Ersetzung ineffizienter maskierter Dense-GEMMs durch optimierte Group-GEMMs für unregelmäßige Gather/Scatter-Muster und dynamische Expertengrößen bei Mixtral-Modellen.
• 7:29 Standalone-Validierung: Vergleich des Flight-DSL-MoE-Kernels mit dem Aether-CK-Baseline-Kernel bei FP8-Präzision zeigt Leistungsparität und in einigen Fällen Überlegenheit.
• 8:12 MaxText-Integration & Performance: End-to-End-Integration im MaxText-Framework (Mixtral 8x7B, BF16) erzielt numerische Parität und massive Beschleunigungen: 2,5x im Prefill und 3,5x im Decoding (Steigerung von 14 auf 50 Token/Sekunde).
• 10:07 Ausblick: Zukünftige Meilensteine umfassen die Unterstützung von Rückwärts-Kerneln (Training), dynamischen Skalaren und die Veröffentlichung vorkompilierter Python-Wheels.

Glossar

• JAX: Ein Python-Framework für Hochleistungs-Berechnungen und automatisches Differenzieren, häufig in der modernen KI-Forschung eingesetzt.
• XLA (Accelerated Linear Algebra): Ein domainspezifischer Compiler für lineare Algebra, der Rechenoperationen optimiert und zu effizienten GPU-Kerneln zusammenlegt.
• Triton: Eine Programmiersprache und ein Compiler, die das Schreiben hocheffizienter GPU-Kernel in Python erleichtern.
• MFMA-Kerne (Matrix Fused Multiply Accumulate): Spezielle Hardware-Einheiten auf AMD-GPUs, die für beschleunigte Matrixmultiplikationen ausgelegt sind.
• Flight DSL: Eine domainspezifische Programmiersprache in Python zur Low-Level-Kontrolle über Speicherlayouts, Threads und Hardware-Pipelines auf AMD-GPUs.
• Cute DSL: Eine domainspezifische Sprache für Tensor-Layout-Manipulationen, die als architektonische Inspiration dient.
• ROCm: Die offene Software-Plattform von AMD für GPU-Computing und -Programmierung.
• HSACO: Das kompilierte Binärformat (Heterogeneous System Architecture Code Object), das direkt auf AMD-GPUs ausgeführt wird.
• FLIR / HCCS: Interne Compiler-Zwischenrepräsentationen und Kompilierungsstufen in der Flight-DSL-Infrastruktur.
• DLPack: Ein offener Industriestandard zum effizienten Teilen von Tensor-Daten im Speicher zwischen verschiedenen Frameworks (z. B. JAX und PyTorch).
• Memref: Eine Speicherreferenz-Datenstruktur in Compiler-Infrastrukturen (MLIR), die Zeiger, Dimensionen und Speicherabstände (Strides) verwaltet.
• GEMM (General Matrix Multiply): Eine grundlegende mathematische Operation (Matrixmultiplikation), die das Rückgrat von Deep-Learning-Berechnungen bildet.
• MoE (Mixture of Experts): Eine KI-Architektur, bei der Eingaben dynamisch an verschiedene Teilnetzwerke ("Experten") geroutet werden, um Rechenleistung zu sparen.
• Swiglu: Eine in modernen Sprachmodellen populäre Aktivierungsfunktion (Swish-Gated Linear Unit).
• FP8 / BF16: Gleitkommaformate mit niedriger Präzision (8-Bit bzw. 16-Bit Bfloat), die den Speicherbedarf und die Rechenzeit drastisch reduzieren.
• Aether CK (Composability Kernel): Ein hochoptimierter Open-Source-Kernel-Baseline für AMD-GPUs.
• MaxText: Ein Open-Source-LLM-Inferenz- und Trainings-Framework, das primär für TPUs entwickelt wurde, aber auch auf GPUs läuft.
• Prefill & Decode: Die zwei Hauptphasen der Ausführung in Sprachmodellen; "Prefill" verarbeitet den Eingabetext, während "Decode" Token für Token sequenziell generiert.
• Shard: Ein Teilbereich (Fragment) eines Modells oder Datensatzes, der auf eine spezifische GPU in einem Cluster verteilt ist.
• Kernel-Fusions: Das Zusammenlegen mehrerer aufeinanderfolgender Rechenschritte in einen einzigen GPU-Kernel, um teure Speicherzugriffe zu minimieren.

Zusammenfassung

Dieser Vortrag von Yangi Guan (Google) stellt Grain vor, ein effizientes Datenverarbeitungswerkzeug für das Jax-Modelltraining, das das sogenannte „Last-Mile-Problem“ der Datenzufuhr zu Machine-Learning-Modellen löst. Grain bietet flexible, deterministische und fehlertolerante Pipelines mit nativer Unterstützung für Python- und C++-Transformationen. Das Framework integriert sich nahtlos in gängige Datensätze (wie TFDS, Hugging Face, PyTorch) und Speicherformate (ArrayRecord, Parquet, TFRecords, Zarr) und bietet fortschrittliche Parallelisierungs- sowie Skalierungsmechanismen für TPUs und GPUs.

Wichtigste Punkte & Zeitstempel

• 0:00 Einführung in Grain: Ein von Google entwickeltes Framework für das letzte Glied der Datenverarbeitung vor dem ML-Training, das Daten von lokalen und Cloud-Speichern einliest, transformiert und an TPUs oder GPUs übermittelt.
• 1:06 Kernvorteile von Grain: Bietet Python-Flexibilität, absolute Deterministik über Pipeline-Läufe hinweg, Ausfallsicherheit durch Checkpoints bei gleichbleibender Deterministik nach Fehlern sowie Hochleistungsparallelisierung (Multi-Threading und Multi-Processing).
• 3:10 Vielseitige Operationen: Unterstützt Operationen wie map, random_map, shuffle, repeat, filter, batch und flat_map sowie parallele Prefetch-Methoden zur Erhöhung der Auslastung (Duty-Cycle) der TPU.
• 4:49 Kompatibilität von Datenquellen: Unterstützt populäre Framework-Datensätze wie TensorFlow Datasets (TFDS), Hugging Face und PyTorch sowie Formate wie Apache Parquet, ArrayRecord, TFRecords, Zarr und In-Memory-Quellen über GCS, AWS oder lokale Speicher.
• 5:47 Jax-Integration: Konkretes Anwendungsbeispiel zeigt das Laden des MNIST-Datensatzes, Vorverarbeitung (Normalisierung), Batching (Größe 256) und die Verteilung über JAX-Sharding (device_put) in einer Single-Host-Multi-Device-Architektur.
• 9:01 Geplante Erweiterungen: Umfasst horizontale Skalierung für ressourcenintensive Workloads wie Video- und Audioprozessierung, Serialisierung von Transformationen zur Vermeidung von Trainings-Inferenz-Abweichungen (Training-Serving Skew) sowie einen nativen C++-Backend-Modus zur Eliminierung von Python-Overhead.

Glossar

• Jax: Ein von Google entwickeltes Hochleistungs-Framework für numerische Berechnungen und maschinelles Lernen, das automatische Differenzierung und JIT-Kompilierung unterstützt.
• TPU (Tensor Processing Unit): Ein von Google entwickelter, anwendungsspezifischer integrierter Schaltkreis (ASIC), der speziell zur Beschleunigung von Deep-Learning-Workloads optimiert ist.
• Determinismus: Die Eigenschaft eines Systems, bei identischen Eingaben und Bedingungen jedes Mal exakt dieselbe Datensequenz zu erzeugen, was für die Reproduzierbarkeit von ML-Modellen entscheidend ist.
• ArrayRecord: Ein von Google entwickeltes, hochleistungsfähiges Dateiformat für die effiziente Speicherung und sequenzielle oder zufällige Lesevorgänge großer Datensätze im Machine-Learning-Bereich.
• TFDS (TensorFlow Datasets): Eine Sammlung von vorbereiteten, standardisierten Datensätzen für den direkten Einsatz mit Machine-Learning-Frameworks.
• Parquet (Apache Parquet): Ein spaltenorientiertes Speicherformat für strukturierte Daten, optimiert für hohe Komprimierungs- und Abfragegeschwindigkeiten.
• Zarr: Ein offenes Dateiformat zur Speicherung von großen, n-dimensionalen numerischen Arrays mit Schwerpunkt auf parallelem und chunk-basiertem Zugriff.
• Training-Serving Skew (Trainings-Inferenz-Abweichung): Abweichungen oder Fehler, die auftreten, wenn die Datenvorverarbeitung oder Logik während des Modelltrainings nicht exakt identisch mit der in der produktiven Inferenz-Umgebung ist.
• TFRecords: Ein einfaches Binärformat zur Speicherung einer Sequenz von Datensätzen, das in TensorFlow-Workflows häufig verwendet wird.

Zusammenfassung

Diese Präsentation stellt TokuMX vor, eine kuratierte Bibliothek maßgeschneiderter, hochleistungsfähiger Kernel für das Training und die Inferenz von KI-Modellen. Die Bibliothek unterstützt modernste Hardwareplattformen, darunter NVIDIA-GPUs (B200, H100) sowie Google-TPUs (Ironwood und TPU 8). Sie bietet umfassende Produktionsinfrastruktur wie automatisches Hyperparameter-Tuning, Xprof-basierte Benchmarks, automatisierte Regressionstests und numerische Validierung gegen XLA als Referenz. Zu den zukünftigen Roadmap-Schwerpunkten gehören Kernel für lineare und komprimierte Attention (z. B. Mamba, DeepSeek), Optimierungen für Mixture-of-Experts (MoE), PyTorch-native APIs über JAX/Palace-Backends sowie KI-gestützte Entwicklungs-Workflows.

Wichtigste Punkte & Zeitstempel

• 0:05 TokuMX-Bibliothek: Vorstellung einer zentralisierten Bibliothek für kuratierte, hochleistungsfähige Kernel zur Optimierung von KI-Modelltraining und -inferenz.
• 0:08 Hardware-Unterstützung: Zielsetzung der Unterstützung aktueller Hardwaregenerationen, darunter NVIDIA B200 und H100 sowie Google TPUs (Fokus auf Ironwood und in Entwicklung befindliche TPU 8).
• 2:14 Automatisches Tuning: Durchführung einer exhaustiven Suche im Hyperparameter-Raum für spezifische Eingangsshape-Kombinationen mit Caching zur Laufzeitoptimierung.
• 4:02 Benchmarking und Regressionstests: Einsatz von Xprof zur präzisen Leistungsmessung und automatisierte Erkennung von Leistungsregressionen im Software-Stack.
• 5:09 Numerische Validierung: Regelmäßiger Abgleich der Kernel-Ausgaben gegen XLA als Gold-Referenz innerhalb definierter Toleranzen zur Vermeidung von Konvergenzfehlern.
• 7:13 Kernel-Roadmap: Erweiterung des Katalogs um lineare Attention (Gated Delta Net, Mamba), komprimierte Attention (DeepSeek, Kimi) und Mixture-of-Experts (MoE) Routing-Operationen.
• 11:14 Ecosystem-Integration: Bereitstellung einer idiomatischen, PyTorch-nativen API, während im Backend JAX und Palace als Low-Level-Ausführungsumgebung verborgen werden.
• 12:59 Framework-Anbindung: Native Integration in übergeordnete Frameworks wie VLLM, MaxText und TuneX zur Vermeidung von Fragmentierung.
• 13:50 KI-gestützte Entwicklung: Nutzung von Agenten-Workflows zur Kernel-Optimierung und Erfassung menschlicher Best Practices als Referenzdatenpunkte.
• 15:22 Infrastruktur-Erweiterungen: Integration des XLA-Profilers zur automatischen Erkennung von Eingangsshapes und Bestrebungen zur Open-Source-Öffnung grundlegender Komponenten.

Glossar

• Kernel: Ein hochoptimierter Programmcode zur direkten Ausführung spezifischer, rechenintensiver Operationen (z. B. Matrixmultiplikationen) auf Hardware-Beschleunigern wie GPUs oder TPUs.
• TPU (Tensor Processing Unit): Ein von Google entwickelter, anwendungsspezifischer integrierter Schaltkreis (ASIC), der speziell für die Beschleunigung von Workloads im Bereich des maschinellen Lernens optimiert ist.
• B200 / H100: High-End-Grafikprozessoren (GPUs) von NVIDIA, die primär für das Training und den Betrieb großer Sprachmodelle (LLMs) eingesetzt werden.
• Auto-Tuning: Ein automatisierter Prozess, bei dem verschiedene Hyperparameter-Kombinationen systematisch getestet werden, um die optimale Ausführungsleistung für spezifische Eingabedaten zu ermitteln.
• Xprof: Ein Analyse- und Profiling-Tool von Google zur Untersuchung der Leistungsfähigkeit und Ressourcennutzung von Programmen auf TPUs und GPUs.
• XLA (Accelerated Linear Algebra): Ein domain-spezifischer Compiler für lineare Algebra von Google, der Rechengraphen optimiert und als Gold-Referenz für numerische Korrektheit dient.
• Mamba / Gated Delta Net: Architekturen für lineare Attention und zustandsbasierte Modelle (State Space Models), die eine effizientere Verarbeitung langer Kontextfenster im Vergleich zu klassischen Transformers ermöglichen.
• MoE (Mixture of Experts): Eine Modellarchitektur, bei der verschiedene Teilnetzwerke ("Experten") dynamisch je nach Eingabedatenpunkt aktiviert werden, was die Recheneffizienz erhöht.
• Quantisierung: Ein Verfahren zur Reduzierung der Bitbreite von Modellgewichten (z. B. von 16-Bit auf 4-Bit oder 8-Bit), um Speicherbedarf und Rechenzeit zu verringern.
• Palace / JAX: Low-Level-Programmierwerkzeuge und domain-spezifische Sprachen (DSLs) zur Erstellung von Hochleistungs-Kerneln bzw. zur automatischen Differenzierung auf Google-Hardware.
• VLLM / MaxText / TuneX: Übergeordnete Software-Frameworks und -Stacks für das Serving und Training von Large Language Models (LLMs).

Abstract

Der Vortrag von Rabimbo (PayPal AI Lab und Google Developer Expert) analysiert einen rigorosen, praxisnahen Benchmark zum Vergleich von LLM-Inference-Workloads (Large Language Models) auf GPUs (CUDA auf acht RTX 6000 Karten über PCIe Gen 5) versus TPUs (OpenXLA, V5E/V6E via MaxText) im Jahr 2026. Ziel war es, die Diskrepanz zwischen irreführenden öffentlichen Benchmarks ("Äpfel-mit-Birnen-Vergleiche") und tatsächlichen Produktionsanforderungen zu überwinden. Durch die Verwendung eines einheitlichen OpenAI-kompatiblen Streaming-Clients wurden Modellarchitektur (Llama 3.1 8B und 70B in BF16) und Request-Mix konstant gehalten. Die Untersuchung zeigt, dass der rohe Token-Durchsatz (Throughput) trügerisch ist und stattdessen Goodput (Anfragen pro Sekunde im Rahmen definierter Service Level Objectives, SLOs) als primäre Metrik herangezogen werden muss, insbesondere bei der Berücksichtigung von P99-Latenzen und TPU-Kompilierungskosten.

Wichtigste Punkte & Zeitstempel

• 0:00 Motivation & Problemstellung: Ziel des Benchmarks war es, objektive Daten für die Entscheidung zu gewinnen, ob LLM-Serving von GPUs auf TPUs migriert werden sollte, da öffentliche Vergleiche zu stark variieren und nicht reproduzierbar sind.
• 1:51 Hardware-Setup & Tech-Stack: Fokus auf RTX 6000 GPUs (über PCIe Gen 5, ohne NVLink) im Vergleich zu TPU v5e/v6e unter Verwendung von VLLM, OpenXLA und MaxText für Llama 3.1 (8B und 70B Parameter in BF16).
• 4:34 Service Level Objectives (SLOs) & Metriken: Definition von Latenz- und Kapazitätsmetriken: Time to First Token (TTFT) für die Responsivität und Time Per Output Token (TPOT) für die Streaming-Geschwindigkeit unter Berücksichtigung von Lastgrenzen.
• 8:34 Die Tücke öffentlicher Benchmarks: Öffentliche Zahlen messen oft nur den Peak-Durchsatz bei isolierten, künstlichen Sequenzlängen ohne Berücksichtigung von Kaltstarts, Präzision, Client-Typen oder realistischen Request-Mixes.
• 10:14 Kontrollierter Benchmark-Aufbau: Um eine faire Vergleichbarkeit zu gewährleisten, wurde ein einziger OpenAI-kompatibler Streaming-Client genutzt, während sich nur das Backend (VLLM-CUDA vs. TPU-Backends) änderte.
• 13:03 Vier Kernmetriken: Unterscheidung zwischen Latenz (TTFT, TPOT) und Kapazität (Roher Durchsatz vs. Goodput, wobei Goodput nur Anfragen zählt, die streng die definierten SLOs einhalten).
• 16:05 Workload-Formen: Analyse von Prefill-lastigen Workloads (langer Prompt, kurze Antwort, rechenintensiv) und Decode-lastigen Workloads (kurzer Prompt, lange Antwort, speicherbandbreitenintensiv).
• 18:31 TPU-Kompilierung & Kaltstart: Erklärung der statischen OpenXLA-Kompilierung auf TPUs, die für dynamische Serving-Längen vordefinierte Tensor-Shape-Buckets und zeitaufwändige Kaltstartphasen (bis zu 20–30 Minuten, ca. 5 Minuten aus dem Cache) erfordert.
• 21:28 Benchmark-Ergebnisse (Durchsatz vs. Latenz): Das 8B-Modell läuft stabil innerhalb der SLOs (P99: 33–88 ms). Das 70B-Modell auf mehreren GPUs ohne NVLink zeigt jedoch explodierende P99-Latenzen (bis zu 462 ms), was zu einem massiven Einbruch des Goodputs führt.
• 25:08 Wichtigste Erkenntnisse: Das Festhalten an fixen Clients, einheitlichen SLOs und der Metrik Goodput offenbart Leistungsschwächen in der Skalierbarkeit, die reiner Durchsatz verschleiern würde.

Glossar

• LLM (Large Language Model): Ein auf maschinellem Lernen basierendes Sprachmodell, das darauf trainiert ist, menschliche Sprache zu verstehen und zu generieren.
• TPU (Tensor Processing Unit): Ein von Google entwickelter, spezialisierter ASIC (anwendungsspezifischer integrierter Schaltkreis) zur Beschleunigung von Workloads im Bereich des maschinellen Lernens.
• GPU (Graphics Processing Unit): Ein Grafikprozessor, der dank paralleler Rechenarchitektur universell für rechenintensive Aufgaben, insbesondere KI-Training und -Inference, eingesetzt wird.
• CUDA: Eine von NVIDIA entwickelte Plattform und Programmierschnittstelle (API) zur Nutzung von GPUs für allgemeine Berechnungen und paralleles Rechnen.
• OpenXLA: Ein Open-Source-Compiler und Optimierungsframework für maschinelles Lernen, das rechenintensive Modelle für Hardware-Beschleuniger wie TPUs und GPUs optimiert.
• VLLM: Eine populäre, hochoptimierte Open-Source-Software-Engine zum schnellen und speichereffizienten Serving von Large Language Models.
• TTFT (Time to First Token): Die Latenzzeit, die von der Anfrage eines Nutzers bis zur Generierung des allerersten Tokens durch das Modell vergeht (wesentliches Maß für die gefühlte Responsivität).
• TPOT (Time Per Output Token): Die Zeit, die für die Generierung jedes weiteren Tokens während des Streaming-Prozesses benötigt wird (Maß für die kontinuierliche Ausgabegeschwindigkeit).
• SLO (Service Level Objective): Ein präzise definiertes Leistungsziel oder qualitativer Standard (z. B. maximale Antwortlatenz), den ein System im Betrieb einhalten muss.
• Goodput: Die Anzahl nutzbarer, erfolgreicher Anfragen pro Sekunde, die alle definierten SLOs (wie Latenzgrenzen) einhalten – im Gegensatz zum reinen Gesamtdurchsatz.
• Throughput (Durchsatz): Die Gesamtzahl der verarbeiteten Tokens oder Anfragen pro Sekunde, unabhängig davon, ob vereinbarte Latenzziele (SLOs) verletzt wurden.
• P99 (99. Perzentil): Ein statistischer Wert, der aussagt, dass 99 % aller gemessenen Anfragen schneller oder gleich dieses Wertes waren; dient zur Messung von Verzögerungen im Extrembereich (Tail Latency).
• Prefill: Die rechenintensive Phase am Anfang der Inferenz, in der das gesamte Eingabedokument (der Prompt) parallel verarbeitet wird.
• Decode: Die sequenzielle Phase der Inferenz, in der Token für Token generiert wird; stark abhängig von der verfügbaren Speicherbandbreite.
• PCIe (Peripheral Component Interconnect Express): Ein serieller Computerbus-Standard zur Verbindung von Hardware-Komponenten (wie GPUs oder Netzwerkkarten) mit dem Motherboard.
• NVLink: Eine proprietäre Hochgeschwindigkeits-Verbindungstechnologie von NVIDIA für die direkte, verlustarme Kommunikation zwischen mehreren GPUs.
• BF16 (Bfloat16): Ein 16-Bit-Gleitkommaformat, das häufig im Deep Learning verwendet wird, um den Speicherbedarf zu halbieren, während der dynamische Bereich erhalten bleibt.
• MaxText: Ein hochperformantes, auf JAX basierendes Open-Source-Textmodellierungsframework von Google für das Training und Serving von LLMs auf TPUs.

Zusammenfassung

Dieses Video bietet eine detaillierte technische Analyse der Workload-Optimierung für Google TPUs (Tensor Processing Units) mittels des Profiling-Tools XPRO. Ein Google-Produktmanager erläutert die TPU-Architektur, typische Leistungsengpässe (wie das „Blindfliegen“) und demonstriert anhand einer Mistral 8x22-Fallstudie, wie gezielte Anpassungen der Batch-Größe die Rechenleistung maximieren. Zudem werden neue Funktionen wie kontinuierliche Profiling-Snapshots (ein 2-GB-Host-Zirkularpuffer für 90 Sekunden Verlauf) und die native PyTorch-Integration vorgestellt.

Wichtigste Punkte & Zeitstempel

• 0:00 Einführung in XPRO: Präsentation des TPU-Software-Suiten-Tools zur Vermeidung von „Blindflügen“, die zu ineffizienten Ressourcenkosten und langsamen Iterationszyklen führen.
• 1:16 TPU-Architektur: Systemweites Gleichgewicht aus Tensor Core (dichte Matrizenrechnungen), Sparse Core (Einbettungssuchen), HBM (High Bandwidth Memory mit 35 TB/s) und direktem ICI-Interconnect.
• 3:01 Arbeitsablauf der Ausführung: Die Host-CPU steuert die Datenpipeline und kompiliert Code über XLA zu Hardware-Binärdateien, die über den PCIe-Bus an den TPU-Beschleuniger übertragen werden.
• 5:04 Fallstudie zu Leistungsengpässen: Analyse eines Mistral 8x22-Modells auf Trillium-Hardware; das anfängliche Profil zeigt eine hohe All-Reduce-Kommunikationszeit (31 %) bei geringer FLOP-Auslastung (27 %).
• 7:36 Speicher-Viewer & Optimierung: Eine Erhöhung der Batch-Größe von 1 auf 8 steigert die Speicherauslastung von 12 GB auf 29 GB (von 32 GB) und hebt die FLOP-Auslastung auf 60,74 %, wodurch das System in einen optimalen, rechengebundenen Zustand wechselt.
• 9:51 Kernel-Profiling: Tiefenanalyse für benutzerdefinierte Kernel (Palace, Mosaic, Triton) mit Zugriff auf HLO-Graphen, LLO-Bündel und über 16.000 rohe Hardware-Performance-Zähler.
• 14:36 Kontinuierliche Profiling-Snapshots: Ein permanenter Flight-Recorder auf Basis eines 2-GB-Host-Zirkularpuffers speichert die letzten 90 Sekunden, um transiente Anomalien und unerwartete Leistungsabfälle abzufangen.
• 15:59 PyTorch-Kompatibilität: Native TPU-Profilierung über Kineto/torch_tpis ohne intrusive Code-Änderungen oder benutzerdefinierte Wrapper.

Glossar

• TPU (Tensor Processing Unit): Ein von Google entwickelter, spezialisierter Mikroprozessor, der gezielt für die Beschleunigung von Workloads im Bereich des maschinellen Lernens (KI) optimiert ist.
• XLA (Accelerated Linear Algebra): Ein domänenspezifischer Compiler für lineare Algebra, der maschinelle Lernmodelle optimiert, um die Ausführungsgeschwindigkeit auf Hardware zu maximieren.
• Tensor Core: Ein Hardware-Kern innerhalb der TPU, der für dichte, rechenintensive Matrixmultiplikationen optimiert ist.
• Sparse Core: Ein spezieller Rechenkern der TPU für unregelmäßige, speicherintensive Operationen wie Embedding-Lookups (Einbettungssuchen).
• HBM (High Bandwidth Memory): Ein ultraschneller, dreidimensional gestapelter Arbeitsspeicher, der direkt auf dem Prozessorsubstrat sitzt und extrem hohe Datenübertragungsraten (hier 35 TB/s) ermöglicht.
• ICI Interconnect: Ein proprietäres Hochgeschwindigkeits-Netzwerkprotokoll, das die direkte Kommunikation zwischen mehreren TPU-Chips ohne Einbindung der Host-CPU erlaubt.
• All-Reduce: Eine verteilte Rechenoperation in Cluster-Systemen, bei der Daten (wie Gradienten) von allen Rechenknoten gesammelt, mathematisch zusammengefasst und synchronisiert werden.
• HLO (High-Level Optimizer): Eine maschinenunabhängige Zwischenrepräsentation von Berechnungen innerhalb des XLA-Compilers zur Optimierung von Operationen auf hoher Abstraktionsebene.
• Kineto / PyTorch Profiler: Ein in PyTorch integriertes Werkzeug zur Messung von Laufzeiten, Speicherverbrauch und Ressourcen-Auslastung von Operationen auf Hardware-Beschleunigern.

Abstract

In diesem Vortrag stellt Matt Johnson vom JAX-Core-Team die Einführung und Architektur des mutablen Referenz-APIs (ref) in JAX vor. Während JAX traditionell auf einem rein funktionalen, unveränderlichen (immutablen) Paradigma basierte und Optimierungen vollständig dem XLA-Compiler überließ, bietet das ref-API Entwicklern nun die Möglichkeit, gezielt die Kontrolle über den physischen Speicher zu übernehmen. Der Vortrag behandelt die Ursprünge von refs im Kernel-Framework Palace, statische Schutzmaßnahmen gegen Speicheraliasing, die Kompatibilität mit Transformationen wie vmap und automatischer Differenzierung sowie die effiziente In-Place-Akkumulation von Gradienten zur Bewältigung von Leistungsproblemen bei Sparse-Operationen und Pipelining.

Wichtigste Punkte & Zeitstempel

• 0:00 JAX-Grundlagen: Verwendung von JAX für numerische Berechnungen, neuronale Netze, automatische Differenzierung und die End-to-End-Kompilierung via XLA auf GPUs, TPUs, Tranium und AMD-Hardware.
• 1:49 Das Leistungs-Paradoxon: Verschiebung von rein deklarativem Code, bei dem der XLA-Compiler alle Optimierungen übernimmt, hin zu imperativer Hardware-Kontrolle für maximale Skalierbarkeit und Rechenleistung.
• 2:56 Das ref-API: Einführung veränderbarer Referenzen zur Lösung von Ausdrucksfähigkeits- und Performance-Pragmatiken im funktionalen Paradigma, insbesondere für zustandsbehaftete Operationen wie Batch-Normalization oder Metrik-Sammlungen.
• 5:40 Ursprung und Aliasing-Schutz: Entstehung der Referenzen im Kernel-Framework Palace; statische Verhinderung von Aliasing durch strenge Einschränkungen wie das Verbot des Zurückgebens oder mehrfachen Übergebens von Referenzen.
• 7:04 Live-Demo und Programmiermodell: Demonstration von In-Place-Zuweisungen und JIT-Kompilierung; refs repräsentieren physische Speicherpuffer (keine mathematischen Arrays), bei denen vor mathematischen Operationen explizit Daten geladen werden müssen.
• 9:29 API-Methoden und Restriktionen: Verwendung von new_ref zur Erstellung und freeze zur Entwertung von Referenzen, kombiniert mit vertrauter NumPy-Indizierung unter strikten Erzeugungs- und Gültigkeitsgrenzen.
• 11:06 Transformationen und vmap: Demonstration der Vektorarisierung über veränderbare Referenzen (vmap) zur Erstellung von paarweisen Distanzmatrizen ohne funktionale Rückgabewerte.
• 14:52 Gradienten-Akkumulation bei Sparsity: Lösung von Ineffizienzen bei spärlichen Indizierungen (Sparse Reads), bei denen traditionell große Null-Arrays alloziiert werden müssen; stattdessen direkte In-Place-Akkumulation in übergebene Referenzen.
• 16:38 Pipelining-Anwendungsfall: Praktischer Einsatz von refs in geschachtelten Scan-Schleifen für Layer-Pipelining, wodurch das Beibehalten mehrerer Parameter-Kopien vermieden und die In-Place-Aktualisierung von Gradienten ermöglicht wird.
• 19:16 XLA-Interaktion und Custom Calls: Übersetzung von refs in XLA HLO Custom Calls mit Input-Output-Aliasing und Puffer-Spende (Buffer Donation), wobei Sequenzierungen erzwungen werden, die die Freiheit des Compilers in Consumer-Ketten einschränken können.

Glossar

• JAX: Ein Python-Framework für Hochleistungs-Berechnungen und maschinelles Lernen, das NumPy-ähnliche Syntax mit automatischer Differenzierung und JIT-Kompilierung verbindet.
• XLA (Accelerated Linear Algebra): Ein Domain-Specific Compiler für lineare Algebra, der JAX- und TensorFlow-Graphen optimiert und für Hardware wie GPUs und TPUs übersetzt.
• TPU / GPU / Tranium: Spezialisierte Hardware-Beschleuniger für massiv parallele Rechenoperationen im Bereich des maschinellen Lernens.
• Autodiff (Automatische Differenzierung): Ein Verfahren zur exakten und automatischen Berechnung von Ableitungen in Computercode, essenziell für das Training neuronaler Netze.
• VJP (Vector-Jacobian Product): Ein Baustein der automatischen Differenzierung im Reverse-Mode, der das Produkt aus einem Vektor und der Jacobi-Matrize zur Gradientenberechnung ermittelt.
• VMAP (Vectorizing Map): Eine JAX-Transformation, die eine Funktion über zusätzliche Batch-Dimensionen automatisiert vektorisiert, ohne die Funktion manuell umschreiben zu müssen.
• HLO (High-Level Optimizer / High-Level Operations): Die abstrakte Zwischenrepräsentation (Intermediate Representation) des XLA-Compilers.
• Palace: Eine Low-Level-Kernel-Sprache innerhalb des JAX-Ökosystems, die direkt auf Speicherreferenzen (refs) aufbaut, um hardwarenahe Berechnungen durchzuführen.
• Aliasing: Ein programmiertechnischer Zustand, bei dem zwei oder mehr Variablen oder Zeiger auf denselben physischen Speicherbereich verweisen.
• Buffer Donation (Puffer-Spende): Ein Kompilierungsmechanismus, bei dem Speicherpuffer von Eingaben direkt für Ausgaben wiederverwendet werden, um Speicherallokationen und Kopien zu vermeiden.
• Sparsity (Spärlichkeit): Eigenschaft von Matrizen oder Tensoren, bei denen ein Großteil der Elemente den Wert Null aufweist.
• Custom Call: Ein Mechanismus im XLA-Compiler, der die Einbindung von externem oder benutzerspezifischem Maschinencode in die optimierte Ausführungspipeline erlaubt.

Zusammenfassung

Der Vortrag präsentiert das JAX GCM (ein differenzierbares, atmosphärisches Modell mittlerer Komplexität), entwickelt von Forschern der Scripps Institution of Oceanography, der UC San Diego und des Data Science Institute. Das System ersetzt traditionellen, schwer anpassbaren Fortran-Klimamodellcode durch eine vollständig in JAX implementierte Architektur. Durch die Bereitstellung von analytischen Gradienten (im Vorwärts- und Rückwärts-Modus) ermöglicht das Modell effizientes Online-Lernen, Parameterestimation, Datenassimilation und die nahtlose Kombination aus traditionellen physikalischen Parametrisierungen und neuronalen Netzen (Hybrid-Modellierung). Erste Validierungen zeigen eine hohe Übereinstimmung mit Fortran-Baselines (z. B. eine 15-tägige Simulation in 26 Sekunden auf einem Laptop). Zukünftige Entwicklungen zielen auf ein vollwertiges Erdsystemmodell durch Kopplung mit Ozean- (Baros) und Land-Modellen (differ land) ab.

Wichtigste Punkte & Zeitstempel

• 0:00 Problem herkömmlicher Klimamodelle: Traditionelle Modelle basieren auf Millionen Zeilen Fortran-Code, die für CPU-basierte HPC-Systeme optimiert, schwer lesbar und nicht differenzierbar sind, was die Optimierung von Hunderten Parametern erschwert.
• 2:01 Integration von JAX & Neural GCM: In Anlehnung an Googles Neural GCM (welches den dynamischen Kern dinosaur nutzt) werden physikalische Parametrisierungen (wie Wolkenbildung und Aerosole) in JAX übersetzt, um sie differenzierbar zu machen.
• 3:51 Erweiterte Optimierungsmethoden: Die Verfügbarkeit von Gradienten ermöglicht Online-Lernen, präzise Parameterestimation, Datenassimilation (Optimierung von Anfangsbedingungen) und die kosteneffiziente Generierung von Ensembles zur Unsicherheitsquantifizierung.
• 6:09 Plug-and-Play-Hybrid-Modellierung: Das JAX GCM erlaubt den direkten Vergleich und Austausch zwischen traditionellen, physikbasierten Parametrisierungen und maschinellen Lernmodellen (neuronalen Netzen) innerhalb desselben dynamischen Kerns.
• 6:49 Validierung & API-Design: Das Modell wurde gegen eine Fortran-Baseline validiert und bietet eine moderne API für Lehre und Forschung; eine 15-tägige Simulation auf einem Standard-Laptop benötigt lediglich 26 Sekunden.
• 7:41 Sensitivitätsanalysen: Durch Vorwärts- und Rückwärtsgradienten lassen sich die Auswirkungen von Störungen (z. B. Meeresoberflächentemperaturen oder Albedo-Änderungen) über die Zeit präzise analysieren.
• 8:29 Erdsystem-Kopplung: Es wird an der Kopplung des Atmosphärenmodells mit Ozeanmodellen (Baros) und Landmodellen (differ land) gearbeitet, um ein vollständig gekoppeltes Erdsystemmodell mit gemeinsamen Gradienten zu realisieren.

Glossar

• JAX: Eine High-Performance-Bibliothek für numerisches Computing in Python, die automatische Differentiation (Gradientenberechnung) und JIT-Kompilierung für GPUs und TPUs unterstützt.
• GCM (General Circulation Model): Ein mathematisches Modell der allgemeinen Zirkulation der Atmosphäre oder des Ozeans, das fundamentale physikalische Gesetze (wie die Navier-Stokes-Gleichungen) integriert.
• Fortran: Eine ältere Programmiersprache, die traditionell in der Hochleistungsrechner- und Klimaforschung für rechenintensive numerische Simulationen eingesetzt wird.
• Dynamischer Kern (Dynamical Core): Der rechenintensive Teil eines Klimamodells, der die Grundgleichungen der Fluid- und Thermodynamik (Navier-Stokes-Gleichungen) auf einer Kugel integriert.
• Parametrisierung (Physics Parameterization): Die mathematische Darstellung subskalierter, nichtlinearer Prozesse (z. B. Wolkenbildung, Aerosolwechselwirkungen) in Klimamodellen, die nicht direkt aufgelöst werden können.
• Gradient: Die mathematische Ableitung einer Funktion nach ihren Eingangsvariablen; essenziell für gradientenbasierte Optimierungsverfahren im Machine Learning und der Parameteranpassung.
• Datenassimilation: Ein mathematisches Verfahren, bei dem reale Beobachtungsdaten in ein dynamisches Modell integriert werden, um die optimalen Anfangs- oder Randbedingungen zu bestimmen.
• Ensemble: Eine Gruppe von parallelen Simulationsläufen mit leicht variierten Anfangsbedingungen oder Parametern, um die Unsicherheit von Vorhersagen zu quantifizieren.
• Navier-Stokes-Gleichungen: Ein System partieller Differentialgleichungen, das die physikalische Bewegung von fluiden Medien (wie Luftmassen und Ozeanwasser) beschreibt.
• Albedo: Das Rückstrahlvermögen (Reflexionsvermögen) von Oberflächen, insbesondere von Schnee, Eis, Wolken und der Erdoberfläche.

Zusammenfassung

Dieser Vortrag des JAX-Entwicklungsteams bietet eine technische Einführung und Programmierübersicht zum Schreiben von Palace-Kernen (Kernels) auf TPU-Sparse-Cores (Dünnschichtkernen). Es werden die Hardware-Architektur (Skalar- und Vektor-Subcores, Speicherhierarchie), das Programmiermodell, erweiterte Speicheroperationen wie Gather und Scatter sowie MPMD-Kerne (Multiple Program, Multiple Data) zur Synchronisation paralleler Workloads für Large Language Models (LLMs) behandelt.

Wichtigste Punkte & Zeitstempel

• 0:14 TPU-Hardware-Architektur: TPU-Sparse-Cores bestehen typischerweise aus 1 Skalar-Subcore (zuständig für Kontrollfluss, Instruktionsladen und dynamische Indizierung) und 16 Vektor-Subcores (für SIMD-Operationen mit einer Vektorlänge von 16 für F32 auf Ironwood).
• 2:46 Speicher- und DMA-Hierarchie: Jedes Subcore besitzt lokalen Speicher (SM/VM) sowie einen gemeinsamen VM-Raum; Subcores können unabhängig voneinander lokale und Remote-DMAs (Direct Memory Access) initiieren.
• 5:46 Stärken des Sparse-Cores: Optimiert für nicht-matrixbasierte Vektorarithmetik, unregelmässige parallele Workloads, hochfrequente Datenkommunikation (Kollektive) sowie Spezialoperationen wie Sortierung, Histogramme und Top-K-Sampling für LLMs.
• 7:39 Gather und Scatter Operationen: Bieten spezialisierte, unregelmässige Speicherzugriffe (Index-Lookups), die essenziell für Expert-Token-Routing und Embedding-Lookups sind; unterstützt sowohl auf VM- als auch auf HBM-Ebene.
• 11:03 Palace-Programmiermodell: Baut auf dem JAX-Gerätenetz (Device Mesh) auf; Kernels werden mit dem pl_kernel-API umschlossen, wobei ausschliesslich mit Referenzen (refs) und veränderlichen Puffern gearbeitet wird.
• 15:20 Vektor-Pipelines: Nutzen das image_pipeline-API zur automatischen Parallelisierung und effizienten Ablaufsteuerung über Vektor-Subcores hinweg.
• 20:25 MPMD-Kerne: Ermöglichen das gleichzeitige Ausführen mehrerer Kerne auf unterschiedlichen Subcore-Typen, die über Semaphoren für Barrieren und Kollektive synchronisiert werden.

Glossar

• TPU (Tensor Processing Unit): Ein von Google entwickelter, spezialisierter ASIC zur Beschleunigung von Machine-Learning- und Tensor-Berechnungen.
• Sparse Core: Ein spezialisierter Rechenkern innerhalb der TPU, der für unregelmässige, nicht-matrixbasierte Operationen und dünnbesetzte (sparse) Daten optimiert ist.
• JAX: Ein Python-Framework für Hochleistungs-Numerik und maschinelles Lernen mit automatischer Differenzierung.
• Palace: Eine Low-Level-Programmierschnittstelle innerhalb des JAX-Ökosystems zum direkten Schreiben von TPU-Kernels.
• SIMD (Single Instruction, Multiple Data): Eine Rechenarchitektur, bei der eine einzige Instruktion parallel auf mehrere Datenelemente eines Vektors angewendet wird.
• HBM (High Bandwidth Memory): Ein sehr schneller, speicherbandbreitenstarker 3D-RAM, der direkt neben dem Hauptprozessor integriert ist.
• DMA (Direct Memory Access): Ein Hardware-Mechanismus, der Datenübertragungen zwischen Speicherbereichen ohne ständige Beteiligung der Haupt-CPU ermöglicht.
• Gather / Scatter: Speicheroperationen zum unregelmässigen Abrufen (Gather) oder Schreiben (Scatter) von Daten an nicht zusammenhängenden Speicheradressen anhand von Indizes.
• MPMD (Multiple Program, Multiple Data): Ein paralleles Berechnungsmodell, bei dem verschiedene Prozessorkerne unterschiedliche Programme auf separaten Datensätzen ausführen.
• F32 (Floating-Point 32): Ein 32-Bit-Gleitkommaformat zur Darstellung reeller Zahlen mit einfacher Genauigkeit.
• Tensor Core: Der primäre Rechenkern einer TPU, der auf massgeschneiderte Matrixmultiplikationen spezialisiert ist.

Abstract

Dieses Transkript dokumentiert eine technische Präsentation über Leistungsoptimierung und Debugging für JAX- und OpenXLA-Workloads, vorgetragen von Ingenieuren von Google und NVIDIA. Im Fokus steht das Tool XPro, ein plattformübergreifender ML-Performance-Profiler mit nativer Unterstützung für TPUs und GPUs. Es werden detaillierte Einblicke in fortgeschrittene Funktionen wie Trace-Ansichten, Speicher-Viewer, Roofline-Analysen und die Überwachung von Custom Calls (z. B. via Palace- oder Triton-Kernel) gegeben. Zudem stellen die NVIDIA-Kollaborateure neue PM-Sampling-Funktionen und Range-Profiling-Methoden vor, die eine tiefe Hardware-Metrikenerfassung für Hopper-GPUs und neuere Architekturen ermöglichen.

Wichtigste Punkte & Zeitstempel

• 0:00 Einführung und Agenda: Präsentation von XPro (OpenXLA Performance Profiler) durch Prasana Reangasami (Google) gemeinsam mit NVIDIA-Partnern zur Leistungs- und Energieanalyse von ML-Workloads.
• 1:47 Überblick über XPro: Der ML-Performance-Profiler unterstützt nativ TPUs und GPUs sowie andere Beschleuniger über den PJRT-Standard und bietet Ansichten für Trace, Graph, Speicher und Roofline.
• 4:00 Profile-Erfassung & Konfiguration: Erläuterung der UI-basierten und programmatischen Erfassung via zax.profiler und torch.profiler inklusive erweiterter Konfigurationsmöglichkeiten.
• 6:01 Visualisierungs-Tools: Detaillierte Darstellung von Gantt-Chart-Zeitlinien, Speicherallokationen (HBM/VM) zur Erkennung von Fragmentierungen und Roofline-Modellen zur Identifizierung von Memory- vs. Compute-Bound-Zuständen.
• 9:17 Custom Calls & Analyse-Herausforderungen: Diskussion von Custom Calls (Palace- und Triton-Kernel) und den Schwierigkeiten bei der Verfolgung asynchroner DMAs und dynamischer Tensor-Formen.
• 13:30 TPU Performance-Zähler (PM Sampling): Nutzung von bis zu 60.000 Perf-Stat-Countern auf V6E/V7-TPUs (Ironwood/Trillium) mit Sampling-Raten von bis zu 833 MHz zur instruktionsgenauen Analyse.
• 24:34 NVIDIA GPU-Integration via CUPTI: Chris Ashton (NVIDIA) präsentiert die Integration des CUDA Profiler Tools Interface (CUPTI) in XPro, einschließlich PM-Sampling für Hopper++ Architekturen.
• 30:08 Range-Profiling & Multi-Pass-Ausführung: Einführung des Range-Profilings für NVIDIA-GPUs, das durch Kernel-Replay im Multi-Host-HLO-Runner den Zugriff auf über 6.000 Hardware-Metriken ermöglicht.
• 38:04 Cupy-Optimierungen & Multi-Subscriber: Neue schema-basierte Aktivitätsdatensätze zur Reduzierung von Profiling-Overhead und Unterstützung paralleler Abonnenten (z. B. XPro und NSight Systems gleichzeitig).
• 40:40 Frage-und-Antwort-Runde: Diskussionen über Zähler-Overhead (keiner bei TPUs, CPU-limitiert bei GPUs) und die Abgrenzung zwischen XPro und NVIDIA-eigenen Tools wie Insight Compute.

Glossar

• XPro (OpenXLA Performance Profiler): Ein plattformübergreifendes Analyse- und Profiling-Werkzeug zur Untersuchung von Leistung, Energieverbrauch und thermischem Verhalten von Machine-Learning-Workloads.
• OpenXLA: Ein offener Compiler- und Laufzeit-Stack für maschinelles Lernen, der optimierten Maschinencode für verschiedene Hardware-Beschleuniger generiert.
• JAX: Ein Python-Framework für Hochleistungs-Numerik und automatisches Differenzieren, das eng mit XLA integriert ist.
• PJRT: Eine standardisierte Plugin-Schnittstelle, die Machine-Learning-Frameworks mit unterschiedlichen Hardware-Beschleunigern verbindet.
• HLO (High-Level Optimizer): Eine plattformunabhängige Zwischenrepräsentation (Intermediate Representation) im XLA-Compiler zur Durchführung von Optimierungs-Pässen.
• HBM (High Bandwidth Memory): Ein extrem schneller, vertikal gestapelter Hauptarbeitsspeicher, der direkt auf modernen GPUs und TPUs verbaut ist.
• VM (Vektor- oder Virtueller Speicher): Ein lokaler, schneller Speicherbereich auf TPU-Architekturen.
• Roofline-Analyse: Ein visuelles Leistungsmodell zur Bestimmung, ob ein Rechenkernel durch die Rechenleistung (Compute-Bound) oder den Speicherdurchsatz (Memory-Bound) limitiert ist.
• Custom Calls: Direkt in den Compiler-Graph eingebettete, handoptimierte Kernel (oft in Triton oder Palace geschrieben), die spezifische Hardware-Funktionen nutzen.
• Triton / Palace: Programmiermodelle und Compiler-Dialekte zur Erstellung von Hochleistungs-Beschleunigerkerneln (für GPUs bzw. TPUs).
• PM Sampling (Performance Metric Sampling): Die regelmäßige, zeitbasierte Erfassung von Hardware-Zählern (Performance Countern) während der Ausführung eines Programms.
• CUPTI (CUDA Profiler Tools Interface): Die offizielle Programmierschnittstelle von NVIDIA zur Erstellung von Profiling- und Trace-Werkzeugen für CUDA-Anwendungen.
• Hopper++ / Blackwell (B100): Moderne NVIDIA-GPU-Architekturen mit spezialisierten Tensor-Kernen und erweiterten Leistungsüberwachungsfunktionen.
• MXU (Matrix Multiply Unit): Eine spezialisierte Hardware-Recheneinheit auf TPUs zur hocheffizienten Durchführung von Matrixmultiplikationen.

Zusammenfassung

Diese Aufzeichnung des siebten „Dev Labs“-Events beleuchtet die historische und technologische Entwicklung der Google-Tensor Processing Units (TPUs), der KI-Infrastruktur sowie der Co-Design-Strategien für Hard- und Software. Die Google-Experten Jeff Dean und Bill diskutieren Meilensteine von den Anfängen der Websuche im Jahr 1999 über die Entstehung der TPUv1 im Jahr 2015 bis hin zu modernen Groß-Pods mit über 9.600 Chips. Im Fokus stehen Skalierbarkeit, die Pathways-Architektur, optische rekonfigurierbare Netzwerke (OCS), die Steigerung der Systemzuverlässigkeit (Goodput) auf bis zu 98 %, Open-Source-Initiativen (JAX, Stable HLO, PyTorch-TPU) sowie zukünftige Anforderungen an agentenbasierte Inferenz-Workloads und automatisierte Chip-Entwürfe mittels Reinforcement Learning.

Wichtigste Punkte & Zeitstempel

• 0:00 Skalierung der Anfänge: Jeff Dean beschreibt die Anfänge von Google im Jahr 1999 in einem kleinen Büro über einem T-Mobile-Store in Palo Alto bei einem wöchentlichen Traffic-Wachstum von 7 %.
• 2:01 In-Memory-Suchindex (2001): Verlagerung des gesamten Web-Suchindexes in den Arbeitsspeicher von 1.200 Maschinen, wodurch die Suchgeschwindigkeit innerhalb von drei Tagen vervierfacht wurde.
• 3:44 Aufstieg des Deep Learnings (2013): Training großer Sprach- und Bildmodelle auf 16.000 CPU-Kernen; die Erkenntnis, dass eine sprachgesteuerte Inferenz für 100 Millionen Nutzer die gesamte damalige Rechenleistung von Google verdoppeln würde, führte zur Entwicklung der TPUs.
• 6:34 TPUv1 und Architektur-Meilenstein: Einführung der ersten TPU im Jahr 2015 als reiner Inferenzbeschleuniger mit einer 30- bis 80-fach besseren Leistung pro Watt im Vergleich zu damaligen CPUs und GPUs.
• 9:01 Evolution der TPU-Generationen: Übergang von TPUv2 (Training mit 2D-Torus) und TPUv3 (Einführung von Flüssigkühlung) zu TPUv4 (optische rekonfigurierbare Netzwerke mit mikro-einstellbaren Spiegeln) sowie den Generationen v5, v6 und v7 mit niedrigeren Präzisionsformaten (16-bit, 8-bit, 4-bit).
• 11:15 Pathways-Architektur (2018): Entwicklung von Pathways zur Steuerung unregelmäßiger, sparsamer Modelle und Bereitstellung einer einheitlichen Programmierabstraktion für zehntausende verteilte Geräte über Pod-Grenzen hinweg.
• 15:58 Goodput-Optimierung: Steigerung der Systemzuverlässigkeit ("Goodput") von anfänglich 50–60 % auf 95–98 % bei Clustern mit fast 10.000 Chips durch Flotten-Zustandsüberwachung, automatisches Gewichts-Averaging und Fehlertoleranz.
• 24:37 Open-Source-Strategie: Engagement bei Open-Source-Frameworks wie TensorFlow, JAX, Stable HLO und der Einführung von PyTorch-TPU zur Vermeidung geschlossener Ökosysteme (Walled Gardens).
• 28:01 Automatisiertes Chip-Design: Nutzung von Reinforcement Learning zur Automatisierung des Hardware-Designprozesses, um Entwicklungszyklen von 2 Jahren (150 Personen) auf 3 Monate (10 Personen) zu verkürzen.
• 30:57 Inferenz und Agenten-Workloads: Anpassung der Hardwarearchitekturen für Inferenz (Minimierung von Datenbewegungen für Modellgewichte und KV-Caches) sowie Beschleunigung interner Tools für agentenbasierte Workflows (z. B. automatisierte Code-Übersetzung wie Project Turn).
• 45:12 Netzwerkskalierung (OCS und DCN): Kombination aus internen ICI-Verbindungen (Inter-Core Interconnect), optischen Schaltern (OCS) für Pods mit knapp 9.600 Chips und Rechenzentrumsnetzwerken (Data Center Networks) für synchrone Trainingsläufe über mehrere Rechenzentren hinweg.

Glossar

• TPU (Tensor Processing Unit): Ein von Google entwickelter, hochspezialisierter Prozessor (ASIC), der gezielt für die rechenintensiven Matrizenoperationen des maschinellen Lernens optimiert ist.
• Goodput (Nutzleistung): Das Verhältnis der tatsächlich für das effektive Modelltraining nutzbaren Rechenzeit zur Gesamtrechenzeit unter Ausschluss von Ausfällen, Kommunikationspausen und Reparaturen.
• Shard (Daten-Shard): Ein logisches Teilfragment einer großen Datenbank oder eines Suchindexes, das auf mehrere Server aufgeteilt wird, um Anfragen parallel zu verarbeiten.
• Inferenz (Inference): Die Phase, in der ein fertig trainiertes KI-Modell auf neue Daten angewendet wird, um Vorhersagen, Antworten oder Entscheidungen zu generieren.
• Pathway: Eine Softwareinfrastruktur und Programmierumgebung von Google, die es einem einzigen Python-Prozess ermöglicht, zehntausende verteilte Rechenchips nahtlos als ein einheitliches System anzusteuern.
• OCS (Optical Circuit Switch): Ein optischer Netzwerkschalter, der mithilfe mikro-mechanischer Spiegel dynamische und rekonfigurierbare Lichtleiterverbindungen zwischen Rechenzentrums-Racks herstellt.
• ICI (Inter-Core Interconnect): Ein extrem schnelles proprietäres Netzwerk, das TPU-Kerne direkt miteinander verbindet, um die Kommunikation innerhalb eines Racks oder Pods zu maximieren.
• Stable HLO: Eine standardisierte Zwischenrepräsentation (Intermediate Representation) für den XLA-Compiler, die die Kompatibilität zwischen verschiedenen KI-Frameworks und Hardware-Beschleunigern gewährleistet.
• JAX: Ein von Google entwickeltes Programmier-Framework für Hochleistungs-Numerik und automatisches Differenzieren, das vor allem in der modernen KI-Forschung genutzt wird.
• KV-Cache (Key-Value Cache): Ein temporärer Zwischenspeicher bei der Textgenerierung von Sprachmodellen, der berechnete Schlüssel- und Wertvektoren früherer Wörter sichert, um doppelte Rechenschritte zu verhindern.
• PyTorch-TPU: Eine Software-Integration, die es Entwicklern ermöglicht, das populäre ML-Framework PyTorch direkt auf Google-TPUs auszuführen.
• Project Turn: Ein internes Google-Projekt zur automatisierten Übersetzung von TensorFlow-Modellen und Codedateien in das JAX-Framework inklusive automatisierter Tests und Verifikation.

Zusammenfassung

Der Vortrag eines Ingenieurs des JAX-Teams präsentiert die Entwicklung und die neuesten Leistungsmerkmale von Palace, einem Framework zum Schreiben von Low-Level-Kernels für Google-TPUs. Es wird der Übergang von der Legacy-API (palace_call) zur neuen, modulareren API (pl.kernel) beleuchtet, die mutable Referenzen (refs) aus JAX nahtlos integriert. Zudem werden die TPU-Speicherhierarchie (HBM, VMA, Scalar Memory), die Unterstützung von Mehrkernarchitekturen (Tensor-Kerne und Sparse-Kerne) sowie fortgeschrittene Techniken wie das Überlappen von Kommunikation und Berechnung (z. B. All-Gather mit Matrizenmultiplikation) erläutert.

Wichtigste Punkte & Zeitstempel

• 0:04 Einführung in Palace: Vorstellung eines Frameworks für Google-TPUs durch das JAX-Team, das maximale Hardware-Kontrolle bei hoher Iterationsgeschwindigkeit ermöglicht.
• 0:21 TPU-Speicherarchitektur: Die logische TPU-Struktur unterteilt sich in HBM (High Bandwidth Memory für Gewichte/Daten), VMA (Vector Memory für rechenintensive Matrixoperationen) und SM (Scalar Memory für Kontrollfluss und dynamisches Slicing).
• 0:51 Legacy-API (palace_call): Der ursprüngliche Ansatz bündelt Speicher-Pipelines, IO-Aliasing und Scalar-Prefetch in einer einzigen, komplexen Funktion, was die Anpassbarkeit einschränkt.
• 0:10 Neue API (pl.kernel): Ein modularer, flexibler Ansatz, der standardmäßig mit HBM-Referenzen arbeitet und kein automatisches Pipelining erzwingt, wodurch Entwickler einzelne Features gezielt kombinieren können.
• 0:12 JAX Mutable Refs: Nutzung veränderbarer Referenzen (jax.new_ref) zur direkten In-Place-Modifikation von Speicherpuffern in JAX und Palace ohne unnötige Kopier-Overheads.
• 0:15 Das emit_pipeline-Hilfsprogramm: Ein Tool innerhalb von pl.kernel, das ein optimiertes und speichereffizientes Tiling von HBM nach VMA abstrahiert.
• 0:17 Dynamisches Slicing mit Scalar Memory (SM): Ermöglicht die Ausführung unregelmäßiger oder spärlicher (sparse) Zugriffsmuster ohne Einschränkungen zur Compile-Zeit.
• 0:19 Live-Demonstration (Matmatmul): Schrittweiser Aufbau eines Matrixmultiplikations-Kernels mit Blockgrößen-Tuning, FP32-Akkumulation und Leistungsvergleich gegen XLA-Referenzen.
• 0:40 Unterstützung von Sparse-Kernen: Moderne TPU-Generationen (V5P/V6E) besitzen neben Tensor-Kernen dedizierte Sparse-Kerne, die über das Ändern des Meshs (palace.sparse_core) direkt ansprechbar sind.
• 0:43 Core-Mesh-Kommunikation: Übertragung des bekannten Geräte-Mesh-Konzepts auf die Kern-Ebene zur präzisen Signalisierung und Koordination zwischen Scalar-, Vektor- und Sparse-Kernen.
• 0:48 Asynchrones All-Gather & Overlapping: Ein fortgeschrittenes Anwendungsbeispiel zeigt, wie Tensor-Kern-Berechnungen und Sparse-Kern-Kommunikation (RDMA-basierte All-Gather-Operationen) in einer Schleife zeitlich perfekt überlappt werden, um Leerlaufzeiten zu eliminieren.

Glossar

• JAX: Ein von Google entwickeltes Python-Framework für Hochleistungsrechnen (HPC) und maschinelles Lernen, das automatische Differenzierung und JIT-Kompilierung über XLA unterstützt.
• Palace: Ein Low-Level-Programmier-Framework innerhalb des JAX-Ökosystems, das Entwicklern die direkte Erstellung maßgeschneiderter TPU-Kernels erlaubt.
• TPU (Tensor Processing Unit): Ein von Google entwickelter, anwendungsspezifischer Prozessor (ASIC), der speziell für die Beschleunigung von Deep-Learning- und Matrix-Workloads optimiert ist.
• HBM (High Bandwidth Memory): Ein extrem schneller, direkt an den Prozessor angebundener Arbeitsspeicher mit hoher Bandbreite, der große Datenmengen und Modellgewichte speichert.
• VMA (Vector Memory): Ein dedizierter, schneller On-Chip-Speicher auf der TPU, in dem rechenintensive Vektor- und Matrixoperationen stattfinden.
• Scalar Memory (SM): Ein spezialisierter Speicherbereich auf der TPU, der für Steuerungslogik, Programmverzweigungen und dynamische Indexierung zuständig ist.
• XLA (Accelerated Linear Algebra): Ein domain-spezifischer Compiler von Google, der Berechnungs-Graphen optimiert und für GPUs sowie TPUs übersetzt.
• Mutable Refs (Veränderbare Referenzen): Programmierkonzept in JAX, das Zeiger auf bestehenden Speicher bereitstellt, um diesen direkt (In-Place) zu verändern, anstatt jedes Mal neue, unveränderliche Arrays zu erzeugen.
• All-Gather: Ein kollektiver Kommunikationsvorgang in verteilten Systemen, bei dem Datenfragmente von allen Prozessoren gesammelt und an alle Beteiligten verteilt werden.
• RDMA (Remote Direct Memory Access): Eine Technologie, die den direkten Datenaustausch zwischen den Speichern verschiedener Hardware-Knoten ermöglicht, ohne die Haupt-CPU zu belasten.
• BF16 (Bfloat16): Ein 16-Bit-Gleitkommaformat (Brain Floating Point), das im Deep Learning häufig verwendet wird, um Rechenleistung und Speichereffizienz bei großem Wertebereich zu maximieren.

Abstract

Dieser Vortrag bietet eine technische Einführung in die Entwicklung von Palace-Kerneln auf dem Google TPU Sparse Core im Rahmen des JAX-Ökosystems. Der Sparse Core dient als spezialisierte Ergänzung zum Tensor Core: Während Letzterer auf dichte Matrixmultiplikationen optimiert ist, übernimmt der Sparse Core speicherintensive, unregelmäßige oder hochparallele Vektoroperationen sowie dynamische Speicherzugriffe.

Die Hardware-Architektur des Sparse Cores basiert auf einer Kombination aus einem Skalar-Subcore (zuständig für Befehlssteuerung, Kontrollfluss und Index-Prefetching) und mehreren Vektor-Subcores (in aktuellen Generationen 16 SIMD-Einheiten) mit jeweils eigenem und gemeinsam genutztem Vektorspeicher (VM / Shared VM). Das Palace-Programmiermodell abstrahiert diese Hardware über eine Device-Mesh-Struktur. Der Vortrag demonstriert den praktischen Einsatz von Vektor-Slicing, automatisiertem Pipelining (emit_pipeline) für mehrstufige HBM-zu-VM-Lookups, hardwarebeschleunigte Gather-/Scatter-Primitiven sowie MPMD-Architekturen (Multiple Program Multiple Data), bei denen Skalar- und Vektor-Subcores über Semaphoren synchronisiert parallel agieren.

Wichtigste Punkte & Zeitstempel

• 0:05 Einführung & JAX/Palace-Kontext: Vorstellung des Palace-Frameworks für TPU Sparse Cores als Ergänzung zu Niedrigebenen-XLA-Kompilaten innerhalb des JAX-Ökosystems.
• 1:58 Hardware-Architektur des Sparse Cores: Abgrenzung zum Tensor Core (Matrix-Engines); der Sparse Core verarbeitet nicht-lineare Vektorarithmetik, unregelmäßige Workloads und wahlfreie Speicherzugriffe parallel zum HBM.
• 2:40 Skalar- und Vektor-Subcores: Differenzierung zwischen dem Skalar-Subcore (Instruktionssteuerung, dynamische Indizierung, Prefetching) und 16 Vektor-Subcores (SIMD-Verarbeitung gleicher Befehle auf Vektoren der Länge N).
• 4:50 Speicherhierarchie & DMA-Transfers: Lokale Speicherbereiche (SM für Skalar, VM für Vektoren) ergänzt durch einen Shared-VM-Bereich; DMA-Transfers ermöglichen Ausführungen zwischen HBM, lokalen Speichern und Remote-Chips (über ICI).
• 7:38 Spezialisierte Operationen (Gather & Scatter): Implementierung hardwaregestützter wahlfreier Lookups (Gather/Scatter) auf VM- und HBM-Ebene für Anwendungsfälle wie Mixture-of-Experts-Routing, Embedding-Lookups und Top-K-Sampling.
• 11:03 Palace-Programmiermodell & Mesh-Abstraktion: Nutzung von pltpu.get_tpu_info zur Hardware-Abfrage und Abbildung von Ausführungsnetzen (Grid/Mesh), bei denen Kernel-Code über pl.kernel pro Core orchestriert wird.
• 15:13 Vektor-Slicing & Pipeline-Automatisierung: Anforderungen an exaktes Vektor-Slicing auf Vektor-Subcores sowie Nutzung von emit_pipeline zur Entkopplung von Rechenoperationen und Speicherlatenzen.
• 18:04 Verschachteltes Pipelining für HBM-Gather: Einsatz verschachtelter Pipelines zur Bewältigung mehrstufiger Datenzugriffe (Index-Load aus HBM in VM, anschließendes Gather und Rechenoperation).
• 21:12 MPMD-Kernel & Semaphor-Synchronisation: Parallele Ausführung unterschiedlicher Programme auf Skalar- und Vektor-Subcores mittels extern zugewiesener Semaphoren zur Orchestrierung komplexer Kollektive.
• 24:29 Q&A zu Atomizität & Inter-Core-Kommunikation: Klarstellung zu atomaren Updates bei Scatter-Operationen (add_update_scatter) sowie der theoretischen DMA-Kopplung zwischen Tensor-Core-VEM und Sparse-Core-VM.

Glossar

• TPU (Tensor Processing Unit): Eine von Google entwickelte anwendungsspezifische integrierte Schaltung (ASIC), die speziell für die Beschleunigung von Machine-Learning-Workloads optimiert ist.
• Tensor Core: Die Hauptrecheneinheit einer TPU, spezialisiert auf hochparallele Matrixmultiplikationen (z. B. für Dense Layers in neuronalen Netzen).
• Sparse Core: Eine spezialisierte Sub-Einheit der TPU, optimiert für unregelmäßige Datenstrukturen, spärliche Matrizen (Sparse Data), wahlfreie Speicherzugriffe und Vektoroperationen.
• Subcore (Skalar / Vektor): Untereinheiten des Sparse Cores. Der Skalar-Subcore steuert Kontrollflüsse und Indizes; die Vektor-Subcores führen eigentliche Berechnungen parallel aus.
• SIMD (Single Instruction, Multiple Data): Ein Rechnerarchitektur-Prinzip, bei dem ein einziger Befehl gleichzeitig auf mehrere Datenelemente angewendet wird.
• HBM (High Bandwidth Memory): Hochgeschwindigkeits-Arbeitsspeicher mit sehr hoher Bandbreite, der direkt mit dem TPU-Chip gekoppelt ist.
• VM / SM (Vector Memory / Scalar Memory): Fast, On-Chip liegende On-Die-Speicherbausteine innerhalb des Sparse Cores für Vektor- bzw. Skalar-Daten.
• DMA (Direct Memory Access): Ein Verfahren, das es Hardware-Subsystemen ermöglicht, direkt auf den Hauptspeicher zuzugreifen, ohne die zentrale Recheneinheit stark zu belasten.
• Gather / Scatter: Spezielle Speicherzugriffsoperationen. Gather sammelt Daten aus nicht zusammenhängenden Speicheradressen basierend auf einem Index-Array; Scatter schreibt Daten an unzusammenhängende Speicheradressen.
• MPMD (Multiple Program, Multiple Data): Ein Parallelausführungsmodell, bei dem verschiedene Prozessoren oder Subcores zeitgleich unterschiedliche Programme auf unterschiedlichen Datenbereichen ausführen.
• Semaphor: Ein Synchronisationskonstrukt in der Software- und Hardwareprogrammierung zur Steuerung des Zugriffs auf gemeinsame Ressourcen oder zur Abstimmung paralleler Threads.
• ICI (Inter-Connect Interconnect): Googles proprietäres, extrem schnelles Inter-Chip-Netzwerk zur direkten Kopplung mehrerer TPUs innerhalb eines Pods.
• JAX / Palace: JAX ist ein Python-Framework für hochleistungsorientiertes numerisches Rechnen und Machine Learning. Palace ist eine Programmiersprachen-Erweiterung/Kernel-API innerhalb von JAX zur direkten Programmierung von TPU-Hardwarearchitekturen.

Abstract

Diese Präsentation stellt die Werkzeuge und Methoden zur Profilierung und Fehlerbehebung von JAX- und OpenXLA-Workloads mittels XPro auf Google TPUs und NVIDIA-GPUs vor.

XPro bietet spezialisierte Ansichten zur Leistungsevaluierung (Trace Viewer, Graph Viewer, Memory Viewer und Roofline-Analysen). Während Standard-HLO-Operationen deterministisch analysiert werden können, stellen Custom Calls (z. B. Palace- oder Triton-Kernels) aufgrund ihrer dynamischen Natur eine Blackbox dar. Zur Analyse dieser Kernels nutzt XPro feingranulares Hardware-Sampling: Auf TPUs (z. B. Ironwood, TPU v6E/v7) können aus einem Fundus von bis zu 60.000 Performance-Countern mehr als 1.100 Zähler pro Durchlauf bei Abtastraten von bis zu 833 MHz erfasst werden. Dies ermöglicht die detaillierte Befehls- und Pipeline-Analyse (z. B. Identifikation von MXU-Latenzen).

In Zusammenarbeit mit NVIDIA wurde XPro um GPU-PM-Sampling (zeitbasierte Erfassung von NVLink-, PCIe- und SM-Aktivitäten mit 10–100 kHz auf Hopper- und Blackwell-Architekturen) sowie Multi-Pass Range Profiling über den Multi-Host HLO Runner erweitert. Letzteres ermöglicht das Replay von Kernels zur exakten Messung von über 6.000 Hardware-Metriken (z. B. auf dem GB100-Chip). Neue Entwicklungen in NVIDIAs CUPTI-Schnittstelle reduzieren Zudem den CPU-Profiling-Overhead durch schemabasierte Aktivitätsdatensätze und verhindern Konflikte durch Multi-Subscriber-Unterstützung.

Wichtigste Punkte & Zeitstempel

• 0:03 Überblick und Agenda: Prasanna Ranganathan leitet die Session zur Profilierung von JAX- und OpenXLA-Workloads ein. Die Agenda umfasst Custom-Call-Debugging, TPU/GPU-Performance-Counter und Profiling-Features in Kooperation mit NVIDIA.
• 1:48 Architektur von XPro: XPro dient als herstellerübergreifender Profiler für TPUs, GPUs und PJRT-kompatible Beschleuniger. Er enthält Ansichten wie den Trace Viewer (Gantt-Diagramm), Graph Viewer (HLO-Graphen), Memory Viewer (Puffer-Zuordnungen/Fragmentierung) und Roofline-Analysen.
• 4:02 Erfassungsmethoden: Profile können interaktiv über das XPro-UI (Lokal oder Remote-Host) oder programmatisch via jax.profiler und torch.profiler gestartet werden, um Metriken wie Step-Time, Token-Durchsatz und Flop-Auslastung zu erfassen.
• 6:45 Telemetrie für Leistung und Thermik: Für TPU-Architekturen (z. B. Trillium, Ironwood) integriert XPro Echtzeitmessungen von Stromverbrauch, Modultemperaturen und thermisch bedingtem Throttling in einer einheitlichen Ansicht.
• 9:21 Herausforderungen bei Custom Calls: Palace- und Triton-Kernels werden direkt in Maschinencode abgesenkt und umgehen HLO-Optimierungen. Durch dynamische Kontrollflüsse und Formate erfordern sie feingranulare Befehlstraces anstelle statischer Analysen.
• 13:40 TPU Hardware-Performance-Counter: TPU-Generationen (v6E, v7) bieten bis zu 60.000 Hardware-Zähler. Bei Abtastraten von bis zu 833 MHz (1 µs) können über 1.100 Zähler pro Durchlauf erfasst werden, um MXU-Auslastungslücken (z. B. Behebung von Stalls durch Erhöhung der Puffer-Tiefe von 1 auf 3) zu analysieren.
• 22:26 GPU PM-Sampling (NVIDIA): Chris Ashton demonstriert die Integration von NVIDIA CUPTI in XPro. Das PM-Sampling erfasst zeitbasierte Hardware-Metriken (SM-Aktivität, NVLink, PCIe) auf GPU-Architekturen wie Hopper und Blackwell (GB100) bei Frequenzen von 10–100 kHz.
• 29:52 Range Profiling via HLO Runner: Um tiefere Hardware-Metriken (über 6.000 auf GB100) ohne Zeitabtastungsfehler zu messen, führt der Multi-Host HLO Runner Multi-Pass-Replays (z. B. 6–7 Durchläufe) aus und übergibt präzise Befehls- und Cache-Zähler an XPro.
• 38:05 CUPTI-Optimierungen und Multi-Subscriber: Neue CUPTI-Funktionen führen schemabasierte Aktivitätsdatensätze zur Reduzierung des Host-CPU-Overheads ein und ermöglichen über Multi-Subscriber-Support den parallelen Zugriff von XPro, JAX PGLE und Nsight Systems ohne Kollisionen.

Glossar

• XPro: Ein von Google entwickeltes, plattformübergreifendes Profiling-Werkzeug für Machine-Learning-Workloads zur Performance-, Speicher- und Hardwareanalyse.
• JAX: Ein Python-Framework für hochleistungsfähiges numerisches Rechnen und Machine Learning, das automatische Differenzierung und XLA-Kompilierung kombiniert.
• OpenXLA / HLO: OpenXLA ist ein domänenspezifischer Compiler für ML-Frameworks; HLO (High-Level Optimizer) ist die Zwischendarstellung (IR) des Compilers zur Graph-Optimierung.
• Custom Call (Palace / Triton): Benutzerdefinierte Code-Kerne, die direkt in maschinennahen Code kompiliert werden und Standard-Compiler-Optimierungen umgehen.
• PJRT (Plugin Joint Runtime): Eine einheitliche Laufzeitschnittstelle in OpenXLA, die es verschiedenen Hardware-Beschleunigern ermöglicht, sich transparent in ML-Frameworks einzubinden.
• Roofline-Modell: Ein anschauliches Leistungsmodell, das aufzeigt, ob die Performance einer Berechnung durch die Speicherbandbreite (memory-bound) oder die Rechenleistung (compute-bound) begrenzt wird.
• MXU (Matrix Multiply Unit): Die spezialisierte Hardware-Recheneinheit auf Google TPUs, die für die Ausführung von Matrixmultiplikationen verantwortlich ist.
• CUPTI (CUDA Profiling Tools Interface): Eine Software-Schnittstelle von NVIDIA, die externen Profiling-Werkzeugen Zugriff auf interne Leistungsinformationen und Hardware-Zähler von NVIDIA-GPUs gewährt.
• PM-Sampling (Performance Monitor Sampling): Das periodische Abtasten von Hardware-Zählern (z. B. Auslastung, Taktrate) während der Laufzeit eines Programms zur Erstellung eines Zeitdiagramms.
• Range Profiling: Ein Testverfahren, bei dem ein bestimmter Codeabschnitt mehrfach durchlaufen wird, um tausende detaillierte Hardware-Zähler ohne Genauigkeitsverlust zu erfassen.
• Multi-Subscriber Support: Eine Eigenschaft von Profiling-Schnittstellen, die es mehreren Diagnose-Werkzeugen gleichzeitig erlaubt, Daten vom selben Treiber/Chip abzugreifen, ohne sich gegenseitig zu blockieren.

Abstract

In diesem Vortrag wird die Weiterentwicklung von Pallas, der Low-Level-Kernel-Programmierschnittstelle für TPUs innerhalb des JAX-Ökosystems, detailliert vorgestellt. Der Fokus liegt auf der spezifischen TPU-Hardwarearchitektur (bestehend aus High Bandwidth Memory / HBM, Vector Memory / VM und Scalar Memory / SM) sowie dem Übergang von der bisherigen Schnittstelle pallas_call zur neuen API pl.kernel. Durch die Integration veränderlicher Speicherreferenzen (jax.Ref) ermöglicht pl.kernel speichereffiziente In-Place-Mutationen, vereinfachtes Alias-Management und eine entkoppelte Steuerung von Speicher-Pipelines über den Helper emit_pipeline. Darüber hinaus wird die Unterstützung heterogener Multi-Core-Systeme behandelt: Entwickler können neben den Tensor Cores nun auch Sparse Cores (z. B. auf TPU v5p, v6e, v7) direkt programmieren. Praktische Demonstrationen belegen die Implementierung gekachelter Matrixmultiplikations-Kernels sowie das Überlappen von Vektor-/Tensor-Berechnungen mit asynchroner RDMA-Kollektivkommunikation (All-Gather) auf den Sparse Cores.

Wichtigste Punkte & Zeitstempel

• 0:04 Einführung in Pallas: Überblick über Pallas als JAX-Erweiterung für die hardwarenahe Kernel-Programmierung auf TPU-Beschleunigern.
• 2:45 TPU-Speicherarchitektur: Analyse des dreistufigen Speichermodells aus High Bandwidth Memory (HBM), Vector Memory (VM) für arithmetische Operationen und Scalar Memory (SM) für Kontrollfluss und Adressierungslogik.
• 3:52 Legacy-API pallas_call: Funktionsweise der ursprünglichen API, die automatisierte Speicher-Pipelines über statische Block-Spezifikationen (block_spec) zur Datenübertragung zwischen HBM und VM nutzt.
• 6:35 Dynamische Steuerung via Scalar Memory: Einsatz von prefetch_scalar_grid zum Laden dynamischer Metadaten und Sparsity-Masken in das SM für laufzeitabhängiges Slicing.
• 9:21 Direkter HBM-Zugriff für Kollektive: Deaktivierung automatischer VM-Pipelines zur Arbeit mit direkten HBM-Referenzen für Inter-Chip-Kommunikationsprimitive.
• 11:08 Die neue pl.kernel-API: Vorstellung der modularisierten API-Schnittstelle, die Kernel-Logik von automatischen Pipelines entkoppelt und direkte Hardware-Kontrolle bietet.
• 12:48 Mutable Referenzen (jax.Ref): Vermeidung von Speicherallokations-Overhead und Pufferkopien durch veränderliche Referenzen und direkte In-Place-Mutationen.
• 15:57 Explizites Pipelining via emit_pipeline: Manuelle Steuerung von überlappten Datenübertragungen zwischen HBM und VM direkt innerhalb des Kernel-Bodys.
• 19:00 Live-Demonstration eines Matmul-Kernels: Schrittweise Implementierung eines gekachelten Matrixmultiplikations-Kernels unter Verwendung von Float32-Akkumulation und Grid-Tiling.
• 31:19 Chunked Matmul mit In-Place-Updates: Umstrukturierung des Matmul-Kernels auf Referenzbasis zur blockweisen In-Place-Verarbeitung ohne Speicher-Overhead.
• 40:12 Multi-Core & Sparse Core Support: Ansteuerung heterogener TPU-Architekturen (v5p/v7) mit dedizierter Programmierung von Tensor Cores und Sparse Cores.
• 42:44 Core-Mesh-Kommunikation: Übertragung des SPMD-Geräte-Mesh-Konzepts auf Inter-Core-Ebene zur feingranularen Signalisierung über Semaphore.
• 46:44 Overlapping von Compute und Kollektiven: Demonstration eines asynchronen Sparse-Core All-Gather-Kernels zur parallelen Datenübertragung während laufender Tensor-Core-Berechnungen.

Glossar

• Pallas: Eine Low-Level-Programmierschnittstelle innerhalb von JAX zur Erstellung benutzerdefinierter, hochoptimierter Hardware-Kernels für TPUs und GPUs.
• TPU (Tensor Processing Unit): Von Google entwickelter anwendungsspezifischer Integrierter Schaltkreis (ASIC) zur Beschleunigung von Machine-Learning-Workloads.
• HBM (High Bandwidth Memory): Der primäre, kapazitätsstarke Hauptspeicher einer TPU, in dem Tensoren, Modellparameter und Aktivierungen abgelegt sind.
• VM (Vector Memory): Schneller, lokaler On-Chip-Speicher der TPU-Vektoreinheiten, der als direkte Datenbasis für mathematische Operationen dient.
• SM (Scalar Memory): Ein lokaler On-Chip-Speicherbereich der TPU zur Aufbewahrung von Skalaren, Adresszeigern und Kontrollflussstrukturen.
• pallas_call: Die ursprüngliche Einstiegsfunktion in Pallas, welche Speicher-Pipelines automatisch auf Basis von Blockspezifikationen verwaltet.
• pl.kernel: Die neue, erweiterte Einstiegsfunktion in Pallas zur direkten Steuerung von Speicherreferenzen, Multi-Core-Topologien und manuellen Pipelines.
• jax.Ref: Eine veränderbare (mutable) Speicherreferenz in JAX, die speichereffiziente In-Place-Operationen ohne Duplizierung von Puffern erlaubt.
• emit_pipeline: Eine Helper-Funktion in Pallas zur expliziten Generierung von speicherüberlappenden Prefetch- und Store-Pipelines zwischen HBM und VM.
• Tensor Core (TC): Die Haupt-Recheneinheit der TPU, optimiert für dichte Matrixmultiplikationen und Tensor-Arithmetik.
• Sparse Core (SC): Eine spezialisierte TPU-Hardwarekomponente zur beschleunigten Verarbeitung von dünnbesetzten (sparsamen) Datenstrukturen, Embedding-Lookups und Kollektiv-Kommunikation.
• RDMA (Remote Direct Memory Access): Ein Kommunikationsprotokoll für den direkten Speicherzugriff zwischen mehreren Beschleuniger-Chips ohne Einbindung der Host-CPU.
• All-Gather: Ein kollektives Kommunikationsmuster, bei dem jeder beteiligte Knoten seine lokalen Daten an alle anderen Knoten des Netzwerks sendet.

Zusammenfassung

Dieses Transkript dokumentiert eine Ausstrahlung des Programms „Sterre en planete“ von Radio Sonder Grense (RSG) vom 21. Dezember 2025. Die Sendung präsentiert traditionelle Sternengeschichten der San (Boesmen), die ursprünglich zwischen 1919 und 1921 von dem Landmesser und Sprachaktivisten Gideon Retief von Wielligh (1895–1932) in einer viermedialen Sammlung („Boesman stories“) publiziert wurden. Im Fokus steht die mythologische Erzählung „Die aandster is die oog van aandskemering“, welche die kosmologischen Ursprünge von Himmelskörpern wie dem Abendstern sowie die Verhaltensweisen bestimmter Tiere (wie des Uhus und des Stachelschweins) durch eine dramatische Familiengeschichte der Sonnensöhne erklärt.

Wichtigste Punkte & Zeitstempel

• 0:00 Historischer und Rundfunk-Kontext: Ausstrahlung der San-Sternengeschichten im RSG-Programm „Sterre en planete“ am 21. Dezember 2025 zur Feier des 100-jährigen Jubiläums der Afrikaans-Sprachbewegung, basierend auf den Aufzeichnungen von Gideon Retief von Wielligh.
• 1:00 Leben und Werk von G.R. von Wielligh: Der 1895 in Paarl geborene und 1932 verstorbene von Wielligh begleitete als Kind seinen Vater auf Transportreisen nach Namakoland, Boesmanland und Hantam, erlernte die Sprache der San und veröffentlichte 1919 bis 1921 vier Bände mit rund 30 Geschichten, von denen etwa 10 % einen astronomischen Bezug aufweisen.
• 2:17 Mythologischer Ursprung von Sonne und Söhnen: Die Sonne war ursprünglich ein Mensch der alten Generation mit strahlendem Gesicht; nach seiner Tötung durch das Volk wurde sein Kopf in den Himmel geworfen. Seine Söhne waren Daraat (der morgens die Sterne verlöschte und den blauen Himmel aufhängte) und Aandskemmering (der abends den Himmel senkte und die Sterne enthüllte).
• 4:02 Intrige und Verwandlung in der Tiermythologie: Wolfs Tochter stiehlt mit Hilfe von Zauberwurzeln das Kind von Aandskemmerings Frau; die vergiftete Frau verliert ihr Gedächtnis, verfällt in Wahn und verwandelt sich in einen Uhu (uil), der nachts auf Dächern nach seinem verlorenen Kind ruft.
• 6:21 Tod von Aandskemmering und Entstehung des Abendsterns: Nachdem ein Hammelkopf den Verrat aufdeckt, wird Aandskemmering von Wolf von einer Klippe gestürzt; sein eines Auge wird zum hellen Abendstern (Aandster), während das trauernde Kind sich in ein nächtliches Stachelschwein (ystervark) verwandelt.

Glossar

• San / Boesman: Indigene Volksgruppen des südlichen Afrikas. Im Text wird angemerkt, dass „Boesman“ von traditionellen Gruppen als die akzeptierte Eigenbezeichnung angesehen wird.
• RSG Sterre en planete: Eine populäre Wissenschafts- und Astronomiesendung des südafrikanischen öffentlich-rechtlichen Rundfunksenders Radio Sonder Grense (RSG).
• Gideon Retief von Wielligh (1895–1932): Ein südafrikanischer Landmesser, Autor und Pionier der Afrikaans-Sprache, der maßgeblich zur Bewahrung der mündlichen Überlieferungen der San beitrug.
• Aandskemmering (Abenddämmerung): In der Mythologie des Transkripts der Name des Sohnes der Sonne, dessen leuchtendes Auge nach seinem Sturz den Abendstern (Aandster) bildet.
• Hammerkop (Scopus umbretta): Ein afrikanischer Vogel, der in der Mythologie als allwissender Beobachter und Aufklärer fungiert, da er Ereignisse wie in einem Spiegel im Wasser erkennen kann.
• Ystervark: Das südafrikanische Stachelschwein, das in der Erzählung als Metamorphose des trauernden Kindes dargestellt wird, welches nachts unter Felsvorsprüngen nach Nahrung scharrt.

Abstract

Dieses Transkript präsentiert eine Episode der RSG-Sendung „Sterre en planete“ vom 21. Dezember 2025, in der traditionelle Sternenmythen der indigenen Urbevölkerung des südlichen Afrikas vorgestellt werden. Die Erzählungen basieren auf dem historischen Werk des Vermessungsingenieurs Gideon Retief von Willig (1895–1932), der diese Mythen zwischen 1919 und 1921 in vier Bänden publizierte, um die mündliche Folklore und die damals junge afrikaanse Schriftsprache zu bewahren. Im Zentrum steht der Mythos „Die Môrester ist das Herz von Daraad“, eine detailreiche mündliche Überlieferung über Eifersucht, Metamorphosen in Raubtiere (Löwin und Karakal) und die ätiologische Erklärung astronomischer Phänomene sowie anatomischer Merkmale von Tieren.

Wichtigste Punkte & Zeitstempel

• 0:00 Historischer Ursprung: Vorstellung der afrikaanssprachigen Adaption indigener Sternengeschichten, aufgezeichnet von Gideon Retief von Willig zwischen 1919 und 1921, ausgestrahlt im RSG-Programm „Sterre en planete“ am 21. Dezember 2025 im Rahmen des 100-jährigen Jubiläums der Afrikaans-Sprachbewegungen.
• 0:31 Ethnologische Terminologie: Erörterung der Bezeichnung „Boesman“ als bevorzugter Begriff traditioneller Gruppen im Gegensatz zu der von manchen als pejorativ empfundenen Bezeichnung „San“, sowie von Willigs Kindheitsprägung in Namaqualand und Boesmanland.
• 1:57 Methodischer Rahmen: Die Publikation umfasst rund 30 Geschichten in vier Bänden, von denen etwa 10 % einen astronomischen Bezug aufweisen; visuell illustriert durch Felszeichnungsmotive von Carol Botha und Martin Kenny.
• 2:17 Die Zerstörung der Harmonie: Beginn der Erzählung „Die Môrester ist das Herz von Daraad“, in der die eifersüchtige Tochter des Wolfes das Essen von Daraads Ehefrau mit magischem Schweiss kontaminiert, um sie in eine Raubkatze zu verwandeln.
• 3:35 Ätiologie des Wolfsangs: Die verkleidete Wölfin wird beinahe von Daraad getötet; bei ihrer Flucht verbrennt sie sich eine Hinterpfote, was den ungleichmäßigen, hinkenden Gang heutiger Wölfe und Schakale erklärt.
• 4:21 Metamorphose zum Karakal: Daraads echte Frau wird aus ihrer Löwengestalt befreit, indem ihr das Fell abgezogen wird; sie verwandelt sich in einen Karakal (rooi kat), der seine charakteristischen Ohrhaare behält.
• 4:45 Kosmologische Nachkommenschaft: Daraads Kinder und Nachkommenschaft steigen nach ihrem Tod in den Nachthimmel auf und bilden als helle Himmelskörper die Verfolgungsschar ihrer Eltern.
• 5:02 Ursprung des Morgensterns: Ein rachsüchtiger Wolf tötet Daraad; aus dessen geöffnetem Körper entflieht das strahlende Herz und steigt als der große Morgenstern (môrester) an den Firmament auf, wo er die Familie an der Sonnenseite vereint.

Glossar

• Sterre en planete: Eine populärwissenschaftliche Wissenschafts- und Astronomiesendung des südafrikanischen öffentlich-rechtlichen Hörfunksenders RSG (Radio Sonder Grense).
• Boesman / San: Historische und ethnologische Bezeichnungen für die indigenen Jäger- und Sammlerkulturen (Khoisan) im südlichen Afrika; der Text weist auf die präferierte Selbstbezeichnung als „Boesman“ in bestimmten traditionellen Gemeinschaften hin.
• Rooi kat: Die afrikaanse Bezeichnung für den Karakal (Caracal caracal), eine kleine bis mittelgroße Raubkatzenart Afrikas mit markanten, schwarz behaarten Pinselohren.
• Karos: Ein traditioneller Umhang oder eine Decke aus gegerbten Tierfällen, die von den indigenen Völkern des südlichen Afrika als Kleidung genutzt wurde.
• Assegaai: Ein traditioneller afrikanischer Wurfspeer mit eiserner Spitze und Holzschaft, der zur Jagd und Verteidigung dient.
• Ätiologie: Die mythologische oder volkskundliche Erklärung für den Ursprung von Naturphänomenen, geografischen Strukturen oder anatomischen Merkmalen von Tieren.