AWS erklärt zwei Kernel-Optimierungen, die den Durchsatz von SageMaker AI auf H200 steigern
AWS hat zwei Optimierungen auf Kernel-Ebene für Fused-Mixture-of-Experts-Modelle (MoE) auf Amazon SageMaker AI beschrieben. Ein Trick mit Matrixtransposition bringt demnach 2 bis 4 Prozent mehr End-to-End-Ausgabedurchsatz, ein anderes Backend für FP8-Modelle weitere 10 Prozent. Beide Ergebnisse wurden laut AWS auf NVIDIA-H200-GPUs gemessen.
- AWS beschreibt zwei Optimierungen auf Kernel-Ebene für Fused-MoE-Modelle auf Amazon SageMaker AI.
- Swap-AB transponiert die Matrixmultiplikation, um sie an das feste 64-Zeilen-WGMMA-Layout von Hopper anzupassen, und bringt auf H200 2 bis 4 Prozent mehr End-to-End-Ausgabedurchsatz.
- Die Swap-AB-Optimierung wurde in vLLM übernommen und ist ab Version 0.24.0 verfügbar.
- Das Routing FP8-Tensor-quantisierter Modelle auf das Triton-Backend statt auf FLASHINFER_CUTLASS steigerte den Durchsatz auf H200 um 10 Prozent.
- Der Backend-Wechsel ist eine reine Konfigurationsänderung, ohne dass Kernel-Code angepasst wird.
Warum MoE-Kernel so wichtig sind
Mixture-of-Experts ist eine dünn aktivierte Architektur, die die Kapazität großer Sprachmodelle erhöht, ohne dass der Rechenaufwand pro Token steigt. Laut AWS entfällt auf MoE-Operationen ein großer Teil der gesamten Rechenzeit, sodass Ineffizienzen dort Latenz und Kosten pro Anfrage direkt in die Höhe treiben. Kernel-Engineering an diesen Operationen biete deshalb das größte Potenzial für Leistungsgewinne, so das Unternehmen.
Eine zentrale Technik ist Fused MoE. Sie fasst Token-Routing, Datenumverteilung und spezialisierte Netzwerkberechnungen in einer einzigen gebündelten Operation auf der GPU zusammen, statt Routing und Expertenberechnungen als getrennte, nacheinander laufende Schleifen auszuführen. AWS weist darauf hin, dass dies zwar Latenzvorteile bringt, die Implementierung und das Debugging der Kernel aber komplexer macht und eine hohe GPU-Auslastung hardwarespezifische Optimierungen erfordert. Als Hauptbeispiel dient im Beitrag ein repräsentatives MoE-Modell.
Swap-AB: Ein Missverhältnis bei Hopper-Tensor-Cores beheben
Die erste Optimierung zielt auf die Warp Group Matrix Multiply Accumulate (WGMMA) Instruktion der NVIDIA-Architektur Hopper. AWS erklärt, dass die Hardware die Zeilendimension der Matrix (M) auf 64 festlegt, während die Spaltendimension (N) variabel ist. Die maximale Auslastung setzt daher voraus, dass M mindestens 64 beträgt.
Das ist in der Decode-Phase ein Problem, in der das System pro Nutzersequenz jeweils ein Token verarbeitet. Die Zeilenzahl M der Tokens liegt dann oft unter 64 und häufig bei nur 1. Im Standardlayout wird Matrix A zu einem schmalen Vektor, etwa 1 x 4096, und kann die 64 Zeilen nicht füllen. Laut AWS bleibt dadurch ein großer Teil der Tensor-Core-Kapazität ungenutzt.
In MoE-Schichten fällt der Effekt stärker aus. In einem Fused-MoE-Kernel ergibt sich N aus der Zahl der gewählten Experten multipliziert mit der Zwischengröße der Experten. Bei einem Modell, das eine topk=8-Routine verwendet, wächst die effektive N-Dimension nach Angaben von AWS auf Zehntausende bis Hunderttausende Parameter.
Die Lösung beruht auf Matrixtransposition: Statt A x B wird BT x AT berechnet, was das mathematische Ergebnis unverändert lässt. Dadurch rückt die große Expertendimension in den ersten Operanden-Slot. Laut AWS füllt das die 64-Zeilen-Schwelle und erlaubt es zudem den Hardware-Copy-Engines von Hopper, Daten asynchron aus dem Shared Memory in die Register vorzuladen. Das Unternehmen berichtet von einer Verbesserung des End-to-End-Ausgabedurchsatzes um 2 bis 4 Prozent bei großen MoE-Modellen auf H200. Die Änderung wurde in vLLM übernommen und ist ab Version 0.24.0 verfügbar.
Backend-Wechsel für FP8-Modelle
Der zweite Gewinn stammt aus der Wahl eines anderen Backends. Das vLLM-MoE-Oracle leitet Quantisierungsschemata an Backends wie Triton oder FlashInfer+CUTLASS weiter. AWS zufolge hat das Profiling gezeigt, dass bei FP8-Tensor-quantisierten Modellen auf Hopper das Triton-Backend den Standardpfad FLASHINFER_CUTLASS durchgehend übertrifft.
AWS führt das darauf zurück, wie die beiden Backends die Berechnung handhaben. CUTLASS behandelt FP8-Tensor-Skalierung, Matrixmultiplikation und Expertenaggregation als getrennte Operationen. Triton fasst die Pipeline dagegen in einem einzigen Hardware-Durchlauf zusammen: Die Daten werden einmal aus dem globalen Speicher mit hoher Bandbreite gelesen, in lokalen SRAM-Registern verarbeitet und dann an das Ziel geschrieben. AWS ergänzt, dass CUTLASS nativ für dichte, große und vorhersehbare Matrixdimensionen optimiert ist. Seine Standard-Kachelung und das Thread-Mapping seien daher bei den extrem dünnen Batches (M < 64 oder M = 1), wie sie beim Decoding typisch sind, ineffizient. Triton verteile kleine, fragmentierte Token-Arrays dagegen dynamisch auf die aktiven Expertengewichte.
Die Anpassung der Routing-Richtlinie ist eine reine Konfigurationsänderung, ohne dass Kernel-Code verändert wird. Dennoch meldet AWS eine Verbesserung des End-to-End-Durchsatzes um 10 Prozent für FP8-quantisierte Fused-MoE-Modelle auf H200.
Was das für Anwender bedeutet
Für Teams, die MoE-Modelle auf GPUs der Hopper-Klasse betreiben, nennt die Quelle zwei praktische Schritte. Die Transpositionsoptimierung steht ab vLLM 0.24.0 zur Verfügung, der Backend-Wechsel erfordert lediglich eine Konfigurationsänderung. Alle Zahlen sind eigene Messungen von AWS auf H200 und gelten für die beschriebenen Workloads. Bei anderen Modellen oder anderer Hardware können die Ergebnisse abweichen.
