Model2Kernel: Model-Aware Symbolic Execution For Safe CUDA Kernels

Model2Kernel is het eerste praktische systeem dat modelbewuste dynamische analyse en gespecialiseerde symbolische uitvoering combineert om automatisch geheugenveiligheidsbugs in CUDA-kernen voor LLM-inferentie te detecteren, wat resulteerde in het vinden van 353 onbekende bugs met slechts negen vals-positieven.

De kern

Stel je voor dat je een gigantische, superkrachtige robot hebt die miljoenen vragen tegelijk kan beantwoorden. Dit is een Grote Taalmodel (LLM), zoals de slimme AI's die we vandaag de dag gebruiken. Maar deze robot is niet gemaakt van vlees en bloed; hij draait op speciale computerchips genaamd GPU's.

Om deze robot snel te laten werken, gebruiken programmeurs kleine, razendsnelle instructiesheets die CUDA-kernels heten. Je kunt deze kernels zien als de specifieke gereedschappen in de gereedschapskist van de robot. Als een gereedschap slecht is gemaakt (bijvoorbeeld een schroevendraaier die uit elkaar valt), kan de hele robot crashen, de verkeerde antwoorden geven, of zelfs gehackt worden.

Het probleem is dat deze gereedschappen steeds complexer worden en vaak fouten bevatten die moeilijk te vinden zijn. Model2Kernel is een nieuw, slim hulpmiddel dat deze fouten automatisch opspoort.

Hier is hoe het werkt, vertaald in alledaagse taal:

1. Het Probleem: De "Onzichtbare" Fouten

Stel je voor dat je een fabriek hebt waar robots worden gebouwd. De ontwerpers (de AI-modellen) zeggen: "Gebruik dit gereedschap voor deze taak." Maar de fabrieksarbeiders (de CUDA-kernels) weten niet altijd precies hoe groot de materialen zijn die ze moeten verwerken.

• Soms denken ze dat ze 100 blokken hebben, maar er zijn er 10.000.
• Soms proberen ze een vakje te openen dat niet bestaat.
• Soms proberen duizenden arbeiders tegelijk hetzelfde vakje te openen, wat tot chaos leidt.

Deze fouten zijn gevaarlijk. Ze kunnen de hele fabriek platleggen of kwaadaardige hackers de kans geven om de robot te manipuleren.

2. De Oplossing: Model2Kernel (De Slimme Inspecteur)

Model2Kernel is als een super-inspecteur die twee dingen tegelijk doet: hij kijkt naar het ontwerp van de fabriek én hij test de gereedschappen in een virtuele wereld.

Het bestaat uit twee hoofdonderdelen:

Deel A: HFProbe (De Ontwerper die de Blauwdruk Leest)

Stel je voor dat de inspecteur eerst naar de blauwdrukken van de fabriek kijkt. Hij leest de instructies van de AI-ontwerper (het model) om te begrijpen:

• "Welke gereedschappen worden gebruikt?"
• "Welke maten zijn vastgelegd door het ontwerp?" (Bijvoorbeeld: de breedte van de robot is altijd 7168).
• "Welke maten worden door de klant bepaald?" (Bijvoorbeeld: hoeveel vragen de klant stelt).

Dit is belangrijk omdat de inspecteur niet elke mogelijke situatie hoeft te testen. Hij weet nu precies welke maten "vast" zijn en welke "variabel" zijn. Hij maakt een virtuele proefomgeving zonder dat hij echt een dure GPU nodig heeft. Hij "speelt" de fabriek na op papier.

Deel B: cuKLEE (De Virtuele Testpiloot)

Nu komt de tweede helft: een virtuele testpiloot die de gereedschappen (kernels) in een virtuele wereld uitprobeert.

• In plaats van te zeggen: "Test met 10 vragen," zegt hij: "Test met elke mogelijke hoeveelheid vragen, van 1 tot 1 miljard."
• Hij gebruikt symbolische uitvoering. Dat klinkt ingewikkeld, maar het is simpel: in plaats van een getal te gebruiken, gebruikt hij een variabele (zoals een doosje met een vraagteken erop). Hij vraagt aan een wiskundige computer: "Is er een combinatie van getallen waarbij dit gereedschap uit elkaar valt?"
• Als het antwoord "ja" is, heeft hij een fout gevonden! Hij kan dan precies zeggen: "Als je 128.000 vragen stelt, breekt dit gereedschap."

3. Waarom is dit zo speciaal?

Vroeger hadden inspecteurs twee keuzes:

1. Testen op echte hardware: Dit was traag, duur en je kon maar een paar situaties testen. Je miste de zeldzame fouten die alleen bij extreme omstandigheden optraden.
2. Kijken naar de code: Dit was te moeilijk omdat de code te complex was en vaak afhankelijk van de grootte van de data.

Model2Kernel combineert het beste van beide werelden. Het begrijpt de context van de AI (via HFProbe) en gebruikt slimme wiskunde om alle mogelijke scenario's in één keer te testen (via cuKLEE).

4. Het Resultaat: Een Veiligere Toekomst

De onderzoekers hebben Model2Kernel laten draaien op de populairste AI-systemen ter wereld (zoals die van vLLM en Hugging Face). Het resultaat?

• Ze vonden 353 nieuwe fouten die niemand eerder zag.
• De meeste waren "integer overflows": rekenfouten waarbij getallen te groot werden voor de computer.
• Ze vonden ook fouten waarbij het gereedschap probeerde naar een plek te kijken die niet bestond (out-of-bounds).

De metafoor samengevat:
Stel je voor dat je een auto bouwt. Normaal gesproken test je of de remmen werken door er een keer mee te remmen. Model2Kernel is als een computer die duizenden simulaties draait in één seconde: "Wat gebeurt er als de remmen falen bij 200 km/u? Wat als de banden lek zijn? Wat als de bestuurder een vreemd signaal geeft?" En hij doet dit zonder de auto ooit fysiek te hoeven bouwen.

Conclusie

Model2Kernel is een game-changer. Het zorgt ervoor dat de AI's van de toekomst veiliger, betrouwbaarder en minder vatbaar zijn voor crashes of hackers. Het is als het hebben van een onzichtbare, super-snelle veiligheidscontrole die elke schroef in de machine controleert voordat de machine überhaupt aan de weg wordt gezet.

Technische samenvatting

Probleemstelling

De snelle adoptie van Large Language Models (LLM's) heeft GPU-versnelde inferentie cruciaal gemaakt voor moderne computing. Deze systemen vertrouwen op CUDA-kernels om de kernoperaties van transformer-architecturen uit te voeren. Deze kernels zijn echter uiterst kwetsbaar voor geheugensicherheidsbugs (zoals buffer-overflows, integer-overflows en data-races) vanwege:

1. Modelafhankelijke tensorlay-outs: De grootte en vorm van tensors hangen af van de specifieke modelarchitectuur en de invoer van de gebruiker.
2. Complexe indexering: Fijnmazige berekeningen over grote tensors en duizenden threads leiden vaak tot out-of-bounds toegang.
3. Variabele invoer: In tegenstelling tot traditionele software, hebben CUDA-kernels in LLM-systemen vaak invoer met variabele lengtes (bijv. sequentielengte en batchgrootte), wat statische analyse bemoeilijkt.

Bestaande technieken zijn ontoereikend: dynamische methoden (zoals Compute Sanitizer) hebben hoge runtime-overhead of vereisen specifieke hardware, terwijl statische tools moeite hebben met buffers waarvan de grootte pas tijdens runtime bekend is. Er ontbreekt een praktische oplossing die specifiek is ontworpen voor de complexe stack van LLM-inferentie.

Methodologie: Model2Kernel

Model2Kernel is een geautomatiseerd systeem dat modelbewuste dynamische analyse combineert met CUDA-specialistische symbolische uitvoering. Het systeem bestaat uit twee hoofdcomponenten:

1. HFProbe (Dynamic Model Profiler)

HFProbe analyseert hoe een LLM-model (bijv. van Hugging Face) CUDA-kernels aanroept zonder daadwerkelijke GPU-hardware nodig te hebben.

• Statische analyse: Bouwt een call-graph van de Python-modelcode om alle mogelijke kernels te identificeren.
• Dynamische profilering: Voert het model uit met "fake" CUDA-kernels om te registreren hoe kernels worden aangeroepen, welke tensorvormen worden gebruikt en welke parameters door het model worden vastgesteld versus door de gebruiker.
• Configuratie-mutatie: Gebruikt een LLM-agent om config.json-bestanden te muteren en framework-instellingen te wijzigen, zodat meer kernels worden getriggerd die niet standaard worden aangeroepen.
• Contextgeneratie: Levert concrete constraints (bijv. vaste hidden sizes) en variabele grenzen (bijv. token-aantallen) aan de symbolische engine, wat de zoekruimte drastisch verkleint.

2. cuKLEE (CUDA-Specialized Symbolic Execution Engine)

cuKLEE is een aangepaste versie van de KLEE-symbolische uitvoering, specifiek ontworpen voor CUDA-kernels.

• Geheugenmodel voor Tensors: In plaats van een contiguus geheugenmodel te gebruiken, modelleert cuKLEE elke tensor als een discrete, gescheiden geheugenregio. Dit vereenvoudigt het volgen van pointers en het controleren van grenzen, zelfs bij dynamische tensorvormen.
• Symbolische Variabelen:
  • Tensor-eigenschappen: Basisadres, aantal elementen, dimensies en datatypen worden gemodelleerd als symbolische variabelen.
  • Thread-identificatie: Thread- en block-ID's (threadIdx.x, blockIdx.x) worden behandeld als symbolische variabelen in plaats van geconcretiseerd. Hierdoor kan één symbolische uitvoering het gedrag van duizenden parallelle threads vertegenwoordigen.
• Data-race detectie: Door een tweede set symbolische thread-ID's te introduceren, kan het systeem controleren of twee verschillende threads conflicterende schrijfbewerkingen op hetzelfde adres kunnen uitvoeren.
• Constraint Oplossing: Gebruikt de Z3-solver om te bepalen of er invoerwaarden bestaan die tot geheugenviolaties leiden.

Belangrijkste Bijdragen

1. cuKLEE: Een nieuwe symbolische uitvoering-engine die dynamisch gevormde tensors en duizenden CUDA-threads efficiënt kan analyseren door middel van een gespecialiseerd geheugenmodel en thread-symbolisatie.
2. HFProbe: Een dynamische profiler die modelarchitectuur-informatie extrahert en gebruikt om de analyse van kernels te sturen, waardoor onrealistische invoer wordt uitgesloten en de schaalbaarheid wordt verbeterd.
3. Model2Kernel: Het geïntegreerde framework dat voor het eerst automatisch geheugensicherheidsbugs in CUDA-kernels voor LLM-inferentie detecteert, met ondersteuning voor modellen in Hugging Face-formaat.

Resultaten

Het team heeft Model2Kernel geëvalueerd op een breed scala aan bronnen:

• Datasets: Alle tekstgeneratiemodellen in vLLM, alle Hugging Face-modellen met aangepaste kernels, en modellen uit recente onderzoekspublicaties.
• Gevonden Bugs: In totaal werden 353 unieke geheugenbugs ontdekt:
  • 328 integer-overflows (vaak leidend tot buffer-overflows).
  • 25 out-of-bounds toegangsfouten.
  • Slechts 9 false positives (een foutpercentage van ~2,5%).
• Vergelijking: Bij een test met 20 bekende vLLM-bugs detecteerde Model2Kernel er 15, terwijl bestaande tools (Honeycomb, GKLEE, ESBMC-GPU) er slechts 0 tot 5 vonden. De bestaande tools faalden vaak door beperkingen in het hanteren van moderne C++-functies, dynamische array-groottes of complexe thread-configuraties.
• Kosten: De analyse kostte ongeveer 986 uur (grotendeels door cuKLEE) en $442 aan LLM-kosten voor configuratiemutatie, wat acceptabel is gezien de omvang van de codebases (>100k regels).

Betekenis en Impact

• Veiligheid: De paper toont aan dat geheugenbugs in CUDA-kernels voor LLM's niet alleen crashes veroorzaken, maar ook kunnen worden uitgebuit voor adversariale aanvallen (zoals het wijzigen van modelgewichten of arbitraire code-uitvoering).
• Praktische Toepasbaarheid: Model2Kernel biedt een praktische oplossing voor ontwikkelaars om de veiligheid van nieuwe kernel-optimalisaties te verifiëren voordat ze in productie gaan.
• Technologische Vooruitgang: Door de combinatie van dynamische profilering en gespecialiseerde symbolische uitvoering, overwint Model2Kernel de beperkingen van eerdere statische en dynamische analysemethoden voor GPU-code. Het stelt een nieuwe standaard voor het testen van high-performance computing componenten in AI-systemen.

Kortom, Model2Kernel is een doorbraak in het waarborgen van de integriteit van de infrastructuur die de huidige AI-revolutie aandrijft.