LLM : PagedAttention et RadixAttention optimisent le cache KV

Deux avancées techniques s’attaquent au talon d’Achille du service de LLM: la gestion du cache KV. PagedAttention recompose l’allocation mémoire, RadixAttention capitalise sur la réutilisation de préfixes. Ensemble, elles visent un service plus concurrent, plus sobre en GPU et plus réactif, sans modifier l’algorithme d’attention lui-même.

RadixAttention réduit le calcul en réutilisant les préfixes

RadixAttention transforme le cache KV en un index réutilisable plutôt qu’en un simple tampon temporaire. Les préfixes sont stockés dans un arbre radix, un trie compressé où chaque arête représente une séquence de tokens, ce qui permet de ne conserver chaque préfixe de prompt qu’une seule fois et de ne bifurquer que là où les tokens diffèrent. À l’arrivée d’une nouvelle requête, le système recherche le plus long préfixe déjà présent, réutilise les tenseurs KV correspondants et ne calcule que le suffixe manquant, qu’il insère dans l’arbre pour les requêtes futures. Plus le préfixe partagé est long, plus la phase de préremplissage est réduite, ce qui diminue directement le Time to First Token, notamment pour les conversations étendues et les applications d’agents. Contrairement à PagedAttention, qui cible l’utilisation mémoire, cette approche agit sur l’efficacité de calcul et convertit des prompts répétés en accès cache, évitant des milliers de calculs identiques de transformeur. Un exemple illustratif évoque trois requêtes partageant le même prompt système, consolidées en un tronc commun avec des branches uniquement sur la question utilisateur.

PagedAttention traite la mémoire : allocation par blocs et table d’adresses

PagedAttention repose sur une idée simple : n’allouer la mémoire KV qu’à la demande. Le cache est découpé en blocs de taille fixe, souvent 16 ou 32 tokens, et de nouveaux blocs ne sont réservés qu’une fois le précédent rempli, ce qui rend la croissance réellement incrémentale. Chaque séquence est vue comme une suite de blocs logiques pouvant résider n’importe où en mémoire GPU. Une table de blocs, propre à chaque requête, associe identifiants logiques et emplacements physiques ; le noyau d’attention la consulte pour rassembler les clés et valeurs et donner l’illusion d’une continuité malgré la dispersion. Inspirée des tables de pages des systèmes d’exploitation, cette architecture remplace l’allocation contiguë par une mise en pages par blocs, sans modifier l’algorithme d’attention ni les sorties du modèle. Les effets attendus portent sur la baisse du gaspillage mémoire, une meilleure utilisation GPU et davantage de requêtes concurrentes sur le même matériel.

Deux problèmes indépendants : fragmentation et recomputation

Deux goulots d’étranglement distincts freinent les charges de LLM : la fragmentation mémoire et la recomputation de préfixes. PagedAttention cible spécifiquement l’allocation efficace pour limiter la fragmentation, tandis que RadixAttention s’attaque à la réutilisation inter-requêtes des préfixes déjà encodés. Même après l’adoption d’une mise en pages mémoire plus rationnelle, la recomputation persistait : en production, nombre de requêtes partagent un prompt système, les conversations réinjectent leur historique et des agents ajoutent sans cesse du contexte. Une part coûteuse du préremplissage re-générait ainsi des tenseurs KV existants, appelant une solution dédiée côté calcul.

L’ampleur du cache KV fixe la concurrence des services

Le décodage auto-régressif s’effectue token par token, chaque nouveau token consultant toutes les clés et valeurs du passé. Recalculer ces vecteurs à chaque pas serait prohibitif, d’où le cache KV, qui supprime des redondances et rend la génération praticable. Ce bénéfice a une contrepartie : la mémoire utilisée croît linéairement avec la longueur de séquence. Dans les modèles à long contexte, le cache KV devient souvent le principal poste dynamique de mémoire GPU et borne le nombre de requêtes simultanées. Le besoin par token dépend du nombre de couches, des têtes KV, de la dimension de tête et des octets par valeur, avec par exemple 2 octets en FP16. Pour un modèle de la classe Llama‑3 8B avec 32 couches, 8 têtes KV et des têtes de dimension 128 en FP16, l’occupation est estimée à environ 128 KiB par token, et un contexte de 100000 tokens approche 12,8 GiB, avant même tout batching. En pratique, la performance en production dépend donc souvent d’abord de la gestion de ce cache, qui pèse sur la concurrence, le débit et la latence.

Pourquoi l’allocation contiguë limite les charges

En 2023, la communauté a pointé la méthode de stockage du cache KV comme première source d’inefficacité du service, plus que l’attention elle-même. Les moteurs allouaient alors de grands blocs contigus par requête, tout en réservant souvent une capacité proche du contexte maximal, faute de connaître la longueur finale de réponse. Une large part de cette mémoire restait inemployée, réduisant le nombre de séquences parallèles. Deux formes de fragmentation s’ensuivaient : interne, quand des milliers d’emplacements réservés restaient vides pour de petites réponses ; et externe, quand les fins de requêtes de longueurs variables laissaient des trous épars rendant impossible une nouvelle grande allocation contiguë malgré une somme de mémoire libre suffisante. L’effet combiné était une mauvaise utilisation du GPU et un débit moindre, alors même que de la mémoire restait disponible.

Partage de blocs et copy-on-write pour les prompts communs

En plus de l’allocation à la demande par blocs fixes, PagedAttention autorise le partage de blocs : des requêtes démarrant par un prompt identique pointent vers les mêmes blocs physiques, évitant les doublons. Quand deux séquences divergent, une copie n’intervient qu’au point de modification, selon un mécanisme de copy-on-write. Ce partage de préfixe se révèle particulièrement économe en mémoire pour la recherche en faisceau, l’échantillonnage en parallèle ou des requêtes concurrentes dotées d’un prompt système commun. Dans ce modèle, de nouveaux blocs ne sont réservés qu’après remplissage des précédents, et une requête qui ne produit que 60 tokens n’occupe que les blocs nécessaires à ces 60 tokens, limitant la fragmentation interne.