KernelEvolve: Escalando la Codificación de Kernels Agentes para Aceleradores de IA Heterogéneos en Meta

KernelEvolve: Scaling Agentic Kernel Coding for Heterogeneous AI Accelerators at Meta

December 29, 2025

Autores: 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

Resumen

Lograr que el entrenamiento y la inferencia de los modelos de recomendación de aprendizaje profundo (DLRM) sean rápidos y eficientes es importante. Sin embargo, esto presenta tres desafíos clave a nivel de sistema: la diversidad de arquitecturas de modelos, la diversidad de primitivas de kernels y la heterogeneidad de generaciones de hardware y arquitecturas. Este artículo presenta KernelEvolve, un marco de trabajo de codificación de kernels basado en agentes, para abordar la heterogeneidad a gran escala en los DLRM. KernelEvolve está diseñado para tomar especificaciones de kernels como entrada y automatizar el proceso de generación y optimización de kernels para modelos de recomendación en diversas arquitecturas de hardware heterogéneas. KernelEvolve lo hace operando en múltiples niveles de abstracción de programación, desde DSLs como Triton y CuTe hasta lenguajes de bajo nivel independientes del hardware, abarcando toda la pila de optimización hardware-software. El proceso de optimización de kernels se describe como una búsqueda basada en grafos con una política de selección, un operador universal, una función de aptitud y una regla de terminación, que se adapta dinámicamente al contexto de ejecución en tiempo de ejecución mediante la síntesis de prompts aumentada con recuperación de información. Hemos diseñado, implementado y desplegado KernelEvolve para optimizar una amplia variedad de modelos de recomendación en producción a lo largo de generaciones de GPUs de NVIDIA y AMD, así como de los aceleradores de IA de Meta. Validamos KernelEvolve en el conjunto de pruebas públicamente disponible KernelBench, logrando una tasa de aprobación del 100% en los 250 problemas de los tres niveles de dificultad, y en 160 operadores ATen de PyTorch en tres plataformas de hardware heterogéneas, demostrando un 100% de corrección. KernelEvolve reduce el tiempo de desarrollo de semanas a horas y logra mejoras sustanciales de rendimiento respecto a las líneas base de PyTorch en diversos casos de uso en producción y para sistemas de IA heterogéneos a gran escala. Más allá de las mejoras en eficiencia de rendimiento, KernelEvolve mitifica significativamente la barrera de programabilidad para nuevo hardware de IA al permitir la generación automatizada de kernels para hardware de IA desarrollado 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.