CUDA Agent entraîne des LLM à écrire des noyaux GPU, surpassant torch.compile en vitesse

Introduction

L'équipe Seed de ByteDance, en collaboration avec l'Institut de recherche sur l'industrie de l'IA (AIR) de l'Université Tsinghua, a lancé CUDA Agent, un système d'apprentissage par renforcement à grande échelle pour agents intelligents, conçu pour apprendre aux modèles de langage à générer des kernels CUDA hautes performances.

Ce travail cible une faiblesse de longue date dans la génération de code de bas niveau par l'IA. Les modèles de langage de pointe parviennent souvent à générer du code CUDA compilable et fonctionnellement correct, mais pour les charges de travail GPU de production, la simple exactitude ne suffit pas. Les kernels générés doivent également être compétitifs face aux implémentations produites par des compilateurs hautement optimisés.

CUDA Agent est précisément construit autour de ce second problème : non seulement générer du code CUDA valide, mais aussi apprendre à analyser, valider et optimiser itérativement les kernels par profilage, jusqu'à obtenir des améliorations de performance significatives.

Pourquoi générer du code CUDA correct ne suffit pas

Avant l'introduction de l'apprentissage par renforcement pour agents, le modèle de base Seed1.6 démontrait déjà de solides capacités d'écriture CUDA sur KernelBench.

Sur les 250 tâches KernelBench de niveaux 1 à 3 utilisées par les chercheurs, Seed1.6 atteignait un taux de réussite de 74,0 %. En d'autres termes, il pouvait générer des solutions fonctionnellement correctes pour la majorité des tâches de référence du modèle de base.

Mais les performances étaient une tout autre histoire.

Seulement 27,2 % des kernels générés par le modèle de base étaient plus rapides que torch.compile, avec une vitesse géométrique moyenne de seulement 0,69× par rapport à la ligne de base du compilateur.

Cet écart met en évidence le principal défi de la génération automatique de kernels GPU.

Les modèles de langage peuvent comprendre la syntaxe CUDA, mais peuvent néanmoins produire des accès mémoire inefficaces, des lancements de kernels excessifs, de mauvais choix de tuilage, des intermédiaires inutiles...

Profilage des performances

• Directives d'optimisation spécifiques à CUDA
• Environnement sandbox restreint
• Signaux de récompense liés aux résultats d'exécution réels

Le fichier SKILL.md restreint également les raccourcis susceptibles de fausser les résultats des benchmarks. Par exemple, l'agent ne peut pas simplement recourir à des opérations PyTorch arbitraires dans ses implémentations de kernels personnalisés.

Cela est important car les systèmes d'apprentissage par renforcement pourraient trouver des moyens de maximiser la récompense sans résoudre le problème d'optimisation attendu.

L'entraînement PPO étend l'agent aux trajectoires d'optimisation longues

Les chercheurs utilisent l'optimisation de politique proximale (PPO) pour entraîner CUDA Agent.

Un défi majeur est que l'optimisation GPU est intrinsèquement une tâche à long horizon. Le modèle peut devoir traverser plusieurs cycles d'édition de code, compilation, débogage, profilage et optimisation supplémentaire avant d'atteindre le résultat souhaité.

Par conséquent, l'article adopte différentes longueurs de contexte et phases d'entraînement pour stabiliser l'apprentissage.

Le modèle de base est Seed1.6, un modèle à mélange d'experts avec 230 milliards de paramètres totaux et 23 milliards de paramètres activés.

Pour l'apprentissage par renforcement de l'agent :

• La fenêtre de contexte est de 131 072 tokens.
• La relecture d'entraînement permet jusqu'à 150 tours d'interaction de l'agent.
• L'évaluation permet jusqu'à 200 tours d'interaction de l'agent.
• Le modèle est entraîné sur 150 étapes de RL.
• Les modèles acteur et critique adoptent une stratégie de préchauffage multi-phases avant l'optimisation complète à long horizon.

Cela est plus complexe qu'un simple affinage du modèle sur des paires de requêtes et réponses CUDA.

L'objectif est d'apprendre au modèle comment améliorer ses propres kernels grâce à des interactions répétées avec le retour d'exécution.

Le sandbox sépare la compilation du profilage GPU

Une mesure fiable des performances est cruciale, car la vitesse d'exécution fait partie de la récompense.

Par conséquent, les chercheurs ont construit une architecture sandbox avec découplage des ressources CPU-GPU.

Un sandbox terminal basé sur Docker gère les tâches orientées CPU (comme la compilation), tandis que la validation et le profilage des performances sont répartis vers un pool de sandbox GPU dédié contenant 128 GPU NVIDIA H20.

Cette séparation a un double objectif.

Premièrement, elle isole la compilation et les interactions de l'agent des tests de référence GPU. Deuxièmement, l'allocation GPU dédiée réduit les interférences entre les charges de travail concurrentes, rendant les mesures de latence plus stables.

C'est important pour l'apprentissage par renforcement : si les mesures de timing sont bruitées, le modèle reçoit des signaux de récompense peu fiables et pourrait apprendre des comportements d'optimisation incorrects.

CUDA Agent atteint un taux de réussite de 98,8 % et dépasse le compilateur sur 96,8 % des tâches

Le système final a été évalué sur 250 tâches de KernelBench niveaux 1 à 3.

CUDA Agent atteint :

• Un taux de réussite global de 98,8 %
• 98,4 % plus rapide que l'exécution immédiate de PyTorch
• 96,8 % plus rapide que torch.compile
• Une accélération géométrique moyenne de 2,60× par rapport à l'exécution immédiate
• Une accélération géométrique moyenne de 2,11× par rapport à torch.compile

Ces résultats représentent une amélioration considérable par rapport au point de départ Seed1.6.

Le modèle non seulement s'améliore pour générer du code passant les tests de correction, mais ses kernels générés sont également plus susceptibles de surpasser les lignes de base du compilateur.

Résultats par niveau de difficulté KernelBench

Les gains de performance restent solides sur les trois niveaux de difficulté de KernelBench.

La performance au niveau 2 est particulièrement remarquable.

Ces tâches impliquent des séquences d'opérateurs où les opportunités d'optimisation par fusion et déplacements mémoire ne sont pas toujours efficacement exploitées par les heuristiques statiques des compilateurs.

Les auteurs de l'article estiment que l'optimiseur par apprentissage itératif peut explorer un espace plus large de stratégies d'implémentation, y compris les modes d'accès mémoire spécifiques au matériel, le tuilage et les choix de fusion.

Les données d'entraînement seront bientôt publiques

Actuellement, les poids complets du modèle entraîné ne sont pas inclus dans la version publique du projet.

Cependant, les chercheurs ont publié plusieurs composants importants de la pile d'entraînement.

CUDA-Agent-Ops-6K

Le jeu de données CUDA-Agent-Ops-6K contient 6 000 tâches d'entraînement synthétiques au niveau opérateur.

Son processus de construction comprend :

1. La collecte d'opérateurs sources à partir de bibliothèques comme torch, transformers, etc.
2. L'utilisation d'un LLM pour combiner plusieurs opérateurs en tâches de fusion plus complexes.
3. Le filtrage des tâches générées par des vérifications basées sur les résultats d'exécution.

La phase de filtrage élimine les cas présentant de l'aléatoire, trop simples, invalides, ou trop similaires aux tâches d'évaluation.

Ce jeu de données est publié sur Hugging Face sous licence CC BY 4.0.

SKILL.md et environnement agent

Le dépôt GitHub contient également le fichier d'instructions SKILL.md orienté CUDA ainsi que l'environnement de travail de l'agent.

Ces documents documentent le flux de travail d'optimisation, les outils disponibles, les contraintes et la configuration de l'environnement d'exécution utilisée pour guider l'agent.

Mécanisme de récompense et schéma d'entraînement

L'article et le dépôt décrivent la conception des récompenses, la phase de préchauffage, l'initialisation du critique et le schéma d'entraînement basé sur PPO pour stabiliser le processus d'optimisation à long cycle.

Cela est particulièrement utile pour les chercheurs intéressés par l'entraînement d'agents pour la programmation système, au-delà de la simple reproduction des scores finaux des benchmarks.

Pourquoi cela va au-delà de KernelBench

L'optimisation des kernels CUDA sous-tend une grande partie de l'infrastructure moderne de l'IA.

Des kernels plus rapides peuvent améliorer :

• Le débit d'entraînement des modèles
• Les performances de latence d'inférence des LLM
• L'utilisation des GPU
• L'efficacité des coûts de service
• Les pipelines IA en temps réel
• Les charges de calcul numérique hautes performances

L'article original d'AIBase souligne son importance dans l'IA...

Applications potentielles dans les domaines des infrastructures, des services d'inférence pour grands modèles, de la conduite autonome et du trading quantitatif — où de minuscules différences de latence peuvent s'avérer cruciales.

Les résultats de CUDA Agent ne signifient pas que les compilateurs traditionnels sont obsolètes.

torch.compile reste une baseline automatique puissante, et l'évaluation de CUDA Agent repose elle-même sur un environnement de benchmark contrôlé et un environnement GPU spécifique.

Les conclusions démontrées par cette recherche, bien que plus restreintes, n'en demeurent pas moins importantes : avec un retour d'exécution suffisant et un apprentissage par renforcement dédié, les modèles de langage peuvent apprendre des stratégies d'optimisation qui surpassent la baseline du compilateur statique sur la grande majorité des tâches de kernels GPU testées.

Cela implique de transformer l'ingénierie de performance sous-jacente en un domaine pertinent pour l'apprentissage par agent, et non plus seulement en une génération de code en une seule passe.

Questions fréquentes

Qu'est-ce que CUDA Agent ?

CUDA Agent est un système d'apprentissage par renforcement multi-agents à grande échelle développé par ByteDance Seed, l'Institut de recherche sur l'industrie intelligente (AIR) de l'Université Tsinghua et leur SIA-Lab conjoint. Il entraîne des modèles de langage à écrire, valider, analyser et optimiser itérativement des kernels CUDA.

Sur quel modèle CUDA Agent est-il basé ?

La recherche utilise Seed1.6 comme modèle de base. L'article le décrit comme un modèle à mélange d'experts avec 230 milliards de paramètres totaux et 23 milliards de paramètres activés.

Qu'est-ce que KernelBench ?

KernelBench est un benchmark ouvert qui évalue si les modèles de langage peuvent générer des kernels GPU corrects et efficaces pour des programmes PyTorch. CUDA Agent a été évalué sur 250 tâches couvrant les niveaux 1 à 3 de KernelBench.

De combien CUDA Agent est-il plus rapide que torch.compile ?

Sur l'ensemble du benchmark, CUDA Agent atteint une accélération géométrique moyenne de 2,11 fois par rapport à torch.compile. Ses kernels générés sont plus rapides que la baseline du compilateur sur 96,8 % des tâches évaluées.

CUDA Agent utilise-t-il l'apprentissage par renforcement ?

Oui. Le système utilise PPO, combiné à un préchauffage multi-phases, un pré-entraînement du modèle de valeur, une conception de récompense robuste et des trajectoires d'agent multi-tours de longue durée.

Quelle est la durée maximale d'une trajectoire d'optimisation de CUDA Agent ?

Le roll-out d'entraînement permet jusqu'à 150 tours d'agent, tandis que l'évaluation permet jusqu'à 200 tours. La fenêtre de contexte d'entraînement de l'agent est de 131 072 tokens.

CUDA Agent est-il open source ?

Le projet a publié son dépôt GitHub, l'environnement d'agent, SKILL.md, la recette d'entraînement et le jeu de données CUDA-Agent-Ops-6K. Les poids complets du modèle entraîné ne figurent pas dans les éléments publiés publiquement actuellement.

Qu'est-ce que CUDA-Agent-Ops-6K ?

CUDA-Agent-Ops-6K est un jeu de données de 6 000 tâches synthétiques d'entraînement CUDA au niveau des opérateurs. Il est spécialement conçu pour l'apprentissage par renforcement multi-agents à grande échelle et comprend des étapes de filtrage pour garantir que les tâches sont exécutables, déterministes, non triviales et séparées de l'ensemble d'évaluation KernelBench.

Outils associés

• CUDA Agent : Page officielle du projet du système de génération de kernels CUDA de ByteDance Seed et de l'Institut de recherche sur l'industrie intelligente de l'Université Tsinghua.
• CUDA Agent GitHub : Dépôt officiel contenant l'environnement d'agent, SKILL.md et les matériaux de projet publiés.
• CUDA-Agent-Ops-6K : Jeu de données officiel d'entraînement à 6 000 échantillons publié avec le projet.
• KernelBench : Benchmark ouvert pour évaluer les kernels GPU générés par les grands modèles de langage.
• PyTorch : Framework d'apprentissage profond dont le mode eager et l'exécution torch.compile sont des baselines clés dans la recherche.
• NVIDIA CUDA : Documentation officielle de NVIDIA sur la programmation CUDA et le développement de kernels GPU.

Liens associés

• Page du projet CUDA Agent : Présentation officielle, résultats de benchmark, graphiques du projet, liens vers l'article, le code et les jeux de données.
• Article CUDA Agent (arXiv) : Article de recherche complet, « Apprentissage par renforcement multi-agents à grande échelle pour la génération de kernels CUDA haute performance. »
• Dépôt GitHub CUDA Agent : Dépôt source de l'environnement publié et des fichiers du projet.
• Jeu de données CUDA-Agent-Ops-6K : Jeu de données Hugging Face officiel comprenant 6 000 tâches d'opérateurs synthétiques.
• Dépôt GitHub KernelBench : Environnement de benchmark et d'évaluation pour la génération efficace de kernels GPU.
• Document officielle PyTorch torch.compile : Baseline du compilateur utilisée dans la recherche.
• Guide de programmation NVIDIA CUDA C++ : Guide officiel de NVIDIA pour la programmation et l'optimisation CUDA.

Résumé

CUDA Agent résout une faiblesse spécifique du code système généré par les grands modèles de langage : le code CUDA généré doit non seulement être correct, mais aussi réellement plus rapide que les alternatives générées par compilateur.

À partir de Seed1.6, les chercheurs ont utilisé un environnement d'agent spécialisé CUDA, un retour d'exécution, une conception de récompense robuste et un entraînement PPO à long horizon pour faire passer le taux de réussite de 74,0 % à 98,8 %, et la proportion de kernels plus rapides que torch.compile de 27,2 % à 96,8 %.

Le projet publie également un jeu de données d'entraînement de 6 000 tâches, SKILL.md, des matériaux d'environnement d'agent et une recette d'entraînement, offrant aux chercheurs une base concrète pour poursuivre les travaux sur l'optimisation GPU apprise.

La principale contribution de CUDA Agent est de démontrer que l'ingénierie de performance elle-même peut être apprise par une boucle d'agent itérative — et non pas seulement approximée par une génération de code en une seule passe.