KernelEvolve: Escalonando a Codificação de Kernels Agentes para Aceleradores de IA Heterogêneos na 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
Resumo
Tornar o treinamento e a inferência do modelo de recomendação por aprendizado profundo (DLRM) rápidos e eficientes é importante. No entanto, isso apresenta três desafios principais de sistema: a diversidade de arquiteturas de modelo, a diversidade de primitivas de kernel e a heterogeneidade de gerações e arquiteturas de hardware. Este artigo apresenta o KernelEvolve - uma estrutura de codificação de kernel agentiva - para enfrentar a heterogeneidade em larga escala para o DLRM. O KernelEvolve foi projetado para receber especificações de kernel como entrada e automatizar o processo de geração e otimização de kernel para modelos de recomendação em diversas arquiteturas de hardware heterogêneas. O KernelEvolve faz isso operando em múltiplos níveis de abstração de programação, desde as DSLs Triton e CuTe até linguagens de baixo nível independentes de hardware, abrangendo toda a pilha de otimização hardware-software. O processo de otimização de kernel é descrito como uma busca baseada em grafo com política de seleção, operador universal, função de aptidão e regra de terminação, adaptando-se dinamicamente ao contexto de execução em tempo de execução por meio de síntese de *prompts* aumentada por recuperação. Nós projetamos, implementamos e implantamos o KernelEvolve para otimizar uma ampla variedade de modelos de recomendação em produção em diversas gerações de GPUs da NVIDIA e AMD, bem como nos aceleradores de IA da Meta. Validamos o KernelEvolve na suíte publicamente disponível KernelBench, alcançando uma taxa de aprovação de 100% em todos os 250 problemas em três níveis de dificuldade, e em 160 operadores ATen do PyTorch em três plataformas de hardware heterogêneas, demonstrando 100% de correção. O KernelEvolve reduz o tempo de desenvolvimento de semanas para horas e alcança melhorias substanciais de desempenho em relação às linhas de base do PyTorch em diversos casos de uso em produção e para sistemas de IA heterogêneos em larga escala. Além das melhorias de eficiência de desempenho, o KernelEvolve mitiga significativamente a barreira de programabilidade para novo hardware de IA, permitindo a geração automatizada de kernel para hardware de IA desenvolvido 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.