AccelOpt : Un système agentique à LLM auto-améliorant pour l'optimisation de noyaux d'accélérateurs d'IA

Résumé

Nous présentons AccelOpt, un système agentique de grands modèles de langage (LLM) auto-améliorant qui optimise de manière autonome les noyaux pour les accélérateurs d'IA émergents, éliminant le besoin d'une expertise d'optimisation matérielle spécifique fournie par des experts. AccelOpt explore l'espace d'optimisation des noyaux par une génération itérative, guidée par une mémoire d'optimisation qui conserve les expériences et les enseignements tirés de paires de noyaux lents-rapides précédemment rencontrées. Nous avons développé NKIBench, une nouvelle suite de référence de noyaux pour l'accélérateur AWS Trainium, de complexité variable, extraits de charges de travail réelles de LLM, afin d'évaluer l'efficacité d'AccelOpt. Notre évaluation confirme que la capacité d'AccelOpt s'améliore avec le temps, augmentant le pourcentage moyen du débit de crête de 49 % à 61 % sur Trainium 1 et de 45 % à 59 % sur Trainium 2 pour les noyaux NKIBench. De plus, AccelOpt est très rentable : en utilisant des modèles open source, il atteint les mêmes améliorations de noyaux que Claude Sonnet 4 tout en étant 26 fois moins cher. Le code est open-source à l'adresse https://github.com/zhang677/AccelOpt.

English

We present AccelOpt, a self-improving large language model (LLM) agentic system that autonomously optimizes kernels for emerging AI acclerators, eliminating the need for expert-provided hardware-specific optimization knowledge. AccelOpt explores the kernel optimization space through iterative generation, informed by an optimization memory that curates experiences and insights from previously encountered slow-fast kernel pairs. We build NKIBench, a new benchmark suite of AWS Trainium accelerator kernels with varying complexity extracted from real-world LLM workloads to evaluate the effectiveness of AccelOpt. Our evaluation confirms that AccelOpt's capability improves over time, boosting the average percentage of peak throughput from 49% to 61% on Trainium 1 and from 45% to 59% on Trainium 2 for NKIBench kernels. Moreover, AccelOpt is highly cost-effective: using open-source models, it matches the kernel improvements of Claude Sonnet 4 while being 26times cheaper. The code is open-sourced at https://github.com/zhang677/AccelOpt.