KV 缓存策略

键值 (KV) 向量用于计算注意力分数。对于自回归模型，KV 分数是*每次*计算的，因为模型一次预测一个 token。每次预测都依赖于之前的 token，这意味着模型每次都会执行相同的计算。

KV *缓存*存储这些计算，以便可以重复使用它们而无需重新计算。高效缓存对于优化模型性能至关重要，因为它可以减少计算时间并提高响应速度。有关缓存工作原理的更详细解释，请参阅缓存文档。

Transformers 提供了几个实现不同缓存机制的 Cache 类。其中一些 Cache 类经过优化以节省内存，而另一些则旨在最大限度地提高生成速度。请参阅下表以比较缓存类型，并使用它来帮助您为您的用例选择最佳缓存。

缓存类型	内存效率	支持 torch.compile()	建议初始化	延迟	长上下文生成
动态缓存	否	否	否	中	否
静态缓存	否	是	是	高	否
卸载缓存	是	否	否	低	是
卸载静态缓存	否	是	是	高	是
量化缓存	是	否	否	低	是
滑动窗口缓存	否	是	是	高	否

本指南向您介绍不同的 Cache 类，并展示如何将它们用于生成。

默认缓存

DynamicCache 是大多数模型的默认缓存类。它允许缓存大小动态增长，以便在生成过程中存储越来越多的键和值。

通过在 generate() 中配置 use_cache=False 来禁用缓存。

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", torch_dtype=torch.float16).to("cuda:0")
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)

缓存类也可以在调用并将其传递给模型的 past_key_values 参数之前进行初始化。此缓存初始化策略仅推荐用于某些缓存类型。

在大多数其他情况下，在 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", torch_dtype=torch.float16).to("cuda:0")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

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

内存高效缓存

KV 缓存会占用大量内存，并可能成为长上下文生成的瓶颈。内存高效缓存侧重于牺牲速度以减少内存使用。这对于大型语言模型 (LLM) 以及硬件内存受限的情况尤其重要。

卸载缓存

OffloadedCache 通过将大部分模型层的 KV 缓存移动到 CPU 来节省 GPU 内存。在模型对层进行 forward 迭代期间，只有当前层缓存保留在 GPU 上。OffloadedCache 异步预取下一层缓存并将上一层缓存发送回 CPU。

此缓存策略始终生成与 DynamicCache 相同的结果，并可作为即插即用替代品或备用方案。如果您有 GPU 并且出现内存不足 (OOM) 错误，您可能希望使用 OffloadedCache。

根据您的模型和生成选择（上下文大小、生成的 token 数量、束的数量等），您可能会注意到与 DynamicCache 相比，生成吞吐量略有下降。

通过在 GenerationConfig 或 generate() 中配置 cache_implementation="offloaded" 来启用 OffloadedCache。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16).to("cuda:0")
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.

以下示例展示了如何在内存不足时回退到 OffloadedCache。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

def resilient_generate(model, *args, **kwargs):
    oom = False
    try:
        return model.generate(*args, **kwargs)
    except torch.cuda.OutOfMemoryError as e:
        print(e)
        print("retrying with cache_implementation='offloaded'")
        oom = True
    if oom:
        torch.cuda.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, torch_dtype=torch.float16).to("cuda:0")
prompt = ["okay "*1000 + "Fun fact: The most"]
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
beams = { "num_beams": 40, "num_beam_groups": 40, "num_return_sequences": 40, "diversity_penalty": 1.0, "max_new_tokens": 23, "early_stopping": True, }
out = resilient_generate(model, **inputs, **beams)
responses = tokenizer.batch_decode(out[:,-28:], skip_special_tokens=True)

量化缓存

QuantizedCache 通过将 KV 值量化为较低精度来减少内存需求。QuantizedCache 目前支持两种量化后端。

HQQQuantizedCache 支持 int2、int4 和 int8 数据类型。
QuantoQuantizedCache 支持 int2 和 int4 数据类型。这是默认的量化后端。

如果上下文长度较短且有足够的 GPU 内存可用于生成而无需启用缓存量化，则量化缓存可能会损害延迟。尝试在内存效率和延迟之间找到平衡点。

通过在 GenerationConfig 中配置 cache_implementation="quantized" 并指示 QuantizedCacheConfig 中的量化后端来启用 QuantizedCache。任何其他量化相关参数也应作为字典或 QuantizedCacheConfig 实例传递。您应该使用这些额外参数的默认值，除非您内存不足。在这种情况下，请考虑减小残差长度。

HQQQuantizedCache

Quanto

速度优化缓存

默认的 DynamicCache 阻止您利用即时 (JIT) 优化，因为缓存大小不是固定的。JIT 优化使您能够以牺牲内存使用为代价来最大限度地提高延迟。以下所有缓存类型都与 JIT 优化兼容，例如 torch.compile 以加速生成。

静态缓存

StaticCache 为 KV 对预分配特定的最大缓存大小。您可以在不修改缓存的情况下生成最多达到最大缓存大小的内容。

通过在 generate() 中配置 cache_implementation="static" 来启用 StaticCache。

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

卸载静态缓存

OffloadedStaticCache 与 OffloadedCache 非常相似，只不过缓存大小设置为最大缓存大小。否则，OffloadedStaticCache 只在 GPU 上保留当前层缓存，其余的则移动到 CPU。

通过在 generate() 中配置 cache_implementation="offloaded_static" 来启用 OffloadedStaticCache。

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", torch_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="offloaded_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"

缓存卸载需要 CUDA GPU。

滑动窗口缓存

SlidingWindowCache 在之前的 kv 对上实现了一个滑动窗口，并且只保留最后的 sliding_window tokens。这种缓存类型旨在仅与支持*滑动窗口注意力*的模型一起使用，例如 Mistral。较旧的 kv 状态被丢弃并替换为新的 kv 状态。

通过在 generate() 中配置 cache_implementation="sliding_window" 来启用 SlidingWindowCache。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", torch_dtype=torch.float16).to("cuda:0")
inputs = tokenizer("Yesterday I was on a rock concert and.", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=30, cache_implementation="sliding_window")
tokenizer.batch_decode(out, skip_special_tokens=True)[0]

模型缓存

某些模型类型，例如编码器-解码器模型或 Gemma2 和 Mamba，具有专用的缓存类。

编码器-解码器缓存

EncoderDecoderCache 专为编码器-解码器模型设计。它管理自注意力缓存和交叉注意力缓存，以确保存储和检索以前的 kv 对。可以分别为编码器和解码器设置不同的缓存类型。

此缓存类型不需要任何设置。它可以在调用 generate() 或模型的 forward 方法时使用。

EncoderDecoderCache 目前仅支持 Whisper。

模型专用缓存

有些模型存储过去 kv 对或状态的方式是独有的，与任何其他缓存类不兼容。

Gemma2 需要 HybridCache，它在底层结合使用了 SlidingWindowCache 用于滑动窗口注意力，以及 StaticCache 用于全局注意力。

Mamba 需要 MambaCache，因为该模型没有注意力机制或 kv 状态。

迭代生成

缓存还可以在模型（聊天机器人）之间进行来回交互的迭代生成设置中工作。像常规生成一样，带有缓存的迭代生成允许模型有效地处理正在进行的对话，而无需在每个步骤重新计算整个上下文。

对于带缓存的迭代生成，首先初始化一个空的缓存类，然后可以输入新的提示。使用聊天模板跟踪对话历史记录。

以下示例演示了 Llama-2-7b-chat-hf。如果您正在使用不同的聊天风格模型，apply_chat_template() 可能会以不同的方式处理消息。根据 Jinja 模板的编写方式，它可能会剪切掉重要的 token。

例如，一些模型在推理过程中使用特殊的 <think> ... </think> tokens。这些 tokens 在重新编码时可能会丢失，导致索引问题。您可能需要手动删除或调整完成中的额外 tokens 以保持稳定。

import torch
from transformers import AutoTokenizer,AutoModelForCausalLM
from transformers.cache_utils import (
    DynamicCache,
    StaticCache,
    SlidingWindowCache,
    QuantoQuantizedCache,
    QuantizedCacheConfig,
)

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_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()

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})

预填充缓存

在某些情况下，您可能希望用特定前缀提示的 kv 对填充 Cache，并重复使用它来生成不同的序列。

以下示例初始化一个 StaticCache，然后缓存一个初始提示。现在，您可以从预填充的提示生成多个序列。

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, torch_dtype=torch.bfloat16, device_map="cuda")
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_batch_size=1, max_cache_len=1024, device="cuda", dtype=torch.bfloat16)

INITIAL_PROMPT = "You are a helpful assistant. "
inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to("cuda")
# 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("cuda")
    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)

< > 在 GitHub 上更新