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

TL;DR

KernelEvolve presents an agentic, production-grade framework for automated kernel generation across heterogeneous AI accelerators used in Meta's DLRM workloads. It combines a graph-based search with retrieval-augmented prompting and a persistent knowledge base to synthesize, validate, and optimize Triton kernels across NVIDIA, AMD, and MTIA hardware, addressing both compute and data-preprocessing operators. The approach achieves 100% correctness across 480 operator-platform configurations and delivers 1.25×–17× speedups over PyTorch baselines on diverse production use cases, while dramatically reducing development time from weeks to hours. By enabling automated kernel generation for proprietary accelerators and integrating end-to-end evaluation, KernelEvolve demonstrates a scalable path toward automated, hardware-aware optimization in fragmented AI infrastructure.

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.

Figures (18)

• Figure 1: Meta's MTIA (Meta Training and Inference Accelerator) is a custom-designed chip optimized for AI workloads. The figure illustrates the MTIA hardware from four perspectives: its integration within a data center or server farm, highlighting the overall facility layout and environment; its deployment within a rack-mounted system for high-bandwidth applications; a close-up of the chip’s circuitry and board connections; and a detailed view of the chip core. Together, these images highlight MTIA’s advanced design, connectivity, and its role in enhancing the performance and efficiency of AI tasks across Meta’s platforms.
• Figure 2: Triton multi-target compilation architecture. Source code transforms through progressive MLIR lowering stages—platform-independent Triton-MLIR, hardware-specific GPU/AMDGPU/MTIA dialects, LLVM-IR—generating native binaries for NVIDIA (PTX/CUBIN), AMD (AMDGCN/HSACO) [amd2025triton], and MTIA (RISC-V) platforms.
• Figure 3: Triton overtakes CUDA as dominant kernel programming model at Meta. Left: Triton has grown to over 8,000 kernels, surpassing CUDA's stagnant legacy codebase, while emerging DSLs (CuTe, TLX, Helion) remain under 600. Right: Growth trajectories show Triton's 60% expansion rate driving this transition, with CuTe at 50% following November deployment. This shift toward higher-level DSLs—while maintaining legacy CUDA and introducing new abstraction (TLX)—creates programming model fragmentation across 5+ languages, motivating KernelEvolve's automated synthesis approach.
• Figure 4: KernelEvolve achieves 1.25-17× speedups across Meta LLMs and production use cases, spanning convolutional Transformers, data preprocessing operators, and recommendation systems, over heterogeneous AI hardware.
• Figure 5: KernelEvolve System Architecture (top) and Execution Workflow (down). KernelEvolve employs a self-improving state machine with tree search to explore and validate kernel optimizations. The system integrates evaluation tooling (accuracy, performance, profiling), specialized sub-agents for context management and deep search, and AI hardware interpreters for MTIA, GPU, and AMD platforms. An LLM synthesizer generates dynamic prompts, which are then used by external (Claude 4.5, GPT-5) or internal (Meta's CWM) LLM backends to generate Triton kernel candidates. Persistent storage includes a metadata store tracking execution scores and parent-child relationships in the search tree (connected to the object store via path references), an object store for kernel files, and a knowledge base that serves as a retrieval system for hardware constraints and optimization guidance to support LLM context augmentation.