Technische Samenvatting: KForge – LLM-gestuurde Cross-Platform Kernel Generatie voor AI-Acceleratoren

Probleemstelling

De proliferatie van agentische AI-applicaties heeft de inferentie-workloads verschoven van monolithische executie naar heterogene pipelines. Deze pipelines weven redenering, tool calls en multi-agent coördinatie door elkaar, waarbij elke fase een eigen reken- en geheugenprofiel heeft. Bijgevolg vereist optimale implementatie het mappen van specifieke pipeline-fasen aan de meest geschikte acceleratoren. Deze heterogeniteit creëert een aanzienlijke systeemuitdaging: ontwikkelaars moeten high-performance kernels produceren voor een groeiend scala aan hardware-backends (NVIDIA, AMD, Intel, Apple) en programmeermodellen (CUDA, Triton, SYCL, Metal, etc.).

Het handmatig schrijven en porteren van deze kernels is tijdrovend, vereist diepe expertise op laag niveau en schaalt niet naarmate de kernelcomplexiteit toeneemt. Hoewel Large Language Models (LLMs) veelbelovend zijn gebleken voor codegeneratie, blijft het toepassen ervan op low-level kerneloptimalisatie moeilijk vanwege de broosheid van performance-code (waar kleine wijzigingen in geheugenlayout of precisie de correctheid kunnen breken) en de schaarste aan trainingsdata voor niet-CUDA platformen. Bovendien evalueren bestaande benchmarks kernels vaak in isolatie, waardoor ze er niet in slagen de impact van kerneloptimalisatie op de volledige end-to-end executie binnen productie-runtimes te vangen.

Methodologie: Het KForge Framework

KForge is een cross-platform, multi-stage autonoom programma-synthese framework dat ontworpen is om de iteratieve workflow van een kernel engineer te spiegelen. Het maakt gebruik van een collaboratieve loop tussen twee gespecialiseerde LLM-gebaseerde agents:

1. Generation Agent: Verantwoordelijk voor het produceren en progressief verfijnen van kernel-kandidaten. Het opereert binnen een prompt-gestuurde loop die het volgende incorporeert:

   • Iteratieve Verfijning: Het gebruik van compilatiefouten, runtime-fouten en numerieke mismatches uit vorige iteraties om de code te corrigeren.
   • Cross-Platform Translatie: Het benutten van bestaande functionele implementaties (bijv. CUDA) als referentie om code te genereren voor ondervertegenwoordigde platformen (bijv. Metal of SYCL).
   • Candidate Guardrails: Het afdwingen van beperkingen om ongewenst gedrag zoals data type downcasting of precisieveranderingen te voorkomen.

2. Performance-Analysis Agent: Een gespecialiseerde agent, afzonderlijk van de generator, om de schaarste aan optimalisatiesignalen binnen grote profiling-datasets te behandelen. Het interpreteert diverse profiling-inputs (programmatische metrieken van tools zoals Nsight Systems of visuele outputs van Xcode Instruments) om concrete code-niveau aanbevelingen te genereren. Deze aanbevelingen richten zich op hardware-benuttingsmetrieken zoals geheugenbandbreedte, warp occupancy en rekenintensiteit.

De Synthese Loop:
Het systeem wisselt af tussen twee fasen:

• Functionele Passes: De generation agent iterereert totdat de code compileert, zonder fouten executeert en numeriek correcte outputs produceert.
• Optimalisatie Passes: Zodra de correctheid is vastgesteld, biedt de performance-analysis agent feedback om de volgende ronde van synthese te sturen richting het dichten van de performance-kloof met hand-getunede baselines.

Het framework ondersteunt een uniforme interface voor vier vendors (NVIDIA, AMD, Intel, Apple) en zes programmeermodellen, waarbij een model registry wordt gebruikt om agnostisch te blijven voor de onderliggende LLM.

Belangrijkste Bijdragen

• KForge Framework: Introductie van een multi-stage autonoom synthese systeem waarbij generation en performance-analysis agents samenwerken met behulp van compilatie-, correctheids- en profiling-feedback.
• Cross-Platform Architectuur: Een systeemontwerp dat vier accelerator-vendors en zes programmeermodellen ondersteunt via een uniforme interface, wat de transfer van algoritmische ideeën over platformen heen mogelijk maakt.
• Dual-Mode Evaluatie: Case studies die de effectiviteit van het framework demonstreren in twee verschillende scenario's: het concurreren met volwassen, door de vendor getunede runtimes en het dienen als een bring-up tool voor opkomende hardware die nog geen referentie-implementaties heeft.

Experimentele Resultaten

De auteurs evalueerden KForge op twee backends met verschillende beschikbaarheid van baselines:

1. NVIDIA B200 (Concurreren met Vendor-getunede Runtimes)

• Context: Optimalisatie van decode-path kernels voor het gpt-oss-20b model tegenover TensorRT-LLM v1.3.0rc9.
• Micro-benchmarks: KForge-gegenereerde kernels presteerden consistent beter dan de baseline bij diverse batch sizes. Met name de "MoE finalize" kernel vertoonde een versnelling die toenam van 1.06× bij kleine batch sizes naar 1.43× bij een batch size van 128.
• End-to-End Performance: Het integreren van deze geoptimaliseerde kernels in de volledige inferentie-pipeline leverde een 2.12% verbetering op in end-to-end throughput (2656.61 vs. 2601.55 tokens/sec) en een 2.07% reductie in wall-clock tijd. De auteurs merken op dat hoewel deze winsten marginaal lijken, ze significant zijn gezien de jaren aan engineering-inspanning die al in de vendor baseline is geïnvesteerd.

2. Intel Arc B580 (Opkomende Hardware Bring-Up)

• Context: Genereren van Triton kernels voor de Intel Arc B580 (Battlemage) waar vergelijkbare hand-getunede referenties niet bestaan.
• Benchmark: Geëvalueerd op 37 GEMM + tail-ops workloads uit KernelBench Level 2.
• Resultaten: KForge bereikte een 5.13× geometrisch gemiddelde versnelling ten opzichte van de snellere van PyTorch eager mode en torch.compile.
• Optimalisatietechnieken: Analyse van de gegenereerde kernels toonde de succesvolle toepassing van:
  • Layer-wise Fusion: Het samenvoegen van post-GEMM operaties in enkele kernels om volledige tensor round-trips te elimineren.
  • Mixed-Precision Execution: Het gebruik van FP16 voor matmuls op XMX-units met FP32 accumulatie.
  • Single-Pass Reductions: Het vervangen van multi-pass GroupNorm en LogSumExp door streaming online algoritmen.

Betekenis en Claims

Het paper positioneert KForge als een praktische tool voor zowel productie-waardige optimalisatie als snelle hardware bring-up. De auteurs beweren dat:

• Agentic Synthese Viabel is: De scheiding tussen generatie en performance analyse stelt LLM's in staat om effectief te navigeren door de complexe trade-offs van low-level kerneloptimalisatie.
• Schaalbaarheid van Winsten: Zelfs marginale procentuele verbeteringen (bijv. ~2%) tegenover hoog-geoptimaliseerde vendor runtimes zijn betekenisvol bij grootschalige inzet en cumuleren tot significante efficiëntiewinsten.
• Het Gat Overbruggen: Het framework adresseert succesvol het "cold start" probleem voor nieuwe hardware (zoals Intel Arc) door performante kernels te genereren zonder dat er vooraf bestaande hand-getunede referenties nodig zijn, terwijl het tegelijkert je evenredig concurreert met gevestigde runtimes op volwassen hardware.

De conclusie van het werk is dat agentic kernel synthese transformeert van een onderzoeksconcept naar een functioneel onderdeel van de AI-system stack, in staat om de heterogeniteit van moderne inferentie-workloads aan te pakken.