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.