🤖 machine learning

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

AutoMegaKernel is een statisch gecontroleerde agent-harness die automatisch een enkele persistente CUDA-megakernel synthetiseert voor HuggingFace Llama-familie modellen, waarbij deadlock-vrije executie over diverse NVIDIA-architecturen wordt gegarandeerd terwijl een versnelling van tot 1,72x ten opzichte van baselines op inference-klasse GPU's wordt bereikt via zelfverbeterende agent-loops.

Oorspronkelijke auteurs: Jaber Jaber, Osama Jaber

Gepubliceerd 2026-06-09

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Jaber Jaber, Osama Jaber

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een enorme, complexe fabriek probeert te runnen (een Large Language Model) om één enkel product te produceren (één woord tekst).

De Oude Manier (Standaard AI):
Het draaien van deze fabriek is momenteel alsof je een arbeider naar het magazijn stuurt om een specifiek gereedschap te pakken, hem terug te brengen naar de assemblagelijn, een klein beetje werk te doen, en de arbeider weer terug naar het magazijn te sturen voor het volgende gereedschap, en dit tientallen keren te herhalen voor elk woord.

Het Probleem: De arbeider brengt de meeste tijd door met heen en weer lopen (wachten op data uit het geheugen) in plaats van met bouwen. Dit wordt "bandwidth-bound" genoemd. De fabriek wordt beperkt door hoe snel de arbeider kan lopen, niet door hoe snel hij kan bouwen.
De Kosten: Elke keer dat de arbeider stopt om gereedschap te pakken, is er een kleine vertraging (launch latency). Het doen van dit proces voor elke stap zorgt ervoor dat het proces traag wordt.

De Nieuwe Manier (AutoMegaKernel):
De auteurs van dit artikel hebben een systeem gebouwd genaamd AutoMegaKernel (AMK). Zie AMK als een revolutionaire nieuwe fabrieksmanager die de regels volledig verandert.

1. De "One-Shot" Fabriek

In plaats van arbeiders heen en weer te sturen, organiseert AMK de hele fabriek zo dat één enkele lancering het hele werk doet.

De Analogie: Stel je voor dat in plaats van een arbeider die één voor één gereedschap haalt, je een team van arbeiders hebt die bij elke machine gepositioneerd zijn. Ze zijn allemaal verbonden door een directe, interne pijplijn. Wanneer de "Start"-knop wordt ingedrukt, draait de hele assemblagelijn van begin tot eind in één continue beweging, zonder dat iemand ooit de vloer verlaat om naar het magazijn te gaan.
Het Resultaat: Dit elimineert de "looptijd" en de "stop-en-start"-vertragingen, wat het proces theoretisch veel sneller maakt.

2. De "Veiligheidsinspecteur" (De Statisch Gecontroleerde Agent)

Een fabriek bouwen waar iedereen tegelijkertijd werkt, is gevaarlijk. Als de arbeiders niet perfect gecoördineerd zijn, kunnen ze tegen elkaar opbotsen (race conditions) of eeuwig wachten op een gereedschap dat nooit aankomt (deadlocks).

De Innovatie: Normaal gesproken heb je hiervoor een menselijke expert nodig om de code handmatig te schrijven, wat moeilijk en foutgevoelig is. AMK gebruikt een AI Agent om het fabrieksontwerp automatisch te ontwerpen.
Het Veiligheidsnet: Voordat de fabriek überhaupt aangaat, controleert een bevroren, onveranderlijke Veiligheidsinspecteur (een validator) het blauwdruk van de AI. Deze bewijst wiskundig dat:
- Niemand eeuwig zal wachten.
- Geen twee arbeiders tegelijkertijd hetzelfde gereedschap zullen gebruiken.
- De volgorde van operaties perfect is.
De Claim: De auteurs hebben deze inspecteur getest tegen 7.160 "lastige" of defecte blauwdrukken. De inspecteur ontdekte elke onveilige blauwdruk en wees deze af voordat de fabriek begon. Hij heeft nooit een gevaarlijk plan doorgelaten.

3. De "Universele Adapter" (Zelf-Retargeting)

Meestal werkt een fabriek die ontworpen is voor een specifiek gebouw (een specifieke computerchip) niet in een ander gebouw.

De Magie: AMK is "zelf-retargeting". Je schrijft de blauwdruk één keer, en het systeem past deze automatisch aan om perfect te werken op verschillende soorten computerchips (zoals de RTX 5090, A100 of H100) zonder dat een mens de code voor elk exemplaar opnieuw hoeft te schrijven.

4. De "Lichtgewicht" Versie (Quantization)

Het systeem heeft ook uitgevogeld hoe het de fabriek met "lichtere" gereedschappen kan laten draaien.

De Analogie: In plaats van zware, volledige gereedschappen te dragen (hoge precisie wiskunde), gebruiken de arbeiders lichtgewicht, compacte gereedschappen (int8 of int4 precisie).
Het Voordeel: Omdat de gereedschappen lichter zijn, kunnen ze er meer tegelijk dragen, en bewegen de arbeiders sneller.
- Int8: Werkt net zo goed als de zware gereedschappen (lossless), maar is 12% sneller.
- Int4: Beweegt nog sneller (het vermindert de "looptijd" met meer dan 2x), maar het eindproduct is misschien iets minder perfect (lossy), hoewel het nog steeds leesbaar is.

5. De Eerlijke Resultaten (De "Real Talk")

De auteurs zijn zeer transparant over waar dit systeem wel en niet wint.

Waar het wint: Op "Inference-Class" chips (zoals de L4, L40S, A10G en consumenten-RTX 5090) is het AMK-systeem sneller dan de huidige industriestandaard (cuBLAS) bij het genereren van één woord per keer. Het verslaat de concurrentie met 1.1x tot 1.3x.
- Waarom? Het heeft succesvol de "heen en weer loop"-vertragingen geëlimineerd die andere systemen teisteren.
Waar het verliest: Op de enorme, hogesnelheid "Training-Class" chips (zoals de A100 en H100) is de huidige versie trager dan het standaard systeem.
- Waarom? De "veiligheidscoördinatie" tussen de arbeiders (synchronisatie) kost een klein beetje tijd. Op de snelste chips wordt deze kleine vertraging de bottleneck. De auteurs geven dit openlijk toe: "We hebben de bottleneck gelokaliseerd, en we zijn eerlijk over het feit."

Samenvatting

AutoMegaKernel is een systeem dat een AI-agent gebruikt om een "one-shot" fabriek te ontwerpen voor het draaien van AI-modellen. Het bevat een strikte veiligheidsinspecteur om te garanderen dat het ontwerp nooit crasht. Het slaagt erin om automatisch op veel verschillende computerchips te draaien.

De Grote Overwinning: Het is momenteel de snelste manier om tekst één woord per keer te genereren op veel populaire consumenten- en datacenterchips, waarbij het de standaardtools verslaat door onnodige vertragingen te verwijderen.
De Catch: Het is nog niet perfect. Op de absoluut snelste supercomputerchips maakt de coördinatie-overhead het systeem nog steeds iets trager dan de gevestigde reuzen, maar de auteurs hebben bewezen dat het systeem veilig, correct en in staat is tot zelfverbetering is.

Technische Samenvatting: AutoMegaKernel (AMK)

Probleemstelling

Single-stream autoregressieve decodering bij een batchgrootte van 1 is fundamenteel bandbreedtegebonden. De latentie-ondergrens wordt bepaald door de ratio tussen het totaal aantal gewichtbytes en de HBM-bandbreedte ( $t_{min} = \text{gewicht\_bytes} / \text{HBM\_bandbreedte}$ ). Conventionele executie-frameworks (bijv. PyTorch, cuBLAS) slagen er niet in deze ondergrens te benaderen omdat ze voor elke operator een aparte kernel lanceren, wat leidt tot aanzienlijke CPU-launch-latentie en ervoor zorgt dat activaties tussen operaties door naar de HBM moeten rondreizen. Hoewel CUDA Graphs de launch-overhead amenuiseren, behouden ze de inter-op kernelgrenzen en HBM-rondreizen.

Bestaande "megakernel"-benaderingen (bijv. Mirage Persistent Kernel) of handmatig geoptimaliseerde kernels pakken de latentie-ondergrens aan door de volledige forward pass te fuseren in één enkele persistente kernel. Deze systemen missen echter vertrouwen en portabiliteit: ze laten vaak statische veiligheidscontroles weg (met risico op deadlocks of racecondities), richten zich handmatig op specifieke modellen/architecturen en bieden geen gestructureerde interface voor autonome agenten om het systeem veilig te optimaliseren.

Methodologie

AutoMegaKernel (AMK) is een autonoom, door agenten aangestuurd systeem dat HuggingFace Llama-familie modellen compileert naar een enkele persistente coöperatieve kernel zonder handgeschreven CUDA per model. Het systeem is gestructureerd in vier lagen en twee autoresearch-loops:

Laag 0 (De VM): Een handgeschreven, bevroren, per-architectuur persistente megakernel (vm/scheduler.cu). Deze lanceert één coöperatieve threadblock per Streaming Multiprocessor (SM). Elke block bezit een lokale wachtrij en doorloopt een globale topologische volgorde van taken, waarbij wait → dispatch → signal instructies worden uitgevoerd. Het vertrouwt op monotone tellers voor synchronisatie, waardoor locks worden vermeden.
Laag 1 (Instructies): ABI-conforme micro-kernels (GEMV, attention, RMSNorm, etc.) die in isolatie op correctheid zijn gecontroleerd.
Laag 2 (Scheduler): Het door de agent toegankelijke bewerkingsvlak. Het importeert een model, verlaagt dit naar een getypeerde Schedule IR (een taak-DAG gesynchroniseerd door monotone tellers), en stelt structurele configuraties voor (tiling, fusie, SM-toewijzing, pipelining-diepte) via een ScheduleConfig object.
Laag 3 (Roadmap): Placeholder voor toekomstige dynamiek (bijv. continuous batching).

De Safety Gate (Veiligheidshek):
De kerninnovatie is een statisch gecontroleerde validator (validate() in schedule/ir.py). Voordat een GPU-launch plaatsvindt, voert de validator een reeks statische graafcontroles uit op de Schedule IR om te certificeren:

Deadlock-vrijheid: Het waarborgen dat de producer-consumer DAG acyclisch is, de wait-thresholds voldoen aan ( $1 \le t \le \#producers$ ), en de per-SM wachtrijvolgordes geldig zijn.
Race-vrijheid: Het waarborgen dat gedeelde tellers correct worden samengevoegd (alle producers moeten klaar zijn voordat een consumer wacht op een multi-producer teller) en dat geheugentoegang de happens-before relaties respecteert.
Correctheid: Het verifiëren van output-bereikbaarheid en referentiële integriteit.

Als een schema faalt voor deze controles, wordt het onmiddellijk afgewezen. Dit voorkomt dat onveilige door agenten voorgestelde schema's hangende GPU's veroorzaken.

Autoresearch Loop:
AMK bevat een onbeheerde, door agenten aanstuurbare loop die iteratief wijzigingen in ScheduleConfig voorstelt, deze valideert, correctheidscontroles uitvoert tegen een CPU-referentie, latentie meet, en wijzigingen behoudt of terugdraait. Deze loop verbetert autonoom de megakernel ten opzichte van de eigen baseline.

Belangrijkste Bijdragen

Statisch Gecontroleerde Schedule IR: Een bevroren validator die deadlock- en race-vrijheid certificeert via statische graafcontroles (niet mechanische bewijzen) over een vertrouwde basis. Het bereikte nul false-accepts over 7.160 adversariële geconstrueerde schema's (waarvan 6.091 onveilig) terwijl het alle 360 echte model-lowering processen accepteerde.
Automatische Generatie Dekking: AMK genereert automatisch correcte megakernels voor 10 van de 10 ondersteunde modellen (inclusief echte checkpoints zoals SmolLM2-135M en TinyLlama-1.1B) zonder handgeschreven CUDA. Het wijst onondersteunde varianten (bijv. biased projections) succesvol af of documenteert blind spots eerlijk.
Zelf-Retargeting: Eén enkele broncode genereert correcte megakernels voor sm_80 (A100), sm_90 (H100), en sm_120 (RTX 5090), waarbij de eager execution bit-voor-bit of binnen nauwe fp32 toleranties ( $\le 3.8 \times 10^{-5}$ ) wordt benaderd.
Gekwantiseerde Megakernels: Het systeem genereert automatisch int8 (greedy-lossless, 1.12× sneller) en int4 (lossy, 2.42× minder gewichtverkeer) weight-only gekwantiseerde kernels vanuit hetzelfde IR-pad.
Agent-Drivable Compiler: Het systeem behandelt correctheid als een structurele invariant, waardoor een agent veilig de schema-ruimte kan doorzoeken zonder het risico op runtime hangs.
Eerlijke Performance Studie: Het paper rapporteert resultaten over zowel inference-klasse als training-klasse GPU's, waarbij expliciet wordt vermeld waar het systeem wint of achterblijft bij gevestigde baselines.

Experimentele Resultaten

Correctheid & Portabiliteit:

Op de RTX 5090, A100 en H100 reproduceert AMK de HuggingFace greedy decode token-voor-token en komt de teacher-forced perplexity overeen tot $2.5 \times 10^{-7}$ op de SmolLM2-135M checkpoint.
De validator wees alle 6.091 onveilige schema's af tijdens de stress-test, wat een 0% false-accept rate demonstreerde.

Prestaties op Inference-Class GPU's (De Winst):

RTX 5090 (Consumer): Een auto-getunede int8 weight-only (W8A16) megakernel verslaat CUDA-graphed cuBLAS (bf16) met 1.19–1.23× bij batch-1 decode. Dit is een precisie-asymmetrische vergelijking (int8 vs. bf16) waarbij de winst voortkomt uit het streamen van minder gewichtbytes (0.61× de bytes), niet uit een snellere per-byte kernel.
Datacenter Inference Fleet: De int8 megakernel verslaat CUDA-graphed cuBLAS bf16 over de gehele inference-class fleet:
- NVIDIA L4: Tot 1.33× sneller (groeiend met modelgrootte).
- L40S: 1.25–1.27× sneller.
- A10G: Tot 1.08× sneller op schaal.
Zelfverbetering: De autoresearch loop verbeterde de megakernel autonoom ten opzichte van de standaardversie met 1.25–1.72× op de RTX 5090.

Prestaties op Training-Class GPU's (Het Verlies):

Op high-bandwidth training-class GPU's (A100, H100) blijft de auto-gegenereerde megakernel achter bij CUDA-graphed cuBLAS en vLLM.
- Op de H100 was AMK 1.65× trager dan het standaard CUDA-graph pad van vLLM.
- Op de A100 bereikte de geoptimaliseerde bf16 GEMV slechts ~63% van de gemeten HBM-piek (vs. ~90% voor cuBLAS).
Diagnose: De bottleneck is de per-tile cross-SM synchronisatiekosten (een vaste kostenpost die de megakernel betaalt, maar cuBLAS niet). Op high-bandwidth silicon domineert deze synchronisatie-overhead, terwijl op inference-class GPU's de bandbreedtebesparingen door fusie en kwantisatie deze kosten compenseren.

Betekenis en Claims

Het paper claimt dat AutoMegaKernel een verschuiving vertegenwoordigt van "ruwe snelheid" naar systematische veiligheid en portabiliteit in megakernel-synthese. De primaire bijdrage is een statisch gecontroleerde compiler-architectuur die autonome agenten in staat stelt om veilig de ruimte van volledige model-schema's te verkennen.

Vertrouwen: Het vervangt het risico op hangende GPU's door een statische veiligheidscontrole die onveilige schema's afwijst vóór de lancering.
Generaliteit: Het bereikt zelf-retargeting over architecturen (sm_80/90/120) en modelfamilies zonder handmatige CUDA-tuning.
Eerlijkheid: De auteurs kaderen de prestatieresultaten expliciet als regime-afhankelijk. Het systeem wint op inference-class GPU's door operaties te fuseren en kwantisatie te gebruiken om het geheugentransport te verminderen, maar verliest momenteel op training-class GPU's door synchronisatie-overheads die grovere-grained scheduling vereisen (toekomstig werk) in plaats van alleen betere micro-kernels.

Het paper concludeert dat hoewel de huidige kwaliteit van de kernels nog niet overeenkomt met hand-geoptimaliseerde baselines op high-end hardware voor training, de statisch gecontroleerde harness de noodzakelijke fundering biedt om de toekomstige verbeteringen (bijv. grovere synchronisatie, diepere pipelining) die nodig zijn om dit gat te dichten, veilig te automatiseren.