Aller au contenu

6 · MoE et GEMM groupée

Le mélange d'experts est l'architecture dominante des grands modèles depuis 2024. Du point de vue du GPU, c'est un problème d'indirection, de déséquilibre de charge et de communication — trois choses que le GPU fait mal.

6.1 Le principe

Dans une couche MoE, le MLP dense est remplacé par \(E\) experts, dont seuls \(k\) sont activés par jeton.

jeton x
  │
  ├─→ routeur (petite matrice d × E) ──→ scores
  │                                        │
  │                            top-k des experts
  │                                        │
  ├─→ expert i₁ ──┐
  ├─→ expert i₂ ──┤ combinaison pondérée ──→ sortie
  └─→ ...       ──┘

L'intérêt : le nombre de paramètres croît avec \(E\), le calcul avec \(k\). Un modèle à 512 experts dont 8 sont actifs a 64 fois plus de paramètres qu'un dense équivalent en calcul.

Le coût côté GPU :

Aspect Dense MoE
Calcul par jeton \(2\Theta\) \(2\Theta_{\text{actifs}}\)
Poids lus par jeton (lot 1) \(2\Theta\) \(2\Theta_{\text{actifs}}\)
Poids lus (grand lot) \(2\Theta\) jusqu'à \(2\Theta\) (tous les experts sollicités)
Régularité parfaite indirection
Communication multi-GPU collectives régulières all-to-all irrégulier

Pourquoi le MoE est excellent en décodage à petit lot

À lot 1, on ne lit que les experts activés. Un modèle de 400 milliards de paramètres dont 30 sont actifs lit 60 Go au lieu de 800 en BF16 — le régime limité par la mémoire est massivement amélioré.

À grand lot, en revanche, tous les experts finissent par être sollicités : on lit tout, et l'avantage disparaît. Le MoE optimise donc précisément le régime où l'inférence est le plus contrainte.

6.2 Le problème algorithmique : de l'indirection à la GEMM groupée

Après le routage, on a une liste de couples (jeton, expert). Deux stratégies.

Stratégie 1 — masquage

Chaque expert traite tous les jetons, et on masque ceux qui ne lui sont pas assignés.

Simple à écrire, et catastrophique : on fait \(E\) fois le travail au lieu de \(k\) fois. Utilisable seulement pour du prototypage.

Stratégie 2 — permutation puis GEMM groupée

C'est la bonne approche.

1. Router : chaque jeton → k experts
2. Trier les jetons par expert  (tri par clés, ou scan sur des compteurs)
3. Rassembler (gather) : construire des tableaux contigus par expert
4. GEMM groupée : une GEMM par expert, tailles différentes
5. Disperser (scatter) : replacer les résultats
6. Combiner : somme pondérée des k contributions

Les étapes 2, 3 et 5 sont du rangement de données pur — pas de calcul utile, seulement du trafic mémoire. Elles peuvent être coûteuses.

Chiffre concret : Mirage MPK mesure que le prétraitement du MoE (le rassemblement avant la GEMM) représente jusqu'à 11 % du temps d'exécution total du MoE, et le supprime en fusionnant le gather dans la GEMM — l'expert lit directement à travers la table d'indices, au lieu de travailler sur une copie contiguë.

6.3 La GEMM groupée

\(E\) multiplications matricielles indépendantes, de tailles différentes :

\[ \mathbf{C}_e = \mathbf{A}_e \mathbf{B}_e, \qquad e = 1 \dots E \]

où \(\mathbf{A}_e\) a \(n_e\) lignes (\(n_e\) = nombre de jetons routés vers l'expert \(e\)), variable et connu seulement à l'exécution.

Pourquoi ce n'est pas trivial :

  • lancer \(E\) GEMM séparées coûte \(E\) lancements et laisse le GPU vide sur les petits experts ;
  • les tailles ne sont connues qu'après le routage, donc les paramètres de lancement doivent être calculés sur le GPU ;
  • le déséquilibre entre experts crée un déséquilibre de charge entre SM.

Les solutions :

Solution Disponibilité
cublasLtMatmul en mode groupé CUDA 13.1+
CUTLASS GroupedGemm depuis longtemps
DeepGEMM (DeepSeek) GEMM groupée FP8, à échelle fine
Triton, avec un ordonnanceur de tuiles manuel à écrire
Marlin MoE pour les poids quantifiés

L'ajout de la GEMM groupée dans cuBLAS avec CUDA 13.1 est significatif : c'était jusque-là une lacune qui obligeait à passer par des bibliothèques tierces.

Le motif de l'ordonnanceur de tuiles

L'implémentation efficace ne lance pas \(E\) noyaux mais un seul noyau persistant qui consomme une file de tuiles :

file de travail = [(expert 0, tuile 0), (expert 0, tuile 1),
                   (expert 1, tuile 0), ...]

chaque bloc :
    tant qu'il reste du travail :
        (e, t) = prochaine tuile   ← atomicAdd sur un compteur
        calculer la tuile t de l'expert e

C'est un noyau persistant avec une file de tâches — exactement la structure d'un megakernel, appliquée à une seule opération. Le lien avec la partie 8 est direct : un megakernel généralise ce motif à toutes les opérations du modèle.

6.4 Le déséquilibre de charge

Le routage n'est pas uniforme : certains experts reçoivent bien plus de jetons que d'autres.

Les conséquences :

  • les SM traitant un expert chargé finissent bien après les autres ;
  • avec une capacité fixe par expert, les jetons excédentaires sont jetés (token dropping), ce qui dégrade la qualité.

Les remèdes, côté modèle et côté système :

Remède Nature
Perte auxiliaire d'équilibrage entraînement
Bruit sur le routeur entraînement
Capacité par expert avec dépassement système, dégrade la qualité
Ordonnancement dynamique (file de tuiles) système, sans dégradation
Experts partagés (toujours actifs) architecture

L'ordonnancement dynamique est la bonne réponse système : au lieu d'affecter statiquement des SM aux experts, on laisse les blocs prendre du travail dans une file globale. C'est ce que fait le mécanisme de distribution juste-à-temps de Mirage MPK, et c'est aussi ce que le Cluster Launch Control de Blackwell fournit en matériel.

6.5 Le MoE distribué : le dispatch all-to-all

Sur plusieurs GPU, les experts sont répartis (parallélisme d'experts). Chaque jeton doit donc être envoyé au GPU qui héberge son expert, puis le résultat renvoyé.

C'est un all-to-all irrégulier : chaque GPU envoie un nombre différent de jetons à chaque autre GPU, et ce nombre change à chaque lot.

GPU 0 : jetons pour experts {0,1} localement, {2,3} sur GPU 1, ...
   │
   ├─ dispatch (all-to-all)   ← communication irrégulière
   ├─ calcul des experts locaux
   └─ combine (all-to-all inverse)

DeepEP (DeepSeek) est la bibliothèque de référence pour cette opération. Ses caractéristiques :

  • un chemin haut débit (HT) pour l'entraînement et le préremplissage ;
  • un chemin basse latence (LL) pour le décodage ;
  • une communication initiée depuis le noyau via NVSHMEM et IBGDA, ce qui permet de recouvrir le dispatch avec le calcul ;
  • un transport multi-étages propre, NVSHMEM n'étant utilisé qu'aux points critiques du chemin RDMA inter-nœuds : échange de métadonnées, écritures RDMA par blocs, et mises à jour atomiques de compteurs distants.

L'évolution de DeepEP

DeepEP V2 (l'upstream courant à la mi-2026) utilise par défaut un back-end NCCL appelé « Gin », et ne dépend plus de NVSHMEM que pour des chemins de code hérités.

C'est une information à connaître : beaucoup de documentation en ligne décrit encore l'architecture V1 fondée sur NVSHMEM.

6.6 Les optimisations de noyau spécifiques

Trois techniques qui reviennent dans tous les travaux récents :

1. Fusionner le gather dans la GEMM. Au lieu de matérialiser une copie contiguë des jetons par expert, la GEMM lit à travers la table d'indices. Gain mesuré par MPK : élimine un prétraitement représentant jusqu'à 11 % du temps du MoE.

2. Recouvrir le dispatch avec le calcul. Envoyer les jetons du premier groupe d'experts pendant qu'on calcule ceux du deuxième. Exige une communication initiée depuis le noyau (NVSHMEM ou équivalent).

3. Trier finement. Trier les jetons par expert améliore la localité pour la lecture des poids d'expert : un bloc qui traite des jetons du même expert lit les mêmes poids.

Des travaux récents comme SonicMoE (optimisations conscientes des E/S et des tuiles) explorent systématiquement cet espace.

6.7 L'arithmétique du MoE

Reprenons le raisonnement du roofline pour un MoE.

Notons \(\Theta_{\text{tot}}\) les paramètres totaux, \(\Theta_{\text{act}}\) les paramètres actifs par jeton, \(E\) le nombre d'experts, \(k\) les experts activés.

Décodage, lot 1 :

  • FLOP : \(2\Theta_{\text{act}}\)
  • octets : \(2\Theta_{\text{act}}\) (seuls les experts activés sont lus)
  • \(I \approx 1\)

Même intensité arithmétique qu'un dense, mais sur beaucoup moins d'octets. Le temps par jeton est divisé par \(\Theta_{\text{tot}}/\Theta_{\text{act}}\).

Décodage, lot \(b\) grand :

Si \(b \cdot k \gg E\), tous les experts sont sollicités :

  • FLOP : \(2\Theta_{\text{act}} b\)
  • octets : \(2\Theta_{\text{tot}}\) (tout est lu)
  • \(I \approx b \cdot \Theta_{\text{act}} / \Theta_{\text{tot}}\)

L'intensité arithmétique croît avec \(b\) mais est pénalisée par le ratio de parcimonie. Un MoE à 5 % d'activation a besoin d'un lot 20 fois plus grand qu'un dense pour atteindre le même régime.

La conséquence pratique

Le MoE est optimisé pour le petit lot. C'est exactement le régime de l'inférence latence-critique, et c'est pourquoi il domine les modèles récents.

C'est aussi pourquoi les MoE et les megakernels se rencontrent : les deux ciblent le même régime, et les travaux de la partie 8 (MPK, Fleet, Ada-MK) traitent tous des modèles MoE.

Résumé du chapitre

À retenir

  • Un MoE dissocie paramètres totaux et paramètres actifs. À lot 1, on ne lit que les experts activés : le temps est divisé par le ratio de parcimonie.
  • L'implémentation correcte est permutation + GEMM groupée, pas masquage.
  • La GEMM groupée s'implémente par un noyau persistant à file de tuiles — la même structure qu'un megakernel.
  • Le prétraitement (gather/scatter) coûte jusqu'à 11 % du temps du MoE ; le fusionner dans la GEMM l'élimine.
  • Le déséquilibre de charge se traite par ordonnancement dynamique, pas par capacité fixe (qui jette des jetons).
  • En multi-GPU, le dispatch est un all-to-all irrégulier ; DeepEP est la référence, avec un chemin haut débit et un chemin basse latence. DeepEP V2 utilise NCCL par défaut.
  • cuBLAS supporte la GEMM groupée depuis CUDA 13.1.

Vérifiez que vous avez compris

Un MoE de 400 milliards de paramètres dont 30 sont actifs. Combien de jetons/s en décodage à lot 1 sur un nœud 8×H100 ?
  • Poids actifs en BF16 : \(2 \times 30\times10^9 = 60\) Go ;
  • avec un parallélisme sur 8 GPU : 7,5 Go par carte, mais chaque carte lit ses propres experts actifs, donc le volume pertinent est ~7,5 Go par carte, lus en parallèle ;
  • \(t_{\min} \approx 7{,}5 / 3\,350 = 2{,}2\) ms.

Soit ~450 jetons/s en théorie. En pratique, on observe bien moins à cause : du cache KV, de la communication all-to-all (deux fois par couche), du surcoût de lancement des noyaux, et du déséquilibre entre experts.

C'est justement l'écart entre ces 2,2 ms théoriques et le temps réel que les megakernels attaquent.

Pourquoi le masquage est-il si mauvais, alors qu'il est parfaitement régulier ?

Parce qu'il fait \(E/k\) fois trop de travail.

Avec \(E = 64\) experts et \(k = 8\) activés, chaque jeton passe par les 64 experts et on jette 56 résultats sur 64. On effectue 8 fois le calcul nécessaire, et on lit 8 fois trop de poids.

Le MoE perd alors tout son intérêt : on paie le calcul d'un dense de \(\Theta_{\text{tot}}\) paramètres tout en n'ayant la qualité que de \(k\) experts.

La régularité ne compense jamais un facteur 8 de travail inutile — c'est le point 1.5 de la checklist : ne pas calculer ce qui sera jeté.

Pourquoi le all-to-all du MoE est-il plus difficile qu'un all-reduce classique ?

Trois raisons.

  1. Les tailles sont irrégulières et dynamiques. Un all-reduce échange toujours le même volume ; le dispatch MoE échange un volume différent entre chaque paire de GPU, différent à chaque lot. Impossible de pré-planifier.
  2. Le motif est all-to-all, pas en anneau ni en arbre. NCCL optimise très bien les seconds ; l'all-to-all sature les liens différemment.
  3. Il y en a deux par couche (dispatch et combine), donc \(2L\) par passe avant. Pour \(L = 60\), cela fait 120 opérations collectives par jeton.

D'où l'importance du recouvrement : une communication initiée depuis le noyau (NVSHMEM, IBGDA) permet d'envoyer les premiers jetons pendant que les suivants se calculent — impossible avec des collectives à granularité de noyau.

Chapitre suivant : 7 · Multi-GPU et collectives

Sources de ce chapitre