🤖 machine learning

AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis

AutoMegaKernel est un harnais d'agent à vérification statique qui synthétise automatiquement un méga-noyau CUDA persistant unique pour les modèles de la famille Llama de HuggingFace, garantissant une exécution sans interblocage sur diverses architectures NVIDIA tout en atteignant jusqu'à 1,72x d'accélération par rapport aux références sur les GPU de classe inférence grâce à des boucles d'agents auto-améliorées.

Auteurs originaux : Jaber Jaber, Osama Jaber

Publié 2026-06-09

📖 6 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Jaber Jaber, Osama Jaber

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez de faire fonctionner une usine massive et complexe (un Grand Modèle de Langage) pour produire un seul produit (un mot de texte).

L'ancienne méthode (IA standard) :
Actuellement, faire fonctionner cette usine revient à envoyer un ouvrier au dépôt pour chercher un outil spécifique, le ramener sur la ligne d'assemblage, effectuer un travail minuscule, puis renvoyer l'ouvrier au dépôt pour l'outil suivant, et répéter cela des dizaines de fois pour chaque mot.

Le problème : L'ouvrier passe la majeure partie de son temps à faire des allers-retours (attendre les données de la mémoire) plutôt qu'à construire. C'est ce qu'on appelle être « limité par la bande passante » (bandwidth-bound). L'usine est limitée par la vitesse à laquelle l'ouvrier marche, et non par sa vitesse de construction.
Le coût : Chaque fois que l'ouvier s'arrête pour prendre un outil, il y a un minuscule délai (latence de lancement). Faire cela à chaque étape s'accumule, ce qui ralentit le processus.

La nouvelle méthode (AutoMegaKernel) :
Les auteurs de ce document ont construit un système appelé AutoMegaKernel (AMK). Considérez l'AMK comme un nouveau gestionnaire d'usine révolutionnaire qui change totalement les règles du jeu.

1. L'usine « en une seule fois » (One-Shot)

Au lieu d'envoyer des ouvriers faire des allers-retours, l'AMK organise toute l'usine de sorte qu'un seul lancement accomplisse tout le travail.

L'analogie : Imaginez qu'au lieu d'un ouvrier qui va chercher les outils un par un, vous avez une équipe d'ouvriers postés à chaque machine. Ils sont tous connectés par un pipeline interne direct. Lorsque l'on appuie sur le bouton « Démarrer », toute la ligne d'assemblage fonctionne du début à la fin en un mouvement continu, sans que personne ne quitte jamais l'atelier pour aller au dépôt.
Le résultat : Cela élimine le « temps de marche » et les délais de « stop-and-start », ce qui, théoriquement, rend le processus beaucoup plus rapide.

2. L'« Inspecteur de Sécurité » (L'agent vérifié statiquement)

Construire une usine où tout le monde travaille simultanément est dangereux. Si les ouvriers ne sont pas parfaitement coordonnés, ils pourraient s'entrechoquer (conflits de type « race conditions ») ou attendre indéfiniment un outil qui n'arrive jamais (blocages de type « deadlocks »).

L'innovation : Habituellement, pour faire fonctionner cela, vous avez besoin d'un expert humain pour écrire le code manuellement, ce qui est difficile et sujet aux erreurs. L'AMK utilise un Agent IA pour concevoir automatiquement la disposition de l'usine.
Le filet de sécurité : Avant même que l'usine ne soit mise en marche, un Inspecteur de Sécurité immuable et figé (un validateur) vérifie le plan de l'IA. Il prouve mathématiquement que :
- Personne ne restera bloqué à attendre indéfiniment.
- Deux ouvriers n'essaieront pas d'utiliser le même outil au même moment.
- L'ordre des opérations est parfait.
L'affirmation : Les auteurs ont testé cet inspecteur contre 7 160 plans « piégés » ou défectueux. L'inspecteur a détecté chaque plan dangereux et les a rejetés avant que l'usine ne démarre. Il n'a jamais laissé passer un plan dangereux.

3. L'« Adaptateur Universel » (Auto-réadressage)

Habituellement, une usine conçue pour un bâtiment spécifique (une puce informatique spécifique) ne fonctionnera pas dans un autre bâtiment.

La magie : L'AMK est capable de « self-retargeting » (auto-réadressage). Vous écrivez le plan une seule fois, et le système l'adapte automatiquement pour qu'il fonctionne parfaitement sur différents types de puces informatiques (comme la RTX 5090, l'A100 ou l'H100) sans avoir besoin qu'un humain réécrive le code pour chacune d'elles.

4. La version « Légère » (Quantification)

Le système a également trouvé comment faire fonctionner l'usine avec des outils plus « légers ».

L'analogie : Au lieu de transporter des outils lourds et de grande taille (calcul haute précision), les ouvriers utilisent des outils légers et compacts (précision int8 ou int4).
Le bénéfice : Comme les outils sont plus légers, ils peuvent en transporter davantage à la fois, et les ouvriers se déplacent plus vite.
- Int8 : Fonctionne aussi bien que les outils lourds (sans perte de précision/lossless) mais est 12 % plus rapide.
- Int4 : Se déplace encore plus vite (réduisant le « temps de marche » de plus de 2 fois), mais le produit final peut être légèrement moins parfait (avec perte/lossy), bien qu'il reste lisible.

5. Les résultats honnêtes (Le « Parler Vrai »)

Les auteurs sont très transparents sur les domaines où leur système gagne et là où il perd.

Là où il gagne : Sur les puces de classe « Inférence » (comme la L4, L40S, A10G et la RTX 5090 grand public), le système AMK est plus rapide que la norme actuelle de l'industrie (cuBLAS) lors de la génération d'un mot à la fois. Il bat la concurrence de 1,1x à 1,3x.
- Pourquoi ? Il a réussi à éliminer les délais d'« allers-retours » qui handicapent les autres systèmes.
Là où il perd : Sur les puces massives et ultra-rapides de classe « Entraînement » (comme l'A100 et l'H100), la version actuelle est plus lente que le système standard.
- Pourquoi ? La « coordination de sécurité » entre les ouvriers (synchronisation) prend un peu de temps. Sur les puces les plus rapides, ce minuscule délai devient le goulot d'étranglement. Les auteurs l'admettent ouvertement : « Nous avons localisé le goulot d'étranglement, et nous sommes honnêtes à ce sujet. »

Résumé

AutoMegaKernel est un système qui utilise un agent d'IA pour concevoir une usine de type « en une seule fois » pour exécuter des modèles d'IA. Il inclut un inspecteur de sécurité rigoureux pour garantir que la conception ne plante jamais. Il réussit à fonctionner sur de nombreuses puces informatiques différentes automatiquement.

Le grand succès : C'est actuellement le moyen le plus rapide de générer du texte mot par mot sur de nombreuses puces populaires (grand public et centres de données), battant les outils standards en éliminant les délais inutiles.
Le bémol : Ce n'est pas encore parfait. Sur les puces de supercalculateurs les plus rapides, la surcharge de coordination rend encore le système légèrement plus lent que les géants établis, mais les auteurs ont prouvé que le système est sûr, correct et capable de s'améliorer de lui-même.

Résumé Technique : AutoMegaKernel (AMK)

Énoncé du Problème

Le décodage autorégressif à flux unique (single-stream) avec une taille de lot (batch size) de 1 est fondamentalement limité par la bande passante. Le plancher de latence est déterminé par le ratio entre le nombre total d'octets de poids et la bande passante de la HBM ( $t_{min} = \text{octets\_poids} / \text{bande\_passante\_HBM}$ ). Les frameworks d'exécution conventionnels (ex. PyTorch, cuBLAS) ne parviennent pas à approcher ce plancher car ils lancent un noyau (kernel) distinct pour chaque opérateur, ce qui induit une latence de lancement CPU significative et force les activations à faire des allers-retours dans la HBM entre les opérations. Bien que les CUDA Graphs amortissent le surcoût de lancement, ils conservent les frontières inter-opérateurs et les allers-retours dans la HBM.

Les approches existantes de type « méga-noyau » (ex. Mirage Persistent Kernel) ou les noyaux optimisés à la main traitent le problème du plancher de latence en fusionnant l'intégralité de la passe avant en un seul noyau persistant. Cependant, ces systèmes manquent de confiance et de portabilité : ils omettent souvent des barrières de sécurité statiques (risquant des blocages ou des conditions de concurrence) et ciblent manuellement des modèles ou architectures spécifiques, sans fournir d'interface structurée permettant à des agents autonomes d'optimiser le système en toute sécurité.

Méthodologie

AutoMegaKernel (AMK) est un système autonome piloté par un agent qui compile des modèles de la famille Llama de HuggingFace en un seul noyau coopératif persistant sans nécessiter d'écriture manuelle de CUDA par modèle. Le système est structuré en quatre couches et deux boucles de recherche automatique :

Couche 0 (La VM) : Un méga-noyau persistant écrit à la main, figé et par architecture (vm/scheduler.cu). Il lance un bloc de threads coopératif par Streaming Multiprocessor (SM). Chaque bloc possède une file d'attente locale et parcourt un ordre topologique global de tâches, exécutant des instructions wait → dispatch → signal. Il repose sur des compteurs monotones pour la synchronisation, évitant ainsi les verrous (locks).
Couche 1 (Instructions) : Des micro-noyaux conformes à l'ABI (GEMV, attention, RMSNorm, etc.) qui sont vérifiés pour leur exactitude de manière isolée.
Couche 2 (Scheduler) : La surface d'édition accessible à l'agent. Il importe un modèle, le réduit en un IR de Schéma typé (un DAG de tâches synchronisé par des compteurs monotones) et propose des configurations structurelles (tiling, fusion, assignation aux SM, profondeur de pipeline) via un objet ScheduleConfig.
Couche 3 (Roadmap) : Un emplacement pour la dynamique future (ex. batching continu).

La Barrière de Sécurité (Safety Gate) :
L'innovation centrale est un validateur à vérification statique (validate() dans schedule/ir.py). Avant tout lancement GPU, le validateur effectue un ensemble de vérifications de graphes statiques sur l'IR du Schéma pour certifier :

L'absence de blocage (Deadlock-freedom) : Garantir que le DAG producteur-consommateur est acyclique, que les seuils d'attente sont satisfaisables ( $1 \le t \le \#producteurs$ ) et que les ordres de files d'attente par SM sont valides.
L'absence de concurrence (Race-freedom) : Garantir que les compteurs partagés sont joints correctement (tous les producteurs doivent finir avant qu'un consommateur n'attende sur un compteur multi-producteurs) et que les accès mémoire respectent les relations de type « se produit avant » (happens-before).
L'exactitude (Correctness) : Vérifier l'accessibilité des sorties et l'intégrité référentielle.

Si un schéma échoue à ces vérifications, il est immédiatement rejeté. Cela empêche des schémas proposés par des agents et potentiellement dangereux de provoquer des blocages de GPU.

Boucle de Recherche Automatique (Autoresearch Loop) :
AMK inclut une boucle non supervisée et pilotable par agent qui propose de manière itérative des changements de ScheduleConfig, les valide, les soumet à des portes de correction par rapport à une référence CPU, mesure la latence, et conserve ou annule les changements. Cette boucle améliore de manière autonome le méga-noyau par rapport à sa propre ligne de base.

Contributions Clés

IR de Schéma à Vérification Statique : Un validateur figé qui certifie l'absence de blocage et de concurrence via des vérifications de graphes statiques (et non des preuves mécanisées) sur une base de confiance. Il a atteint un taux de faux-positifs nul sur 7 160 schémas construits de manière adverse (dont 6 091 étaient dangereux) tout en acceptant les 360 réductions de modèles réels.
Couverture de Génération Automatique : AMK génère automatiquement des méga-noyaux corrects pour 10 modèles sur 10 supportés (incluant des checkpoints réels comme SmolLM2-135M et TinyLlama-1.1B) sans aucune écriture manuelle de CUDA. Il rejette avec succès les variantes non supportées (ex. projections biaisées) ou documente honnêtement ses lacunes.
Auto-rétargeting : Une base de code unique génère des méa-noyaux corrects pour sm_80 (A100), sm_90 (H100) et sm_120 (RTX 5090), correspondant bit à bit ou restant dans des tolérances strictes de fp32 ( $\le 3.8 \times 10^{-5}$ ) par rapport à l'exécution eager.
Noyaux Quantifiés : Le système génère automatiquement des noyaux de poids quantifiés int8 (greedy-lossless, 1.12× plus rapide) et int4 (lossy, 2.42× réduction du trafic de poids) à partir du même chemin d'IR.
Compilateur Pilotable par Agent : Le système traite la correction comme une invariant structurel, permettant à un agent d'explorer l'espace des schémas en toute sécurité sans risquer de blocage du GPU.
Étude de Performance Honnête : L'article rapporte des résultats sur des GPU de classe inférence et de classe entraînement, en divulguant explicitement les domaines où le système gagne ou perd face aux bases de référence établies.

Résultats Expérimentaux

Exactitude et Portabilité :

Sur la RTX 5090, l'A100 et l'H100, AMK reproduit le décodage greedy de HuggingFace jeton par jeton et correspond à la perplexité de l'enseignant (teacher-forced) à $2.5 \times 10^{-7}$ sur le checkpoint SmolLM2-135M.
Le validateur a rejeté les 6 091 schémas dangereux lors du test de stress, démontant un taux de faux-accept de 0 %.

Performance sur les GPU de Classe Inférence (Le Gain) :

RTX 5090 (Grand Public) : Un méga-noyau auto-optimisé en poids uniquement int8 (W8A16) bat le cuBLAS (bf16) avec CUDA-graphs de 1.19–1.23× lors du décodage à batch-1. Il s'agit d'une comparaison asymétrique de précision (int8 vs bf16) où le gain provient du streaming de moins d'octets de poids (0.61× les octets), et non d'un noyau par octet plus rapide.
Flotte d'Inférence de Datacenter : Le méga-noyau int8 bat le cuBLAS (bf16) avec CUDA-graphs sur l'ensemble de la flotte de classe inférence :
- NVIDIA L4 : Jusqu'à 1.33× plus rapide (le gain croît avec la taille du modèle).
- L40S : 1.25–1.27× plus rapide.
- A10G : Jusqu'à 1.08× plus rapide à grande échelle.
Auto-amélioration : La boucle de recherche automatique a amélioré le méga-noyau par rapport à sa version par défaut de 1.25–1.72× sur la RTX 5090.

Performance sur les GPU de Classe Entraînement (La Perte) :

Sur les GPU de classe entraînement à haute bande passante (A100, H100), le méga-noyau auto-généré est en retrait par rapport à cuBLAS avec CUDA-graphs et vLLM.
- Sur H100, AMK était 1.65× plus lent que le chemin CUDA-graph par défaut de vLLM.
- Sur A100, le GEMV bf16 optimisé n'a atteint que ~63 % du pic de bande passante HBM mesuré (contre ~90 % pour cuBLAS).
Diagnostic : Le goulot d'étranglement est le coût de la synchronisation inter-SM par tile (un coût fixe que le méga-noyau paie et que cuBLAS ne paie pas). Sur le silicium à haute bande passante, ce surcoût de synchronisation domine, alors que sur les GPU de classe inférence, les économies de bande passante issues de la fusion et de la quantification amortissent ce coût.

Signification et Revendications
L'article affirme qu'AutoMegaKernel représente un passage de la « vitesse brute » vers la sécurité systématique et la portabilité dans la synthèse de méga-noyaux. Sa principale contribution est une architecture de compilateur à vérification statique qui permet à des agents autonomes d'explorer en toute sécurité l'espace des schémas de modèles complets.

Confiance : Il remplace le risque de blocage de GPU par une barrière de sécurité statique qui rejette les schémas dangereux avant le lancement.
Généralité : Il atteint un auto-retargeting à travers différentes architectures (sm_80/90/120) et familles de modèles sans optimisation CUDA manuelle.
Honnêteté : Les auteurs cadrent explicitement les résultats de performance en fonction des régimes. Le système gagne sur les GPU de classe inférence en fusionnant les opérations et en utilisant la quantification pour réduire le trafic mémoire, mais il perd actuellement sur les GPU de classe entraînement en raison des surcoûts de synchronisation qui nécessitent un ordonnancement plus grossier (travail futur) plutôt que de simples micro-noyaux meilleurs.

L'article conclut que bien que la qualité actuelle des noyaux ne corresponde pas encore aux références optimées à la main sur le matériel haut de gamme d'entraînement, le harnais à vérification statique fournit la fondation nécessaire pour automatiser en toute sécurité les améliorations futures (ex. synchronisation plus grossière, pipeline plus profond) requises pour combler cet écart.