Le cache KV et PagedAttention : pourquoi votre serveur LLM ralentit

Votre serveur LLM est lent.

Vous avez déployé un modèle 70B sur quatre GPU A100. Tout semble fonctionner normalement à 8h00. À l'heure du déjeuner, la latence a doublé. Vous vérifiez votre mémoire. La majeure partie est occupée par des « buffers de tenseurs ». Il s'agit en réalité d'états mis en cache provenant d'anciennes conversations.

C'est le problème du cache KV. C'est le principal goulot d'étranglement dans le service de LLM en production.

Qu'est-ce que le cache KV ?

Chaque modèle transformer génère des tokens un par un. Pour créer un nouveau token, le modèle a besoin des tenseurs Key et Value de tous les tokens précédents. Les recalculer à chaque fois est trop lent. À la place, le moteur les stocke. Ce stockage est le cache KV.

Le problème de mémoire :

Pour un modèle Llama 3.1 70B, une seule séquence de 4096 tokens nécessite environ 1,3 Go de mémoire.

Si vous avez 256 utilisateurs simultanés, vous avez besoin de 336 Go de mémoire. C'est plus que ce que quatre GPU A100 peuvent contenir. Le cache KV croît si rapidement qu'il utilise souvent plus de mémoire que les poids du modèle eux-mêmes.

La gestion traditionnelle de la mémoire échoue car :

  • Fragmentation interne : vous allouez de l'espace pour 4096 tokens mais n'en utilisez que 300. Vous gaspillez 93 % de cet espace.
  • Absence de partage : deux utilisateurs ayant le même prompt système stockent chacun leur propre copie de ce prompt.
  • Éviction « tout ou rien » : lorsque la mémoire vient à manquer, vous devez déplacer toute la séquence vers le CPU. Cela bloque le GPU.

Comment PagedAttention résout ce problème :

PagedAttention fonctionne comme un système d'exploitation. Il divise le cache KV en petits blocs de taille fixe appelés pages.

Cela résout trois problèmes principaux :

  • Allocation à la demande : une séquence ne consomme des pages qu'au fur et à mesure de sa croissance. Vous ne gaspillez pas de mémoire pour une capacité inutilisée.
  • Prise en charge des préfixes partagés : plusieurs utilisateurs peuvent partager les mêmes pages physiques pour un prompt système commun. Cela utilise une logique de « copy-on-write » pour économiser d'énormes quantités de mémoire.
  • Éviction à granularité fine : lorsque la mémoire est pleine, le système déplace de petites pages vers le CPU au lieu de séquences massives.

Le résultat :

L'utilisation de PagedAttention (la technologie utilisée dans vLLM) peut augmenter le débit de 2 à 4 fois par rapport aux méthodes traditionnelles.

Quand l'utiliser :

  • Haute concurrence.
  • Séquences de longueurs différentes.
  • Prompts partageant le même début.

Quand l'éviter :

  • Inférence locale pour un utilisateur unique.
  • Modèles très petits.
  • Tâches où chaque séquence a exactement la même longueur.

KV Cache et PagedAttention : Ce qu'ils font et pourquoi ils sont importants

L'inférence des grands modèles de langage (LLM) est un processus extrêmement gourmand en ressources, en particulier en ce qui concerne la bande passante de la mémoire et la capacité de stockage. À mesure que les modèles deviennent plus grands et que les demandes augmentent, l'efficacité de la gestion de la mémoire devient cruciale.

Deux concepts clés sont au cœur de cette efficacité : le KV Cache et PagedAttention. Dans cet article, nous allons explorer ce qu'ils sont, comment ils fonctionnent et pourquoi ils sont essentiels pour l'inférence moderne des LLM.

Qu'est-ce que le KV Cache ?

Pour comprendre le KV Cache, nous devons d'abord comprendre comment un LLM génère du texte. Les LLM utilisent l'architecture Transformer, qui repose sur le mécanisme d'attention. Lors de la génération de chaque nouveau jeton (token), le modèle doit examiner tous les jetons précédents pour comprendre le contexte.

Le mécanisme d'attention calcule des vecteurs pour chaque jeton : les vecteurs Query (Q), Key (K) et Value (V).

Lors de la génération de texte (inférence auto-régressive), le modèle produit un jeton à la fois. Pour générer le jeton $n+1$, il a besoin des vecteurs $K$ et $V$ de tous les jetons précédents ($1$ à $n$). Sans optimisation, le modèle devrait recalculer ces vecteurs $K$ et $V$ à chaque étape, ce qui serait incroyablement inefficace et coûteux en termes de calcul.

C'est là qu'intervient le KV Cache. Au lieu de recalculer les vecteurs $K$ et $V$ pour les jetons déjà traités, nous les stockons en mémoire. Lors de l'étape suivante, nous ne calculons que les nouveaux vecteurs $K$ et $V$ pour le jeton actuel et nous les ajoutons au cache existant.

Le problème : La gestion de la mémoire

Bien que le KV Cache accélère considérablement l'inférence, il introduit un nouveau défi : la consommation massive de mémoire. La taille du KV Cache augmente linéairement avec la longueur de la séquence et la taille du modèle.

Le problème majeur réside dans la manière dont cette mémoire est allouée. Traditionnellement, pour chaque requête, un bloc de mémoire contigu est alloué pour stocker le KV Cache. Cela pose deux problèmes majeurs :

  1. Fragmentation interne : On alloue souvent une quantité de mémoire basée sur la longueur maximale de la séquence prévue (par exemple, 2048 jetons). Si la réponse générée ne fait que 100 jetons, le reste de l'espace alloué est gaspillé.
  2. Fragmentation externe : Comme les blocs doivent être contigus, il peut y avoir des espaces libres entre les blocs alloués qui sont trop petits pour accueillir une nouvelle requête, même si la mémoire totale disponible est suffisante.

Cela mène à ce que l'on appelle le "mur de la mémoire" (memory wall), où la capacité de l'inférence est limitée par la mémoire disponible plutôt que par la puissance de calcul.

Qu'est-ce que PagedAttention ?

C'est ici que PagedAttention entre en jeu. Introduit par l'équipe derrière vLLM, PagedAttention s'inspire de la gestion de la mémoire virtuelle utilisée dans les systèmes d'exploitation modernes.

Dans un système d'exploitation, la mémoire n'est pas nécessairement allouée de manière contiguë. Elle est divisée en petites unités appelées "pages". Ces pages peuvent être dispersées dans la mémoire physique, mais l'application les voit comme un espace d'adressage continu grâce à une table de pages.

PagedAttention applique ce même concept au KV Cache.

Au lieu d'allouer un bloc de mémoire contigu et massif pour chaque requête, PagedAttention divise le KV Cache en petits blocs de taille fixe.

Comment cela fonctionne-t-il ?

  1. Blocs non contigus : Les blocs du KV Cache pour une seule requête peuvent être stockés n'importe où dans la mémoire disponible. Ils ne sont pas obligés d'être côte à côte.
  2. Table de blocs (Block Table) : Pour suivre l'emplacement de ces blocs, PagedAttention utilise une table de blocs qui fait le lien entre la séquence logique de la requête et ses blocs physiques en mémoire.
  3. Allocation dynamique : La mémoire n'est allouée que lorsqu'elle est réellement nécessaire. Lorsqu'un nouveau jeton est généré, un nouveau bloc est extrait de la liste des blocs libres et utilisé.

Les avantages de PagedAttention

L'implémentation de PagedAttention apporte des améliorations significatives :

  • Réduction drastique de la fragmentation : La fragmentation interne est presque éliminée, car les blocs sont petits et l'allocation est dynamique. La fragmentation externe est également minimisée car n'importe quel bloc libre peut être utilisé pour n'importe quelle requête.
  • Augmentation du débit (throughput) : En utilisant la mémoire de manière beaucoup plus efficace, on peut stocker beaucoup plus de requêtes simultanées (batching) dans la même quantité de VRAM. Cela permet de traiter davantage de jetons par seconde sur un même matériel.
  • Partage de mémoire efficace : PagedAttention facilite le partage de données entre différentes requêtes. Par exemple, lors de l'utilisation du Parallel Sampling ou du Beam Search, plusieurs requêtes peuvent partager les mêmes blocs de KV Cache pour les jetons de base (le prompt), économisant ainsi une quantité énorme de mémoire.

Conclusion

Le KV Cache est une technique indispensable pour rendre l'inférence des LLM viable, mais sa gestion mémoire traditionnelle est inefficace. PagedAttention résout ce problème en appliquant les principes de la pagination des systèmes d'exploitation à la gestion du cache.

En optimisant l'utilisation de la mémoire, PagedAttention permet de passer d'une gestion rigide et gaspilleuse à une gestion flexible et dynamique, ouvrant la voie à des systèmes d'inférence à haut débit capables de servir de nombreuses requêtes simultanément.

Source : https://dev.to/tech_nuggets/kv-cache-and-pagedattention-what-they-do-and-why-they-matter-jce

Communauté d'apprentissage optionnelle : https://t.me/GyaanSetuAi