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.