Technische Samenvatting: CUDA-L2

Probleemstelling

Matrixvermenigvuldiging (matmul) is een fundamentele operatie in Large Language Models (LLMs) en vormt een aanzienlijk deel van de rekentijd tijdens zowel training als inferentie. Hoewel handmatig geoptimaliseerde CUDA-kernels bestaan, blijft het bereiken van optimale prestaties over diverse matrixdimensies ($M, N, K$) en GPU-architecturen een aanzienlijke uitdaging. Verschillende dimensie-triplets vereisen vaak afzonderlijke optimalisatiestrategieën die niet goed overdraagbaar zijn naar andere hardware. Bovendien variëren optimale strategieën, zelfs voor identieke dimensies, op basis van details zoals de precisie van de accumulator (FP16 versus FP32).

Bestaande geautomatiseerde benaderingen, zoals die gebruikmaken van Large Language Models (LLMs) voor kernelgeneratie, hebben zich primair gericht op benchmarks met enkelvoudige, vaste configuraties (bijv. KernelBench). Het blijft onduidelijk hoe deze methoden vertalen naar productieomgevingen waar ze moeten concurreren met hoogwaardig geoptimaliseerde, closed-source bibliotheken zoals NVIDIA's cuBLAS en cuBLASLt. Er is een gebrek aan geautomatiseerde systemen die in staat zijn om de enorme configuratieruimte van Half-precision General Matrix Multiply (HGEMM) systematisch te verkennen om menselijk-expert-geoptimaliseerde baselines te overtreffen.

Methodologie: CUDA-L2

De auteurs stellen CUDA-L2 voor, een systeem dat Large Language Models (LLMs) en Reinforcement Learning (RL) combineert om HGEMM CUDA-kernels automatisch te optimaliseren. Het systeem is ontworpen om 1.000 afzonderlijke configuraties te verwerken, die alle combinaties van $M, N, K$-waarden vertegenwoordigen uit $\{64, 128, 256, 512, 1024, 2048, 4096, 8192, 12288, 16384\}$.

Systeemarchitectuur en Training

CUDA-L2 breidt het eerdere CUDA-L1-framework uit met vier belangrijke technische verbeteringen:

1. Continued Pretraining: Het model wordt voorgetraind op een diverse corpus van CUDA-code die verder gaat dan standaard benchmarks, inclusief bronnen van PyTorch, ATen, CUTLASS en NVIDIA-tutorials. Instructiebeschrijvingen worden gegenereerd voor ruwe codefragmenten, en retrieval-augmented context wordt toegevoegd om nieuwe kennis en architecturale kenmerken te integreren.
2. Multi-Stage RL Training: De training verloopt van brede algemene kerneltypen naar gespecialiseerde domeinen (specifiek matmul).
   • General Kernel RL: Traint het model op een breed scala aan operaties (lineaire algebra, convolutie, reductie, etc.) met gebruik van executiesnelheid als beloning.
   • HGEMM RL: Richt zich specifiek op HGEMM-kernels met variërende $M, N, K$-configuraties.
3. Verrijkte Context met Profiling Metrics: De RL-context bevat uitgebreide NVIDIA Nsight Compute (NCU) profiling metrics (bijv. geheugendoorvoer, SM-bezetting, cache-efficiëntie) om betere optimalisatiebeslissingen mogelijk te maken dan enkel end-to-end executietijd.
4. Ontwerp van de Beloningsfunctie: De beloning voor een gegenereerde kernel is gebaseerd op de gemiddelde versnellingsscore over testiteraties. Om robuustheid te garanderen, wordt de beloning gestraft voor:
   • Numerieke afwijking van een FP32 CPU ground truth.
   • Excessieve codelengte (om beknoptheid te stimuleren).
   • Het systeem handhaaft strikte correctheidscriteria, inclusclusief exacte overeenkomsten voor binaire inputs en begrensde afwijkingen tegen betrouwbare baselines (cuBLAS, cuBLASLt).

Evaluatieprotocol

Het systeem wordt geëvalueerd onder twee scenario's om verschillende implementatiereliteiten te weerspiegelen:

• Offline Mode: Kernels worden opeenvolgend uitgevoerd zonder tijdsintervallen, wat een piekdoorvoer simuleert.
• Server Mode: Kernels worden op willekeurige intervallen uitgevoerd om real-time inferentie te simuleren, waarbij rekening wordt gehouden met GPU-thermische dynamica (koelperiodes en clock boosting).

De evaluatie vergelijkt CUDA-L2 met vier baselines:

1. torch.matmul (PyTorch's standaard dispatch).
2. cuBLAS (met optimale layoutselectie tussen NN en TN).
3. cuBLASLt-heuristic (selectie van de hoogst gerangschikte algoritme uit NVIDIA's heuristische API).
4. cuBLASLt-AutoTuning (uitputtend testen van tot 100 kandidaten om de snelste te vinden).

Belangrijkste Bijdragen en Ontdekte Technieken

CUDA-L2 laat zien dat LLM-gestuurde RL superieure implementaties kan ontdekken door systematisch configuratieruimtes te verkennen die onpraktisch zijn voor handmatige tuning. Het systeem ontdekt en past autonoom verschillende geavanceerde optimalisatietechnieken toe:

• Abstractie Selectie: Het model selecteert automatisch tussen ruwe WMMA intrinsics voor kleine matrices en CuTe-abstracties voor grotere, complexe getegelde operaties.
• Input Padding: Het systeem ontdekt strategieën waarbij inputmatrices worden aangevuld met nullen voor niet-deelbare dimensies (bijv. het aanvullen van $M=8192$ naar $8320$) om efficiëntere tile-groottes ($B_M$) mogelijk te maken die anders ongeldig zouden zijn, waarbij de overhead van extra berekeningen wordt afgewogen tegen de prestatiewinst.
• Geavanceerde Memory en Compute Overlap:
  • Double-Buffered Register Fragments: Afwisselen tussen twee sets registerfragmenten om datamovement te overlappen met berekeningen (ping-pong executie).
  • Aggressieve Prefetching: Verder gaan dan single-step prefetching naar multi-step prefetching (het laden van meerdere iteraties vooruit) wanneer de registerdruk dit toelaat.
  • Staggered Prefetch Scheduling: Het afwisselen van de prefetching van matrices A en B rond de matrix-multiply-accumulate (MMA) operatie om instructie-niveau parallellisme te maximalen.
  • Direct Register-to-Shared-Memory Copy: Het elimineren van tussenliggende tijdelijke tensors in de epiloog wanneer registerlayouts overeenkomen met shared memory layouts, gebruikmakend van brede datatypen (bijv. uint128_t) om kopieeroperaties te verminderen.
  • Block Swizzling: Het herordenen van thread block executie om L2-cache hits te verbeteren, met name voor grote probleemomvang.

Resultaten

CUDA-L2 presteert systematisch beter dan alle belangrijke baselines over 1.000 configuraties op A100 GPU's.

Offline Mode Prestaties (Gemiddelde versnelling t.o.v. baselines):

• +22,0% t.o.v. torch.matmul.
• +19,2% t.o.v. cuBLAS (met optimale layout).
• +16,8% t.o.v. cuBLASLt-heuristic.
• +11,4% t.o.v. cuBLASLt-AutoTuning (de sterkste baseline, die uitputtend 100 kandidaten test).

Server Mode Prestaties:
De versnellingen nemen toe in het server-scenario, waarschijnlijk door het vermogen van het systeem om thermische dynamica en cache-koelingseffecten beter te beheren:

• +28,7% t.o.v. torch.matmul.
• +26,0% t.o.v. cuBLAS.
• +22,4% t.o.v. cuBLASLt-heuristic.
• +15,9% t.o.v. cuBLASLt-AutoTuning.

De win rates variëren van 79,3% tot 95,7% over alle baselines, wat aangeeft dat de verbeteringen systematisch zijn en niet gedreven worden door uitschieters. Het systeem laat ook zien dat de versnellingen over het algemeen hoger zijn voor kleinere matrixgroottes waar de GPU-benutting lager is, terwijl grotere matrices de doorvoer van de GPU verzadigen, waardoor de kloof met de baselines kleiner wordt.

Betekenis

Het artikel stelt dat CUDA-L2 aantoont dat zelfs de meest prestatie-kritische en zwaar geoptimaliseerde kernels, zoals HGEMM, kunnen worden verbeterd door LLM-gestuurde RL-automatisering. Door optimalisatie te schalen over configuratieruimtes die moeilijk handmatig te tunen zijn, bereikt het systeem prestatiewinsten die superieur zijn aan state-of-the-art closed-source bibliotheken.

Hoewel de huidige versie beperkt is tot de A100 (Ampere) architectuur, is het framework ontworpen voor brede toepasbaarheid. De auteurs merken op dat zij voortdurend werken aan het uitbreiden van het systeem naar andere GPU-architecturen, waaronder Ada Lovelace, Hopper en Blackwell. Het werk suggereert dat geautomatiseerde methoden systematisch enorme configuratieruimtes kunnen verkennen om optimalisaties te vinden die handmatige tuning of heuristiek-gebaseerde bibliotheken mogelijk missen.