Un agent guidé par des questions amène un noyau CUDA à 1,13× la vitesse de FlashAttention-3, contre 0,50

Des chercheurs de la Carnegie Mellon University proposent une recette inhabituelle pour les agents qui optimisent des noyaux GPU. Au lieu de fournir au modèle des rapports de profileur déjà dépouillés, une recherche dans la documentation et un stock de « leçons apprises », ils lui posent de courtes questions aux moments clés. Sur la même tâche, avec le même évaluateur, sur le même nœud H200 et la même semaine, cet agent atteint 1,13 fois la vitesse de FlashAttention-3 en moyenne géométrique sur 26 configurations. L'agent autonome, construit selon le schéma dominant aujourd'hui, s'arrête à 0,50 après 133 étapes.

Entre le script et la recherche libre

Les auteurs décrivent deux camps. Les agents scriptés fixent à l'avance ce que le modèle fait à chaque étape : politique entraînée, menu d'optimisations, rôles rigides. Ils dépensent leurs tentatives avec parcimonie là où la recette convient, et patinent là où elle ne convient pas. Chez Autocomp, par exemple, les performances s'effondrent sans le menu d'optimisations, et une approche RL à boucle fixe, compétitive sur des opérateurs isolés, échoue sur les modèles fusionnés (fused).

Les agents ouverts comme AVO laissent au modèle la liberté de faire muter une population de noyaux, et le paient en nombre de tentatives. Selon l'article, l'agent AVO a exploré plus de 500 directions en sept jours pour retenir 40 versions, et il a eu besoin d'un superviseur pour sortir des boucles de modifications ratées. Environ 30 % des lignes ajoutées par de tels agents sont identiques, octet pour octet, à des lignes supprimées auparavant, et une partie des candidats CUDA trouvés par évolution se révèlent incorrects.

Les deux camps réagissent de la même façon face aux objectifs complexes : ils donnent plus d'informations à l'agent. Les auteurs estiment que ce qui manque est autre chose : une direction, c'est-à-dire une hypothèse sur l'endroit où il reste un écart. Les LLM gèrent mal l'optimisation ouverte : ils cherchent de façon gloutonne, s'accrochent à la solution présente dans le prompt, s'arrêtent trop tôt et déclarent avoir terminé contre les données de l'environnement.

Comment fonctionne la méthode

Le nom de l'agent lui-même ne figure pas dans le texte disponible, mais sa conception est décrite clairement. On part d'un noyau CUDA existant, d'un ensemble de scénarios et d'un objectif. Trois éléments s'y ajoutent : un harnais d'évaluation minimal, un agent de code généraliste sans scénario prédéfini et une source de questions.

L'évaluateur compile le candidat en une seule commande, compare la sortie à la référence et à la référence fp32 sur chaque scénario, puis mesure le temps par rapport au noyau initial sur le même GPU. Les contraintes sont vérifiées de façon mécanique : l'interdiction de CUTLASS, CuTe, cuBLAS et cuDNN est contrôlée d'après les symboles du binaire compilé. L'objectif prend la forme d'un hook Stop : quand l'agent tente de terminer son tour, un prédicat est évalué, et si les conditions ne sont pas remplies, l'agent est renvoyé au travail.

Le prédicat doit couvrir chaque scénario, et non la moyenne, faute de quoi l'agent ajusterait la moyenne sur les formes les plus favorables. La seule fois où le verdict a été confié à un modèle juge, celui-ci a accepté une déclaration de fin prématurée, qu'un humain a ensuite annulée.

Les questions sont posées à cinq moments types : au début de la campagne, sur un plateau, quand l'agent annonce une limite ou la fin du travail, après un gain suspect par son ampleur, et quand il explique une erreur de correction pour l'écarter. Exemple tiré de l'article : « Le noyau de référence gagne-t-il en cycles d'horloge ou seulement en temps réel (wall-clock) ? » Une telle question est bornée et vérifiable, et pour y répondre, le modèle peut réécrire une partie du noyau qu'il n'aurait jamais touchée de lui-même.

Parmi les 162 messages humains, les plus fréquents portaient sur les vérifications de mesures (13), la comparaison avec une estimation de la borne (10), la validité du benchmark (7) et l'analyse des écarts (7). Entre les questions, l'agent travaille seul : environ 71 appels d'outils par message. Selon les auteurs, tous les gains importants ont suivi une question qui faisait passer l'agent à un autre niveau du système ou contestait une limite annoncée, et non une question qui apportait une information nouvelle.

Résultats et réserves

Les noyaux finaux de la passe avant de FlashAttention pour Hopper et Blackwell sont des fichiers CUDA uniques, sans CUTLASS, CuTe, cuBLAS ni cuDNN. Ils sont corrects sur les 182 scénarios pour Hopper (pour Blackwell, sur tous les scénarios exécutables) et atteignent en moyenne géométrique 1,16 fois FlashAttention-3 et 1,46 fois FlashAttention-4, dont 1,10 fois FA4 sur les formes de production. L'ensemble a demandé 16 sessions et environ 11,5 mille appels d'outils.

Les auteurs ont regroupé les questions récurrentes dans une banque : 101 entrées principales réparties en dix groupes, cinq pour le cas où une autre implémentation de comparaison existe, et 18 liées à des architectures. Mais seules 29 entrées sont confirmées par des traces, tandis que 72, tirées de sources publiques, n'ont pas encore été vérifiées. La première version de la banque, de 48 questions, était biaisée en faveur d'un seul cas : 21 questions renvoyaient au noyau de référence et 13 à la structure de l'attention. Ces éléments ont été retirés de la version publiée.

Les réserves sont évidentes. Il s'agit d'un seul cas, l'attention, pour lequel il existe une référence solide. Dans les expériences, c'est un humain qui posait les questions, et l'usage autonome de la banque est moins bien vérifié. La comparaison avec l'agent « outillé » a été construite par les auteurs eux-mêmes, même si c'est dans le style d'AVO. En revanche, le code des noyaux, l'évaluateur, la banque de questions et toutes les traces doivent être publiés, de sorte que le résultat pourra être revérifié.