Kernel-Smith: A Unified Recipe for Evolutionary Kernel Optimization

Probleemstelling

Het genereren van hoogpresterende GPU-kernels (zoals voor Triton of CUDA) is cruciaal voor de efficiëntie van grote taalmodellen (LLMs) en wetenschappelijke toepassingen. Hoewel Large Language Models (LLMs) goed zijn in het genereren van code, kampen bestaande benaderingen met twee fundamentele beperkingen:

1. Beperkte exploratie: Bestaande systemen vertrouwen vaak op "one-shot" generatie of multi-turn conversaties die te snel vastlopen in lokale optimums. Ze missen de diversiteit die nodig is om de enorme zoekruimte van kernel-implementaties (zoals fusion-patronen, tiling-strategieën) effectief te verkennen.
2. Onbetrouwbare evaluatie: De evaluatie van kernelprestaties is gevoelig voor ruis (variabiliteit in uitvoeringstijden). Dit leidt tot het behouden van suboptimale kernels of het verwijderen van veelbelovende kandidaten.
3. Training mismatch: Modellen worden vaak getraind om één correcte en snelle kernel in één keer te genereren, in plaats van geoptimaliseerd te worden voor iteratieve verbetering binnen een zoekproces.

Methodologie: Kernel-Smith Framework

Kernel-Smith introduceert een unificerend framework dat een evolutionaire agent combineert met een evolutie-gerichte trainingsspecificatie.

1. Evolutionaire Agent Framework

In plaats van een enkele lijn van code te verfijnen, onderhoudt Kernel-Smith een populatie van uitvoerbare kandidaten.

• Zoekproces: Het systeem gebruikt een MAP-Elites-achtige structuur (een archief van top-performerende en diverse programma's).
• Stabiele Evaluatie: Om ruis te minimaliseren, gebruikt het een gespecialiseerde evaluatie-backend voor NVIDIA GPUs (Triton) en MetaX GPUs (MACA).
  • Stabilisatie: Warm-up runs, meervoudige metingen, uitsluiting van outliers en het gebruik van CUDAGraph zorgen voor een tijdsvariatie van minder dan 1%.
  • Hack-detectie: Het systeem detecteert en straft modellen die native PyTorch-oproepen gebruiken in plaats van de gegenereerde kernel (wat een valse snelheid van ~1x zou geven).
• Feedback: De agent ontvangt gestructureerde feedback over compilatie, correctheid en snelheidswinst, niet slechts een scalar beloning.

2. Training Recipe (Post-Training)

Het model wordt niet getraind als een eenmalige generator, maar als een sterke lokale verbeteraar binnen de evolutionaire lus.

• Data Synthese: Er wordt een dataset van 59k PyTorch-modules gegenereerd uit GitHub-repositories. Trajecten van evolutie worden gegenereerd met sterke "teacher" modellen.
• Stap-gerichte Supervisie: In plaats van het hele traject te imiteren, filtert het systeem alleen de hoogste-winst revisies (stappen die correctheid behouden en snelheid verbeteren).
  • SFT (Supervised Fine-Tuning): Gebruikt gefilterde data voor correctheid en initiële snelheidswinst.
  • RL (Reinforcement Learning): Gebruikt GRPO (Group Relative Policy Optimization) waarbij de snelheidswinst ten opzichte van de ouder-kernel als beloningssignaal dient. Het model leert hierdoor de "atomaire verbeteringen" die het meest bijdragen aan de uiteindelijke snelheid.

Belangrijkste Bijdragen

1. Unificatie van Zoek en Trainen: Een framework dat de stabiliteit van evolutionaire zoekalgoritmen koppelt aan een trainingsstrategie die specifiek is ontworpen voor iteratieve optimalisatie.
2. Robuuste Evaluatie-Backend: Een gedistribueerd evaluatiesysteem dat ruis onderdrukt en "hacking" (omzeilen van kernel-generatie) detecteert, werkend op zowel NVIDIA als MetaX hardware.
3. Stap-gerichte RL-strategie: Een innovatieve aanpak om lange evolutie-trajecten te comprimeren tot de meest informatieve verbeterstappen, waardoor het model beter generaliseert dan bij het trainen op volledige trajecten.
4. Heterogene Adaptatie: Het framework is platform-onafhankelijk ontworpen en wordt succesvol gedemonstreerd op zowel NVIDIA (Triton) als MetaX (MACA) backends.

Resultaten

De resultaten tonen aan dat Kernel-Smith state-of-the-art prestaties levert, zowel op benchmarks als in de praktijk.

• KernelBench (NVIDIA Backend):

  • Het model Kernel-Smith-235B-RL behaalde de beste gemiddelde snelheidswinst (3.70x) en presteerde beter dan geavanceerde gesloten modellen zoals Gemini-3.0-pro en Claude-4.6-opus.
  • Het model behield een hoge correctheid (96.33%) terwijl het de snelheidswinst maximaliseerde, vooral op moeilijke taken (Level 3).
  • De groeicurve van de beste scores toont aan dat het model effectiever gebruikmaakt van extra "test-time compute" dan concurrenten.

• MetaX Backend:

  • Kernel-Smith-MACA-30B versloeg grote concurrenten zoals DeepSeek-V3.2-think en Qwen3-235B-2507-think in het genereren van MACA-kernels, wat de schaalbaarheid naar verschillende hardware-platforms bewijst.

• Real-world Applicaties:

  • SGLang & LMDeploy: Het framework produceerde kernels die succesvol zijn gemerged in productiestacks (upstream pull requests), met meetbare snelheidswinsten in end-to-end inferentie (bijv. 1.36x tot 4.78x op kernel-niveau).
  • DeepSeek Engram: Op een nieuw, niet-benchmark PyTorch-module werd een 14.59x snelheidswinst behaald, wat aantoont dat het systeem generaliseert naar recente onderzoekscode.

Betekenis en Impact

Kernel-Smith bewijst dat LLM-gedreven kernel-optimalisatie kan evolueren van gecontroleerde benchmarks naar praktische, productieve implementaties. De kerninzichten zijn:

• Betrouwbare feedback is essentieel voor stabiele zoekdynamiek; zonder dit faalt de evolutie door ruis.
• Stap-gerichte training is superieur aan het imiteren van volledige trajecten; het leert het model om specifieke, hoogwaardige verbeteringen te maken in plaats van alleen het eindresultaat na te bootsen.
• Het framework biedt een blauwdruk voor het automatiseren van hardware-optimalisatie over diverse accelerators, wat essentieel is voor de toekomstige schaalbaarheid van AI-systemen en wetenschappelijke computing.

De paper concludeert dat de combinatie van een evolutionaire agent en een specifieke post-training strategie de drempel verlaagt voor het produceren van hoogwaardige, productieklaar GPU-code door AI.