4 · ThunderKittens

Une bibliothèque C++ d'une équipe de recherche (Hazy Research, Stanford) dont le pari est simple : la bonne unité de programmation n'est ni le thread ni le bloc, c'est la tuile de 16×16.

C'est aussi la bibliothèque avec laquelle les megakernels de la partie 8 ont été écrits, ce qui justifie ce chapitre.

---

4.1 L'observation de départ

Le matériel des tensor cores opère sur des fragments de taille fixe. Une instruction mma.m16n8k16 manipule des tuiles 16×16 et 16×8. Une instruction wgmma manipule des multiples de ces tuiles.

Donc : au lieu de manipuler des threads et de reconstruire péniblement des tuiles, prenons la tuile comme type primitif.

#include "kittens.cuh"
using namespace kittens;

// Une tuile de registres : 64×64 en bf16, disposition ligne
rt_bf<64, 64> a_reg;

// Une tuile de mémoire partagée : 128×64
st_bf<128, 64> a_smem;

Les types encodent : l'emplacement (registres rt_, mémoire partagée st_), le type d'élément, les dimensions, et la disposition. Les opérations sont typées en conséquence.

---

4.2 Ce à quoi ressemble le code

__global__ void attention_kernel(const bf16* Q, const bf16* K,
                                 const bf16* V, bf16* O) {
    // Tuiles en mémoire partagée
    __shared__ st_bf<64, 64> q_s, k_s, v_s;

    // Tuiles en registres
    rt_fl<16, 64> att_block;
    rt_bf<16, 64> att_block_bf;
    rt_fl<16, 64> o_reg;

    zero(o_reg);

    for (int bloc = 0; bloc < n_blocs; ++bloc) {
        // Chargement asynchrone (TMA sous le capot)
        load_async(k_s, K + offset(bloc), barriere);
        load_async(v_s, V + offset(bloc), barriere);
        wait(barriere);

        // Produit matriciel : Q·Kᵀ
        zero(att_block);
        mma_ABt(att_block, q_reg, k_s, att_block);

        // Softmax en ligne
        row_max(max_vec, att_block, max_vec);
        sub_row(att_block, att_block, max_vec);
        exp(att_block, att_block);
        row_sum(norm_vec, att_block, norm_vec);

        copy(att_block_bf, att_block);

        // Accumulation : att·V
        mma_AB(o_reg, att_block_bf, v_s, o_reg);
    }

    div_row(o_reg, o_reg, norm_vec);
    store(O + offset_sortie, o_reg);
}

Ce qui saute aux yeux :

• aucun threadIdx dans la logique — il n'apparaît que pour les identifiants de warp ;
• les opérations (mma_ABt, row_max, exp) prennent des tuiles ;
• load_async masque TMA, les descripteurs et le swizzling ;
• le code se lit comme l'algorithme mathématique.

L'argument central

ThunderKittens produit du code aussi rapide que du CUDA expert tout en restant lisible par quelqu'un qui connaît l'algorithme mais pas le matériel. C'est ce qui compte pour une équipe de recherche, où le code doit être modifié souvent.

---

4.3 Ce que la bibliothèque couvre

Catégorie	Opérations
Types de tuiles	rt_* (registres), st_* (mémoire partagée), vecteurs associés
Mouvement	load, store, load_async (TMA), store_async
Matriciel	mma_AB, mma_ABt, mma_AtB, mma_AtBt
Élémentaire	exp, mul, add, max, copy, zero…
Réductions	row_max, row_sum, col_max, col_sum
Diffusion	sub_row, div_row, mul_col…
Synchronisation	barrières, sémaphores, gestion de pages
Multi-GPU	Parallel Global Layout pour les tenseurs distribués

Le support matériel : wgmma sur H100, tcgen05 sur B200. La bibliothèque choisit l'instruction selon la cible.

---

4.4 Le lien avec les megakernels

C'est la raison principale de ce chapitre.

Les deux megakernels de Stanford — le latence-optimisé sur Llama-1B et le débit-optimisé sur Llama-70B en tensor-parallèle — sont construits sur ThunderKittens.

Le motif : chaque « instruction » du megakernel est un gabarit C++ commun, avec trois sections (load, compute, store), écrites en primitives de tuiles. L'interpréteur sur GPU exécute ces instructions à la chaîne.

// Structure conceptuelle d'une instruction de megakernel
template <typename Config>
struct MonInstruction {
    static __device__ void loader(state<Config>& s, /* args */) {
        // émettre les copies TMA vers les pages de mémoire partagée
    }
    static __device__ void compute(state<Config>& s, /* args */) {
        // les MMA et l'arithmétique, en tuiles ThunderKittens
    }
    static __device__ void storer(state<Config>& s, /* args */) {
        // écrire les résultats, incrémenter les compteurs
    }
};

Un dépôt dédié, HazyResearch/Megakernels, contient ces instructions pour Llama.

Pourquoi ThunderKittens plutôt que CUTLASS pour un megakernel

Trois raisons pratiques :

1. La flexibilité de composition. Un megakernel enchaîne des opérations hétérogènes (norme, projection, attention, MLP). CUTLASS est structuré autour de la GEMM ; ThunderKittens fournit des briques.
2. Le contrôle de la mémoire partagée. Le megakernel a besoin d'un allocateur par pages ; ThunderKittens n'impose pas d'allocation fixe.
3. Le temps de compilation. Un megakernel est un seul noyau énorme ; les templates CUTLASS le rendraient impraticable.

---

4.5 Les limites

Ce que ThunderKittens n'est pas

• Ce n'est pas une bibliothèque de production généraliste. C'est un projet de recherche, avec les conséquences habituelles : API en évolution, documentation inégale, couverture matérielle ciblée sur les GPU de centre de données récents.
• NVIDIA uniquement. Pas de back-end AMD.
• La granularité de 16×16 est imposée. Les problèmes qui ne s'y prêtent pas (structures irrégulières, tailles très petites) n'en tirent rien.
• Le débogage reste du débogage CUDA. L'abstraction aide à écrire, pas à diagnostiquer.

Pour un usage en production sur une GEMM ou une attention standard, CUTLASS, cuDNN ou FlashAttention restent les bons choix.

---

4.6 Où le situer

Contrôle ↑
         │  CUDA C++ brut
         │  ThunderKittens   ← tuiles 16×16, lisible, TMA/wgmma enveloppés
         │  CuTe DSL         ← layouts explicites, Python
         │  CUTLASS C++      ← composants GEMM
         │  Gluon
         │  Triton
         ↓  Helion

ThunderKittens et CuTe DSL occupent une niche similaire — contrôle proche du matériel avec une abstraction utilisable — mais avec des philosophies opposées :

	CuTe	ThunderKittens
Abstraction centrale	le layout	la tuile
Généralité	maximale (tout layout exprimable)	volontairement restreinte à 16×16
Courbe d'apprentissage	raide	douce
Vérification	types C++ / IR	types C++
Soutien	NVIDIA	académique

---

Résumé du chapitre

À retenir

• ThunderKittens prend la tuile 16×16 comme type primitif, parce que c'est la granularité du matériel des tensor cores.
• Types rt_* (registres) et st_* (mémoire partagée), opérations typées (mma_AB, row_max, exp), TMA enveloppée dans load_async.
• Le code se lit comme l'algorithme, sans threadIdx.
• C'est la base des megakernels de Stanford : chaque « instruction » est un gabarit load / compute / store écrit en primitives de tuiles.
• Projet de recherche : NVIDIA seul, API en évolution, granularité imposée.
• Pour de la production sur une opération standard, préférez CUTLASS, cuDNN ou FlashAttention.

---

Vérifiez que vous avez compris

Pourquoi 16×16 et pas 32×32 ou 8×8 ?

Parce que c'est la granularité des instructions matérielles. Une mma.m16n8k16 produit une tuile 16×8 ; les fragments d'entrée sont des 16×16 et 16×8. Les instructions ldmatrix chargent des tuiles 16×16 depuis la mémoire partagée.

Choisir 16×16 signifie que chaque opération de la bibliothèque correspond à un petit nombre d'instructions matérielles, sans reconstruction ni permutation coûteuse. Une granularité plus fine gaspillerait les instructions ; plus grosse forcerait des découpages.

Peut-on écrire une attention avec masque causal en ThunderKittens sans perdre le facteur 2 ?

Oui, et c'est même naturel : puisqu'on raisonne en tuiles, il suffit de ne pas boucler sur les tuiles entièrement masquées.

for (int bloc = 0; bloc <= bloc_courant; ++bloc) {   // ← borne causale
    // ... les tuiles au-dessus de la diagonale ne sont jamais chargées
}
// Pour la tuile diagonale, appliquer le masque triangulaire
make_causal(att_block, att_block, -INFINITY);

C'est exactement ce que fait FlashAttention. Le raisonnement par tuiles rend ce genre d'optimisation évident, là où un raisonnement par thread le rend pénible.

Votre équipe doit maintenir un noyau d'attention custom pendant trois ans. CUTLASS ou ThunderKittens ?

La question porte sur la stabilité, pas sur la performance.

CUTLASS est soutenu par NVIDIA, versionné, et suivra les futures architectures. Son coût est la complexité des templates et des temps de compilation.

ThunderKittens est plus lisible et plus rapide à modifier, mais c'est un projet académique : rien ne garantit son support sur Rubin ou au-delà.

Une troisième voie mérite considération : CuTe DSL, qui combine le soutien NVIDIA et une syntaxe Python lisible. C'est le choix qu'a fait l'équipe FlashAttention pour la version 4 — un signal fort sur ce qui sera maintenable à moyen terme.

---

Chapitre suivant : 5 · Helion et torch.compile

---

Sources de ce chapitre

• ThunderKittens sur GitHub
• HazyResearch/Megakernels
• Hazy Research, Look Ma, No Bubbles!
• Hazy Research, We Bought the Whole GPU