KernelEvolve: Scaling Agentic Kernel Coding for Heterogeneous AI Accelerators at Meta (2512.23236v2)

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.

• The paper introduces an agentic kernel coding system that leverages LLM-driven synthesis and multi-phase tree search to reduce total cost of ownership and accelerate production workloads.
• It employs dynamic retrieval-based prompt engineering and persistent metadata stores to achieve 100% correctness across 480 operator-platform combinations.
• Empirical results demonstrate up to 17× speedup for production models, enabling seamless deployment on heterogeneous AI accelerators without LLM retraining.

Agentic Kernel Generation for Heterogeneous AI Accelerators: Review of "KernelEvolve"

Introduction and Motivation

"KernelEvolve: Scaling Agentic Kernel Coding for Heterogeneous AI Accelerators at Meta" (2512.23236) presents black, an agentic AI-powered kernel generation and optimization system deployed for production-scale recommendation model serving and training. The work addresses the combinatorial complexity intrinsic to modern ML infrastructure: diverse model architectures, richly varied kernel primitives, and rapidly proliferating hardware accelerators spanning Meta's custom MTIA, NVIDIA, and AMD platforms.

The paper provides a comprehensive characterization of both the economic and architectural imperatives driving kernel optimization in high-throughput environments. Marginal performance improvements directly yield multi-million dollar TCO reduction and enable monolithic serving architectures by removing critical deployment blockers. Notably, the paper identifies that missing preprocessing kernel implementations create binary constraints on production rollout, introducing substantial architectural penalties (network latency, reliability degradation) far surpassing the incremental costs due to suboptimal GEMM primitives.

System Architecture

black employs a persistent, self-improving state machine manifest as a tree search algorithm over the kernel solution space, integrating LLM-driven kernel synthesis with multi-phase, retrieval-augmented prompting. The system's architecture is detailed in

Figure 2: Agentic kernel generation system architecture, featuring knowledge-augmented LLM prompt synthesis, specialized hardware interpreters, persistent storage for artifacts and metadata, and multi-agent search orchestration.

Key features include:

• Dynamic retrieval-based prompt engineering: Specialized sub-agents extract relevant hardware constraints, optimization guidance, and historical profiling data from a hierarchical knowledge base, circumventing LLM training limitations (particularly for proprietary accelerators such as MTIA).
• Persistent metadata/object stores: The execution context for each tree node (kernel candidate) is persistently tracked, enabling distributed concurrent exploration, complex contextual queries, automatic checkpointing, and historical knowledge leverage for warm starts and fault tolerance.
• Unified evaluation and profiling infrastructure: Automated correctness validation, multi-granularity profiling (system, kernel, intra-kernel), interpreter-based execution on target accelerators, and streamlined deployment through continuous integration.
• Universal operator abstraction: The search process utilizes a single context-augmented operator, supporting holistic kernel improvements (correctness, performance, architectural tuning) unconstrained by fixed operator semantics, a critical advance over traditional Debug/Improve prompt splitting.

Kernel Synthesis and Optimization Methodology

The system formalizes kernel generation as graph-based search $(\mathcal{F}, \pi_{\text{sel}}, \mathcal{O}, \tau)$, with $\mathcal{F}$ as the fitness function (speedup over PyTorch baseline), $\pi_{\text{sel}}$ providing search node selection, $\mathcal{O}$ as the universal transformation operator conditioned by runtime profiling and retrieved knowledge, and $\tau$ defining termination. This abstraction supports not only greedy and MCTS strategies but also evolutionary algorithms for population-based search.

Execution feedback (correctness, profiling), runtime introspection, and dynamic context augmentation guide optimization, enabling automatic discovery of hardware-specific tuning strategies and fusion opportunities otherwise requiring extensive expert engineering.

Empirical Results: Correctness and Production Speedup

The paper reports rigorous validation across 160 ATen operators on three hardware platforms (NVIDIA, AMD, MTIA), achieving 100% correctness in 480 operator-platform combinations and full pass rate on KernelBench difficulty levels. Fitness trajectories indicate consistent improvement via feedback-driven search:

Figure 4: Fitness score trajectories for 6 ATen operators, reflecting feedback-driven refinement and optimization convergence.

Production-level use cases demonstrate that black can autonomously synthesize kernels for complex, multi-operator workloads, achieving speedups ranging from 1.2$\times$ to 17$\times$ over PyTorch references. These include:

• Optimized FM for recommendation models: Operator fusion and bespoke tiling provide 2--4$\times$ speedup for feature counts $N \leq 64$; speedup smoothly degrades for larger $N$ where tiling overhead dominates.

Figure 6: Speedup analysis for Wukong-style Optimized FM kernels, with robustness for production feature counts and degradation at large $N$ due to memory tiling overhead.

• InterFormer Personalized FeedForward Network (PFFN): Single-pass kernel generation with cross-operator tile reuse achieves peak speedups of 2.0--2.4$\times$ for small batches, converging to 1.2--1.4$\times$ as kernel launch overhead is amortized.

Figure 8: PFFN kernel speedup as a function of batch size and input dimension, highlighting SRAM-dependent non-monotonic scaling.

• MTIA-specific preprocessing (MBDT): Data preprocessing operator fusion and register/SIMD optimizations yield speedups ranging from 2.31--9.25$\times$ depending on batch/feature/border scaling, with marked gains on larger batch sizes.

Figure 1: MBDT kernel latency comparison, showing strong scaling of speedup with batch size across MTIA hardware generations.

The system supports shape-specific dispatch to avoid regressions for non-representative input shapes, ensuring safe production deployment.

Practical and Theoretical Implications

The work establishes several critical advances:

• Production-level heterogeneous enablement: Automated kernel synthesis bridges the gap between accelerator deployment and software ecosystem maturity. Proprietary architectures (MTIA) can be targeted without LLM retraining via knowledge injection.
• Elimination of architectural deployment blockers: Comprehensive kernel coverage enables unified monolithic accelerator topologies, eliminating latency and reliability penalties imposed by disaggregated preprocessing tiers.
• Scalability of kernel development: Weeks-to-hours reduction in kernel enablement time is achieved, critical given the frequency of model and hardware churn in large-scale enterprise systems.
• Inference-time scaling and distributed search: The design supports massive, parallel exploration of the kernel search space, suggesting that aggregate compute investment can predictably yield kernel quality improvements, opening avenues for further scaling research.

Theoretically, black demonstrates that LLM-agentic systems, when coupled with persistent, retrieval-augmented knowledge engineering and multi-phase profiling, can approach or exceed expert-level performance for deeply specialized, hardware-bound programming tasks. The universal operator abstraction enables contextually optimal reasoning unconstrained by pre-specified prompt boundaries.

Future Prospects and Speculation

The authors articulate several key trajectories for future research:

• Scale to new hardware and richer abstractions: Enabling rapid adaptation to emerging accelerator classes (ARM, next-generation GPUs, ASICs).
• Model-wide, cross-layer fusion and full-graph optimization: Integration with compiler suites (AOT Inductor), model serving stacks, and end-to-end continuous deployment.
• Vertical stack integration: Extending beyond the Triton DSL to MLIR, PTX, diagnostic languages for ultimate performance tuning.
• Sustainable, efficient agentic inference: Quantification and reduction of token consumption and carbon footprint in large-scale inference-driven kernel synthesis.

Conclusion

KernelEvolve (black) introduces an agentic kernel coding framework that fundamentally transforms kernel development for heterogeneous AI accelerators. The system achieves 100% correctness across a wide operator-platform matrix, delivers robust speedups (up to 17$\times$ on production workloads), and dramatically reduces time-to-solution for novel hardware targets. By systematizing context retrieval, profiling integration, and universal operator agentics, black addresses both the software ecosystem maturity gap and architectural constraints of state-of-the-art AI serving infrastructures.

This work points to the emergence of universal agentic compilation layers, where foundation model agents, augmented by dynamic knowledge injection and distributed optimization, continuously generate, validate, and deploy kernels as hardware and model architectures rapidly evolve. The established design principles—persistent search, multi-granularity profiling, hardware-aware context augmentation—will be essential as the space of accelerators, models, and associated kernels continues to expand.

---