Auto-attention, noyaux CUDA et FlashAttention

Les Transformers modernes reposent sur le calcul matriciel de l'auto-attention, des compilateurs GPU de haut niveau comme Triton et des noyaux CUDA tenant compte de la mémoire comme FlashAttention pour traiter d'immenses séquences de tokens à la vitesse maximale du matériel.

1. Le fonctionnement mathématique de l'auto-attention dans un transformeur de vision

Dans un Vision Transformer (ViT), une image (par exemple, 224×224 pixels) est découpée en une grille de zones de 16×16, produisant N = 196 zones aplaties. Chaque zone est projetée linéairement dans un vecteur de dimension D (par exemple, 768) et combinée à un plongement positionnel, formant une matrice d'entrée X ∈ R^(N × D).

• Création des requêtes, des clés et des valeurs. Le modèle multiplie la matrice des zones X par trois matrices de poids apprises (W_Q, W_K, W_V ∈ R^(D × d_k)) pour créer trois représentations distinctes de chaque zone d'image :

Q=XW_(Q), K=XW_(K), V=XW_(V)

• Requête (Q) : représente ce qu'une zone d'image donnée recherche (par exemple, "Je suis une texture de fourrure et je cherche des yeux ou des oreilles à proximité").
• Clé (K) : indique les informations contenues dans une zone (par exemple, "Je contiens un contour courbe net").
• Valeur (V) : contient les caractéristiques visuelles effectives qui seront transmises si une correspondance est trouvée.

• Calcul des scores d'attention. Pour mesurer la force du lien entre chaque zone et toutes les autres, le modèle calcule le produit scalaire entre toutes les requêtes et toutes les clés :

Scores=(QK^(T))/(sqrt(d_(k)))

La division par sqrt(d_(k)) empêche les produits scalaires élevés de pousser la fonction softmax qui suit vers des régions aux gradients minuscules. On obtient ainsi une matrice d'attention N × N dont l'entrée (i, j) évalue l'attention que la zone i doit accorder à la zone j.

• Pondération et agrégation par Softmax. L'application de Softmax ligne par ligne convertit les scores bruts en probabilités dont la somme vaut 1. La multiplication de ces probabilités par la matrice des valeurs V produit une représentation combinée de chaque zone :

Attention(Q,K,V)=softmax((QK^(T))/(sqrt(d_(k))))V

Dans l'auto-attention multi-tête, ce même processus s'exécute en parallèle dans plusieurs têtes (par exemple, 12), permettant à différentes têtes de se concentrer simultanément sur différentes distances spatiales, couleurs ou parties d'objets.

2. Comparaison pratique du code : CUDA C++ brut et OpenAI Triton

Pour voir comment la programmation GPU a évolué, prenons une simple addition vectorielle élément par élément (C = A + B) pour N éléments.

CUDA C++ brut (gestion au niveau des threads)

En CUDA C++, vous programmez du point de vue d'un seul thread scalaire. Vous devez calculer manuellement les identifiants de threads, vérifier les limites des tableaux et écrire les configurations de lancement répétitives côté hôte.

#include <cuda_runtime.h>

// Device kernel: executed by thousands of individual threads
__global__ void vector_add_cuda(const float* A, const float* B, float* C, int N) {
    // Manually compute global thread index from block and thread IDs
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Guard against out-of-bounds memory access
    if (idx < N) {
        C[idx] = A[idx] + B[idx];
    }
}

// Host launch setup
void launch_cuda(const float* d_A, const float* d_B, float* d_C, int N) {
    int threads_per_block = 256;
    int blocks_per_grid = (N + threads_per_block - 1) / threads_per_block;
    vector_add_cuda<<<blocks_per_grid, threads_per_block>>>(d_A, d_B, d_C, N);
}

OpenAI Triton (programmation Python au niveau des blocs)

Dans Triton, vous programmez du point de vue d'un bloc (tuile) de données directement en Python. Le compilateur Triton gère automatiquement en interne la synchronisation des threads, la coalescence mémoire et la vectorisation.

import triton
import triton.language as tl

@triton.jit
def vector_add_triton(A_ptr, B_ptr, C_ptr, N, BLOCK_SIZE: tl.constexpr):
    # Identify which block of data this program instance handles
    pid = tl.program_id(axis=0)
    
    # Create a vector of offsets for the entire block at once
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    
    # Load chunks of memory into fast SRAM, add them, and store back
    a = tl.load(A_ptr + offsets, mask=mask)
    b = tl.load(B_ptr + offsets, mask=mask)
    tl.store(C_ptr + offsets, a + b, mask=mask)

# Launch directly from PyTorch tensors
def launch_triton(A, B, C):
    N = A.numel()
    grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
    vector_add_triton[grid](A, B, C, N, BLOCK_SIZE=1024)

3. FlashAttention : résoudre le goulot d'étranglement mémoire des Transformers

L'auto-attention standard se heurte à un obstacle majeur lors du traitement d'images haute résolution ou de longues séquences de texte, car la matrice d'attention S = QK^T croît de manière quadratique (N × N) avec la longueur N de la séquence.

• Le goulot d'étranglement de la bande passante mémoire (HBM et SRAM). Les GPU NVIDIA modernes disposent de deux principaux niveaux de mémoire :

• HBM (mémoire à haute bande passante / VRAM) : volumineuse (par exemple, 80 Go), mais relativement lente en lecture et en écriture.
• SRAM (mémoire partagée sur puce) : minuscule (~20 Mo au total pour tous les Streaming Multiprocessors), mais environ 10 à 20 fois plus rapide que la HBM.

Dans l'attention standard, le GPU calcule S = QK^T, écrit l'immense matrice N × N dans la HBM lente, la relit pour appliquer Softmax, réécrit les probabilités dans la HBM, puis les relit pour les multiplier par V. Les Tensor Cores passent l'essentiel de leur temps inactifs à attendre les transferts mémoire.

• Découpage en tuiles (sans jamais matérialiser la matrice N × N). Introduit par Tri Dao et ses collègues (2022), FlashAttention utilise la fusion de noyaux pour maintenir tout le calcul de l'attention dans la SRAM rapide sur puce. Il découpe Q, K et V en petits blocs tenant dans la SRAM, calcule l'attention tuile par tuile et réécrit uniquement la matrice de sortie finale dans la HBM. L'immense matrice intermédiaire N × N n'est jamais enregistrée dans la VRAM.

• L'astuce de Softmax en ligne. Normalement, Softmax exige de connaître la valeur maximale et la somme des exponentielles de toute une ligne avant de normaliser le moindre élément. FlashAttention utilise un algorithme Softmax en ligne qui suit les maxima courants et les facteurs de normalisation bloc par bloc, en remettant mathématiquement à l'échelle les sorties partielles à chaque nouvelle tuile traitée.

• Recalcul pendant la rétropropagation. Comme FlashAttention n'enregistre pas la matrice d'attention N × N pendant la passe avant, il adopte un compromis astucieux lors de l'entraînement : il ne conserve que les petites statistiques de normalisation dans la HBM et recalcule les tuiles d'attention à la volée dans la SRAM rapide pendant la rétropropagation. Effectuer des calculs supplémentaires dans la SRAM se révèle beaucoup plus rapide que lire de grandes matrices depuis la HBM, accélérant le temps d'exécution réel d'un facteur 2 à 4 tout en faisant passer l'utilisation mémoire d'une croissance quadratique O(N²) à une croissance linéaire O(N).