LLMs in einen MegaKernel kompilieren, um Low-Latency-Inferenz zu ermöglichen
(zhihaojia.medium.com)- Es wurde ein Compiler entwickelt, der die LLM-Inferenz automatisch in einen einzelnen MegaKernel umwandelt
- Der MegaKernel-(Persistent-Kernel-)Ansatz integriert bei der LLM-Inferenz Berechnung und Kommunikation vollständig in einen einzigen GPU-Kernel und ermöglicht dadurch eine sehr niedrige Latenz
- Aufgrund der verteilten Struktur bestehender ML-Frameworks und Kernel-Bibliotheken ist eine Single-Kernelisierung der gesamten Pipeline äußerst schwierig
- Mirage Persistent Kernel (MPK) wandelt mithilfe eines Compiler- und Laufzeitsystems Multi-GPU-LLM-Inferenz automatisch in einen hochperformanten Megakernel um
- MPK transformiert den Berechnungsgraphen in einen fein granularen Task-Graphen und maximiert so Software-Pipelining sowie die Überlappung von Berechnung und Kommunikation
- Mit MPK sinkt im Vergleich zu bestehenden Systemen die Verzögerung bei der Token-Generierung, und der Leistungsgewinn wächst mit der Anzahl der GPUs weiter an
Überblick und Vorteile des MegaKernel-Ansatzes
- Bei der Inferenz großer Sprachmodelle (LLMs) ist eine wirksame Methode zur Reduzierung der Latenz, alle Berechnungs- und Kommunikationsprozesse in einen einzelnen Megakernel (konsistenten Kernel) zu fusionieren
- Bei diesem Ansatz übernimmt ein einziger GPU-Kernel ohne Unterbrechung die gesamte Verarbeitung, einschließlich der schichtweisen Operationen des Modells und der Kommunikation zwischen GPUs
- Die wichtigsten Vorteile sind:
- Wegfall wiederholter Kernel-Aufrufe und damit Eliminierung des Kernel-Launch-Overheads
- Software-Pipelining über mehrere Layer hinweg wird möglich
- Berechnung und Kommunikation laufen gleichzeitig, wodurch Latenz verborgen wird
Bisherige Grenzen und das Aufkommen von MPK
- Bestehende ML-Frameworks wie PyTorch, Triton und TVM unterstützen die automatische end-to-end-Erzeugung eines Megakernels grundsätzlich nicht
- Reale LLM-Systeme bestehen aus einer Kombination verschiedener Kernel-Bibliotheken wie NCCL/NVSHMEM (Kommunikation), FlashInfer/FlashAttention (Attention), CUDA/Triton (benutzerdefinierte Operationen), was die Integration in einen einzigen Kernel erschwert
- Vor diesem Hintergrund entwickelten Forschende von CMU, UW, Berkeley, NVIDIA und Tsinghua den Mirage Persistent Kernel (MPK)
- MPK kombiniert Compiler und Laufzeit und wandelt die gesamte LLM-Inferenz-Pipeline automatisch in einen hochperformanten Megakernel um
Der zentrale Wert von MPK
- MPK eliminiert den Kernel-Launch-Overhead vollständig und maximiert die Überlappung von schichtübergreifender Berechnung, Datenladen und Kommunikation, um eine LLM-Inferenz mit extrem niedriger Latenz zu ermöglichen
- In realen Tests (39-Token-Prompt, Generierung von 512 Tokens, ohne speculative decoding)
- verkürzt MPK auf einer einzelnen NVIDIA A100 40GB GPU die Decoding-Latenz pro Token auf 12,5 ms, verglichen mit 14,5 ms bei bestehenden optimierten Systemen wie vLLM/SGLang
- dieser Wert liegt nahe an der theoretischen Untergrenze von 10 ms (bei 1,6 TB/s Speicherbandbreite und Laden von 16 GB Gewichten)
- In Multi-GPU-Umgebungen werden Berechnung und Kommunikation vollständig integriert, wodurch der Leistungsvorteil von MPK mit steigender GPU-Zahl noch deutlicher wird
Details zur Funktionsweise von MPK
Part 1. Compiler – Umwandlung des LLM-Berechnungsgraphen in einen Task-Graphen
- Im Allgemeinen wird die LLM-Berechnung als Computational Graph dargestellt, in dem jede Operation (z. B. Matrixmultiplikation, Attention) oder Kommunikationsoperation (z. B. all-reduce) ein Knoten und Datenabhängigkeiten die Kanten sind
- In herkömmlichen Designs ist die Ausführung eines separaten Kernels pro Operator üblich, doch das bildet nur Abhängigkeiten auf Kernel-Ebene und nicht auf Ebene der tatsächlich abhängigen Dateneinheiten ab, wodurch Pipelining-Möglichkeiten begrenzt sind
- Beispiel: Folgt auf eine Matrixmultiplikation ein all-reduce, kann all-reduce erst starten, wenn die gesamte Matrixmultiplikation abgeschlossen ist. Tatsächlich wäre es möglich, die Daten aufzuteilen und partielle Ausführung sowie Abhängigkeiten auszunutzen
- Der MPK-Compiler verfeinert den Berechnungsgraphen und wandelt ihn automatisch in einen fein granularen Task-Graphen um, der zur tatsächlichen Dateneinheit passt
- Jeder Task (Rechteck) ist eine Berechnungs- oder Kommunikationseinheit, die einem einzelnen GPU-SM zugewiesen wird
- Jedes Event (Kreis) ist ein Synchronisationspunkt zwischen Tasks
- Kanten zwischen Tasks und Events drücken Daten- und Steuerabhängigkeiten effizient aus
- Dank dieses Task-Graphen kann MPK Berechnung und Kommunikation stärker teilweise oder parallel überlappen lassen
- Mit dem Mirage kernel superoptimizer werden zudem automatisch hochperformante CUDA-Implementierungen für jeden Task erzeugt
Part 2. Laufzeit – Ausführung des Task-Graphen innerhalb des Megakernels
- Die MPK-Laufzeit führt den Task-Graphen vollständig innerhalb eines einzigen GPU-Kernels (Megakernels) aus
- Alle SMs (Streaming Multiprocessors) der GPU werden statisch in Worker- und Scheduler-Rollen aufgeteilt
Worker
- Jeder Worker arbeitet auf SM-Ebene und verwaltet eine eigene Task-Queue
- In einer Schleife
- holt er den nächsten Task aus der Queue
- führt ihn aus (z. B. matmul, attention, Datentransfer)
- meldet bei Abschluss ein Event
- und wiederholt den Vorgang
- Dadurch werden Ressourcennutzung pro Worker optimiert und asynchrone schichtübergreifende Operationen ermöglicht
Scheduler
- Ein verteilter Scheduler arbeitet pro SM auf Ebene eines einzelnen Warps; bis zu 4 Scheduler können gleichzeitig laufen
- Jeder Scheduler verwaltet eine Queue aktivierter Events und weist Tasks, deren Bedingungen erfüllt sind, den Workern zu
- Dadurch ist eine großskalige verteilte Task-Verarbeitung ohne zentralisierten Synchronisations-Overhead möglich
Ereignisbasierte Ausführung
- Wenn ein Task abgeschlossen ist, wird ein bestimmter Event-Zähler erhöht. Erreicht der Zähler einen Schwellenwert, wird das Event aktiviert und in die Scheduler-Queue eingefügt
- Der Scheduler führt anschließend die nachfolgenden Tasks aus, die von diesem Event abhängen
- Dadurch entstehen fein granulareres Software-Pipelining und Überlappung von Berechnung und Kommunikation auf natürliche Weise
- Beispiel: matmul in einem Layer und attention in einem anderen Layer können gleichzeitig laufen
- Sobald teilweise abgeschlossene matmul-Ergebnisse vorliegen, kann die all-reduce-Kommunikation beginnen
- Da sämtliches Scheduling und alle Task-Wechsel innerhalb eines einzigen Kernel-Kontexts stattfinden, ist der Overhead zwischen Tasks mit 1–2 Mikrosekunden (μs) sehr gering
Ausblick
-
Ziel von MPK: Entwickler sollen mit wenig Python-Code (nur einige Dutzend Zeilen) LLMs einfach in einen Megakernel kompilieren und maximale Leistung erzielen können
-
Wichtige Entwicklungsrichtungen
- Unterstützung aktueller GPU-Architekturen: etwa für NVIDIA Blackwell und warp-spezialisierte Ansätze
- Verarbeitung dynamischer Workloads: Forschung an Kompilierungsstrategien für Modelle mit dynamischem Kontrollfluss wie mixture-of-experts (MoE)
- Fortgeschrittenes Task-Scheduling: Untersuchung und potenzielle Anwendung moderner Richtlinien wie Prioritätsbasiertheit und Durchsatzoptimierung
-
MPK markiert einen grundlegenden Wendepunkt in der Art, wie GPU-basierte LLM-Inferenz kompiliert und ausgeführt wird, und wünscht sich eine engere Zusammenarbeit mit der Community
Zusätzliche Materialien
- Code und Dokumentation zu MPK (Mirage Persistent Kernel) sowie aktuelle Forschungsergebnisse sind auf GitHub verfügbar: https://github.com/mirage-project/mirage
1 Kommentare
Hacker-News-Kommentare
An den Autor: Es ist interessant, dass der On-GPU-Interpreter-Ansatz wie eine sehr vielversprechende Richtung für die Zukunft wirkt. Es gibt auch andere Arbeiten mit einem fast identischen Ansatz, daher empfehle ich diesen verwandten Beitrag. Das grundlegende Programmiermodell von CUDA (z. B. Kernel-Launches) wird für feingranulare, auf kleinen Aufgaben basierende Parallelisierung umgangen, und ich habe selbst gesehen, dass dieser Ansatz die Hardware besser auslastet. Ich frage mich, ob CUDA uns in vielerlei Hinsicht nicht ausgebremst hat. Ich bin gespannt, ob die Forschung des Autors als experimentelles Backend in PyTorch landen könnte. Und noch ein kleiner Hinweis: Die beiden Absätze im ersten Teil sind fast identisch, also wohl ein kleiner Tippfehler.
Ich habe eine Zeit lang eng an vLLM und SGLang gearbeitet und bin überzeugt, dass dieses Projekt genau so aussehen sollte wie ein ideales Nachfolgeprojekt. Beeindruckend ist die Analyse des Berechnungs-Abhängigkeitsgraphen sowie das Fusionieren von Operationen und die intelligentere Aufgabenplanung. Glückwunsch an das Team.
Ich habe den Artikel und das GitHub-README überflogen und finde, dass das ein wirklich großartiges Projekt ist. Ich frage mich, ob sich solche Optimierungsansätze nicht nur für Inferenz, sondern auch für das Training anwenden lassen. Mir ist bewusst, dass insbesondere das Fusionieren von Backward-Operationen und Gradientenkommunikation eine Herausforderung wäre. Soweit ich weiß, wird derzeit kein dynamischer Workload (z. B. MoE) unterstützt, aber es gibt kürzlich das Paper FlashDMoE: Fast Distributed MoE in a Single Kernel zum Verarbeiten von MoE in einem einzigen Kernel.
Vielen Dank, dass du sogar den Artikel und das README gelesen hast. Unterstützung für den Trainingsschritt ist ebenfalls möglich, aber Trainingskernel sind im Allgemeinen größer, sodass der Kernel-Launch-Overhead dort kein so großes Problem ist. Deshalb profitiert vor allem die Inferenz, insbesondere bei niedriger Latenz. Das von dir geteilte FlashDMoE-Paper fanden wir ebenfalls sehr interessant, und die Unterstützung von MoE-Modellen ist eines unserer nächsten Ziele.
Persönlich bin ich etwas skeptisch, ob es sinnvoll ist, viel Zeit in gradientenbasierte Optimierung für Training zu investieren. Viele reale Trainingsaufgaben haben diskrete Eigenschaften und lassen sich meiner Meinung nach mit gradientenbasiertem Lernen nicht gut behandeln.
Der nächste Schritt wäre dann wohl, direkt nach Verilog zu kompilieren und sich bei AliExpress einfach eigene LLM-Hardware zu kaufen.
Hier ist ein Beitrag zur Einführung in Hardware-Beschreibung mit Chisel usw.. Vor dem Aufkommen von AI und GPUs war die Idee, Software direkt in Hardware zu überführen, ein vielversprechender Ansatz. Die CPU-Entwicklung stagniert, und der Wunsch, die Zwischenschicht zwischen Software und Hardware stärker zu optimieren, bleibt bestehen, aber GPU-artiges paralleles Rechnen wird wahrscheinlich der dominierende Beschleunigungsansatz bleiben. Allgemeine CPUs werden am Ende wohl nur noch als kleines Gehirn übrig bleiben, das die GPUs verwaltet. Trotzdem dürfte sich ein direkter Übergang von Software zu Hardware kaum zum Mainstream entwickeln.
Wenn sich die Struktur von LLMs in 5 bis 10 Jahren stabilisiert, könnte eine direkte Abbildung auf Hardware praktisch werden. Mit heutiger Technologie könnten sogar Modelle mit mehreren zehn Milliarden Parametern auf einen einzelnen Wafer passen, wenn nur ultraniedrigpräzise Logikgatter im Bereich von etwa 1,5 Bit verwendet werden. Mit steigender Präzision wächst die Zahl der Gatter exponentiell, daher ist es derzeit effizienter, den Gewichtsspeicher beizubehalten und Recheneinheiten gemeinsam zu nutzen. In Zukunft wird die Entwicklung ultraniedrigpräziser LLMs wohl eine zentrale Aufgabe sein.
Ein Scherz darüber, dass die Trainingskosten ohnehin schon hoch sind und zusätzliche Maskenkosten die Lage nur noch verschlimmern würden, sowie die nüchterne Einschätzung, dass AI-Hardware-Startups solche Richtungen im Grunde schon seit Langem verfolgen.
Ein LLM-in-a-box-Ansatz wäre tatsächlich ziemlich attraktiv. Ich werde bald wohl die Gelegenheit haben, in einer Offline-Umgebung (air-gapped) zu arbeiten, und so eine Lösung wäre dort sehr nützlich.
Ich habe den Code selbst in einer Modal-GPU-Umgebung ausgeführt, und die in der Forschung behaupteten Leistungssteigerungen lassen sich tatsächlich reproduzieren. Hier ist der geteilte Ergebniscodes des mirage-Projekts. Mit einer Triton- + FlashInfer-Kombination lag die Latenz bei etwa 19,2 ms pro Token, bei MPK unter denselben Bedingungen bei 7,7 ms, also eine massive Verbesserung.
Ich habe früher einmal an einem kleinen CUDA-Wettbewerb teilgenommen. Es ging um einen parallelen Algorithmus im Bild- oder Vision-Bereich, und ich wollte besonders clever sein und habe Zwischenergebnisse im Speicher gecacht. Als die Ergebnisse des Wettbewerbs veröffentlicht wurden, war ich überrascht, dass andere viel schnelleren Code als ich eingereicht hatten. Der Grund war, dass sie solche Zwischenergebnisse gar nicht gecacht, sondern einfach immer wieder neu berechnet hatten. Die Rechenkosten waren viel geringer als die Speicherzugriffe. Ich vermute, bei diesem Projekt ist es ähnlich. Durch das Kompilieren zu einem Megakernel verschwinden die Layer-Grenzen, wodurch weniger Zwischenergebnisse geteilt werden und die Rechenmenge steigt, aber insgesamt gibt es einen großen Gewinn, weil Speicherzugriffe entfallen. Besonders bei Convolutional Networks dürfte es irgendeinen sweet spot geben, aber ich weiß nicht, wie der Megakernel damit umgeht.
Es tauchen immer noch neue Metaphern für LLMs auf. Vielleicht könnte man LLMs sogar wie Transistoren betrachten. Ich stelle mir vor, wir befinden uns gerade in einer Phase wie bei raumgroßen Computern, die mit Lochkarten nur Multiplikation ausführen konnten. Es macht Spaß, sich vorzustellen, was passieren würde, wenn man 1 Million o3-pro-Anfragen gleichzeitig ausführen könnte.
Dieses Projekt stammt aus CMU (Carnegie Mellon). Erwähnenswert ist auch der Blog No Bubbles von Hazy Research in Stanford über Megakernel. Es ist beeindruckend zu sehen, wie lebhaft der Wettbewerb in diesem Bereich ist. (Nachtrag) Es gibt auch ein Paper zum größeren Gesamtbild des Projekts „mirage“, das den Megakernel-Ansatz jedoch nicht behandelt: Paper-Link
Der Verfasser des Beitrags antwortet selbst. Ich stimme zu, dass die Forschung parallel zu Stanford stattfindet. Der Hauptunterschied ist, dass wir uns auf einen Compiler zur automatisierten Megakernel-Erzeugung konzentrieren.
Auch ThunderKittens von Hazy Research ist eine sehr coole Bibliothek. In letzter Zeit wird viel Aufwand in Formalisierung, Pipelining, Divide-and-Conquer, Effizienzmaximierung sowie die Entwicklung spezialisierter Compiler und DSLs gesteckt, um aktuelle NVIDIA-GPUs maximal auszunutzen.
Die Leistungszahlen für Qwen 8B wären, falls sie sich bestätigen, ziemlich beeindruckend. Es wirkt praktischer als frühere Megakernel-Ansätze. Diese Art von Kernel, bei der pro SM einer bestehen bleibt, erinnert an Larrabee. Ich frage mich, wie die Welt heute aussehen würde, wenn man statt des bestehenden CUDA-Wegs den traditionelleren Pfad Prozess-Thread-SIMD eingeschlagen hätte.
Die Idee, statt softwarebasierter Inferenz ein fest verdrahtetes LLM als reines ASIC zu bauen. Gibt es Kostenvorteile? Könnte man zusätzliche Schichten vorsehen, die softwareseitig noch angepasst oder feinjustiert werden können? Da wir faktisch fast an einem „gut genug“-Niveau angekommen sind, ist es denkbar, dass man sich in den nächsten 2 bis 4 Jahren dazu entschließt, spezialisierte Chips fest darauf auszulegen. Ich frage mich, ab welchem Punkt die Vorteile extrem spezialisierter Hardware wirklich zum Tragen kommen.