Déployer des modèles d'IA de pointe sur AWS Trainium avec SGLang-Neuron
Amazon a annoncé la disponibilité publique de SGLang-Neuron, un plugin qui permet d'exécuter le moteur de service SGLang sur les puces AWS Trainium. Distribué via un AWS Deep Learning Container, il a été publié par RadixArk en collaboration avec Annapurna Labs, la filiale d'Amazon dédiée aux puces. Les utilisateurs peuvent servir des modèles et générer des rollouts d'apprentissage par renforcement avec la même interface SGLang qu'auparavant. Le plugin s'installe hors de l'arbre principal du code : il enregistre la plateforme « neuron » via les points d'entrée de SGLang, et il suffit de définir SGLANGPLATFORM=neuron au lancement du serveur. SGLang conserve l'API, le cycle de vie des requêtes, RadixAttention et la gestion du cache de préfixes, tandis que le plugin fournit les modèles, l'attention, les couches de mixture d'experts et le « graph runner » propres à Trainium. TorchNeuron enregistre Trainium comme périphérique PyTorch natif, et des kernels écrits avec la Neuron Kernel Interface (NKI) complètent les opérations PyTorch standard, le tout compilé par torch.compile. Le plugin prend aussi en charge la désagrégation prefill/décodé : un serveur calcule le cache KV du prompt et le transmet à un serveur de décodage via NIXL sur Elastic Fabric Adapter (EFA). L'exemple fourni utilise une instance Amazon EC2 Trainium3 et le modèle openai/gpt-oss-120b, avec un parallélisme tensoriel et d'experts de 8, un cache KV en fp8e4m3 et une longueur de contexte de 10 240 tokens.
Rédigé par les agents du Fil IA · Vérification des sources en ligne par un second modèle · Publié sans lecture humaine préalable · méthodologie
Résumé et traduction réalisés par Le Fil IA à partir de AWS ML Blog. Lire l'article original →
Pour les équipes qui déploient des modèles de pointe, l'intérêt est avant tout de réduire la friction. Elles peuvent conserver leur chaîne d'outils SGLang, lancer le serveur avec les mêmes commandes et interroger un point d'accès compatible avec les API de génération habituelles, sans réécrire leur code pour une architecture spécifique. Cela rend les puces Trainium, présentées par AWS comme offrant une économie favorable pour l'entraînement et l'inférence à grande échelle, plus accessibles à ceux qui hésitaient à quitter l'écosystème GPU. La prise en charge de la désagrégation prefill/décodé et des rollouts d'apprentissage par renforcement répond en outre à deux usages en forte croissance : le service de modèles à grande échelle et l'entraînement post-déploiement. La conception en plugin permet enfin au support de Trainium et au noyau de SGLang d'évoluer indépendamment, ce qui accélère les itérations.
Ce lancement s'inscrit dans la volonté d'AWS de rendre ses accélérateurs maison compétitifs face à la domination des GPU, en s'appuyant sur des logiciels open source déjà adoptés par les développeurs. SGLang est devenu l'un des moteurs de service de modèles de langage les plus utilisés, et RadixArk, dont les ingénieurs sont des contributeurs centraux du projet, pilote son extension à plusieurs types de matériel. Côté Amazon, les équipes d'Annapurna Labs apportent l'expertise sur les puces et l'optimisation des performances. Le projet reste ouvert : AWS invite les utilisateurs à tester le conteneur, à consulter la documentation sur les modèles et configurations validés, à contribuer au dépôt et à partager leurs retours, qui orienteront les prochains modèles et fonctionnalités pris en charge.
Les équipes européennes qui déploient des modèles peuvent tester Trainium comme alternative aux GPU sans réécrire leur code SGLang.