Aller au contenu

3 · CUTLASS et CuTe DSL

La bibliothèque qui définit l'état de l'art de la multiplication matricielle sur GPU NVIDIA, et son algèbre de layouts — probablement l'abstraction la plus profonde de tout le domaine.


3.1 Ce qu'est CUTLASS

CUTLASS (CUDA Templates for Linear Algebra Subroutines) est une bibliothèque de gabarits C++ publiée par NVIDIA. Elle n'est pas une boîte noire comme cuBLAS : c'est un jeu de composants qu'on assemble pour construire sa propre GEMM.

Elle sert trois publics :

  1. ceux qui ont besoin d'une GEMM avec un épilogue custom (activation, quantification, biais fusionné) ;
  2. ceux qui ont besoin d'une forme inhabituelle (GEMM groupée, batched, convolution implicite) ;
  3. ceux qui veulent lire comment on écrit une GEMM optimale.

Le troisième usage est le plus répandu. CUTLASS est la documentation vivante des techniques décrites en partie 4.


3.2 CuTe : l'algèbre de layouts

Depuis CUTLASS 3.x, tout repose sur CuTe (CUDA Tensors), une couche qui formalise la question : où se trouve l'élément \((i, j)\) ?

La définition

Un layout est une paire (forme, pas) :

\[ \text{Layout} = (\text{Shape}, \text{Stride}) \]

qui définit une fonction des coordonnées vers un indice linéaire :

\[ \text{idx}(i_0, i_1, \dots) = \sum_k i_k \cdot d_k \]

où \(d_k\) est le pas de la dimension \(k\).

using namespace cute;

// Matrice 8×4 stockée par colonnes (column-major)
auto layout = make_layout(make_shape(Int<8>{}, Int<4>{}),
                          make_stride(Int<1>{}, Int<8>{}));
// idx(i, j) = i * 1 + j * 8

print_layout(layout);
//       0   1   2   3
//   0   0   8  16  24
//   1   1   9  17  25
//   ...

Ce n'est pas nouveau. Ce qui est nouveau, c'est que les layouts se composent algébriquement.

La composition

\[ (A \circ B)(x) = A(B(x)) \]

Si \(B\) décrit « comment un thread parcourt sa portion » et \(A\) « où la portion est en mémoire », alors \(A \circ B\) décrit directement « quelle adresse ce thread touche ».

Les opérations de CuTe :

Opération Sens
composition(A, B) \(A \circ B\)
logical_divide(A, B) découper \(A\) en tuiles selon \(B\)
logical_product(A, B) répliquer \(A\) selon \(B\)
complement(A, M) le layout « complémentaire » dans un espace de taille \(M\)
coalesce(A) simplifier un layout
Swizzle<B,M,S> permutation XOR anti-conflits

Pourquoi c'est puissant

Toutes les questions pénibles de la programmation GPU deviennent des manipulations de layouts :

  • « comment répartir une tuile 128×64 sur 256 threads ? » → une division logique ;
  • « comment éviter les conflits de banc ? » → composer avec un Swizzle ;
  • « quelle disposition de registres attend wgmma ? » → un layout fourni par l'atome ;
  • « comment charger cette tuile avec TMA ? » → une TiledCopy.

Et surtout : ces manipulations sont vérifiables à la compilation, parce que les formes et les pas sont des types.

Les atomes et les opérations pavées

CuTe distingue :

  • un atome : la plus petite opération matérielle (une instruction mma, une copie cp.async) avec le layout de threads qu'elle exige ;
  • une opération pavée (tiled operation) : l'atome répliqué sur une tuile plus grande.
// Atome MMA : une instruction wgmma
using MmaAtom = SM90_64x128x16_F32BF16BF16_SS;

// Pavage : comment répliquer l'atome sur le warp group
using TiledMma = decltype(make_tiled_mma(MmaAtom{},
                                         Layout<Shape<_2,_1,_1>>{}));

// Atome de copie : TMA
using CopyAtom = SM90_TMA_LOAD;

Le programmeur choisit les atomes et le pavage ; CuTe garantit la cohérence des layouts.


3.3 La structure d'un noyau CUTLASS 3.x

Kernel (sm90_gemm_tma_warpspecialized_pingpong)
 ├─ CollectiveMainloop
 │   ├─ TiledMma          : quelle instruction MMA, quel pavage
 │   ├─ TiledCopy A, B    : TMA, avec quel swizzle
 │   ├─ SmemLayout A, B   : disposition en mémoire partagée
 │   └─ PipelineTmaAsync  : mbarriers, nombre d'étages
 ├─ CollectiveEpilogue
 │   ├─ TiledCopy C, D    : écriture, éventuellement TMA
 │   └─ ThreadEpilogueOp  : activation, biais, quantification
 └─ TileScheduler         : quel bloc traite quelle tuile

Le CollectiveBuilder assemble tout cela à partir de spécifications de haut niveau :

using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
    cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
    ElementA, LayoutA, AlignmentA,
    ElementB, LayoutB, AlignmentB,
    ElementAccumulator,
    TileShape, ClusterShape,
    cutlass::gemm::collective::StageCountAuto,
    cutlass::gemm::collective::KernelScheduleAuto
>::CollectiveOp;

StageCountAuto calcule le nombre d'étages qui tient en mémoire partagée. KernelScheduleAuto choisit entre pingpong et cooperative.

Le coût d'entrée est réel

CUTLASS est du C++ à métaprogrammation intensive. Un message d'erreur de template peut faire plusieurs milliers de lignes. Les temps de compilation se comptent en minutes par instanciation.

C'est ce qui a motivé le CuTe DSL en Python.


3.4 Le CuTe DSL

CUTLASS 4.x introduit CuTe DSL, une interface Python native aux mêmes concepts.

A Python-native interface for writing high-performance CUDA kernels based on core CUTLASS and CuTe concepts without performance compromises.

Le fonctionnement : le code Python est traduit en une représentation intermédiaire propre, puis compilé en JIT vers du CUDA optimisé via MLIR et ptxas.

# CuTe DSL — structure indicative
import cutlass
import cutlass.cute as cute

@cute.kernel
def mon_noyau(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, _, _ = cute.arch.block_idx()

    # Découper les tenseurs en tuiles
    gA = cute.local_tile(mA, tiler, coord)
    sA = cute.make_tensor(smem_ptr, sA_layout)

    # Copie TMA
    cute.copy(tma_atom, tAgA, tAsA)

    # MMA
    cute.gemm(tiled_mma, tCrC, tCrA, tCrB, tCrC)

Ce que ça change

CUTLASS C++ CuTe DSL
Langage C++ templates Python
Compilation minutes secondes
Messages d'erreur illisibles lisibles
Vérification des layouts à la compilation à la compilation (dans l'IR)
Performance référence équivalente
Intégration CMake pip install nvidia-cutlass-dsl

La preuve par l'usage : FlashAttention-4 est écrit entièrement en CuTe DSL. Il atteint jusqu'à 1 605 TFLOPS sur B200 (71 % d'utilisation matérielle), soit 1,3× plus rapide que cuDNN 9.13 et 2,7× plus rapide que les implémentations Triton. Ses auteurs notent que l'installation et la compilation prennent désormais des secondes au lieu de minutes ou d'heures.

C'est l'argument le plus solide possible : le code le plus rapide du domaine est écrit dans un DSL Python.

État en août 2026

La version courante est CUTLASS 4.5.2 (mai 2026). Le CuTe DSL est en bêta publique et doit sortir de bêta « d'ici la fin de l'été 2026 ».

Une intégration avec JAX est documentée, permettant d'appeler des noyaux CuTe DSL depuis des programmes JAX.


3.5 Quand utiliser quoi

Situation Choix
GEMM standard, formes courantes cuBLAS
GEMM avec épilogue fusionné CUTLASS C++ (ou CuTe DSL)
GEMM groupée (MoE) CUTLASS, ou cuBLAS depuis CUDA 13.1
Attention custom de pointe CuTe DSL
Format de quantification maison CUTLASS C++ (ElementA custom)
Prototypage rapide Triton, puis CuTe DSL si insuffisant
Comprendre comment on fait lire CUTLASS

Le conseil pratique

Même si vous n'écrivez jamais une ligne de CUTLASS, lisez ses exemples. Le répertoire examples/ du dépôt contient des implémentations commentées de tous les motifs de la partie 4 : TMA, warp specialization, ping-pong, clusters, tcgen05.

C'est la meilleure source pour voir ces techniques assemblées correctement.


3.6 Le tutoriel Colfax

Une ressource à connaître : Colfax Research publie une série de tutoriels CUTLASS qui sont, de loin, la meilleure documentation pédagogique du domaine.

Tutoriel Sujet
Mastering the NVIDIA TMA descripteurs, copies, multicast
Fast Matrix-Multiplication with WGMMA l'instruction, ses descripteurs
Efficient GEMM kernel designs with Pipelining étages, mbarriers
GEMM with Thread Block Clusters on Blackwell clusters et tcgen05
Persistent Kernels and Stream-K l'ordonnancement

Ces documents expliquent le matériel et l'API, avec du code complet. Ils sont référencés dans les sources de la partie 4 pour cette raison.


Résumé du chapitre

À retenir

  • CUTLASS est un jeu de composants C++ pour construire des GEMM, pas une boîte noire. Trois usages : épilogue custom, forme inhabituelle, apprentissage.
  • CuTe formalise les layouts comme des paires (forme, pas) qui se composent algébriquement. Répartition de threads, swizzling, dispositions de tensor cores deviennent des manipulations de layouts vérifiées à la compilation.
  • Un noyau CUTLASS 3.x = CollectiveMainloop + CollectiveEpilogue + TileScheduler.
  • CuTe DSL (CUTLASS 4.x) porte tout cela en Python, avec compilation en secondes et performance équivalente.
  • FlashAttention-4 est écrit entièrement en CuTe DSL : 1 605 TFLOPS sur B200, 1,3× cuDNN, 2,7× Triton.
  • Lisez les exemples CUTLASS et les tutoriels Colfax même si vous n'écrivez pas de CUTLASS.

Vérifiez que vous avez compris

En quoi un layout CuTe diffère-t-il d'une simple paire (shape, stride) NumPy ?

Sur le principe, pas du tout — c'est la même chose. La différence est dans ce qu'on en fait :

  1. Les layouts CuTe sont hiérarchiques : une forme peut être imbriquée, ((4,2),8), ce qui permet d'exprimer « 4 threads × 2 valeurs, sur 8 tuiles ».
  2. Ils se composent : \(A \circ B\) est un layout, avec des règles algébriques (division logique, produit logique, complément).
  3. Ils sont des types C++, donc évalués et vérifiés à la compilation. Un layout incompatible avec un atome MMA est une erreur de compilation, pas un résultat faux.

C'est cette troisième propriété qui rend CUTLASS sûr malgré la complexité de ce qu'il génère.

Pourquoi FlashAttention-4 est-il en CuTe DSL plutôt qu'en Triton ?

Parce que ses optimisations exigent un contrôle que Triton n'expose pas :

  • une approximation polynomiale de exp() sur les unités FMA plutôt que sur la SFU, ce qui demande de contrôler les instructions émises ;
  • un pipeline de remise à l'échelle conditionnelle du softmax en ligne, réduisant d'environ 10× le nombre de rescalings ;
  • un ordonnancement fin entre tcgen05, TMA et les unités vectorielles sur Blackwell.

Les auteurs mesurent 2,7× de mieux que les implémentations Triton — ce qui quantifie précisément l'écart entre les deux niveaux d'abstraction sur ce problème.

Vous devez ajouter une activation GELU fusionnée à une GEMM. Quelle voie ?

Trois options, par ordre d'effort croissant :

  1. torch.compile : Inductor sait fusionner l'épilogue d'une GEMM avec une activation dans certains cas. À essayer en premier.
  2. CUTLASS avec un EpilogueOp custom : c'est le cas d'usage nominal. On fournit un foncteur appliqué aux accumulateurs avant l'écriture, et tout le reste (TMA, pipeline, MMA) est réutilisé tel quel.
  3. Triton : écrire la GEMM complète avec la GELU dans la boucle. Plus simple à écrire, mais on repart d'une GEMM Triton, donc en dessous du pic.

L'option 2 est la bonne si la performance compte : on garde une GEMM de référence et on ne modifie que l'épilogue.


Chapitre suivant : 4 · ThunderKittens


Sources de ce chapitre