IA : optimiser le cache, l’obsession de tout un secteur

Une meilleure gestion de l’attention et donc du KV Cache

La raison principale de cette évolution est à chercher du côté du mécanisme d’attention.

L’attention vise à calculer la relation entre tous les tokens d’une séquence, d’en comprendre le sens, avant de produire une réponse. Or cette opération est normalement quadratique : il faut calculer le score d’attention entre chaque token en entrée. Des optimisations de ce mécanisme ont permis de réduire ce fardeau computationnel et d’allonger les fenêtres de contexte. Néanmoins, les gains les plus probants concernent le KV Cache.

 « Le KV Cache est un composant standard dans l’implémentation de l’attention et stocke les paires clés-valeurs pour tous les tokens précédents », explique Philip Kiely, directeur de l’éducation sur l’IA chez Baseten, dans son ouvrage Inference Engineering.

Les clés (K) représentent l’identité de chaque token (ce qu’il est), tandis que les valeurs (V) encodent l’information qu’il peut apporter. Ces représentations sont obtenues en transformant les embeddings des tokens via des matrices de poids apprises lors de l’entraînement. Or, lorsque ces deux informations sont déjà connues, il n’est pas nécessaire de les recalculer. En les stockant en cache, l’opération devient linéaire par étape de génération : seule la requête (Q) du token à générer est calculée et comparée aux clés en cache pour obtenir les scores d’attention.

Des algorithmes d’optimisation comme FlashAttention et PagedAttention ont permis d’adapter l’attention aux subtilités des GPU. Le premier minimise les écritures et lectures entre la HBM et la SRAM (mémoire partagée intégrée à la puce) en vue d’accélérer la génération. Le second s’inspire de la pagination mémoire des systèmes d’exploitation pour partitionner et allouer dynamiquement le KV Cache afin de traiter davantage de requêtes simultanément. Ces deux algorithmes sont implémentés dans les frameworks d’inférence vLLM, TensorRT LLM, Dynamo ou encore SGLang.

Il faut également prendre en compte la compression (quantization) des poids des modèles, ainsi que du contenu du KV Cache lui-même. Abaisser la précision des calculs en virgule flottante peut réduire de « manière significative » l’empreinte mémoire GPU d’un modèle en affectant modérément la qualité des résultats. Il en va de même pour le KV Cache, mais attention aux effets cumulés des erreurs de précision entre le traitement de chaque token, souligne Philip Kiely. La technique est publiquement utilisée par les porteurs de modèles open weight, dont Mistral AI.

La montée en puissance du prefix caching…

Une autre technique, plus sûre, a su s’imposer comme un moyen d’optimiser les coûts d’inférence : le « prefix caching ».

« L’idée centrale est simple : nous mettons en cache les blocs de KV cache des requêtes traitées, et réutilisons ces blocs lorsqu’une nouvelle requête arrive avec le même préfixe que les requêtes précédentes », expliquent les mainteneurs de vLLM, dans la documentation du projet. « Étant donné que la mise en cache des préfixes n’a presque que des avantages et ne modifie pas les sorties du modèle, elle a été largement adoptée par de nombreux fournisseurs de LLM (OpenAI, Anthropic, etc.) et la plupart des frameworks d’inférence LLM open source (SGLang) ».

En clair, la mise en cache des préfixes permet d’étendre à des suites de requêtes, alors que le KV Cache ne servait qu’à stocker temporairement des suites de tokens dans une même requête.

Cette approche est commercialement nommée « prompt caching » ou « context caching ». Elle est appliquée plus ou moins automatiquement.

Dans le principe, cette technique n’est efficace que si l’on ne modifie pas la première (grande) partie d’un prompt. Plus le prompt ou le contexte initial répété est long (comme c’est le cas pour le system prompt d’un chatbot, pour un agent IA, dans un mécanisme RAG, ou dans la révision de code), plus sa mise en cache est intéressante. Ces gains sont plus ou moins répercutés sur les clients finaux : les rabais vont de 50 % à 90 % du coût initial au million de tokens.

La plupart des fournisseurs proposent également des fonctions explicites de mise en cache des préfixes, parfois à un coût supérieur. Des balises permettent d’identifier des éléments d’un prompt à mettre en cache, d’étendre sa durée de vie (de 5 minutes à une heure en général) ou encore, d’appeler ces éléments depuis une autre session. Signalons que les fournisseurs surfacturent les tokens écrits en cache. La lecture, elle, est généralement facturée le dixième du prix de référence au million de tokens.

Dans une étude parue au mois de mai, Datadog considérait que la grande majorité de ses clients méconnaissaient ces fonctions de prompt caching.

Or, signale DDN, Nvidia a prouvé que les agents de programmation peuvent réutiliser « 85 à 97 % du contexte mis en cache après le premier appel et les agents multitours atteignant des taux de lecture réussie du cache de 99 % sur des prompts de 30 000 tokens ». ProjectDiscovery, une startup spécialisée dans les pentesting, explique qu’une gestion explicite de la mise en cache avec Claude d’Anthropic permet de faire augmenter le taux de lecture réussie de 7 % à 84 %, pour un coût inférieur de 59 %.

Néanmoins, cette mise en cache par préfixe n’est pas si simple à gérer pour les usagers finaux. Le moindre changement dans le prompt étiqueté provoque une invalidation du cache. Les préfixes légèrement différents peuvent s’entasser, tandis que la première requête et l’ensemble des essais pour obtenir le prompt qui sera exploité par défaut paient le coût complet du prefill.

De plus, le comportement des agents IA et des charges de travail exécutés par les LLM ne sont pas aussi prévisibles que le laissent entendre les évaluations de Nvidia. La manière dont les agents consomment leur contexte dépend aussi des outils qui les encadrent.

… et de ses alternatives

Des projets comme CacheBlend et LMCache introduisent un moyen d’appliquer la réutilisation à tous les blocs KV Cache, peu importe leur position dans un prompt. Cela reste un domaine actif de recherche, signale Philip Kiely. Il faut tout de même relancer le prefill : l’approche nécessite de recalculer une partie des entrées du KV Cache et le positionnement des embeddings (10-15 % du total des tokens dans le meilleur des cas).

 La prédiction, la compression adaptative du KV Cache selon l’importance des tokens, le partage du cache entre différentes variantes d’un LLM et la création d’une couche de contexte sémantique sont autant de moyens pour optimiser l’inférence. Reste que la majorité de ces initiatives sont encore à l’état exploratoire.

Sous le capot : moins compter sur les GPU que sur la RAM et les SSD

Cela ne règle pas le problème de la mémoire : pour justifier ces gains, les fournisseurs doivent pouvoir stocker plus efficacement ses longs prompts répétés.

Alors que la stratégie initiale consistait à stocker l’intégralité du KV Cache dans la mémoire vive vidéo (VRAM) de GPU, les équipementiers et certains fournisseurs entendent mettre en place jusqu’à quatre niveaux de stockage. Les blocs ou les pages de KV Cache peuvent être déchargés en mémoire vive (RAM) du système associé au GPU, sur des SSD locaux et des serveurs SSD partagés à travers le réseau.

Une autre approche consiste à créer un pool de KV Cache commun entre les serveurs d’inférence en s’appuyant sur le stockage SSD partagé.

Il faut tout de même les réinjecter en VRAM au moment de générer chaque nouveau token, ce qui génère de la latence. Selon DDN, la HBM serait 10 000 fois plus rapide que les (0,001 ms contre 10 ms) que les SSD en réseau. Nvidia ainsi que les spécialistes du stockage (Weka, DDN, Dell, etc.) tentent de s’entendre pour la réduire au maximum.

« Grâce à cette technique de déchargement et de récupération du cache, plusieurs avantages majeurs apparaissent », affirme Dylan Patel de SemiAnalysis. « D’abord, le délai avant la génération du premier token se réduit considérablement, puisqu’il n’est plus nécessaire de recalculer le contexte. Ensuite, une fois le cache récupéré et la génération lancée, le besoin de recalculer le prefill devient extrêmement rare ».

Pour Dylan Patel, cela veut aussi dire qu’il n’est pas nécessaire de s’appuyer sur les GPU Nvidia. Les instances AMD, mieux dotées en mémoire, peuvent s’en sortir avec certains LLM, juge-t-il. « Ou bien vous avez de moins bons GPU, mais une meilleure solution de stockage et vous pouvez obtenir un meilleur rapport performance-prix ».

Il n’est pas évident que tous les fournisseurs de LLM exploitent déjà les quatre niveaux de stockage. Par exemple, Philip Kiely de Baseten signale que l’interconnexion GPU-CPU proposée par les systèmes Nvidia GB200 en fait de meilleurs candidats pour un déchargement à deux niveaux (sur la DRAM et les SSD locaux) que d’autres serveurs. Toutefois, le fournisseur d’inférence dit utiliser elle-même la librairie Dynamo pour distribuer le KV Cache sur SSD locaux et en réseau.

Le 8 juillet dernier, Samsung a présenté une autre tactique qui l’arrangerait : utiliser CXL, sa technologie de partage de DRAM entre les serveurs pour y décharger le KV Cache plus efficacement.

OpenAI, Google et Anthropic utilisent de manière plus documentée des techniques comme le routage du KV cache à travers les répliques des serveurs d’inférence (en s’appuyant notamment sur la technologie NIXL de Nvidia). Il s’agit de distribuer les requêtes auprès des répliques qui contiennent le bon contenu en cache. Enfin, ils regroupent les requêtes d’inférence en lots similaires pour les traiter en parallèle.