KV Cache

Le cache KV (Key-Value Cache) est une technique d'optimisation essentielle utilisée principalement dans les Large Language Models (LLMs) et d'autres architectures basées sur les Transformer pour accélérer la latence d'inférence et réduire les coûts de calcul. À la base, le cache KV stocke les matrices Key et Value générées par le mécanisme d'attention pour les tokens précédents d'une séquence. En sauvegardant ces calculs intermédiaires, le modèle évite de recalculer les états d'attention pour l'ensemble de l'historique de la conversation à chaque génération d'un nouveau token. Ce processus transforme le flux de génération de texte d'une opération à complexité quadratique en 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, la génération du mot suivant nécessite 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 pour toute la séquence à chaque étape. Le cache KV résout ce problème en agissant comme une banque de mémoire.

• Amélioration de la vitesse : En récupérant les clés et les valeurs précalculées en mémoire, le système accélère considérablement le moteur d'inférence. C'est essentiel pour les applications nécessitant une faible latence, telles que l'inférence en temps réel dans les bots de service client.
• Efficacité des ressources : Bien qu'il augmente l'utilisation de la mémoire (VRAM), il réduit considérablement les calculs (FLOPs) requis par token. Ce compromis est souvent géré par des techniques telles que la quantification de modèle ou le paging, de la même manière que les systèmes d'exploitation gèrent la RAM.
• Contexte étendu : Une gestion efficace du cache KV permet aux modèles de traiter une fenêtre de contexte plus large, leur permettant de traiter de longs documents ou de maintenir des conversations cohérentes sur de longues périodes.

Applications concrètes#

Le cache KV est un composant fondamental dans le déploiement de l'IA générative moderne, mais ses principes s'étendent également à la computer vision (CV).

1. Chatbots génératifs : Des services comme ChatGPT ou Claude reposent fortement sur la mise en cache KV. Lorsqu'un utilisateur pose une question de suivi, le modèle ne relit pas tout l'historique du chat à partir de zéro. Au lieu de cela, il ajoute la nouvelle entrée aux états mis en cache du tour précédent, permettant des réponses quasi instantanées.

2. Compréhension vidéo : Dans les tâches de compréhension vidéo, les modèles traitent les images de manière séquentielle. Tout comme les tokens de texte, les caractéristiques visuelles des images passées peuvent être mises en cache pour aider le modèle à suivre des objets ou à reconnaître des actions sans retraiter l'ensemble de l'historique vidéo. Ceci 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 grandissent, la taille de la KV cache peut devenir un goulot d'étranglement, consommant des gigaoctets de mémoire GPU. Des avancées récentes se concentrent sur l'optimisation de 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 de mémoire non contigus. Cela réduit la fragmentation et permet des tailles de batch plus élevées lors du serving de modèle.
• Quantification du cache KV : Pour économiser de l'espace, les développeurs appliquent souvent une précision mixte ou une quantification int8 spécifiquement aux valeurs mises en cache. Cela réduit l'empreinte mémoire, permettant aux appareils d'edge AI dotés d'une RAM limitée d'exécuter des modèles performants.
• Mise en cache de prompt : Une technique apparentée où les états KV d'un prompt système statique (par exemple, "Tu es un assistant de codage utile") sont calculés une seule fois et réutilisés dans de nombreuses sessions utilisateur différentes. C'est une fonctionnalité essentielle pour optimiser les flux de travail de prompt engineering à grande échelle.

Distinguer les concepts apparentés#

Il est utile de distinguer la KV Cache d'autres termes de mise en cache et d'optimisation :

• Cache KV vs Mise en cache de prompt : Le cache KV fait généralement référence à la mémoire dynamique, token par token, utilisée lors d'un unique flux de génération. La mise en cache de prompt fait spécifiquement référence au stockage de l'état traité d'une instruction d'entrée fixe à réutiliser sur plusieurs appels d'inférence indépendants.
• Cache KV vs Embeddings : Les embeddings sont des représentations vectorielles de données d'entrée (texte ou images) qui 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, spécifiquement dans le but de générer des séquences.
• Cache KV vs Poids du modèle : Les poids du modèle sont les paramètres statiques et appris du réseau de neurones. Le cache KV se compose de données dynamiques et temporaires générées lors de la passe avant d'une séquence d'entrée spécifique.

Exemple : Contexte dans les modèles de vision#

Bien que la mise en cache KV soit plus célèbre en NLP, le concept de maintien d'état s'applique aux modèles de vision avancés. Dans l'exemple ci-dessous, nous simulons l'idée de passer l'état (le contexte) dans un scénario de suivi vidéo en utilisant Ultralytics YOLO26. Ici, le tracker maintient l'identité des objets à travers les images, conceptuellement similaire à la façon dont un cache maintient le contexte à travers les tokens.

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 datasets et déployer des modèles optimisés peuvent utiliser la plateforme Ultralytics, qui simplifie le pipeline depuis l'annotation des données jusqu'au déploiement du modèle efficace. Pour ceux qui s'intéressent aux mécanismes plus profonds de l'attention, des bibliothèques comme PyTorch fournissent les blocs fondamentaux où ces mécanismes de mise en cache sont implémentés.