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

2025-12-29

AI Summary

Deep Learning Recommendation Models (DLRMs) face severe kernel development bottlenecks due to triple heterogeneity—diverse model architectures, operator primitives, and hardware platforms (NVIDIA/AMD GPUs and Meta’s custom accelerators). Method: This paper introduces the first proxy-based kernel programming framework that synergistically integrates graph search with retrieval-augmented generation (RAG) for prompt synthesis, enabling multi-level abstraction language co-optimization and dynamic runtime context adaptation. Built upon domain-specific languages (DSLs) such as Triton and CuTe, it supports hardware-agnostic operator modeling and fitness-driven automatic kernel generation. Contribution/Results: Evaluated on KernelBench (250 tests), the framework achieves 100% pass rate, correctly implements all 160 PyTorch ATen operators across platforms with 100% functional correctness, reduces kernel development time from weeks to hours, and consistently outperforms PyTorch baselines in performance.

Technology Category

• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Machine Learning: Kernel Methods
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

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.

Problem

• Automates kernel generation for diverse DLRM models on heterogeneous hardware
• Reduces development time from weeks to hours by optimizing kernel coding
• Mitigates programmability barriers for new AI accelerators through automation

Innovation

• Automated kernel generation for heterogeneous AI accelerators
• Graph-based search with retrieval-augmented prompt synthesis
• Multi-abstraction programming from DSL to hardware-agnostic languages