KernelEvolve: 메타의 이기종 AI 가속기를 위한 에이전트 기반 커널 코딩 확장

초록

딥러닝 추천 모델(DLRM)의 학습과 추론을 빠르고 효율적으로 만드는 것은 중요합니다. 그러나 이는 세 가지 주요 시스템 과제를 제시합니다. 바로 모델 아키텍처의 다양성, 커널 프리미티브의 다양성, 그리고 하드웨어 세대 및 아키텍처의 이질성입니다. 본 논문은 DLRM을 위한 대규모 이질성 문제를 해결하기 위해 KernelEvolve라는 에이전트 기반 커널 코딩 프레임워크를 제안합니다. KernelEvolve는 커널 명세를 입력으로 받아 이기종 하드웨어 아키텍처 전반에 걸쳐 추천 모델을 위한 커널 생성 및 최적화 과정을 자동화하도록 설계되었습니다. KernelEvolve는 Triton 및 CuTe DSL부터 하위 수준의 하드웨어 독립적 언어에 이르기까지 여러 프로그래밍 추상화 수준에서 작동하여 전체 하드웨어-소프트웨어 최적화 스택을 아우름으로써 이를 실현합니다. 커널 최적화 과정은 선택 정책, 유니버설 연산자, 적합도 함수, 종료 규칙을 갖춘 그래프 기반 탐색으로 설명되며, 검색 증강 프롬프트 합성을 통해 런타임 실행 컨텍스트에 동적으로 적응합니다. 우리는 여러 세대의 NVIDIA 및 AMD GPU와 Meta의 AI 가속기를 아우르는 다양한 프로덕션 추천 모델을 최적화하기 위해 KernelEvolve를 설계, 구현 및 배포했습니다. 공개된 KernelBench suite에서 KernelEvolve를 검증한 결과, 세 가지 난이도에 걸친 250개 문제 전체와 세 가지 이기종 하드웨어 플랫폼의 160개 PyTorch ATen 연산자에 대해 100% 통과율을 달성하여 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: 메타의 이기종 AI 가속기를 위한 에이전트 기반 커널 코딩 확장

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

초록

Support