PyTorch erhält FlyDSL-Backend für schnellere Operatoren auf AMD-MI350-GPUs

Was unterstützt wird und wie der Zugriff funktioniert

FlyDSL deckt dichtes und gruppiertes GEMM, fused Multi-Head-Attention, MXFP8/MXFP4-skaliertes GEMM, RMSNorm und TopK ab. Laut PyTorch zielt die aktuelle Unterstützung auf GPUs der AMD-MI350-Serie mit der Architektur gfx950.

Wie die einzelnen Operationen bei FlyDSL landen, hängt vom Ausführungsmodus ab:

• Eager-Modus: Geeignete RMSNorm- und TopK-Aufrufe werden an FlyDSL weitergeleitet, wenn die optionale Runtime installiert und aktiviert ist. Andere Aufrufe bleiben bei ATen.
• torch.compile: TorchInductor filtert geeignete Kandidaten für dichtes, gruppiertes und MXFP-GEMM, vergleicht FlyDSL per Benchmark mit ATen und, wo unterstützt, Triton, speichert den Sieger je Workload im Cache und führt den schnellsten gemessenen Kernel aus.
• Attention: F.scaled_dot_product_attention ruft AOTriton auf, das nun anhand seiner Offline-Operator-Tuning-Datenbank für jeden Workload zwischen FlyDSL-, Triton- und AITER-ASM-Kerneln wählt.

FlyDSL-Attention-Kernel werden ab der Serie 0.14b als zusätzliches AOTriton-Backend ausgeliefert. FP32-Eingaben werden von den FlyDSL-Attention-Kerneln nicht unterstützt und laufen weiter über die Triton- oder AITER-ASM-Kernel von AOTriton. Zudem liefert AOTriton FlyDSL-Attention-Kernel für RDNA4 (gfx1201) mit.

Worin sich FlyDSL von Triton unterscheidet

PyTorch beschreibt die beiden Ansätze als komplementär. Das Block-Tensor-Modell von Triton erlaubt es Autoren, Tiles auf Programmebene zu beschreiben, während der Compiler einen Großteil der Zuordnung zu Threads, Waves und Instruktionen übernimmt. FlyDSL macht diese Zuordnungen über Layout-Algebra explizit. Kernel-Autoren können Tiles auf Blöcke, Waves, Threads und Werte aufteilen, Tiled-Copy- und MFMA-Atome auswählen, LDS-Swizzles und Staging definieren sowie Synchronisations- oder Compiler-Scheduling-Grenzen setzen.

FlyDSL drückt diese Steuerungsmöglichkeiten in Python aus und überführt sie über MLIR in Code, wobei die explizite Hardware-Zuordnung erhalten bleibt. Laut dem Beitrag eignet sich das für architekturspezifische Templates, bei denen Datenbewegung und Instruktionszuständigkeit Teil des Algorithmus sind.

Als Beispiel nennt der Beitrag Attention. Die FlyDSL-Attention-Kernel nutzen eine Dual-Wave-Software-Pipeline, bei der die Softmax-VALU-Arbeit einer KV-Tile mit der MFMA-Arbeit einer anderen überlappt. Bei den größten Head-Dimensionen teilen sie außerdem die Head-Dimension auf, zeitlich über LDS und über Waves hinweg für den Ausgabe-Akkumulator. Da beide Schedules im Kernel selbst geschrieben sind, lassen sie sich laut PyTorch ohne Compiler-Änderungen pro Architektur abstimmen.

Gemessene Leistung

Die Benchmarks auf Kernel-Ebene wurden auf GPUs vom Typ AMD Instinct MI355X (gfx950) gemessen. Laut PyTorch konzentrieren sie sich auf die Ausführung im eingeschwungenen Zustand, schließen Kompilierung und anfängliches Autotuning aus und stehen nicht für die End-to-End-Latenz von Modellen. Die MXFP-Messungen beginnen mit bereits quantisierten Operanden.

• Dichtes GEMM: 1,10-fache Beschleunigung im geometrischen Mittel für die BF16-Suite mit 15 Shapes gegenüber der jeweils schnelleren Variante aus ATen und Triton.
• MXFP8 und MXFP4: 1,58-fache bzw. 1,68-fache Beschleunigung im geometrischen Mittel gegenüber ATen über je 17 Shapes pro Format.
• Gruppiertes GEMM: 1,20-faches geometrisches Mittel gegenüber Triton über alle 24 Fälle.
• RMSNorm: das 1,17- bis 3,66-Fache von ATen über 22 Fälle.
• TopK: 4,82-faches geometrisches Mittel gegenüber ATen bei 10 Fällen mit kleinem K, bei deaktivierten deterministischen Algorithmen.

Bei Attention liegt der veröffentlichte gfx950-Durchsatz von AOTriton 0.14.2b je nach Head-Dimension und Causal-Modus zwischen dem 0,92- und 2,34-Fachen von 0.13b im Forward- sowie dem 0,93- und 2,47-Fachen im Backward-Pass. In End-to-End-A/B-Tests mit vLLM erreicht die Beschleunigung ganzer Anfragen je nach Modell und Parallelität das 1,13-Fache bei BF16, das 2,04-Fache bei MXFP8 und das 3,01-Fache bei MXFP4.

Grenzen und Folgen für Anwender

Jede Operation hat Anforderungen an Shape und Alignment. MXFP-skaliertes GEMM akzeptiert FP8-E4M3- oder gepackte FP4-E2M1-Operanden mit einer E8M0-Skala pro 32 Werte entlang der Reduktionsdimension, und der aktuelle Pfad verlangt, dass K ein Vielfaches von 128 ist. Der RMSNorm-Backward-Pass läuft weiterhin über die bestehende PyTorch-Implementierung. Unter torch.compile ist flex_attention eine eigene API, die Inductor in eigene Triton-Kernel kompiliert und die daher AOTriton nicht nutzt.

Für Teams, die PyTorch auf Hardware der MI350-Serie betreiben, ist die Änderung additiv: Unterstützte Aufrufe können FlyDSL-Kernel ohne Codeänderungen nutzen, nicht unterstützte fallen zurück. PyTorch weist darauf hin, dass die Gewinne vor allem dann ins Gewicht fallen, wenn die unterstützten Operationen einen erheblichen Anteil der Modelllaufzeit ausmachen. Das Update AOTriton 0.14.2b bringt die Attention-Kernel in PyTorch.