Variantes de l’attention et compression du cache KV

Lors du décodage autorégressif, chaque nouveau jeton porte attention aux clés et aux valeurs des jetons précédents. Mettre ces tenseurs en cache évite de les recalculer, mais la mémoire du cache croît linéairement avec la longueur de séquence et la taille de lot. La lecture du cache peut devenir un coût majeur de bande passante, en particulier avec de longs contextes et de petits lots de décodage.

Les principales approches réduisent l’un de ces trois éléments : le nombre de têtes KV, la largeur de la représentation mise en cache, ou le nombre ou la précision des jetons mis en cache.

Voir l’image en grand

Partez des requêtes bleues. Au centre, deux requêtes partagent une paire clé-valeur ; à droite, toutes partagent la même. Le nombre de requêtes reste constant. À dimensions égales, le cache diminue parce qu’il contient moins de clés et de valeurs.

Ainslie et al., GQA (2023), Figure 2. · CC BY 4.0

MHA et GQA

Soit un modèle ayant $L$ couches, une longueur de séquence $S$, $n_h$ têtes de requête, $n_{kv}$ têtes KV, une dimension de tête $d_h$ et $b$ octets par élément mis en cache. En ignorant les surcoûts d’allocateur et de métadonnées, le cache KV pour une séquence est approximativement

$$2L S n_{kv} d_h b.$$

Le facteur deux tient compte des clés et des valeurs.

• L’attention multi-tête (MHA) utilise normalement $n_{kv}=n_h$.
• L’attention à requête multiple (MQA) utilise une tête KV partagée par toutes les têtes de requête.
• L’attention à requêtes groupées (GQA) se situe entre les deux : plusieurs têtes de requête partagent chaque tête KV.

Toutes les autres dimensions étant fixes, passer de MHA à GQA réduit cette part de la mémoire cache proportionnellement à $n_{kv}/n_h$. Que la qualité, le débit et la prise en charge par les noyaux restent acceptables dépend du modèle entraîné et de la pile de service ; GQA est un choix architectural fait pendant l’entraînement, pas un commutateur sans perte applicable à des poids arbitraires.

Pour un exemple hypothétique avec $L=32$, $S=4096$, $n_h=32$, $n_{kv}=8$, $d_h=128$ et deux octets par élément, le cache occupe $536{,}870{,}912$ octets, soit 512 Mio par séquence. MHA avec 32 têtes KV occuperait 2 Gio ; quatre séquences GQA de même longueur occupent 2 Gio avant surcoûts. Il s’agit du cache, pas de la mémoire totale du modèle.

Quand la réutilisation est-elle exacte ?

Dans un modèle causal en mode évaluation, ajouter un token ne modifie pas les représentations antérieures. Chaque couche conserve ses propres K/V ; seules les requêtes des nouvelles positions sont calculées. La réutilisation exacte exige le même préfixe tokenisé, les mêmes poids (adaptateurs compris), indices de position, règles d’attention et représentation numérique. Modifier un token ancien invalide les états en aval. L’attention bidirectionnelle ne permet généralement pas cette réutilisation après ajout, puisque les anciens états peuvent dépendre du nouveau token.

Avec $S$ tokens en cache et un nouveau segment de $q$ tokens, la ligne de requête $i$ (comptée depuis zéro dans le segment) peut voir jusqu’à la position absolue de clé $S+i$. Un masque triangulaire aligné naïvement en haut à gauche d’une matrice $q\times(S+q)$ cacherait des clés valides. Conserver les indices de position d’origine après éviction : renuméroter les tokens change l’attention positionnelle. Comparer d’abord les logits avec et sans cache sur le même préfixe, puis tester la compression approchée, en tolérant les écarts flottants dus aux formes des noyaux ou à l’ordre des réductions.

La spéculation en arbre étend ce problème aux branches candidates : les isoler et conserver les états KV et comprimés correspondant au chemin accepté.

Attention latente multi-tête

DeepSeek-V2 a introduit l’attention latente multi-tête (MLA). Au lieu de mettre en cache les tenseurs complets de clés et de valeurs par tête, le modèle apprend une représentation latente de dimension plus faible. Des parties du chemin des clés utilisées pour l’information de position rotative restent distinctes, car la rotation dépendante de la position ne peut pas simplement être incorporée dans une projection fixe.

Une comparaison simplifiée de la largeur du cache est :

$$\begin{aligned} \text{Largeur MHA/GQA par jeton et par couche} &\approx 2n_{kv}d_h,\\ \text{Largeur MLA par jeton et par couche} &\approx d_c+d_R, \end{aligned}$$

où $d_c$ est la largeur latente compressée et $d_R$ la largeur découplée des clés rotatives. Le chemin réel de mémoire et de calcul dépend des projections exactes du modèle et de l’implémentation d’inférence. MLA peut réduire fortement le trafic du cache, mais les résultats publiés concernent des modèles entraînés avec cette architecture.

Compresser un cache existant

Des méthodes telles que SnapKV et PyramidKV agissent sur les jetons retenus dans le cache d’un modèle existant au lieu de modifier l’architecture d’attention entraînée.

• SnapKV note les positions du prompt à partir des motifs d’attention observés dans une fenêtre d’observation et conserve un sous-ensemble choisi de clés et de valeurs passées.
• PyramidKV attribue à différentes couches différents budgets de cache, en s’appuyant sur les différences observées dans la façon dont les couches répartissent l’attention sur le prompt.
• La quantification du cache KV stocke les éléments du cache à plus faible précision.
• L’attention fenêtrée ou glissante ne conserve qu’une région récente bornée dans les couches conçues pour la prendre en charge.

Ces techniques ne sont pas interchangeables. L’éviction de jetons peut perdre un passage bref mais décisif ; la quantification peut modifier les scores d’attention ; les dispositions personnalisées peuvent exiger des noyaux spécialisés. Les résultats d’exactitude ou de perplexité au niveau d’un article s’appliquent aux modèles, tâches, budgets et implémentations évalués, et non à toute charge de travail à long contexte.

Ce qu’il faut mesurer

Pour un déploiement concret, consignez :

1. le modèle exact, l’architecture d’attention et le dtype du cache ;
2. la longueur du prompt, la longueur générée et le nombre de séquences simultanées ;
3. la mémoire GPU et hôte de pointe ;
4. le temps de préremplissage, le temps jusqu’au premier jeton et le débit de décodage ;
5. la réussite de la tâche sur des cas à long contexte exigeant des informations situées à différentes positions ;
6. le comportement après éviction du cache, quantification ou déport vers le CPU.

Un modèle qui annonce une longue fenêtre de contexte peut rester inutilisable sur le matériel disponible, tandis qu’une méthode économisant le cache qui tient en mémoire peut encore échouer à la tâche. L’adéquation mémoire et la qualité des réponses préservée nécessitent des tests distincts.

Mes lectures · Les notes ouvertes sont mémorisées dans ce navigateur.