第12章：推理 I：KV缓存（KV-Cache）#
12.1 推理过程概述#
在构建故事讲述AI大语言模型的过程中，除了训练阶段，推理阶段同样至关重要。推理是指使用训练好的模型生成文本的过程，这是用户最终与模型交互的环节。对于大型语言模型（LLM）来说，推理过程的效率和质量直接影响用户体验。
大语言模型的推理过程通常是自回归的，即模型基于已生成的文本序列逐个生成下一个词或标记。这个过程可以表示为：
$$P(x_1, x_2, …, x_n) = \prod_{i=1}^{n} P(x_i | x_1, x_2, …, x_{i-1})$$
其中，$P(x_1, x_2, …, x_n)$是整个序列的概率，$P(x_i | x_1, x_2, …, x_{i-1})$是在已知前面所有标记的条件下，生成第$i$个标记的概率。
在实际应用中，推理过程面临几个主要挑战：
计算效率：大语言模型通常有数十亿甚至数千亿参数，每次推理都需要大量计算资源。
内存消耗：模型参数和中间状态需要占用大量内存，特别是在生成长文本时。
推理速度：自回归生成过程是顺序的，难以并行化，这限制了生成速度。
生成质量：需要平衡多样性和连贯性，避免生成重复或无意义的内容。
为了解决这些挑战，研究人员提出了多种优化技术，其中KV缓存（Key-Value Cache）是最重要的优化技术之一，它显著提高了推理效率。在本章中，我们将深入探讨KV缓存的原理、实现和优化方法。
12.2 Transformer架构回顾#
在深入KV缓存之前，我们需要回顾Transformer架构，特别是自注意力机制，因为KV缓存主要优化的就是这一部分。
12.2.1 Transformer基本结构#
Transformer模型由编码器（Encoder）和解码器（Decoder）组成，对于像GPT这样的自回归语言模型，通常只使用解码器部分。Transformer解码器由多个相同的层堆叠而成，每一层包含以下子层：
自注意力层（Self-Attention）：允许模型关注输入序列的不同位置，捕捉序列内的依赖关系。
前馈神经网络（Feed-Forward Network）：由两个线性变换和一个非线性激活函数组成，对每个位置的表示进行独立处理。
此外，每个子层都使用残差连接（Residual Connection）和层归一化（Layer Normalization）。
12.2.2 自注意力机制详解#
自注意力机制是Transformer的核心，它允许模型在处理序列中的每个位置时，考虑整个序列的信息。自注意力的计算过程如下：
对输入序列$X$进行线性变换，得到查询（Query）、键（Key）和值（Value）： $$Q = XW_Q, K = XW_K, V = XW_V$$ 其中，$W_Q$、$W_K$和$W_V$是可学习的权重矩阵。
计算注意力权重： $$Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$$ 其中，$d_k$是键向量的维度，用于缩放点积，防止梯度消失。
在多头注意力（Multi-Head Attention）中，上述过程被并行执行多次，然后将结果拼接： $$MultiHead(Q, K, V) = Concat(head_1, head_2, …, head_h)W_O$$ 其中，$head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)$，$W_O$是输出投影矩阵。
在自回归生成过程中，为了保持因果关系（即当前位置只能看到之前的位置），通常使用掩码自注意力（Masked Self-Attention），通过掩码矩阵将未来位置的注意力权重设为负无穷，使其在softmax后接近于零。
12.2.3 自回归生成中的计算冗余#
在自回归生成过程中，模型逐个生成标记。对于每个新生成的标记，模型需要重新计算整个序列的表示。具体来说，如果已经生成了$t$个标记，要生成第$t+1$个标记，需要：
将前$t$个标记输入模型。
计算每一层的自注意力和前馈网络。
基于最后一个位置的输出预测下一个标记。
这个过程中存在明显的计算冗余：每次生成新标记时，都要重新计算前面所有标记的表示，而这些表示在之前的步骤中已经计算过了。这种冗余在生成长文本时尤为明显，严重影响推理效率。
KV缓存正是为了解决这个问题而设计的，它通过缓存之前计算的键（Key）和值（Value）向量，避免重复计算，从而显著提高推理效率。
12.3 KV缓存原理#
KV缓存（Key-Value Cache）是一种优化自回归生成过程的技术，通过缓存已计算的键（Key）和值（Value）向量，避免在生成新标记时重复计算，从而提高推理效率。
12.3.1 基本原理#
在自回归生成过程中，对于每个新生成的标记，模型需要计算其与之前所有标记的注意力权重。这涉及到计算查询（Query）、键（Key）和值（Value）向量，以及它们之间的交互。
KV缓存的核心思想是：对于已经生成的标记，其键和值向量在不同的生成步骤中是不变的，因此可以计算一次并缓存起来，在后续步骤中直接使用。
具体来说，如果已经生成了$t$个标记，要生成第$t+1$个标记，使用KV缓存的过程如下：
只对第$t$个标记计算查询、键和值向量（$Q_t$、$K_t$和$V_t$）。
从缓存中获取前$t-1$个标记的键和值向量（$K_{1:t-1}$和$V_{1:t-1}$）。
将当前的键和值向量与缓存中的拼接：$K_{1:t} = [K_{1:t-1}; K_t]$，$V_{1:t} = [V_{1:t-1}; V_t]$。
计算注意力：$Attention(Q_t, K_{1:t}, V_{1:t})$。
更新缓存，将$K_t$和$V_t$添加到缓存中。
通过这种方式，每次只需要计算一个新标记的查询、键和值向量，而不是重新计算整个序列，从而显著减少计算量。
12.3.2 数学表示#
为了更清晰地理解KV缓存，我们可以用数学公式表示这个过程。
在没有KV缓存的情况下，生成第$t+1$个标记的自注意力计算如下：
$$Q_{1:t} = X_{1:t}W_Q, K_{1:t} = X_{1:t}W_K, V_{1:t} = X_{1:t}W_V$$ $$A_t = softmax(\frac{Q_t K_{1:t}^T}{\sqrt{d_k}})V_{1:t}$$
其中，$X_{1:t}$是前$t$个标记的输入表示，$Q_t$是第$t$个位置的查询向量，$A_t$是第$t$个位置的注意力输出。
使用KV缓存后，计算过程变为：
$$Q_t = X_t W_Q$$ $$K_t = X_t W_K, V_t = X_t W_V$$ $$K_{1:t} = [K_{1:t-1}; K_t], V_{1:t} = [V_{1:t-1}; V_t]$$ $$A_t = softmax(\frac{Q_t K_{1:t}^T}{\sqrt{d_k}})V_{1:t}$$
其中，$K_{1:t-1}$和$V_{1:t-1}$是从缓存中获取的前$t-1$个位置的键和值向量。
通过比较这两个过程，我们可以看到KV缓存的主要优势：
减少了矩阵乘法的计算量：从$O(t \times d^2)$减少到$O(d^2)$，其中$d$是模型维度。
避免了重复计算：每个位置的键和值向量只计算一次。
这些优势在生成长文本时尤为明显，可以显著提高推理速度。
12.3.3 多层多头注意力中的KV缓存#
在实际的Transformer模型中，通常有多个注意力层和多个注意力头。KV缓存需要为每个层和每个头维护单独的缓存。
假设模型有$L$层，每层有$h$个注意力头，则KV缓存的结构如下：
$$Cache = {(K_{l,h}, V_{l,h}) | l \in [1, L], h \in [1, H]}$$
其中，$K_{l,h}$和$V_{l,h}$分别是第$l$层第$h$个头的键和值缓存。
在实现中，这通常表示为形状为$[L, 2, H, T, D/H]$的张量，其中：
$L$是层数
$2$表示键和值两种缓存
$H$是头数
$T$是当前序列长度
$D/H$是每个头的维度
随着生成过程的进行，$T$会不断增加，缓存需要相应地扩展。
12.4 KV缓存实现#
了解了KV缓存的原理后，我们来看如何在实际代码中实现它。我们将使用PyTorch作为示例，展示如何在Transformer模型中实现KV缓存。
12.4.1 基本实现#
首先，我们定义一个简化的Transformer层，包含自注意力和前馈网络：
import torch
import torch.nn as nn
import torch.nn.functional as F
class TransformerLayer(nn.Module):
def __init__(self, d_model, nhead, d_ff):
super().__init__()
self.d_model = d_model
self.nhead = nhead
self.head_dim = d_model // nhead
# 自注意力参数
self.q_proj = nn.Linear(d_model, d_model)
self.k_proj = nn.Linear(d_model, d_model)
self.v_proj = nn.Linear(d_model, d_model)
self.o_proj = nn.Linear(d_model, d_model)
# 前馈网络
self.ff1 = nn.Linear(d_model, d_ff)
self.ff2 = nn.Linear(d_ff, d_model)
# 层归一化
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
def forward(self, x, kv_cache=None, return_cache=False):
# 自注意力
residual = x
x = self.norm1(x)
# 计算查询、键、值
q = self.q_proj(x)
k = self.k_proj(x)
v = self.v_proj(x)
# 重塑为多头形式
batch_size, seq_len, _ = q.shape
q = q.view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)
k = k.view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)
v = v.view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)
# 使用KV缓存
if kv_cache is not None:
past_k, past_v = kv_cache
k = torch.cat([past_k, k], dim=2)
v = torch.cat([past_v, v], dim=2)
# 计算注意力
scores = torch.matmul(q, k.transpose(2, 3)) / (self.head_dim ** 0.5)
# 掩码自注意力（因果关系）
seq_len_k = k.size(2)
mask = torch.triu(torch.ones(seq_len, seq_len_k, dtype=torch.bool, device=x.device), diagonal=1)
scores.masked_fill_(mask.unsqueeze(0).unsqueeze(0), float('-inf'))
attn_weights = F.softmax(scores, dim=-1)
attn_output = torch.matmul(attn_weights, v)
# 重塑回原始形状
attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
attn_output = self.o_proj(attn_output)
# 残差连接
x = residual + attn_output
# 前馈网络
residual = x
x = self.norm2(x)
x = self.ff2(F.gelu(self.ff1(x)))
x = residual + x
# 返回输出和更新的缓存
if return_cache:
return x, (k, v)
return x
接下来，我们定义一个使用KV缓存的Transformer模型：
class TransformerWithKVCache(nn.Module):
def __init__(self, vocab_size, d_model, nhead, d_ff, num_layers):
super().__init__()
self.token_embedding = nn.Embedding(vocab_size, d_model)
self.position_embedding = nn.Embedding(1024, d_model) # 假设最大长度为1024
self.layers = nn.ModuleList([
TransformerLayer(d_model, nhead, d_ff) for _ in range(num_layers)
])
self.norm = nn.LayerNorm(d_model)
self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
self.num_layers = num_layers
self.nhead = nhead
self.d_model = d_model
self.head_dim = d_model // nhead
def forward(self, input_ids, kv_cache=None):
batch_size, seq_len = input_ids.shape
# 获取词嵌入和位置嵌入
token_emb = self.token_embedding(input_ids)
pos_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)
pos_emb = self.position_embedding(pos_ids)
x = token_emb + pos_emb
# 初始化KV缓存
if kv_cache is None:
kv_cache = [(None, None) for _ in range(self.num_layers)]
# 更新后的KV缓存
new_kv_cache = []
# 通过每一层
for i, layer in enumerate(self.layers):
layer_cache = kv_cache[i] if kv_cache else None
x, new_cache = layer(x, layer_cache, return_cache=True)
new_kv_cache.append(new_cache)
x = self.norm(x)
logits = self.lm_head(x)
return logits, new_kv_cache
def generate(self, input_ids, max_length, temperature=1.0):
"""使用KV缓存生成文本"""
device = input_ids.device
batch_size = input_ids.shape[0]
# 初始输入和KV缓存
current_ids = input_ids
kv_cache = None
generated_ids = []
for _ in range(max_length):
# 前向传播
logits, kv_cache = self.forward(current_ids, kv_cache)
# 只关注最后一个位置的预测
next_token_logits = logits[:, -1, :] / temperature
# 采样下一个标记
probs = F.softmax(next_token_logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
# 添加到生成的序列中
generated_ids.append(next_token)
# 更新当前输入（只包含新生成的标记）
current_ids = next_token
# 拼接所有生成的标记
return torch.cat(generated_ids, dim=1)
这个实现展示了KV缓存的基本原理：
在每一层的前向传播中，计算并返回更新后的KV缓存。
在生成过程中，每次只输入新生成的标记，并使用缓存中的键和值向量。
随着生成过程的进行，KV缓存不断扩展，包含所有已生成标记的信息。
12.4.2 高效实现考虑因素#
在实际应用中，KV缓存的实现需要考虑更多因素，以实现最佳性能：
内存布局：为了提高内存访问效率，KV缓存的张量应该是连续的（contiguous）。在PyTorch中，可以使用
contiguous()方法确保张量是连续的。预分配内存：为了避免频繁的内存分配和释放，可以预先分配足够大的缓存空间，然后在生成过程中逐步填充。
批处理支持：在处理批量输入时，每个样本可能有不同的序列长度，需要适当处理掩码和缓存。
内存优化：对于非常长的序列，KV缓存可能占用大量内存。可以考虑使用较低精度（如float16或int8）存储缓存，或者实现缓存清理策略。
以下是一个考虑这些因素的更高效实现：
class EfficientTransformerWithKVCache(nn.Module):
def __init__(self, vocab_size, d_model, nhead, d_ff, num_layers, max_seq_len=1024):
super().__init__()
# 模型定义与之前相同
# ...
self.max_seq_len = max_seq_len
def _create_kv_cache(self, batch_size, device, dtype=torch.float32):
"""预分配KV缓存空间"""
head_dim = self.d_model // self.nhead
# 为每一层创建缓存
kv_cache = []
for _ in range(self.num_layers):
# 键缓存：[batch_size, nhead, max_seq_len, head_dim]
k_cache = torch.zeros(
(batch_size, self.nhead, self.max_seq_len, head_dim),
device=device, dtype=dtype
)
# 值缓存：[batch_size, nhead, max_seq_len, head_dim]
v_cache = torch.zeros(
(batch_size, self.nhead, self.max_seq_len, head_dim),
device=device, dtype=dtype
)
kv_cache.append((k_cache, v_cache))
return kv_cache
def forward(self, input_ids, kv_cache=None, cache_position=0):
batch_size, seq_len = input_ids.shape
device = input_ids.device
# 获取词嵌入和位置嵌入
token_emb = self.token_embedding(input_ids)
pos_ids = torch.arange(cache_position, cache_position + seq_len, device=device).unsqueeze(0).expand(batch_size, -1)
pos_emb = self.position_embedding(pos_ids)
x = token_emb + pos_emb
# 初始化或使用现有的KV缓存
if kv_cache is None:
kv_cache = self._create_kv_cache(batch_size, device, x.dtype)
# 通过每一层
for i, layer in enumerate(self.layers):
k_cache, v_cache = kv_cache[i]
# 自注意力
residual = x
x = self.norm1(x)
# 计算查询、键、值
q = self.q_proj(x)
k = self.k_proj(x)
v = self.v_proj(x)
# 重塑为多头形式
q = q.view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)
k = k.view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)
v = v.view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)
# 更新KV缓存
k_cache[:, :, cache_position:cache_position+seq_len] = k
v_cache[:, :, cache_position:cache_position+seq_len] = v
# 使用完整的KV缓存计算注意力
k = k_cache[:, :, :cache_position+seq_len]
v = v_cache[:, :, :cache_position+seq_len]
# 计算注意力
scores = torch.matmul(q, k.transpose(2, 3)) / (self.head_dim ** 0.5)
# 掩码自注意力（因果关系）
seq_len_k = k.size(2)
mask = torch.triu(torch.ones(seq_len, seq_len_k, dtype=torch.bool, device=device), diagonal=1+cache_position)
scores.masked_fill_(mask.unsqueeze(0).unsqueeze(0), float('-inf'))
attn_weights = F.softmax(scores, dim=-1)
attn_output = torch.matmul(attn_weights, v)
# 重塑回原始形状
attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
attn_output = self.o_proj(attn_output)
# 残差连接
x = residual + attn_output
# 前馈网络
residual = x
x = self.norm2(x)
x = self.ff2(F.gelu(self.ff1(x)))
x = residual + x
x = self.norm(x)
logits = self.lm_head(x)
return logits, kv_cache, cache_position + seq_len
def generate(self, input_ids, max_length, temperature=1.0):
"""使用KV缓存生成文本"""
device = input_ids.device
batch_size = input_ids.shape[0]
# 初始输入和KV缓存
current_ids = input_ids
kv_cache = None
cache_position = 0
generated_ids = [input_ids]
# 首先处理整个输入序列
logits, kv_cache, cache_position = self.forward(current_ids, kv_cache, cache_position)
# 然后逐个生成新标记
for _ in range(max_length):
# 只取最后一个标记的预测
next_token_logits = logits[:, -1, :] / temperature
# 采样下一个标记
probs = F.softmax(next_token_logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
# 添加到生成的序列中
generated_ids.append(next_token)
# 使用KV缓存进行下一步预测
logits, kv_cache, cache_position = self.forward(next_token, kv_cache, cache_position)
# 检查是否需要扩展缓存
if cache_position >= self.max_seq_len:
# 简单策略：丢弃前一半的缓存
for i in range(len(kv_cache)):
k_cache, v_cache = kv_cache[i]
half_len = self.max_seq_len // 2
k_cache[:, :, :half_len] = k_cache[:, :, half_len:self.max_seq_len]
v_cache[:, :, :half_len] = v_cache[:, :, half_len:self.max_seq_len]
# 清空后半部分
k_cache[:, :, half_len:].zero_()
v_cache[:, :, half_len:].zero_()
cache_position = half_len
# 拼接所有生成的标记
return torch.cat(generated_ids, dim=1)
这个实现增加了几个重要的优化：
预分配固定大小的KV缓存，避免频繁的内存分配。
跟踪缓存位置（cache_position），只更新和使用需要的部分。
当缓存接近满时，实现了一个简单的缓存清理策略（丢弃前一半）。
使用连续的内存布局，提高内存访问效率。
12.4.3 主流框架中的KV缓存实现#
在实际应用中，我们通常使用成熟的深度学习框架和库，如PyTorch、TensorFlow、Hugging Face Transformers等。这些框架已经实现了高效的KV缓存。
以Hugging Face Transformers库为例，它在GPT-2、GPT-J、LLaMA等模型中都实现了KV缓存。以下是使用Hugging Face Transformers库进行文本生成的示例：
from transformers import GPT2LMHeadModel, GPT2Tokenizer
# 加载预训练模型和分词器
model = GPT2LMHeadModel.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
# 准备输入
input_text = "Once upon a time"
input_ids = tokenizer.encode(input_text, return_tensors="pt")
# 使用KV缓存生成文本
output = model.generate(
input_ids,
max_length=50,
num_return_sequences=1,
temperature=0.7,
use_cache=True # 启用KV缓存
)
# 解码输出
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
print(generated_text)
在这个例子中，use_cache=True参数启用了KV缓存。Hugging Face Transformers库在内部实现了高效的KV缓存管理，包括预分配内存、批处理支持和内存优化等。
12.5 KV缓存优化技术#
虽然基本的KV缓存已经可以显著提高推理效率，但在实际应用中，我们可以使用更多优化技术来进一步提高性能。
12.5.1 内存优化#
KV缓存的主要挑战之一是内存消耗。对于长序列生成，缓存可能占用大量内存。以下是一些内存优化技术：
低精度存储：使用较低精度（如float16或int8）存储KV缓存，可以减少内存占用。例如：
def _create_kv_cache(self, batch_size, device, dtype=torch.float16): # 使用float16
# 创建缓存的代码...
稀疏注意力：对于非常长的序列，可以实现稀疏注意力机制，只关
To save on context only part of this file has been shown to you. You should retry this tool after you have searched inside the file with grep -nin order to find the line numbers of what you are looking for.