KernelEvolve: Metaにおける多様なAIアクセラレータ向けエージェント型カーネルコーディングのスケーリング

要旨

深層学習推薦モデル（DLRM）の学習と推論を高速かつ効率的にすることは重要である。しかし、これには3つの主要なシステム課題が存在する―モデルアーキテクチャの多様性、カーネルプリミティブの多様性、そしてハードウェア世代およびアーキテクチャの異種混合性である。本論文は、DLRMにおける大規模な異種混合性に対処するため、エージェント型カーネルコーディングフレームワーク「KernelEvolve」を提案する。KernelEvolveは、カーネル仕様を入力として受け取り、異種ハードウェアアーキテクチャにわたる推薦モデルのためのカーネル生成と最適化のプロセスを自動化するように設計されている。KernelEvolveは、TritonやCuTe DSLから低水準のハードウェア非依存言語まで、複数のプログラミング抽象化レベルで動作し、ハードウェアとソフトウェアの最適化スタック全体をカバーすることでこれを実現する。カーネル最適化プロセスは、選択ポリシー、普遍演算子、適合度関数、終了規則を備えたグラフベースの探索として記述され、検索強化型プロンプト合成を通じて実行時コンテキストに動的に適応する。我々はKernelEvolveを設計、実装、導入し、NVIDIAおよびAMDの複数世代のGPU、ならびにMetaのAIアクセラレータにわたる多種多様な本番推薦モデルを最適化した。公開されているKernelBenchスイートでKernelEvolveを検証した結果、3つの難易度レベルにわたる全250問題で100%の合格率を達成し、3つの異種ハードウェアプラットフォームにわたる160のPyTorch ATen演算子において100%の正確性を実証した。KernelEvolveは、開発時間を数週間から数時間に短縮し、多様な本番ユースケースおよび大規模な異種AIシステムにおいて、PyTorchベースラインを大幅に上回る性能向上を達成する。性能効率の改善を超えて、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.

KernelEvolve: Metaにおける多様なAIアクセラレータ向けエージェント型カーネルコーディングのスケーリング

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

要旨

Support