# Fonctionnement de la mise en cache

Imaginez que vous êtes en train de discuter avec quelqu'un et que, au lieu de retenir ce qu'il a dit précédemment, il doit repartir de zéro à chaque fois que vous répondez. Ce serait lent et inefficace, n'est-ce pas ?

Vous pouvez étendre cette analogie aux modèles transformers. La génération par un modèle autorégressif peut être lente car elle effectue une prédiction un jeton à la fois. Chaque nouvelle prédiction dépend de tout le contexte précédent.

Pour prédire le 1000e jeton, le modèle a besoin d'informations issues des 999 jetons précédents. Ces informations sont représentées sous forme de multiplications de matrices entre les représentations des jetons.

Pour prédire le 1001e jeton, vous avez besoin des mêmes informations issues des 999 jetons précédents, en plus de toute information provenant du 1000e jeton. Cela représente un grand nombre de multiplications de matrices que le modèle doit calculer à nouveau pour chaque jeton !

Un cache clé-valeur (KV) élimine cette inefficacité en stockant les paires kv dérivées des couches d'attention des tokens précédemment traités. Les paires kv stockées sont récupérées depuis le cache et réutilisées pour les tokens suivants, évitant ainsi le besoin de les recalculer.

> [!WARNING]
> La mise en cache ne doit être utilisée que pour l'**inférence**. Son activation pendant l'entraînement peut provoquer des erreurs inattendues.

Pour mieux comprendre le fonctionnement et l'utilité de la mise en cache, examinons de plus près la structure des matrices d'attention.

## Matrices d'attention

L'**attention par produit scalaire mis à l'échelle** est calculée comme indiqué ci-dessous pour un lot de taille `b`, un nombre de têtes d'attention `h`, une longueur de séquence jusqu'à présent `T` et une dimension par tête d'attention `d_head`.

$$ \text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V $$

Les matrices de requête (`Q`), de clé (`K`) et de valeur (`V`) sont des projections des embeddings d'entrée de forme `(b, h, T, d_head)`.

Pour l'attention causale, le masque empêche le modèle de prêter attention aux jetons futurs. Une fois un jeton traité, sa représentation ne change plus par rapport aux jetons futurs, ce qui signifie que $ K_{\text{past}} $ et $ V_{\text{past}} $ peuvent être mis en cache et réutilisés pour calculer la représentation du dernier jeton.

$$ \text{Attention}(q_t, [\underbrace{k_1, k_2, \dots, k_{t-1}}_{\text{en cache}}, k_{t}], [\underbrace{v_1, v_2, \dots, v_{t-1}}_{\text{en cache}}, v_{t}]) $$

Lors de l'inférence, il ne faut que la requête du dernier jeton pour calculer la représentation $ x_t $ qui prédit le jeton suivant $ t+1 $. À chaque étape, les nouveaux vecteurs clé et valeur sont **stockés** dans le cache et **ajoutés** aux clés et valeurs passées.

$$ K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t) $$

L'attention est calculée indépendamment dans chaque couche du modèle, et la mise en cache est effectuée couche par couche.

Reportez-vous au tableau ci-dessous pour comparer comment le cache améliore l'efficacité.

| sans mise en cache | avec mise en cache |
|---|---|
| pour chaque étape, recalculer toutes les valeurs précédentes de `K` et `V`  | pour chaque étape, ne calculer que les `K` et `V` actuels |
| le coût de l'attention par étape est **quadratique** par rapport à la longueur de la séquence | le coût de l'attention par étape est **linéaire** par rapport à la longueur de la séquence (la mémoire augmente de manière linéaire, mais le calcul par jeton reste faible) |

## Classe de cache

Une interface de cache KV de base prend un tenseur de clé et un tenseur de valeur pour le jeton courant et renvoie les tenseurs `K` et `V` mis à jour. Cette gestion est assurée en interne par la méthode `forward` du modèle.

```py
new_K, new_V = cache.update(k_t, v_t, layer_idx)
attn_output = attn_layer_idx_fn(q_t, new_K, new_V)
```

Lorsque vous utilisez la classe [Cache](/docs/transformers/main/fr/internal/generation_utils#transformers.Cache) de Transformers, le module d'auto-attention effectue plusieurs étapes essentielles pour intégrer les informations passées et présentes.

1. Le module d'attention concatène les paires kv actuelles avec les paires kv passées stockées dans le cache. Cela crée des poids d'attention de forme `(new_tokens_length, past_kv_length + new_tokens_length)`. Les paires kv actuelles et passées sont essentiellement combinées pour calculer les scores d'attention, garantissant que le modèle prend en compte le contexte précédent et l'entrée actuelle.

2. Lorsque la méthode `forward` est appelée de manière itérative, il est essentiel que la forme du masque d'attention corresponde à la longueur combinée des paires kv passées et actuelles. Le masque d'attention doit avoir la forme `(batch_size, past_kv_length + new_tokens_length)`. Cela est généralement géré en interne par [generate()](/docs/transformers/main/fr/main_classes/text_generation#transformers.GenerationMixin.generate), mais si vous souhaitez implémenter votre propre boucle de génération avec [Cache](/docs/transformers/main/fr/internal/generation_utils#transformers.Cache), gardez ce point à l'esprit ! Le masque d'attention doit contenir les valeurs des jetons passés et actuels.

## Implémentation du stockage de cache

Les caches sont structurés sous forme d'une liste de couches, où chaque couche contient un cache de clés et un cache de valeurs. Les caches de clés et de valeurs sont des tenseurs de forme `[batch_size, num_heads, seq_len, head_dim]`.

Les couches peuvent être de différents types (par exemple `DynamicLayer`, `StaticLayer`, `StaticSlidingWindowLayer`), ce qui modifie principalement la manière dont la longueur de la séquence est gérée et dont le cache est mis à jour.

La plus simple est un `DynamicLayer` qui grandit à mesure que de nouveaux jetons sont traités. La dimension de la longueur de séquence (`seq_len`) augmente avec chaque nouveau jeton :

```py
cache.layers[idx].keys = torch.cat([cache.layers[idx].keys, key_states], dim=-2)
cache.layers[idx].values = torch.cat([cache.layers[idx].values, value_states], dim=-2)
```

D'autres types de couches, comme `StaticLayer` et `StaticSlidingWindowLayer`, ont une longueur de séquence fixe définie lors de la création du cache. Cela les rend compatibles avec `torch.compile`. Dans le cas de `StaticSlidingWindowLayer`, les jetons existants sont décalés hors du cache lorsqu'un nouveau jeton est ajouté.

L'exemple ci-dessous montre comment créer une boucle de génération avec [DynamicCache](/docs/transformers/main/fr/internal/generation_utils#transformers.DynamicCache). Comme discuté, le masque d'attention est une concaténation des valeurs des jetons passés et actuels.

```py
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache
from accelerate import Accelerator

device = Accelerator().device

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

past_key_values = DynamicCache(config=model.config)
messages = [{"role": "user", "content": "Hello, what's your name."}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)

generated_ids = inputs.input_ids
max_new_tokens = 10

for _ in range(max_new_tokens):
    outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
    # Greedily sample one next token
    next_token_ids = outputs.logits[:, -1:].argmax(-1)
    generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1)
    # Prepare inputs for the next generation step by leaving unprocessed tokens, in our case we have only one new token
    # and expanding attn mask for the new token, as explained above
    attention_mask = inputs["attention_mask"]
    attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
    inputs = {"input_ids": next_token_ids, "attention_mask": attention_mask}

print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])
"[INST] Hello, what's your name. [/INST]  Hello! My name is LLaMA,"
```