Optimisation du KV Cache dans les LLMs

Le KV cache (cache des vecteurs Key et Value) est une technique d'optimisation essentielle pour accélérer l'inférence des grands modèles de langage. Elle peut diviser le temps d'inférence par cinq, mais elle introduit des défis significatifs en termes de consommation mémoire. Cet article explore son fonctionnement, ses limites, et les pistes d'optimisation au niveau du moteur d'inférence, y compris une approche originale exploitant la microarchitecture CPU.

Comprendre le KV Cache

Le mécanisme de base des LLMs

Les LLMs utilisent un transformer pour produire des états cachés (hidden states) à partir des tokens d'entrée. Ces états sont projetés dans l'espace du vocabulaire pour générer des logits, et seul le logit du dernier token est utilisé pour prédire le suivant. Ce processus se répète pour chaque nouveau token généré.

Le rôle de l'attention

Dans le mécanisme d'attention, la génération d'un nouveau token ne nécessite que le vecteur query du dernier token et tous les vecteurs key et value précédents. Point crucial : les vecteurs key et value des tokens précédents restent constants une fois calculés.

Le principe du KV Cache

Plutôt que de recalculer les vecteurs key et value pour chaque token précédent à chaque étape, ils sont stockés dans un cache. Pour générer un nouveau token, on calcule uniquement le vecteur key/value du token précédent et on récupère les autres depuis le cache, réduisant ainsi le temps de calcul.

C'est ce qui explique pourquoi des modeles comme ChatGPT prennent plus de temps pour le premier token (le fameux "time-to-first-token" ou TTFT, correspondant au calcul du cache KV pour le prompt) et sont plus rapides ensuite.

L'ampleur du problème mémoire

Pour un modèle comme Llama3-70B (80 couches, taille cachée de 8k, max 4k tokens), chaque token occupe environ 2,5 Mo en cache, soit 10,5 Go pour 4k tokens. Avec plusieurs utilisateurs simultanés, la mémoire GPU est vite saturée.

Optimisations au niveau du moteur d'inférence

Il est important de distinguer les optimisations liées au modèle (comme FlashAttention ou Grouped-Query Attention) de celles liées au moteur d'inférence. Les optimisations suivantes concernent spécifiquement le moteur, c'est-à-dire la gestion du KV cache, son allocation mémoire, et son utilisation pour accélérer la génération de tokens.

PagedAttention

Le moteur divise le KV cache en pages de taille fixe, allouées dynamiquement selon les besoins. Les pages inutilisées peuvent être déchargées vers la mémoire CPU ou le disque, tandis que seules les pages actives restent en mémoire GPU. Cette technique, utilisée notamment dans vLLM, est idéale pour des applications avec des prompts longs ou des sessions multi-utilisateurs.

RadixAttention

Le moteur utilise une structure de données basée sur des arbres radix pour organiser et accéder rapidement aux vecteurs KV, réduisant les coûts de recherche dans le cache. Cette approche améliore la vitesse d'accès au cache, surtout pour des contextes très longs.

Déchargement partiel (Offloading)

Le moteur déplace les parties moins utilisées du KV cache (par exemple les vecteurs de tokens anciens) vers la mémoire CPU ou le disque, ne conservant en GPU que les données critiques. Cette technique est supportée par des frameworks comme HuggingFace Accelerate, DeepSpeed-Inference et FlexGen.

Gestion dynamique de la taille du cache

Le moteur ajuste la taille du KV cache en temps réel en fonction de la longueur du contexte et de la disponibilité mémoire, en supprimant les entrées obsolètes. Cela évite les dépassements de mémoire tout en priorisant les données pertinentes.

Distribution multi-GPU

Le moteur répartit le KV cache entre plusieurs GPU, synchronisant les accès pour générer des tokens en parallèle ou gérer des contextes massifs. Cette approche, utilisée dans TensorRT-LLM de NVIDIA, permet d'échelonner les performances avec des contextes de centaines de milliers de tokens.

Une approche originale : exploiter la microarchitecture CPU

Au-delà des techniques classiques, une piste d'optimisation prometteuse consiste à repenser la façon dont le moteur d'inférence utilise le matériel, en particulier sur CPU.

Le problème des moteurs actuels

Les moteurs d'inférence actuels ont une gestion du cache relativement directe :

Charger le modèle en mémoire
Pour chaque requête : charger le résultat du tokenizer, appliquer les vecteurs aux matrices KMV jusqu'au prochain token
Répéter

À chaque itération, les poids du modèle sont potentiellement évictés du cache CPU (L1/L2/L3), ce qui génère des cache misses coûteux (100 à 200 cycles par accès RAM contre 1 à 2 cycles pour un accès cache).

Traitement simultané de requêtes multiples

L'idée est de traiter 3 ou 4 requêtes simultanément pour amortir le coût des évictions de cache. Les poids du modèle, qui sont communs à toutes les requêtes, restent en cache plus longtemps lorsqu'ils sont réutilisés pour plusieurs requêtes consécutives. Sur des milliards d'opérations, cette optimisation peut se révéler significative.

Répartition des poids sur CPU multicœur

Plutôt que de charger l'intégralité du modèle sur chaque cœur, on peut diviser les poids en blocs et les assigner à différents cœurs. Chaque cœur conserve sa portion des poids en cache L1/L2, et les données intermédiaires (activations, KV cache) transitent d'un cœur à l'autre. Cela maximise l'utilisation du cache privé de chaque cœur et réduit les conflits dans le cache partagé L3.

Utilisation d'AVX512

Les instructions SIMD AVX512 permettent de traiter les vecteurs de plusieurs requêtes simultanément dans un seul registre de 512 bits. Combinées avec le pre-fetching (_mm_prefetch) pour anticiper les données en cache, ces instructions peuvent diviser le temps de calcul de manière significative.

Pages géantes pour le TLB

Sur Linux, l'utilisation de huge pages de 1 Go (via l'option hugepagesz=1G au boot du noyau) permet de réduire drastiquement les évictions de cache TLB. Pour un modèle de 10,5 Go, on passe de 2,6 millions d'entrées TLB (pages de 4 Ko) à seulement 10-11 entrées (pages de 1 Go), ce qui élimine quasiment les TLB misses.

Configuration du noyau Linux :

GRUB_CMDLINE_LINUX_DEFAULT="quiet splash hugepagesz=1G hugepages=10"

En Rust, l'allocation se fait via mmap avec le flag MAP_HUGETLB :

use libc::{mmap, MAP_HUGETLB, PROT_READ, PROT_WRITE,
           MAP_PRIVATE, MAP_ANONYMOUS};

let size = 1 << 30; // 1 Go
let ptr = unsafe {
    mmap(
        std::ptr::null_mut(),
        size,
        PROT_READ | PROT_WRITE,
        MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB,
        -1,
        0,
    )
};

Mesurer l'impact

Pour valider ces optimisations, plusieurs metriques sont essentielles :

Cache misses : perf stat -e cache-misses pour comparer avec et sans batching
TLB misses : perf stat -e tlb_load_misses pour mesurer l'impact des huge pages
TTFT (Time-to-first-token) : temps avant le premier token généré
Throughput : nombre de tokens générés par seconde

L'approche idéale est de prototyper avec un modèle réduit (100M paramètres) sur un CPU multicœur avec AVX512, puis d'extrapoler vers des modèles plus grands.

Conclusion

Le KV cache est indispensable pour l'inférence des LLMs, mais sa gestion optimale reste un défi ouvert. Les techniques classiques (PagedAttention, offloading, distribution multi-GPU) répondent aux problèmes de mémoire à grande échelle. L'exploitation fine de la microarchitecture CPU, via le batching intelligent, la répartition des poids sur les cœurs, les instructions AVX512 et les huge pages, ouvre une voie complémentaire pour des gains significatifs, en particulier sur des déploiements CPU où les ressources sont contraintes. Ces optimisations, codées en Rust et assembleur, restent à valider par des benchmarks rigoureux, mais le potentiel est réel.