Atteindre une efficacité extrême grâce à la génération de kernels GPU spécialisés

Traditionnellement, les systèmes d'inférence en production s'appuient sur des kernels génériques pour gérer divers modèles et charges de travail. C'est sous-optimal car la forme des opérations GPU est déterminée par une combinaison de paramètres de modèles statiques et de facteurs dynamiques au moment de la requête ; par exemple, alors qu'un modèle définit l'une des dimensions pour une multiplication de matrices, l'autre dimension fluctue en fonction du nombre de tokens spécifique de chaque requête. Il y a un intérêt croissant pour la génération de kernels GPU par des agents, et les efforts récents se révèlent prometteurs. Nous avons exploré une question fondamentale : si la génération de kernels peut être automatisée, pourquoi des modèles de tailles extrêmement différentes (de 1 milliard à 1 000 milliards de paramètres) devraient-ils s'appuyer sur le même kernel ? En spécialisant les kernels pour les formes spécifiques rencontrées lors de l'exécution, nous pouvons atteindre une efficacité extrême.

Dans ce blog, nous partageons nos succès et nos enseignements tirés de l'utilisation d'agents pour générer des kernels GPU. Nous avons conçu Proteus, un système développé pour atteindre une spécialisation extrême, ce qui nécessite un harnais adapté pour une optimisation, une validation et une gestion du contexte rigoureuses.

Les harnais de codage conventionnels échouent souvent ici car les agents ont tendance à contourner le système de récompense (reward-hack) : ils suivent la lettre de la règle plutôt que l'esprit. Si vous donnez un benchmark à un agent, il risque d'optimiser le benchmark et non l'opération prévue.

Pour y remédier, Proteus propose des kernels, les vérifie par rapport à une implémentation de référence contrôlée, mesure le temps d'exécution de ceux qui réussissent et améliore de manière itérative les meilleurs résultats. Bien que le processus soit simple, son succès dépend entièrement de la résolution de deux défis fondamentaux. La figure 1 montre l'architecture simplifiée de notre conception. En utilisant notre harnais Proteus, nous avons généré des kernels Qwen 3.5 122B qui étaient 1,8 à 5,2 fois plus rapides que les meilleurs disponibles dans vLLM.

Validation

À l'origine, nous considérions la recherche de kernels comme la partie difficile : comment explorer un vaste espace de programmes sans rester bloqué sur un plateau sans amélioration ? En pratique, la première question était plus fondamentale. Mesurons-nous vraiment ce que nous pensons mesurer ?

Un modèle optimise le score que vous lui donnez. Il n'a pas besoin d'une faille complexe : il se peut simplement que l'évaluation repose sur une hypothèse erronée. Un exemple concernait les kernels pour les rotary position embeddings (RoPE), une étape courante dans les couches d'attention. Un candidat pouvait réutiliser du code compilé restant d'une tentative précédente et sembler moins coûteux qu'une reconstruction équitable à partir de zéro. Un autre pouvait enregistrer un lot de lancements GPU dans un graphe (par exemple, un graphe CUDA) et les rejouer comme une seule unité, tandis que la référence à laquelle nous le comparions lançait toujours chaque élément séparément, de sorte que les deux côtés ne faisaient pas le même travail. Un autre était performant sur les tailles d'entrée que nous avions placées dans le jeu de test visible et faible sur les tailles qui ne lui avaient pas été présentées.

Nous avons donc consacré nos premiers efforts de conception au vérificateur, et non au prompt. Mesurez le temps des deux côtés de la même manière, y compris avec plus d'un minuteur (par exemple, le minuteur d'événements CUDA, le temps réel et le minuteur CUPTI) lorsque nous avons besoin d'une double vérification. Effacez l'état compilé restant qui ne devrait pas persister, et maintenez la cohérence de l'ordre de configuration et de démontage afin qu'un côté ne puisse pas sauter un travail que l'autre doit encore effectuer. Mesurez à nouveau le temps des gagnants avant de les utiliser comme point de départ pour le tour suivant. Conservez certains tests invisibles pour le candidat, afin qu'il ne puisse pas s'adapter uniquement à l'examen. Pour éviter la « triche » lors de l'évaluation avec des performances artificiellement gonflées, nous implémentons des vérifications de cohérence automatisées pour signaler les accélérations théoriquement impossibles (par exemple, >100x) qui dépassent la bande passante physique du GPU et les limites de calcul. Cela protège contre les mêmes pièges de contournement de récompense observés dans des cas passés de l'industrie, où les agents optimisaient pour les métriques du harnais plutôt que pour de réels gains de performance. Sans ces contraintes, générer plus de kernels produisait principalement plus de bruit.

L'accent mis sur le vérificateur modifie également le goulot d'étranglement de la génération de kernels par des agents. Dans le simple travail de recherche de programmes (c'est-à-dire l'optimisation itérative où le système recherche des programmes en générant de manière répétée des variantes), les bons candidats sont rares, de sorte que leur écriture domine le coût. Nous pouvons produire de nombreux brouillons en parallèle, mais nous ne pouvons pas sauter la validation. Nous devons concevoir la validation avec soin, et la vérification doit s'exécuter sur de vrais GPU, de manière isolée, et plus d'une fois. Le système progresse aussi vite qu'il peut faire confiance à un kernel, et non aussi vite qu'il peut en écrire un.

Gestion du contexte

Un autre défi consiste à déterminer ce que le modèle de génération de kernels est autorisé à voir. C'est un compromis. Donnez au modèle un prompt plus grand et il disposera de plus d'informations : le meilleur kernel actuel, les échecs récents, les indices du profileur, les notes des exécutions précédentes. Cela peut aider. Cela coûte également plus cher, car nous payons pour chaque token lu par le modèle. Et à mesure que le prompt grandit, il est plus facile pour la tentative suivante de dériver. Les signaux utiles sont mélangés à des conseils obsolètes, des astuces contradictoires et des détails qui s'appliquent à une taille d'entrée différente ou à une opération différente. Le modèle ne sait pas toujours à quelles phrases faire confiance, il suit donc les plus évidentes, ou toutes un peu.

Donnez-lui trop peu et c'est l'inverse qui se produit. Chaque tentative repart de zéro. Les mêmes impasses réapparaissent. Rien ne se transmet de la dernière exécution, ou d'une opération connexe, et la boucle n'avance pas.

Nous voulions une couche de connaissances pour nous y aider : mémoriser ce qui a fonctionné, le réutiliser plus tard, et ce, sans intervention humaine. Cette couche présente un second compromis, entre le niveau de détail d'une leçon stockée et sa portée d'application.

Une note très spécifique (« sur ce kernel, avec cette taille d'entrée, déroulez cette boucle ») peut être exactement ce dont la tentative suivante a besoin. Elle est également facile à mal utiliser lors de l'opération suivante, sur le GPU suivant ou avec une taille d'entrée différente. Une note très générale (« faites un meilleur usage de la mémoire sur puce ») s'applique presque partout et n'indique presque rien au modèle sur ce qu'il doit faire. Nous avons constaté ces deux modes d'échec. Lorsque les leçons étaient trop générales, elles reformulaient un échec sans proposer d'action. Lorsque nous stockions plus de détails, ils étaient souvent trop liés à une seule exécution pour aider la suivante. Sur une longue exécution, la majeure partie de ce que le modèle lisait et écrivait était consacrée à la récupération et à l'acheminement de cette mémoire plutôt qu'à l'écriture de kernels. La couche de mémoire faisait beaucoup de travail. Elle ne rendait pas le candidat suivant meilleur. La figure 2 montre la répartition du coût en tokens d'un tel système. Le coût est dominé par les couches de connaissances.

La version des connaissances qui vaut la peine d'être conservée est plus restreinte et trouve un équilibre entre généralité et spécificité. Lorsque le modèle est sur le point d'écrire un kernel, son prompt ne doit inclure qu'un contexte de haute confiance : des enseignements exploitables qui associent des situations spécifiques à des actions (synthétisées à partir de la cartographie passée des modifications par rapport aux impacts) et des notes d'échec concises provenant d'exécutions parentes étroitement liées. Récupérées via un filtrage hiérarchique par balises combiné à une recherche hybride (mots-clés + sémantique), les leçons doivent être suffisamment spécifiques pour être exploitées et suffisamment délimitées pour clarifier les cas où elles ne s'appliquent pas. Les opérations plus profondes, comme la réorganisation et la distillation continue de la base de leçons, relèvent de tâches en arrière-plan, et non d'un parcours synchrone multi-sauts des exécutions passées à chaque tentative. Si un enseignement ne peut pas nommer la situation et l'action, cela ne vaut pas la peine de l'inclure dans le prompt. La figure 3 montre la répartition du coût en tokens après la correction de la couche de connaissances ; la plupart des tokens sont consacrés à la génération de candidats après cette correction.

Étude de cas : décodage compacté Gated DeltaNet

Un exemple concret est le kernel de décodage compacté sur le chemin Gated DeltaNet dans Qwen 3.5 122B. L'opération met à jour un état récurrent et écrit la sortie de décodage à partir d'entrées QKV compactées, de paramètres de porte et d'indices d'état. Nous avons utilisé cette tâche pour tester l'ensemble de la boucle Proteus sur des GPU NVIDIA B200 avec un backend Triton : valider le contrat de tâche, mesurer l'implémentation de référence, demander des candidats de kernels aux agents, exécuter des vérifications statiques et des compilations, vérifier l'exactitude par rapport à la référence contrôlée, benchmarquer uniquement les candidats vérifiés, puis mesurer à nouveau les meilleurs candidats.

La figure 4 se lit de gauche à droite. Le nœud de référence ancre le benchmark à 0,025 ms. Le candidat 0000 est la graine sécurisée : il a reproduit la structure de décodage compacté et a réussi la validation, mais il était plus lent que la référence, donc Proteus l'a conservé comme parent mesuré plutôt que de le traiter comme une victoire. À partir de là, Proteus a cessé d'optimiser un kernel générique pour chaque forme et a divisé la recherche en chemins spécifiques à chaque forme.

Le chemin de réparation Batch-1 a produit un kernel spécifique à la forme pour le candidat 012, atteignant 1,5x sur la forme de décodage à lot unique (single-batch). Les résultats les plus solides ont été obtenus sur le chemin de décodage de service (serving-decode) : le candidat 030 a trouvé la latence de kernel mesurée la plus faible à 0,018 ms, et le candidat 036 a produit la meilleure accélération de forme à 1,6x. Ce fragment de service gagnant s'est spécialisé pour la disposition Batch=4, Key=128, Value=128 et a traité la dimension de valeur par blocs de 64 de large, il s'agit donc d'un kernel sûr pour cette forme spécifique plutôt que d'un remplacement universel.

Le dernier détour de la chronologie montre pourquoi la trace est importante. Les tentatives de générations ultérieures en C++ (plutôt qu'en Triton) ont rencontré des échecs de build et de génération, et la longue exécution s'est terminée après avoir épuisé le budget de tentatives de cette branche. L'artefact utile n'est donc pas seulement le candidat le plus rapide. Il s'agit du chemin complet illustré par la figure : les échecs sémantiques ont été rejetés, les kernels corrects mais plus lents ont été mesurés, et les gains de performance réels sont restés associés à la forme qui permet de les composer en toute sécurité dans un kernel de production.

Ce sur quoi nous travaillons ensuite

La boucle évolutive que nous avons construite est stricte dans son rôle de rédacteur. Elle appelle souvent le modèle selon un schéma fixe : prendre le meilleur kernel actuel, essayer une petite modification, vérifier, répéter. Cela prive l'agent de l'autonomie dont il a besoin. Il ne peut pas facilement changer de structure, changer de langage ou abandonner une conception sans issue.

La boucle reste nécessaire. Non pas pour rédiger le kernel, mais pour donner à l'agent un indice fiable pour l'étape suivante. Cet indice doit provenir de deux sources.

Premièrement, la communication avec la couche de connaissances : quelques enseignements suffisamment précis pour être exploités, et délimités de manière à savoir où ils ne s'appliquent pas. Sans cela, chaque tentative repart de zéro.

Deuxièmement, les résultats d'un vérificateur de confiance : la correction et le chronométrage que l'agent n'a pas mesurés lui-même. Ces chiffres sont des indices pour la tentative suivante. Ce sont également les seuls scores auxquels nous devrions nous fier. Si l'agent mesure lui-même son propre temps d'exécution, nous retombons dans les problèmes de caches résiduels, de comparaisons non correspondantes et de tests visibles par lui.

La distinction que nous souhaitons faire est donc plus subtile que « agent contre boucle ». Donnez à l'agent l'autonomie sur la façon dont un kernel est écrit. Conservez la boucle comme canal pour la mémoire et l'évaluation. L'agent propose. La boucle renvoie ce qu'elle est autorisée à voir, et si la dernière proposition a effectivement été retenue.

Proteus a construit des kernels spécialisés pour des parties de Qwen 3.5 122B sur son chemin Gated DeltaNet (un bloc de style attention linéaire) s'exécutant sur des GPU NVIDIA B200. Les accélérations sur les kernels individuels se situaient dans une fourchette de 1,8x à 5,2x.

La leçon à retenir est que la génération est l'étape la moins coûteuse. La validation et la gestion du contexte constituent la partie difficile. C'est là qu'un temps de conception minutieux et des innovations sont nécessaires.

La génération agentique de kernels GPU a libéré le potentiel incroyable de la spécialisation extrême, mais la construction de harnais de test fiables et prêts pour la production reste un défi de taille. Nous relevons les défis les plus complexes à l'intersection de l'IA et des systèmes, et nous recherchons des ingénieurs audacieux pour nous aider à façonner l'avenir de l'inférence efficace. Si vous êtes passionné par l'idée de repousser les limites du possible, nous recrutons !

(Cet article de blog a été traduit à l'aide d'outils basés sur l'intelligence artificielle) Article original