MHA, MQA et GQA : comment l’attention modifie la consommation de mémoire KV
Traduction automatique
Cet article a été traduit automatiquement depuis la version originale en anglais.
MHA attribue à chaque tête de requête sa propre tête de clé et sa propre tête de valeur. MQA partage une seule tête de clé/valeur entre toutes les têtes de requête. GQA partage une tête de clé/valeur au sein de chaque groupe de têtes de requête. Un nombre réduit de têtes KV diminue le cache d’attention classique et la quantité de données que le décodage y lit.
Pour estimer la taille du cache, vérifiez num_key_value_heads. Le seul nombre de têtes de requête peut conduire à une forte surestimation de sa taille.
Comparer directement le partage des têtes
Supposons huit têtes de requête, avec une dimension par tête, un nombre de couches, une longueur de séquence et une précision du cache identiques :
| Attention | Têtes de requête | Têtes KV | Cache par rapport à MHA |
|---|---|---|---|
| Attention multi-têtes, MHA | 8 | 8 | 100% |
| Attention multi-requêtes, MQA | 8 | 1 | 12.5% |
| Attention à requêtes groupées, GQA | 8 | 2 | 25% |
Une tête KV comprend une tête de clé et une tête de valeur. Les projections des requêtes restent distinctes. L’article de Shazeer sur MQA et l’article d’Ainslie et al. sur GQA définissent ces changements.
Pour Llama 3 70B, 64 têtes de requête partagent huit têtes KV : cela représente huit fois moins de données KV brutes que pour un modèle par ailleurs identique possédant 64 têtes KV. Llama 3.1 405B possède 128 têtes de requête et huit têtes KV, soit une réduction d’un facteur seize selon la même comparaison. Le tableau d’architecture de Meta fournit ces nombres.
Ce que la réduction ne garantit pas
Un KV cache huit fois plus petit ne signifie ni huit fois moins de mémoire GPU totale, ni une génération huit fois plus rapide. Les poids, les activations, l’exécution des kernels et la communication restent nécessaires. Les implémentations distribuées peuvent aussi répliquer les têtes KV entre les partitions.
L’article sur GQA a observé un compromis utile entre qualité et vitesse dans les modèles de la famille T5 testés. Convertir un checkpoint MHA entraîné en GQA constitue un changement d’architecture qui nécessite une adaptation et une évaluation ; ce n’est pas un paramètre d’allocation mémoire.
Vérifiez la configuration d’attention du checkpoint, utilisez le nombre de têtes KV dans la formule du cache et mesurez les performances du serveur avec les longueurs de contexte et le nombre de requêtes simultanées dont vous avez besoin. Incluez la qualité des résultats sur vos tâches dans la comparaison entre checkpoints.
Guide d’ingénierie : GQA et MQA illustre les modes de partage des têtes.