In dieser Präsentation stellt Arthi vom AMD JAX-Team die Integration von FlyDSL-Kerneln in das JAX-Framework auf AMD-GPUs vor. FlyDSL ist eine Python-basierte Kernel-DSL (Domain-Specific Language), die Entwicklern feingranulare Kontrolle über GPU-Hardwarestrukturen wie Kachelgrößen (Tiling), Speicherlayouts, MFMA-Instruktionen (Matrix Fused Multiply-Add) und Epilog-Fusionen bietet. Dies geht über die automatischen Optimierungsmöglichkeiten des XLA-Compilers hinaus, der bei komplexen Workloads wie Fused Attention oder Mixture-of-Experts (MoE) an seine Grenzen stößt.
Die vorgestellte Architektur nutzt eine Brückenschicht, die JAX-Tensoren via DLPack in memref-Formate konvertiert, XLA Custom Calls registriert und zur Laufzeit einen C-Trampoline-Layer verwendet, um Puffer an GPU-Shards zu übergeben. Am Beispiel eines zwei-stufigen gerouteten GEMM-Kernels für MoE-Modelle (Mixtral 8x7B in FP8) zeigt der Ansatz deutliche Effizienzgewinne: Durch die Fusion von Gather-, Group-GEMM-, SwiGLU- und Scatter-Operationen wird unnötiger Rechenaufwand vermieden. In End-to-End-Inferenz-Testläufen auf MaxText wurde eine 2,5-fache Beschleunigung im Prefill und eine 3,5-fache Beschleunigung im Decode erzielt (Durchsatzsteigerung von 14 auf 50 Tokens/Sekunde). Zukünftige Erweiterungen umfassen die Unterstützung von Rückwärts-Kerneln (Custom VJP) für das Training, dynamische Skalare sowie vorgefertigte Python-Wheels.
Wichtigste Punkte & Zeitstempel
0:05 Motivation und Grenzen von XLA: XLA erzeugt zwar effizienten Code für elementweise und Reduktions-Operationen (z. B. Triton-Softmax), stößt jedoch bei leistungskritischen Workloads (MoE GEMMs, Fused Attention) an Grenzen, die explizite Kontrolle über Tiling, Layouts und MFMA-Instruktions-Scheduling erfordern.
1:38 FlyDSL-Architektur: FlyDSL bietet ein Python-basiertes „Layout-first“-Programmiermodell (ähnlich CuTe) mit Kontrolle über Block-, Wavelet-, Thread- und Atom/MFMA-Ebenen. Der Kompilierungspfad verläuft von Python über den Fly-Dialekt und AMDOGDL bis zum HSA Code Object.
2:41 Kernel-Autorisierung und JAX-Einbindung: Trennung in GPU-Ausführung (@flyc.kernel) und Grid/Block-Launcher (@flyc.jit). Der Wrapper from_jax konvertiert JAX-Tensoren via DLPack in speicherkompatible memref-Strukturen für FlyDSL.
4:07 Parameterisierung der fly_call-Schnittstelle: Bietet Übergabe von Eingaben, Launcher-Kernel, Ausgabespezifikationen (Form/Dtype für XLA-Allokation) und statischen Kompilierzeit-Skalaren (aktuell noch keine dynamischen Laufzeit-Skalare).
5:05 Tracing-, Kompilierungs- und Laufzeitpfad: Bei der Tracing-Phase wird der Kernel in Fly-IR überführt, als HSA-Code-Objekt kompiliert und als XLA Custom Call registriert. Zur Laufzeit entpackt ein C-Trampolin-Layer Puffer, Skalare und Streams je GPU-Shard.
5:56 Optimierung gerouteter MoE-GEMMs: MoE-Workloads beinhalten dynamische Expertengrößen. Der AMD-Standardpfad nutzt rechenintensive maskierte dichte GEMMs; FlyDSL eliminiert Overhead durch die Fusion von Gather-, Group-GEMM-, SwiGLU- und Scatter-Operationen in einen einzigen Kernel.
7:31 Standalone-Benchmarks: Die Validierung des FlyDSL-MoE-Kernels gegen die Composable Kernel (CK)-Baseline (FP8, Mixtral 8x7B) belegt Leistungsparität sowie leichte Vorteile in ausgewählten Testfällen.
8:14 End-to-End-Integration in MaxText: Einbindung des gerouteten GEMMs in das Mixtral-Inferenzmodell (Top-2-Routing, 16 Shards) inklusive einmaligem Load-Time-Preshuffle; erzielt vollständige numerische Parität mit HuggingFace-Referenzgewichten.
9:14 End-to-End-Performance-Ergebnisse: Erreichen einer 2,5-fachen Beschleunigung im Prefill und einer 3,5-fachen Beschleunigung im Decode. Der Decode-Durchsatz steigt von 14 auf 50 Tokens/Sekunde (insbesondere da GFX950/BF16 ohne Group-GEMM-Lowering sonst auf rocBLAS zurückfällt).
10:07 Ausblick und Roadmap: Zukünftige Arbeitsschwerpunkte umfassen Custom VJP für Trainingsunterstützung (Rückwärts-Pass), dynamische Laufzeitskalare sowie die Veröffentlichung vorgefertigter Python-Wheels.
Diskussions-Highlights
11:31 Fusionskontrolle in FlyDSL: Kernel-Fusionen werden manuell vom Entwickler im Python-Code geschrieben und definiert; der FlyDSL-Compiler führt keine automatischen Fusionsoptimierungen durch.
Glossar
FlyDSL: Eine von AMD entwickelte domänenspezifische Sprache (DSL) in Python zur Erstellung hochoptimierter GPU-Kernel für AMD-Hardware.
JAX: Ein von Google entwickeltes Python-Framework für automatisches Differenzieren und hochleistungsfähiges numerisches Rechnen auf GPUs/TPUs.
XLA (Accelerated Linear Algebra): Ein domänenspezifischer Compiler für maschinelles Lernen, der JAX- oder TensorFlow-Graphen in optimierten Maschinencode übersetzt.
MFMA (Matrix Fused Multiply-Add): Hardware-Spezialinstruktionen und Recheneinheiten auf AMD-CDNA-GPUs zur beschleunigten Ausführung von Matrixmultiplikationen.
GEMM (General Matrix Multiply): Die grundlegende mathematische Operation der allgemeinen Matrixmultiplikation in Deep-Learning-Workloads.
MoE (Mixture of Experts): Eine Modellarchitektur im maschinellen Lernen, bei der Eingaben (Tokens) dynamisch an spezialisierte Teilnetzwerke („Experten“) geroutet werden.
DLPack: Ein sprach- und framework-übergreifender In-Memory-Tensor-Transfer-Standard zur speichereffizienten Übergabe von Daten zwischen Frameworks ohne Kopieren.
HSA (Heterogeneous System Architecture) Code Object: Das binäre Ausführungsformat für Kernel, die auf AMD-ROCm-Laufzeitumgebungen ausgeführt werden.
SwiGLU: Eine mathematische Aktivierungsfunktion (Kombination aus Swish und Gated Linear Unit), die häufig in modernen Large Language Models (LLMs) verwendet wird.
Group GEMM: Eine Matrixmultiplikationsvariante, die mehrere unabhängige Matrixoperationen unterschiedlicher Dimensionen effizient in einem einzigen GPU-Aufruf gebündelt ausführt.
CK (Composable Kernel): AMDs C++-basiertes Performance-Framework zur Erstellung hochoptimierter GPGPU-Kernel.
GFX950: Die interne GPU-Architekturbezeichnung für AMD Instinct-Beschleuniger der neuesten Generation.