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

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Liao, Gang, Qin, Hongsen, Wang, Ying, Golden, Alicia, Kuchnik, Michael, Yetim, Yavuz, Ang, Jia Jiunn, Fu, Chunli, He, Yihan, Hsia, Samuel, Jiang, Zewei, Li, Dianshi, Pashkevich, Uladzimir, Puvvada, Varna, Shi, Feng, Steiner, Matt, Xiao, Ruichao, Yan, Nathan, Yu, Xiayu, Fang, Zhou, Levenstein, Roman, Ho, Kunming, Zhu, Haishan, Hammond, Alec, Li, Richard, Mathews, Ajit, Gondkar, Kaustubh, Zainul-Abedin, Abdul, Singh, Ketan, Yu, Hongtao, Chi, Wenyuan, Huang, Barney, Zhang, Sean, Weller, Noah, Marine, Zach, Cook, Wyatt, Wu, Carole-Jean, Liu, Gaoxiang
Format: Preprint
Veröffentlicht: 2025
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866912829767942144
author Liao, Gang
Qin, Hongsen
Wang, Ying
Golden, Alicia
Kuchnik, Michael
Yetim, Yavuz
Ang, Jia Jiunn
Fu, Chunli
He, Yihan
Hsia, Samuel
Jiang, Zewei
Li, Dianshi
Pashkevich, Uladzimir
Puvvada, Varna
Shi, Feng
Steiner, Matt
Xiao, Ruichao
Yan, Nathan
Yu, Xiayu
Fang, Zhou
Levenstein, Roman
Ho, Kunming
Zhu, Haishan
Hammond, Alec
Li, Richard
Mathews, Ajit
Gondkar, Kaustubh
Zainul-Abedin, Abdul
Singh, Ketan
Yu, Hongtao
Chi, Wenyuan
Huang, Barney
Zhang, Sean
Weller, Noah
Marine, Zach
Cook, Wyatt
Wu, Carole-Jean
Liu, Gaoxiang
author_facet Liao, Gang
Qin, Hongsen
Wang, Ying
Golden, Alicia
Kuchnik, Michael
Yetim, Yavuz
Ang, Jia Jiunn
Fu, Chunli
He, Yihan
Hsia, Samuel
Jiang, Zewei
Li, Dianshi
Pashkevich, Uladzimir
Puvvada, Varna
Shi, Feng
Steiner, Matt
Xiao, Ruichao
Yan, Nathan
Yu, Xiayu
Fang, Zhou
Levenstein, Roman
Ho, Kunming
Zhu, Haishan
Hammond, Alec
Li, Richard
Mathews, Ajit
Gondkar, Kaustubh
Zainul-Abedin, Abdul
Singh, Ketan
Yu, Hongtao
Chi, Wenyuan
Huang, Barney
Zhang, Sean
Weller, Noah
Marine, Zach
Cook, Wyatt
Wu, Carole-Jean
Liu, Gaoxiang
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2512_23236
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle KernelEvolve: Scaling Agentic Kernel Coding for Heterogeneous AI Accelerators at Meta
Liao, Gang
Qin, Hongsen
Wang, Ying
Golden, Alicia
Kuchnik, Michael
Yetim, Yavuz
Ang, Jia Jiunn
Fu, Chunli
He, Yihan
Hsia, Samuel
Jiang, Zewei
Li, Dianshi
Pashkevich, Uladzimir
Puvvada, Varna
Shi, Feng
Steiner, Matt
Xiao, Ruichao
Yan, Nathan
Yu, Xiayu
Fang, Zhou
Levenstein, Roman
Ho, Kunming
Zhu, Haishan
Hammond, Alec
Li, Richard
Mathews, Ajit
Gondkar, Kaustubh
Zainul-Abedin, Abdul
Singh, Ketan
Yu, Hongtao
Chi, Wenyuan
Huang, Barney
Zhang, Sean
Weller, Noah
Marine, Zach
Cook, Wyatt
Wu, Carole-Jean
Liu, Gaoxiang
Machine Learning
Artificial Intelligence
Hardware Architecture
Multiagent Systems
Performance
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.
title KernelEvolve: Scaling Agentic Kernel Coding for Heterogeneous AI Accelerators at Meta
topic Machine Learning
Artificial Intelligence
Hardware Architecture
Multiagent Systems
Performance
url https://arxiv.org/abs/2512.23236