Résumé Technique : HTAM - Mémoire Hiérarchique à Attention sur les Transitions pour l'Optimisation d'Opérateurs

1. Énoncé du Problème

L'augmentation rapide de l'échelle des modèles de langage de grande taille (LLM) et l'évolution des architectures d'accélérateurs spécialisés (par exemple, les GPU) ont créé une demande critique pour le développement d'opérateurs haute performance. Bien que les LLM aient démontré de solides capacités de génération de code, l'optimisation des noyaux GPU reste un problème de recherche conscient du matériel, difficile à résoudre par une génération unique.

Le défi central réside dans le décalage de granularité des méthodes d'optimisation basées sur les LLM existantes :

• Les approches à granularité large (par exemple, indices textuels ou diagnostics de goulots d'étranglement) sont réutilisables et faciles à récupérer, mais souvent trop abstraites pour guider la génération de code CUDA exécutable.
• Les approches à granularité fine (par exemple, traces d'implémentation détaillées) fournissent des détails actionnables mais créent un vaste espace de recherche bruyant, rendant la récupération inefficace et obscurcissant les goulots d'étranglement de l'optimisation.

De plus, l'optimisation des noyaux nécessite des décisions couplées multi-étapes (sélection de tailles de tuiles, mappages de threads, motifs de mémoire, etc.) et la traduction non triviale de ces stratégies en code correct et conscient du matériel. Les méthodes existantes peinent à organiser l'expérience d'optimisation pour guider cette recherche multi-étapes tout en fournissant des preuves concrètes pour les modifications de code.

2. Méthodologie : Cadre HTAM

Pour remédier au décalage de granularité, l'article propose HTAM (Mémoire Hiérarchique à Attention sur les Transitions), un cadre de grossier à fin qui organise l'expérience d'optimisation dans une banque de mémoire structurée. HTAM reflète le processus décisionnel des experts humains : d'abord sélectionner une direction globale de haut niveau, puis l'affiner avec des stratégies locales concrètes.

2.1 Graphe de Transitions Hiérarchique (HTG)

HTAM structure sa mémoire sous la forme d'un Graphe de Transitions Hiérarchique $M = (V_G, V_{GL}, E_G)$ :

• Nœuds Globaux ($V_G$) : Représentent des directions d'optimisation larges (par exemple, "Accès Mémoire", "Réutilisation des Données", "Débit"). Ces nœuds stockent l'intention de haut niveau, des a priori statiques (objectifs, déclencheurs, risques) et des statistiques d'exécution (taux de réussite, tendances d'amélioration).
• Nœuds Locaux ($V_{GL}$) : Représentent des stratégies spécifiques et exécutables sous un nœud global parent (par exemple, "chargements vectorisés alignés" sous "Accès Mémoire"). Ces nœuds stockent des connaissances au niveau de l'implémentation, y compris des motifs de modification, des vérifications de validation et des preuves positives/négatives issues d'exécutions passées.
• Arêtes de Transition ($E_G$) : Relient les nœuds globaux, stockant des connaissances dépendantes du chemin sur la transition d'une direction globale à une autre. Ces arêtes enregistrent quand une transition est applicable, sa justification, ses risques et des preuves au niveau de la trajectoire (gains de performance, résultats de compilation).

2.2 Processus d'Optimisation

L'optimisation se déroule comme un processus d'évolution en $T$ étapes où l'implémentation $c_t$ est mise à jour de manière itérative :

1. Résumé de l'État : L'implémentation actuelle et l'historique sont résumés en un état compact $s_t$ contenant le contexte de la tâche, les retours (compilation, correction, temps d'exécution) et les chemins d'optimisation récents.
2. Sélection de la Direction Globale : HTAM sélectionne une direction globale $g_t$ basée sur l'état actuel et l'historique récent. Elle utilise un mécanisme de scoring conscient des transitions qui agrège les preuves du préfixe récent des directions globales ($p^g_t$) en utilisant des poids sensibles à la position (inspirés des mécanismes d'attention). Cela évite le coût élevé du stockage de mémoires de chaînes complètes tout en capturant les dépendances multi-étapes.
   • Le score est calculé comme une somme pondérée des utilités de transition des nœuds précédents vers le nœud candidat.
   • Une politique basée sur Softmax équilibre l'exploitation des transitions à haute récompense avec l'exploration.
3. Sélection de la Stratégie Locale : Une fois une direction globale choisie, l'espace de recherche est restreint aux nœuds locaux associés. Un LLM évalue les stratégies locales candidates en fonction de l'état, de l'implémentation actuelle et des mémoires spécifiques des nœuds globaux/locaux pour sélectionner la meilleure stratégie $l_t$.
4. Implémentation et Mise à Jour : Le LLM génère la prochaine implémentation $c_{t+1}$ guidée par la stratégie locale sélectionnée. Les retours d'évaluation sont ensuite utilisés pour mettre à jour le HTG : les statistiques des nœuds globaux, les preuves des nœuds locaux et les retours des arêtes de transition sont actualisés pour informer les étapes futures.

3. Contributions Clés

• Cadre de Grossier à Fin : HTAM décompose l'optimisation d'opérateurs basée sur les LLM en sélection de direction globale et raffinement de stratégie locale, remédiant au décalage de granularité des méthodes existantes.
• Organisation de Mémoire Structurée : Il introduit un Graphe de Transitions Hiérarchique qui organise l'expérience en nœuds globaux, nœuds locaux et arêtes de transition. Cette structure réduit l'espace de recherche et la complexité de stockage par rapport à une mémoire de chaîne complète ou une récupération plate.
• Prise de Décision Consciente des Transitions : Le cadre emploie un mécanisme de scoring qui agrège l'historique récent d'optimisation via des mémoires d'arêtes, guidant les décisions multi-étapes sans le coût exponentiel du stockage de trajectoires complètes.
• Validation Empirique : De nombreuses expériences démontrent que HTAM surpasse significativement les références en termes de correction, de vitesse de découverte de solutions rapides et de gain de vitesse moyen géométrique.

4. Résultats Expérimentaux

Les expériences ont été menées sur KernelBench, une suite de 250 tâches (Niveau-1 opérateurs simples, Niveau-2 opérateurs fusionnés, Niveau-3 tâches au niveau du modèle), utilisant DeepSeek-R1 comme modèle de base sur des GPU NVIDIA A100.

• Gains de Performance : HTAM a atteint 98,4 % de correction, 84,0 % de Fast@1 (pourcentage de tâches avec une solution correcte et plus rapide dès la première tentative), et un gain de vitesse moyen géométrique de 1,978×.
• Comparaison avec les Références :
  • Par rapport à la référence la plus forte reproduite (Récupération de Mémoire Plate), HTAM a amélioré la correction de 8,8 points de pourcentage, Fast@1 de 25,2 points de pourcentage, et le gain de vitesse de 1,464× à 1,978×.
  • Il a significativement surpassé la génération brute de LLM et les méthodes de raffinement basées uniquement sur les retours.
• Généralisation :
  • Transfert de Backend : HTAM est resté efficace sur différents backends de LLM (DeepSeek-R1, GPT-4o, DeepSeek-V4-Flash), bien que les performances aient été corrélées à la capacité de génération de code du backend.
  • Transfert Inter-Benchmark : Lorsqu'initialisé avec une mémoire construite à partir de KernelBench et appliqué aux tâches Robust-KBench, HTAM a montré des gains de vitesse améliorés (par exemple, +1,76× pour LayerNorm), indiquant une expérience d'optimisation transférable, bien que les gains dépendent de l'opérateur.
• Études d'Ablation : La suppression de la structure hiérarchique ou de la mémoire locale a provoqué des baisses significatives de performance (par exemple, le gain de vitesse chutant à 0,974× sans mémoire locale), confirmant que l'expérience structurée global-local est critique. De même, la suppression du scoring de transition sensible au préfixe a réduit les performances, validant la nécessité d'une prise de décision dépendante de l'historique.

5. Importance et Revendications

L'article revendique que la mémoire structurée sert non seulement de référentiel de cas passés, mais aussi de priori décisionnel pour sélectionner la prochaine direction d'optimisation utile. En organisant l'expérience de manière hiérarchique et en modélisant les transitions entre les étapes d'optimisation, HTAM permet aux LLM de naviguer plus efficacement dans l'espace de recherche complexe et multi-étapes de l'optimisation des noyaux GPU que la récupération plate ou la génération unique.

Les auteurs soulignent que HTAM fournit une voie pratique pour automatiser le développement d'opérateurs haute performance, atteignant des niveaux de correction et de gain de vitesse proches de ceux des experts tout en maintenant un coût d'API en ligne gérable (environ 0,21 $ par tâche pour une optimisation en 6 étapes utilisant DeepSeek-R1).

Limites Reconues :

• L'évaluation est limitée par les coûts de calcul et d'API, se concentrant principalement sur KernelBench et un sous-ensemble de Robust-KBench.
• Le système repose sur une validation exécutable (compilation, correction, temps d'exécution) plutôt que sur une vérification formelle, ce qui signifie que la correction n'est pas garantie pour toutes les formes ou distributions de données possibles.
• La généralisation à des matériels ou familles d'opérateurs substantiellement différents peut nécessiter l'adaptation du schéma de mémoire.