Dieser Vortrag von Rutaj Vidya (Software Engineer bei AMD) stellt Jaxator vor, ein Integrations-Framework, das hochoptimierte AMD-AITer-Kernel in den JAX/XLA-Softwarestack für die AMD ROCm-Plattform einbettet. AITer ist AMDs KI-Operator-Bibliothek, die optimierte und gefügte Kernel (z. B. Attention, GEMM, MoE) über verschiedene Backends wie Composable Kernel (CK), Triton, HIP und Hand-Assembly bereitstellt. Da AITer historisch eng an PyTorch gebunden war, dient Jaxator als entkoppelte Brücke, die JAX-APIs über XLA Custom Calls anbindet. Jaxator kümmert sich um Lowering-Regeln, Custom VJP für automatische Differentiation, Sharding-Regeln für verteilte Multi-GPU-Systeme und die HIP-Laufzeitintegration. Die Leistungsfähigkeit wird anhand zweier Fallstudien demonstriert: Die Integration von Multi-Head Attention (MHA) erzielte einen medianen Geschwindigkeitsgewinn von 4,39 %, während ein experimentelles End-to-End MXFP4-Trainingsrezept (4-Bit-Fließkomma-Quantisierung) in MaxText auf LLaMA 3.0/3.1 bis zu 30 % Durchsatzsteigerung gegenüber einer FP8-Baseline erreichte. Jaxator wird als Open-Source-Wheel bereitgestellt, um native, hochperformante KI-Trainings- und Inferenz-Pipelines auf ROCm zu ermöglichen.
Wichtigste Punkte & Zeitstempel
00:05 Zielsetzung von Jaxator: Anbindung von AMD AITer-Kerneln an den JAX/XLA-Workload auf ROCm zur Bereitstellung performanter End-to-End-Pipelines inklusive Rückwärtspfad-Unterstützung und verteiltem Training.
02:07 Architektur von AMD AITer: High-Performance-Operator-Bibliothek für ROCm mit optimierten, gefügten Kerneln (Attention, GEMM, MoE, Norms) über Backends wie CK, Triton, HIP und Hand-Assembly, primär mit PyTorch-Abhängigkeiten entwickelt.
03:21 Leistungsgewinne durch AITer: Performance-Steigerungen zwischen 2x und 17x auf Operator-Ebene sowie eine Verdopplung des Durchsatzes bei Realwelt-Modellen wie DeepSeek via SGLang.
04:26 Funktionsweise des Jaxator-Frameworks: Schnittstelle zwischen JAX-APIs und XLA Custom Calls zur Integration hardwarenaher Kernel mit korrekten Lowering-Regeln, Layouts und Gradienten.
06:36 Benutzererfahrung & Integration: Bereitstellung als installierbares Python-Wheel; direkte Nutzung mit JAX-Arrays, Kompatibilität mit JIT-Kompilierung und custom VJP für Differenzierbarkeit.
08:00 Niedrigschwellige Systemarchitektur: XLA behandelt den Kernel als opaque Black-Box; Jaxator liefert explizite Custom Calls, VJP-Regeln für das Training sowie Sharding-Regeln für Multi-GPU-Kommunikation.
10:04 Fallstudie 1: Multi-Head Attention (MHA): Fokussierung auf MHA als zentralen Rechen-Hotspot beim LLM-Training; Erzielung eines medianen Performance-Gewinns von 4,39 % durch entkoppelte C++-APIs.
12:03 Fallstudie 2: MXFP4 End-to-End-Stress-Test: Validierung eines vollständigen 4-Bit-Fließkomma-Trainingsrezepts (MXFP4) in MaxText auf LLaMA 3.0/3.1 inklusive Quantisierung von BF16, Hadamard-Transformation und Stochastic Rounding.
13:40 Kommunikation & Sharding bei MXFP4: Notwendigkeit expliziter Kollektiv-Kommunikationsregeln während des Weight-Gradient-Passes (DB), da die Kontraktionsachse M über mehrere GPUs geshardet wird.
16:32 MXFP4-Performance-Ergebnisse: Erreichen von 20 % Durchsatzsteigerung bei reinen MLP-Schichten und 30 % bei voller Integration (MLP + QKV/O-Projektionen) gegenüber einer FP8-Baseline bei vergleichbarer Konvergenz.
18:09 Erkenntnisse & Zukünftige Entwicklungen: Entkopplung weiterer PyTorch-Abhängigkeiten in reine C++-APIs, Bereitstellung neuer Rezepte, MLPerf-Konvergenzvalidierung und Upstreaming der Quantisierungsmodi in MaxText.
23:37 Q&A - Framework-Entkopplung & Backend-Auswahl: Erläuterung der Entfernung von PyTorch-Abhängigkeiten, Funktionsweise des gefügten Quantisierungs-Kernels und automatische Selektion der optimalen AITer-Backends (z. B. Assembly vs. CK) basierend auf Tensor-Shapes.
Glossar
AITer / AIT: AMD AI Operator Library; eine hochoptimierte Bibliothek von KI-Operatoren für die AMD ROCm-Plattform.
ROCm: Radeon Open Compute platform; AMDs Open-Source-Software-Stack für GPU-Computing und parallele Programmierung.
JAX / XLA: JAX ist ein Framework für transformierbaren Python-/NumPy-Code; XLA (Accelerated Linear Algebra) ist der zugrundeliegende Domänen-spezifische Compiler für Machine Learning.
Custom Call: Eine Schnittstelle in XLA, die es ermöglicht, externe, native C++/GPU-Kernel direkt aufzurufen, ohne dass der XLA-Compiler deren innere Mathematik analysieren muss.
VJP (Vector-Jacobian Product): Ein mathematischer Operator der automatischen Differentiation, der im Rückwärtspass von neuronalen Netzen zur Gradientenberechnung verwendet wird.
MHA (Multi-Head Attention): Ein zentraler Aufmerksamkeitsmechanismus in Transformer-Architekturen, der die parallele Verarbeitung verschiedener Sequenzabschnitte ermöglicht.
MXFP4: Microscaling Floating-Point 4-bit; ein extrem niedrigpräzises 4-Bit-Fließkommaformat mit blockweisem Skalierungsfaktor zur Beschleunigung von KI-Training und Inferenz.
BF16 / FP8: Bfloat16 (16-Bit) und Float8 (8-Bit); numerische Fließkomma-Formate zur Optimierung von Speicherbedarf und Rechenleistung in Deep-Learning-Workloads.
Sharding: Die Aufteilung von Daten oder Tensor-Dimensionen auf mehrere Beschleunigerkarten (GPUs) zur verteilten Berechnung.
CK (Composable Kernel): Eine C++-GPU-Bibliothek von AMD zur Erstellung performanter Rechen-Kernel.
Hadamard-Transformation & Stochastic Rounding: Mathematische Verfahren zur Reduzierung von Quantisierungsfehlern und Ausreißern bei der Umwandlung von Daten in Niedrigpräzisionsformate wie MXFP4.