4 · wgmma et tcgen05¶

Les deux instructions matricielles qui définissent la performance des GPU modernes. Ce chapitre ne vous apprendra pas à les écrire à la main — presque personne ne le fait — mais à comprendre ce que fait le code qui les utilise.

4.1 D'où l'on vient¶

Rappel de l'évolution (Fondations 5) :

Instruction	Portée	Opérandes	Synchronisme
`mma.sync.m16n8k16`	warp (32 threads)	registres	synchrone
`wgmma.mma_async`	warp group (128 threads)	mém. partagée + registres	asynchrone
`tcgen05.mma`	1 ou 2 SM	mém. partagée + TMEM	asynchrone

Chaque étape déplace les opérandes plus loin des registres et rend l'opération plus asynchrone. Ce n'est pas un hasard : les registres sont la ressource la plus rare du SM.

4.2 `wgmma` : le produit matriciel du warp group¶

Ce qu'est un warp group¶

Quatre warps consécutifs, soit 128 threads, alignés sur threadIdx.x / 128. C'est l'unité d'exécution de wgmma.

La forme de l'instruction¶

wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16
    {%f0, %f1, ..., %f127},     // accumulateur, dans les registres
    %desc_a,                     // descripteur mémoire partagée pour A
    %desc_b,                     // descripteur mémoire partagée pour B
    1,                           // scale-D : 0 = écraser, 1 = accumuler
    1, 1,                        // scale-A, scale-B (±1)
    0, 0;                        // transposition de A, de B

Les tailles supportées sur Hopper : \(m = 64\) toujours, \(n \in \{8, 16, \dots, 256\}\), \(k = 16\) pour le BF16/FP16, \(k = 32\) pour le FP8.

Un wgmma.m64n256k16 effectue \(64 \times 256 \times 16 \times 2 = 524\,288\) opérations flottantes en une instruction.

Le descripteur de matrice¶

C'est un entier 64 bits qui encode l'adresse en mémoire partagée et sa disposition :

Bits	Champ
0-13	adresse de départ en mémoire partagée (÷16)
16-29	leading byte offset
32-45	stride byte offset
49-51	mode de swizzle

Le matériel lit directement les opérandes depuis la mémoire partagée en suivant ce descripteur. Aucun ldmatrix, aucun registre consommé pour \(\mathbf{A}\) et \(\mathbf{B}\).

L'asynchronisme¶

wgmma.fence.sync.aligned;              // les registres d'accumulation sont prêts
wgmma.mma_async... ;                    // lancer (plusieurs fois)
wgmma.mma_async... ;
wgmma.commit_group.sync.aligned;        // clore le groupe
wgmma.wait_group.sync.aligned 0;        // attendre

Entre le lancement et l'attente, le warp group peut faire autre chose : émettre des copies TMA, calculer un softmax, préparer l'épilogue. C'est ce qui permet le recouvrement calcul/mémoire à l'intérieur d'un warp.

Pourquoi wgmma atteint 95 % du pic

Trois raisons cumulées :

Les opérandes ne passent pas par les registres : plus de ldmatrix, donc plus d'instructions de mouvement dans la boucle chaude.
Les tuiles sont grandes (\(n\) jusqu'à 256), ce qui amortit le coût fixe de l'instruction.
L'asynchronisme permet de superposer la MMA aux copies TMA.

Mesure : 95 % du pic théorique contre 62,9 % pour la mma classique rétro-compatible (arXiv:2501.12084).

4.3 `tcgen05` : la génération Blackwell¶

Les trois ruptures¶

1. La Tensor Memory (TMEM).

256 Ko par SM, structurés en 512 colonnes × 128 lignes de cellules 32 bits, avec 16 To/s en lecture et 8 To/s en écriture. C'est une mémoire distincte des registres et de la mémoire partagée.

L'accumulateur d'une MMA y réside désormais, ce qui libère complètement le banc de registres.

Instructions dédiées :

tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%addr], 128;
tcgen05.ld.sync.aligned.32x32b.x64.b32 {%r0, ...}, [%tmem];
tcgen05.st.sync.aligned.32x32b.x64.b32 [%tmem], {%r0, ...};
tcgen05.cp.cta_group::1.128x256b [%tmem], %desc_smem;
tcgen05.dealloc.cta_group::1.sync.aligned.b32 [%addr], 128;

2. La paire de CTA.

Deux blocs d'un même cluster peuvent exécuter une seule MMA collective, portant la tuile de \(128\times128\) à \(256\times128\) (voire \(256\times256\times16\) selon la configuration).

tcgen05.mma.cta_group::2.kind::f16 [%tmem_d], %desc_a, %desc_b, %idesc, 1;

3. L'émission par un thread unique.

Contrairement à wgmma qui exige les 128 threads du warp group, tcgen05.mma est émise par un seul thread. C'est cohérent avec TMA : le mouvement de données et le calcul matriciel sont tous deux devenus des opérations « déclenchées puis attendues ».

La latence constante¶

Le fait le plus important pour la conception de noyaux :

Instruction	Tuile	Latence (cycles)
`wgmma`	m64n64k16	32,0
`wgmma`	m64n128k16	64,0
`wgmma`	m64n256k16	128,0
`tcgen05.mma`	m64n64k16	11,0
`tcgen05.mma`	m128n128k16	11,3
`tcgen05.mma`	m256n256k16	11,4

Source : arXiv:2512.02189.

Sur Hopper, doubler la tuile doublait la latence. Sur Blackwell, elle est constante. Conséquence directe : on prend toujours la plus grande tuile qui tient, et le goulot se déplace vers l'alimentation en données.

C'est exactement ce que FlashAttention-4 exploite et ce qui l'oblige à traiter séparément le problème des exponentielles, désormais limitantes.

4.4 L'anatomie d'un noyau GEMM Hopper moderne¶

Assemblons tout ce qui précède. Le squelette conceptuel d'une GEMM CUTLASS sm90_gemm_tma_warpspecialized_pingpong :

┌──────────────────────────────────────────────────────────────┐
│ Bloc de 384 threads = 3 warp groups                          │
├──────────────────────────────────────────────────────────────┤
│ Warp group 0 : PRODUCTEUR                                    │
│   boucle :                                                   │
│     vide[e].wait()                                           │
│     un thread émet TMA(A) et TMA(B) vers tampon[e]           │
│     mbarrier expect_tx                                       │
├──────────────────────────────────────────────────────────────┤
│ Warp groups 1 et 2 : CONSOMMATEURS (ping-pong)               │
│   boucle :                                                   │
│     plein[e].wait()                                          │
│     wgmma.fence                                              │
│     wgmma.mma_async(acc, desc_a[e], desc_b[e])  × K/16       │
│     wgmma.commit_group ; wgmma.wait_group                    │
│     vide[e].arrive()                                         │
│   épilogue : écrire acc dans C via TMA store                 │
└──────────────────────────────────────────────────────────────┘

Le ping-pong : les deux warp groups consommateurs travaillent sur des tuiles de sortie différentes, décalés d'une demi-période. Pendant que l'un fait son épilogue (écriture de \(\mathbf{C}\), opérations élémentaires), l'autre fait ses MMA. Les tensor cores ne s'arrêtent jamais.

Il existe une variante cooperative où les deux warp groups collaborent sur la même tuile de sortie. Elle convient mieux aux grandes tuiles, le ping-pong convient mieux aux épilogues coûteux.

4.5 Ce qu'il faut vraiment en retenir¶

N'écrivez pas ces instructions à la main

Le descripteur wgmma est un champ de bits de 64 bits. La disposition des accumulateurs en registres est une table de correspondance de plusieurs pages. L'ordre des fence, commit_group et wait_group est subtil. Une erreur donne un résultat faux, sans message.

Utilisez :

Outil	Niveau
cuBLAS, cuDNN	ne rien écrire
CUTLASS C++ (`CollectiveMma`)	assembler des composants
CuTe DSL (Python)	contrôle fin, syntaxe Python
ThunderKittens	tuiles 16×16, très lisible
Triton / Gluon	génération automatique

FlashAttention-4, l'état de l'art de l'attention sur Blackwell, est écrit entièrement en CuTe DSL, et ses auteurs notent que la compilation prend désormais des secondes au lieu de minutes ou d'heures.

Ce qu'il faut savoir, en revanche :

Reconnaître ces instructions dans le SASS (QGMMA, UTCHMMA, HMMA) pour vérifier que votre noyau utilise bien les tensor cores.
Comprendre pourquoi les opérandes vont en mémoire partagée : pour libérer les registres.
Comprendre l'asynchronisme : c'est ce qui permet la spécialisation des warps.
Savoir que la latence est constante sur Blackwell : cela change le dimensionnement des tuiles.

4.6 Le côté AMD¶

Pas d'équivalent asynchrone. MFMA sur CDNA reste :

synchrone (le wavefront attend le résultat) ;
avec des opérandes en registres vectoriels (VGPR) ;
à l'échelle du wavefront de 64.

// CDNA3, tuile 16×16×16 en FP16, accumulation FP32
using f32x4 = __attribute__((__vector_size__(4 * sizeof(float)))) float;
f32x4 d = __builtin_amdgcn_mfma_f32_16x16x16f16(a, b, c, 0, 0, 0);

CDNA 4 (gfx950, MI350X/MI355X) double le débit pour les types ≤ 16 bits et ajoute des instructions à échelle par blocs d'exposant (MXFP8/6/4), mais conserve le modèle synchrone.

Conséquence : le pipelining sur AMD repose davantage sur l'ordonnancement des instructions et le préchargement dans le LDS que sur des mécanismes asynchrones matériels. C'est une des raisons pour lesquelles les megakernels AMD (Kog) mettent l'accent sur le streaming continu des poids plutôt que sur des barrières asynchrones.

Résumé du chapitre¶

À retenir

wgmma (Hopper) : MMA asynchrone à l'échelle du warp group (128 threads), opérandes lus directement en mémoire partagée via un descripteur 64 bits. 95 % du pic contre 62,9 % pour mma.
tcgen05 (Blackwell) : accumulateur en Tensor Memory (256 Ko/SM, 16 To/s), coopération sur une paire de CTA (tuiles jusqu'à 256×256×16), émission par un seul thread.
Latence Blackwell constante à ~11 cycles quelle que soit la tuile — contre 32 à 128 cycles sur Hopper. On prend donc toujours la plus grande tuile possible.
Le noyau GEMM moderne : 1 warp group producteur (TMA), 2 consommateurs (wgmma) en ping-pong, tampon circulaire avec mbarrier.
N'écrivez pas ces instructions à la main. Utilisez CUTLASS, CuTe DSL, ThunderKittens ou Triton.
AMD : MFMA reste synchrone, opérandes en registres, wavefront de 64.

Vérifiez que vous avez compris¶

Pourquoi le passage des opérandes des registres vers la mémoire partagée est-il un gain, alors que la mémoire partagée est plus lente que les registres ?

Parce que la ressource critique n'est pas la vitesse, c'est la capacité.

Un SM a 256 Ko de registres pour jusqu'à 2 048 threads, soit 128 octets par thread. Une tuile \(\mathbf{A}\) de \(64\times16\) en BF16 fait 2 Ko ; la charger en registres consomme 16 octets par thread pour un warp group, à quoi s'ajoutent les instructions ldmatrix.

En laissant les opérandes en mémoire partagée, ces registres deviennent disponibles pour l'accumulateur, qui est ce qui limite réellement la taille de tuile. Et le tensor core lit la mémoire partagée à un débit suffisant pour ne pas être le goulot.

Blackwell pousse la logique plus loin en sortant aussi l'accumulateur des registres, vers la TMEM.

Pourquoi le schéma ping-pong utilise-t-il deux warp groups consommateurs plutôt qu'un seul plus large ?

Parce qu'une GEMM alterne entre deux phases aux besoins matériels différents : les MMA (tensor cores) et l'épilogue (unités vectorielles, écritures mémoire, opérations élémentaires comme un biais ou une activation).

Avec un seul consommateur, les tensor cores sont inactifs pendant tout l'épilogue. Avec deux consommateurs décalés, l'épilogue de l'un recouvre les MMA de l'autre : les tensor cores ne s'arrêtent jamais.

C'est le même principe que le pipelining producteur/consommateur, appliqué cette fois entre deux ressources de calcul plutôt qu'entre mémoire et calcul.

Sur Blackwell, la latence de MMA est constante. Quel nouveau goulot cela crée-t-il ?

L'alimentation en données. Si l'instruction MMA prend toujours 11 cycles quelle que soit la tuile, mais qu'une tuile 256×256 exige 16× plus d'octets qu'une tuile 64×64, le rapport calcul/mémoire de l'instruction se dégrade — ou plutôt, il devient impératif que TMA suive.

Et pour des noyaux comme l'attention, un autre goulot apparaît : les unités de fonctions spéciales. Le débit tensor core a doublé de Hopper à Blackwell, celui de la SFU non. FlashAttention-4 répond en calculant exp() par approximation polynomiale sur les unités FMA, dont Blackwell a de l'abondance, et en réduisant d'environ 10× le nombre de remises à l'échelle du softmax en ligne.

Chapitre suivant : 5 · La spécialisation des warps