Le KV Cache : Pourquoi les LLMs Génèrent Vite (et Mangent la Mémoire)

Quand vous discutez avec un LLM, le modèle génère un token, puis le suivant, puis le suivant. À chaque étape, il doit tenir compte de tout ce qui précède. Fait naïvement, ce processus devient de plus en plus lent à mesure que la réponse s'allonge. Le KV cache est l'astuce qui empêche ça, et c'est aussi la première chose qui sature la mémoire d'un GPU en production.

Pourquoi il faut un cache

La génération est autorégressive

Un LLM ne produit pas une réponse d'un coup. Il produit un token, l'ajoute à la séquence, et recommence avec la séquence allongée.

À chaque étape, le mécanisme d'attention compare le nouveau token à tous les tokens précédents. Sans précaution, l'étape 100 refait tout le travail des étapes 1 à 99.

Ce que ça coûte sans cache

Pour générer une réponse de N tokens, le nombre de calculs d'attention ressemble à 1 + 2 + 3 + ... + N, soit environ N²/2. Pour 1 000 tokens, c'est 500 000 comparaisons au lieu de 1 000.

Ce qu'on met en cache exactement

Rappel : Q, K et V

Dans une couche d'attention, chaque token est projeté en trois vecteurs :

• Query (Q) : « qu'est-ce que je cherche ? »
• Key (K) : « qu'est-ce que je contiens ? »
• Value (V) : « qu'est-ce que je transmets si on me sélectionne ? »

Le nouveau token calcule sa Query, la compare aux Keys de tous les tokens précédents, et récupère une moyenne pondérée de leurs Values.

L'observation clé

Dans un modèle causal (GPT, Llama, Mistral), un token ne voit que ceux qui le précèdent. Donc les K et V d'un token déjà traité ne changent jamais quand on ajoute des tokens après lui.

D'où le nom : on garde les K et les V, pour chaque token, pour chaque couche.

Combien de mémoire ça prend

La formule

Pour chaque token, le cache stocke un K et un V par couche, chacun composé de plusieurs têtes :

mémoire par token = 2 × nb_couches × nb_têtes_KV × dim_tête × taille_dtype

Le 2 correspond à K + V. La taille du dtype est de 2 octets en FP16/BF16.

Deux exemples chiffrés

Les deux modèles ont une taille comparable, mais Llama 3 utilise 4 fois moins de mémoire de cache. La différence vient du nombre de têtes KV, on y revient plus bas.

Micro-exercice 1 : calculer le cache d'un modèle depuis sa config

Le fichier config.json de n'importe quel modèle Hugging Face contient tout ce qu'il faut.

from transformers import AutoConfig

def taille_kv_cache(nom_modele: str, nb_tokens: int, octets_dtype: int = 2) -> float:
    """Retourne la taille du KV cache en Go pour une séquence donnée."""
    cfg = AutoConfig.from_pretrained(nom_modele)

    nb_couches = cfg.num_hidden_layers
    # Si le modèle utilise GQA/MQA, num_key_value_heads < num_attention_heads
    nb_tetes_kv = getattr(cfg, "num_key_value_heads", cfg.num_attention_heads)
    dim_tete = cfg.hidden_size // cfg.num_attention_heads

    par_token = 2 * nb_couches * nb_tetes_kv * dim_tete * octets_dtype
    return par_token * nb_tokens / 1e9

for modele in ["meta-llama/Llama-2-7b-hf", "meta-llama/Meta-Llama-3-8B", "Qwen/Qwen2.5-7B"]:
    go = taille_kv_cache(modele, nb_tokens=8192)
    print(f"{modele:35s} → {go:.2f} Go pour 8k tokens")

Voir le cache en action

Micro-exercice 2 : mesurer l'effet du cache

Avec Hugging Face Transformers, le cache est activé par défaut. On peut le désactiver pour comparer. Le modèle gpt2 tourne sur CPU sans problème.

import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2").eval()

entrees = tokenizer("Le KV cache est", return_tensors="pt")

def generer(use_cache: bool, nb_tokens: int = 200) -> float:
    debut = time.perf_counter()
    with torch.no_grad():
        model.generate(
            **entrees,
            max_new_tokens=nb_tokens,
            do_sample=False,
            use_cache=use_cache,
            pad_token_id=tokenizer.eos_token_id,
        )
    return time.perf_counter() - debut

print(f"Sans cache : {generer(False):.2f} s")
print(f"Avec cache : {generer(True):.2f} s")

Inspecter le contenu du cache

with torch.no_grad():
    sortie = model(**entrees, use_cache=True)

cache = sortie.past_key_values  # objet DynamicCache dans les versions récentes
k_couche0, v_couche0 = cache[0]  # première couche

# Forme : (batch, nb_têtes, nb_tokens, dim_tête)
print("K couche 0 :", tuple(k_couche0.shape))
print("V couche 0 :", tuple(v_couche0.shape))
print("Nombre de couches en cache :", len(cache))

On retrouve exactement les dimensions de la formule : une entrée par couche, avec nb_têtes × nb_tokens × dim_tête valeurs pour K et autant pour V.

Prefill et decode : deux phases très différentes

Le KV cache explique pourquoi une requête LLM se déroule en deux temps.

En decode, le GPU passe l'essentiel de son temps à lire les poids et le cache, pas à calculer. C'est pour ça qu'un GPU seul génère pour un utilisateur à quelques pour cent de sa capacité de calcul, et c'est ce que le cours sur le continuous batching cherche à corriger.

Réduire la taille du cache

Comme le cache est la ressource limitante, toute l'industrie a travaillé à le compresser.

Moins de têtes KV : MQA et GQA

Dans l'attention classique (MHA), chaque tête de Query a sa propre tête K et V. Grouped-Query Attention partage une paire K/V entre plusieurs Queries. Multi-Query Attention pousse à l'extrême : une seule paire pour toutes.

C'est la raison de l'écart de 4× entre Llama 2 et Llama 3 vu plus haut. La qualité baisse très peu avec GQA, ce qui en fait le standard actuel.

Les autres leviers

Le prefix caching en pratique

Si toutes vos requêtes commencent par le même system prompt de 2 000 tokens, le prefill de ces 2 000 tokens est identique à chaque fois. Le prefix caching le calcule une fois et le partage.

# vLLM : activer le partage de préfixe
from vllm import LLM, SamplingParams

llm = LLM(model="Qwen/Qwen2.5-7B-Instruct", enable_prefix_caching=True)

system = "Tu es un assistant juridique. " * 200  # long préfixe commun
prompts = [system + f"Question {i} : que dit le Code civil sur le bail ?" for i in range(20)]

sorties = llm.generate(prompts, SamplingParams(max_tokens=64))

La première requête paie le prefill complet. Les 19 suivantes réutilisent le cache du préfixe et ne calculent que la partie qui diffère.

Les pièges à éviter

1. Dimensionner sur la longueur du prompt seul

La mémoire du cache dépend de prompt + réponse, et de nombre de requêtes simultanées. Un modèle qui tient pour un utilisateur à 4k tokens ne tiendra pas pour dix.

2. Oublier que le cache grossit pendant la génération

Une requête qui commence avec 500 tokens de cache peut en avoir 4 500 après une longue réponse. Les moteurs de serving réservent ou préemptent pour ça, un script maison ne le fait pas.

3. Confondre KV cache et prompt caching

Le KV cache existe dans toutes les générations, c'est un mécanisme interne du modèle. Le prompt caching (ou prefix caching) est une optimisation au-dessus, qui réutilise ce cache entre requêtes. Les API commerciales facturent le second, pas le premier.

4. Beam search et cache

Avec beam search, chaque faisceau a son propre cache. Quatre faisceaux, quatre fois la mémoire. Sur un GPU serré, c'est souvent ce qui provoque l'OOM plutôt que le modèle lui-même.

À retenir

Pour aller plus loin

• Exercice bonus : reprenez le micro-exercice 1 et tracez la mémoire du cache en fonction de la longueur de contexte pour trois modèles. À partir de quelle longueur un GPU de 24 Go ne tient plus le modèle 7B en FP16 ?
• Cours suivant : le continuous batching, ou comment servir des dizaines d'utilisateurs sur un GPU en exploitant le fait que le decode est limité par la mémoire.
• Cours lié : gérer le context window et la mémoire des LLMs, pour la partie prompt caching côté API.