Prérequis et compatibilité

• Version CUDA et PyTorch : Assurez-vous que votre version de PyTorch est compatible avec la version de FlashAttention que vous installez. La version 2.0 de FlashAttention requiert des GPU NVIDIA de la série 30 ou supérieure (les cartes comme la 2080 ne sont pas supportées).
• Version Python : Les versions Python 3.10 et 3.11 sont généralement bien supportées.

Installation sous Linux

L'installation directe via pip peut parfois échouer ou être très longue. Une méthode recommandée consiste à télécharger le fichier Wheel (.whl) précompilé correspondant à votre configuration.

Téléchargement du fichier Wheel

Vous pouvez trouver les versions appropriées sur la page des releases de FlashAttention. Choisissez la version qui correspond à votre version de PyTorch, CUDA et Python.

Exemple de commande pour installer un fichier Wheel spécifique (remplacez le nom du fichier par celui que vous avez téléchargé) :

pip install flash_attn-2.7.4.post1+cu12torch2.5cxx11abiFALSE-cp311-cp311-linux_x86_64.whl
 

Configuration de PyTorch avec des miroirs

Pour accélérer le téléchargement de PyTorch, vous pouvez utiliser des miroirs comme ceux d'Alibaba :

pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 \
   --index-url https://mirrors.aliyun.com/pytorch-wheels/cu124/ \
   --extra-index-url https://mirrors.aliyun.com/pypi/simple/ \
   --no-cache-dir
 

Compatibilité des versions

Il est crucial de faire correspondre les versions de FlashAttention et de PyTorch :

• v2.8.2 supporte jusqu'à PyTorch 2.7.
• v2.8.0.post2 supporte jusqu'à PyTorch 2.7.
• v2.8.1 supporte PyTorch 2.8 et 2.9, mais requiert Python 3.12.
• v2.8.3 supporte PyTorch 2.9.

Notez que certaines versions de bibliothèques comme vLLM peuvent avoir des contraintes sur la version de FlashAttention supportée (par exemple, vLLM 0.11 pourrait ne pas supporter v2.8.3).

Vérification de l'installation

Une fois l'installation terminée, vous pouvez vérifier qu'elle a réussi.

Test d'importation

Exécutez la commande suivante dans votre terminal Python :

import flash_attn
 

Test d'utilisation

Ce script crée des tenseurs aléatoires et utilise la fonction FlashAttention pour un test rapide.

import torch
 import flash_attn
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 batch_size = 2
 seq_len = 16
 embed_dim = 128
 
 # Génération des tenseurs Q, K, V
 q = torch.randn(batch_size, seq_len, embed_dim, device=device, dtype=torch.float16)
 k = torch.randn(batch_size, seq_len, embed_dim, device=device, dtype=torch.float16)
 v = torch.randn(batch_size, seq_len, embed_dim, device=device, dtype=torch.float16)
 
 # Concaténation pour la fonction qkvpacked
 qkv_feats = torch.cat([q, k, v], dim=-1).half()
 
 # Appel de la fonction d'attention
 out = flash_attn.flash_attn_qkvpacked_func(qkv_feats)
 
 print('Output shape:', out.shape)
 

Si le script s'exécute sans erreur et affiche la forme de sortie attendue, l'installation est réussie.

Fonctionnement de FlashAttention

FlashAttention est une implémentation optimisée du mécanisme d'attention. Elle améliore l'efficacité en termes de mémoire et de calcul, particulièrement pour les longues séquences, grâce à des techniques telles que :

• Calcul par blocs : L'attention est calculée sur des sous-ensembles de la séquence, réduisant ainsi la consommation mémoire maximale.
• Rétrocalcul : Certaines valeurs intermédiaires sont recalculées à la demande plutôt que d'être stockées, économisant de la mémoire.
• Optimisation des accès mémoire : Réduction des transferts de données entre les différents niveaux de mémoire de l'unité de calcul.

Cela permet de traiter des séquences beaucoup plus longues avec des ressources matérielles plus limitées par rapport aux implémentations d'attention standard.