缓存的工作原理

想象一下，你正在和某人交谈，但对方每次回应你时，都必须从头开始，完全不记得之前说过的话。这显然会很慢且低效，对吧？

你可以将这一类比扩展到 Transformer 模型。自回归模型的生成过程可能较慢，因为它每次只预测一个 token。每个新的预测都依赖于之前所有的上下文。

要预测第 1000 个 token，模型需要前 999 个 token 的信息。这些信息通过 token 表示之间的矩阵乘法来表示。

要预测第 1001 个 token，除了第 1000 个 token 的信息外，还需要前 999 个 token 的相同信息。这意味着模型必须为每个 token 反复执行大量的矩阵乘法运算！

键值（KV）缓存通过存储从先前处理的 token 的注意力层中派生的 kv 对来消除这种低效。存储的 kv 对从缓存中检索并复用于后续 token，从而避免了重新计算的需要。

缓存仅应用于推理。如果在训练期间启用缓存，可能会导致意外错误。

为了更好地理解缓存是如何以及为何起作用的，让我们更仔细地看一下注意力矩阵的结构。

注意力矩阵

对于批量大小为 缩放点积注意力 的计算方式如下，其中 b, 为注意力头数 h, 为当前序列长度 T, 为每个注意力头的维度 d_head. $$\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V$$

查询（Q）、键（K）和值（V）矩阵是输入嵌入的投影，其形状为 (b, h, T, d_head)。

对于因果注意力，掩码会阻止模型关注未来的 token。一旦某个 token 被处理，其表示相对于未来的 token 将不再发生变化，这意味着 $K_{\text{past}}$ 和 $V_{\text{past}}$ 可以被缓存并复用，以计算最后一个 token 的表示。 $$\text{Attention}(q_t, [\underbrace{k_1, k_2, \dots, k_{t-1}}_{\text{cached}}, k_{t}], [\underbrace{v_1, v_2, \dots, v_{t-1}}_{\text{cached}}, v_{t}])$$

在推理阶段，你只需要最后一个 token 的 query 来计算表示 $x_t$ ，该表示用于预测下一个 token $ t+1 $。在每一步中，新的 key 和 value 向量会 存储 到缓存中，并 追加 到过去的 key 和 value 中。 $$K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t)$$

注意力在模型的每一层中独立计算，缓存也是按层进行的。

请参阅下表，以比较缓存如何提升效率。

缓存类

基本的 KV cache 接口接收当前 token 的 key 和 value 张量，并返回更新后的 K 和 V 张量。该过程由模型的 forward 方法在内部进行管理。

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

使用 Transformers 的 Cache 类时，自注意力模块会执行几个关键步骤，以整合过去和当前的信息。

1. 注意力模块将当前的 kv 对与缓存中存储的历史 kv 对进行拼接。这会生成形状为 (new_tokens_length, past_kv_length + new_tokens_length) 的注意力权重。当前和历史 kv 对本质上被合并起来以计算注意力分数，从而确保模型能够感知到之前的上下文以及当前的输入。

2. 当迭代调用 forward 方法时，注意力掩码的形状必须与过去和当前 kv 对的合并长度相匹配。注意力掩码的形状应为 (batch_size, past_kv_length + new_tokens_length)。这通常由 generate() 内部处理，但如果您想使用 Cache 自行实现生成循环，请务必注意这一点！注意力掩码应包含过去和当前 token 的值。

缓存存储实现

缓存被组织为一个层列表，其中每一层都包含一个键缓存和一个值缓存。键缓存和值缓存是形状为 [batch_size, num_heads, seq_len, head_dim] 的张量。

层可以有不同的类型（例如 DynamicLayer、StaticLayer、StaticSlidingWindowLayer），这些类型主要影响序列长度的处理方式以及缓存的更新方式。

最简单的方式是一个随着处理更多 token 而增长的 DynamicLayer。序列长度维度（seq_len）会随着每个新 token 的增加而增大：

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)

StaticLayer 和 StaticSlidingWindowLayer 等其他层类型具有在创建缓存时设定的固定序列长度，因此与 torch.compile 兼容。对于 StaticSlidingWindowLayer，当添加新 token 时，已有的 token 会从缓存中移出。

以下示例演示了如何使用 DynamicCache 创建生成循环。如前所述，注意力掩码是过去和当前 token 值的拼接。

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,"