KernelEvolve : Mise à l'échelle du codage agentique des noyaux pour les accélérateurs d'IA hétérogènes chez Meta

papers.abstract

Rendre l'apprentissage et l'inférence des modèles de recommandation par apprentissage profond (DLRM) rapides et efficaces est crucial. Cependant, cela pose trois défis systémiques majeurs : la diversité des architectures de modèles, la diversité des primitives de noyaux, et l'hétérogénéité des générations de matériel et des architectures. Cet article présente KernelEvolve, un framework agentiel de codage de noyaux, pour relever le défi de l'hétérogénéité à grande échelle pour les DLRM. KernelEvolve est conçu pour prendre des spécifications de noyaux en entrée et automatiser le processus de génération et d'optimisation de noyaux pour les modèles de recommandation sur des architectures matérielles hétérogènes. Pour ce faire, KernelEvolve opère à plusieurs niveaux d'abstraction de programmation, des DSL Triton et CuTe aux langages de bas niveau indépendants du matériel, couvrant ainsi toute la pile d'optimisation matérielle-logicielle. Le processus d'optimisation des noyaux est décrit comme une recherche basée sur des graphes avec une politique de sélection, un opérateur universel, une fonction d'aptitude et une règle de terminaison, s'adaptant dynamiquement au contexte d'exécution grâce à une synthèse d'invites augmentée par retrieval. Nous avons conçu, implémenté et déployé KernelEvolve pour optimiser une grande variété de modèles de recommandation en production sur plusieurs générations de GPU NVIDIA et AMD, ainsi que sur les accélérateurs IA de Meta. Nous validons KernelEvolve sur la suite KernelBench, accessible publiquement, en atteignant un taux de réussite de 100% sur les 250 problèmes de trois niveaux de difficulté, et sur 160 opérateurs ATen de PyTorch sur trois plates-formes matérielles hétérogènes, démontrant une exactitude de 100%. KernelEvolve réduit le temps de développement de plusieurs semaines à quelques heures et permet des améliorations substantielles des performances par rapport aux lignes de base PyTorch dans divers cas d'usage en production et pour des systèmes IA hétérogènes à grande échelle. Au-delà des gains d'efficacité des performances, KernelEvolve atténue significativement la barrière de la programmabilité pour les nouveaux matériels IA en permettant la génération automatisée de noyaux pour les matériels IA développés en interne.

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 : Mise à l'échelle du codage agentique des noyaux pour les accélérateurs d'IA hétérogènes chez Meta

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

papers.abstract

Support