Contexte long et gestion du KV-cache

Donner à un grand modèle de langage 4 000 tokens de contexte est facile ; lui en donner un million sans le ruiner en mémoire ni l'écrouler en latence est l'un des problèmes d'ingénierie les plus profonds de l'inférence moderne. Le coupable est double : l'attention est quadratique en longueur de séquence, et le KV-cache — l'astuce même qui rend la génération rapide — voit sa mémoire croître linéairement à chaque token. Cet article explique pourquoi le contexte long est difficile, ce que stocke réellement le KV-cache, et l'arsenal d'idées (MQA/GQA, FlashAttention, PagedAttention, RoPE/YaRN, attention sinks, ring attention) qui rendent les fenêtres de contexte géantes possibles — ainsi que leurs limites bien réelles.

Pourquoi le contexte long est difficile : le mur du O(n²)

Le cœur d'un Transformer est l'auto-attention : chaque token « regarde » tous les autres. Pour une séquence de longueur n, on calcule une matrice de scores de taille n × n : chaque token est comparé à chaque autre token. Le coût en calcul et en mémoire de cette matrice est donc en O(n²).

Doubler le contexte quadruple le travail d'attention. À 1 000 tokens, la matrice a un million d'entrées. À 100 000 tokens, elle en a dix milliards. C'est le mur fondamental du contexte long, et aucune astuce d'ingénierie ne l'efface entièrement.

Ce mur a deux faces distinctes :

• À l'entraînement (et au prefill) — le traitement initial du prompt — le coût est le calcul quadratique brut de l'attention sur toute la séquence d'un seul coup.
• À la génération token par token, c'est différent : on ne recalcule pas toute la matrice à chaque pas, justement grâce au KV-cache. Mais ce cache a son propre coût mémoire, linéaire et implacable.

Comprendre cette dualité est essentiel : le prefill est limité par le calcul, tandis que le décodage autorégressif est limité par la mémoire et la bande passante. Les optimisations efficaces ciblent l'une ou l'autre de ces phases — rarement les deux à la fois.

Le KV-cache : ce qu'il stocke et pourquoi il enfle

Quand un modèle génère du texte de façon autorégressive, il produit un token à la fois. Sans optimisation, chaque nouveau token rejouerait l'attention sur tout le préfixe déjà généré, recalculant les clés (K) et valeurs (V) de tous les tokens passés à chaque pas. C'est un gaspillage colossal : le travail croît en O(n²) sur l'ensemble de la génération.

Le KV-cache supprime cette redondance par une idée simple :

• on calcule la clé et la valeur de chaque token une seule fois ;
• on les conserve en mémoire (le « cache ») ;
• au pas suivant, on ne calcule que le K/V du nouveau token, qui interroge le cache existant.

C'est ce qui rend la génération viable en pratique. La contrepartie est la mémoire : il faut stocker, pour chaque token, chaque couche et chaque tête d'attention, une paire (clé, valeur). La taille du cache vaut donc, en ordre de grandeur :

mémoire_KV ≈ 2 (K et V) × n_couches × n_têtes_kv × dim_tête × longueur × batch × octets_par_valeur

Le facteur décisif est longueur_séquence : à architecture fixe, la mémoire du cache croît linéairement avec le contexte. Chaque token de plus ajoute une tranche constante.

Prenons un ordre de grandeur concret. Pour un modèle de quelques dizaines de couches, en 16 bits (2 octets), avec des têtes K/V de dimension modeste, le KV-cache d'une seule séquence longue peut atteindre plusieurs dizaines de gigaoctets — souvent davantage que les poids actifs nécessaires à un pas de décodage. Multipliez par la taille du batch et la mémoire explose. C'est le KV-cache, pas le calcul, qui limite alors :

• la taille du batch (combien de requêtes en parallèle) ;
• la longueur de contexte maximale réellement servable ;
• et donc le débit (throughput) global du serveur.

Figure : à chaque token, une nouvelle paire (K, V) s'ajoute par couche et par tête ; la mémoire occupée croît linéairement avec la longueur de la séquence.

La conséquence pratique est centrale : servir du contexte long est le plus souvent un problème de mémoire, pas de FLOPS. Toutes les techniques qui suivent attaquent l'un de ces trois leviers :

1. réduire ce que le cache stocke (MQA, GQA) ;
2. réduire les transferts mémoire pendant le calcul (FlashAttention) ;
3. mieux gérer l'espace mémoire occupé (PagedAttention, fenêtre glissante, distribution).

Rétrécir le cache à la source : MQA et GQA

La piste la plus directe est de réduire le facteur n_têtes_kv. Dans l'attention multi-tête classique (MHA), chaque tête de requête (Q) possède sa propre tête de clé/valeur (K/V) : 32 têtes de requête signifient 32 paires K/V à mettre en cache.

L'attention multi-requête (Multi-Query Attention, MQA, Shazeer, 2019) pousse à l'extrême :

• toutes les têtes de requête partagent une seule paire K/V ;
• le cache fond d'un facteur égal au nombre de têtes (par exemple ÷32) ;
• mais ce partage agressif coûte en qualité et rend parfois l'entraînement instable, car le modèle perd de la diversité de « points de vue ».

L'attention par groupes de requêtes (Grouped-Query Attention, GQA, Ainslie et al., Google, 2023) est le compromis devenu standard :

• on regroupe les têtes de requête en G groupes ; chaque groupe partage une paire K/V ;
• GQA interpole continûment entre MHA (G = nombre de têtes) et MQA (G = 1) ;
• le cache est divisé par têtes / G tout en conservant l'essentiel de la qualité.

C'est pourquoi GQA équipe la plupart des grands modèles ouverts récents. Point crucial : c'est un choix d'architecture, donc figé à l'entraînement — on ne convertit pas un modèle MHA en GQA après coup sans le réentraîner ou l'adapter.

Figure : de MHA (une paire K/V par tête, cache = référence) à MQA (une seule paire partagée, cache ÷4 ici), GQA occupe le juste milieu en partageant par groupes.

À noter : MQA/GQA réduisent uniquement la taille du cache, indépendamment de tout le reste. On peut donc les combiner avec FlashAttention (calcul) et PagedAttention (gestion) — ces techniques sont orthogonales et se cumulent en production.

Accélérer le calcul sans le tronquer : FlashAttention

MQA/GQA réduisent la taille du cache ; FlashAttention (Dao et al., 2022) attaque un autre goulot : les transferts mémoire pendant le calcul de l'attention.

L'idée clé tient à la hiérarchie mémoire d'un GPU :

• la SRAM sur puce est minuscule mais d'une bande passante énorme ;
• la HBM (mémoire principale du GPU) est vaste mais environ 10× plus lente ;
• sur un GPU moderne, l'attention est donc memory-bound, pas compute-bound.

Le problème de l'attention naïve : elle matérialise la matrice n × n complète en HBM, puis la relit pour appliquer le softmax. Ces allers-retours saturent les transferts lents, et le GPU passe son temps à attendre la mémoire plutôt qu'à calculer.

FlashAttention est IO-aware : il ne matérialise jamais la matrice complète en HBM. Sa recette tient en trois ingrédients :

• Tuilage (tiling) : on découpe Q, K, V en blocs qui tiennent en SRAM.
• Online softmax : on calcule le softmax de façon incrémentale, en ré-échelonnant les résultats partiels au fur et à mesure, sans jamais voir toute la ligne d'un coup.
• Recalcul (recomputation) : à la rétropropagation, plutôt que de stocker la grande matrice, on la recalcule à la volée — un peu plus de FLOPS contre beaucoup moins de trafic mémoire.

Le tout est fusionné en un seul kernel GPU, ce qui supprime les écritures intermédiaires. Résultat : un calcul exact (aucune approximation, contrairement à l'attention creuse), une mémoire linéaire en longueur de séquence, et des accélérations mesurées de l'ordre de 3× sur GPT-2 (séquence 1K) dans le papier original. FlashAttention est aujourd'hui un composant de base de l'entraînement et de l'inférence à contexte long, et ses successeurs (v2, v3) ont encore amélioré le parallélisme et l'occupation du GPU.

Gérer la mémoire comme un OS : PagedAttention et vLLM

Même avec un cache rétréci, la gestion de l'espace mémoire est un problème en soi. Les serveurs naïfs réservent à chaque requête un bloc contigu dimensionné pour la longueur maximale possible. D'où un gâchis massif, de trois natures :

• fragmentation interne : place réservée mais inutilisée (la requête s'arrête tôt) ;
• fragmentation externe : des trous entre blocs, inutilisables car trop petits ;
• réservation excessive : on bloque la mémoire d'une longueur qu'on n'atteindra peut-être jamais.

Ce gaspillage fait chuter le nombre de requêtes servies en parallèle, donc le débit.

PagedAttention (Kwon et al., SOSP 2023), au cœur de vLLM, importe une idée des systèmes d'exploitation : la pagination. Le mécanisme :

• le KV-cache d'une requête est découpé en blocs (pages) de taille fixe ;
• ces blocs sont alloués dynamiquement et non contigus en mémoire physique ;
• une table de correspondance relie blocs logiques et blocs physiques, exactement comme la mémoire virtuelle d'un OS.

Le gaspillage tombe ainsi à quasi zéro, ce qui permet d'augmenter fortement la taille de batch.

Mieux : la pagination débloque le partage de mémoire, sous deux formes :

• Prefix caching : deux requêtes qui partagent un préfixe commun (même system prompt, même document) peuvent partager les mêmes pages physiques au lieu de les recalculer.
• Copie sur écriture (copy-on-write) : lors d'un échantillonnage parallèle (plusieurs continuations d'un même prompt), les pages communes sont partagées et seules les divergences sont dupliquées.

Le papier annonce un débit 2 à 4× supérieur aux serveurs de l'état de l'art à latence équivalente, gains d'autant plus marqués que les séquences sont longues, les modèles gros, et les algorithmes de décodage complexes.

Étendre la fenêtre : encodage positionnel et RoPE / PI / YaRN

Réduire la mémoire ne suffit pas : encore faut-il que le modèle comprenne des positions qu'il n'a jamais vues à l'entraînement. La plupart des grands modèles utilisent le RoPE (Rotary Position Embedding) :

• on encode la position en faisant tourner les vecteurs Q et K d'un angle proportionnel à la position ;
• comme l'angle dépend de la position, le produit scalaire entre deux tokens ne dépend que de leur distance relative ;
• on obtient ainsi une position relative « gratuite », sans paramètre entraînable supplémentaire.

Le problème : entraîné sur 4K tokens, un modèle extrapole mal au-delà. Aux positions inédites, les composantes haute fréquence ont tourné dans des plages d'angles jamais apprises, et la qualité s'effondre. Plusieurs méthodes corrigent cela :

• Position Interpolation (PI) (Chen et al., 2023) : au lieu d'extrapoler, on compresse les positions dans la plage connue (la position p d'un contexte doublé est traitée comme p / 2). Un court fine-tuning suffit ensuite à adapter le modèle.
• NTK-aware scaling : interpolation non uniforme selon la fréquence, qui préserve mieux les hautes fréquences (importantes pour la position fine).
• YaRN (Peng et al., 2023) : raffine l'interpolation par bandes de fréquences, sans surcoût de calcul ou de mémoire par rapport aux méthodes précédentes, et converge le plus vite au fine-tuning.

YaRN est l'une des recettes les plus utilisées pour passer, par exemple, de 4K à 128K tokens avec un budget de réentraînement modeste. La leçon générale : étendre le contexte n'est pas qu'une affaire de mémoire — c'est aussi une affaire d'interprétation des positions.

Oublier intelligemment : fenêtre glissante, StreamingLLM et attention sinks

Une autre famille d'approches accepte de ne pas tout garder. L'attention à fenêtre glissante (sliding window) ne conserve en cache que les k derniers tokens : la mémoire devient constante au lieu de linéaire. Hélas, dès que le texte dépasse la fenêtre et que les tout premiers tokens sont évincés du cache, la qualité s'effondre brutalement.

StreamingLLM (Xiao et al., ICLR 2024) explique pourquoi et corrige. L'observation surprenante :

• les modèles accordent une part énorme de leur attention aux tout premiers tokens, même sémantiquement insignifiants ;
• ces tokens servent de « puits d'attention » (attention sinks) où le softmax « dépose » sa masse de probabilité ;
• évincer ces tokens déstabilise toute la distribution d'attention, d'où l'effondrement.

La solution est minime : garder quelques tokens initiaux (typiquement 4) comme puits, plus une fenêtre glissante des tokens récents. Avec cette astuce, le modèle reste stable sur des séquences de millions de tokens, sans aucun fine-tuning, avec jusqu'à 22× d'accélération sur la ligne de base à recalcul.

Attention toutefois à ne pas confondre les promesses : c'est de la génération à mémoire bornée, pas une vraie compréhension de tout le passé. L'information évincée de la fenêtre est bel et bien perdue — StreamingLLM résout la stabilité du streaming, pas le rappel longue distance.

Répartir la séquence : Ring Attention

Quand une seule séquence dépasse la mémoire d'un seul GPU, on distribue. Ring Attention (Liu et al., 2023) procède ainsi :

• on découpe la séquence en blocs répartis sur plusieurs appareils disposés en anneau ;
• chaque appareil calcule l'attention sur son bloc local ;
• puis il fait circuler ses blocs K/V au voisin suivant dans l'anneau ;
• la communication est intégralement recouverte par le calcul des blocs.

Aucune approximation : l'attention reste exacte. Le contexte adressable croît avec le nombre d'appareils, permettant des séquences dépassant la centaine de millions de tokens — au prix d'une infrastructure multi-GPU et d'une orchestration soignée des transferts.

Le piège du « lost in the middle »

Un contexte long techniquement disponible ne garantit pas qu'il soit réellement exploité. L'étude « Lost in the Middle » (Liu et al., 2023) a révélé une courbe de performance en U :

• les modèles retrouvent bien une information placée au début ou à la fin du contexte ;
• mais leur précision chute fortement quand elle est enfouie au milieu.

Le test canonique est l'aiguille dans la botte de foin (needle in a haystack) : on cache un fait précis à différentes profondeurs d'un long document et on demande au modèle de le retrouver. La courbe en U montre que la simple présence d'une information dans la fenêtre ne suffit pas à garantir son utilisation.

La leçon : la taille de la fenêtre est une borne supérieure, pas une garantie d'usage. Quelques conséquences pratiques :

• placez l'essentiel aux extrémités du prompt, pas au centre ;
• ne supposez jamais qu'un modèle « a lu » tout son contexte sous prétexte qu'il y tient ;
• évaluez le contexte long au-delà du needle trivial, vers du raisonnement multi-documents.

Les modèles les plus récents progressent nettement sur ce front, mais le placement de l'information dans le prompt reste un levier concret de qualité.

Contexte long ou RAG ? Le bon arbitrage

Faut-il alors tout entasser dans une immense fenêtre, ou récupérer dynamiquement les bons passages (RAG) ? Le compromis est concret.

Contexte long :

• zéro infrastructure de récupération à maintenir ;
• raisonnement global sur tout le matériel d'un coup, idéal quand l'information pertinente est dense et imbriquée ;
• mais chaque token coûte en mémoire et en latence (le prefill quadratique) ;
• le « lost in the middle » guette, et la facture grimpe vite.

RAG :

• on n'injecte que les quelques passages pertinents, ce qui garde le prompt court et peu coûteux ;
• la base de connaissances se met à jour sans réentraîner le modèle ;
• mais la qualité dépend entièrement du retriever : un passage manqué est une réponse manquée ;
• le raisonnement inter-documents est plus difficile à orchestrer.

En pratique, les deux se combinent souvent : RAG pour cibler les bons documents, contexte long pour raisonner dessus confortablement. La règle saine est de ne pas confondre capacité et usage : une fenêtre de 1M tokens n'est utile que si le modèle l'exploite vraiment et si l'on est prêt à en payer la mémoire et la latence.

Une feuille de route pratique

Pour servir du contexte long efficacement, ces leviers se cumulent :

1. Choisir une architecture GQA (cache plus petit dès la conception).
2. Utiliser FlashAttention pour le prefill et l'entraînement (calcul memory-bound éliminé).
3. Servir avec un moteur paginé (vLLM / PagedAttention) pour le prefix caching et un batch dense.
4. Étendre la fenêtre par YaRN/PI si le contexte cible dépasse l'entraînement de base.
5. Pour le streaming sans fin, ajouter des attention sinks (StreamingLLM).
6. Pour une séquence unique gigantesque, distribuer via Ring Attention.
7. Toujours mesurer le rappel réel (needle, multi-documents) avant de promettre une capacité.

En résumé

Le contexte long bute sur deux murs : l'attention O(n²) et le KV-cache linéaire. On les attaque sur trois fronts complémentaires.

• À la source, MQA/GQA réduisent le nombre de têtes K/V mises en cache.
• Au calcul, FlashAttention élimine les transferts mémoire inutiles par tuilage et recalcul exact, tandis que RoPE/PI/YaRN étendent la fenêtre positionnelle.
• À la gestion, PagedAttention pagine et partage le cache comme un OS, StreamingLLM borne la mémoire via les attention sinks, et Ring Attention distribue la séquence sur plusieurs GPU.

Au-dessus de tout cela plane le « lost in the middle » : disposer d'un contexte géant n'est pas savoir s'en servir — et c'est souvent le choix raisonné entre contexte long et RAG qui fait la différence en production.