OptiML: An End-to-End Framework for Program Synthesis and CUDA Kernel Optimization

Probleemstelling

Het genereren van hoogwaardige, prestatie-geoptimaliseerde CUDA-kernels blijft een grote uitdaging. Hoewel Large Language Models (LLM's) vaak functioneel correcte CUDA-code kunnen synthetiseren, missen deze modellen vaak de diepgaande expertise die nodig is voor prestatieoptimalisatie. De kloof tussen functioneel correcte code en prestatie-geoptimaliseerde code is groot, omdat kleine beslissingen in de code (zoals toegangspatronen, hergebruikstrategieën en instructiemix) de kernel kunnen verschuiven van reken-gebonden naar geheugen-gebonden en vice versa.

Bestaande benaderingen behandelen code-generatie en optimalisatie vaak als gescheiden problemen:

1. Generatie: LLM's schrijven code, maar deze is vaak niet geoptimaliseerd, compileert niet, of voldoet niet aan numerieke constraints.
2. Optimalisatie: Bestaande kernels worden geoptimaliseerd via autotuning of heuristieken, maar dit vereist vaak een goed startpunt.

De uitdagingen bij het optimaliseren van GPU-kernels zijn:

• Combinatorische zoekruimte: Veranderingen zoals tiling, vectorisatie en loop-herstructurering interageren complex met registerdruk, bezetting (occupancy) en geheugencoalescing.
• Duur en ruis: Prestatie-feedback is duur (profiling) en ruisgevoelig. Het vertrouwen op eind-timing alleen maakt zoekprocessen fragiel.
• Fragiliteit van naieve strategieën: Simpele methoden zoals willekeurige perturbaties of vaste regels convergeren vaak traag of generaliseren slecht.

Methodologie: Het OptiML Framework

OptiML is een end-to-end framework dat programmasynthese en kernel-optimalisatie unificeert. Het werkt in twee ontkoppelde fasen en kan ingang vinden via natuurlijke taal (intentie) of bestaande CUDA-code.

1. OptiML-G: Mixture-of-Thoughts Generator (Synthese)

Wanneer de input een natuurlijke taalbeschrijving is, fungeert OptiML-G als een voorstelbeleid (proposal policy) om een eerste uitvoerbaar programma te genereren.

• Architectuur: Het gebruikt een "Mixture-of-Thoughts" (MoT) aanpak. Er is een pool van meerdere bevroren, voorgetrainde code-generatie-experts (bijv. Qwen2.5-Coder, HPC-Coder-V2, StarCoder2).
• Routing: Een lichte router selecteert de top-K experts relevant voor de input.
• Latente Samenwerking: In tegenstelling tot traditionele ensembling, laten deze experts in OptiML-G op meerdere diepten in hun transformer-architectuur samenwerken in een gedeelde latente ruimte via kruis-attention (cross-attention). De primaire expert integreert "gedachten" (verborgen representaties) van de andere experts.
• Doel: Het genereren van een structuursterke, compilatie-veilige kernel die al goede kansen biedt voor latere optimalisatie, in plaats van een willekeurige start.

2. OptiML-X: Zoekgebaseerde Optimizer (Optimalisatie)

OptiML-X optimaliseert zowel de gegenereerde kernels als door gebruikers aangeleverde code. Het modelleert optimalisatie als een zoekproces onder verificatie.

• Monte Carlo Tree Search (MCTS): OptiML-X gebruikt MCTS om multi-stap optimalisatietrajecten te verkennen. Dit is cruciaal omdat tussentijdse stappen de prestaties tijdelijk kunnen verslechteren om later grotere winsten te behalen.
• LLM-as-a-Judge: Een LLM fungeert als rechter die:
  • Bottlenecks diagnoseert op basis van profiler-data (Nsight Compute).
  • Concrete code-patches (edits) voorstelt (bijv. tiling, shared-memory staging).
  • Beoordeelt of een voorgestelde wijziging plausibel is en de diagnose ondersteunt.
• Hardware-bewuste Beloning (Reward Model): De zoektocht wordt geleid door een samengestelde beloning die bestaat uit:
  • Tijdsbeloning: Verbetering in runtime.
  • Proxy-metrics: Gebaseerd op Nsight Compute data (bijv. SM/DRAM/TEX utilization, L1-sectoren, instructie-aantal).
  • Guardrails: Straffe straffen voor fouten (compilatiefouten, regressies in werklast, registers die de limiet overschrijden).
  • Bottleneck-awareness: De beloning weegt metrics dynamisch af afhankelijk van of de kernel geheugen-gebonden of reken-gebonden is.
• Validatie: Elke kandidaat wordt gecompileerd, getest (L0/L1 correctheidstests) en geprofild. Alleen kandidaten die aan alle constraints voldoen, worden in de zoekboom opgenomen.

Belangrijkste Bijdragen

1. Unificatie van Synthese en Optimalisatie: OptiML is het eerste framework dat intentie-gebaseerde synthese en post-hoc optimalisatie in één end-to-end pipeline combineert, waarbij de output direct een geoptimaliseerde CUDA-kernel is.
2. Mixture-of-Thoughts Generator (OptiML-G): Een nieuwe architectuur die latent-space samenwerking tussen diverse code-experts mogelijk maakt om robuustere startkernels te genereren dan enkelvoudige LLM's.
3. Search-Based Optimizer (OptiML-X): Een MCTS-gebaseerde optimizer die gebruikmaakt van een "LLM-as-a-Judge" en hardware-proxy metrics om gerichte, interpreteerbare optimalisaties uit te voeren die specifiek de gemeten bottlenecks aanpakken.
4. Interpretabele Evaluatie: Het systeem biedt niet alleen snellere code, maar verklaart waarom de code sneller is (bijv. "reductie van geheugentransacties" of "verbeterde SM-utilisatie") door middel van profiler-data.

Resultaten

De evaluatie is uitgevoerd op een diverse set van CUDA-kernels (o.a. Matrix Multiplication, Max Pooling 3D, Multi-Head Self-Attention, ReLU) op een NVIDIA A100 GPU.

• Prestatieverbetering: OptiML overtreft consistent sterke LLM-baselines (zoals GPT-5, Qwen2.5-Coder, HPC-Coder-V2) en varianten die alleen optimalisatie toepassen.
  • Voor Matrix Multiplication bereikte OptiML een snelheidswinst van 1.64x ten opzichte van de beste LLM-baseline.
  • Voor ReLU was de winst 1.67x.
  • Voor Multi-Head Self-Attention (waar baselines vaak faalden) bereikte OptiML een snelheidswinst van 1.15x.
• Compilatie en Correctheid: Veel LLM-only baselines faalden bij het compileren (CF - Compile Failed) voor complexe kernels. OptiML-G leverde een veel hogere "Pass@k" (kans op correcte code) en OptiML-X zorgde ervoor dat zelfs moeilijke kernels (zoals Max Pooling 3D) succesvol werden geoptimaliseerd.
• Efficiëntie: OptiML bereikt de beste resultaten in minder zoekrondes dan OptiML-X alleen, omdat OptiML-G een sterkere startpositie biedt. Dit vermindert de noodzaak om tijd te besteden aan het repareren van basisfouten.
• Kwalitatieve Analyse: De resultaten tonen aan dat OptiML niet alleen sneller is, maar ook de juiste hardware-metrics verbetert (hoger SM/TEX utilization, lager aantal L1-sectoren en instructies), wat aantoont dat de optimalisatie fundamenteel is en niet toevallig.

Betekenis en Impact

OptiML vertegenwoordigt een verschuiving in de manier waarop accelerator-programmering wordt benaderd. In plaats van te vertrouwen op één-shot code-generatie of losse optimalisatietools, biedt OptiML een systematische, end-to-end oplossing die:

• De drempel voor het schrijven van hoogwaardige GPU-code verlaagt voor niet-experts.
• De betrouwbaarheid van LLM-generatie voor kritieke systemen verhoogt door integratie van strenge verificatie en hardware-feedback.
• Een nieuwe standaard zet voor "search under verification", waarbij zoekprocessen worden geleid door fysieke hardware-metrics in plaats van alleen statistische waarschijnlijkheid.

Het framework toont aan dat het combineren van geavanceerde generatiemodellen (MoT) met zoekalgoritmen (MCTS) en hardware-aware feedback de beste weg is om de prestatiekloof tussen generieke code en expert-geoptimaliseerde kernels te overbruggen.