想要掌握 LLM 缓存管理？从这些资源开始。

KV Cache： https://huggingface.co/docs/transformers/main/en/cache_explanation… https://hamzaelshafie.bearblog.dev/paged-attention-from-first-principles-a-view-inside-vllm…

Prefix Caching https://docs.vllm.ai/en/latest/features/automatic_prefix_caching.html…

Continuous Batching https://anyscale.com/blog/continuous-batching-llm-inference…

Speculative Decoding https://research.google/blog/looking-back-at-speculative-decoding/…

KV Cache Quantization https://huggingface.co/blog/kv-cache-quantization…

研究论文：

1. FlashInfer (Attention Engine) https://arxiv.org/abs/2501.01005
2. Zipage (Compressed PagedAttention) https://arxiv.org/abs/2603.08743
3. IceCache (Memory-Efficient KV Cache) https://arxiv.org/abs/2604.10539

缓存 · Hugging Face

来源：https://huggingface.co/docs/transformers/main/en/cache_explanation 想象一下，你正在与某人交谈，而对方每次回应时都不记得你之前说过什么，必须从头开始。这将会变得缓慢且低效，对吧？

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

为了预测第 1000 个 token，模型需要来自前 999 个 token 的信息。这些信息以 token 表示之间的矩阵乘法形式呈现。

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

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

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

为了更好地理解缓存的工作原理以及为何有效，让我们更仔细地观察注意力矩阵的结构。

https://huggingface.co/docs/transformers/main/en/cache_explanation#attention-matrices 注意力矩阵

缩放点积注意力的计算公式如下，其中批次大小为 b，注意力头数为 h，当前序列长度为 T，每个注意力头的维度为 d_head。 Attention(Q,K,V)=softmax(QK⊤dhead×mask)V\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_past 和 V_past 可以被缓存并复用于计算最后一个 token 的表示。 Attention(qt,[k1,k2,…,kt−1⏟cached,kt],[v1,v2,…,vt−1⏟cached,vt])\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 的查询来计算预测下一个 token $ t+1 $ 的表示 $ x_t $。在每一步，新的键和值向量会被存储到缓存中，并追加到过去的键和值之后。 Kcache←concat(Kpast,kt),Vcache←concat(Vpast,vt)K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t)

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

参考下表，比较缓存如何提高效率。

https://huggingface.co/docs/transformers/main/en/cache_explanation#cache-class 缓存类

一个基本的 KV 缓存接口接收当前 token 的键和值张量，并返回更新后的 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 (https://huggingface.co/docs/transformers/main/en/internal/generation_utils#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() (https://huggingface.co/docs/transformers/main/en/main_classes/text_generation#transformers.GenerationMixin.generate) 内部处理，但如果你想结合 Cache (https://huggingface.co/docs/transformers/main/en/internal/generation_utils#transformers.Cache) 实现自己的生成循环，请注意这一点！注意力掩码应包含过去和当前 token 的值。

https://huggingface.co/docs/transformers/main/en/cache_explanation#cache-storage-implementation 缓存存储实现

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

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

最简单的是 DynamicLayer，它会随着更多 token 的处理而增长。序列长度维度（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 (https://huggingface.co/docs/transformers/main/en/internal/generation_utils#transformers.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)
    
    next_token_ids = outputs.logits[:, -1:].argmax(-1)
    generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1)
    
    # 更新注意力掩码，连接一个全1的列表示新 token
    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,"

更新于 GitHub (https://github.com/huggingface/transformers/blob/main/docs/source/en/cache_explanation.md)