Comment la cache KV est-elle implémentée lors de l'inférence d'un transformeur autoregressif?

Dans les transformateurs autorégressifs, lors de l'inférence, il est important d'utiliser efficacement les calculs pour ne pas recalculer toute la séquence à chaque fois. Pour cela, on utilise le cache KV, qui stocke les clés (Key) et les valeurs (Value) des couches d'attention (auto-attention).

L'idée est qu'en générant le token suivant, nous avons déjà calculé les clés et les valeurs pour tous les tokens précédents. Au lieu de les recalculer, nous les stockons en mémoire (cache) et, lors du calcul de l'attention pour le nouveau token, nous utilisons ce cache.

Cela fonctionne approximativement comme suit :

• Pour chaque couche du transformateur, lors du traitement d'un token, les clés et les valeurs sont stockées dans des buffers séparés.
• Lors de la génération du token suivant, seules la clé et la valeur pour ce token sont calculées, puis combinées avec celles déjà stockées.
• Le mécanisme d'attention utilise l'ensemble combiné de clés et de valeurs pour calculer le contexte.

Cela accélère considérablement l'inférence, en réduisant la complexité computationnelle de quadratique en longueur de séquence à linéaire.

Exemple de pseudocode pour une couche :

# kv_cache stocke les clés et valeurs des tokens précédents
new_key, new_value = compute_kv(new_token)
kv_cache.keys = concatenate(kv_cache.keys, new_key)
kv_cache.values = concatenate(kv_cache.values, new_value)
output = attention(query=new_key, keys=kv_cache.keys, values=kv_cache.values)

Ainsi, la cache KV est un mécanisme pour sauvegarder les clés et valeurs intermédiaires afin d'accélérer la génération séquentielle dans les transformateurs.