KAIST supprime la calibration du décodage MoE et accélère les blocs jusqu'à 2,69× par rapport à vLLM

Le décodage des Mixture-of-Experts sur petits batchs est limité non par le calcul, mais par la mémoire : pour chaque token, il faut charger les poids des experts routés. Geunhyeon Kim et Sung Ju Hwang, chercheurs à KAIST AI, montrent qu'une technique d'accélération répandue, la sélection des canaux importants au sein d'un expert, peut se passer de données de calibration. Leur lecteur, construit uniquement à partir des poids du modèle, conserve 98,1% de la précision du modèle complet en ne lisant que 17,8% des poids des experts. Les blocs MoE en noyaux Triton tournent 1,60 à 1,90 fois plus vite que vLLM optimisé en batch one.

Le seuil promet un budget, mais ne le tient pas

L'idée de la sélection de canaux est simple. Dans chaque expert, la majorité des canaux intermédiaires influence très peu un token donné. Un sélecteur prédit quels canaux comptent et ne lit en mémoire que leurs poids. En général, on retient les canaux dont le score dépasse un seuil, lequel est ajusté sur un échantillon hors ligne.

Les auteurs montrent que ce schéma présente deux faiblesses. D'abord, le seuil ne fixe pas le volume de lecture : le nombre de canaux qui passent la coupure varie d'une entrée à l'autre. Prox, avec des seuils calibrés pour un budget de 18%, lisait en réalité 12,9% en moyenne sur leurs benchmarks. L'ingénieur qui planifie latence et débit obtient donc un chiffre différent de celui sur lequel il comptait.

Ensuite, gagner en précision à budget donné n'est pas garanti non plus.

Top-k plutôt que seuil

L'alternative ne nécessite aucune donnée : pour chaque token, on conserve les k canaux au score le plus élevé. Le budget de lecture est connu à l'avance et correspond à celui annoncé. À scores identiques, la version top-k du sélecteur a égalé ou dépassé la version à seuil dans 40 cas sur 43, pour des budgets de 13% et 18%. Sur les tâches de génération, elle l'a emporté dans les 23 comparaisons.

Il y a aussi un effet qualitatif. Les seuils calibrés faisaient sombrer Gemma 4 dans des boucles de répétition, alors que les versions top-k les évitent. La conclusion des auteurs : la calibration du seuil n'apporte aucune des deux choses pour lesquelles on la pratique, ni budget de lecture connu, ni gain de précision stable à ce budget.

Un lecteur construit à partir des seuls poids

La seconde partie du travail est un lecteur maison, qui n'a besoin ni de données de calibration ni de fine-tuning. Il procède en deux étapes :

• les canaux sont d'abord écartés grossièrement à l'aide d'une esquisse ternaire (ternary sketch) des poids ;
• les canaux ignorés sont compensés par une application de bas rang, elle aussi construite à partir des poids du modèle.

Chiffres clés sur trois modèles MoE post-entraînés :

• précision : 98,1% de celle du modèle complet ;
• volume de lecture des poids des experts routés : 17,8% ;
• vitesse des blocs MoE en noyaux Triton en batch one : 1,60 à 1,90 fois supérieure à celle de vLLM optimisé ;
• sur les gros batchs, l'accélération atteint 2,69 fois sur les mêmes trois modèles ;
• Prox, calibré pour 18%, lisait en réalité 12,9% ;
• top-k contre seuil : 40 comparaisons sur 43 pour des budgets de 13% et 18%, 23 sur 23 en génération.

À qui cela sert

Les principaux bénéficiaires sont ceux qui font tourner de gros modèles MoE sur du matériel à mémoire limitée. Dans une discussion sur Hugging Face, l'un des auteurs a qualifié sans détour le problème de l'élagage d'activations d'aigu pour les GPU à mémoire lente, comme le DGX Spark. Sur ces machines, c'est la bande passante mémoire, et non les FLOPS, qui détermine la vitesse de génération des tokens. Selon lui, le prochain régime de sparsité viendra précisément de la sparsité d'activation pour le serving sur des appareils soumis à une limite stricte de bande passante.

L'absence de calibration présente un avantage pratique : la méthode s'applique à n'importe quel modèle open-weights récent dès sa sortie. Inutile de chercher un échantillon représentatif, et il n'y a pas de risque qu'un seuil ajusté sur un domaine fasse s'effondrer la génération sur un autre, comme cela s'est produit avec Gemma 4. Un budget de lecture fixe simplifie les calculs : les 17,8% retenus se traduisent par une économie prévisible de trafic mémoire et de latence.

Le code est publié dans le dépôt calibration-free-moe sur GitHub. Pour les équipes qui écrivent elles-mêmes leurs noyaux Triton ou optimisent l'inférence locale, c'est un point de départ prêt à l'emploi. Il vaut la peine de le tester sur ses propres modèles, en particulier dans les scénarios interactifs en batch one, où la mémoire est le plus fortement sollicitée.