Résumé technique : GQLA – Attention Latente par Groupes de Requêtes pour le Décodage Adaptatif au Matériel des Modèles de Langage de Grande Taille

1. Énoncé du problème

Le décodage autoregressif moderne des modèles de langage de grande taille (LLM) est fondamentalement limité par le trafic du cache Key-Value (KV), car chaque token généré nécessite la lecture de l'historique complet des clés et valeurs mises en cache depuis la mémoire hors puce. Bien que des architectures récentes comme l'Attention Latente Multi-têtes (MLA), utilisée dans DeepSeek-V2/V3, atteignent une réduction de l'état de l'art du cache KV en compressant conjointement les clés et les valeurs dans un latent de faible rang, elles souffrent de trois inconvénients matériels couplés lors du déploiement sur divers matériels d'inférence :

1. Couplage matériel : Les poids entraînés de la MLA n'exposent qu'un seul chemin de décodage (une forme Multi-Query Attention ou MQA « absorbée »). Ce chemin est parfaitement ajusté au rapport calcul/bande passante des GPU haut de gamme comme le NVIDIA H100 (crête $\approx$ 295 FLOPs/byte). Cependant, sur les GPU restreints à l'exportation ou grand public comme le H20 (qui conserve la bande passante mais réduit le calcul d'un facteur $\sim$7, abaissant la crête à $\approx$ 37 FLOPs/byte), ce chemin devient gravement limité par le calcul, entraînant un faible débit.
2. Perte de parallélisme tensoriel (TP) : La forme MQA absorbée canalise toutes les têtes de requêtes vers un seul latent KV partagé. Cette structure empêche le parallélisme tensoriel selon l'axe des têtes, car chaque dispositif doit répliquer l'intégralité du latent, limitant ainsi l'évolutivité.
3. Gain nul de la prédiction multi-token (MTP) : La MTP, qui génère plusieurs tokens par étape pour augmenter l'intensité arithmétique, repousse le décodage déjà limité par le calcul du H20 encore plus loin au-delà de sa crête d'efficacité, produisant aucun gain de débit. Inversement, sur le H100, la MTP dépasse la crête, réduisant l'efficacité du gain.

2. Méthodologie : Attention Latente par Groupes de Requêtes (GQLA)

Les auteurs proposent la GQLA, une modification architecturale minimale de la MLA qui préserve la compression latente conjointe de faible rang tout en exposant deux chemins de décodage algébriquement équivalents sur les mêmes paramètres entraînés.

Architecture

La GQLA indexe les matrices de projection ascendante ($W_{UK}, W_{UV}$) par $g$ groupes plutôt que de les répliquer sur toutes les $h_q$ têtes de requêtes. Cette conception permet deux chemins d'exécution distincts :

1. Chemin MQA-Absorbé (Favorable à la mémoire) :

   • Mécanisme : Identique à la MLA. Les projections ascendantes sont absorbées dans les projections de requête et de sortie. Toutes les $h_q$ têtes s'attendent directement au cache latent compact ($cKV$).
   • Cache : Stocke uniquement le latent ($r_{kv}$) et une clé RoPE partagée ($d^R_h$).
   • Matériel cible : Optimisé pour le matériel de classe H100 où l'intensité arithmétique ($\approx$ 242 FLOPs/byte pour $s_q=1$) se situe juste en dessous de la crête mémoire-calcul.

2. Chemin GQA (Favorable au calcul) :

   • Mécanisme : Le latent est étendu en $g$ groupes de clés/valeurs. Le décodage exécute une Attention par Groupes de Requêtes (GQA) standard sans expansion latente à chaque étape.
   • Cache : Stocke les clés et valeurs étendues par groupe ($2gd_h$) plus la clé RoPE partagée.
   • Matériel cible : Optimisé pour le matériel de classe H20. En ajustant le nombre de groupes ($g$) et le nombre de tokens de requête par étape ($s_q$), l'intensité arithmétique peut être calée sur la crête du H20 ($\approx$ 37 FLOPs/byte).
   • Évolutivité : Prend en charge jusqu'à un parallélisme tensoriel à 8 voies sans redondance selon l'axe des groupes.

Sélection du chemin

L'exécution sélectionne le chemin en fonction du matériel cible sans réentraînement ni noyaux personnalisés :

• H100 : Utilise le chemin MQA-absorbé ($s_q=1$).
• H20 : Utilise le chemin GQA (par exemple, $g=8, s_q=2$ ou $g=4, s_q=1$).
• MTP : Sur le H20, le chemin GQA avec $s_q=2$ récupère des gains de débit proches de l'idéal, alors que la MLA ne gagne rien.

TransGQLA : Conversion à partir de modèles pré-entraînés

Pour éviter un pré-entraînement à partir de zéro, les auteurs introduisent TransGQLA, une extension du pipeline TransMLA. Il convertit un point de contrôle GQA pré-entraîné en un modèle GQLA via une modification ciblée unique :

• Fusion de têtes : Au lieu de fusionner les têtes groupées en un seul latent et de répliquer les projections ascendantes sur toutes les têtes de requêtes (comme dans TransMLA), TransGQLA conserve les projections ascendantes indexées par groupe.
• Résultat : Le modèle conserve structurellement le chemin d'exécution GQA tout en permettant le chemin MQA-absorbé via l'opération d'absorption standard. Cela préserve le parallélisme tensoriel selon l'axe des groupes.

GQLA Sparse

La conception s'étend à l'attention sparse fine. Étant donné que la configuration canonique GQLA ($h_q/g = 16$) correspond à la dimension de tuile MMA des Tensor-Cores ($m=16$), la GQLA sparse peut préserver le chemin GQA sur le matériel de classe H20. En revanche, la MLA sparse est structurellement verrouillée sur le chemin MQA absorbé en raison des contraintes de tuiles, héritant des mêmes problèmes de surcharge de calcul.

3. Contributions clés

1. Identification des inconvénients de la MLA : L'article identifie que la conception mono-chemin de la MLA la couple aux ratios de la classe H100, renonce au parallélisme tensoriel selon l'axe des têtes et produit des gains MTP nuls sur les GPU grand public.
2. Architecture GQLA : Une variante minimale de la MLA qui expose deux chemins algébriquement équivalents (MQA-absorbé et GQA) sur les mêmes poids. Une seule configuration ($h_q=128, g=8$) avec une tête MTP fixe les diagrammes de toit (rooflines) du H100 et du H20 simultanément.
3. Pipeline TransGQLA : Une méthode de conversion qui transforme les points de contrôle GQA pré-entraînés en modèles GQLA avec des changements architecturaux minimaux, conservant le parallélisme tensoriel et permettant une transition fluide vers la forme latente.
4. Analyse Roofline : Une vérification théorique montrant que les poids GQLA peuvent atteindre un débit proche du pic sur le H100 et le H20 en changeant de chemin, alors que la MLA est inefficace sur le H20.

4. Résultats

Les auteurs ont évalué TransGQLA sur le modèle LLaMA-3-8B :

• Préservation des capacités : La conversion des poids GQA vers la forme latente GQLA (sans entraînement supplémentaire) a entraîné une perte d'environ 9,7 points en moyenne sur les benchmarks de raisonnement de bon sens (MMLU, ARC, etc.) par rapport au modèle pré-entraîné original de 15T tokens.
• Compression du cache : Sur le chemin MQA-absorbé, le cache KV par token a été compressé à 28,125 % de la référence GQA originale (une réduction de 71,875 %).
• Projections de débit :
  • H100 : Le chemin MQA-absorbé atteint 354K tokens/seconde.
  • H20 : Le chemin GQA (avec $g=8, s_q=2$) atteint 221K tokens/seconde, soit une amélioration de 3,4x par rapport à la MLA sur le même dispositif (qui est bloqué à 65K tokens/seconde en raison d'une limitation par le calcul).
• Pré-entraînement continu : Les auteurs prévoient que le pré-entraînement continu (suivant la trajectoire TransMLA) permettra de récupérer les capacités perdues jusqu'à 0,5 point du modèle original après 30 milliards de tokens, tout en maintenant la compression du cache KV.

5. Importance et revendications

L'article affirme que la GQLA offre un principe de conception pratique et agnostique au matériel pour l'inférence LLM. En exposant plusieurs chemins de décodage algébriquement équivalents sur un seul ensemble de poids entraînés, la GQLA permet à un modèle de s'adapter à différentes contraintes matérielles (riche en calcul vs riche en mémoire) au moment du déploiement sans réentraînement ni noyaux personnalisés.

Les auteurs positionnent la GQLA comme une solution au problème de « couplage matériel », permettant un déploiement efficace à la fois sur des accélérateurs phares (H100) et sur du matériel grand public/restreint à l'exportation (H20). Ils soulignent que cette approche prend en charge la prédiction multi-token et le parallélisme tensoriel sur du matériel où la MLA échoue à le faire efficacement.

Limites reconnues :

• Les affirmations matérielles sont basées sur la modélisation Roofline ; des benchmarks complets au niveau des noyaux sur du matériel H20/H100 réel sont en attente.
• Les résultats de pré-entraînement continu pour TransGQLA sont actuellement en cours et pas encore finalisés dans l'article.
• L'évaluation est actuellement limitée à l'épine dorsale LLaMA-3-8B et aux benchmarks de raisonnement de bon sens.
• La GQLA sparse est décrite structurellement mais n'est pas encore évaluée par rapport aux méthodes d'attention sparse existantes.