自动化生成适配多硬件的推荐模型加速核,大幅缩短开发周期。
KernelEvolve: Scaling Agentic Kernel Coding for Heterogeneous AI Accelerators at Meta
- 用图搜索+动态提示合成实现跨平台内核自动优化
- 在250个测试问题上100%通过率,3种硬件上160个算子正确率100%
- 适合需要快速部署新硬件的工业级AI系统团队
深度学习推荐模型(DLRM)的训练与推理需高效快速,但面临模型架构多样、内核原语差异及硬件代际与架构异构三大挑战。本文提出KernelEvolve——一个代理式内核编码框架,用于在大规模异构AI加速器上解决上述问题。该框架以内核规范为输入,自动完成推荐模型在多种硬件上的内核生成与优化。其通过多层级编程抽象(从Triton、CuTe DSL到硬件无关低层语言),覆盖全栈软硬件优化。内核优化过程被建模为基于图的搜索,包含选择策略、通用算子、适应性评估函数与终止规则,并通过检索增强提示合成动态适配运行时上下文。我们设计、实现并部署了KernelEvolve,在NVIDIA和AMD GPU以及Meta自研AI加速器上优化了多种生产级推荐模型。在公开的KernelBench基准测试中,对250个不同难度的问题达成100%通过率,对160个PyTorch ATen算子在三种异构硬件平台上实现100%正确性验证。相比PyTorch基线,KernelEvolve将开发时间从数周缩短至数小时,并在多样化的实际场景中显著提升性能表现。此外,它有效降低了新硬件的可编程门槛,实现了对自研硬件的自动化内核生成。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。