Dieser Vortrag, präsentiert von Arthi aus dem JAX-Team bei AMD, beschreibt die Integration von FlyDSL-Kerneln in JAX für AMD-GPUs. FlyDSL ist eine Python-basierte Domain-Specific Language (DSL), die eine layout-erste, feingranulare Steuerung von GPU-Berechnungen ermöglicht, um leistungskritische Workloads wie Mixture-of-Experts (MoE) GEMMs und Fused Attention zu optimieren. Benchmarks des Mixtral 8x7B-Modells in FP8 zeigen im Vergleich zum Standard-XLA-Pfad eine 2,5-fache Leistungssteigerung im Prä-Fill und eine 3,5-fache im Dekodieren, wobei der Durchsatz von 14 auf 50 Token pro Sekunde steigt.
Wichtigste Punkte & Zeitstempel
0:00 Motivation für FlyDSL: Während XLA Standard-Operationen und einfache Fusionen automatisiert, erfordern leistungskritische Workloads wie Fused Attention oder Mixture-of-Experts (MoE) GEMMs eine explizite Hardware-Steuerung von Layouts, Tiling und Speicher-Pipelines.
0:18 FlyDSL-Architektur: FlyDSL bietet ein layout-erstes Python-Programmiermodell ähnlich wie Cute, das vollständige Kontrolle über die Rechenhierarchie von Block-, Warp- und Thread-Ebene bis zu den MFMA-Atomen (Matrix-Hardware-Einheiten) gestattet.
0:38 JAX-Integration: Ein Brückenlayer konvertiert JAX-Tensoren über DLPAC-Metadaten (Form, Strides, Datentyp) in ein von FlyDSL lesbares Speicherformat, da FlyDSL nativ primär PyTorch unterstützt.
0:41 Kompilierungs-Pipeline: JAX-Code wird getraced, über FlyDSL in LLVM-ähnliche Dialekte (fl) und amdogdl übersetzt, zu hsa co kompiliert und via C-Trampolines im Runtime-Verhalten pro GPU-Sharding ausgeführt.
0:56 Ziel-Workload (MoE GEMMs): Bewältigung irregulärer Speicherzugriffe (Gather/Scatter) und dynamischer Token-Verteilungen in Mixture-of-Experts-Modellen durch gruppierte GEMMs statt ineffizienter maskierter dichter GEMMs.
0:7:31 Standalone-Validierung: Tests des isolierten FlyDSL-Kernels gegen AMDs CK-Kernel (Composable Kernel) beim Mixtral 8x7B-Modell in FP8 zeigen volle Leistungs-Parität und stellenweise Überlegenheit.
0:8:11 End-to-End-Inferenz: Die Integration in das Max-Inferenz-Framework für Mixtral 8x7B liefert eine 2,5-fache Beschleunigung im Prä-Fill und 3,5-fache im Dekodieren; der Inferenz-Durchsatz erhöht sich von 14 auf 50 Token/Sekunde.
0:10:07 Zukunftsaussichten: Die Roadmap umfasst die Unterstützung von Rückwärtspass-Kerneln (Training), dynamischen Skalaren zur Laufzeit sowie die Bereitstellung vorgefertigter Python-Wheels.
0:11:31 Fusionslogik (Q&A): Kernel-Fusionen müssen vom Entwickler explizit im FlyDSL-Code authorisiert werden und sind keine automatische Optimierung durch den Compiler.
Glossar
JAX: Ein von Google entwickeltes Python-Framework für hochperformantes numerisches Rechnen und maschinelles Lernen mit automatischer Differenzierung.
XLA (Accelerated Linear Algebra): Ein domainspezifischer Compiler für lineare Algebra, der Rechengraphen in JAX analysiert und optimierten Maschinencode generiert.
FlyDSL: Eine Python-basierte Domain-Specific Language (DSL) für AMD-GPUs, die Entwicklern die direkte Kontrolle über Thread-Layouts und Hardware-Einheiten auf niedriger Ebene gibt.
GEMM (General Matrix Multiply): Eine grundlegende mathematische Operation zur Matrixmultiplikation und -addition, die das Fundament neuronaler Netzwerke bildet.
Mixture of Experts (MoE): Eine Transformer-Architektur, bei der eingehende Daten nur an eine dynamische Teilmenge von spezialisierten Teil-Netzwerken ("Experten") geleitet werden, um Rechenleistung zu sparen.
MFMA (Matrix-Fused Multiply-Add): Spezifische Hardware-Recheneinheiten auf AMD-Grafikkarten, die für beschleunigte Matrixmultiplikationen bei Deep-Learning-Workloads optimiert sind.
FP8 (8-bit Floating Point): Ein extrem komprimiertes Gleitkommaformat mit 8 Bit Genauigkeit, das den Speicherbedarf halbiert und die Verarbeitungsgeschwindigkeit in KI-Modellen drastisch erhöht.
DLPAC / DLPack: Ein offener Standard, der den direkten und schnellen Austausch von Tensor-Daten im Arbeitsspeicher zwischen verschiedenen KI-Frameworks ermöglicht.
Composable Kernel (CK): Eine Hochleistungs-Bibliothek von AMD zur Erstellung optimierter GPU-Kernel für Deep-Learning-Operationen.
Triton: Eine alternative Programmiersprache und ein Compiler, die es ermöglichen, CUDA-ähnliche Kernel direkt in Python zu schreiben, ohne C++ oder CUDA beherrschen zu müssen.