KernelArc: A Multi-Agent Framework for GPU Kernel Optimization
KernelArc est un framework multi-agents qui optimise de manière autonome les noyaux GPU grâce à des agents parallèles spécialisés par stratégie se coordonnant via une mémoire partagée et des gardes déterministes, atteignant les premiers rangs du classement SOL-ExecBench pour divers types de charges de travail sur les GPU NVIDIA H100 et B200.
Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète
Les ordinateurs modernes s'appuient sur des puces spécialisées appelées unités de traitement graphique, ou GPU, pour gérer les calculs massifs requis par l'intelligence artificielle. Ces puces sont incroyablement puissantes, mais ce sont aussi des machines complexes avec de nombreuses parties différentes qui doivent fonctionner en parfaite synchronisation. Pour tirer le meilleur parti d'elles, les ingénieurs doivent écrire de minuscules programmes très spécifiques appelés « kernels » (noyaux) qui indiquent à la puce exactement comment déplacer les données et effectuer les calculs mathématiques. Pendant des années, cela a été le travail d'experts humains qui passent des heures innombrables à peaufiner le code, à équilibrer l'utilisation de la mémoire et à synchroniser les opérations pour extraire chaque once de vitesse. À mesure que ces puces sont devenues plus sophistiquées, le travail manuel est devenu plus difficile, et l'écart entre ce que le matériel peut faire et ce que le logiciel réalise s'est creusé.
Des chercheurs ont commencé à utiliser des modèles de langage étendus — des programmes informatiques entraînés sur de vastes quantités de texte — pour aider à écrire et à améliorer ce code. Ces modèles peuvent suggérer des modifications, les tester et apprendre des résultats, agissant comme un ingénieur automatisé. Cependant, un agent automatisé unique finit souvent par stagner. Il peut trouver une bonne solution et continuer à la polir, manquant ainsi d'autres approches meilleures qui se trouvent dans une direction différente. C'est comme un randonneur qui trouve un joli sentier pour monter une montagne et continue de marcher sur celui-ci, sans réaliser qu'un autre sentier à proximité mène à un sommet bien plus élevé. Pour résoudre ce problème, une équipe de chercheurs d'IMEC a développé un nouveau système appelé KernelArc, qui utilise un groupe de ces agents automatisés travaillant ensemble pour explorer de nombreux chemins à la fois.
Le système KernelArc fonctionne en assignant différentes stratégies à plusieurs agents s'exécutant en parallèle. Au lieu qu'un seul agent tente de résoudre un problème seul, chaque agent du groupe se concentre sur un angle d'attaque différent, tel que modifier la façon dont les données sont stockées, altérer la précision mathématique utilisée ou fusionner plusieurs étapes en une seule. Ces agents ne partagent pas l'intégralité de leur processus de pensée ni chaque étape qu'ils franchissent, ce qui serait trop désordonné et confus. Au lieu de cela, ils ne partagent que leurs conclusions finales : ce qui a fonctionné et ce qui a échoué. Ils écrivent ces résultats dans un espace de mémoire partagée qui agit comme un tableau d'affichage. Si un agent découvre une astuce qui accélère un calcul, il publie le résultat. Les autres agents peuvent lire cette publication et utiliser l'enseignement pour guider leur propre travail, évitant ainsi les impasses et s'appuyant sur le succès.
Pour s'assurer que les agents ne perdent pas de temps sur du code défectueux, le système inclut un arbitre automatisé strict. Cet arbitre soumet chaque nouvelle suggestion de code à une série de tests pour vérifier si elle est correcte et quelle est sa vitesse. Si une suggestion échoue au test, elle est immédiatement écartée. Si elle fonctionne mais n'est pas plus rapide que la meilleure version actuelle, elle est conservée comme sauvegarde mais ne remplace pas le leader. Ce n'est que lorsqu'une nouvelle version est à la fois correcte et plus rapide qu'elle devient le nouveau standard. Ce processus permet à l'équipe d'agents d'explorer un large éventail de possibilités sans se perdre dans les détails de chaque tentative. Le système dispose également d'un mécanisme de sécurité qui s'active si un agent reste bloqué sur un plateau, là où plus aucune amélioration ne peut être trouvée. Lorsque cela arrive, le système force l'agent à essayer une approche complètement différente, garantissant que la recherche continue d'avancer.
Les chercheurs ont testé ce système sur deux des GPU les plus avancés disponibles, le NVIDIA H100 et le B200. Ils se sont concentrés sur un ensemble de tâches standards utilisées pour mesurer la performance, qui incluent des opérations comme la multiplication de matrices, les mécanismes d'attention utilisés dans les modèles de langage et diverses formes de fusion de données. Dans un test spécifique impliquant un agent travaillant seul avec un guide détaillé, le système a réussi à atteindre une vitesse de 766 téraflops, soit environ 3,2 % de plus que la meilleure bibliothèque existante pour cette tâche spécifique. Cela a montré qu'un agent unique pouvait approfondir considérablement un problème spécifique s'il disposait d'un chemin clair. Cependant, lorsque les chercheurs sont passés au système multi-agents KernelArc pour aborder un ensemble de tâches plus large, les résultats ont été encore plus impressionnants. Le système a produit des implémentations personnalisées pour diverses opérations complexes, incluant des mécanismes d'attention spécialisés et des couches fusionnées pour les grands modèles de langage.
Sur un classement public qui hiérarchise la vitesse de ces opérations à travers de nombreuses formes et tailles de données, les soumissions de KernelArc ont pris la première place dans plusieurs catégories. Pour une tâche impliquant l'attention et l'addition résiduelle, un agent seul est resté bloqué à un score de performance de 0,441. Le système multi-agents, en partageant les enseignements et en explorant différentes directions, a brisé cette barrière pour atteindre un score de 0,481. Dans un autre test impliquant une tâche d'attention complexe, le système multi-agents a obtenu une accélération de près de 291 fois par rapport à une implémentation de référence standard, et plus de 143 fois plus vite qu'une base de référence hautement optimisée. Ces résultats suggèrent que le fait d'avoir plusieurs agents partageant leurs découvertes permet au système de trouver de meilleures solutions plus rapidement qu'un agent seul, surtout lorsque l'espace des problèmes est vaste et varié.
L'étude a également examiné le comportement du système lorsque le nombre d'agents et la quantité de mémoire partagée changent. Ils ont constaté que donner aux agents une mémoire partagée où ils pouvaient publier leurs succès et leurs échecs les aidait à atteindre des résultats plus solides dans un nombre fixe de tentatives. Lorsque la mémoire était illimitée, le système a obtenu les meilleures performances, atteignant une accélération de plus de 290 fois par rapport au point de départ. Cela indique que la capacité à apprendre de l'expérience collective du groupe est un facteur clé de succès. Les chercheurs ont noté que la valeur de chaque caractéristique, telle que la mémoire partagée ou la spécialisation des stratégies, dépend de la tâche spécifique et de l'étape de la recherche. Parfois, les agents doivent explorer largement, et d'autres fois, ils doivent affiner une solution spécifique.
Ce travail ne prétend pas avoir résolu tous les problèmes d'optimisation informatique, ni suggère qu'on n'a plus besoin d'ingénieurs humains. Les résultats sont spécifiques aux tâches et au matériel testés, et le système repose toujours sur un cadre conçu par l'homme pour guider les agents. Cependant, les conclusions démontrent qu'un groupe coordonné d'agents automatisés peut explorer une gamme de solutions plus large qu'un agent unique. En ne partageant que leurs conclusions les plus précieuses, ces agents peuvent éviter de répéter les erreurs et construire sur les découvertes des autres. Cette approche offre une voie prometteuse pour gérer la complexité croissante de l'informatique moderne, où l'écart entre le potentiel du matériel et la performance du logiciel continue de croître. Le système montre qu'avec une coordination adéquate, les outils automatisés peuvent aider à libérer toute la puissance des puces qui pilotent la prochaine génération d'intelligence artificielle.
Noyé(e) sous les articles dans votre domaine ?
Recevez des digests quotidiens des articles les plus récents correspondant à vos mots-clés de recherche — avec des résumés techniques, dans votre langue.