AWS Trainium : des IA qui écrivent elles-mêmes leurs optimisations ?

Imaginez un développeur qui n’a plus besoin de passer des semaines à écrire et optimiser manuellement des noyaux (ces petits programmes ultra-rapides qui font tourner les calculs sur les puces). Avec les nouvelles capacités de Neuron Agentic Development, c’est désormais possible : des agents d’intelligence artificielle s’en chargent à sa place, en quelques minutes seulement.

LES LIMITES DU Développement MANUEL DE NOYAUX

Les modèles d’IA les plus avancés grandissent en taille et en complexité chaque jour. Résultat : les développeurs doivent extraire le maximum de performances de leurs puces, comme les AWS Trainium ou AWS Inferentia. Pourtant, la différence entre les performances théoriques d’une puce et ce que la plupart des équipes obtiennent reste énorme. Traditionnellement, la solution passait par le développement manuel de noyaux, une tâche qui demande une expertise pointue en architecture matérielle, des cycles d’optimisation interminables et des tests répétés. Peu d’équipes peuvent se le permettre.

DES AGENTS IA POUR AUTOMATISER LA CRÉATION DE NOYAUX

Et si chaque ingénieur en machine learning pouvait devenir un expert en optimisation matérielle ? Et si un développeur habitué à une architecture pouvait maîtriser une autre en quelques jours, au lieu de mois ? C’est exactement ce que propose Neuron Agentic Development : une collection d’agents et de compétences d’IA qui automatisent la création, le débogage et le profilage des noyaux pour les puces AWS Trainium et Inferentia.

CINQ COMPÉTENCES POUR UNE OPTIMISATION COMPLÈTE

Le package Neuron Agentic Development propose cinq compétences spécialisées qui suivent le processus naturel de développement d’un noyau : écrire, déboguer, profiler, analyser. Chaque compétence peut être utilisée individuellement pour des tâches ciblées, ou combinée avec l’agent neuron-nki-agent, qui sélectionne automatiquement le bon flux de travail en fonction de la demande.

Pour les utiliser, il suffit d’ajouter les compétences à votre répertoire de compétences dans votre IDE agentique (comme VS Code, Cursor ou Kiro). Par exemple, dans Kiro, placez-les dans le dossier .kiro/skills. Ces compétences doivent s’exécuter sur une instance Amazon EC2 basée sur Trainium.

ÉCRIRE UN NOYAU : TRADUIRE PYTORCH EN CODE OPTIMISÉ

La compétence neuron-nki-writing est votre point de départ pour créer des noyaux NKI (Neuron Kernel Interface). Elle transforme du code PyTorch, NumPy ou des descriptions en langage naturel en code NKI correct et optimisé. Par exemple, elle gère les stratégies de tuilage qui respectent les contraintes matérielles, comme les dimensions de partition (128) ou les dimensions libres (512/4096). Elle optimise aussi les accès mémoire, les opérations de calcul avec des paramètres explicites, et les directives d’efficacité pour la taille du DMA et la réutilisation du SBUF.

Cette compétence classe votre tâche par complexité et ne charge que les références nécessaires. Résultat : un noyau prêt à l’emploi, sans erreur de syntaxe ni problème de compatibilité.

DÉBOGUER UN NOYAU : RÉSOUDRE LES ERREURS EN QUELQUES CLICS

La compétence neuron-nki-debugging offre un flux de travail systématique pour résoudre les erreurs de compilation et d’exécution des noyaux NKI sur les puces Trainium et Inferentia. Elle couvre la configuration de l’environnement avec les bons drapeaux --target, la résolution des erreurs de compilation avec un index catégorisé de tous les 28 codes d’erreur NCC, et la validation numérique par rapport à des références calculées sur CPU.

Plus besoin de passer des heures à chercher pourquoi un noyau ne compile pas : l’agent vous guide étape par étape, avec des explications claires et des solutions concrètes.

PROFILER UN NOYAU : MESURER LES PERFORMANCES EN TEMPS RÉEL

La compétence neuron-nki-profiling capture les profils d’exécution directement sur le matériel. Elle configure les variables d’inspection du runtime, exécute le noyau, identifie le bon format de fichier d’exécution Neuron (NEFF), et capture la trace avec neuron-explorer, y compris les notifications DGE (DMA Graph Engine) pour des détails au niveau du DMA. Elle extrait ensuite des métriques au format JSON et produit les fichiers NEFF consommés par la compétence suivante.

En résumé : vous obtenez une cartographie précise de l’utilisation de votre noyau, avec des données exploitables pour l’optimisation.

ANALYSER UN PROFIL : TROUVER LES GASPILLAGES DE PERFORMANCE

La compétence neuron-nki-profile-querying ingère les fichiers NEFF et NTFF, puis exécute des requêtes SQL pour calculer les limites de performance, identifier les moteurs bottlenecks et localiser les inefficacités au niveau des lignes de code source NKI. Trois approches d’analyse sont possibles : l’API serveur de neuron-explorer, DuckDB directement sur Parquet, ou pandas pour des calculs personnalisés.

L’agent vous permet de creuser dans les données pour comprendre exactement où se situent les problèmes, sans avoir à analyser manuellement des logs complexes.

DOCUMENTATION INTÉGRÉE : UNE AIDE À CHAQUE ÉTAPE

La compétence neuron-nki-docs est disponible tout au long du développement. Pendant l’écriture, elle fournit les signatures d’API et des tutoriels. Pendant le débogage, elle explique les codes d’erreur. Pendant le profilage, elle clarifie les détails de l’architecture matérielle. Vous pouvez demander des informations sur une API spécifique comme nisa.* ou nl.*, consulter les codes d’erreur, trouver des tutoriels ou explorer les guides d’architecture pour Trainium 1, 2 et 3.

DES AGENTS POUR DES FLUX DE TRAVAIL AUTONOMES

Si les compétences offrent des blocs de construction pour des tâches individuelles, les agents combinent plusieurs compétences en flux de travail autonomes. Chaque agent est une persona spécialisée qui gère des scénarios de développement multi-étapes de bout en bout. Par exemple, l’agent neuron-nki-agent peut écrire un noyau, le déboguer, le profiler et analyser ses performances, le tout sans intervention humaine.

TUTORIEL : CRÉER UN NOYAU SOFTMAX EN 4 ÉTAPES

Pour illustrer le fonctionnement de ces capacités agentiques, voici un tutoriel pas à pas. Nous allons créer deux noyaux : un softmax et un SwiGLU, puis analyser leurs performances sur une charge de travail réelle.

Première étape : installer les compétences. Exécutez la commande suivante pour installer l’outil Kiro et ses compétences Neuron Agentic Development :

curl -fsSL https://cli.kiro.dev/install | bash

Note : les instances trn2.3xlarge génèrent des frais horaires pendant leur utilisation. Pensez à les arrêter une fois le tutoriel terminé pour éviter des coûts inutiles.

Pour configurer l’instance en détail, consultez le guide de configuration de l’AWS Neuron DLAMI.

VÉRIFIER LA CONFIGURATION DE L’INSTANCE

Depuis le terminal distant, vérifiez que les périphériques Neuron sont visibles :

# Confirmer que les périphériques Neuron sont visibles
neuron-ls

# Confirmer que neuron-explorer est disponible
which neuron-explorer && neuron-explorer --version

L’image AMI AWS Neuron est pré-équipée d’un environnement virtuel à l’emplacement suivant :

~opt/awsneuronxvenvpytorch2_9

Activez-le avec :

source ~opt/awsneuronxvenvpytorch2_9/bin/activate

Avec l’environnement configuré, vous pouvez commencer à développer des noyaux en exécutant :

kiro-cli --agent neuron-nki-agent

ÉTAPE 1 : ÉCRIRE UN NOYAU SOFTMAX

Dans la session interactive Kiro CLI, entrez la commande suivante :

“Write an NKI kernel that computes scaled softmax: softmax(x * scale) along the last dimension, for input shape [batch, seqlen, hiddendim] in bfloat16.”

L’agent produit un noyau complet en trois passes (max par ligne, somme des exponentielles, normalisation) utilisant nisa.activation(np.exp, .) pour l’exponentielle accélérée par le matériel, une accumulation en float32 pour la stabilité numérique, et un tuilage adapté à la dimension libre. Il explique ses choix de conception : une instance de programme par ligne, PMAX=128 (correspondant à la limite matérielle de 128 partitions), FMAX=2048 (correspondant à la limite de 2048 éléments sur la dimension libre de Trainium), et une sortie en bfloat16.

Demandez à l’agent d’exécuter le noyau et de vérifier la parité numérique par rapport à une référence PyTorch.

VALIDER LE NOYAU SUR MATÉRIEL RÉEL

Après synchronisation du noyau corrigé sur l’instance et exécution sur le périphérique :

PASS: shape=(2, 128, 512), max_diff=0.000008
PASS: shape=(4, 256, 1024), max_diff=0.000004
PASS: shape=(1, 1, 64), max_diff=0.000061
PASS: shape=(2, 300, 768), max_diff=0.000007

All tests passed.

Tous les quatre cas passent avec une erreur maximale bien inférieure à la tolérance du bfloat16, confirmant que le noyau est numériquement correct sur du matériel Trainium réel.

ÉTAPE 2 : PROFILER LE NOYAU SOFTMAX

Une fois le noyau compilé et produisant des résultats numériquement corrects, l’étape suivante consiste à profiler son exécution sur le matériel pour identifier les goulots d’étranglement et guider les optimisations.

Pour démontrer le profilage et l’analyse sur une charge de travail réelle, cette étape utilise un noyau SwiGLU MLP, un module courant dans les grands modèles de langage (LLM).

POINTER L’AGENT SUR LE NOYAU SWIGLU

Demandez à l’agent d’analyser le profil du noyau SwiGLU. L’agent compile d’abord le noyau en un fichier NEFF et capture une trace NTFF via neuron-explorer. Ensuite, il mène une investigation en deux parties : d’abord une analyse des statistiques et des limites de performance au niveau du noyau, puis une plongée dans les inefficacités spécifiques en interrogeant le profil au niveau de l’exécution des instructions.

ANALYSE DES LIMITES DE PERFORMANCE

L’agent exécute une analyse complète des limites sur le profil capturé et identifie plusieurs écarts à investiguer :

Les investigations aident à auditer les écarts et à prioriser les directions d’optimisation actionnables. Alors que le moteur bottleneck (TE, ou Tensor Engine) serait la cible principale d’optimisation, l’agent découvre que les instructions de multiplication matricielle (matmul) du NKI fonctionnent déjà près de leur efficacité maximale. En revanche, les instructions DMA sont bien en dessous de leur taille cible (inefficaces) et les entrées sont rechargées huit fois (redondant). L’agent identifie même les trois lignes exactes de code NKI responsables des transferts sous-optimaux.

Corriger ces lignes pourrait réduire les périodes d’inactivité du TE et améliorer la latence du noyau.

CONSIDÉRATIONS PRATIQUES POUR L’UTILISATION

Pour éviter des frais horaires inutiles, terminez l’instance trn2.3xlarge que vous avez créée au début du tutoriel. Vous pouvez le faire via la console AWS (EC2 > Instances, sélectionnez l’instance, puis choisissez État de l’instance > Terminer), ou en exécutant :

aws ec2 terminate-instances --instance-ids <your-instance-id>

Confirmez que l’état de l’instance affiche « terminated » avant de fermer la console.

AU-DELÀ DE L’AUTOMATISATION : UNE VISION PLUS LARGE

Les compétences d’écriture et de profilage de noyaux réduisent la barrière à l’entrée pour l’écriture de noyaux haute performance sur Trainium. Mais elles ne représentent qu’une partie d’une vision plus large. Aujourd’hui, les développeurs utilisent les insights de profilage pour guider leur prochaine série d’éditions de noyaux. Ce cycle itératif (profiler, diagnostiquer, refactoriser, re-profiler) est celui où le plus de temps est passé. L’objectif est de rendre ce cycle entièrement agentique : des agents qui itèrent automatiquement sur un noyau jusqu’à ce qu’il atteigne sa cible de performance, sans nécessiter que le développeur interprète chaque résultat de profilage et ne fabrique manuellement la correction suivante.

Les développeurs de performances soulignent aussi que les noyaux personnalisés ne sont qu’une partie d’un défi plus large. Ils veulent que leurs modèles fonctionnent sur Trainium sans avoir à se soucier du portage du code, de la résolution des écarts d’opérateurs, de l’application des optimisations au niveau du modèle ou de la validation de la correction à grande échelle. L’approche agentique vise à résoudre ce problème plus large.

LA VISION : DES MODÈLES DE FRONTIÈRE PLUS PERFORMANTS SUR TRAINIUM

En résumé, l’objectif est de soutenir la prochaine vague d’innovations pour les modèles de frontière en utilisant Trainium et le SDK Neuron, et d’utiliser la suite de capacités Neuron Agentic Development pour atteindre des performances coût-leader pour des cas d’usage allant de l’expérimentation avec de nouvelles architectures de modèles à l’exécution de modèles de production à grande échelle.

Plus de détails seront partagés à mesure que ces capacités mûrissent. Pour commencer avec ce qui est disponible aujourd’hui, consultez le dépôt GitHub Neuron Agentic Development.

COMMENT DÉBUTER ?

Les capacités Neuron Agentic Development sont disponibles dès aujourd’hui. Pour commencer, clonez le dépôt neuron-agentic-development et écrivez votre premier noyau NKI en quelques minutes.