PTXBench : une nouvelle façon de mesurer la capacité des LLM à optimiser les noyaux GPU

3 septembre 202617 vues

Le benchmark évalue les modèles sur des tâches GEMM et d'attention pour les GPU H100 et B200 — de la correction à l'amélioration réelle de la vitesse. Les résultats montrent que même une exécution précise des instructions ciblées ne garantit pas des performances compétitives.

PTXBench : une nouvelle façon de mesurer la capacité des LLM à optimiser les noyaux GPU

L'optimisation des cœurs GPU nécessite généralement une compréhension approfondie de l'architecture matérielle et des instructions de bas niveau. Les modèles de langage gèrent cela de manière encore instable : ils peuvent générer un code plausible, mais il n'est pas toujours correct ni rapide. PTXBench est un nouveau benchmark qui tente d'apporter de la mesurabilité et de la clarté à ce processus.

Qu'est-ce que PTXBench et pourquoi en a-t-on besoin ?

PTXBench est un banc d'essai d'évaluation pour les grands modèles de langage qui tentent d'utiliser le PTX (Parallel Thread Execution) spécifique à l'architecture pour optimiser les cœurs GPU. Le PTX est une représentation intermédiaire des instructions pour les GPU de NVIDIA, et c'est à ce niveau que l'on peut contrôler finement les registres, la mémoire et le planificateur. Les benchmarks classiques ne vérifient que le code final, mais PTXBench regarde plus loin : le modèle applique-t-il réellement les instructions de bas niveau ciblées, ou se contente-t-il de générer un code « ressemblant à du C » ?

Les auteurs se sont concentrés sur deux charges de travail classiques : la multiplication de matrices (GEMM) et le mécanisme d'attention, essentiel pour les transformeurs. Les tests ont été effectués sur des accélérateurs récents — H100 et B200. Ce choix permet de vérifier non seulement les compétences de base du modèle, mais aussi sa capacité à s'adapter aux architectures modernes et à leurs spécificités.

Que mesure le benchmark ?

PTXBench évalue les modèles selon trois paramètres clés :

  • Correction fonctionnelle — le résultat du noyau généré doit correspondre à la référence.
  • Exécution des instructions ciblées — on vérifie si les instructions PTX que le modèle était censé utiliser apparaissent bien à l'exécution.
  • Accélération — dans quelle mesure le noyau obtenu est plus rapide que les bibliothèques frontier, c'est-à-dire les meilleures implémentations prêtes à l'emploi.

Cette approche sépare trois problèmes distincts. Un modèle peut écrire un code correct sans pour autant utiliser les instructions requises. Ou bien les utiliser, mais obtenir des performances inférieures à celles d'une bibliothèque standard. PTXBench permet de voir à quelle étape exacte l'échec se produit.

Résultats : l'optimisation reste inégale pour les modèles

L'évaluation a montré que la capacité des LLM à travailler avec le PTX dépend fortement de la tâche. Sur les cas simples, les modèles obtiennent de bons résultats, mais la proportion de solutions réussies chute fortement sur les tâches backward complexes pour l'attention. C'est attendu : le passage inverse exige une gestion plus minutieuse des gradients et de la mémoire.

Un autre constat important : l'exécution des instructions ciblées ne garantit pas en soi la vitesse. Un modèle peut « tomber » sur les bonnes instructions PTX, mais choisir une mauvaise stratégie de parallélisation ou un mauvais placement des synchronisations. Au final, aucun des modèles testés n'a réussi à concurrencer de manière stable les bibliothèques frontier sur l'ensemble des tâches.

Comment le modèle a été affiné et ce qui a fonctionné

Pour comprendre si ces compétences peuvent être améliorées, les auteurs ont procédé à un affinage supervisé du modèle Qwen3.6-27B. Ils ont également utilisé une technique que l'on pourrait appeler l'apprentissage avec condition de correction : le modèle apprend non seulement à générer du code, mais aussi à le corriger en recevant un retour sur les étapes incorrectes.

Cette approche a effectivement amélioré plusieurs tâches, mais sans garantir un progrès uniforme. La généralisation est restée faible : le modèle pouvait bien gérer certains types de noyaux et échouer sur d'autres. L'analyse a montré que la taille du jeu de données d'entraînement n'est pas le seul facteur décisif. Ce qui compte davantage, c'est sa couverture, l'équilibre entre les différents scénarios et la qualité du « professeur » qui génère les exemples pour les chaînes de raisonnement. Si les données sont biaisées vers les cas simples, le modèle n'apprendra pas la complexité.

Le principal enseignement

PTXBench offre à la communauté un banc d'essai vérifiable pour suivre les progrès des LLM dans l'optimisation des cœurs GPU. Il ne se contente pas de donner une note « rapide/lent », mais montre à quel niveau les plans du modèle s'effondrent : sur la syntaxe, sur l'application d'instructions spécifiques ou sur le choix de la stratégie. C'est une étape importante pour rendre l'optimisation automatique des noyaux prévisible et pratique.

Foire aux questions

PTXBench : une nouvelle façon de mesurer la capacité des LLM à optimiser les noyaux GPU