Technische Samenvatting: GPU Forecasters: Taalmodellen als Selectieve Surrogaten voor Kernel Runtime Optimalisatie

Probleemstelling

GPU-kernels zijn de fundamentele werkpaarden van moderne deep learning, waarbij kleine efficiëntiewinsten vertalen naar significante reducties in trainings- en inferentiekosten. Het optimaliseren van deze kernels omvat doorgaans een generatie-evaluatiecyclus waarbij Large Language Models (LLM's) kandidaat-kernels voorstellen, die vervolgens worden gecompileerd en uitgevoerd op fysieke hardware om de prestaties te meten. Echter, dit evaluatieproces vormt een ernstige bottleneck. Elke evaluatie vereist compilatie (wat 60–120 seconden duurt voor CUDA) en uitvoering op het apparaat, waardoor de doorvoer van evaluaties orders van grootte langzamer is dan die van generatie. Naarmate de zoekbudgetten schalen, beperkt de kosten van herhaalde fysieke metingen de mogelijkheid om de enorme ruimte van potentiële kernel-implementaties te verkennen.

Terwijl eerder werk zich richtte op het gebruik van LLM's om betere kernels te genereren, richt dit artikel zich op het complementaire probleem van het verminderen van de evaluatiekosten. De auteurs stellen voor om LLM's niet alleen als generatoren te gebruiken, maar als selectieve surrogaten voor de fysieke GPU. Het doel is om de relatieve prestaties van een kandidaat-kernel ten opzien van een referentie-implementatie te voorspellen en om te identificeren wanneer het model onzeker is, zodat alleen de onzekere gevallen worden uitgesteld naar de dure fysieke GPU.

Methodologie

Taakformulering: Selectieve Voorspelling

De auteurs formuleren het probleem als het voorspellen van de relatieve versnelling ($S_i = T_{ref} / T_i$) van een kandidaat-kernel ten opzien van een referentie-kernel, gegeven de hardware-specificaties. Het voorspellen van de absolute runtime in seconden wordt als onhandelbaar beschouwd vanwege de complexiteit van het modelleren van compilergedrag en micro-architecturale details. In plaats daarvan wordt het doel gediscretiseerd in acht ordinale versnellingsbins (variërend van "Ernstige Vertraging" tot "Extreme Versnelling").

De surrogaat (een LLM) neemt de referentie-kernel, de kandidaat-kernel en de hardware-specificaties als input. Het genereert een keten van redeneringen (chain of reasoning) en geeft een waarschijnlijkheidsverdeling over de acht bins output. De geëstateerde betrouwbaarheid van het model wordt gedefinieerd als de maximale waarschijnlijkheid toegewezen aan een enkele bin.

Selectieve Voorspellingsstrategie

Het systeem werkt onder een selectieve voorspellingsframework. De surrogaat wordt gebruikt om kandidaten te rangschikken. Als de betrouwbaarheid van de surrogaat een bepaalde drempel overschrijdt, wordt de voorspelling geaccepteerd en wordt de kandidaat niet gemeten op de GPU. Als de betrouwbaarheid laag is (wat wijst op onzekerheid), wordt de kandidaat uitgesteld naar de fysieke GPU voor grondwaarheid-meting. Dit stelt het zoekproces in staat om veel meer kandidaten te overwegen binnen een vast GPU-meetbudget.

Reinforcement Learning (RL) voor Kalibratie

Om de nauwkeurigheid van de surrogaat te verbeteren en, cruciaal, de kalibratie (de afstemming tussen de geëstateerde betrouwbaarheid en de werkelijke nauwkeurigheid), passen de auteurs Reinforcement Learning toe op een open-weights model (GPT-OSS-20B). Ze definiëren beloningen gebaseerd op:

1. Correctheid: Of de voorspelde bin overeenkomt met de gemeten bin.
2. Categorische Kalibratie (Brier Score): Het straffen van de verdeling op basis van de afstand tot de ware one-hot target.
3. Afstand-bewuste Kalibratie (CRPS): Het straffen van de cumulatieve verdeling, rekening houdend met het ordinale karakter van de versnellingsbins (bijv. het voorspellen van een 2× versnelling wanneer de waarheid 1.5× is, is minder ernstig dan het voorspellen van 4×).

De trainingsdata wordt geconstrueerd als een bijproduct van de kernelzoektocht zelf, waarbij succesvolle kandidaten uit eerdere zoekruns als trainingsvoorbeelden worden gebruikt.

Belangrijkste Bijdragen

1. LLM's als GPU-Surrogaten: Het artikel demonstreert dat LLM's een niet-trivial nauwkeurigheid kunnen bereiken in het voorspellen van hoe een fysieke GPU reageert op kandidaat-kernels, waarbij ze een aanzienlijk deel van de beschikbare versnelling terugwinnen door kandidaten te rangschikken op basis van voorspellingen in plaats van metingen.
2. Selectief Voorspellingsframework: De auteurs presenteren GPU-voorspelling als een selectief voorspellingsprobleem, waarbij zij aantonen dat surrogaten gekalibreerde onzekerheid kunnen rapporteren. Dit maakt een zoekstrategie mogelijk die betrouwbare voorspellingen vertrouwt en fysieke metingen reserveert voor onzekere gevallen.
3. Impact van RL op het Gedrag van de Surrogaat: Door middel van experimenten karakteriseert het artikel hoe RL-training de eigenschappen van de surrogaat verandert. Specifiek zorgt het toevoegen van kalibratiebeloningen (Brier of CRPS) voor een verbetering van de betermetbaarheid van de betrouwbaarheidsschattingen, terwijl training op enkel correctheid de relatie tussen betrouwbaarheid en nauwkeurigheid kan verslechteren.
4. End-to-End Zoekverbetering: De studie toont aan dat het integreren van de surrogaat in een kernelzoeklus (specifiek een PUCT-gebaseerde zoektocht) de zoektocht in staat stelt om aanzienlijk meer kandidaten te verkennen per GPU-meting. Op de meeste taken leidt dit tot het vinden van snellere kernels dan een baseline die elke voorstel onder hetzelfde GPU-budget meet.
5. Dataset Release: De auteurs brengen een herbruikbare collectie van 12.388 door LLM's gegenereerde GPU-kernels met gemeten runtimes vrij, verspreid over 118 problemen, meerdere GPU-typen (A100, L40S, H100) en diverse generatiemethoden. Deze dataset ondersteunt breder onderzoek naar GPU-prestatievoorspelling en kernelzoektochten.

Resultaten

• Voorspellingsnauwkeurigheid: Off-the-shelf LLM's (bijv. GPT-OSS-20B) herstellen ongeveer 83% van de best beschikbare versnelling wanneer slechts de bovenste 1% van de kandidaten op de GPU wordt gemeten. Deze prestatie verbetert naar 94% wanneer de bovenste 50% wordt gemeten.
• Kalibratie en RL: RL-training verandert het gedrag van de surrogaat aanzienlijk.
  • Training op enkel correctheid verbetert de exacte-bin nauwkeurigheid maar verslechtert de correlatie tussen betrouwbaarheid en fout (het model wordt overmoedig bij foutieve voorspellingen).
  • Kalibratie-training (Brier/CRPS) verbetert de monotoniciteit van de fout versus de betrouwbaarheid. Opvallend genoeg verbetert Brier-gekalibreerde training de metriek "recovered speedup" met ~2,1 procentpunt ten opzien van het basemodel, wat suggereert dat beter gekalibreerde verdelingen leiden tot betere rangschikkingssignalen voor het zoekalgoritme.
• Zoakefficiëntie: In een end-to-end kernelzoektocht vindt de door de surrogaat ondersteunde zoektocht snellere kernels dan de baseline op vier van de zes taken, waarbij vaak de beste versnelling van de baseline wordt bereikt met aanzienlijk minder GPU-metingen. Op taken waar de baseline wint, lijkt de zoektocht vroeg verzadigd te zijn, wat aangeeft dat de surrogaat het meest effectief is wanneer de zoekruimte groot en onverkend is.
• Ontdekkingsmomenten: De surrogaat kan "ontdekkingsmomenten" (mutaties waarbij een kind-kernel aanzienlijk sneller is dan zijn ouder) beter identificeren dan door toeval, hoewel precisie een uitdaging blijft, wat GPU-bevestiging noodzakelijk maakt voor beslissingen met hoge inzet.

Betekenis en Claims

Het artikel claimt dat LLM's een bredere rol kunnen spelen in kerneloptimalisatie dan louter codegeneratie. Door te fungeren als virtuele modellen van een GPU, kunnen ze dienen als hoog-doorvoers, laag-kosten filters die de schaarste aan fysieke hardwarebronnen beheren.

De auteurs benadrukken dat de primaire waarde van deze aanpak niet alleen snelheid is, maar de mogelijkheid om model-gebaseerde optimalisatie uit te voeren zonder bij elke stap de oorspronkelijke system aan te roepen. De surrogaat biedt een hogere doorvoer en lagere kosten per query dan een GPU, waarbij een afruil wordt gemaakt met voorspellingsruis die beheerd kan worden via gekalibreerde onzekerheidsschattingen. Het werk verbindt kerneloptimalisatie met het bredere veld van het bouwen van neurale simulatoren voor complexe fysieke systemen, en suggereert een pad naar optimizers die hardwaremeting behandelen als een schaars middel dat selectief moet worden ingezet.