L'agent à questions dépasse l'agent à outils : 113% contre 50% de la vitesse de FlashAttention-3

Les kernels GPU rapides pour l'attention sont la base de la vitesse d'entraînement et d'exécution des grands modèles. Mais les écrire à la main prend du temps : chaque génération de cartes graphiques impose de les réécrire, et des équipes d'experts y consacrent des mois. Des chercheurs de Carnegie Mellon University ont testé si cette tâche pouvait être confiée à un agent de codage, et ils aboutissent à une conclusion peu intuitive.

Ce qui aide un agent bloqué, ce ne sont ni de nouveaux outils ni plus de liberté, mais une courte question posée au bon moment. Sur une même tâche, l'agent doté de l'ensemble des outils s'est arrêté à 50% de la vitesse de FlashAttention-3 (le kernel d'attention de référence pour les cartes Hopper) après 133 étapes. L'agent auquel on posait des questions a atteint 113% sur les mêmes 26 configurations.

Pour un développeur, cela signifie que si un agent patine, il lui manque probablement une hypothèse sur l'endroit où chercher le problème, et non une nouvelle dose de données.

Pourquoi les outils ne suffisent pas

Les agents actuels d'optimisation de kernels se situent aux deux extrémités d'un spectre. Les agents à scénario fixent rigidement chaque étape : politique entraînée, menu d'optimisations, rôles figés. Tant que la recette convient, ils sont économes ; quand elle ne convient plus, ils se bloquent.

Les agents libres font eux-mêmes muter un ensemble de kernels, mais paient cette liberté en milliers d'exécutions, en corrections répétées et en candidats erronés. Dans un travail antérieur, un tel agent a exploré plus de 500 pistes en sept jours et produit 40 versions ; il a fallu un superviseur pour le sortir des boucles de modifications ratées. Environ 30% des lignes ajoutées par ce type d'agents sont identiques, octet pour octet, à des lignes supprimées auparavant.

Face à un objectif difficile, les deux écoles répondent de la même façon : donner plus à l'agent. Sortie de profileur analysée, étude du désassembleur, recherche dans la documentation, stockage de « leçons apprises ». Les auteurs estiment que l'agent dispose déjà de suffisamment d'informations. Les modèles s'accrochent avidement à la première solution, s'arrêtent prématurément, tournent en rond parmi des variantes quasi identiques et déclarent avoir terminé alors que la vérification prouve le contraire.

Comment une question change la tâche

L'instruction « rends-le plus rapide » est ouverte, et le modèle y répond de façon gloutonne, en retouchant ce qui se trouve à portée de main. La question « gagnes-tu en cycles du kernel ou seulement en temps réel (wall-clock) ? » est bornée. Pour y répondre, l'agent doit décortiquer et réécrire une partie du kernel qu'il n'aurait pas touchée de lui-même. Une bonne question ne contient pas de réponse toute faite : elle nomme une hypothèse que l'agent n'avait pas envisagée, et celle-ci se vérifie par des moyens ordinaires, une mesure.

La méthode n'a qu'un seul code spécifique : le harnais de vérification. Une seule commande compile le candidat, compare ses sorties à celles du kernel initial et à une référence en fp32, mesure la vitesse sur chaque scénario par rapport au kernel initial, puis calcule l'accélération moyenne géométrique. L'objectif est défini comme une condition d'arrêt vérifiable, et non comme une simple demande au modèle.

Les questions arrivent lorsque la progression est à l'arrêt, lorsque l'agent déclare « c'est la limite » ou « c'est terminé », et après un gain suspect par son ampleur. Au départ, c'est un humain qui les pose ; les questions récurrentes sont ensuite réunies dans une banque que l'agent exploite lui-même. La méthode a aussi été intégrée à une recherche évolutionnaire, où les consignes de mutation sont formulées comme des questions.

Ce que cela a donné

Les campagnes complètes ont produit des kernels de passe avant de l'attention pour Hopper et Blackwell : un seul fichier CUDA, sans CUTLASS, CuTe, cuBLAS ni cuDNN. Ils sont corrects sur les 182 scénarios de chaque architecture (pour Blackwell, sur tous ceux qui peuvent être lancés). En moyenne géométrique, ils sont 1,16 fois plus rapides que FlashAttention-3, soit 16% de plus, et 1,46 fois plus rapides que FlashAttention-4, soit 46% de plus. Sur les formes de requêtes de production, le gain face à FlashAttention-4 est de 1,10 fois. Les auteurs les présentent comme les kernels ouverts les plus rapides de cette classe sur les scénarios testés.

Les statistiques sur les interventions sont instructives : 162 messages humains et 11,5 mille appels d'outils sur 16 sessions. Les questions ne signifiaient pas que l'agent utilisait moins les outils ; il les employait simplement autrement. Chaque gain important a suivi une question qui déplaçait l'agent vers une autre couche du système ou qui récusait la limite qu'il avait déclarée, et aucun n'a suivi un message qui se contentait de fournir une information.

Les réserves sont évidentes. Les auteurs ont eux-mêmes construit l'agent autonome de comparaison, sur le modèle de travaux antérieurs. L'expérience est unique : c'est un cas où existe une référence solide pour mesurer. Une partie du travail d'orientation est assurée par un humain. Les chercheurs ont publié les kernels, le harnais, la banque de questions et toutes les traces.

La conclusion pratique : avant de brancher un outil de plus ou une couche de recherche supplémentaire, définissez un objectif vérifiable sur tous les scénarios et préparez des questions pour les moments de stagnation, de limite déclarée et de victoire trop belle.