Grouped Query Attention (GQA)

Le Grouped Query Attention (GQA) est une conception d'attention dans laquelle plusieurs têtes de requête partagent un ensemble plus restreint de têtes de clé et de valeur. Il préserve les perspectives diversifiées des requêtes multi-têtes tout en réduisant la mémoire et le mouvement de données requis pour les clés et les valeurs. Le GQA est particulièrement précieux lors de l'inférence autorégressive, où un modèle génère un token à la fois et lit de manière répétée les états d'attention stockés.

Comment fonctionne le GQA#

Au sein d'un mécanisme d'attention, chaque token est projeté en trois représentations :

• Requête (Query) : Décrit les informations que le token actuel recherche.
• Clé (Key) : Décrit ce que représente chaque token disponible.
• Valeur (Value) : Contient les informations récupérées lorsqu'une requête correspond à une clé.

Dans le self-attention conventionnel, ces projections sont divisées en têtes afin que le modèle puisse apprendre différentes relations en parallèle. Le GQA conserve de nombreuses têtes de requête mais assigne des groupes d'entre elles à des têtes de clé et de valeur partagées. Par exemple, une couche d'attention peut utiliser huit têtes de requête et deux têtes de clé-valeur. Chaque tête de clé-valeur dessert alors quatre têtes de requête.

Le calcul de l'attention en soi reste une attention par produit scalaire mise à l'échelle (scaled dot-product attention). Ce qui change, c'est l'agencement des têtes et la quantité de données de clé-valeur produites, stockées et lues. Des frameworks tels que le PyTorch scaled dot-product attention exigent donc que le nombre de têtes de requête soit divisible par le nombre de têtes de clé-valeur lorsque le GQA est activé. (docs.pytorch.org)

GQA vs Multi-Head et Multi-Query Attention#

Le GQA se situe entre le multi-head attention et le multi-query attention :

• Multi-head attention : Chaque tête de requête possède ses propres têtes de clé et de valeur. Cela offre une indépendance maximale des têtes mais crée le besoin en mémoire clé-valeur le plus élevé.
• Grouped query attention : Plusieurs têtes de requête partagent chaque tête de clé-valeur. Cela équilibre la capacité de représentation et l'efficacité mémoire.
• Multi-query attention : Toutes les têtes de requête partagent une tête de clé et une tête de valeur. Cela minimise le stockage clé-valeur mais offre une moindre diversité de clés et de valeurs.

Si une couche possède 32 têtes de requête, le multi-head attention peut également utiliser 32 têtes de clé-valeur, le GQA peut en utiliser 8, et le multi-query attention en utilise 1. Le GQA ne remplace pas l'architecture Transformer dans son ensemble ; il s'agit d'une configuration possible au sein d'une couche d'attention Transformer. La documentation sur le multi-head, multi-query et grouped-query attention de NVIDIA décrit ces variantes comme différant principalement par le nombre de têtes de clé-valeur qui desservent les têtes de requête. (nvidia.github.io)

Pourquoi le GQA est important#

Lors de la génération autorégressive, un modèle stocke les clés et valeurs précédentes dans un KV cache. Avec le multi-head attention standard, chaque tête d'attention contribue à des clés et valeurs mises en cache séparément. Des séquences longues, de grands lots (batches) et de nombreuses couches de modèle peuvent ainsi consommer une mémoire GPU considérable.

Parce que le GQA utilise moins de têtes de clé-valeur, son cache est plus petit proportionnellement à la réduction de ces têtes. Cela peut améliorer :

• Capacité mémoire : Des prompts plus longs 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é-valeur doivent être lues pour chaque token généré.
• Latence d'inférence : Un trafic mémoire réduit peut raccourcir le temps de traitement de token à token.
• Évolutivité de la fenêtre de contexte : Le traitement de séquences longues devient plus pratique.

Ces gains dépendent du matériel, de la longueur de la séquence, de la taille du lot et de la prise en charge des noyaux (kernels). Les environnements de production ont toujours besoin d'une allocation de cache efficace ; le système de cache KV TensorRT-LLM de NVIDIA, par exemple, associe la prise en charge du GQA à un cache basé sur des blocs, à la réutilisation et au déchargement (offloading). (nvidia.github.io)

Applications concrètes#

Assistants à contexte long : Un assistant de code peut traiter des milliers de tokens de code source avant de générer une réponse. Le GQA réduit l'état de clé-valeur mis en cache associé à cet historique, permettant à 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 images et vidéos : Un modèle vision-langage peut représenter des patchs d'image ou des trames vidéo sous forme de longues séquences de tokens. Le GQA réduit la pression du cache d'attention après l'entrée de ces tokens visuels dans le décodeur linguistique, libérant potentiellement plus de mémoire pour des images, des trames ou des requêtes supplémentaires. Cela diffère d'un Vision Transformer standard, où l'attention peut traiter tous les patchs d'image en parallèle sans mise en cache KV autorégressive.

Mise en œuvre pratique et compromis#

PyTorch expose le GQA via enable_gqa=True. Cet exemple CUDA documenté utilise huit têtes de requête et deux têtes de clé-valeur 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ête, tandis que les clés et les valeurs n'utilisent que deux têtes. Le sélecteur de backend d'attention PyTorch peut contrôler quel noyau pris en charge exécute l'opération, tandis que le guide des blocs de construction Transformer PyTorch fournit un contexte de mise en œuvre plus large. (docs.pytorch.org)

Le GQA doit être choisi lors de la conception ou de l'adaptation d'un modèle ; il ne s'agit généralement pas d'un paramètre commutable au moment de l'inférence pour un point de contrôle (checkpoint) existant. Les développeurs doivent vérifier la prise en charge par le framework, la divisibilité des 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 flux de travail complémentaires tels que l'intégration TensorRT d'Ultralytics et le mode benchmark d'Ultralytics aident à mesurer si les optimisations architecturales et d'exécution apportent des améliorations significatives sur le matériel cible.