PagedAttention : comment optimiser la mémoire GPU lors de l’inférence des LLM

Quand on commence à déployer des modèles de langage à grande échelle, on se heurte très vite à un mur invisible mais extrêmement rigide : la gestion de la mémoire du GPU. Lors de mes premières expérimentations, je voyais les ressources s’épuiser à une vitesse alarmante, sans toujours comprendre pourquoi. C’est en me plongeant dans le fonctionnement de PagedAttention que j’ai enfin compris comment surmonter cette limite.

Voici les points clés que j’ai retirés de cette technologie :

• PagedAttention résout le gaspillage de la mémoire du cache Key-Value (KV) en le divisant en pages de taille fixe, stockées de manière non contiguë.
• Cette approche réduit le gaspillage de mémoire à presque zéro et multiplie par 2 à 4 le débit d’inférence (throughput).
• Elle s’associe parfaitement au batching continu et s’intègre avec FlashAttention pour combiner capacité et vitesse de calcul.
• Des frameworks comme SGLang poussent le concept encore plus loin grâce à la mise en cache et au partage de préfixes à long terme.

Le problème de départ : le gaspillage du KV Cache

Lors de la génération de texte par un grand modèle de langage, le système conserve en mémoire les états des tokens précédents. C’est ce qu’on appelle le cache Key-Value (KV cache), indispensable pour s’éviter des calculs redondants et coûteux à chaque nouveau token. Cependant, les moteurs d’inférence traditionnels gèrent très mal cet espace pour plusieurs raisons :

• La pré-allocation excessive : Les systèmes classiques réservent un bloc de mémoire GPU unique et continu pour chaque requête, en se basant sur sa longueur maximale possible (par exemple, 2048 tokens). C’est un peu comme réserver une place de parking de semi-remorque pour une simple bicyclette.
• La fragmentation interne : Si le modèle génère une réponse très courte, tout l’espace restant dans ce bloc réservé reste inutilisé et totalement gaspillé.
• La fragmentation externe : Comme les requêtes ont des tailles dynamiques qui varient constamment, la mémoire globale finit par ressembler à un gruyère. Il devient alors impossible d’intégrer de nouvelles requêtes par lots (batching), même s’il reste théoriquement de la mémoire libre.

Selon l’article de recherche original publié par Kwon et al. en 2023 (disponible ici), ces méthodes d’allocation traditionnelles finissent par gaspiller entre 60 % et 80 % de la mémoire allouée au cache KV. C’est ce problème majeur que PagedAttention vient résoudre.

Comment fonctionne PagedAttention

Au lieu de maintenir un tampon continu gigantesque pour chaque requête, PagedAttention gère la mémoire de manière dynamique (une idée d’ailleurs directement inspirée de la mémoire virtuelle de nos systèmes d’exploitation) :

• Des blocs KV de taille fixe : Le cache KV d’une requête est découpé en blocs logiques (ou pages) de taille identique, contenant chacun un nombre précis de tokens.
• Une table de blocs (Block Table) : Cette table légère sert d’indirection pour mapper ces blocs logiques vers des emplacements physiques totalement arbitraires et non contigus dans la mémoire du GPU.
• Une allocation à la demande : À mesure que le modèle génère du texte, le système n’alloue de nouvelles pages physiques que lorsque la page en cours est entièrement remplie.

Un exemple concret d’allocation

J’ai mis un peu de temps à visualiser le mécanisme au départ, mais un exemple simple permet de bien comprendre la différence avec l’allocation traditionnelle.

Prenons la configuration suivante :

• Taille de bloc : 4 tokens.
• Prompt soumis : “The weather is” (3 tokens).
• Tokens générés par le modèle : “nice”, “today”, puis “!”.

Voici comment la mémoire se remplit étape par étape :

1. Traitement du prompt
Le prompt de 3 tokens [“The”, “weather”, “is”] arrive.
Méthode traditionnelle : Elle alloue d’emblée un espace contigu gigantesque (par exemple pour 2048 tokens), espérant que cela suffira. La majorité de cet espace reste vide.
PagedAttention : Le système voit 3 tokens. Puisque la taille de bloc est de 4, il alloue exactement un bloc physique (le Bloc A) en mémoire. Ce bloc contient alors : ["The", "weather", "is", (Slot vide)].

2. Génération du premier token (“nice”)
Le modèle produit le token “nice”.
PagedAttention consulte la table de blocs et constate qu’il reste un emplacement vide dans le Bloc A. Il y écrit directement le token. Le Bloc A est désormais plein : ["The", "weather", "is", "nice"].

3. Génération du deuxième token (“today”)
Le modèle produit le token “today”.
PagedAttention constate que le Bloc A est saturé. Il demande alors dynamiquement un nouveau bloc au pool global de mémoire GPU. Le système lui attribue le Bloc B, situé n’importe où ailleurs sur la puce (non contigu). La table de blocs est mise à jour pour associer la requête aux Blocs A et B. Le token “today” est écrit dans le premier slot du Bloc B : ["today", (Vide), (Vide), (Vide)].

4. Génération du troisième token (“!”)
Le modèle produit le token “!”.
PagedAttention consulte la table de blocs, cible le Bloc B et y trouve 3 emplacements vides. Il écrit le token dans le deuxième emplacement. Le Bloc B contient maintenant : ["today", "!", (Vide), (Vide)].

La table de blocs en arrière-plan

Pour que le modèle ait l’illusion de lire une séquence continue alors que les données sont éparpillées sur le GPU, PagedAttention maintient cette table de correspondance :

Bloc logique (vu par le modèle)	Bloc physique (emplacement réel sur le GPU)	Tokens stockés
Bloc Logique 0	Bloc Physique #42	[“The”, “weather”, “is”, “nice”]
Bloc Logique 1	Bloc Physique #7	[“today”, “!”, (Vide), (Vide)]

Lors des calculs d’attention, le système utilise ces adresses physiques pour récupérer les bons tokens à la volée.

Dans un moteur traditionnel, si nous avions réservé de l’espace pour 100 tokens pour n’en générer que 6, l’équivalent de 94 tokens de mémoire aurait été gaspillé. Avec PagedAttention, seuls 2 slots (dans le Bloc B) restent temporairement inutilisés. Dès que la requête prend fin, ces blocs sont libérés, ramenant le gaspillage à un niveau proche de zéro.

Les bénéfices concrets à l’usage

• Gaspillage presque nul : En supprimant les marges de sécurité pré-allouées, la quasi-totalité de la mémoire GPU est exploitée pour stocker de vrais tokens.
• Tailles de lots accrues : Cette économie drastique de mémoire permet de traiter beaucoup plus de requêtes d’utilisateurs en parallèle (batching continu) sur un seul GPU.
• Partage flexible du cache KV : Comme les pages sont modulaires, plusieurs requêtes peuvent partager les mêmes blocs physiques. Par exemple, si plusieurs utilisateurs soumettent des requêtes basées sur un même et long prompt système, ce préfixe commun n’est stocké qu’une seule fois en mémoire.
• Support des algorithmes de décodage complexes : PagedAttention optimise nativement le comportement des méthodes gourmandes en mémoire comme l’échantillonnage parallèle ou la recherche par faisceaux (beam search) via des mécanismes de copie sur écriture (Copy-on-Write).

Développé à l’origine pour le moteur d’inférence open-source vLLM, PagedAttention est devenu un standard de l’industrie. On le retrouve aujourd’hui intégré dans presque toutes les infrastructures majeures d’inférence, comme Hugging Face TGI, TensorRT-LLM ou BentoML.

La synergie avec le batching continu (Continuous Batching)

Le batching continu et PagedAttention fonctionnent de concert : le premier optimise la planification des calculs du GPU, tandis que le second optimise l’usage de sa mémoire. Sans PagedAttention, le batching continu à grande échelle est tout simplement impossible à maintenir.

Le batching continu travaille au niveau du token individuel. Au lieu d’attendre que l’intégralité d’un lot de requêtes soit finalisée pour en accepter de nouvelles, le moteur d’inférence injecte ou retire des requêtes à chaque étape de calcul du modèle.

[ Lot actif ] ---> Étape de génération de token ---> Requête terminée ?
      |                                                    |
      |<--- Injecter nouvelle requête (Prefill) <--------- OUI (Libérer les blocs)
      |                                                    |
      └<--- Continuer la génération <────────────────────── NON (Allouer nouveau bloc)

PagedAttention agit comme le socle permettant de rendre ce flux fluide :

1. Libération instantanée de la mémoire : Dès qu’une requête se termine (lorsqu’elle rencontre un token de fin de séquence), le batching continu l’éjecte. PagedAttention restitue immédiatement ses blocs physiques au pool global. Ces blocs peuvent être réattribués dès l’itération suivante à un prompt entrant.
2. Densité maximale des lots : Les systèmes traditionnels imposent de réserver l’espace maximal requis pour chaque utilisateur, ce qui sature rapidement la VRAM. PagedAttention n’allouant de la mémoire que pour les tokens existants, l’empreinte mémoire par utilisateur chute, ce qui permet de multiplier par 2 ou 4 le nombre de requêtes simultanées.
3. Fusion des phases de Prefill et de Decode : Le batching continu mélange le traitement des nouveaux prompts (l’étape de “Prefill”, très calculatoire) et la génération du token suivant pour les requêtes en cours (l’étape de “Decode”, limitée par la bande passante mémoire). PagedAttention permet au modèle de lire les pages de mémoire des anciennes requêtes tout en écrivant de nouvelles pages pour les nouveaux prompts au cours d’une même étape de calcul, sans conflit d’accès.

Si la mémoire VRAM venait à manquer malgré tout en raison d’un trop grand nombre de requêtes simultanées, le planificateur gère alors une préemption pour suspendre temporairement certaines tâches.

PagedAttention vs. FlashAttention

Il ne faut pas confondre ces deux technologies car elles s’attaquent à des goulots d’étranglement différents. FlashAttention accélère la vitesse de calcul, tandis que PagedAttention optimise la capacité de stockage en mémoire. Les moteurs d’inférence modernes les exploitent d’ailleurs simultanément.

Caractéristique	PagedAttention	FlashAttention
Objectif principal	Réduire la fragmentation et le gaspillage de la VRAM.	Accélérer les calculs mathématiques de la matrice d’attention.
Goulot d’étranglement résolu	Capacité mémoire (empreinte VRAM).	Vitesse matérielle (bande passante de la mémoire).
Niveau d’exécution	Mémoire globale du GPU (HBM).	Cache interne ultra-rapide du GPU (SRAM).
Méthode clé	Découpage du cache KV en pages virtuelles.	Calcul de l’attention par tuiles/blocs adaptés à la SRAM.
Bénéfice principal	Tailles de lots élevées et gaspillage mémoire quasi nul.	Temps de génération de tokens nettement plus rapides.

Pendant la génération de texte, le GPU passe beaucoup de temps à transférer des données entre sa mémoire globale lente (HBM) et son cache interne ultra-rapide mais très limité (SRAM). FlashAttention évite d’écrire et de relire de grandes matrices d’attention intermédiaires dans la HBM en effectuant les calculs de softmax de manière incrémentale directement dans la SRAM.

Cependant, FlashAttention part du principe que le cache KV de chaque requête est stocké sous forme d’un unique bloc contigu en VRAM globale. C’est là que PagedAttention intervient pour segmenter ce cache en pages non contiguës et compacter l’espace disponible en VRAM.

Dans un moteur d’inférence optimisé, le flux d’exécution combine le meilleur des deux mondes :

1. PagedAttention consulte sa table de blocs et récupère les pages de mémoire éparpillées dans la mémoire globale (HBM) pour les requêtes actives.
2. FlashAttention prend ces pages, les charge sous forme de blocs optimisés dans le cache SRAM ultra-rapide et exécute les calculs d’attention à la vitesse maximale du matériel.

L’étape suivante : SGLang et RadixAttention

SGLang utilise PagedAttention comme fondation, mais y ajoute un mécanisme appelé RadixAttention. Alors que PagedAttention classique gère la mémoire à l’échelle d’une requête ou d’un lot de requêtes en cours, SGLang permet de partager la mémoire et de conserver un cache à long terme sur des requêtes totalement distinctes arrivant à des moments différents.

Cette approche est idéale pour les architectures d’agents, les pipelines RAG ou les discussions à plusieurs tours où de longs prompts système sont réutilisés en permanence.

La différence d’architecture

Pour comprendre l’apport de SGLang par rapport à un moteur comme vLLM, il faut observer comment la mémoire est structurée :

• Allocation virtuelle plate (vLLM) : vLLM gère un ensemble à plat de pages mémoire. Lorsqu’une requête utilisateur prend fin, les pages physiques correspondantes sont libérées et remises dans le pool vide.
• Structure arborescente hiérarchique (SGLang) : SGLang mappe les blocs de mémoire paginée sur un arbre de préfixes compressé (un Radix Tree). Les branches de cet arbre représentent des séquences de tokens, et les nœuds correspondent directement aux blocs physiques de cache KV.

            [Nœud racine : "You are a coding assistant..."] (Prompt système partagé)
                               /          \
                              /            \
  [Utilisateur A : "Write a loop..."]     [Utilisateur B : "Fix this SQL..."]
             |                                          |
  (Tokens de code générés)                    (Tokens de requête générés)

Le cache de préfixes automatique (Automatic Prefix Caching)

Si une centaine d’utilisateurs interrogent un serveur avec le même prompt système ou un document de contexte RAG très volumineux :

• Avec PagedAttention classique, le système doit traiter le prompt et allouer de la mémoire pour chaque requête séparément si elles n’arrivent pas en même temps.
• Avec SGLang, le prompt système est traité une seule fois. Les pages du cache KV obtenues sont stockées dans un nœud du Radix Tree. Dès qu’un nouvel utilisateur arrive avec un prompt commençant par la même séquence, SGLang détecte la correspondance dans l’arbre et fait pointer la requête vers ces pages déjà calculées. Le temps de génération du premier token (TTFT) tombe à pratiquement zéro.

La politique d’éviction LRU

Pour éviter de perdre des calculs potentiellement réutilisables, SGLang traite la VRAM du GPU comme un cache intelligent. À la fin d’une requête, les pages restent conservées et associées au Radix Tree. Si l’espace vient à manquer pour une nouvelle requête, SGLang applique un algorithme d’éviction de type LRU (Least Recently Used) pour libérer les branches les plus anciennes ou les moins consultées de l’arbre.

Stratégie de configuration de taille de page dans SGLang

Le paramètre page_size peut être ajusté dans SGLang pour s’adapter à des besoins spécifiques :

• Taille de page élevée (par exemple page_size = 16 ou 32) : Ce choix favorise le débit matériel brut. Les gros blocs contigus sont en effet plus rapides à traiter pour les noyaux d’exécution de FlashAttention dans le GPU.
• Taille de page au niveau du token (page_size = 1) : Cette option offre une flexibilité maximale pour la réutilisation des préfixes. Si les prompts se recoupent de manière irrégulière ou dynamique (comme dans le cas de branchements complexes d’agents), configurer la taille à 1 permet un alignement précis au token près sur le Radix Tree.

Quel outil choisir selon l’usage ?

• Je recommande d’utiliser vLLM si les charges de travail se composent principalement de requêtes indépendantes, sans état et à tour unique (par exemple pour de la classification de texte simple ou des requêtes rapides et isolées).
• Mon choix se porte plutôt sur SGLang lors du développement d’architectures multi-agents, de chatbots avec de longs historiques de conversation ou de pipelines générant des structures complexes (comme du JSON) où les mêmes contraintes de formatage sont réutilisées de manière répétitive.