Cet article démontre comment profiler les mécanismes d'attention dans PyTorch en utilisant les traces du profileur pour identifier les goulots d'étranglement de performance et les opportunités d'optimisation. Il compare les implémentations d'attention naïves à l'Attention par Produit Scalé (SDPA) intégrée de PyTorch pour illustrer le comportement des noyaux.
- L'attention causale naïve implique plusieurs noyaux, notamment matmul, mise à l'échelle, masquage, softmax et une copie mémoire inattendue causée par des opérations hors place.
- Le remplacement de masked_fill par masked_fill_ in-place supprime le noyau de copie mémoire inutile de la trace GPU.
- La fonction SDPA de PyTorch dispatche automatiquement vers le backend disponible le plus rapide (tel que Flash Attention ou cuDNN) en fonction des types d'entrée et du matériel.
- Le profilage des backends SDPA individuels révèle des signatures de noyaux distinctes, permettant aux développeurs de vérifier quelle implémentation est active.
Comprendre ces traces de profileur aide les utilisateurs à optimiser les modèles de transformeurs en éliminant les opérations mémoire redondantes et en exploitant les backends d'attention efficaces de PyTorch.