Ingénierie de kernels Mixture-of-Experts sur Amazon SageMaker AI
Amazon détaille les optimisations de noyaux GPU qu'elle a mises au point pour les modèles Mixture-of-Experts (MoE) sur Amazon SageMaker AI. Cette architecture n'active qu'une partie de ses paramètres pour chaque token, ce qui permet d'augmenter la capacité d'un grand modèle de langage sans faire croître proportionnellement le calcul. Les opérations MoE représentent toutefois une part importante du temps de calcul total, si bien que toute inefficacité se répercute directement sur la latence et le coût des requêtes. La technique centrale est le « fused MoE », qui fusionne le routage des tokens, le réordonnancement des données et les calculs des experts en une seule opération groupée sur le GPU, au lieu de boucles séquentielles distinctes. Le premier levier décrit s'appelle « Swap-AB » et vise l'architecture NVIDIA Hopper. Celle-ci repose sur l'instruction WGMMA, exécutée par un groupe coopératif de 128 threads. Au niveau de l'assembleur (wgmma.mma_async), la dimension de lignes M de l'opérande A est fixée à 64, tandis que la dimension de colonnes N reste variable. Les deux emplacements d'opérandes ne sont donc pas interchangeables.
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 →
Cette contrainte devient un goulot d'étranglement sévère pendant la phase de décodage. Dans le calcul standard A[M,K] x B[K,N], M correspond au nombre de tokens, K à la dimension cachée et N à la dimension de sortie de l'expert. En décodage, chaque séquence utilisateur ne traite qu'un token à la fois, de sorte que M tombe souvent sous 64 et descend fréquemment à 1. La matrice A devient alors un mince vecteur ligne, par exemple 1 x 4096, tandis que les poids d'expert B forment un grand bloc dense, par exemple 4096 x 8192. Comme A ne peut pas remplir les 64 lignes exigées, une grande partie de la capacité des Tensor Cores reste inutilisée et le silicium consomme des cycles sur des emplacements vides. Le gain visé par l'inversion des opérandes est donc une meilleure utilisation du matériel, avec à la clé des réponses plus rapides et moins coûteuses pour les utilisateurs de modèles MoE en production.
Ce travail s'inscrit dans la montée en puissance des architectures MoE, devenues courantes pour déployer de très grands modèles à un coût d'inférence maîtrisé. L'enjeu est que le gain théorique de la sparsité ne se concrétise que si les noyaux exploitent finement le matériel. Dans une couche dense classique, les matrices de poids sont statiques, unifiées et contiguës. Dans une couche MoE, la charge est au contraire fragmentée entre de nombreux experts, chacun ne recevant qu'un petit nombre de tokens, ce qui aggrave le problème des matrices trop minces. Les noyaux fusionnés complexifient en retour l'implémentation et le débogage, et exigent des optimisations propres à chaque génération de GPU. L'article présente Swap-AB comme la première d'une série d'optimisations, et ses auteurs entendent montrer comment les clients de SageMaker AI peuvent en profiter. Le texte fourni s'arrête avant la description des autres techniques et des résultats chiffrés.
Pas d\'impact direct sur la France/UE