Aller au contenu

7 · AMD et le monokernel

Une approche structurellement différente, imposée par un matériel différent. Sans TMA, sans wgmma, sans clusters — mais avec 8 XCD et un LDS de 160 Ko.


7.1 Le travail de Kog

L'entreprise Kog a publié le récit de la construction d'un moteur d'inférence à noyau unique optimisé pour la latence sur GPU AMD MI300X.

Ils appellent leur objet un monokernel — un noyau GPU persistant exécutant l'intégralité de la passe de décodage : préremplissage, décodage et échantillonnage, sans interruption.

Le résultat annoncé

Plus de 3 000 jetons/s par requête sur un nœud unique à 8 GPU AMD MI300X, avec un modèle de 2 milliards de paramètres en FP16.

Pour situer, les points de comparaison qu'ils donnent :

Système Jetons/s par requête
Systèmes GPU typiques (2-8 B) 100 – 300
Cerebras WSE (GPT-OSS-120B) ~3 000
Kog sur 8× MI300X (2 B) > 3 000
Kog sur 8× H200 2 100

Comment lire ces chiffres

La comparaison avec Cerebras porte sur des modèles de tailles très différentes (2 B contre 120 B), ce qui la rend peu informative en soi. Ce que les auteurs veulent montrer est qu'un GPU peut atteindre un régime de latence jusque-là associé au matériel à l'échelle de la galette.

Ils ne fournissent pas de comparaison directe avec vLLM ou SGLang, en expliquant que ces systèmes optimisent le débit à grand lot, un régime d'optimisation fondamentalement différent du leur (lot 1). C'est une réserve honnête, et elle limite la comparabilité des chiffres.


7.2 Le diagnostic

Leur analyse de départ identifie quatre coûts : « surcoût de lancement de noyaux, synchronisations, déséquilibres de charge entre unités de calcul, et effets de queue ».

Deux chiffres précis en ressortent :

Mesure Valeur
Surcoût d'un lancement de noyau (HIP, MI300X) ~4,5 µs
Part de la synchronisation de grille dans le temps de génération d'un jeton ~35 %

Le second est le plus frappant. Sur AMD, la barrière inter-blocs est proportionnellement bien plus coûteuse que sur NVIDIA — conséquence directe de la topologie à 8 dies : une barrière globale doit faire communiquer des CU répartis sur 8 XCD, sans chemin direct entre eux.


7.3 La technique de la sentinelle

C'est la contribution la plus élégante du travail.

Le problème

Une synchronisation par atomiques exige :

  1. le producteur écrit les données ;
  2. le producteur émet une barrière mémoire ;
  3. le producteur incrémente un compteur atomique ;
  4. le consommateur sonde le compteur ;
  5. le consommateur lit les données.

Deux lectures côté consommateur (le compteur, puis les données), et des opérations atomiques coûteuses.

La solution

Initialiser les tampons de sortie à une valeur impossible — ici NaN. Le consommateur lit directement les données et attend qu'elles ne soient plus NaN.

// Producteur : écrit simplement
sortie[i] = valeur;          // remplace le NaN

// Consommateur : attend que la donnée soit valide
float v;
do {
    v = lecture_non_cachee(&sortie[i]);
} while (isnan(v));

Une seule lecture, aucun atomique.

Le gain mesuré

Approche Latence de synchronisation
Synchronisation de grille naïve (atomiques) 7,59 – 7,88 µs
Sentinelle NaN 0,80 – 0,93 µs

Soit une amélioration d'environ 9×.

Les conditions de validité

Cette technique n'est pas universelle :

  • il faut une valeur impossible dans le domaine de sortie. NaN convient pour des activations, pas pour des données arbitraires ;
  • il faut réinitialiser les tampons à chaque itération, ce qui coûte de la bande passante ;
  • la lecture doit contourner le cache L1 (ld.global.cv, volatile, ou l'équivalent HIP), faute de quoi le consommateur peut lire indéfiniment une valeur mise en cache ;
  • elle ne fournit pas de garantie d'ordre entre plusieurs écritures : si le producteur écrit deux valeurs et que seule la seconde est vérifiée, rien ne garantit la visibilité de la première sans barrière mémoire.

C'est une optimisation à manier avec précaution, et ils la documentent comme telle.


7.4 L'optimisation des chiplets

Le MI300X a 8 Accelerator Compute Dies, chacun avec son propre L2.

La réponse de Kog : dupliquer les tenseurs par die d'entrée/sortie pour éviter les pénalités de communication inter-chiplets.

C'est le même arbitrage que celui analysé au chapitre 6 : payer \(N\) fois l'espace mémoire pour transformer des accès HBM en accès L2 locaux, rentable dès que le tenseur est relu plus d'une dizaine de fois par die.


7.5 Le streaming continu des poids

Un mécanisme qu'on ne trouve pas sous cette forme dans les megakernels NVIDIA.

Les poids sont préchargés dans le LDS (l'équivalent AMD de la mémoire partagée) et dans les registres pendant les phases de calcul précédentes, en utilisant des hints non temporels — ce qui évite de polluer les caches avec des données lues une seule fois.

L'objectif : que le flux de poids depuis la HBM ne s'interrompe jamais.

Pourquoi cette approche plutôt que la spécialisation des warps

Sur Hopper, le motif rentable est producteur/consommateur : un warp dédié émet des copies TMA, les autres calculent. Le gain vient de ce qu'un seul thread suffit à alimenter tout le bloc.

Sans TMA, ce gain disparaît : le producteur doit calculer toutes ses adresses à la main, ce qui consomme registres et instructions.

D'où un autre motif : plutôt que de dédier des warps, on entrelace les chargements dans le code de calcul de manière à ce qu'il y ait toujours des requêtes en vol. C'est du pipelining logiciel classique, appliqué avec rigueur.

Le LDS de 160 Ko par CU sur CDNA 4 (contre 64 Ko sur CDNA 3) rend cette approche nettement plus praticable.


7.6 Le parallélisme de tenseurs différé

Kog introduit une variante de transformeur qu'ils appellent DTP (Delayed Tensor Parallelism) : elle retarde les opérations de réduction entre GPU pour cacher la communication derrière le calcul, réduisant le surcoût de synchronisation à presque zéro.

C'est le même principe que la « transposition distribuée » de Hazy Research : quand on contrôle la communication au niveau du noyau, on peut restructurer l'algorithme pour qu'elle tombe à un moment où elle est masquable, au lieu de subir le placement imposé par les collectives standard.


7.7 Le tableau des différences NVIDIA / AMD

Aspect NVIDIA (Hopper/Blackwell) AMD (CDNA 3/4)
Chargement asynchrone TMA (mono-thread, multi-D) manuel, buffer_load
Instruction matricielle wgmma, tcgen05 (asynchrones) MFMA (synchrone)
Coopération multi-SM clusters + DSMEM aucune
Redistribution de registres setmaxnreg aucune
Mémoire rapide par unité 227 Ko (partagée) 160 Ko (LDS, CDNA 4)
Dies 1 (H100), 2 (B200) 8 XCD
Coût d'une barrière de grille µs 7,6-7,9 µs (naïve)
Motif de pipelining dominant spécialisation des warps streaming continu

La conclusion pratique : un megakernel NVIDIA ne se porte pas sur AMD. Il se reconçoit.


7.8 Ce que l'approche AMD apporte au domaine

Trois idées transférables, indépendamment du matériel.

1. La synchronisation par sentinelle. L'idée d'utiliser la donnée elle-même comme signal de disponibilité est générale et sous-utilisée. Elle vaut partout où un domaine de valeurs impossibles existe.

2. La duplication consciente des dies. Elle devient pertinente sur B200 également, et le sera davantage à mesure que le nombre de chiplets augmente.

3. La restructuration algorithmique pour la communication. DTP et la transposition distribuée sont deux instances de la même idée : le contrôle au niveau du noyau permet de déplacer la communication dans l'algorithme, pas seulement de l'optimiser.


Résumé du chapitre

À retenir

  • Kog a construit un monokernel exécutant préremplissage, décodage et échantillonnage en un seul noyau persistant sur MI300X.
  • Résultat annoncé : > 3 000 jetons/s par requête sur 8 MI300X avec un modèle de 2 B en FP16. Pas de comparaison directe avec vLLM/SGLang, les régimes d'optimisation étant différents.
  • Diagnostic : lancement à ~4,5 µs, et la synchronisation de grille représentait ~35 % du temps de génération d'un jeton.
  • La sentinelle NaN remplace les atomiques : synchronisation de 7,59-7,88 µs → 0,80-0,93 µs, soit ~9×. Conditions : valeur impossible disponible, réinitialisation, lecture non cachée, pas de garantie d'ordre.
  • Duplication des tenseurs par XCD pour éviter le trafic inter-chiplets.
  • Streaming continu des poids vers le LDS avec hints non temporels, à la place de la spécialisation des warps — faute de TMA.
  • DTP : retarder les réductions inter-GPU pour les masquer derrière le calcul.
  • Un megakernel NVIDIA ne se porte pas sur AMD, il se reconçoit.

Vérifiez que vous avez compris

Pourquoi la synchronisation de grille coûte-t-elle 35 % du temps sur MI300X, alors qu'elle est bien moins coûteuse sur H100 ?

À cause de la topologie à 8 dies.

Une barrière globale exige que tous les blocs communiquent via un point de cohérence commun. Sur H100, un seul die et un L2 unifié : la communication est interne à la puce.

Sur MI300X, les 8 XCD ont chacun leur L2. Une barrière doit donc traverser l'interconnexion inter-die, avec une latence et une contention bien supérieures. Le nombre de participants est aussi plus élevé (304 CU contre 132 SM).

Résultat : 7,6-7,9 µs par barrière. Avec plusieurs barrières par couche et des dizaines de couches, on atteint facilement 35 % d'un pas de décodage qui devrait durer quelques millisecondes.

C'est aussi ce qui rend le gain de la sentinelle si important sur cette plateforme.

La sentinelle NaN ne fournit pas de garantie d'ordre. Pourquoi est-ce dangereux ?

Parce que le consommateur peut voir la donnée « signal » sans voir les données qui l'ont précédée.

Scénario : le producteur écrit a[0] = x puis a[1] = y. Le consommateur attend que a[1] ne soit plus NaN, puis lit a[0].

Rien ne garantit que l'écriture de a[0] soit visible : le matériel peut réordonner les écritures, et les caches ne sont pas cohérents entre CU.

La règle : soit on vérifie chaque valeur individuellement (ce qui est le cas si le tampon entier est initialisé à NaN et que chaque consommateur n'attend que les valeurs qu'il lit), soit on ajoute une barrière mémoire avant l'écriture du signal — ce qui réintroduit une partie du coût qu'on voulait éviter.

La première option est celle qui fonctionne ici, et elle explique pourquoi la technique convient bien à des tampons d'activation où chaque consommateur lit exactement ce qu'il attend.

Sans TMA ni wgmma, comment AMD peut-il être compétitif ?

En jouant sur d'autres forces.

  1. La bande passante mémoire : 5,3 To/s sur MI300X, 8,0 To/s sur MI355X — comparable ou supérieur à Hopper. Sur un problème limité par la mémoire, c'est la ressource qui compte.
  2. La capacité : 192 Go (MI300X) et 288 Go (MI355X) contre 80 Go sur H100. Moins de GPU nécessaires, donc moins de communication.
  3. Le LDS de 160 Ko par CU sur CDNA 4, ce qui autorise de grandes tuiles et un pipelining profond en logiciel.
  4. Le FP64 (81,7 TFLOPS), sans rapport avec l'inférence mais décisif en HPC.

Ce qui manque est le confort de programmation : sans TMA ni instructions asynchrones, il faut plus de travail pour atteindre le même pourcentage du pic. C'est un coût d'ingénierie, pas une limite physique — et c'est exactement pourquoi les travaux comme Kog et Fleet ont de la valeur : ils montrent que le pic est atteignable, par d'autres chemins.


Chapitre suivant : 8 · Écrire son megakernel


Sources de ce chapitre