Trois optimisations pour une inférence MoE plus rapide sur NVIDIA DGX Spark

Les ingénieurs qui optimisent les grands modèles de langage se heurtent souvent à un mur où l'ajout de puissance de calcul n'accélère pas la génération. Une récente évaluation technique du modèle Ling-3.0-flash sur le matériel NVIDIA DGX Spark a identifié que les performances de décodage étaient limitées par la bande passante mémoire plutôt que par la puissance de calcul brute. En ciblant des goulets d'étranglement spécifiques dans la lecture des poids et la vérification des tokens, les chercheurs ont obtenu des réductions significatives de latence sans sacrifier la précision.

Ce qui s'est passé

L'évaluation s'est concentrée sur Ling-3.0-flash, un modèle mixture-of-experts (MoE) d'Inclusion AI comptant 124 milliards de paramètres au total, mais seulement 5,1 milliards de paramètres actifs par token. Lors de l'exécution avec une taille de lot de un, les outils de profilage ont montré que le temps consacré à la génération des tokens était dominé par deux opérations : le calcul MoE lui-même et la projection de la couche de sortie, connue sous le nom de tête LM. Ces étapes nécessitaient la lecture de grandes quantités de données de poids depuis la mémoire, créant un goulot d'étranglement que les techniques d'optimisation standard ne parvenaient pas à résoudre entièrement.

Pour y remédier, l'équipe a testé trois optimisations distinctes. Premièrement, ils ont remplacé le backend FlashInfer MXFP4 MoE par défaut par la bibliothèque de noyaux Humming. Deuxièmement, ils ont appliqué une quantization FP8 en ligne spécifiquement à la tête LM pour réduire son empreinte mémoire. Troisièmement, ils ont mis en œuvre un décodage spéculatif à prédiction multi-tokens utilisant la couche NEXTN intégrée au modèle. Chaque modification ciblait un aspect différent du pipeline d'inférence, allant de l'efficacité de l'exécution des noyaux au volume de données et à la réduction des étapes.

Les résultats ont été mesurés sur des charges de travail synthétiques et sur l'ensemble complet de tests GSM8K. Alors que les tests synthétiques montraient des gains massifs dus aux motifs de sortie répétitifs, l'évaluation réelle sur 1 319 questions GSM8K a démontré une augmentation de 42,5 % du débit de décodage et une réduction de 24,5 % de la latence moyenne. Crucialement, ces améliorations se sont accompagnées d'aucun changement statistiquement détectable dans la précision du modèle, maintenant un score de 96,82 % par rapport à la référence de 96,89 %.

Comment cela fonctionne

La première optimisation exploite Humming, une bibliothèque de noyaux indépendante qui utilise la compilation à l'exécution CUDA de NVIDIA pour générer du code adapté aux dimensions spécifiques du modèle. Contrairement au chemin FlashInfer de base, qui compile de nombreuses cibles génériques, Humming crée des binaires spécialisés pour les formes matricielles réelles utilisées par Ling-3.0-flash. Cette spécialisation réduit le temps par token de sortie en garantissant que le GPU exécute uniquement les instructions nécessaires pour la taille intermédiaire étroite du modèle.

La deuxième technique concerne la tête LM, qui projette les états cachés vers un vocabulaire de plus de 157 000 tokens. Dans la configuration de base, cette couche utilise une précision BF16, nécessitant une bande passante mémoire importante pour chaque token généré. En passant à une quantization FP8 en ligne pour cette couche uniquement, le système réduit de moitié le volume de données de poids lu lors de chaque étape. Ce changement est mutuellement exclusif avec l'option de tête LM FP32 de base et nécessite une configuration soignée pour éviter les problèmes de compatibilité avec d'autres parties du modèle.

La troisième méthode utilise le décodage spéculatif pour générer plusieurs tokens en une seule ronde de vérification. Ling-3.0-flash inclut une couche native NEXTN qui agit comme un modèle de brouillon, prédisant les tokens suivants que le modèle principal vérifie ensuite. Au lieu de générer un token à la fois, le système rédige trois tokens et vérifie quatre positions simultanément. Cette approche réduit le nombre total d'étapes de décodage requises, bien qu'elle augmente la complexité de chaque étape de vérification en raison d'un nombre plus élevé d'activations d'experts.

Détails clés

• Ling-3.0-flash possède 124 milliards de paramètres au total avec 5,1 milliards de paramètres actifs par token répartis sur 42 couches.
• Le remplacement du backend FlashInfer par Humming a réduit le temps par token de sortie de 27,71 ms à 25,35 ms, soit un gain de 9,3 %.
• La quantization FP8 en ligne de la tête LM a réduit de moitié le volume de lecture des poids pour cette couche, contribuant à une amélioration cumulative du débit de 21,7 %.
• La prédiction multi-tokens avec NEXTN a atteint 3,90 tokens acceptés par ronde sur des données synthétiques, augmentant le débit de 185,1 % dans ce test spécifique.
• Sur le benchmark GSM8K, les optimisations combinées ont augmenté le débit de décodage de 42,5 % et réduit la latence moyenne de 24,5 %.
• La précision est restée stable, ne chutant que de 0,08 point de pourcentage, passant de 96,89 % à 96,82 %, ce qui n'est pas statistiquement significatif.

Pourquoi c'est important

Pour les développeurs construisant des workflows agentiques ou des applications en temps réel, la latence est souvent la contrainte principale. Cette analyse confirme que la bande passante mémoire, et non seulement la puissance de calcul, est le facteur limitant pour de nombreux modèles MoE à faible taille de lot. En se concentrant sur la réduction du trafic de poids et l'augmentation du nombre de tokens par étape, les ingénieurs peuvent obtenir des gains de performance substantiels sans avoir besoin de clusters matériels plus importants. Cela est particulièrement pertinent pour les déploiements en périphérie ou les environnements cloud sensibles aux coûts, où chaque milliseconde compte.

Cependant, l'étude met également en lumière les risques liés à la combinaison d'optimisations sans une connaissance approfondie de l'intégration. Les chercheurs ont découvert que l'activation conjointe de la quantization FP8 de la tête LM et de la prédiction multi-tokens exigeait que le modèle de brouillon partage le module entier de la tête LM. Sans cette interface spécifique, le système basculait silencieusement vers une configuration plus lente, interprétant les poids FP8 comme étant en BF16. Cela souligne l'importance de valider chaque optimisation au démarrage plutôt que de supposer la compatibilité basée sur des tests individuels des composants.

Ce que vous pouvez faire

• Profilez votre pipeline d'inférence MoE actuel pour identifier si le décodage est limité par la mémoire ou par le calcul avant d'appliquer des optimisations.
• Testez la bibliothèque de noyaux Humming comme remplacement direct de FlashInfer si votre modèle présente des dimensions matricielles spécialisées.
• Activez la quantization FP8 en ligne pour la tête LM en utilisant la variable d'environnement SGLANG_ENABLE_FP8_LM_HEAD pour réduire l'utilisation de la bande passante mémoire.
• Mettez en œuvre la prédiction multi-tokens en utilisant des couches de brouillon natives comme NEXTN pour réduire le nombre d'étapes séquentielles de décodage.
• Vérifiez que les modèles de brouillon et cible partagent les modules corrects lorsque vous combinez la quantization avec le décodage spéculatif pour éviter une dégradation silencieuse des performances.
• Effectuez des benchmarks des modifications sur des charges de travail réalistes comme GSM8K plutôt que de compter uniquement sur des tests synthétiques, qui peuvent surestimer les bénéfices du décodage spéculatif.