KernelEvolve: Масштабирование агентного кодирования ядер для гетерогенных AI-ускорителей в Meta
KernelEvolve: Scaling Agentic Kernel Coding for Heterogeneous AI Accelerators at Meta
December 29, 2025
Авторы: 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
Аннотация
Обеспечение высокой скорости и эффективности обучения и вывода рекомендательных моделей глубокого обучения (DLRM) является важной задачей. Однако это порождает три ключевые системные проблемы: разнообразие архитектур моделей, разнообразие ядерных примитивов, а также неоднородность поколений и архитектур аппаратного обеспечения. В данной статье представлен KernelEvolve — агентный фреймворк для программирования ядер, предназначенный для решения проблемы неоднородности в масштабе для DLRM. KernelEvolve разработан для приема спецификаций ядер на вход и автоматизации процесса генерации и оптимизации ядер для рекомендательных моделей на различных гетерогенных аппаратных архитектурах. Это достигается за счет работы на нескольких уровнях программных абстракций — от предметно-ориентированных языков (DSL) Triton и CuTe до низкоуровневых аппаратно-независимых языков, охватывая весь стек оптимизации программно-аппаратного обеспечения. Процесс оптимизации ядер описывается как графовый поиск с политикой выбора, универсальным оператором, функцией пригодности и правилом остановки, который динамически адаптируется к контексту выполнения во время работы через синтез промптов, дополненный поиском. Мы разработали, реализовали и внедрили KernelEvolve для оптимизации широкого спектра промышленных рекомендательных моделей на различных поколениях GPU NVIDIA и AMD, а также на AI-ускорителях Meta. Мы проверили работу KernelEvolve на общедоступном наборе тестов KernelBench, достигнув 100% прохождения всех 250 задач трех уровней сложности и 160 операторов PyTorch ATen на трех гетерогенных аппаратных платформах, что демонстрирует 100% корректность. KernelEvolve сокращает время разработки с недель до часов и обеспечивает значительное повышение производительности по сравнению с базовыми реализациями PyTorch в разнообразных промышленных сценариях использования и для масштабных гетерогенных AI-систем. Помимо улучшения эффективности производительности, KernelEvolve существенно снижает порог программируемости для нового AI-оборудования, обеспечивая автоматическую генерацию ядер для внутренне разрабатываемых AI-ускорителей.
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.