KernelEvolve: Scalabilità della Programmazione Kernel Agente per Acceleratori IA Eterogenei presso Meta
KernelEvolve: Scaling Agentic Kernel Coding for Heterogeneous AI Accelerators at Meta
December 29, 2025
Autori: 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
Abstract
Rendere veloce ed efficiente l'addestramento e l'inferenza dei modelli di raccomandazione basati sul deep learning (DLRM) è di fondamentale importanza. Tuttavia, ciò presenta tre principali sfide a livello di sistema: la diversità delle architetture dei modelli, la diversità dei kernel primitivi e l'eterogeneità delle generazioni e delle architetture hardware. Questo articolo presenta KernelEvolve, un framework agent-based per la codifica dei kernel, concepito per affrontare l'eterogeneità su larga scala per i DLRM. KernelEvolve è progettato per ricevere specifiche dei kernel come input e automatizzare il processo di generazione e ottimizzazione dei kernel per i modelli di raccomandazione su architetture hardware eterogenee. KernelEvolve opera a più livelli di astrazione di programmazione, dai DSL Triton e CuTe fino a linguaggi di basso livello indipendenti dall'hardware, coprendo l'intero stack di ottimizzazione hardware-software. Il processo di ottimizzazione del kernel è descritto come una ricerca basata su grafi con criteri di selezione, operatori universali, funzioni di fitness e regole di terminazione, che si adatta dinamicamente al contesto di esecuzione runtime attraverso una sintesi prompt aumentata da retrieval. Abbiamo progettato, implementato e distribuito KernelEvolve per ottimizzare un'ampia varietà di modelli di raccomandazione in produzione su diverse generazioni di GPU NVIDIA e AMD, nonché sugli acceleratori AI interni di Meta. Convalidiamo KernelEvolve sulla suite pubblicamente disponibile KernelBench, raggiungendo un tasso di successo del 100% su tutti i 250 problemi attraverso tre livelli di difficoltà e su 160 operatori ATen di PyTorch su tre piattaforme hardware eterogenee, dimostrando una correttezza del 100%. KernelEvolve riduce il tempo di sviluppo da settimane a ore e consegue sostanziali miglioramenti delle prestazioni rispetto ai baseline PyTorch in svariati casi d'uso produttivi e per sistemi di IA eterogenei su larga scala. Oltre ai miglioramenti dell'efficienza prestazionale, KernelEvolve mitiga significativamente la barriera della programmabilità per i nuovi hardware di IA, abilitando la generazione automatizzata di kernel per hardware di IA sviluppati internamente.
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.