Aller au contenu

8 · Écrire son megakernel

Le guide pratique. Structure, mécanismes, ordre de construction, et les pièges que tout le monde rencontre.

8.1 Avant de commencer

Trois questions à trancher d'abord

1. Le gain est-il là ? Calculez la part du surcoût :

\[\frac{t_{\text{mesuré}} - t_{\min}}{t_{\text{mesuré}}}\]

où \(t_{\min} = \text{octets lus} / B\). Si c'est en dessous de 20 %, arrêtez.

2. Avez-vous épuisé les alternatives ? CUDA Graphs, fusion, PDL, quantification. Toutes sont infiniment plus simples.

3. Existe-t-il un compilateur pour votre cas ? Essayez Mirage MPK avant d'écrire quoi que ce soit. Il compile des modèles Hugging Face « en quelques dizaines de lignes de Python ».

Si vous écrivez quand même, c'est probablement que votre architecture est inhabituelle, ou que vous faites de la recherche. Les deux sont de bonnes raisons.

8.2 La structure générale

// ═══════════════ Types ═══════════════

struct Instruction {
    int  opcode;
    int  compteur_entree;      // à attendre
    int  valeur_attendue;
    int  compteur_sortie;      // à incrémenter
    int  args[16];             // indices de couche, offsets, dimensions
};

struct GlobalState {
    Instruction* instructions;   // liste plate, par SM ou file globale
    int*         debut_par_sm;   // offsets
    unsigned*    compteurs;      // dépendances
    float*       activations;    // tampons intermédiaires
    // pointeurs vers les poids, le cache KV, etc.
};

// ═══════════════ Le noyau ═══════════════

__global__ __launch_bounds__(NB_THREADS, 1)
void megakernel(GlobalState g) {
    extern __shared__ char smem[];
    PageAllocator pages(smem, NB_PAGES, TAILLE_PAGE);

    const int sm  = blockIdx.x;
    const int tid = threadIdx.x;
    const int role = classer_role(tid);     // loader / compute / storer

    int pc = g.debut_par_sm[sm];            // compteur ordinal du SM

    while (true) {
        Instruction ins = g.instructions[pc];
        if (ins.opcode == OP_FIN) break;

        // ── Attendre les dépendances ──
        if (tid == 0 && ins.compteur_entree >= 0) {
            attendre(g.compteurs, ins.compteur_entree, ins.valeur_attendue);
        }
        __syncthreads();

        // ── Exécuter ──
        switch (ins.opcode) {
            case OP_RMSNORM_QKV_ROPE: executer<RmsNormQkvRope>(g, ins, pages, role); break;
            case OP_ATTENTION:        executer<Attention>     (g, ins, pages, role); break;
            case OP_ATTN_REDUCE:      executer<AttnReduce>    (g, ins, pages, role); break;
            case OP_O_PROJ:           executer<OProj>         (g, ins, pages, role); break;
            case OP_MLP_UP_GATE:      executer<MlpUpGate>     (g, ins, pages, role); break;
            case OP_MLP_DOWN:         executer<MlpDown>       (g, ins, pages, role); break;
            case OP_LM_HEAD:          executer<LmHead>        (g, ins, pages, role); break;
        }

        // ── Signaler ──
        __threadfence();
        __syncthreads();
        if (tid == 0 && ins.compteur_sortie >= 0) {
            atomicAdd(&g.compteurs[ins.compteur_sortie], 1u);
        }

        ++pc;
    }
}

Ce code n'a pas été compilé

Il illustre la structure, pas une implémentation fonctionnelle. L'environnement de rédaction n'a ni GPU ni nvcc. Pour du code réel, voyez HazyResearch/Megakernels et mirage-project/mirage.

8.3 Le lancement persistant

void lancer(GlobalState& g, size_t smem_octets) {
    // 1. S'assurer d'avoir accès à toute la mémoire partagée
    CUDA_CHECK(cudaFuncSetAttribute(
        megakernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_octets));

    // 2. Combien de blocs peuvent résider par SM ?
    int blocs_par_sm = 0;
    CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &blocs_par_sm, megakernel, NB_THREADS, smem_octets));
    assert(blocs_par_sm >= 1 && "le noyau ne tient pas sur un SM");

    // 3. Combien de SM ?
    int nb_sm = 0;
    CUDA_CHECK(cudaDeviceGetAttribute(&nb_sm, cudaDevAttrMultiProcessorCount, 0));

    // 4. Lancer EXACTEMENT une vague
    int blocs = nb_sm * blocs_par_sm;
    megakernel<<<blocs, NB_THREADS, smem_octets>>>(g);
    CUDA_CHECK(cudaGetLastError());
}

L'erreur numéro un

Lancer plus de blocs que le GPU ne peut en héberger. Résultat : interblocage silencieux. Le programme se fige, sans message.

L'assertion et le calcul par cudaOccupancyMaxActiveBlocksPerMultiprocessor sont obligatoires, pas optionnels.

8.4 L'allocateur de pages

class PageAllocator {
    char*     base_;
    unsigned* libres_;        // masque de bits en mémoire partagée
    int       nb_pages_;
    int       taille_page_;

public:
    __device__ PageAllocator(char* smem, int n, int taille);

    // Un seul thread appelle ; le résultat est diffusé
    __device__ int allouer() {
        unsigned m = *libres_;
        int p = __ffs(m) - 1;              // première page libre
        if (p < 0) return -1;
        *libres_ = m & ~(1u << p);
        return p;
    }

    __device__ void liberer(int p) { *libres_ |= (1u << p); }
    __device__ char* adresse(int p) const { return base_ + p * taille_page_; }
};

Le point important : libérer une page dès qu'elle n'est plus lue, pas à la fin de l'instruction. C'est ce qui permet à l'instruction suivante de démarrer ses chargements.

// Dans une instruction
int p_poids = pages.allouer();
charger_tma(pages.adresse(p_poids), poids_ptr, barriere);
attendre(barriere);
calculer_avec(pages.adresse(p_poids));
pages.liberer(p_poids);          // ← ICI, pas à la fin de la fonction
// ... suite du calcul avec d'autres pages ...

Dimensionnement : Hazy Research utilise 13 pages de 16 Ko sur H100 (213 Ko), MPK des pages de 32 Ko. Le compromis : petites pages = libération plus fine mais copies TMA moins efficaces.

8.5 Les compteurs de dépendance

L'attente

__device__ void attendre(unsigned* compteurs, int idx, unsigned cible) {
    // Lecture non cachée obligatoire
    volatile unsigned* c = &compteurs[idx];
    int essais = 0;
    while (*c < cible) {
        __nanosleep(min(20 << min(essais++, 8), 1024));  // recul exponentiel
    }
}

Le __nanosleep avec recul exponentiel réduit la contention sur le L2. Sans lui, des dizaines de blocs martèlent la même ligne de cache.

La signalisation

__threadfence();               // rendre les écritures visibles à TOUT le device
__syncthreads();               // tous les threads du bloc ont fini
if (threadIdx.x == 0) {
    atomicAdd(&compteurs[idx], 1u);
}

L'ordre est critique : __threadfence() avant l'incrément. Sinon, le consommateur peut voir le compteur avancer sans voir les données.

La version moderne

#include <cuda/atomic>

using Compteur = cuda::atomic<unsigned, cuda::thread_scope_device>;

// Producteur
compteur.fetch_add(1, cuda::memory_order_release);

// Consommateur
while (compteur.load(cuda::memory_order_acquire) < cible) { __nanosleep(50); }

Plus sûr : la paire release/acquire garantit la visibilité sans __threadfence() explicite, et le modèle mémoire est formellement défini.

8.6 L'ordre de construction

Ne construisez pas tout d'un coup. Sept étapes, chacune vérifiable.

Étape Ce que vous construisez Test de sortie
1 Un noyau persistant qui exécute une instruction triviale il se termine
2 La boucle de dispatch avec 2 opcodes l'ordre est respecté
3 Les compteurs entre deux instructions dépendantes résultat correct, pas d'interblocage
4 L'allocateur de pages pas de corruption, pas de fuite
5 Une instruction réelle (RMSNorm) résultat identique à PyTorch
6 Toutes les instructions d'une couche une couche correcte
7 Le modèle complet sortie identique jeton à jeton

Le critère du niveau 7

AutoMegaKernel se donne comme critère la parité jeton à jeton avec la référence Hugging Face, mesurée par une variance de perplexité de \(2{,}5 \times 10^{-7}\) sur SmolLM2-135M.

C'est le bon critère : pas « les résultats sont proches » mais « le modèle produit exactement la même séquence ». Toute divergence signale un bug de synchronisation ou de numérique.

8.7 Le débogage

C'est la partie la plus difficile, et il faut s'y préparer.

Les symptômes et leurs causes

Symptôme Cause probable
Le noyau ne se termine jamais trop de blocs (non-résidence), ou compteur jamais incrémenté
Résultats faux et non déterministes __threadfence() manquant, ou lecture cachée
Résultats faux mais déterministes erreur d'indexation ou de logique
Correct en -G, faux en -O3 variable non volatile, ou course de données
Correct à petite taille, faux à grande dépassement de pages, ou débordement de compteur
Lent sans raison apparente spilling de registres, contention atomique

Les outils

compute-sanitizer --tool memcheck  ./prog     # accès invalides
compute-sanitizer --tool racecheck ./prog     # courses en mémoire partagée
compute-sanitizer --tool synccheck ./prog     # __syncthreads() divergents

nvcc -Xptxas -v ...                            # registres et spills
ncu --set full --section WarpStateStats ./prog # où le temps est perdu

La technique du journal sur GPU

Puisqu'on ne peut pas déboguer pas à pas, on instrumente :

struct Trace { int sm, pc, opcode; unsigned long long horodatage; };

__device__ void journaliser(Trace* journal, unsigned* idx,
                            int sm, int pc, int opcode) {
    if (threadIdx.x == 0) {
        unsigned i = atomicAdd(idx, 1u);
        if (i < TAILLE_JOURNAL) {
            journal[i] = {sm, pc, opcode, clock64()};
        }
    }
}

Puis on analyse le journal côté hôte : ordre d'exécution, durées par instruction, SM qui attendent. C'est laborieux et c'est souvent le seul moyen.

L'approche par vérification statique

AutoMegaKernel (juin 2026) propose une alternative : un validateur d'IR d'ordonnancement figé qui certifie statiquement des propriétés de sûreté — absence d'interblocage et de course — sans preuve mécanisée.

Un mécanisme piloté par agent propose des ordonnancements, qui sont rejetés avant exécution s'ils sont dangereux. Les auteurs annoncent 7 160 ordonnancements adverses validés avec zéro faux positif d'acceptation.

C'est une direction prometteuse : plutôt que de déboguer les courses, les rendre impossibles par construction.

8.8 Les optimisations, dans l'ordre

Une fois que c'est correct :

# Optimisation Gain typique
1 Libérer les pages au plus tôt 1,2-1,3×
2 Découper les dépendances (4 morceaux au lieu de 1) 1,1-1,3×
3 Recul exponentiel dans l'attente 1,05-1,2×
4 Spécialiser les threads (loader/compute/storer) 1,02-1,06×
5 File de travail dynamique au lieu d'un ordre fixe 1,1-1,15×
6 Recouvrir la communication (multi-GPU) 1,1×
7 Localité de chiplet (B200, MI300X) 1,27-1,30×
8 Éliminer les spills 1,1-1,5×

Les chiffres sont ceux rapportés dans les travaux cités (MPK pour 1 et 6, Hazy Research pour 4 et 5, Fleet pour 7). Ils dépendent fortement du contexte.

8.9 La checklist

STRUCTURE
□  Grille = nb_SM × blocs_par_SM, calculé, pas deviné
□  cudaFuncSetAttribute pour la mémoire partagée > 48 Ko
□  __launch_bounds__ pour contrôler les registres
□  Assertion : blocs_par_sm >= 1

SYNCHRONISATION
□  __threadfence() AVANT chaque incrément de compteur
□  volatile ou cuda::atomic pour toutes les lectures de compteur
□  __syncthreads() après l'attente, avant l'exécution
□  Recul exponentiel dans les boucles d'attente
□  Aucun __syncthreads() dans une branche divergente

MÉMOIRE PARTAGÉE
□  Pages libérées dès la dernière lecture
□  Pas de fuite (toute page allouée est libérée)
□  Alignement 128 octets pour les destinations TMA

CORRECTION
□  compute-sanitizer memcheck : propre
□  compute-sanitizer racecheck : propre
□  Parité jeton à jeton avec la référence
□  Test à plusieurs tailles de lot et longueurs de contexte

PERFORMANCE
□  Aucun spill (nvcc -Xptxas -v)
□  Bande passante mesurée vs plancher théorique
□  Warp State Statistics : pas de Stall Barrier dominant

Résumé du chapitre

À retenir

  • Vérifiez d'abord que le gain est là (part du surcoût > 20 %) et essayez MPK avant d'écrire.
  • Structure : instructions, compteur ordinal par SM, dispatch par switch, allocateur de pages, compteurs de dépendance.
  • Grille = nb_SM × blocs_par_SM, calculé par cudaOccupancyMaxActiveBlocksPerMultiprocessor. Dépasser = interblocage.
  • Libérer les pages dès la dernière lecture, pas à la fin de l'instruction.
  • __threadfence() avant l'incrément, lectures volatile ou cuda::atomic avec acquire/release.
  • Construire en sept étapes vérifiables, avec pour critère final la parité jeton à jeton.
  • Déboguer avec compute-sanitizer et un journal sur GPU. La vérification statique (AutoMegaKernel) est une alternative prometteuse.

Vérifiez que vous avez compris

Pourquoi le recul exponentiel dans la boucle d'attente améliore-t-il les performances ?

Parce que la boucle d'attente naïve martèle une adresse en mémoire globale.

Chaque lecture volatile d'un compteur est une transaction qui atteint le L2. Avec 100 blocs sondant le même compteur des milliers de fois par microseconde, on sature la file des requêtes L2 — ce qui ralentit le producteur qui essaie d'écrire ses données au même endroit.

Le recul exponentiel (__nanosleep avec un délai croissant) réduit drastiquement le nombre de sondages, libérant la bande passante pour le travail utile.

C'est le même raisonnement que le pause d'un verrou tournant sur CPU.

Vous obtenez des résultats corrects en -G et faux en -O3. Que chercher ?

Presque toujours une variable de synchronisation non volatile.

En -G, le compilateur ne garde rien en registre : chaque lecture va réellement en mémoire, et la boucle d'attente fonctionne par accident.

En -O3, le compilateur constate que rien dans la boucle ne modifie le compteur, charge sa valeur une fois en registre, et produit soit une boucle infinie, soit une sortie immédiate.

Le remède : volatile, ou mieux, cuda::atomic avec un ordre mémoire explicite — qui exprime l'intention au compilateur et au matériel.

Cherchez aussi les __threadfence() manquants : en -G, la fenêtre de course est plus large et le bug se manifeste moins.

Votre megakernel est correct mais 20 % plus lent que la version à noyaux multiples. Où chercher ?

Quatre suspects, par ordre de probabilité :

  1. Spilling de registres. Le noyau unique contient tout le code ; le compilateur alloue les registres pour le pire cas. Vérifiez nvcc -Xptxas -v. C'est la cause la plus fréquente, et Hazy Research la cite explicitement comme un problème restant.
  2. Occupancy effondrée. Un seul bloc par SM, avec beaucoup de mémoire partagée réservée pour les pages. Si votre pipelining ne compense pas, vous perdez.
  3. Contention atomique. Trop de blocs sondant les mêmes compteurs. Recul exponentiel, ou réduction du nombre d'événements (c'est ce que fait la fusion d'événements de MPK).
  4. Instructions mal fusionnées. Si vos instructions sont plus lentes que les noyaux spécialisés qu'elles remplacent (cuBLAS, FlashAttention), le gain sur les frontières ne compense pas.

Le point 4 est le plus sérieux conceptuellement : un megakernel ne vaut que si ses instructions sont au niveau des noyaux qu'il remplace. C'est aussi pourquoi ThunderKittens ou CUTLASS sont indispensables pour les écrire.

Chapitre suivant : 9 · Quand ne pas le faire

Sources de ce chapitre