Contexte technique

Le blog décrit la construction d’un noyau d’attention dense pour le GPU NVIDIA Blackwell B200, écrit en CUDA pur avec un peu de PTX. L’objectif est d’atteindre 94,4 % des performances de FlashAttention‑4 (FA4) sur des séquences de 4 K, 8 K et 16 K, en conservant un format BF16, une dimension de tête de 128 et un comportement non‑causal. La démarche s’appuie sur 60 schémas illustrant chaque optimisation, ce qui en fait un guide visuel détaillé pour les développeurs GPU.

Parallélisation et architecture des CTA

Les tenseurs Q, K, V et la sortie O partagent la forme B × num_heads × seq_len × head_dim. Ils sont découpés en tuiles de [128, 128] afin de tenir dans la mémoire partagée. Chaque CTA (Cooperative Thread Array) est responsable d’une tuile de sortie O. Le nombre total de CTA correspond au nombre de tuiles O, ce qui garantit une couverture complète du résultat. À l’intérieur d’un CTA, la boucle sur les tuiles K/V reste séquentielle : le CTA charge une tuile Q une fois, puis itère sur toutes les paires de tuiles K/V, calcule S_tile = Q_tile @ K_tile.T, applique le softmax en ligne, puis accumule P_tile @ V_tile dans un tampon privé. Après la dernière itération, le CTA normalise le tampon et le stocke en mémoire globale.

Goulot d'étranglement du softmax et optimisations appliquées

Sur le B200, le débit des Tensor‑Cores a presque doublé, alors que les unités d’exponentielle (ALU/MUFU) n’ont guère évolué. Ainsi, le calcul du softmax, qui s’exécute sur ces unités, consomme un nombre de cycles comparable à celui des multiplications‑matrices (MMAs). Cette asymétrie constitue le principal goulot d’étranglement. Les optimisations successives visent à réduire le coût du softmax et à le superposer aux MMA :

• Réduction du coût du softmax : utilisation d’une implémentation en ligne qui évite les passages de données entre les registres et la mémoire partagée.

• Superposition du softmax avec les MMA : les calculs de softmax sont déclenchés pendant que les MMA traitent les tuiles précédentes, cachant ainsi la latence.

• Exploitation du PTX : insertion de sections PTX pour forcer l’utilisation de registres spécifiques et éviter les stalls du scheduler.

Chaque optimisation est introduite par un diagramme, puis le code correspondant est présenté. Le processus itératif permet de mesurer l’impact de chaque changement, ce qui explique la progression graduelle vers 94,4 % de la performance de FA4.

Résultats et portée

Le noyau final, testé sur des séquences de 4 K, 8 K et 16 K, atteint 94,4 % du débit de FlashAttention‑4 tout en restant entièrement implémenté en CUDA/PTX, sans recourir à CuTe. Cette approche démontre que les performances de pointe sont accessibles avec un code plus lisible et plus modifiable, ce qui facilite la recherche d’optimisations supplémentaires. Le guide s’applique également à d’autres kernels d’attention ou de multiplication de matrices, car les concepts de tiling, de parallélisation par CTA et de superposition des opérations sont généraux. Le code complet est disponible dans le dépôt GitHub du projet.

// Exemple de kernel de base (simplifié)
extern "C" __global__ void b200_attention(
    const half* __restrict__ Q,
    const half* __restrict__ K,
    const half* __restrict__ V,
    half* __restrict__ O,
    int seq_len, int head_dim) {
    // Chargement des tuiles Q, K, V dans la mémoire partagée
    __shared__ half Q_tile[128][128];
    __shared__ half K_tile[128][128];
    __shared__ half V_tile[128][128];
    // Calcul du produit Q·K^T
    half S_tile[128][128];
    // ... (code de multiplication matrix‑matrix) ...
    // Softmax en ligne
    // ... (code de réduction et exponentiation) ...
    // Multiplication S·V et accumulation
    // ... (code de MMA) ...
    // Écriture du résultat O
    // ...
}