Grouped Query Attention (GQA)

L’attention à requêtes groupées (GQA) est une conception de l’attention dans laquelle plusieurs têtes de requêtes partagent un ensemble plus restreint de têtes de clés et de valeurs. Elle préserve la diversité des perspectives des requêtes multi-têtes tout en réduisant la mémoire et les déplacements de données nécessaires pour les clés et les valeurs. La GQA est particulièrement utile lors de l’inférence autorégressive, lorsqu’un modèle génère un token à la fois et relit régulièrement les états d’attention stockés.

Fonctionnement de la GQA#

Dans un mécanisme d’attention, chaque token est projeté en trois représentations :

• Requête : décrit les informations recherchées par le token actuel.
• Clé : décrit ce que représente chaque token disponible.
• Valeur : contient les informations récupérées lorsqu’une requête correspond à une clé.

Dans l’auto-attention conventionnelle, ces projections sont réparties entre des têtes afin que le modèle puisse apprendre différentes relations en parallèle. La GQA conserve de nombreuses têtes de requêtes, mais les répartit en groupes associés à des têtes de clés et de valeurs partagées. Par exemple, une couche d’attention peut utiliser huit têtes de requêtes et deux têtes de clés-valeurs. Chaque tête de clés-valeurs sert alors quatre têtes de requêtes.

Le calcul de l’attention lui-même reste une attention à produit scalaire mis à l’échelle. Ce qui change, c’est la disposition des têtes et la quantité de données de clés-valeurs produites, stockées et lues. Les frameworks tels que l’attention à produit scalaire mis à l’échelle de PyTorch exigent donc que le nombre de têtes de requêtes soit divisible par le nombre de têtes de clés-valeurs lorsque la GQA est activée. (docs.pytorch.org)

GQA, attention multi-têtes et attention multi-requêtes#

La GQA se situe entre l’attention multi-têtes et l’attention multi-requêtes :

• Attention multi-têtes : chaque tête de requêtes possède ses propres têtes de clés et de valeurs. Cela offre une indépendance maximale entre les têtes, mais crée le besoin en mémoire de clés-valeurs le plus élevé.
• Attention à requêtes groupées : plusieurs têtes de requêtes partagent chaque tête de clés-valeurs. Elle équilibre la capacité de représentation et l’efficacité mémoire.
• Attention multi-requêtes : toutes les têtes de requêtes partagent une seule tête de clés et une seule tête de valeurs. Cela réduit au minimum le stockage des clés-valeurs, mais offre moins de diversité pour les clés-valeurs.

Si une couche possède 32 têtes de requêtes, l’attention multi-têtes peut également utiliser 32 têtes de clés-valeurs, la GQA peut en utiliser 8 et l’attention multi-requêtes en utilise 1. La GQA ne remplace pas l’architecture Transformer au sens large ; il s’agit d’une configuration possible au sein d’une couche d’attention Transformer. La documentation de NVIDIA sur l’attention multi-têtes, multi-requêtes et à requêtes groupées décrit ces variantes comme différant principalement par le nombre de têtes de clés-valeurs qui servent les têtes de requêtes. (nvidia.github.io)

Pourquoi la GQA est importante#

Lors de la génération autorégressive, un modèle stocke les clés et les valeurs précédentes dans un cache KV. Avec l’attention multi-têtes standard, chaque tête d’attention contribue des clés et des valeurs mises en cache distinctes. Les longues séquences, les lots volumineux et les nombreuses couches du modèle peuvent donc consommer une quantité importante de mémoire GPU.

Comme la GQA utilise moins de têtes de clés-valeurs, son cache est réduit proportionnellement à la diminution du nombre de ces têtes. Cela peut améliorer :

• Capacité mémoire : des invites plus longues ou davantage de requêtes simultanées peuvent tenir dans la mémoire disponible.
• Débit de génération : moins de données de clés-valeurs doivent être lues pour chaque token généré.
• Latence d’inférence : la réduction du trafic mémoire peut raccourcir le temps de traitement entre deux tokens.
• Évolutivité de la fenêtre de contexte : la gestion de longues séquences devient plus pratique.

Ces gains dépendent du matériel, de la longueur des séquences, de la taille des lots et de la prise en charge des noyaux. Les runtimes de production doivent toujours assurer une allocation efficace du cache ; le système de cache KV de TensorRT-LLM de NVIDIA, par exemple, associe la prise en charge de la GQA à une mise en cache par blocs, à la réutilisation et au déchargement. (nvidia.github.io)

Applications concrètes#

Assistants à long contexte : un assistant de programmation peut traiter des milliers de tokens de code source avant de générer une réponse. La GQA réduit l’état de clés-valeurs mis en cache associé à cet historique, ce qui permet à un serveur de gérer des fichiers plus longs ou davantage d’utilisateurs simultanés sans augmenter proportionnellement la mémoire GPU.

Assistants multimodaux pour les images et les vidéos : un modèle vision-langage peut représenter des patches d’image ou des images vidéo sous forme de longues séquences de tokens. La GQA réduit la pression exercée sur le cache d’attention une fois ces tokens visuels transmis au décodeur de langage, ce qui peut libérer davantage de mémoire pour des images, des trames ou des requêtes supplémentaires. Cela diffère d’un Vision Transformer standard, dans lequel l’attention peut traiter tous les patches d’image en parallèle sans mise en cache KV autorégressive.

Mise en œuvre pratique et compromis#

PyTorch expose la GQA via enable_gqa=True. Cet exemple CUDA documenté utilise huit têtes de requêtes et deux têtes de clés-valeurs partagées :

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

assert torch.cuda.is_available(), "A CUDA device is required."

query = torch.randn(1, 8, 16, 32, device="cuda")
key = torch.randn(1, 2, 16, 32, device="cuda")
value = torch.randn(1, 2, 16, 32, device="cuda")

# Four query heads share each key-value head.
with sdpa_kernel(SDPBackend.MATH):
    output = F.scaled_dot_product_attention(
        query,
        key,
        value,
        enable_gqa=True,
    )

print(output.shape)

La sortie conserve huit têtes de requêtes, tandis que les clés et les valeurs n’utilisent que deux têtes. Le sélecteur de backend d’attention de PyTorch peut contrôler quel noyau pris en charge effectue l’opération, tandis que le guide des composants de base Transformer de PyTorch fournit un contexte plus large sur la mise en œuvre. (docs.pytorch.org)

La GQA doit être choisie lors de la conception ou de l’adaptation d’un modèle ; il ne s’agit généralement pas d’un commutateur activable au moment de l’inférence pour un checkpoint existant. Les développeurs doivent vérifier la prise en charge par le framework, la divisibilité du nombre de têtes, le comportement numérique, l’utilisation de la mémoire et la qualité de la tâche. Pour le déploiement en vision par ordinateur, des workflows complémentaires tels que l’intégration TensorRT d’Ultralytics et le mode benchmark d’Ultralytics aident à mesurer si les optimisations de l’architecture et du runtime apportent des améliorations significatives sur le matériel cible.