Aller au contenu

2 · Le Tensor Memory Accelerator

Le mécanisme le plus important de Hopper. Une unité matérielle dédiée qui copie des tenseurs multidimensionnels entre mémoire globale et mémoire partagée, sur ordre d'un seul thread.

2.1 Le problème que ça résout

Avec cp.async, charger une tuile \(128 \times 64\) de bf16 demande :

  • que chaque thread calcule son adresse source et destination ;
  • de gérer les bords (que faire si la tuile dépasse la matrice ?) ;
  • de gérer le swizzling pour éviter les conflits de banc ;
  • d'émettre une instruction par thread, soit 256 instructions pour un bloc.

Tout ce calcul d'adresses consomme des registres et des cycles d'exécution, dans une boucle qui devrait ne faire que du produit matriciel.

TMA supprime tout cela.

2.2 Le principe

TMA repose sur un descripteur de tenseur (CUtensorMap), une structure de 128 octets construite côté hôte, qui décrit une fois pour toutes :

  • l'adresse de base du tenseur en mémoire globale ;
  • son nombre de dimensions (jusqu'à 5) ;
  • la taille de chaque dimension ;
  • les pas (strides) de chaque dimension ;
  • la forme de la tuile à copier ;
  • le mode de swizzle appliqué à la destination ;
  • le mode de remplissage des bords (zéros ou NaN) ;
  • le type d'élément.

Ensuite, dans le noyau, un seul thread émet la copie en donnant juste des coordonnées :

cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes
    [%dst_smem], [%tensor_map, {%coord_x, %coord_y}], [%mbarrier];

Le matériel fait le reste : calcul des adresses, gestion des bords, swizzling, et incrémentation du compteur d'octets de la mbarrier.

L'idée à retenir

TMA transforme « 256 threads calculent 256 adresses » en « 1 thread donne 2 coordonnées ». Le calcul d'adresses passe du logiciel au matériel.

C'est ce que la documentation appelle un modèle mono-thread : a single thread in a warp is elected to issue an asynchronous TMA operation, et tout le bloc continue à travailler pendant que les données sont en vol.

2.3 Construire un descripteur

Côté hôte, avec l'API pilote :

#include <cuda.h>

CUtensorMap creer_descripteur(void* base, int lignes, int colonnes) {
    CUtensorMap desc;

    uint64_t taille[2]  = { (uint64_t)colonnes, (uint64_t)lignes };
    uint64_t pas[1]     = { (uint64_t)colonnes * sizeof(__nv_bfloat16) };
    uint32_t tuile[2]   = { 64, 128 };      // forme de la tuile en mém. partagée
    uint32_t pas_elem[2]= { 1, 1 };

    cuTensorMapEncodeTiled(
        &desc,
        CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
        2,                                  // rang
        base,
        taille,
        pas,
        tuile,
        pas_elem,
        CU_TENSOR_MAP_INTERLEAVE_NONE,
        CU_TENSOR_MAP_SWIZZLE_128B,         // ← swizzle géré par le matériel
        CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
        CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE   // bords remplis de zéros
    );
    return desc;
}

Le descripteur est ensuite passé au noyau, soit par argument (il doit être aligné sur 64 octets et déclaré const __grid_constant__), soit via la mémoire constante.

__global__ void noyau(const __grid_constant__ CUtensorMap desc_A,
                      const __grid_constant__ CUtensorMap desc_B) { ... }

Les contraintes du descripteur

  • L'adresse de base doit être alignée sur 16 octets ;
  • le pas de la dimension la plus interne doit être un multiple de 16 octets ;
  • la destination en mémoire partagée doit être alignée sur 128 octets ;
  • le mode de swizzle contraint la largeur de la tuile (32, 64 ou 128 octets).

Une violation ne produit pas d'erreur claire : elle produit des données fausses. C'est la principale raison pour laquelle on utilise CUTLASS.

2.4 Émettre une copie dans le noyau

#include <cuda/barrier>
#include <cudaTypedefs.h>

__global__ void gemm_tma(const __grid_constant__ CUtensorMap desc_A, ...) {
    __shared__ alignas(128) __nv_bfloat16 As[128][64];
    __shared__ cuda::barrier<cuda::thread_scope_block> bar;

    if (threadIdx.x == 0) {
        init(&bar, blockDim.x);
        cuda::device::experimental::fence_proxy_async_shared_cta();
    }
    __syncthreads();

    // Un seul thread émet la copie
    if (threadIdx.x == 0) {
        cuda::device::experimental::cp_async_bulk_tensor_2d_global_to_shared(
            &As[0][0], &desc_A, coord_x, coord_y, bar);
        // Annoncer les octets attendus
        cuda::device::barrier_expect_tx(bar, sizeof(As));
    }

    // Tout le bloc attend l'arrivée des données
    bar.arrive_and_wait();

    // Utiliser As[][]
}

Le point crucial est barrier_expect_tx : il indique à la mbarrier combien d'octets doivent arriver. Le matériel TMA incrémente le compteur au fur et à mesure ; l'attente se débloque quand le compte est atteint.

C'est exactement l'usage pour lequel mbarrier a été conçue.

2.5 Ce que TMA sait faire d'autre

Copies dans les deux sens

TMA gère aussi mémoire partagée → mémoire globale (cp.async.bulk.tensor...global.shared), avec les mêmes descripteurs. Utile pour l'épilogue d'une GEMM.

Multicast vers un cluster

C'est la fonctionnalité la plus sous-estimée. Une seule copie TMA peut écrire dans la mémoire partagée de plusieurs blocs d'un cluster :

cp.async.bulk.tensor.2d.shared::cluster.global.multicast::cluster ...

Dans une GEMM, la même tuile de \(\mathbf{A}\) est nécessaire à tous les blocs d'une même ligne de tuiles. Sans multicast, chaque bloc la lit depuis la HBM. Avec multicast, elle est lue une fois et distribuée.

Gain : divise le trafic HBM par le nombre de blocs du cluster (typiquement 2 ou 4).

Remplissage des bords

Le mode out-of-bounds fill remplit automatiquement de zéros les parties de la tuile qui dépassent le tenseur. Cela supprime toutes les gardes de bord du noyau — un gain de code et de divergence considérable dans les convolutions et les GEMM de dimensions non multiples de la tuile.

im2col

TMA a un mode dédié à la convolution qui effectue la transformation im2col à la volée pendant la copie. La convolution devient une GEMM sans matérialiser la matrice intermédiaire.

2.6 Ce que TMA change dans le raisonnement

1. Le calcul d'adresses disparaît du budget d'instructions. Dans une GEMM classique, une fraction non négligeable des instructions sert à calculer des adresses. Elles disparaissent.

2. Les registres se libèrent. Plus d'adresses ni de valeurs en transit dans les registres. Ces registres vont aux accumulateurs, donc à des tuiles plus grandes, donc à une intensité arithmétique supérieure.

3. L'occupancy cesse d'être le levier. Avec un seul thread qui alimente le bloc en permanence, il n'est plus nécessaire d'avoir 64 warps pour maintenir la mémoire occupée. On peut descendre à 4 ou 8 warps et tout donner aux registres et à la mémoire partagée.

4. Le modèle producteur/consommateur devient naturel. Un warp dédié émet les TMA, les autres calculent. C'est la spécialisation des warps.

Le lien avec wgmma

TMA et wgmma sont conçus ensemble. wgmma lit ses opérandes directement en mémoire partagée via un descripteur ; TMA les y dépose avec le bon swizzle. La combinaison donne un pipeline où les données ne touchent jamais les registres entre la HBM et le tensor core.

C'est ce qui permet à wgmma d'atteindre 95 % du pic Hopper contre 62,9 % pour l'instruction mma classique (arXiv:2501.12084).

2.7 Le prix à payer

TMA est difficile

  • Les descripteurs doivent être construits côté hôte, ce qui complique le code dynamique (formes variables).
  • Les contraintes d'alignement et de swizzle sont nombreuses et leur violation est silencieuse.
  • Le débogage est pénible : la copie est asynchrone et effectuée par une unité matérielle séparée.
  • Le code n'est pas portable : sm_90a minimum, pas de repli automatique.

En pratique, on utilise :

  • CUTLASS / CuTe (SM90_TMA_LOAD, TiledCopy) ;
  • ThunderKittens, qui expose tma::load_async(tuile, ...) ;
  • Triton, dont le compilateur émet TMA depuis tl.load sur des descripteurs (tl.make_tensor_descriptor) ;
  • CuTe DSL en Python.

2.8 Blackwell : ce qui s'y ajoute

Blackwell conserve TMA et lui ajoute :

  • le support des copies vers la Tensor Memory (tcgen05.cp) ;
  • le Cluster Launch Control, qui permet à un cluster de demander dynamiquement le prochain lot de travail au lieu d'un mapping statique — utile pour l'équilibrage de charge dans les noyaux persistants ;
  • des modes de swizzle supplémentaires adaptés aux formats à échelle par blocs (MXFP).

Le chapitre wgmma et tcgen05 détaille l'articulation.

Résumé du chapitre

À retenir

  • TMA copie des tenseurs jusqu'à 5D entre mémoire globale et mémoire partagée, sur ordre d'un seul thread.
  • Le descripteur (CUtensorMap, 128 octets) est construit côté hôte et encode tailles, pas, forme de tuile, swizzle et remplissage des bords.
  • Il se synchronise avec une mbarrier en mode expect_tx : on attend l'arrivée effective des octets.
  • Le multicast vers un cluster divise le trafic HBM par le nombre de blocs.
  • Le remplissage automatique des bords supprime les gardes du noyau.
  • Conséquence stratégique : les registres se libèrent, les tuiles grandissent, et l'occupancy cesse d'être le levier principal.
  • À utiliser via CUTLASS, ThunderKittens ou Triton, pas à la main.

Vérifiez que vous avez compris

Pourquoi TMA a-t-il besoin d'une mbarrier plutôt que d'un simple __syncthreads() ?

Parce que __syncthreads() synchronise les threads, et qu'ici il n'y a qu'un seul thread impliqué dans l'émission. Après le __syncthreads(), tous les threads auraient franchi la barrière alors que les données ne sont pas encore arrivées : la copie est effectuée par une unité matérielle asynchrone.

La mbarrier en mode expect_tx attend un nombre d'octets, ce qui est la condition réelle. C'est précisément pour TMA que ce mécanisme a été conçu.

Dans une GEMM avec un cluster de 2 blocs, le multicast TMA divise le trafic HBM par combien ?

Par 2 sur les opérandes partagés. Concrètement, si les deux blocs du cluster calculent des tuiles de \(\mathbf{C}\) sur la même bande de lignes, ils ont besoin de la même tuile de \(\mathbf{A}\). Sans multicast, elle est lue deux fois depuis la HBM ; avec, une seule fois.

Sur une GEMM où \(\mathbf{A}\) et \(\mathbf{B}\) contribuent également au trafic, le gain global est de ~25 % (moitié du trafic divisée par 2). Avec des clusters de 4, ~37 %.

Pourquoi le mode « remplissage des bords » de TMA est-il plus intéressant qu'il n'y paraît ?

Trois gains cumulés.

  1. Suppression du code de garde : plus de if (i < N) dans la boucle chaude.
  2. Suppression de la divergence : les gardes de bord font diverger le dernier warp de chaque bloc.
  3. Uniformité des tuiles : toutes les tuiles ont la même forme, donc le même chemin de code, donc pas de spécialisation ni de branchement.

C'est aussi ce qui permet aux noyaux wgmma d'avoir une boucle interne entièrement déroulée et sans branchement — condition nécessaire pour approcher le pic.

Chapitre suivant : 3 · Clusters et mémoire partagée distribuée

Sources de ce chapitre