PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX

Binnen de moderne computer fungeert de graphics processing unit, of GPU, als een krachtige motor die is ontworpen om enorme hoeveelheden gegevens gelijktijdig te verwerken. Om het maximale uit deze motor te halen, schrijven softwareontwikkelaars gespecialiseerde instructies die de hardware precies vertellen hoe gegevens moeten worden verplaatst en berekeningen moeten worden uitgevoerd. Deze instructies worden geschreven in een laag-niveau taal genaamd PTX, wat de programmeur directe controle geeft over de interne werking van de machine. Hoक hoger-niveau hulpmiddelen bestaan om dit proces gemakkelijker te maken, maar zij missen soms de unieke kenmerken van de nieuwste hardwaregeneraties. Terwijl deze machines snel evolueren, wordt het een moeilijke uitdaging om software perfect af te stemmen op hun specifieke capaciteiten. De vraag rijst: kan kunstmatige intelligentie, specifiek grote taalmodellen, leren om deze precieze, laag-niveau instructies zelfstandig te schrijven, of vertrouwt het nog steeds op oudere, generieke methoden die prestaties onbenut laten?

Een team van onderzoekers zette zich in om dit te beantwoorden door een nieuwe testomgeving genaamd PTXBench te creëren. Ze wilden zien of deze AI-modellen niet alleen code konden schrijven die werkt, maar ook code die actief gebruikmaakt van de specifieke, geavanceerde functies van de nieuwste GPU's. De onderzoekers richtten zich op twee van de krachtigste grafische kaarten die beschikbaar zijn, de H100 en de B200, en testten de modellen op taken die complexe wiskundige operaties omvatten die worden gebruikt in kunstmatige intelligentie, zoals matrixvermenigvuldiging en attention-mechanismen. Het doel was om drie dingen te meten: of de code zonder fouten draaide, of de code daadwerkelijk de specifieke hardware-instructies activeerde waar de onderzoekers om vroegen, en of de code sneller draaide dan de beste bestaande softwarebibliotheken.

De resultaten onthulden een aanzienlijke kloof tussen wat de modellen konden en wat vereist was voor topprestaties. Hoewel de AI-modellen vaak succesvol waren in het schrijven van code die correct werkte voor eenvoudigere taken, hadden ze aanzienlijk moeite met complexere operaties, met name die die de backward pass van attention-mechanismen betreffen, die cruciaal zijn voor het trainen van AI-systemen. Zelfs wanneer de modellen erin slaagden code te schrijven die de gevraagde hardware-instructies uitvoerde, kwam die code zelden overeen met de snelheid van de leidende professionele bibliotheken. In veel gevallen was de door AI gegenereerde code trager, wat suggereert dat het simpelweg laten uitvoeren van de instructies niet genoeg was om echte optimalisatie te bereiken. De studie vond dat de modellen vaak terugvielen op generieke codepatronen in plaats van de unieke, architectie-specifieke kenmerken van de nieuwere chips te benutten, zelfs wanneer die kenmerken expliciet aan hen werden beschreven.

Om te begrijpen hoe deze modellen verbeterd kunnen worden, voerden de onderzoekers een gecontroleerd experiment uit waarbij ze een specifiek AI-model onderwezen met een methode genaamd supervised fine-tuning. Ze lieten het model voorbeelden zien van zijn eigen fouten, samen met gecorrigeerde versies van de code en uitleg waarom de correcties werkten. Deze aanpak, die ze Fixit noemden, hielp het model te verbeteren op verschillende taken, met name bij het genereren van functioneel correcte code. De verbetering was echter niet uniform. Het model bleef moeite hebben met het generaliseren van zijn nieuwe vaardigheden naar verschillende probleemgroottes of variaties van de taken die het niet tijdens de training had gezien. De onderzoekers ontdekten dat de kwaliteit van de correcties en de balans van de trainingsdata er net zo belangrijk waren als de enorme hoeveelheid data zelf. Een model dat getraind was op een diverse set problemen presteerde beter dan een model dat getraind was op een grotere maar minder gevarieerde set, wat aangeeft dat het type leerervaring belangrijker was dan het volume.

De studie benadrukte ook het belang van het voorzien van de juiste context voor de AI. Wanneer de onderzoekers de modellen gedetailleerde informatie gaven over de specifieke mogelijkheden van de hardware, waren de modellen veel eerder geneigd om de gevraagde instructies te gebruiken. Zonder deze specifieke kennis gebruikte het model zelden de gevraagde PTX. Dit suggereert dat hoewel AI kan leren programmeren op dit lage niveau, het nauwkeurige begeleiding en hoogwaardige voorbeelden vereist om dit effectief te doen. De onderzoekers concludeerden dat de huidige AI-modellen weliswaar veelbelovend zijn, maar nog niet klaar zijn om ervaren menselijke ingenieurs te vervangen bij het optimaliseren van code voor de nieuwste hardware. De weg vooruit omvat betere trainingsdata, meer gebalanceerde leerervaringen en een dieper begrip van hoe deze modellen te leren redeneren over de specifieke beperkingen van evoluerende computerarchitecturen.