Stratégies de cache

Les vecteurs clé-valeur (KV) sont utilisés pour calculer les scores d'attention. Pour les modèles autorégressifs, les scores KV sont calculés à chaque itération, car le modèle prédit un jeton à la fois. Chaque prédiction dépend des jetons précédents, ce qui signifie que le modèle effectue les mêmes calculs à chaque fois.

Un cache KV stocke ces calculs afin qu'ils puissent être réutilisés sans avoir à les recalculer. Un cache efficace est crucial pour optimiser les performances du modèle, car il réduit le temps de calcul et améliore les taux de réponse. Consultez la documentation Caching pour une explication plus détaillée du fonctionnement d'un cache.

Transformers propose plusieurs classes Cache qui implémentent différents mécanismes de mise en cache. Certaines de ces classes Cache sont optimisées pour économiser la mémoire, tandis que d'autres sont conçues pour maximiser la vitesse de génération. Reportez-vous au tableau ci-dessous pour comparer les types de cache et l'utiliser afin de sélectionner le cache le plus adapté à votre cas d'usage.

Type de cache	Prend en charge les couches glissantes	Prend en charge le déchargement	Prend en charge torch.compile()	Mémoire attendue
Cache dynamique	Oui	Oui	Non	Moyen
Cache statique	Oui	Oui	Oui	Élevé
Cache quantifié	Non	Non	Non	Faible

Ce guide vous présente les différentes classes Cache et vous montre comment les utiliser pour la génération.

Cache par défaut

Le DynamicCache est la classe de cache par défaut pour tous les modèles. Il permet à la taille du cache de croître dynamiquement afin de stocker un nombre croissant de clés et de valeurs au fur et à mesure que la génération progresse.

Notez que pour les modèles utilisant l’attention à fenêtre glissante (Mistral, Gemma2,…) ou l’attention par morceaux (Llama4), le cache cessera de croître lorsque les couches utilisant ces types d’attention auront atteint leur taille maximale (la taille de la fenêtre glissante ou du morceau).

Désactivez le cache en configurant use_cache=False dans generate().

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

model.generate(**inputs, do_sample=False, max_new_tokens=20, use_cache=False)

Les classes de cache peuvent également être initialisées avant d’être appelées et passées au paramètre past_key_values des modèles. Cela peut être utile pour un contrôle plus fin ou pour des usages plus avancés, tels que la mise en cache du contexte.

Dans la plupart des cas, il est plus simple de définir la stratégie de cache dans le paramètre cache_implementation.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

past_key_values = DynamicCache(config=model.config)
out = model.generate(**inputs, do_sample=False, max_new_tokens=20, past_key_values=past_key_values)

Cache de taille fixe

Le DynamicCache par défaut vous empêche de tirer parti de la plupart des optimisations just-in-time (JIT), car la taille du cache n’est pas fixe. Les optimisations JIT vous permettent de minimiser la latence au détriment de l’utilisation de la mémoire. Tous les types de cache suivants sont compatibles avec des optimisations JIT telles que torch.compile pour accélérer la génération.

Un cache de taille fixe (StaticCache) pré-alloue une taille de cache maximale spécifique pour les paires kv. Vous pouvez générer jusqu’à la taille de cache maximale sans avoir besoin de la modifier. Cependant, le fait d’avoir une taille fixe (généralement grande) pour les états clé/valeur signifie qu’au cours de la génération, de nombreux jetons seront en réalité masqués car ils ne devraient pas participer à l’attention. Cette astuce permet donc de compile facilement l’étape de décodage, mais elle entraîne un gaspillage de jetons dans le calcul de l’attention. Comme pour tout, il s’agit d’un compromis qui sera très avantageux si vous générez avec plusieurs séquences de longueurs plus ou moins similaires, mais qui peut s’avérer sous-optimal si vous avez par exemple une très grande séquence, suivie uniquement de séquences courtes (la taille fixe du cache étant grande, beaucoup serait gaspillé pour les séquences courtes). Assurez-vous de bien comprendre l’impact si vous l’utilisez !

En ce qui concerne DynamicCache, notez que pour les modèles utilisant l'attention à fenêtre glissante (Mistral, Gemma2,…) ou l'attention par morceaux (Llama4), le cache ne dépassera jamais la taille de la fenêtre glissante/des morceaux sur les couches utilisant ces types d'attention, même si la longueur maximale spécifiée est supérieure.

Vous pouvez activer StaticCache en configurant cache_implementation="static" dans generate(). Cela activera également la compilation automatique de l'étape de décodage pour les stratégies de décodage glouton et par échantillonnage.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("Hello, my name is", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="static")
tokenizer.batch_decode(out, skip_special_tokens=True)[0]
"Hello, my name is [Your Name], and I am a [Your Profession] with [Number of Years] of"

Garder les tenseurs de génération sur le CPU

Les backends de compilation tels que Neuron et TPU tracent votre modèle en un graphe de calcul fixe. La boucle de génération maintient des tenseurs (séquence de sortie, attention_mask, position_ids) qui augmentent d’un jeton à chaque étape. Le compilateur retrace le graphe chaque fois que la forme de ces tenseurs change sur l’accélérateur, ce qui ralentit la génération.

generate() ne déplace que les tenseurs consommés par forward vers le périphérique du modèle, juste avant chaque appel à forward. La sortie est ensuite déplacée pour correspondre au périphérique de l’entrée. Passez vos entrées sur CPU afin de maintenir la gestion des tenseurs croissants de la boucle hors de l’accélérateur. Le graphe compilé reste stable, et comme la sortie suit le périphérique de l’entrée, la sortie générée reste sur CPU.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", device_map="auto")

# leave the inputs on CPU instead of calling .to(model.device).
inputs = tokenizer("The French Bread Law states", return_tensors="pt")

# generate runs forward on the model device but returns the output on the input_ids device.
output = model.generate(**inputs, do_sample=False, max_new_tokens=20)
print(output.device)
cpu

Déchargement du cache

Le cache KV peut occuper une part importante de la mémoire et devenir un goulot d'étranglement pour la génération de contexte long. Les caches optimisés pour la mémoire privilégient une réduction de l'empreinte mémoire au détriment de la vitesse. C'est particulièrement important pour les grands modèles de langage (LLM) et si votre matériel est limité en mémoire.

Le déchargement du cache permet d’économiser de la mémoire GPU en déplaçant le cache KV des couches du modèle, à l’exception d’une seule, vers le CPU. Seule la couche courante est maintenue sur le GPU pendant l’itération forward du modèle sur les couches. Le cache de la couche suivante est préchargé de manière asynchrone, et le cache de la couche courante est renvoyé vers le CPU après le calcul d’attention.

Vous pouvez envisager l’offloading si vous disposez d’un GPU de petite taille et que vous rencontrez des erreurs de mémoire insuffisante (OOM).

Vous remarquerez peut-être une légère dégradation du débit de génération par rapport à un cache entièrement sur l’appareil, en fonction de votre modèle et de vos choix de génération (taille du contexte, nombre de jetons générés, nombre de faisceaux, etc.). Cela est dû au fait que le déplacement des états clé/valeur de part et d’autre nécessite un certain travail.

Le déchargement (offloading) est disponible pour DynamicCache et StaticCache. Vous pouvez l’activer en configurant cache_implementation="offloaded" pour la version dynamique, ou cache_implementation="offloaded_static" pour la version statique, dans GenerationConfig ou generate(). De plus, vous pouvez instancier votre propre DynamicCache ou StaticCache avec l’option offloading=True, et passer ce cache dans generate ou la méthode forward de votre modèle (par exemple, past_key_values=DynamicCache(config=model.config, offloading=True) pour un cache dynamique).

Notez que les deux classes Cache mentionnées ci-dessus disposent d’une option supplémentaire lors de leur instanciation directe, offload_only_non_sliding. Cet argument supplémentaire détermine si les couches utilisant l’attention à fenêtre glissante ou par morceaux (le cas échéant) seront également déchargées. Étant donné que ces couches sont généralement courtes de toute façon, il peut être préférable d’éviter de les décharger, car le déchargement peut entraîner une pénalité en termes de vitesse. Par défaut, cette option est False pour DynamicCache, et True pour StaticCache.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, dtype=torch.float16, device_map="auto")
inputs = tokenizer("Fun fact: The shortest", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=23, cache_implementation="offloaded")
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
Fun fact: The shortest war in history was between Britain and Zanzibar on August 27, 1896.

L'exemple ci-dessous montre comment vous pouvez basculer sur un cache déchargé si vous épuisez la mémoire :

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from accelerate import Accelerator

def resilient_generate(model, *args, **kwargs):
    oom = False
    device = Accelerator().device
    torch_device_module = getattr(torch, device, torch.cuda)
    try:
        return model.generate(*args, **kwargs)
    except torch.OutOfMemoryError as e:
        print(e)
        print("retrying with cache_implementation='offloaded'")
        oom = True
    if oom:
        torch_device_module.empty_cache()
        kwargs["cache_implementation"] = "offloaded"
        return model.generate(*args, **kwargs)

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, dtype=torch.float16, device_map="auto")
prompt = ["okay "*1000 + "Fun fact: The most"]
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
beams = { "num_beams": 40, "num_return_sequences": 20, "max_new_tokens": 23, "early_stopping": True, }
out = resilient_generate(model, **inputs, **beams)
responses = tokenizer.batch_decode(out[:,-28:], skip_special_tokens=True)

Cache quantifié

QuantizedCache réduit les besoins en mémoire en quantifiant les valeurs KV à une précision inférieure. QuantizedCache prend actuellement en charge trois backends de quantification :

hqq prend en charge les types de données int1, int2, int3, int4 et int8.
quanto prend en charge les types de données int2 et int4. C'est le backend de quantification par défaut.
fp8 stocke les valeurs KV en virgule flottante sur 8 bits (float8_e4m3fn). Il n’exige aucune dépendance supplémentaire et quantise par tenseur plutôt que par groupe, ce qui le rend incompatible avec nbits, axis_key, axis_value, q_group_size et residual_length.

La quantification du cache peut nuire à la latence si la longueur du contexte est courte et qu’il y a suffisamment de mémoire GPU disponible pour la génération sans activer la quantification du cache. Essayez de trouver un équilibre entre l’efficacité mémoire et la latence.

Activez QuantizedCache en configurant cache_implementation="quantized" dans GenerationConfig. Le backend de quantification, ainsi que les paramètres de quantification supplémentaires, sont spécifiés dans le dictionnaire cache_config passé à generate(). Vous devriez utiliser les valeurs par défaut pour ces paramètres supplémentaires, sauf si vous rencontrez des erreurs de mémoire insuffisante. Dans ce cas, envisagez de réduire la longueur des résidus.

Pour le backend hqq, nous recommandons de définir les paramètres axis_key et axis_value sur 1.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(
    **inputs,
    do_sample=False,
    max_new_tokens=20,
    cache_implementation="quantized",
    cache_config={
        "nbits": 4,
        "backend": "hqq",
        "axis_key": 1,
        "axis_value": 1
    }
)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's loud and energetic. It's a great way to express myself and rel

Pour le backend quanto, nous recommandons de définir les paramètres axis_key et axis_value sur 0.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(
    **inputs,
    do_sample=False,
    max_new_tokens=20,
    cache_implementation="quantized",
    cache_config={
        "nbits": 4,
        "backend": "quanto",
        "axis_key": 0,
        "axis_value": 0
    }
)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's loud and energetic. It's a great way to express myself and rel

Le backend fp8 n’accepte aucun paramètre supplémentaire et lève une erreur si l’un des paramètres group-wise lui est fourni.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B-Instruct-2507")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B-Instruct-2507", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="quantized", cache_config={"backend": "fp8"})
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it is powerful and emotional. I also like classical music because it is elegant and sophisticated. I enjoy

Cache encodeur-décodeur

EncoderDecoderCache est conçu pour les modèles encodeur-décodeur. Il gère à la fois les caches d'auto-attention et d'attention croisée afin d'assurer le stockage et la récupération des paires kv précédentes. Il est possible de définir un type de cache différent pour l'encodeur et le décodeur.

Ce type de cache ne nécessite aucune configuration. Il s’agit d’un simple wrapper autour de 2 Cache tels que décrits ci-dessus, qui seront utilisés indépendamment et directement par le modèle.

Caches spécifiques aux modèles

Certains modèles disposent d'une méthode de stockage des paires kv passées ou des états qui n'est compatible avec aucune autre classe de cache.

Les modèles Mamba, tels que Mamba, nécessitent un cache spécifique car le modèle ne dispose ni d’un mécanisme d’attention ni d’états kv. Ils ne sont donc pas compatibles avec les classes Cache ci-dessus.

Génération itérative

Un cache peut également fonctionner dans des configurations de génération itérative impliquant une interaction aller-retour avec un modèle (chatbots). Tout comme pour la génération standard, la génération itérative avec un cache permet à un modèle de gérer efficacement des conversations en cours sans recalculer l'intégralité du contexte à chaque étape.

Pour une génération itérative avec un cache, commencez par initialiser une classe de cache vide, puis vous pouvez y injecter vos nouveaux prompts. Suivez l’historique de la conversation à l’aide d’un chat template.

L’exemple suivant illustre Llama-2-7b-chat-hf. Si vous utilisez un autre modèle de type chat, apply_chat_template() peut traiter les messages différemment. Il peut supprimer des jetons importants selon la façon dont le modèle Jinja est rédigé. Pour les modèles de chat multimodaux, consultez le guide iterative chatting with cache pour savoir comment traiter les images ou l’audio.

Par exemple, certains modèles utilisent des jetons spéciaux <think> ... </think> pendant le raisonnement. Ils pourraient être perdus lors de la ré-encodage, provoquant des problèmes d’indexation. Il peut être nécessaire de supprimer ou d’ajuster manuellement les jetons supplémentaires des complétions pour maintenir la stabilité.

import torch
from transformers import AutoTokenizer,AutoModelForCausalLM, DynamicCache, StaticCache

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained(model_id)

user_prompts = ["Hello, what's your name?", "Btw, yesterday I was on a rock concert."]

past_key_values = DynamicCache(config=model.config)

messages = []
for prompt in user_prompts:
    messages.append({"role": "user", "content": prompt})
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
    input_length = inputs["input_ids"].shape[1]
    outputs = model.generate(**inputs, do_sample=False, max_new_tokens=256, past_key_values=past_key_values)
    completion = tokenizer.decode(outputs[0, input_length: ], skip_special_tokens=True)
    messages.append({"role": "assistant", "content": completion})

Recadrer un cache

Utilisez crop() pour annuler des jetons depuis un cache, par exemple lorsqu’une continuation candidate est rejetée lors d’une génération itérative. Passez le nombre de jetons à supprimer en tant qu’entier négatif.

# Remove the last three tokens from the cache.
past_key_values.crop(-3)

Les couches d’attention à fenêtre glissante et d’attention linéaire abandonnent les états passés au fil de l’exécution, il n’y a donc par défaut aucun état auquel revenir. Appelez activate_past_recording() avant le passage avant pour qu’elles conservent ces états. Sans cela, crop lève une RuntimeError dès qu’une couche à fenêtre glissante remplit sa fenêtre, et immédiatement pour les couches d’attention linéaire.

# Enable this before the forward pass if you need to roll back tokens.
past_key_values.activate_past_recording()

crop(0) ne vide pas le cache. Il ne supprime aucun jeton de la séquence mise en cache et constitue une opération nulle pour les couches d’attention complète. Pour les couches à fenêtre glissante et à attention linéaire, il peut éliminer des informations mises en cache que l’appel suivant du modèle n’utilisera pas. Cela réduit l’utilisation de la mémoire sans annuler les jetons.

                    cached positions
                    ┌───┬───┬───┬───┬───┬───┐
before crop(-3)     │ 0 │ 1 │ 2 │ 3 │ 4 │ 5 │
                    └───┴───┴───┴───┴───┴───┘
                                ╰─────┬─────╯
                                 rolled back

after crop(-3)
                    ┌───┬───┬───┐
full attention      │ 0 │ 1 │ 2 │    keeps the whole prefix
                    └───┴───┴───┘
                            ┌───┐
sliding_window=2            │ 2 │    keeps sliding_window - 1
                            └───┘

Les deux couches signalent désormais une longueur de séquence de 3, et la génération reprend à partir de la position 3.

Passer 0 ne supprime aucun jeton de la séquence. C’est une opération sans effet pour les couches d’attention complète, mais pour les couches à fenêtre glissante et à attention linéaire, elle supprime toujours les états enregistrés que le prochain passage en avant ne lira pas, ce qui libère de la mémoire sans annuler quoi que ce soit.

Préremplir un cache (cache de préfixe)

Dans certaines situations, vous pouvez vouloir remplir un Cache avec des paires kv pour un certain préfixe de prompt et le réutiliser pour générer différentes séquences.

L'exemple ci-dessous initialise un StaticCache, puis met en cache un prompt initial. Vous pouvez maintenant générer plusieurs séquences à partir du prompt pré-rempli.

import copy
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, StaticCache

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map={"": 0})
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Init StaticCache with big enough max-length (1024 tokens for the below example)
# You can also init a DynamicCache, if that suits you better
prompt_cache = StaticCache(config=model.config, max_cache_len=1024)

INITIAL_PROMPT = "You are a helpful assistant. "
inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to(model.device.type)
# This is the common prompt cached, we need to run forward without grad to be able to copy
with torch.no_grad():
     prompt_cache = model(**inputs_initial_prompt, past_key_values = prompt_cache).past_key_values

prompts = ["Help me to write a blogpost about travelling.", "What is the capital of France?"]
responses = []
for prompt in prompts:
    new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to(model.device.type)
    past_key_values = copy.deepcopy(prompt_cache)
    outputs = model.generate(**new_inputs, past_key_values=past_key_values,max_new_tokens=20)
    response = tokenizer.batch_decode(outputs)[0]
    responses.append(response)

print(responses)

Update on GitHub