KV Cache

Le cache KV (cache clé-valeur) est une technique d’optimisation essentielle, principalement utilisée dans les grands modèles de langage (LLMs) et d’autres architectures basées sur Transformer afin d’accélérer la latence d’inférence et de réduire les coûts de calcul. Essentiellement, le cache KV stocke les matrices Key et Value générées par le mécanisme d’attention pour les jetons précédents d’une séquence. En sauvegardant ces calculs intermédiaires, le modèle évite de recalculer les états d’attention de tout l’historique de la conversation à chaque nouveau jeton généré. Ce processus transforme le flux de travail de génération de texte, qui passe d’une opération à complexité quadratique à une opération linéaire, rendant possibles les interactions en temps réel avec les chatbots et les agents IA.

Mécanisme et avantages#

Dans un modèle Transformer standard, générer le mot suivant exige de prêter attention à tous les mots précédents pour comprendre le contexte. Sans mise en cache, le modèle devrait recalculer les relations mathématiques de toute la séquence à chaque étape. Le cache KV résout ce problème en faisant office de banque mémoire.

• Gain de vitesse : En récupérant en mémoire les clés et valeurs précalculées, le système accélère considérablement le moteur d’inférence. C’est essentiel pour les applications qui exigent une faible latence, comme l’inférence en temps réel dans les robots de service à la clientèle.
• Efficacité des ressources : Bien qu’il augmente l’utilisation de la mémoire (VRAM), le cache réduit considérablement le calcul (FLOPs) nécessaire par jeton. Ce compromis est souvent géré au moyen de techniques comme la quantification de modèle ou la pagination, comme le font les systèmes d’exploitation avec la RAM.
• Contexte étendu : Une gestion efficace du cache KV permet aux modèles de prendre en charge une fenêtre de contexte plus large, afin de traiter de longs documents ou de maintenir des conversations cohérentes pendant de longues périodes.

Applications concrètes#

Le cache KV est un composant fondamental du déploiement de l’IA générative moderne, mais ses principes s’appliquent aussi à la vision par ordinateur (CV).

1. Chatbots génératifs : Des services comme ChatGPT ou Claude s’appuient fortement sur la mise en cache KV. Lorsqu’un utilisateur pose une question complémentaire, le modèle ne relit pas tout l’historique de la conversation depuis le début. Il ajoute plutôt la nouvelle entrée aux états mis en cache du tour précédent, ce qui permet de répondre presque instantanément.

2. Compréhension vidéo : Dans les tâches de compréhension vidéo, les modèles traitent les images séquentiellement. Comme pour les jetons de texte, les caractéristiques visuelles des images précédentes peuvent être mises en cache pour aider le modèle à suivre les objets ou à reconnaître les actions sans retraiter tout l’historique vidéo. C’est particulièrement pertinent pour la reconnaissance d’actions, où le contexte temporel est crucial.

Gestion efficace de la mémoire#

À mesure que les modèles grossissent, la taille du cache KV peut devenir un goulot d’étranglement et consommer des gigaoctets de mémoire GPU. Les avancées récentes visent à optimiser ce stockage.

• PagedAttention : Inspiré de la mémoire virtuelle des systèmes d’exploitation, PagedAttention (introduit par vLLM) permet de stocker le cache KV dans des blocs mémoire non contigus. Cela réduit la fragmentation et permet d’augmenter la taille des lots lors du déploiement de modèles en service.
• Quantification du cache KV : Pour économiser de l’espace, les développeurs appliquent souvent la précision mixte ou la quantification int8 aux valeurs mises en cache. Cela réduit l’empreinte mémoire et permet aux appareils d’IA en périphérie dotés d’une RAM limitée d’exécuter des modèles performants.
• Mise en cache des prompts : Technique connexe où les états KV d’un prompt système statique (p. ex., « Tu es un assistant de programmation serviable ») sont calculés une fois, puis réutilisés dans de nombreuses sessions utilisateur. C’est une fonctionnalité essentielle pour optimiser à grande échelle les flux de travail d’ingénierie des prompts.

Distinguer les concepts associés#

Il est utile de distinguer le cache KV des autres termes liés à la mise en cache et à l’optimisation :

• Cache KV vs. mise en cache des prompts : Le cache KV désigne généralement la mémoire dynamique, jeton par jeton, utilisée au cours d’un même flux de génération. La mise en cache des prompts désigne précisément le stockage de l’état traité d’une instruction fixe afin de le réutiliser dans plusieurs appels d’inférence indépendants.
• Cache KV vs. embeddings : Les embeddings sont des représentations vectorielles des données d’entrée (texte ou images) qui en capturent le sens sémantique. Le cache KV stocke les activations (clés et valeurs) dérivées de ces embeddings dans les couches d’attention, précisément pour générer des séquences.
• Cache KV vs. poids du modèle : Les poids du modèle sont les paramètres statiques appris par le réseau neuronal. Le cache KV se compose de données dynamiques et temporaires générées pendant la passe avant d’une séquence d’entrée donnée.

Exemple : contexte dans les modèles de vision#

Bien que la mise en cache KV soit surtout connue en NLP, le concept de maintien d’un état s’applique aux modèles de vision avancés. Dans l’exemple ci-dessous, nous simulons l’idée de transmettre un état (contexte) dans un scénario de suivi vidéo avec Ultralytics YOLO26. Ici, le suiveur conserve l’identité des objets d’une image à l’autre, de manière conceptuellement similaire à un cache qui maintient le contexte d’un jeton à l’autre.

from ultralytics import YOLO

# Load the Ultralytics YOLO26 model
model = YOLO("yolo26n.pt")

# Track objects in a video, maintaining identity state across frames
# The 'track' mode effectively caches object features to link detections
results = model.track(source="https://ultralytics.com/images/bus.jpg", show=False)

# Print the ID of the tracked objects
if results[0].boxes.id is not None:
    print(f"Tracked IDs: {results[0].boxes.id.numpy()}")

Les développeurs qui souhaitent gérer des jeux de données et déployer des modèles optimisés peuvent utiliser la plateforme Ultralytics, qui simplifie le pipeline, de l’annotation des données au déploiement efficace des modèles. Pour approfondir le fonctionnement de l’attention, des bibliothèques comme PyTorch fournissent les blocs de base où ces mécanismes de mise en cache sont implémentés.