🤖 machine learning

AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis

AutoMegaKernel ist ein statisch geprüfter Agenten-Harness, der automatisch einen einzelnen persistenten CUDA-Megakernel für HuggingFace Llama-Familienmodelle synthetisiert, wobei eine deadlockfreie Ausführung über diverse NVIDIA-Architekturen hinweg sichergestellt wird und durch selbstverbessernde Agenten-Schleifen eine bis zu 1,72-fache Beschleunigung gegenüber Baselines auf Inferenz-GPUs erreicht wird.

Ursprüngliche Autoren: Jaber Jaber, Osama Jaber

Veröffentlicht 2026-06-09

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Jaber Jaber, Osama Jaber

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie versuchen, eine riesige, komplexe Fabrik (ein Large Language Model) zu betreiben, um ein einziges Produkt (ein Wort Text) herzustellen.

Der alte Weg (Standard-KI):
Derzeit ist der Betrieb dieser Fabrik so, als würde man einen Arbeiter ins Lager schicken, um ein bestimmtes Werkzeug zu holen, es zurück zur Montagelinie zu bringen, ein winziges Stück Arbeit zu verrichten, den Arbeiter wieder zurück ins Lager zu schicken, um das nächste Werkzeug zu holen, und dies dutzende Male für jedes einzelne Wort zu wiederholen.

Das Problem: Der Arbeiter verbringt die meiste Zeit mit dem Hin- und Herlaufen (Warten auf Daten aus dem Speicher), anstatt zu bauen. Dies wird als „bandbreitenbeschränkt“ bezeichnet. Die Fabrik ist durch die Geschwindigkeit begrenzt, mit der der Arbeiter laufen kann, nicht durch die Geschwindigkeit, mit der er baut.
Die Kosten: Jedes Mal, wenn der Arbeiter stoppt, um ein Werkzeug zu holen, entsteht eine winzige Verzögerung (Launch-Latenz). Dieses ständige Anhalten bei jedem Schritt summiert sich und macht den Prozess langsam.

Der neue Weg (AutoMegaKernel):
Die Autoren dieses Papers haben ein System namens AutoMegaKernel (AMK) entwickelt. Stellen Sie sich AMK als einen revolutionären neuen Fabrikmanager vor, der die Regeln völlig ändert.

1. Die „One-Shot“-Fabrik

Anstatt Arbeiter immer wieder hin- und herzusenden, organisiert AMK die gesamte Fabrik so, dass ein einziger Start (Single Launch) die ganze Arbeit erledigt.

Die Analogie: Stellen Sie sich vor, anstatt dass ein Arbeiter nach dem anderen Werkzeuge holt, gibt es ein Team von Arbeitern, die an jeder einzelnen Maschine stationiert sind. Sie sind alle durch eine direkte, interne Pipeline miteinander verbunden. Wenn der „Start“-Knopf gedrückt wird, läuft die gesamte Montagelinie in einer kontinuierlichen Bewegung von Anfang bis Ende durch, ohne dass jemals jemand den Boden verlassen muss, um ins Lager zu gehen.
Das Ergebnis: Dies eliminiert die „Laufzeiten“ und die „Stopp-und-Go“-Verzögerungen, was den Prozess theoretisch viel schneller macht.

2. Der „Sicherheitsinspektor“ (Der statisch geprüfte Agent)

Eine Fabrik zu bauen, in der alle gleichzeitig arbeiten, ist gefährlich. Wenn die Arbeiter nicht perfekt koordiniert sind, könnten sie zusammenstoßen (Race Conditions) oder ewig auf ein Werkzeug warten, das nie ankommt (Deadlocks).

Die Innovation: Normalerweise benötigt man dafür einen menschlichen Experten, der den Code manuell schreibt, was schwierig und fehleranfällig ist. AMK nutzt einen KI-Agenten, um das Fabriklayout automatisch zu entwerfen.
Das Sicherheitsnetz: Bevor die Fabrik überhaupt in Betrieb genommen wird, prüft ein eingefrorener, unveränderlicher Sicherheitsinspektor (ein Validator) den Bauplan der KI. Er beweist mathematisch, dass:
- Niemand in einer Endlosschleife stecken bleibt.
- Nicht zwei Arbeiter gleichzeitig versuchen, dasselbe Werkzeug zu benutzen.
- Die Reihenfolge der Abläufe perfekt ist.
Die Behauptung: Die Autoren haben diesen Inspektor gegen 7.160 „tricky“ oder fehlerhafte Baupläne getestet. Der Inspektor hat jeden einzelnen unsicheren Plan erkannt und abgelehnt, bevor die Fabrik startete. Er hat niemals einen gefährlichen Plan durchgelassen.

3. Der „Universaladapter“ (Self-Retargeting)

Normalerweise funktioniert eine Fabrik, die für ein bestimmtes Gebäude (einen spezifischen Computerchip) entworfen wurde, nicht in einem anderen Gebäude.

Die Magie: AMK ist „selbst-retargetingfähig“. Sie schreiben den Bauplan einmal, und das System passt ihn automatisch an, um perfekt auf verschiedenen Arten von Computerchips (wie der RTX 5090, A100 oder H100) zu funktionieren, ohne dass ein Mensch den Code für jeden einzelnen neu schreiben muss.

4. Die „Leichtgewicht“-Version (Quantisierung)

Das System hat auch herausgefunden, wie man die Fabrik mit „leichteren“ Werkzeugen betreibt.

Die Analogie: Anstatt schwere, vollwertige Werkzeuge (hochpräzise Mathematik) zu tragen, nutzen die Arbeiter leichte, kompakte Werkzeuge (int8- oder int4-Präzision).
Der Vorteil: Weil die Werkzeuge leichter sind, können sie mehr davon gleichzeitig tragen, und die Arbeiter bewegen sich schneller.
- Int8: Funktioniert genauso gut wie die schweren Werkzeuge (verlustfrei), ist aber 12 % schneller.
- Int4: Bewegt sich noch schneller (reduziert die „Laufzeit“ um mehr als das Zweifache), aber das fertige Produkt ist vielleicht etwas weniger perfekt (verlustbehaftet), obwohl es immer noch lesbar ist.

5. Die ehrlichen Ergebnisse (Das „Real Talk“)

Die Autoren sind sehr transparent darüber, wo ihr System gewinnt und wo es verliert.

Wo es gewinnt: Auf „Inference-Class“-Chips (wie dem L4, L40S, A10G und der Consumer-RTX 5090) ist das AMK-System schneller als der aktuelle Industriestandard (cuBLAS), wenn es einzelne Wörter generiert. Es schlägt die Konkurrenz um das 1,1- bis 1,3-fache.
- Warum? Es hat erfolgreich die „Hin- und Herlauf“-Verzögerungen eliminiert, die andere Systeme plagen.
Wo es verliert: Auf den massiven, Hochgeschwindigkeits-„Training-Class“-Chips (wie dem A100 und H100) ist die aktuelle Version langsamer als das Standardsystem.
- Warum? Die „Sicherheitskoordination“ zwischen den Arbeitern (Synchronisation) nimmt eine winzige Zeit in Anspruch. Auf den schnellsten Chips wird diese winzige Verzögerung zum Flaschenhals. Die Autoren geben dies offen zu: „Wir haben den Flaschenhals lokalisiert, und wir sind ehrlich darin.“

Zusammenfassung

AutoMegaKernel ist ein System, das einen KI-Agenten nutzt, um eine „One-Shot“-Fabrik für das Ausführen von KI-Modellen zu entwerfen. Es beinhaltet einen strengen Sicherheitsinspektor, um sicherzustellen, dass das Design niemals abstürzt. Es läuft erfolgreich auf vielen verschiedenen Computerchips automatisch.

Der große Gewinn: Es ist derzeit die schnellste Methode, um Text Wort für Wort auf vielen beliebten Consumer- und Datacenter-Chips zu generieren, indem es die Standardwerkzeuge durch das Entfernen unnötiger Verzögerungen übertrifft.
Der Haken: Es ist noch nicht perfekt. Auf den absolut schnellsten Supercomputer-Chips macht der Koordinationsaufwand die Synchronisation immer noch etwas langsamer als die etablierten Giganten, aber die Autoren haben bewiesen, dass das System sicher, korrekt und fähig zur Selbstverbesserung ist.

Technisches Resümee: AutoMegaKernel (AMK)

Problemstellung

Die autoregressive Dekodierung mit Single-Stream bei einer Batch-Größe von 1 ist fundamental bandbreitenbeschränkt. Die Latenzuntergrenze wird durch das Verhältnis der gesamten Gewichtsbahyte zur HBM-Bandbreite bestimmt ( $t_{min} = \text{gewicht\_bytes} / \text{HBM\_bandwidth}$ ). Konventionelle Ausführungsframeworks (z. B. PyTorch, cuBLAS) erreichen diese Untergrenze nicht, da sie für jeden Operator einen separaten Kernel starten, was eine signifikante CPU-Launch-Latenz verursacht und dazu führt, dass Aktivierungen zwischen den Operationen in den HBM zurückgeschrieben werden müssen. Während CUDA Graphs die Launch-Overheads amortisieren, behalten sie dennoch die Inter-Op-Kernel-Grenzen und HBM-Round-Trips bei.

Existierende „Megakernel“-Ansätze (z. B. Mirage Persistent Kernel) oder handoptimierte Kernel adressieren die Latenzuntergrenze, indem sie den gesamten Forward-Pass in einen einzigen persistenten Kernel fusionieren. Diese Systeme mangelt es jedoch an Vertrauen (Trust) und Portabilität: Sie lassen oft statische Sicherheitskontrollen aus (was zu Deadlocks oder Race Conditions führen kann), zielen manuell auf spezifische Modelle/Architekturen ab und bieten keine strukturierte Schnittstelle, damit autonome Agenten das System sicher optimieren können.

Methodik

AutoMegaKernel (AMK) ist ein autonomes, agentengesteuertes System, das HuggingFace Llama-Familienmodelle ohne handgeschriebenes CUDA pro Modell in einen einzigen persistenten kooperativen Kernel kompiliert. Das System ist in vier Schichten und zwei Autoresearch-Schleifen strukturiert:

Layer 0 (Die VM): Ein handgeschriebener, eingefrorener, pro-Architektur-persistenter Megakernel (vm/scheduler.cu). Er startet einen kooperativen Threadblock pro Streaming Multiprocessor (SM). Jeder Block besitzt eine lokale Warteschlange und durchläuft eine globale topologische Ordnung von Aufgaben, wobei er wait → dispatch → signal-Instruktionen ausführt. Er stützt sich auf monotone Zähler zur Synchronisation und vermeidet Locks.
Layer 1 (Instruktionen): ABI-konforme Mikro-Kernel (GEMV, Attention, RMSNorm usw.), die isoliert auf Korrektheit geprüft wurden.
Layer 2 (Scheduler): Die agentenbasierte Editieroberfläche. Er importiert ein Modell, transformiert es in eine typisierte Schedule IR (einen durch monotone Zähler synchronisierten Task-DAG) und schlägt strukturelle Konfigurationen (Tiling, Fusion, SM-Zuweisung, Pipelining-Tiefe) via eines ScheduleConfig-Objekts vor.
Layer 3 (Roadmap): Platzhalter für zukünftige Dynamik (z. B. Continuous Batching).

Das Sicherheitsgate (Safety Gate):
Die zentrale Innovation ist ein statisch geprüfter Validator (validate() in schedule/ir.py). Vor jedem GPU-Launch führt der Validator eine Reihe von statischen Graph-Checks auf die Schedule IR durch, um zu zertifizieren:

Deadlock-Freiheit: Sicherstellung, dass der Producer-Consumer-DAG azyklisch ist, die Wait-Schwellen erfüllbar sind ( $1 \le t \le \#producers$ ) und die Per-SM-Warteschlangen-Ordnungen gültig sind.
Race-Freiheit: Sicherstellung, dass geteilte Zähler korrekt zusammengeführt werden (alle Producer müssen abgeschlossen sein, bevor ein Consumer auf einen Multi-Producer-Zähler wartet) und dass Speicherzugriffe die Happens-before-Beziehungen respektieren.
Korrektheit: Verifizierung der Output-Erreichbarkeit und der referenziellen Integrität.

Wenn ein Schedule diese Prüfungen fehlschlägt, wird er sofort abgelehnt. Dies verhindert, dass unsichere, von Agenten vorgeschlagene Schedules die GPUs zum Einfrieren bringen.

Autoresearch-Schleife:
AMK enthält eine unüberwachte, agentensteuerbare Schleife, die iterativ Änderungen an der ScheduleConfig vorschlägt, diese validiert, gegen eine CPU-Referenz auf Korrektheit prüft, die Latenz misst und Änderungen entweder beibehält oder rückgängig macht. Diese Schleife verbessert den Megakernel autonom gegenüber seinem eigenen Baseline-Modell.

Wichtigste Beiträge

Statisch geprüfte Schedule IR: Ein eingefrorener Validator, der Deadlock- und Race-Freiheit mittels statischer Graph-Checks (statt mechanisierter Beweise) über einer vertrauenswürdigen Basis zertifiziert. Er erreichte null False-Accepts bei 7.160 adversariell konstruierten Schedules (davon 6.091 unsicher), während alle 360 realen Modell-Lowerings akzeptiert wurden.
Automatische Generierungsabdeckung: AMK generiert automatisch korrekte Megakernel für 10 von 10 unterstützten Modellen (einschließlich realer Checkpoints wie SmolLM2-135M und TinyLlama-1.1B) ohne handgeschriebenes CUDA. Es lehnt nicht unterstützte Varianten (z. B. biased Projections) erfolgreich ab oder dokumentiert Blind Spots ehrlich.
Selbst-Retargeting: Eine einzige Quellcodebasis generiert korrekte Megakernel für sm_80 (A100), sm_90 (H100) und sm_120 (RTX 5090), wobei die Eager-Execution bitweise oder innerhalb enger fp32-Toleranzen ( $\le 3.8 \times 10^{-5}$ ) exakt übereinstimmt.
Quantisierte Megakernel: Das System generiert automatisch int8 (greedy-lossless, 1.12× schneller) und int4 (lossy, 2.42× reduzierter Gewichtstraffik) Weight-only-quantisierte Kernel aus demselben IR-Pfad.
Agentensteuerbarer Compiler: Das System behandelt Korrektheit als strukturelle Invariante, was es einem Agenten ermöglicht, den Schedule-Raum sicher zu durchsuchen, ohne das Risiko von Runtime-Hangs einzugehen.
Ehrliche Performance-Studie: Die Arbeit berichtet Ergebnisse über Inference-Class- und Training-Class-GPUs hinweg und legt explizit offen, wo das System gewinnt und wo es etablierten Baselines unterliegt.

Experimentelle Ergebnisse

Korrektheit & Portabilität:

Auf der RTX 5090, A100 und H100 reproduziert AMK das HuggingFace Greedy-Decoding Token für Token und stimmt mit der Teacher-Forced Perplexity bis zu $2.5 \times 10^{-7}$ auf dem SmolLM2-135M Checkpoint überein.
Der Validator lehnte alle 6.091 unsicheren Schedules im Stresstest ab, was eine False-Accept-Rate von 0 % demonstriert.

Performance auf Inference-Class GPUs (Der Gewinn):

RTX 5090 (Consumer): Ein auto-getunter int8 Weight-only (W8A16) Megakernel schlägt CUDA-graphed cuBLAS (bf16) bei Batch-1 Decode um 1.19–1.23×. Dies ist ein präzisionsasymmetrischer Vergleich (int8 vs. bf16), bei dem der Gewinn durch das Streaming weniger Gewichtsbahyte (0.61× die Bytes) erzielt wird, nicht durch einen schnelleren Per-Byte-Kernel.
Datacenter Inference Fleet: Der int8 Megakernel schlägt CUDA-graphed cuBLAS bf16 über die gesamte Inference-Class-Flotte:
- NVIDIA L4: Bis zu 1.33× schneller (steigend mit der Modellgröße).
- L40S: 1.25–1.27× schneller.
- A10G: Bis zu 1.08× schneller bei Skalierung.
Selbstverbesserung: Die Autoresearch-Schleife verbesserte den Megakernel autonom gegenüber seinem Standard um das 1.25–1.72× auf der RTX 5090.

Performance auf Training-Class GPUs (Der Verlust):

Auf High-Bandwidth Training-Class GPUs (A100, H100) liegt der auto-generierte Megakernel hinter CUDA-graphed cuBLAS und vLLM zurück.
- Auf der H100 war AMK 1.65× langsamer als der Standard-CUDA-Graph-Pfad von vLLM.
- Auf der A100 erreichte der optimierte bf16 GEMV nur ~63 % des gemessenen HBM-Peaks (gegenüber ~90 % bei cuBLAS).
Diagnose: Der Flaschenhert ist der Per-Tile Cross-SM Synchronisationsaufwand (ein fixer Kostenfaktor, den der Megakernel zahlt, cuBLAS jedoch nicht). Auf High-Bandwidth-Silizium dominiert dieser Synchronisations-Overhead, während auf Inference-Class GPUs die Bandbreiteneinsparungen durch Fusion und Quantisierung diesen Aufwand amortisieren.

Bedeutung und Ansprüche

Die Arbeit behauptet, dass AutoMegaKernel einen Wechsel von „Rohgeschwindigkeit“ hin zu systematischer Sicherheit und Portabilität in der Megakernel-Synthese darstellt. Ihr primärer Beitrag ist eine statisch geprüfte Compiler-Architektur, die es autonomen Agenten ermöglicht, den Raum ganzheitlicher Modell-Schedules sicher zu explorieren.

Vertrauen: Es ersetzt das Risiko von hängenden GPUs durch ein statisches Sicherheitsgate, das unsichere Schedules vor dem Launch ablehnt.
Generalität: Es erreicht Selbst-Retargeting über Architekturen (sm_80/90/120) und Modellfamilien hinweg ohne manuelles CUDA-Tuning.
Ehrlichkeit: Die Autoren fassen die Performance-Ergebnisse explizit als regimebedingt zusammen. Das System gewinnt auf Inference-Class GPUs durch die Fusion von Operationen und die Nutzung von Quantisierung zur Reduktion des Speicherverkehrs, verliert aber derzeit auf Training-Class GPUs aufgrund von Synchronisations-Overheads, die eine gröbere Skalierung (Future Work) statt nur besserer Mikro-Kernel erfordern.

Die Arbeit schließt mit der Feststellung, dass die aktuelle Kernel-Qualität die handoptimierten Baselines auf High-End-Training-Hardware zwar noch nicht erreicht, das statisch geprüfte Harness jedoch das notwendige Fundament bildet, um die zukünftige Automatisierung (z. B. gröbere Synchronisation, tieferes Pipelining) sicher zu gestalten, die erforderlich ist, um diese Lücke zu schließen.