PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX

Au sein de l'ordinateur moderne, le processeur graphique, ou GPU, agit comme un moteur puissant conçu pour traiter de masses considérables de données simultanément. Pour tirer le meilleur parti de ce moteur, les développants de logiciels écrivent des instructions spécialisées qui indiquent au matériel exactement comment déplacer les données et effectuer des calculs. Ces instructions sont écrites dans un langage de bas niveau appelé PTX, qui donne au programmeur un contrôle direct sur les rouages internes de la machine. Bien que des outils de plus haut niveau existent pour faciliter ce processus, ils passent parfois à côté des fonctionnalités uniques des dernières générations de matériel. Alors que ces machines évoluent rapidement, maintenir un logiciel parfaitement ajusté à leurs capacités spécifiques devient un défi difficile. La question se pose : l'intelligence artificielle, plus précisément les grands modèles de langage, peut-elle apprendre à écrire ces instructions précises et de bas niveau par elle-même, ou repose-t-elle encore sur des méthodes anciennes et génériques qui sacrifient la performance ?

Une équipe de chercheurs s'est donné pour mission de répondre à cette question en créant un nouveau terrain d'essai appelé PTXBench. Ils voulaient voir si ces modèles d'IA pouvaient non seulement écrire du code qui fonctionne, mais aussi du code qui utilise activement les fonctionnalités spécifiques et avancées des derniers GPU. Les chercheurs se sont concentrés sur deux des cartes graphiques les plus puissantes disponibles, la H100 et la B200, et ont testé les modèles sur des tâches impliquant des opérations mathématiques complexes utilisées dans l'intelligence artificielle, telles que la multiplication de matrices et les mécanismes d'attention. L'objectif était de mesurer trois choses : si le code s'exécutait sans erreur, s'il déclenchait réellement les instructions matérielles spécifiques demandées par les chercheurs, et s'il s'exécutait plus rapidement que les meilleures bibliothèques logicielles existantes.

Les résultats ont révélé un écart significatif entre ce que les modèles pouvaient faire et ce qui était requis pour une performance de haut niveau. Si les modèles d'IA réussissaient souvent à écrire du code fonctionnant correctement pour des tâches simples, ils éprouvaient des difficultés considérables avec des opérations plus complexes, particulièrement celles impliquant la passe arrière (backward pass) des mécanismes d'attention, qui sont cruciaux pour l'entraînement des systèmes d'IA. Même lorsque les modèles parvenaient à écrire du code exécutant les instructions matérielles demandées, ce code égalait rarement la vitesse des principales bibliothèques professionnelles. Dans de nombreux cas, le code généré par l'IA était plus lent, suggérant que le simple fait de faire exécuter les instructions ne suffisait pas à atteindre une véritable optimisation. L'étude a révélé que les modèles revenaient fréquemment à des modèles de codage génériques plutôt que de tirer parti des caractatures uniques et spécifiques à l'architecture des nouvelles puces, même lorsque ces fonctionnalités leur étaient explicitement décrites.

Pour comprendre comment améliorer ces modèles, les chercheurs ont mené une expérience contrôlée où ils ont enseigné à un modèle d'IA spécifique une méthode appelée fine-tuning supervisé. Ils ont montré au modèle des exemples de ses propres erreurs, accompagnés de versions corrigées du code et d'explications sur les raisons pour lesquelles les corrections fonctionnaient. Cette approche, qu'ils ont nommée Fixit, a aidé le modèle à s'améliorer sur plusieurs tâches, particulièrement dans la génération de code fonctionnellement correct. Cependant, l'amélioration n'était pas uniforme. Le modèle avait toujours du mal à généraliser ses nouvelles compétences à différentes tailles de problèmes ou à des variations de tâches qu'il n'avait pas rencontrées lors de l'entraînement. Les chercheurs ont découvert que la qualité des corrections et l'équilibre des données d'entraînement importaient tout autant que la quantité brute de données. Un modèle entraîné sur un ensemble diversifié de problèmes était plus performant qu'un modèle entraîné sur un ensemble plus vaste mais moins varié, indiquant que le type d'expérience d'apprentissage était plus important que le volume.

L'étude a également souligné l'importance de fournir le bon contexte à l'IA. Lorsque les chercheurs donnaient aux modèles des informations détaillées sur les capacités spécifiques du matériel, les modèles étaient beaucoup plus susceptibles d'utiliser les instructions demandées. Sans cette connaissance spécifique, le modèle utilise rarement le PTX demandé. Cela suggère que, bien que l'IA puisse apprendre à programmer à ce bas niveau, elle nécessite des directives précises et des exemples de haute qualité pour le faire efficacement. Les chercheurs ont conclu que, bien que les modèles d'IA actuels soient prometteurs, ils ne sont pas encore prêts à remplacer les ingénieurs humains experts dans l'optimisation du code pour le matériel le plus récent. La voie à suivre implique de meilleures données d'entraînement, des expériences d'apprentissage plus équilibrées et une compréhension plus profonde de la manière d'enseigner à ces modèles comment raisonner sur les contraintes spécifiques des architectures informatiques en constante évolution.