Résumé Technique : LYNX – Permettre une Inférence MoE Efficace par une Sélection Dynamique d'Experts Sensible aux Lots

1. Énoncé du Problème

Les modèles fondamentaux modernes utilisent de plus en plus des architectures de Mélange d'Experts (MoE) pour découpler le nombre de paramètres du coût computationnel. Bien que les MoE n'activent qu'un petit sous-ensemble de paramètres par token (par exemple, 14 sur 57 milliards dans Qwen2-57B), cette promesse d'efficacité est rompue dans le déploiement en production en raison de la tension entre la latence et le débit.

Pour atteindre un débit élevé, les systèmes d'inférence regroupent les requêtes en lots. Cependant, comme chaque token d'un lot sélectionne indépendamment ses $k$ meilleurs experts, l'ensemble des experts devant être chargés depuis la mémoire à large bande passante (HBM) croît avec la taille et la composition du lot. Aux tailles de lots modérées requises pour respecter les objectifs de niveau de service (SLO) en matière de latence, ce motif d'accès mémoire dynamique et dépendant des données rend le décodage MoE limité par la bande passante mémoire. Par conséquent, la latence évolue avec le nombre total d'experts distincts activés dans l'ensemble du lot, saturant souvent l'ensemble du pool d'experts (par exemple, un lot de seulement 8 requêtes diverses peut activer les 64 experts de Qwen2-57B). Cela annule les bénéfices de la parcimonie des MoE, les faisant performer de manière similaire à des modèles denses de capacité équivalente, mais avec la surcharge supplémentaire du déplacement dynamique des experts.

Les solutions existantes, telles que l'élagage statique, la quantification ou le regroupement, reposent sur un étalonnage hors ligne et altèrent le modèle de manière permanente, les rendant rigides face aux changements de charge de travail. Les approches dynamiques par token échouent à résoudre le goulot d'étranglement car elles ne réduisent pas le nombre agrégé d'experts récupérés au niveau du lot.

2. Méthodologie : LYNX

LYNX est un système agnostique à la charge de travail conçu pour atténuer le goulot d'étranglement de la bande passante mémoire dans l'inférence MoE sans modifier l'architecture du modèle ni nécessiter d'étalonnage hors ligne. Il exploite une idée clé dérivée des dynamiques d'entraînement des MoE : les pertes d'équilibrage de charge forcent les routeurs à distribuer les tokens largement parmi les experts, même lorsque la confiance dans les sélections secondaires est faible. Cela crée une redondance au niveau du lot où le routeur sélectionne des experts en raison de contraintes d'entraînement plutôt que d'une véritable affinité token-expert.

LYNX exploite cette redondance grâce à une nouvelle technique d'AffinityBinning et à un pipeline à trois composants opérant entièrement au moment de l'exécution :

2.1 Principes Fondamentaux

La conception repose sur trois observations empiriques :

1. Sensibilité de la Phase : Les MoE sont limités par le calcul pendant la phase de préremplissage (prefill) mais limités par la bande passante mémoire pendant la phase de décodage (decode). Le remappage des experts n'est bénéfique que pendant le décodage. De plus, la phase de préremplissage est hautement sensible à la fidélité des experts, tandis que la phase de décodage est résiliente face à la réaffectation.
2. Confiance du Token : Les scores de confiance du routeur (dérivés de l'écart entre les logits des experts les mieux classés) identifient de manière fiable quels mappages token-expert sont critiques. Les tokens à faible confiance peuvent être remappés avec une dégradation minimale de la précision.
3. Hiérarchie du Rang des Experts : L'expert le mieux classé (rang-0) domine la qualité de la sortie. Les experts de rang inférieur parmi les $k$ meilleurs présentent une forte redondance. La sensibilité à la réaffectation varie également selon les couches.

2.2 Composants du Système

LYNX s'intègre au moteur de service (par exemple, vLLM) via un optimiseur sensible à la phase et trois composants de noyaux CUDA légers :

• Optimiseur Sensible à la Phase : Détermine si l'itération actuelle est limitée par la mémoire (décodage) ou par le calcul (préremplissage). Il contourne la logique LYNX pendant les phases de préremplissage ou de préremplissage par blocs pour éviter une surcharge inutile.
• Analyseur de Confiance : Intercepte les logits du routeur à chaque couche. Il discrétise la confiance du routeur de chaque token en classes (bins) en utilisant l'AffinityBinning. La largeur et le nombre de classes sont déterminés uniquement par le ratio de parcimonie du modèle ($k/N$), rendant le système auto-étalonné et indépendant de la charge de travail. Il identifie les tokens ayant une faible affinité avec leurs experts assignés.
• Scoreur Adaptatif d'Experts : Calcule un score à l'échelle du lot pour chaque expert. Il utilise un schéma de pondération exponentielle basé sur les valeurs de classes attribuées par l'Analyseur de Confiance. Cela garantit que les experts préférés par les tokens à forte conviction reçoivent des scores disproportionnellement plus élevés. Il détermine l'ensemble minimal d'experts "critiques" requis pour le lot.
• Remappeur d'Experts : Redirige les affectations token-expert à faible affinité vers l'ensemble d'experts critiques survivants. Crucialement, LYNX préserve la sémantique d'activation des $k$ meilleurs ; chaque token active toujours exactement $k$ experts, mais le lot active collectivement un ensemble plus petit et consolidé d'experts uniques.

2.3 Implémentation

LYNX est implémenté comme un module plug-and-play au-dessus de vLLM. Sa logique centrale est réalisée à l'aide de quatre noyaux fusionnés Triton/CUDA qui interceptent la sortie du routeur. Cette conception élimine le mouvement intermédiaire des données de tenseurs et réduit la surcharge de lancement des noyaux, garantissant que le système ajoute une latence négligeable (<4 % de surcharge) tout en opérant dans le chemin critique.

3. Contributions Clés

1. Premier Système Agnostique à la Charge de Travail : LYNX est le premier système à atténuer le goulot d'étranglement de la bande passante mémoire dans l'inférence MoE sans nécessiter de données d'étalonnage ni de modifications permanentes du modèle.
2. Technique AffinityBinning : Une méthode novatrice pour discrétiser la confiance du routeur basée sur les ratios de parcimonie architecturale, permettant une adaptation automatique à n'importe quelle architecture MoE.
3. Optimisation au Niveau du Lot : Contrairement aux approches dynamiques par token précédentes, LYNX opère au niveau du lot pour réduire le nombre total d'experts uniques récupérés depuis la HBM, s'attaquant directement au goulot d'étranglement de la bande passante mémoire.
4. Complémentarité : LYNX est conçu pour fonctionner en conjonction avec des optimisations existantes, y compris le déchargement (offloading) et la quantification.

4. Résultats Expérimentaux

LYNX a été évalué sur des modèles de pointe de quatre familles (Qwen, DeepSeek, Mixtral, Llama) à travers huit benchmarks (génération de code, mathématiques, raisonnement et vision).

• Réduction de la Latence : LYNX atteint jusqu'à 1,30× de latence inférieure (Temps par Token de Sortie) par rapport aux références. Dans les scénarios de service disjoints (où le préremplissage et le décodage sont séparés), les réductions atteignent jusqu'à 1,35×.
• Impact sur la Précision : Le système entraîne moins de 1 % de perte de précision dans le pire des cas. Notamment, LYNX améliore souvent la précision en moyenne (par exemple, +1,9 % sur HumanEval pour DeepSeek-Coder-V2), suggérant que l'équilibrage de charge lors de l'entraînement force des sélections d'experts sous-optimales que LYNX corrige.
• Débit : Sous des budgets SLO fixes (par exemple, 25 ms ou 50 ms), LYNX améliore le débit du système jusqu'à 2,1×.
• Compatibilité :
  • Déchargement : Combiné à des techniques de déchargement (par exemple, Fiddler), LYNX booste les performances de 31 %.
  • Quantification : Appliqué au-dessus de la quantification INT4, LYNX fournit un gain de vitesse supplémentaire de 10 %.
  • Parallélisme : LYNX apporte des gains de débit additifs au-dessus du Parallélisme d'Experts (EP) et du Parallélisme de Tenseurs (TP).
• Robustesse : Les gains de performance sont cohérents à travers des tailles de lots variables (saturant aux tailles modérées) et des traces réelles avec des taux d'arrivée en rafales.

5. Signification

L'article postule que, à mesure que le calcul GPU continue de dépasser la croissance de la bande passante mémoire, l'écart d'intensité arithmétique pour le décodage MoE s'élargira, rendant les goulots d'étranglement limités par la mémoire de plus en plus critiques. LYNX y répond en changeant fondamentalement la manière dont les experts sont sélectionnés au niveau du lot. En exploitant la redondance introduite par l'équilibrage de charge lors de l'entraînement, LYNX permet aux architectures MoE de réaliser leur promesse d'efficacité sur le matériel futur sans nécessiter de réentraînement du modèle ni de profilage hors ligne. Le système démontre qu'une gestion dynamique et consciente du temps d'exécution des experts peut découpler significativement le débit des contraintes de bande passante mémoire inhérentes aux paradigmes actuels de service MoE.