KernelEvolve: Schaalvergroting van Agent-Gestuurde Kernelprogrammering voor Heterogene AI-Accelerators bij Meta
KernelEvolve: Scaling Agentic Kernel Coding for Heterogeneous AI Accelerators at Meta
December 29, 2025
Auteurs: Gang Liao, Hongsen Qin, Ying Wang, Alicia Golden, Michael Kuchnik, Yavuz Yetim, Jia Jiunn Ang, Chunli Fu, Yihan He, Samuel Hsia, Zewei Jiang, Dianshi Li, Uladzimir Pashkevich, Varna Puvvada, Feng Shi, Matt Steiner, Ruichao Xiao, Nathan Yan, Xiayu Yu, Zhou Fang, Abdul Zainul-Abedin, Ketan Singh, Hongtao Yu, Wenyuan Chi, Barney Huang, Sean Zhang, Noah Weller, Zach Marine, Wyatt Cook, Carole-Jean Wu, Gaoxiang Liu
cs.AI
Samenvatting
Het versnellen en efficiënter maken van training en inferentie voor deep learning-aanbevelingsmodellen (DLRM) is van groot belang. Dit brengt echter drie belangrijke systeemuitdagingen met zich mee: diversiteit in modelarchitectuur, diversiteit in kernelprimitieven, en heterogeniteit in hardwaregeneraties en -architecturen. Dit artikel presenteert KernelEvolve – een agent-gebaseerd kernelcoderingsframework – om heterogeniteit op grote schaal aan te pakken voor DLRM. KernelEvolve is ontworpen om kernelspecificaties als invoer te nemen en het proces van kernelgeneratie en -optimalisatie voor aanbevelingsmodellen te automatiseren across heterogene hardwarearchitecturen. KernelEvolve bereikt dit door te opereren op meerdere programmeerabstractieniveaus, van Triton- en CuTe-DSL's tot low-level hardware-agnostische talen, waarbij de volledige hardware-software-optimalisatiestack wordt bestreken. Het kerneloptimalisatieproces wordt omschreven als een op grafen gebaseerde zoektocht met selectiebeleid, universele operator, fitnessfunctie en stopregel, die zich dynamisch aanpast aan de runtime-uitvoeringscontext via retrieval-augmented prompt-synthese. We hebben KernelEvolve ontworpen, geïmplementeerd en ingezet om een grote verscheidenheid aan productie-aanbevelingsmodellen te optimaliseren across generaties NVIDIA- en AMD-GPU's, evenals Meta's AI-accelerators. We valideren KernelEvolve op de publiek toegankelijke KernelBench-suite, waarbij een slagingspercentage van 100% wordt behaald op alle 250 problemen across drie moeilijkheidsgraden, en 160 PyTorch ATen-operators across drie heterogene hardwareplatforms, wat 100% correctheid aantoont. KernelEvolve reduceert ontwikkeltijd van weken naar uren en behaalt aanzienlijke prestatieverbeteringen ten opzichte van PyTorch-baselines across diverse productie-use cases en voor heterogene AI-systemen op grote schaal. Naast prestatie-efficiëntieverbeteringen, vermindert KernelEvolve de programmeerbaarheidsbarrière voor nieuwe AI-hardware aanzienlijk door geautomatiseerde kernelgeneratie mogelijk te maken voor intern ontwikkelde AI-hardware.
English
Making deep learning recommendation model (DLRM) training and inference fast and efficient is important. However, this presents three key system challenges - model architecture diversity, kernel primitive diversity, and hardware generation and architecture heterogeneity. This paper presents KernelEvolve-an agentic kernel coding framework-to tackle heterogeneity at-scale for DLRM. KernelEvolve is designed to take kernel specifications as input and automate the process of kernel generation and optimization for recommendation model across heterogeneous hardware architectures. KernelEvolve does so by operating at multiple programming abstractions, from Triton and CuTe DSL to low-level hardware agnostic languages, spanning the full hardware-software optimization stack. The kernel optimization process is described as graph-based search with selection policy, universal operator, fitness function, and termination rule, dynamically adapts to runtime execution context through retrieval-augmented prompt synthesis. We designed, implemented, and deployed KernelEvolve to optimize a wide variety of production recommendation models across generations of NVIDIA and AMD GPUs, as well as Meta's AI accelerators. We validate KernelEvolve on the publicly-available KernelBench suite, achieving 100% pass rate on all 250 problems across three difficulty levels, and 160 PyTorch ATen operators across three heterogeneous hardware platforms, demonstrating 100% correctness. KernelEvolve reduces development time from weeks to hours and achieves substantial performance improvements over PyTorch baselines across diverse production use cases and for heterogeneous AI systems at-scale. Beyond performance efficiency improvements, KernelEvolve significantly mitigates the programmability barrier for new AI hardware by enabling automated kernel generation for in-house developed AI hardware.