Am 1. August fand in Peking der 9. Meet AI Compiler Technical Salon statt, veranstaltet von HyperAI. Im Mittelpunkt der Veranstaltung standen die neuesten Fortschritte in der KI-Kompilierungstechnologie. Zahlreiche Experten aus Industrie und Forschungseinrichtungen präsentierten ihre Erkenntnisse zu Programmiersprachen, Operatorentwicklung, Kompilierungsoptimierung und Inferenzausführung und demonstrierten die kollaborative Entwicklung von KI-Compilern von der Hochsprachen-Ausdrucksformulierung bis zur Hardware-Ausführung.

In,Zhu Honglin, ein Ingenieur für Hochleistungsrechnen bei Tencent, stellte in einem Vortrag mit dem Titel „FalconGEMM: Überwindung von Hardware-Grenzen durch Matrixmultiplikation mit geringer Komplexität“ die Algorithmen- und Operatoroptimierungspraktiken seines Teams für die Matrixmultiplikation mit geringer Komplexität vor.

Angesichts der Herausforderung, dass ausgereifte Operatorbibliotheken wie cuBLAS die Leistung der Matrixmultiplikation nahezu an die Hardware-Grenze gebracht haben und der traditionelle Optimierungsraum auf Kernel-Ebene immer kleiner wird, ging das Team erneut von der algorithmischen Komplexität aus. Basierend auf Matrixmultiplikationsalgorithmen mit geringer Komplexität wie Strassen und AlphaTensor entwickelten sie ein einheitliches LCMA-Framework (Low-Complexity Matrix Algorithms). Durch die Kombination von QDSL, Operatorfusion, Persistent Kernel, feingranularer Ablaufplanung und Kostenmodell konnten sie den theoretischen Vorteil der „Reduzierung der Anzahl von Multiplikationen“ in tatsächliche Leistungssteigerungen auf GPUs umsetzen.

In den FP16- und BF16-Tests von NVIDIA H20 übertrifft FalconGEMM cuBLAS bei einer großen Anzahl von Matrixformen, mit Spitzenleistungsverbesserungen von etwa 10%-16%, während die numerische Genauigkeit im Wesentlichen mit der Standard-Matrixmultiplikation im Sprachmodell-Benchmark übereinstimmt.

HyperAI hat die geteilten Inhalte zusammengestellt und zusammengefasst, ohne deren ursprüngliche Bedeutung zu verändern.

Folge dem offiziellen WeChat-Konto "HyperAI" und antworte im Hintergrund mit dem Schlüsselwort "".0801 KI-CompilerSie können die Präsentations-PPT des autorisierten Referenten durch Anklicken von "..." abrufen.

Ausgehend von Strassen untersuchen wir erneut den Optimierungsraum für die Matrixmultiplikation.

Die Matrixmultiplikation ist einer der wichtigsten fundamentalen Operatoren im Deep Learning und beansprucht üblicherweise den größten Teil der Rechenzeit des Modells. Software-Stacks wie CUDA, MKL und cuBLAS sind seit vielen Jahren optimiert, und in vielen Szenarien ist die Leistung eines einzelnen GEMM-Kernels bereits nahezu maximal. Das bedeutet, dass der Spielraum für Verbesserungen zunehmend begrenzt wird, wenn wir weiterhin nur lokale Optimierungen auf Befehls-, Pipeline- und Speicherzugriffsebene vornehmen. Daher wandte sich das Team wieder dem klassischen Strassen-Algorithmus zu.

Der Strassen-Algorithmus wurde 1969 von Volker Strassen vorgeschlagen. Für die einfachste 2×2-Matrixmultiplikation benötigt die traditionelle Methode acht Multiplikationen, während Strassen durch Rekombination der Eingabematrix nur sieben Multiplikationen benötigt und anschließend zusätzliche Additions- und Subtraktionsoperationen verwendet, um das Endergebnis wiederherzustellen. Dies entspricht einer Reduzierung des Rechenaufwands um ein Achtel. Sind die Operanden lediglich Skalare, lohnt sich dieser Kompromiss nicht; werden die Operanden jedoch zu Untermatrizen, wird der Komplexitätsunterschied zwischen Matrixaddition (O(N²)) und Matrixmultiplikation (O(N³)) signifikant.Dadurch gewinnt die Aussage „eine Matrixmultiplikation weniger und mehrere Matrixadditionen mehr durchführen“ einen praktischen Nutzen.

Wird der Strassen-Algorithmus rekursiv angewendet, lässt sich die Anzahl der Multiplikationen weiter reduzieren. Beispielsweise benötigt eine herkömmliche Multiplikation einer 4×4-Blockmatrix 64 Blockmultiplikationen, während zwei Strassen-Ebenen nur 49 benötigen. Allerdings führt eine Erhöhung der Anzahl rekursiver Ebenen auch zu mehr Additionen, einem höheren Aufwand für die Datenorganisation und einem größeren Speicherzugriffsaufwand.Daher verwenden praktische Systeme oft nur eine begrenzte Anzahl von Schichten, um ein Gleichgewicht zwischen Rechenaufwandsreduzierung und zusätzlichem Aufwand zu finden.

Im Jahr 2022 erweiterte DeepMinds AlphaTensor diesen algorithmischen Raum nochmals. Es transformiert die Matrixmultiplikation in ein Tensorzerlegungsproblem und nutzt Reinforcement Learning, um Zerlegungsmethoden mit niedrigerem Rang zu finden. Dies zeigt, dass es neben dem klassischen Strassen-Algorithmus eine Vielzahl unterschiedlicher Matrixmultiplikationsalgorithmen mit geringer Komplexität für verschiedene Formen von M, N und K geben kann.

Allerdings besteht nach wie vor ein praktisches Problem beim Übergang von der Algorithmenentwicklung zu technischen Anwendungen:Wenn für jeden Algorithmus mit geringer Komplexität ein separater, handgeschriebener GPU-Kernel erforderlich wäre, wären die Entwicklungs- und Wartungskosten offensichtlich zu hoch.Um dieses Problem zu lösen, abstrahierte das Team diese Algorithmen in LCMA (Low-Complexity Matrix Algorithms). LCMA beschreibt einheitlich, welche Teilblöcke der Eingabematrix vorab kombiniert werden müssen, wie viele Matrixmultiplikationen tatsächlich durchgeführt werden und wie die Zwischenergebnisse schließlich zur Ausgabematrix zusammengeführt werden. Die entsprechende Implementierung wird dann automatisch mit Codegen generiert.

Daher stellt sich die Frage: „Wie implementiert man einen Strassen-Kernel?“Der Fokus verlagerte sich auf die Frage, „wie man ein einheitliches Framework aufbauen kann, das eine Vielzahl von Matrixalgorithmen mit geringer Komplexität unterstützt und gleichzeitig eine hohe Leistungsfähigkeit gewährleistet“.

Gleichzeitig müssen Algorithmen mit geringer Komplexität auch die numerische Genauigkeit berücksichtigen. Obwohl die Strassen-Multiplikation algebraisch äquivalent zur Standard-Matrixmultiplikation ist, unterliegen Gleitkommaoperationen nicht strikt der Assoziativität, und Änderungen der Berechnungsreihenfolge können zusätzliche Rundungsfehler verursachen. Daher muss LCMA neben der Optimierung der Leistung auch die Fehlerfortpflanzung bei Berechnungen mit geringer Genauigkeit kontrollieren.

Nachdem eine einheitliche Algorithmenbeschreibung vorlag, bestand der nächste Schritt darin, eine geeignete GPU-Implementierung zu finden. Das Team testete CUDA, Triton, TiLang und QDSL. CUDA bot die stärkste Hardwarekontrolle, doch angesichts der Vielzahl unterschiedlicher LCMA-Algorithmen erforderten Register, gemeinsam genutzter Speicher und Zwischensummenstrukturen jeweils spezifische Anpassungen, was zu hohen Erweiterungs- und Wartungskosten führte.

Triton erreicht in einfachen Strassen-Szenarien nahezu CUDA-Performance. Bei der Skalierung des Algorithmus auf größere Blockstrukturen ist jedoch eine präzise Wiederverwendung von Registerpuffern über mehrere Zwischenberechnungen hinweg erforderlich, was Triton anfälliger für zusätzliche Speicherüberläufe macht. TiLang bietet mehr Flexibilität bei der Register- und Shared-Memory-Steuerung, seine Performance lag in Teamtests aber dennoch um etwa 51–101 TP3T unter der von Triton. Für einfaches Strassen, mit einem theoretischen Gewinn von nur 12,51 TP3T, ist dieser Verlust signifikant genug, um die Vorteile des Algorithmus zu schmälern.

Finale,Das Team wählte QDSL als Hauptimplementierungs-Backend für FalconGEMM. QDSL bietet eine Entwicklungsgranularität ähnlich der von CUDA und verfügt gleichzeitig über Codegenerierungsfunktionen sowie Unterstützung für eingebettetes PTX. Dies erleichtert die Migration bestehender Hochleistungsimplementierungen und eignet sich für die Batch-Codegenerierung auf Basis verschiedener LCMA-Beschreibungen, was eine größere Flexibilität für die nachfolgende Integration und kundenspezifische Optimierung ermöglicht.

Von LCMA zu FalconGEMM: Übertragung algorithmischer Vorteile auf GPUs

Die einfachste Strassen-GPU-Implementierung lässt sich in mehrere Schritte unterteilen: Kombination der Teilmatrizen A und B zur Generierung von 7 neuen Eingabepaaren; Ausführung von 7 gebündelten GEMMs; und schließlich Zusammenführung der 7 Mengen von Zwischenergebnissen zur endgültigen Matrix C. Verglichen mit einem regulären GEMM beträgt der rechenintensive Teil der Matrixmultiplikation nur 7/8 des ursprünglichen Aufwands.Solange der zusätzliche Zeitaufwand für die Vor- und Nachbearbeitung weniger als 1/8 der eingesparten Rechenlast beträgt, besteht daher die Möglichkeit, insgesamt Vorteile zu erzielen.

Das Team testete zunächst auf der NVIDIA H20. Die H20 eignet sich mit ihrer hohen Speicherbandbreite und relativ geringen Spitzenrechenleistung gut für diesen Ansatz, „die Datenverarbeitung zu erhöhen und gleichzeitig die Rechenlast zu reduzieren“. Bei Matrixgrößen ab etwa 2048³ zeigte die Basisimplementierung bereits stabile Leistungssteigerungen. Bei kleineren Matrixformen steigt jedoch der Anteil der Eingabekombinationen, der Zwischenergebnisse und der Ausgabekombinationen rapide an und verbraucht so schnell die eingesparte Rechenlast.

Daher verlagerte sich der Fokus der nachfolgenden Optimierung von GEMM selbst auf den Zugriff auf den Zwischenspeicher. Der direkteste Ansatz ist die Operatorfusion.Versuchen Sie, Zwischenergebnisse so weit wie möglich auf dem Chip zu speichern, anstatt sie wiederholt in den globalen Speicher zurückzuschreiben.Die direkte Kombination der Eingaben A/B in GEMM ist jedoch ungeeignet, da derselbe Teilblock von mehreren SMs verwendet werden kann, was leicht zu doppeltem Laden und Summieren führt. Im Gegensatz dazu ist die nachträgliche Fusion von Batched GEMM und Combine H praktikabler.

Die eigentliche Herausforderung besteht darin, dass Strassens sieben Zwischenergebnisse auf unterschiedliche Weise zu den vier finalen Ausgabesubmatrizen beitragen. Wird H als parallele Einheit verwendet, können mehrere SMs gleichzeitig in dieselbe C zurückschreiben, was zu schwerwiegenden atomaren Konflikten führt; wird C als parallele Einheit verwendet, wird ein Teil von H von verschiedenen SMs wiederholt berechnet. Beide Ansätze heben die Vorteile der reduzierten Multiplikation auf.

Das Team gab schließlich die Organisation der Aufgaben anhand von Strassens Zwischenergebnissen auf. Stattdessen gruppierten sie die Aufgaben nach den räumlichen Koordinaten der Matrix: Sieben Multiplikationskacheln an derselben Position in sieben zusammengefassten GEMMs wurden zu einer Gruppe zusammengefasst und auf demselben SM ausgeführt. Dadurch konnten die Ergebnisse nach Abschluss der Berechnungen einer Gruppe direkt im finalen C auf dem Chip akkumuliert werden, wodurch das Zurückschreiben von Zwischenergebnissen in den globalen Speicher entfiel und erhebliche Schreibkonflikte zwischen den SMs vermieden wurden.

Dieser Fusionsansatz reduziert die durch Strassen verursachten zusätzlichen Speicherzugriffe deutlich, die größere Gruppengranularität führt jedoch zu einer Lastungleichverteilung. Beispielsweise kann bei einer 4096³-Matrixmultiplikation die grobkörnige Planung zu einem zusätzlichen Wellenverlust von etwa 211 TP3T führen und damit sogar die durch Strassen selbst erzielte Rechenzeitreduktion von 12,51 TP3T übertreffen.

zu diesem Zweck,Das Team ließ sich von Stream-K inspirieren, um eine Gruppe bei Bedarf in zwei SMs für die Ausführung aufzuteilen.Die Planungsschicht verwendet weiterhin die Gruppe als Basiseinheit, die eigentliche Ausführung kann jedoch weiter auf Kachelebene verfeinert werden, wodurch die Leerlaufzeit des letzten SM reduziert und die Hardwareauslastung verbessert wird, während gleichzeitig die Vorteile der Datenwiederverwendung auf Gruppenebene erhalten bleiben.

Nach der Behebung des Lastverteilungsproblems trat jedoch ein neues Problem auf: L2-Cache-Thrashing. Durch die Aufteilung der Gruppen konnten verschiedene Arten von Zwischenmultiplikationen innerhalb derselben Welle vermischt werden, und die abgerufenen Daten waren unabhängig voneinander, was zu einem deutlichen Rückgang der L2-Cache-Trefferrate führte. Gleichzeitig nutzte GEMM die Tensor-Kerne bereits stark aus, und bei nahezu maximaler Speicherauslastung stieß der H2O an seine Leistungsgrenze. In realen Tests sank die Kernfrequenz von etwa 1,8 GHz auf 1,6 GHz, was zu einem Rückgang der Rechenleistung führte und einige der Vorteile der Fusion erneut zunichtemachte.

Um das Thrashing des L2-Caches zu beheben, passte das Team die Reihenfolge der Split-Gruppen weiter an. Dadurch wurde sichergestellt, dass Zwischenergebnisse desselben Typs in derselben Verarbeitungswelle verarbeitet wurden und gemischte Verarbeitung nur in wenigen Nachbearbeitungswellen stattfand. Dies erhielt die durch die feingranulare Planung erreichte Lastverteilung bei gleichzeitig verbesserter Datenlokalität im L2-Cache und beseitigte letztendlich das Problem der signifikanten Frequenzreduzierung.

Es ist erwähnenswert, dassDie wichtigste Grundlage für diese Scheduling-Optimierungen ist der Persistent Kernel.Anders als bei einem regulären Kernel, bei dem der Task Advisor (CTA) nach Abschluss eines Blocks beendet wird, ermöglicht ein persistenter Kernel dem CTA, über einen längeren Zeitraum im Streaming Service (SM) zu verbleiben und kontinuierlich weitere Aufgaben anzunehmen. Dies bietet Entwicklern mehr Flexibilität bei der Steuerung der Ausführungsreihenfolge von Gruppen und Tiles und ermöglicht die Wiederverwendung von On-Chip-Ressourcen. Aufgabenaufteilung, Umstrukturierung der Ablaufplanung und Optimierung des Caching können somit ebenfalls innerhalb desselben Kernels erfolgen.

Von der Cache-Neuanordnung bis zum Kostenmodell beträgt die maximale Leistungsverbesserung 10%-16%.

Nach Zusammenführung, Lastverteilung und Cache-Neuanordnung kann FalconGEMM die Rechenvorteile von LCMA auf mehr Datenstrukturen anwenden. LCMA ist jedoch nicht in allen Fällen dem herkömmlichen GEMM überlegen:Im Kern bleibt es jedoch dasselbe: Man tauscht zusätzlichen Datenverarbeitungsaufwand gegen weniger Multiplikationsberechnungen.Wenn das ursprüngliche GEMM bereits durch den Speicherzugriff begrenzt ist, bringt eine weitere Reduzierung der Rechenleistung keine ausreichenden Vorteile; Algorithmen mit geringer Komplexität sind nur dann vorteilhafter, wenn die Rechendichte hoch ist.

daher,Das Team entwarf außerdem ein Roofline-ähnliches Kostenmodell, um zu bestimmen, wann LCMA eingesetzt werden soll und welche der verschiedenen LCMA-Methoden ausgewählt werden soll.Da das Ziel darin besteht, "den richtigen Algorithmus auszuwählen" und nicht die Ausführungszeit genau vorherzusagen, analysiert das Modell hauptsächlich die Rechen- und Speicherzugriffskosten verschiedener Verfahren und schätzt den Engpassbereich für Rechenleistung und Speicherzugriff der Ziel-GPU durch die Kombination ihrer Rechenleistung und Bandbreite ab.

In diesem Modell entspricht die reduzierte Anzahl an Multiplikationen im Algorithmus mit geringer Komplexität einem Rechengewinn, während zusätzliche Datenkombinationen und wiederholte Speicherzugriffe einen neuen Speicheraufwand verursachen; die zuvor erwähnte Fusionsoptimierung reduziert diese Speicherzugriffskosten weiter. Daher kann FalconGEMM die Nutzengrenzen von herkömmlichem GEMM und verschiedenen LCMA-Verfahren basierend auf unterschiedlichen M-, N- und K-Formen bestimmen und automatisch eine geeignetere Implementierung auswählen.

Durch die Nutzung der Codegenerierungsfunktionen von QDSL bildet das gesamte Framework letztendlich einen relativ vollständigen Ausführungsablauf: Zunächst wird basierend auf der LCMA-Beschreibung ein entsprechender, fusionierter persistenter Kernel generiert, wodurch durch Fusion die Zwischenspeicherzugriffe reduziert werden. Anschließend wählt das Kostenmodell einen geeigneten Matrixmultiplikationsalgorithmus für die jeweilige Form aus. Schließlich generiert und kompiliert QDSL automatisch den Zielcode. Auf diese Weise ist LCMA nicht mehr nur eine feste Strassen-Implementierung, sondern bildet einen Algorithmenraum, der dynamisch je nach Arbeitslast ausgewählt werden kann.

Die Leistungstests wurden hauptsächlich auf der NVIDIA H20 durchgeführt. Die Ergebnisse zeigen, dass in verschiedenen Szenarien der Matrixmultiplikation mit niedriger GenauigkeitFalconGEMM übertrifft cuBLAS bei einer großen Anzahl von Formen, mit Spitzenleistungsverbesserungen von ungefähr 10%–16%.Nach Abschluss der Gruppenaufteilung und Cache-Neuanordnung kann die große Struktur stabil von den Rechenvorteilen des Algorithmus mit geringer Komplexität profitieren, die Leistung der kleinen Struktur wird ebenfalls verbessert und die durch den Leistungsabfall aufgrund von L2-Cache-Jitter verursachte Frequenzreduzierung wird vermieden.

Auf Grundlage der Ergebnisse der Kostenmodellauswahl konnte das Modell für die meisten Formen, die den LCMA-Rückgabeschwellenwert überschreiten, die leistungsfähigere Implementierung auswählen. Dies deutet darauf hin, dass der Ansatz „Algorithmusauswahl + Fusionskern“ Matrixmultiplikationsszenarien mit unterschiedlichen Rechendichten effektiv abdecken kann.

Neben der Leistung ist auch die numerische Genauigkeit ein entscheidender Faktor, den FalconGEMM überprüfen muss. Das Team stieß in frühen Experimenten mit geringer Präzision auf erhebliche Fehler, hauptsächlich aufgrund von Änderungen in der Reihenfolge der Gleitkommaaddition. Beispielsweise ist A+C+B−C algebraisch gleich A+B, aber bei Gleitkommaberechnungen mit endlicher Genauigkeit sind sie nicht unbedingt exakt gleich.

Weitere Analysen ergaben, dass die wirklich signifikanten Fehler primär auf ungenaue Berechnungen und nicht auf die FP32-Akkumulation selbst zurückzuführen waren. In gängigen FP16/BF16-Eingabeszenarien wird die Matrixmultiplikation typischerweise zunächst mit FP32-Akkumulation durchgeführt und anschließend wieder in eine niedrigere Genauigkeit umgewandelt. Werden häufig Zwischenergebnisse berechnet, gehen die FP32-Mantisseninformationen dabei kontinuierlich verloren.

Der Fusionsansatz behebt dieses Problem. Die Ausgabe von WGMMA behält die FP32-Genauigkeit bei, während FalconGEMM die finalen C-Werte direkt auf dem Chip mit FP32 kombiniert und akkumuliert, um sie erst nach Abschluss der Berechnung wieder auf die Zielgenauigkeit umzurechnen. Im Vergleich zum wiederholten Zurückschreiben von Zwischenergebnissen mit niedriger Genauigkeit zwischen mehreren unabhängigen Kernels reduziert diese Methode eine oder mehrere Genauigkeitsumrechnungen, wodurch Fehler, die durch Änderungen in der Berechnungsreihenfolge entstehen, eher in den unteren Bits von FP32 verbleiben.

Im Sprachmodell-Benchmark sind die mit FalconGEMM und der Standard-Matrixmultiplikation erzielten Endergebnisse nahezu identisch, mit nur sehr geringen Unterschieden, was darauf hindeutet, dass die aktuelle Implementierung nicht zu einer signifikanten Verringerung der Modellgenauigkeit führt.

In der nächsten Phase plant das Team, die Arbeit in zwei Richtungen fortzusetzen: Erstens soll Combine A/B durch Anpassung der Reihenfolge von Batch-Gruppen und K-Schleifen weiter integriert werden, um den Speicherzugriff für Zwischenergebnisse weiter zu reduzieren. Zweitens soll LCMA auf Attention erweitert werden. Flash Attention weist ebenfalls ein hohes Verhältnis von Rechenaufwand zu Speicherbedarf auf, und die weitere Kombination von Matrixalgorithmen mit geringer Komplexität mit Blockzerlegung und Pipeline könnte zusätzliches Leistungspotenzial eröffnen.

Von Strassen und AlphaTensor bis hin zu LCMA und FalconGEMM geht die Bedeutung dieser Arbeit weit über die bloße Beschleunigung des bereits hochoptimierten GEMM um einige Prozentpunkte hinaus. Sie bietet einen anderen Ansatz: Wenn der Kernel selbst bereits an die Hardwaregrenzen stößt, kann die Leistungsoptimierung nicht nur die Befehle und Pipelines weiter optimieren, sondern auch neue Möglichkeiten zur Verbesserung der algorithmischen Komplexität aufzeigen. Durch Kompilierung, Fusion und Scheduling lässt sich die theoretische Reduzierung des Rechenaufwands dann tatsächlich in Laufzeitvorteile umsetzen.