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.