Quelle mémoire consomme un modèle : comment GQA, fenêtres glissantes et MLA transforment le cache KV

Situation classique : vous avez monté un modèle prêt à l'emploi, les poids tiennent sans problème dans la mémoire du GPU, et pourtant le serveur plante en out-of-memory dès la dixième requête parallèle. Le coupable n'est pas la taille du modèle, mais le cache KV. Son volume dépend de la conception de la couche d'attention, et vous ne pouvez plus la choisir : la décision a été prise pendant l'entraînement.

Le checkpoint fige l'architecture

La config et les poids déterminent le nombre de têtes de requêtes et de têtes key-value, la taille de tête, le schéma d'attention locale et les dimensions de la représentation latente. Remplacer MHA par GQA après l'entraînement revient à modifier le modèle, et il faudra sans doute récupérer la qualité par un nouvel entraînement. Remplacer l'attention complète par une fenêtre glissante revient à supprimer un contexte que le modèle a l'habitude de voir.

La pratique est donc la suivante : ouvrir la config (dans Hugging Face, c'est AutoConfig) et examiner le nombre de couches, num_attention_heads, num_key_value_heads, head_dim et les champs relatifs à la sliding window. Une réserve : le champ de la fenêtre peut figurer dans la config tout en étant désactivé par un flag voisin, donc vérifiez à la fois la valeur et le flag. Certains modèles alternent couches locales et couches globales, et la vérité se trouve toujours dans le code de l'implémentation.

Calculer le cache sur un coin de table

La mémoire persistante vit entre les étapes de décodage, et c'est elle qui détermine combien de requêtes vous servez simultanément. Les buffers temporaires, comme les tuiles de scores d'attention, peuvent être libérés après la couche. Pour un cache standard, la formule est : 2 × octets par élément × couches × tokens × nombre de têtes KV × taille de tête. Le 2 correspond aux clés plus les valeurs. Pour un batch, on multiplie par sa taille.

Prenons un modèle théorique : 32 couches, 32 têtes, head_dim 128, fp16, contexte de 8192 tokens. Avec du MHA classique, une requête occupe 4 Gio. Si le modèle utilise GQA avec 8 têtes KV, on passe à 1 Gio. Avec MQA et une seule tête KV, seulement 128 Mio. L'écart entre « 4 requêtes tiennent » et « 16 tiennent » apparaît sans changer le nombre de paramètres.

Les schémas mixtes se calculent couche par couche. Supposons que, pour ce même modèle GQA de 32 couches, 24 soient locales avec une fenêtre de 4096 et 8 complètes, avec un contexte de 32 768 tokens. L'attention complète partout donnerait 4 Gio, contre environ 1,4 Gio pour le schéma mixte. Mais le cache total n'est pas borné pour autant : chaque nouveau token ajoute de la mémoire dans les couches complètes. Il croît plus lentement, mais il ne cesse pas de croître.

Quel type d'attention

• MHA : chaque tête de requêtes a sa propre paire clé-valeur. C'est le cache le plus coûteux, et en décodage aussi la lecture mémoire la plus lourde.
• MQA et GQA : les têtes de requêtes se partagent des têtes KV communes. Le cache et les paramètres des projections K/V diminuent exactement dans la proportion où le nombre de têtes KV a baissé.
• Fenêtre glissante : la couche ne retient que les W derniers tokens, et le cache peut tenir dans un buffer circulaire de taille fixe. Les positions se calculent alors d'après le numéro logique du token, et non d'après l'emplacement dans le buffer.
• MLA : au lieu des clés et valeurs complètes, on stocke un vecteur latent compressé. Le cache est plus petit, mais le décodage demande plus de calcul. C'est généralement avantageux, car le décodage est limité par la bande passante mémoire, et non par l'arithmétique.

On ne peut pas injecter MLA dans la formule MHA via num_key_value_heads : les tenseurs mis en cache sont différents. Calculez d'après ce que l'implémentation stocke réellement, et utilisez des kernels spécialisés.

Ce que le runtime peut changer

Il est utile de distinguer quatre questions : qui regarde qui (attention complète ou fenêtrée), combien de représentations KV sont apprises (MHA/GQA/MQA), ce qui est mis en cache (tenseurs KV ou latent) et comment tout cela s'exécute. Le checkpoint décide des trois premières, le moteur de la quatrième. FlashAttention calcule l'attention exacte par tuiles et n'écrit pas d'énormes matrices de scores en mémoire : les buffers temporaires et le trafic baissent, mais le cache logique et le nombre de têtes KV restent identiques.

Paged attention stocke le cache par blocs qui ne sont pas obligés d'être contigus, et améliore l'allocation et le partage, mais ne réduit pas la taille du cache en soi. L'allocateur doit en outre respecter la règle d'éviction du modèle : on ne peut pas jeter d'anciennes entrées d'une couche complète au seul motif que les couches voisines sont à fenêtre.

Conclusion pratique

Avant de choisir un checkpoint, lisez la config et estimez le cache pour le contexte et le batch visés. Si deux candidats sont proches en qualité, prenez celui dont les têtes KV ou les fenêtres coûtent le moins : cela se traduit directement en nombre de requêtes simultanées par carte. Et n'oubliez pas le padding : avec des batchs denses, vous pouvez payer pour des positions vides, alors qu'une implémentation ragged l'évite.

La limite de la méthode est simple : le calcul donne une estimation de la mémoire persistante, mais ne tient pas compte des poids, des activations ni des surcoûts de l'allocateur. Les chiffres ci-dessus sont illustratifs ; sur un vrai modèle, utilisez les valeurs de sa config.