汤不热吧Contents
引言：为什么KV Cache是LLM推理的生命线？
对于自回归（Autoregressive）的Transformer模型，特别是大型语言模型（LLMs），推理延迟主要发生在解码阶段。每生成一个新的Token，模型必须回顾所有历史Token的信息。如果不对这些历史信息进行缓存，每一步的计算量都会不必要地膨胀，从而严重拖慢推理速度。
KV Cache (Key-Value Cache) 正是解决这一问题的核心技术。它通过存储和重用自注意力机制中计算出的键（Key）和值（Value）向量，将推理延迟从与序列长度相关的重复计算中解放出来。
1. KV Cache的工作原理
在标准的自注意力机制中，给定输入序列 $X$，模型会生成查询（Query, $Q$）、键（Key, $K$）和值（Value, $V$）向量。注意力计算如下：
$$ \text{Attention}(Q, K, V) = \text{Softmax}(\frac{Q K^T}{\sqrt{d_k}}) V $$
在自回归解码过程中，假设我们已经生成了 $N-1$ 个Token，现在要生成第 $N$ 个Token：
- 没有缓存时： 为了计算第 $N$ 个Token的注意力，模型需要重新计算并拼接 Token 1 到 Token $N$ 的所有 $K$ 和 $V$ 向量。这意味着，前 $N-1$ 步已经计算过的 $K$ 和 $V$ 都会被重复计算。
-
使用KV Cache时： 模型将前 $N-1$ 步计算好的 $K_{1:N-1}$ 和 $V_{1:N-1}$ 存储在内存中（即KV Cache）。生成第 $N$ 个Token时，模型只需计算新的 $K_N$ 和 $V_N$，然后将它们追加到缓存中。最终的注意力计算是：
$$ Q_N \to K_{\text{cache}} = [K_{1:N-1}, K_N] \text{ 和 } V_{\text{cache}} = [V_{1:N-1}, V_N] $$
这样，计算成本大大降低，推理延迟主要由内存带宽决定，而不是重复的矩阵乘法。
2. 实践：在Hugging Face中启用KV Cache
对于使用Hugging Face Transformers库部署的模型，KV Cache的管理非常直观，通常通过一个布尔参数控制。
在调用模型的 generate() 方法时，关键是将 use_cache 参数设置为 True。
Python代码示例：启用缓存
以下代码演示了如何利用Hugging Face接口的内置KV Cache功能进行高效推理：
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
from transformers import AutoModelForCausalLM, AutoTokenizer
# 1. 加载模型和分词器
model_name = "facebook/opt-125m"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).cuda()
# 2. 准备输入
prompt = "AI infrastructure deployment is becoming increasingly important because"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
# 3. 启用 KV Cache 进行生成
# 此时，模型在内部会传递和更新 past_key_values
print("--- 开始生成 (使用 KV Cache) ---")
outputs = model.generate(
**inputs,
max_length=100,
do_sample=False,
use_cache=True # 启用 KV Cache，这是关键！
)
# 4. 解码结果
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
# 验证：如果查看 model.forward 的函数签名，你会发现它接受一个名为 past_key_values 的参数，这就是缓存的载体。
当 use_cache=True 时，模型在每次生成新Token后，会将当前层的 K/V 结果打包成一个元组，作为 past_key_values 输出。在下一次调用 forward 时，这个元组会被作为输入传入，直接用于注意力计算。
3. 生产环境的优化：vLLM的PagedAttention
虽然基本的KV Cache解决了单次推理的计算冗余问题，但它带来了新的挑战：内存管理。
在并发高、批处理（Batching）大的生产环境中，KV Cache会占据大量的GPU内存。由于不同请求的序列长度差异大，如果使用传统的连续内存分配（Contiguous Allocation），内存碎片化和浪费会非常严重。
高性能推理引擎如 vLLM 引入了 PagedAttention 机制来解决这个问题，该机制灵感来源于操作系统中的分页（Paging）概念。
PagedAttention的核心优势：
- 非连续存储： KV Cache不再需要连续存储，而是存储在固定大小的“块”（Blocks）中。
- 动态分配： 块可以根据需要动态分配给不同的请求，就像内存页一样。
- 消除碎片化： 提高了GPU内存利用率，显著提高了吞吐量。
部署配置：vLLM示例
使用 vLLM 部署模型时，你主要需要配置用于存储 KV Cache 的 GPU 内存比例，以及块大小（block-size）。
2
3
4
5
6
7
python -m vllm.entrypoints.api_server \
--model meta-llama/Llama-2-7b-hf \
--tensor-parallel-size 4 \
--gpu-memory-utilization 0.90 \
--dtype bfloat16 \
--block-size 16 # 定义KV Cache的块大小（通常是16或32）
- –gpu-memory-utilization 0.90：这是设置 KV Cache 内存池的关键参数。它告诉 vLLM 将多少比例的 GPU 内存用于存储 KV Cache。预留更多空间可以支持更长的序列或更大的并发批次。
- –block-size 16：定义了 PagedAttention 的存储单元大小。选择合适的块大小是吞吐量和内存利用率之间的权衡。
通过配置这些参数，vLLM能够高效地管理成百上千个并发请求的 KV Cache，实现数倍于传统部署框架的吞吐量。