大模型的性能提升：KV-Cache-CSDN博客

本文链接：https://blog.csdn.net/wireless_com/article/details/148753258

大语言模型（LLM）在生成文本时，通常是一个 token 一个 token 地进行。每当模型生成一个新的 token，它就会把这个 token 加入输入序列，作为下一步预测下一个 token 的依据。这一过程不断重复，直到完成整个输出。

然而，这种逐词生成的方式带来了一个问题：每一步的输入几乎与前一步相同，只是多了一个新 token。如果不对计算过程进行优化，模型就不得不反复重复大量相同的计算，造成资源浪费和效率低下。

为了解决这个问题，KV-Cache（Key Value-Cache）应运而生，它是提升 LLM 推理性能的关键技术之一。

KV-Cache 的核心思想是缓存中间计算结果。在 Transformer 架构中，每个 token 在生成下一个 token 时都需要通过自注意力机制计算一系列 Key 和 Value 向量。这些向量描述了当前 token 与前面所有 token 的关系。

1. 关于自注意力机制

要理解KV缓存，首先需要掌握注意力机制的基本原理。在著名的论文《Attention Is All You Need》中，提出了用于Transformer模型的注意力机制公式，该公式帮助模型在生成每个新token时确定应关注哪些先前的token。

注意力机制的核心在于通过一系列计算来量化不同token之间的关联强度。具体来说，当模型生成一个新的token时，它会根据输入序列中的所有token计算出一组Query（查询）、Key（键）和Value（值）向量。这些向量通过特定的公式相互作用，以决定当前上下文中每个token的重要性。关于Transformer的更多秒杀请参考《从零构建大模型之Transformer公式解读》。

让我们一步一步来看看这个方程是怎么在工程上实现的。

2. 提示词阶段（预填充阶段）

我们从一个示例输入开始，称为“提示词”（prompt）。这个输入首先会被分词器（tokenizer）切分为一个个的 token，也即模型可处理的基本单位。例如，短语 "The quick brown fox" 在使用 OpenAI 的 o200kbase 分词器时会被拆分为 4 个 token。随后，每个 token 被转换为一个嵌入向量，记作 x₁ 到 x₄，这些向量的维度由模型决定，通常表示为 d_model。在原始 Transformer 论文中，d_model 设定为 512。

为了高效计算，我们可以将所有 n 个嵌入向量堆叠成一个矩阵 X，其形状为 [n × d_model]。接下来，为了执行注意力机制，我们需要通过三个可学习的投影矩阵 Wq、Wk 和 Wv 来分别生成查询（Query）、键（Key）和值（Value）向量。

具体来说：

Wq 的形状为 [d_model × d_k]，用于将嵌入向量映射到查询空间，得到 q 矩阵，其形状为 [n × d_k]；
Wk 的形状也为 [d_model × d_k]，生成 k 矩阵，形状为 [n × d_k]；
Wv 的形状为 [d_model × d_v]，生成 v 矩阵，形状为 [n × d_v]。

这三个矩阵在训练过程中不断优化，以捕捉不同 token 之间的依赖关系。其中，dk 和 dv 是设计模型结构时设定的超参数，在最初的 Transformer 模型中被设为 64。

这种线性变换不仅将高维嵌入压缩到更易操作的空间，还保留了关键语义信息，为后续的注意力计算奠定了基础。

接下来，我们通过计算查询矩阵 q 与其对应键矩阵 k 的转置的乘积，得到一个大小为 [n × n] 的自注意力分数矩阵。这个矩阵中的每个元素代表了某个查询向量与所有键向量之间的相似性得分，反映了在生成当前 token 时，模型应将多少注意力分配给前面每一个 token。

对于解码器类型的大型语言模型（LLM），为了避免模型在生成过程中“偷看”未来的信息，我们采用一种称为“掩码自注意力”的机制——即将该矩阵的上三角部分设为负无穷（-inf）。这样一来，在后续的 softmax 操作中，这些位置的值会趋近于零，从而确保每个 token 在预测时只能关注到它之前的历史 token，而不会看到未来的输入。

处理后的自注意力分数矩阵因此成为一个下三角结构，体现了严格的因果关系。由于这一操作，模型在生成序列的过程中能够保持逻辑连贯性和时间顺序的正确性。

为了得到最终的注意力输出，我们首先对这个带有掩码的注意力分数矩阵进行缩放，除以 sqrt（d_k)以防止点积结果过大导致梯度饱和；然后应用 softmax 函数将其转化为概率分布；最后将该分布与值矩阵 v 相乘，得到加权聚合后的上下文信息。这一过程即完成了对输入序列 "The quick brown fox" 的注意力输出计算，为下一步的 token 预测提供了基础。

3. 自回归生成阶段（解码阶段）

现在，当我们生成下一个 token（例如 "jumps"）并计算其对应的注意力输出时，在自回归生成阶段（即解码过程中），模型是如何工作的呢？

在生成“jumps”这一新 token 时，模型会为它生成新的查询向量 q₅、键向量 k₅ 和值向量 v₅。但值得注意的是，并不需要重新计算之前所有 token 的 k 和 v，因为这些中间结果已经保存在 KV-Cache 中。

如图所示，我们只需关注注意力矩阵中新增的最后一行，其余部分可以从缓存中直接复用。图中以灰色突出显示的 k₁ 到 k₄ 和 v₁ 到 v₄ 都是之前步骤中已计算并缓存的结果。这意味着，只有与当前 token 相关的 q₅、k₅ 和 v₅ 是新生成的。

最终，第 n 个 token 的注意力输出只需以下计算（忽略 softmax 和缩放因子以便理解）：

这正是 KV-Cache 的价值所在：通过缓存先前 token 的 Key 和 Value 向量，模型无需重复计算整个历史序列，从而显著减少冗余运算，提升推理效率。

随着生成序列变长，KV-Cache 的作用愈发重要——它不仅减少了计算负担，也降低了内存带宽的压力，使得 LLM 在实际应用中能够实现高效、流畅的文本生成。

4. 速度与内存的权衡

KV-Cache 的引入显著提升了大型语言模型（LLM）的推理速度。其核心思想在于缓存每一步计算生成的 Key 和 Value 向量，使得在生成新 token 时，模型无需重复计算历史上下文中的 K 和 V 值，从而大幅减少冗余计算，加快响应生成。

然而，这种性能提升也带来了内存上的代价。由于 KV-Cache 需要为每个已生成的 token 保存对应的 K 和 V 向量，它会持续占用 GPU 显存。对于本身就需要大量资源的 LLM 来说，这进一步加剧了显存压力，尤其是在处理长序列时更为明显。

因此，在实际应用中存在一个权衡：一方面，使用 KV-Cache 可以加速生成过程，但另一方面，它也会增加内存消耗。当显存紧张时，开发者可能需要选择牺牲部分生成速度来节省内存，或接受更高的硬件资源开销以换取更快的推理表现。这种性能与资源之间的平衡，是部署 LLM 时必须仔细考量的关键因素之一。

5. KV-Cache 实践

KV-Cache 是现代大型语言模型（LLM）推理引擎中至关重要的一项优化技术。在 Hugging Face 的 Transformers 库中，当你调用 model.generate() 函数生成文本时，默认会启用 use_cache=True 参数，也就是自动使用 KV-Cache 来提升生成效率。

类似的技术细节在 Hugging Face 官方博客的文章《KV Caching Explained: Optimizing Transformer Inference Efficiency》中也有深入解析。该文通过实验展示了启用 KV-Cache 带来的显著加速效果：在 T4 GPU 上对模型 HuggingFaceTB/SmoLLM2-1.7B 进行测试，使用 KV-Cache 相比不使用时，推理速度提升了 5.21 倍。

这一性能提升充分说明了 KV-Cache 在实际应用中的重要性，测试的参考代码如下：

from transformers import AutoTokenizer, AutoModelForCausaLLM
import torch
import time

# Choose model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausaLLM.from_pretrained("gpt2")

# Move model to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

# Prepare input
input_sentence = "The red cat was"
inputs = tokenizer(input_sentence, return_tensors="pt").to(device)

# Function to measure generation time
def generate_and_time(use_cache: bool):
    torch.cuda.empty_cache()
    start_time = time.time()

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=300,
            use_cache=use_cache
        )

    end_time = time.time()
    duration = end_time - start_time
    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return duration, output_text

# Measure with KV-cache enabled
time_with_cache, text_with_cache = generate_and_time(use_cache=True)
print(f"\n[use_cache=True] Time taken: {time_with_cache:.4f} seconds")
print(f"Generated Text:\n{text_with_cache}\n")

# Measure with KV-cache disabled
time_without_cache, text_without_cache = generate_and_time(use_cache=False)
print(f"[use_cache=False] Time taken: {time_without_cache:.4f} seconds")
print(f"Generated Text:\n{text_without_cache}\n")

# Speedup factor
speedup = time_without_cache / time_with_cache
print(f"Speedup from using KV-cache: {speedup:.2f}×")

KV-Cache的运行速度实际上受到多种因素的综合影响，其中包括模型的规模（具体体现在注意力层数的多少）、输入文本的长度n、所使用的硬件设备以及具体的实现细节等。