Les noyaux d’attention rapides et la fusion d’opérateurs réorganisent le calcul des transformers pour limiter les déplacements de données et les lancements de noyaux sur l’accélérateur. Ces optimisations laissent en principe inchangés les poids du modèle et sa fonction mathématique, même si l’ordre des opérations et la précision peuvent produire de petites différences numériques. Le tutoriel publié par Machine Learning Mastery détaille ces techniques, avec des exemples de code.
- Les noyaux d’attention rapides et la fusion d’opérateurs réduisent les déplacements de données et les lancements de noyaux sur l’accélérateur.
- FlashAttention calcule l’attention par blocs et n’stocke jamais la matrice complète des scores.
- Le cache paginé limite la fragmentation mémoire, car il découpe le stockage en blocs de taille fixe.
- torch.compile fusionne les opérations avec TorchInductor, et le mode « reduce-overhead » rejoue le calcul avec des CUDA Graphs.
Une couche de transformer repose sur des opérations courantes : projections linéaires, attention, normalisation, fonctions d’activation et additions résiduelles. Une implémentation directe lance un noyau de calcul distinct, appelé « kernel », pour chaque opération, et écrit chaque tenseur intermédiaire dans la mémoire de l’accélérateur. Le calcul reste juste, mais le temps d’exécution peut être consacré davantage au déplacement des valeurs et au lancement des noyaux qu’aux calculs utiles.
Fusion d’opérateurs
Le bloc SwiGLU, utilisé par de nombreux modèles à décodeur seul, illustre ce problème. Une implémentation directe lance des multiplications de matrices pour les projections, un noyau pour la fonction SiLU, un autre pour la multiplication élément par élément, puis une dernière multiplication de matrices. À chaque frontière, un tenseur intermédiaire peut être écrit dans la mémoire à haute bande passante (HBM) puis relu. La fusion consiste à combiner ces étapes dans un seul noyau. Elle fait gagner du temps lorsque les déplacements de mémoire et les lancements pèsent lourd, et beaucoup moins lorsqu’une grande multiplication de matrices domine déjà l’exécution. Le choix du noyau reste empirique : il faut profiler l’ensemble de la charge de travail.
D’autres fusions sont courantes dans les transformers, comme l’ajout d’une connexion résiduelle suivi d’une normalisation RMSNorm ou LayerNorm, ou la combinaison de la déquantification avec une multiplication de matrices. La fusion a toutefois ses limites : des opérations aux organisations de données incompatibles ou un état intermédiaire très volumineux se combinent mal. Un noyau trop gros peut augmenter la pression sur les registres, réduire le taux d’occupation et devenir plus lent que plusieurs noyaux bien réglés.
Attention de type FlashAttention et attention paginée
Une attention classique calcule une matrice de scores de taille T² par tête, T étant la longueur de la séquence. Sur un GPU, la produire, l’écrire dans la mémoire globale puis la relire consomme beaucoup de bande passante, ce qui peut devenir un goulot d’étranglement lors des longs préremplissages (« prefill »). FlashAttention traite les requêtes, les clés et les valeurs par blocs. Pour chacun, il calcule une partie des scores, met à jour les statistiques du softmax et accumule la sortie correspondante dans la mémoire située sur la puce, sans stocker la matrice complète. PyTorch expose cette optimisation via la fonction scaled_dot_product_attention(). Le tutoriel précise qu’imposer un backend en dur peut faire échouer le code si aucun backend ne prend en charge les entrées ou le matériel.
Le cache KV, qui conserve les clés et valeurs déjà calculées, grandit d’un token à la fois, pour une longueur finale imprévisible. Une allocation contiguë par requête laisse donc de l’espace réservé et inutilisé. Le cache paginé découpe le stockage en blocs de taille fixe et s’appuie sur une table de blocs qui relie les positions logiques aux blocs physiques. Le noyau d’attention doit alors traduire chaque position à travers cette table. La pagination ne réduit pas le nombre d’éléments logiques, mais elle limite la fragmentation et la capacité réservée.
Dans l’exemple fourni, vLLM utilise le modèle Qwen/Qwen2.5-1.5B-Instruct avec des blocs de 16 tokens (block_size=16) et le backend « TRITON_ATTN ». Selon le tutoriel, vLLM vérifie la compatibilité du backend avec le modèle et le matériel lors de l’initialisation, et signale une erreur en cas d’incompatibilité.
Fusion et CUDA Graphs avec torch.compile
La fonction torch.compile() capture une suite d’opérations avec TorchDynamo, puis la compile avec TorchInductor, son backend par défaut. TorchInductor peut générer un seul noyau Triton pour des opérations compatibles que le mode standard lancerait séparément. Les interruptions de graphe, dues par exemple à un contrôle de flux dépendant des données ou à des conversions de tenseurs en scalaires Python, limitent les zones que le compilateur peut optimiser.
Le mode « reduce-overhead » demande à TorchInductor d’utiliser des CUDA Graphs lorsque le graphe capturé y est éligible. Un CUDA Graph enregistre la suite des opérations GPU et la rejoue avec une seule soumission au processeur. Il ne crée pas de fusion supplémentaire, il complète celle de TorchInductor. Il exige des adresses mémoire stables et des formes compatibles, et il peut consommer davantage de mémoire. Le tutoriel recommande de vérifier les résultats avec des tolérances adaptées au type de données, de profiler l’exécution et de mesurer la performance de bout en bout.
