In diesem Video präsentieren Liz Lee und Ju von AMD die Integration von FlyDSL als neues Backend in TorchInductor zur Optimierung moderner KI-Arbeitslasten. Da kein einzelnes Kernel-Backend über alle Operatoren und Architekturen hinweg optimal ist, verbindet FlyDSL – ein Python-natives, auf maschinellem Lernen basierendes GPU-Kernel-System von AMD – sich mit der Inductor-Pipeline. Dies umfasst die Absenkung (Lowering), Kompilierung, Caching, Benchmarking und das Autotuning, während die gewohnte PyTorch- und Triton-Erfahrung erhalten bleibt. Die Integration bringt hochoptimierte Kernels wie GEMM, Group-GEMM und FlexAttention in PyTorch. Auf AMD MI355-Hardware liefern BF16-GEMMs einen geometrischen Geschwindigkeitszuwachs von 619 % gegenüber Triton; FlyDSL FlexAttention erreicht einen 5,8-fachen Speedup bei MiniMax M3-Dekodierformen sowie nahezu das Doppelte bei DeepSeek V2- und V3-Präfill-Größen.
Wichtigste Punkte & Zeitstempel
0:00 KI-Workload-Herausforderungen: Liz Lee und Ju von AMD betonen, dass kein einzelnes Kernel-Backend über alle Operatoren und Architekturen hinweg optimal für moderne KI-Arbeitslasten ist.
0:18 FlyDSL-Backend-Integration: Einführung von FlyDSL (AMD's Python-nativer, ML-basierter GPU-Kernel-DSL) als neues Backend innerhalb von TorchInductor.
0:32 Pipeline-Kopplung: Nahtlose Verbindung mit den Kompilierungs-, Caching-, Benchmarking- und Autotuning-Pipelines von Inductor unter Beibehaltung der vertrauten PyTorch- und Triton-Programmiererfahrung.
0:44 Kernel-Erweiterungen: Bereitstellung optimierter FlyDSL-Kernels für GEMM, Group-GEMM und FlexAttention in PyTorch.
0:53 Leistungsdaten (MI355): BF16-GEMMs erzielen einen geometrischen Mittelwert von 619 % Beschleunigung gegenüber Triton; FlyDSL FlexAttention erreicht einen 5,8-fachen Speedup bei MiniMax M3-Dekodierformen und einen fast 2-fachen Speedup bei DeepSeek V2- und V3-Präfill-Größen.
Glossar
TorchInductor: Ein tiefes Compiler-Backend- und Codegenerierungs-Framework für PyTorch, das mathematische Operationen optimiert und für GPUs und CPUs kompiliert.
FlyDSL: Ein Python-natives, auf maschinellem Lernen basierendes Domain-Specific Language (DSL)- und Kernel-Entwicklungssystem von AMD zur Erstellung und Optimierung von GPU-Kernels.
Triton: Eine von OpenAI entwickelte Programmiersprache und ein Kompilierer für GPUs, die das Schreiben händisch optimierter CUDA-ähnlicher Kernel in Python erleichtern.
GEMM (General Matrix Multiply): Eine grundlegende mathematische Operation (allgemeine Matrixmultiplikation), die das rechnerische Fundament von neuronalen Netzwerken und Deep-Learning-Workloads bildet.
Group-GEMM: Eine Variante der Matrixmultiplikation, bei der mehrere unabhängige Matrixmultiplikationen unterschiedlicher Größen in einem einzigen Kernel-Aufruf gebündelt und ausgeführt werden.
FlexAttention: Ein flexibler PyTorch-Operator zur Implementierung optimierter Attention-Mechanismen (ähnlich wie FlashAttention) mit anpassbarer Maskierung und Logik.
MI355: Eine Hochleistungs-Beschleuniger-GPU von AMD, die für KI-Training und -Inferenz in Rechenzentren konzipiert ist.
BF16 (Bfloat16): Ein 16-Bit-Gleitkommaformat, das im Deep Learning weit verbreitet ist, um den Speicherbedarf im Vergleich zu FP32 zu halbieren, bei nahezu identischem numerischen Dynamikbereich.
Präfill (Prefill): Die erste Phase bei der Verarbeitung von Large Language Models (LLMs), in der die gesamte Eingabeaufforderung (Prompt) parallel verarbeitet wird, um den Key-Value-Cache aufzubauen.
Dekodierung (Decode Shapes): Die generative Phase bei LLMs, in der Tokens sequenziell (eines nach dem anderen) erzeugt werden.