Ingénierie des noyaux Mixture-of-Experts sur Amazon SageMaker AI
aws ·
Cet article détaille les optimisations d'ingénierie de noyaux pour les modèles Mixture-of-Experts (MoE) sur Amazon SageMaker AI, en mettant l'accent sur les gains de performance pour l'inférence des grands modèles de langage. Les auteurs décrivent trois techniques clés : la transposition Swap-AB pour aligner les dimensions des matrices avec les exigences matérielles WGMMA du processeur NVIDIA Hopper, des ajustements du routage backend pour utiliser Triton au lieu de CUTLASS pour la quantification FP8, et l'élimination des recompilations JIT Triton inutiles afin de réduire la latence. Ces optimisations répondent à des goulots d'étranglement spécifiques aux architectures MoE, tels que les tailles de lot minces lors du décodage et les nombres variables de tokens. La mise en œuvre de ces modifications a permis d'augmenter le débit jusqu'à 10 % et d'améliorer le Time to First Token (TTFT) de 9 % sur les GPU H200. Plusieurs de ces améliorations ont été intégrées dans la communauté open-source vLLM. Ce travail démontre comment un réglage fin spécifique au matériel peut considérablement améliorer l'efficacité des modèles à activation sparse.
aws