Notizie IA Logo

AITalk

Actualités et analyses sur l'intelligence artificielle

PagedAttention et RadixAttention : parlons de KV Cache

ResearchGenerative AIApplications

pagedattention-radixattention.jpg

Dans notre article précédent sur AITalk, nous avions ouvert sur un chiffre inconfortable : Llama-3.1-70B, en précision BF16, accumule environ 0,31 mégaoctet de cache pour chaque token traité, ce qui signifie qu'un contexte de 128 000 tokens coûte déjà 40 gigaoctets de mémoire GPU, avant même d'avoir chargé la moindre requête concurrente. Nous avions présenté trois réponses à ce problème — TurboQuant, OSCAR et EpiCache —, trois façons différentes de réduire la taille de ces données en agissant sur le nombre de bits par valeur ou sur les portions de conversation qu'il vaut la peine de conserver en mémoire.

Cependant, compresser les données ne résout que la moitié du problème. C'est comme réussir à emballer les dossiers d'un d'archive dans des boîtes plus fines sans changer la façon dont l'archive elle-même est organisée : si les étagères restent attribuées en bloc à chaque dossier, quelle que soit son importance réelle, et si chaque nouveau dossier identique au précédent est recopié de zéro au lieu d'être récupéré, l'efficacité gagnée sur le poids des dossiers individuels est tout de même perdue en chemin. Cet article traite des deux techniques qui ont affronté précisément ce second problème : non pas combien pèse un token dans le cache, mais comment on l'alloue et comment on le réutilise. Elles s'appellent PagedAttention et RadixAttention, et ce ne sont pas des alternatives à la compression : elles opèrent à un niveau différent et, dans les systèmes de production les plus mûrs, coexistent avec elle.

Deux problèmes indépendants, et non un seul

Même en disposant d'un cache déjà compressé, un moteur d'inférence se heurte à deux inefficacités qui n'ont rien à voir l'une avec l'autre. La première concerne l'allocation de la mémoire. La méthode traditionnelle, adoptée par les premiers systèmes de serving, réservait pour chaque requête entrante un bloc contigu de mémoire GPU dimensionné sur la longueur maximale de contexte que le modèle pouvait théoriquement atteindre. Le problème est que le système n'a aucun moyen de savoir à l'avance quelle sera la longueur réelle de la réponse, si bien que la marge de sécurité finissait presque toujours par être surdimensionnée. Il en découle deux formes de gaspillage distinctes. Il y a la fragmentation interne, lorsqu'une requête réserve de l'espace pour des milliers de tokens futurs mais n'en génère que quelques dizaines, laissant inutilisée la majeure partie de la mémoire attribuée. Et il y a la fragmentation externe, lorsque des requêtes de longueurs différentes se terminent à des moments différents en laissant des trous dispersés dans la mémoire GPU : la somme de l'espace libre peut être abondante, mais aucun fragment n'est assez grand et contigu pour accueillir une nouvelle requête volumineuse. Le résultat pratique, mesuré par les chercheurs qui ont abordé le problème pour la première fois, était que les systèmes d'alors gaspillaient entre 60 et 80 % de la mémoire allouée, un chiffre qu'il vaut la peine de garder à l'esprit pour évaluer l'ampleur du gain ultérieur.

Le second problème est conceptuellement éloigné du premier et concerne le calcul, non la mémoire. Dans les charges de travail réelles, les requêtes sont rarement indépendantes les unes des autres. Des milliers d'utilisateurs qui interrogent le même assistant partagent le même prompt système, une conversation à plusieurs tours réintroduit à chaque échange l'ensemble de l'historique précédent, et un agent qui raisonne en boucle continue d'ajouter de nouvelles étapes à un contexte qui reste en grande partie identique au tour précédent. Dans tous ces cas, le moteur d'inférence se retrouve à recalculer de zéro, durant la phase de prefill, exactement les mêmes tenseurs clé et valeur qu'il avait déjà produits quelques instants plus tôt pour une autre requête, ou pour le même utilisateur au tour précédent. C'est un travail purement redondant qui, dans des scénarios de chat prolongé ou de systèmes RAG, peut représenter la part la plus lourde du temps de réponse perçu par l'utilisateur.

Ces deux inefficacités sont nées de besoins différents et ont reçu des solutions différentes. PagedAttention traite la première, RadixAttention traite la seconde.

PagedAttention : la mémoire comme un système d'exploitation

L'idée fondamentale de PagedAttention, introduite par le groupe à l'origine de vLLM, emprunte presque directement un concept vieux de plusieurs décennies : la mémoire virtuelle paginée des systèmes d'exploitation. Quiconque a déjà ouvert le capot d'un ordinateur sait que les programmes n'obtiennent pas un bloc physique contigu de mémoire RAM, mais une illusion de contiguïté construite sur des pages dispersées gérées par une table de traduction. PagedAttention applique exactement la même logique au cache d'attention.

Le mécanisme s'articule en quatre étapes. Tout d'abord, le cache de chaque séquence est divisé en blocs logiques de taille fixe, généralement de seize ou trente-deux tokens chacun, au lieu d'être traité comme un conteneur monolithique unique. Deuxièmement, le système maintient une table de pages, une carte qui traduit chaque bloc logique dans son emplacement physique réel dispersé dans la mémoire de la GPU : pendant le calcul de l'attention, le moteur consulte cette table pour rassembler les blocs nécessaires, et pour le modèle, la séquence apparaît continue même si elle ne l'est pas du tout physiquement. La troisième étape est la croissance à la demande : la mémoire n'est plus réservée en bloc au départ, mais allouée bloc par bloc au fur et à mesure que la génération progresse, de sorte qu'une réponse de soixante tokens n'occupe que l'espace de soixante tokens, et non celui prévu pour un contexte maximal qui ne sera peut-être jamais atteint. Le quatrième élément, peut-être le plus élégant, est le partage de blocs avec copy-on-write : si plusieurs requêtes commencent par le même prompt, comme cela arrive quotidiennement avec un prompt système partagé par des milliers d'utilisateurs, les blocs physiques correspondants sont référencés en commun au lieu d'être dupliqués, et ce n'est qu'au moment précis où une requête diverge des autres que le bloc concerné est effectivement copié.

Le résultat de cette architecture est que la fragmentation s'effondre pratiquement à zéro, avec des marges de gaspillage qui, dans les benchmarks publiés, tombent sous la barre des 4 % contre 60 à 80 % pour les systèmes précédents, ce qui permet de servir un nombre beaucoup plus élevé de requêtes concurrentes sur le même matériel. Il est intéressant de noter que PagedAttention ne touche en rien à l'algorithme d'attention et ne modifie pas les sorties du modèle : son innovation est purement architecturale et concerne l'endroit et la manière dont les données sont physiquement placées, non ce que le modèle calcule. immagine1.jpg Image tirée de l'article sur arxiv.org

RadixAttention : un arbre qui n'oublie pas

Si PagedAttention résout le , RadixAttention résout le quoi. L'idée, introduite par l'équipe de SGLang, part d'une observation aussi simple que négligée : dans les systèmes traditionnels, une fois une requête terminée, son cache était simplement rejeté, comme si chaque conversation naissait et mourait sans laisser de trace utile pour les suivantes. RadixAttention renverse cette logique en transformant le cache en une structure persistante et searchable — un radix tree, c'est-à-dire un trie compressé dans lequel chaque préfixe unique de tokens est stocké une seule fois, et des requêtes différentes partageant le même début parcourent la même branche de l'arbre, ne bifurquant qu'au point exact où leurs tokens commencent à différer.

Imaginons trois utilisateurs qui ouvrent une conversation avec le même message système, suivi chacun d'une question différente. Au lieu de conserver trois copies presque identiques du même préfixe, l'arbre le maintient une seule fois et ne crée trois feuilles séparées que pour la dernière partie, celle véritablement spécifique à chaque question. Lorsqu'une nouvelle requête arrive, le moteur parcourt l'arbre token par token à la recherche du préfixe le plus long déjà présent, récupère immédiatement les tenseurs clé et valeur correspondants sans les recalculer, et ne calcule enfin que le suffixe pas encore vu, en insérant le nouveau parcours dans l'arbre pour qu'il puisse être réutilisé à l'avenir. Un aspect souvent sous-estimé est que la structure ne peut pas croître à l'infini : lorsque la mémoire GPU sature, RadixAttention applique une politique d'éviction basée sur le least recently used qui élimine en premier les feuilles les moins récemment utilisées, tout en protégeant les nœuds internes partagés dont dépendent plusieurs requêtes, afin de préserver précisément les préfixes qui rendent l'ensemble du mécanisme rentable.

Le bénéfice principal ne concerne pas la taille mémoire en soi, mais le Time to First Token, c'est-à-dire le temps qui s'écoule avant que l'utilisateur ne voie apparaître le premier mot de la réponse : si mille neuf cent tokens sur deux mille sont déjà présents dans l'arbre, le modèle n'a besoin de traiter que les cent restants, avec un impact direct sur la latence perçue, particulièrement évident dans les chatbots, les systèmes RAG et les workflows d'agents où les contextes croissent de façon incrémentale au lieu d'être réécrits de zéro à chaque fois. Ce mécanisme rappelle, pour ceux qui connaissent la narration vidéoludique de niche, la structure de 80 Days, le jeu narratif d'Inkle dans lequel des milliers de parcours de voyage partagent des pans entiers d'intrigue commune et ne se séparent qu'aux points de choix effectifs : le moteur du jeu ne réécrit pas chaque ramification possible depuis le début, mais la construit en la greffant sur un tronc narratif partagé. immagine2.jpg https://arxiv.org/html/2312.07104v1

Deux couches, pas deux rivaux

Une erreur fréquente, entretenue en partie par la façon dont les deux projets sont souvent présentés en opposition, est de penser à PagedAttention et RadixAttention comme des technologies concurrentes entre lesquelles il faudrait choisir. En réalité, elles répondent à des questions différentes. PagedAttention décide où les blocs de cache résident physiquement dans la mémoire GPU ; RadixAttention décide si ces blocs existent déjà et peuvent être réutilisés. Le radix tree, dans un système intégrant les deux techniques, pointe simplement vers des blocs qui sont eux-mêmes gérés par l'allocateur paginé : ce sont deux couches superposées d'une même pile, et non deux voies alternatives. immagine3.jpg

En revenant à l'exemple des trois utilisateurs ayant le même prompt système, sans RadixAttention le moteur recalculerait ce préfixe trois fois, et sans PagedAttention chacune de ces trois requêtes réserverait quand même une portion de mémoire surdimensionnée. En combinant les deux techniques, le préfixe n'est calculé qu'une seule fois, stocké dans des blocs paginés compacts, et réutilisé par chaque requête qui le partage.

Il convient de préciser que RadixAttention n'est pas la seule voie vers la réutilisation automatique des préfixes. vLLM, le projet qui a introduit PagedAttention, prend également en charge aujourd'hui le prefix caching automatique, mais avec une structure de données différente : au lieu d'un radix tree, il utilise un mécanisme de hachage en chaîne, dans lequel chaque bloc terminé reçoit un hash calculé à partir du hash du bloc parent, des tokens contenus dans le bloc lui-même et d'éventuelles métadonnées supplémentaires. Comme chaque hash dépend du précédent, la chaîne entière représente de manière univoque le préfixe menant à ce bloc ; si une autre requête produit exactement la même séquence de tokens, elle génère la même chaîne de hash, trouvant ainsi immédiatement les blocs déjà calculés. En pratique, pour la majorité des applications, la différence entre les deux implémentations reste plus architecturale que fonctionnelle : les deux moteurs parviennent au même résultat — sauter le prefill des préfixes déjà vus —, quoique avec des structures de données différentes en arrière-plan.

Au-delà des deux techniques : où va la gestion du cache

PagedAttention et RadixAttention ont abordé les deux problèmes fondateurs, mais les contextes continuent d'augmenter, et avec eux sont nés des besoins que ni l'une ni l'autre de ces techniques n'avait été conçue, seule, pour couvrir.

La première direction est la mémoire hiérarchique. Traiter la GPU comme le seul niveau de cache disponible devient insoutenable lorsque les contextes dépassent les centaines de milliers de tokens : des projets comme LMCache et Mooncake organisent désormais le cache sur plusieurs niveaux, avec les blocs les plus chauds dans la mémoire à bande passante élevée de la GPU, ceux utilisés récemment dans la RAM de l'hôte, et les plus anciens mais potentiellement encore utiles basculés sur du stockage distribué ou des disques NVMe, récupérés automatiquement quand ils redeviennent nécessaires. Pour ceux qui ont lu Zafón, c'est une architecture qui rappelle de près le Cimetière des livres oubliés : rien n'est détruit définitivement, mais les volumes les moins consultés sont déplacés vers des étagères de plus en plus distantes, prêts à être récupérés si quelqu'un en a de nouveau besoin.

La deuxième direction est le routage conscient du cache. Dans une infrastructure distribuée sur plusieurs nœuds, un équilibreur de charge classique qui répartit les requêtes de façon purement round-robin risque d'envoyer deux tours d'une même conversation à deux GPU différentes, annulant tout bénéfice du prefix caching même lorsque le préfixe partagé existe. Des systèmes comme le routeur de llm-d traitent le problème en aiguillant les requêtes vers la réplique qui possède déjà le cache du préfixe demandé, en équilibrant non seulement la charge mais aussi la localité des données.

La troisième direction, moins débattue mais pertinente pour ceux qui exploitent des infrastructures multi-tenants, concerne la sécurité. Le prefix caching introduit un canal latéral potentiel : si un préfixe est servi visiblement plus vite parce qu'il est déjà présent en cache, un utilisateur malveillant observant des temps de réponse anormalement bas pourrait en déduire qu'un autre utilisateur a envoyé auparavant le même prompt, sans même que le contenu de la réponse soit directement exposé. La contre-mesure la plus répandue est le cache salting, qui inclut dans le hash servant à identifier les blocs un élément spécifique du tenant, de sorte que des requêtes identiques provenant de clients différents génèrent tout de même des clés de cache différentes, éliminant la possibilité d'un hit croisé entre tenants tout en conservant intacts les bénéfices de latence au sein du même tenant.

Compression et gestion, deux faces d'une même pièce

Pour revenir au point de départ de notre article précédent, la vision complète de la gestion moderne de la KV cache n'émerge qu'en rassemblant ces deux fronts. La compression — dont TurboQuant, OSCAR et EpiCache sont des expressions récentes — réduit la taille du problème en agissant sur le nombre de bits nécessaires pour représenter chaque token et sur les portions d'historique de conversation qu'il vaut la peine de conserver. La gestion — dont PagedAttention et RadixAttention sont les piliers — en optimise la structure, en décidant où ces données résident physiquement et si elles peuvent être réutilisées au lieu d'être recalculées. Un système de production à l'état de l'art intègre aujourd'hui ces deux approches : il compresse ce qui doit rester en mémoire et organise efficacement l'espace et le calcul autour de cette mémoire compressée.

Pour quiconque conçoit ou évalue une infrastructure d'inférence, cela signifie que le choix ne réside plus seulement dans le modèle à servir, mais dans la manière de le servir : maîtriser le fonctionnement de PagedAttention et RadixAttention est devenu aujourd'hui aussi primordial que de connaître l'architecture du modèle que ce cache soutient effectivement.


Note technique : les données citées dans cet article proviennent des papiers originaux de vLLM et SGLang ainsi que de la documentation de leurs projets open source respectifs ; elles n'ont pas été reproduites de façon indépendante par cette rédaction et, comme toujours dans ce domaine, les chiffres de laboratoire ne garantissent pas automatiquement les mêmes performances dans des scénarios de production avec des charges de travail réelles.