Aller au contenu

5 · La spécialisation des warps

Tous les warps d'un bloc exécutent le même code depuis 2006. La programmation Hopper casse cette symétrie : certains warps ne font que déplacer des données, d'autres ne font que calculer. C'est le motif architectural central des noyaux modernes — et celui des megakernels.

5.1 Le modèle homogène et sa limite

Dans un noyau classique, chaque warp fait la même chose :

__global__ void homogene() {
    for (int t = 0; t < N; ++t) {
        charger(t);          // tous les warps chargent
        __syncthreads();
        calculer(t);         // tous les warps calculent
        __syncthreads();
    }
}

Deux problèmes.

1. Les ressources sont sous-utilisées alternativement. Pendant la phase de chargement, les tensor cores sont inactifs. Pendant le calcul, les unités load/store le sont.

2. Les registres sont dimensionnés pour le pire cas. Le compilateur alloue assez de registres pour couvrir les deux phases dans chaque warp, alors qu'aucun warp n'a besoin des deux simultanément.

5.2 Le modèle spécialisé

__global__ void specialise() {
    int wg = threadIdx.x / 128;              // identifiant du warp group

    if (wg == 0) {
        // ── PRODUCTEUR ──
        // Registres : peu (juste des adresses)
        for (int t = 0; t < N; ++t) {
            vide[t % E].wait();
            emettre_tma(t, plein[t % E]);
        }
    } else {
        // ── CONSOMMATEUR ──
        // Registres : beaucoup (accumulateurs)
        for (int t = 0; t < N; ++t) {
            plein[t % E].wait();
            wgmma(acc, tampon[t % E]);
            vide[t % E].arrive();
        }
    }
}

Les deux branches ne se rejoignent jamais. C'est de la divergence assumée à l'échelle du warp, ce qui ne coûte rien : la divergence n'est chère qu'à l'intérieur d'un warp.

Le gain sur les registres

CUDA offre un mécanisme dédié pour redistribuer les registres entre warp groups :

setmaxnreg.dec.sync.aligned.u32 24;     // producteur : rendre des registres
setmaxnreg.inc.sync.aligned.u32 240;    // consommateur : en prendre

En CUDA C++ (Hopper, sm_90a) :

if (wg == 0) {
    asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" :: "n"(24));
    // producteur
} else {
    asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" :: "n"(240));
    // consommateur
}

Un producteur qui n'a besoin que de 24 registres en rend 200 aux consommateurs. Sur un bloc de 3 warp groups (1 producteur + 2 consommateurs), cela permet des accumulateurs beaucoup plus grands, donc des tuiles plus grandes, donc une intensité arithmétique supérieure.

Pourquoi c'est le motif gagnant sur Hopper

Il aligne trois choses :

  1. TMA est mono-thread : un seul warp suffit à alimenter tout le bloc ;
  2. wgmma est asynchrone : le consommateur peut lancer plusieurs MMA puis attendre ;
  3. mbarrier compte des octets : la synchronisation producteur → consommateur est exacte.

Les trois mécanismes de Hopper ont été conçus pour ce motif.

5.3 Les deux schémas canoniques

Ping-pong

Deux warp groups consommateurs travaillent sur des tuiles de sortie différentes, décalés d'une demi-période.

temps →
WG1 : [MMA tuile A][épilogue A][MMA tuile C][épilogue C]
WG2 :              [MMA tuile B][épilogue B][MMA tuile D]
TC  : ■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■  ← jamais inactifs

Adapté quand l'épilogue est coûteux (activation, quantification, écriture avec transformation).

Cooperative

Les deux warp groups collaborent sur la même tuile de sortie, chacun traitant la moitié des lignes.

WG1 : [MMA lignes 0-63  ][épilogue 0-63 ]
WG2 : [MMA lignes 64-127][épilogue 64-127]

Adapté aux grandes tuiles où la mémoire partagée ne permettrait pas deux tuiles simultanées.

CUTLASS nomme ces variantes sm90_gemm_tma_warpspecialized_pingpong et ..._cooperative. Le choix se fait par mesure ; il n'y a pas de règle absolue.

5.4 La forme générale : quatre rôles

Dans les noyaux les plus sophistiqués (megakernels compris), on trouve jusqu'à quatre rôles.

Rôle Ce qu'il fait Ressource utilisée
Loader émet les copies TMA global → partagé unités TMA
Compute exécute les MMA tensor cores
Storer écrit les résultats, gère la communication inter-GPU unités load/store, NVLink
Scheduler distribue le travail, gère les dépendances atomiques, mémoire

Le megakernel de débit de Hazy Research utilise exactement cette décomposition : des threads loader, compute et storer spécialisés au sein de chaque SM, « permettant le pipelining d'instructions — charger les poids de l'opération suivante pendant qu'on termine la courante ». Gain mesuré : 2 à 6 %.

Le gain paraît modeste, mais il s'ajoute à deux autres niveaux de recouvrement :

Niveau Mécanisme Gain mesuré
Dans un SM threads loader/compute/storer 2-6 %
Entre SM file de travail globale, entrelacement d'instructions d'opérations différentes 14,2 % (lot 8 192, vs tourniquet)
Entre GPU threads storer dédiés à la communication asynchrone inclus dans les 22 % totaux

Résultat final : 23 468 jetons/s contre 19 170 pour SGLang sur Llama-70B en tensor-parallèle sur 8 H100, soit +22 %.

5.5 Le lien avec les megakernels

C'est ici que la partie 4 rejoint la partie 8.

Un megakernel est la spécialisation des warps portée à l'échelle du GPU entier :

Concept Dans une GEMM Dans un megakernel
Unité de travail une tuile une « instruction » / « tâche »
Producteur 1 warp group émettant TMA threads loader de chaque SM
Consommateur 2 warp groups en ping-pong threads compute
Tampon mémoire partagée circulaire, E étages pages de mémoire partagée (16 ou 32 Ko)
Synchronisation mbarrier locale au bloc compteurs en mémoire globale
Ordonnancement statique (blockIdx → tuile) file de travail dynamique + SM ordonnanceurs

Mirage MPK pousse la logique jusqu'à dédier des SM entiers au rôle d'ordonnanceur : 16 SM ordonnanceurs et 104/128/144 SM workers selon la carte (A100/H100/B200).

L'idée à emporter de ce chapitre

La spécialisation des warps n'est pas une astuce d'optimisation. C'est un changement de modèle de programmation : on passe d'un modèle SPMD (tout le monde fait la même chose) à un modèle de flux de données (chacun son rôle, communication par tampons).

Les megakernels sont ce modèle appliqué à un réseau de neurones entier.

5.6 Le support dans les outils

Écrire cela à la main est pénible. L'état du support en 2026 :

Outil Support de la spécialisation
CUTLASS C++ complet, c'est son architecture native (CollectiveMma)
CuTe DSL complet, en Python
ThunderKittens oui, via des templates de « workers »
Triton en cours — tl.async_task et des passes automatiques ; voir la feuille de route PyTorch
Gluon oui, c'est un de ses arguments : « expose layouts, shared memory, warp specialization »
Mojo oui, via les structured kernels
HIP / AMD manuel ; pas d'équivalent asynchrone à wgmma

Le fait que Triton ait dû ajouter la spécialisation des warps — et qu'OpenAI ait créé Gluon pour l'exposer explicitement — est le meilleur indicateur de son importance : le modèle de tuiles automatique de Triton ne suffisait plus pour atteindre le pic sur Hopper.

Résumé du chapitre

À retenir

  • La spécialisation des warps donne des rôles différents aux warps d'un même bloc : producteur (TMA), consommateur (MMA), storer, ordonnanceur.
  • La divergence entre warps est gratuite ; seule celle à l'intérieur d'un warp coûte.
  • setmaxnreg redistribue les registres : le producteur en rend, le consommateur en prend. C'est ce qui permet de grandes tuiles.
  • Deux schémas : ping-pong (tuiles différentes, bon si l'épilogue est coûteux) et cooperative (même tuile, bon si elle est grande).
  • Le recouvrement se fait à trois niveaux : dans un SM (2-6 %), entre SM (14,2 %), entre GPU. Cumulés : +22 % sur SGLang.
  • Un megakernel est ce modèle appliqué au GPU entier, avec des compteurs globaux à la place des mbarrier locales.

Vérifiez que vous avez compris

Pourquoi la divergence entre warp groups est-elle gratuite alors que la divergence intra-warp coûte cher ?

Parce que l'unité d'ordonnancement est le warp, pas le bloc. Chaque warp a son propre flux d'instructions et est planifié indépendamment par l'ordonnanceur de sa partition.

Quand deux warps exécutent des branches différentes, ils occupent simplement des créneaux d'exécution différents — c'est le fonctionnement normal. Quand deux threads d'un même warp divergent, le matériel doit exécuter les deux chemins successivement avec des masques, puisqu'il n'y a qu'un seul flux d'instructions par warp.

Pourquoi le gain de la spécialisation intra-SM (2-6 %) est-il si faible comparé au gain inter-SM (14,2 %) ?

Parce qu'ils corrigent des problèmes d'échelles différentes.

La spécialisation intra-SM récupère les bulles courtes : les quelques dizaines de cycles où un SM attend une donnée. Avec un pipeline déjà profond, ces bulles sont rares.

L'ordonnancement inter-SM corrige le déséquilibre de charge : certains SM finissent leur travail bien avant d'autres et restent inactifs. À lot 8 192, avec des opérations hétérogènes (GEMM compute-bound, attention memory-bound), ce déséquilibre est structurel et coûte beaucoup plus.

C'est cohérent avec l'analyse générale : les frontières et les déséquilibres coûtent plus que les micro-inefficacités. C'est le raisonnement même des megakernels.

Vous portez un noyau warp-specialized de H100 vers MI300X. Quelles difficultés ?

Quatre, au minimum :

  1. Pas de TMA. Le producteur doit calculer ses adresses à la main, ce qui consomme des registres et des instructions — le gain principal de la spécialisation disparaît en partie.
  2. MFMA est synchrone. Le consommateur ne peut pas lancer plusieurs MMA puis attendre ; il faut structurer le recouvrement différemment.
  3. Pas de setmaxnreg. L'allocation de registres est uniforme sur le bloc.
  4. Wavefront de 64. Toute la granularité change.

En pratique, les noyaux AMD performants utilisent un autre motif : le streaming continu des poids vers le LDS avec des hints non-temporels, tel que décrit par Kog. Le portage n'est pas une traduction, c'est une reconception.

Chapitre suivant : 6 · CUDA Tile

Sources de ce chapitre