Qu'est-ce que FlashAttention du point de vue de la programmation de kernels?

FlashAttention est un noyau CUDA optimisé pour le calcul du mécanisme d’attention dans les transformateurs, ce qui accélère considérablement les calculs et réduit la consommation de mémoire.

Du point de vue de la programmation de noyaux, FlashAttention implémente le calcul de l’attention en utilisant une gestion efficace de la mémoire et du parallélisme sur le GPU. Les principales caractéristiques :

• Utilisation de blocs de mémoire partagée pour stocker les données intermédiaires, ce qui réduit les accès à la mémoire globale lente.
• Calcul de l’attention dans des threads avec une duplication minimale des données.
• Optimisation du calcul du softmax en tenant compte de la stabilité numérique.
• Traitement de séquences longues sans avoir besoin de charger toutes les données en mémoire en même temps.

Cela permet d’effectuer l’attention plus rapidement et avec une consommation de mémoire moindre par rapport aux implémentations classiques, ce qui est critique pour l’entraînement de grands modèles.

Par exemple, le noyau charge des blocs de clés et de requêtes, calcule des produits scalaires, applique softmax et multiplie par les valeurs, en utilisant la mémoire partagée pour accélérer et minimiser les accès à la mémoire globale.