Transformer模型在推理时，其自回归生成过程本质上是串行的，但每个生成步骤（time step）内部的自注意力计算可以借助KV缓存（Key-Value Cache）技术进行大幅优化，将计算复杂度从 O(n²) 降低至 O(n)，从而极大地提升了单步内的计算效率和并行度[ref_1][ref_5]。 ### 一、KV缓存的核心原理：避免重复计算 在自回归解码的每一步，模型都需要基于当前已生成的全部token序列（长度为 `t`）来计算自注意力。如果不使用缓存，每次计算都需要为序列中的每个token重新计算其对应的Key（K）和Value（V）向量。这导致了大量的重复计算，因为对于第 `t` 步，前 `t-1` 个token的K、V向量在之前的步骤中已经计算过了[ref_1][ref_3]。 KV缓存的核心思想是：**将每个token在自注意力层中计算出的K、V向量存储下来，供后续生成步骤直接复用**。这样，在生成新token时，只需计算当前新token的K、V向量，并将其与缓存中所有历史token的K、V向量拼接，再进行注意力计算[ref_1][ref_5]。 下表对比了使用KV缓存前后的计算差异： | 生成步骤 | 序列长度 | **无KV缓存的计算** | **有KV缓存的计算** | | :--- | :--- | :--- | :--- | | **Step 1** | 1 | 计算 token₁ 的 Q₁, K₁, V₁ | 计算 token₁ 的 Q₁, K₁, V₁，并缓存 K₁, V₁ | | **Step 2** | 2 | 为 token₁, token₂ **重新计算** K₁, V₁, K₂, V₂ | 从缓存读取 K₁, V₁，仅计算 token₂ 的 K₂, V₂，然后拼接 [K₁, K₂], [V₁, V₂] | | **Step 3** | 3 | 为 token₁, token₂, token₃ **重新计算** K₁, V₁, K₂, V₂, K₃, V₃ | 从缓存读取 K₁, V₁, K₂, V₂，仅计算 token₃ 的 K₃, V₃，然后拼接 [K₁, K₂, K₃], [V₁, V₂, V₃] | | **...** | ... | ... | ... | | **Step t** | t | 需计算 t 个token的K, V，复杂度 O(t²) | 仅计算第t个token的Kₜ, Vₜ，复杂度 O(t) | 通过这种机制，**单步自注意力计算的理论复杂度从序列长度的平方级 O(t²) 降低到了线性级 O(t)**。这使得模型在生成长序列时，每一步的计算开销仅线性增长，而非平方增长，从而显著加速了推理过程[ref_1][ref_5]。 ### 二、KV缓存的工作流程与代码实现 以下是一个结合PyTorch的简化代码示例，清晰地展示了KV缓存如何在Decoder的自注意力层中工作。 ```python import torch import torch.nn as nn import torch.nn.functional as F class CausalSelfAttentionWithKVCache(nn.Module): def __init__(self, d_model, num_heads): super().__init__() self.d_model = d_model self.num_heads = num_heads self.head_dim = d_model // num_heads assert self.head_dim * num_heads == d_model, "d_model must be divisible by num_heads" # 线性投影层，用于计算Q, K, V 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.out_proj = nn.Linear(d_model, d_model) def forward(self, x, cache=None): """ x: 当前步的输入，形状为 [batch_size, seq_len, d_model] cache: 可选的KV缓存，是一个元组 (past_keys, past_values) 每个元素的形状为 [batch_size, num_heads, past_seq_len, head_dim] """ batch_size, seq_len, _ = x.shape # 1. 计算当前输入x的Q, K, V Q = self.q_proj(x) # [batch_size, seq_len, d_model] K = self.k_proj(x) # [batch_size, seq_len, d_model] V = self.v_proj(x) # [batch_size, seq_len, d_model] # 2. 重塑为多头格式 Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # [B, H, L, D] K = K.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) V = V.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # 3. KV缓存处理：拼接历史K/V与当前K/V if cache is not None: past_keys, past_values = cache # 将当前步的K, V与缓存的K, V在序列维度上拼接 K = torch.cat([past_keys, K], dim=2) # dim=2 对应序列长度维度 V = torch.cat([past_values, V], dim=2) # 4. 更新缓存（用于下一步） # 注意：在自回归生成中，当前步的输入x通常只包含最新生成的token（seq_len=1） # 因此，K, V 就是需要缓存的新内容。如果seq_len>1（如预填充），则需要缓存整个序列。 new_cache = (K, V) # 5. 计算缩放点积注意力 # Q: [B, H, L_q, D], K: [B, H, L_kv, D], V: [B, H, L_kv, D] attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5) # [B, H, L_q, L_kv] # 应用因果掩码，确保当前位置不能看到未来的信息 if cache is None: # 第一次调用，没有缓存，需要为整个输入序列生成掩码 mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool() attn_scores = attn_scores.masked_fill(mask[None, None, :, :], float('-inf')) else: # 有缓存时，当前查询（通常是最后一个token）可以关注所有历史token及自身，无需额外掩码 # 因为缓存中的K/V已经按时间顺序排列，且新的Q只与当前位置及之前的K/V计算注意力 pass attn_weights = F.softmax(attn_scores, dim=-1) # [B, H, L_q, L_kv] attn_output = torch.matmul(attn_weights, V) # [B, H, L_q, D] # 6. 将多头输出合并回原始维度 attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) output = self.out_proj(attn_output) return output, new_cache # 返回输出和更新后的缓存 # 模拟自回归推理循环，展示KV缓存的使用 def autoregressive_generation(model, initial_input, max_len=20): """ model: 包含上述自注意力层的Transformer Decoder initial_input: 起始token，形状 [batch_size, 1] """ generated = initial_input cache = None # 初始化缓存为空 for step in range(max_len): # 当前输入是上一步生成的最后一个token（或起始token） current_input = generated[:, -1:] # [batch_size, 1] # 前向传播，传入当前的KV缓存 output, cache = model(current_input, cache=cache) # 从output中获取下一个token的logits（假设模型输出层已包含） next_token_logits = output[:, -1, :] # [batch_size, vocab_size] next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) # [batch_size, 1] # 将新生成的token拼接到已生成序列中 generated = torch.cat([generated, next_token], dim=1) # 检查是否生成结束符（假设END_TOKEN_ID=2） if (next_token == 2).all(): break return generated ``` *代码解析*： 1. **缓存初始化**：在生成开始时，`cache` 被初始化为 `None`。 2. **缓存拼接**：在自注意力层的 `forward` 方法中，如果 `cache` 不为空，则将缓存的 `past_keys` 和 `past_values` 与当前步新计算的 `K`、`V` 在序列长度维度进行拼接（`torch.cat`）。这使得当前的查询（`Q`）能够同时关注所有历史token和当前token的信息[ref_1][ref_3]。 3. **缓存更新**：每次前向传播后，函数返回更新后的 `new_cache`（即拼接后的 `K` 和 `V`），供下一步使用。 4. **掩码处理**：当没有缓存（第一步）时，需要为整个初始序列创建因果掩码。当有缓存时，因为缓存中的K/V是按时间顺序排列的历史信息，新的Q（对应最新token）在计算注意力时，其与K/V的注意力分数矩阵天然是“下三角”形式的一部分，所以通常无需显式应用额外的因果掩码[ref_5]。 5. **单步计算**：在循环中，每一步的输入 `current_input` 通常是单个新token（形状为 `[batch_size, 1]`）。模型只需为该token计算 `Q`，并从缓存中读取所有历史 `K`、`V`，然后计算注意力。这避免了为整个历史序列重新计算 `K`、`V`，实现了计算量的线性增长[ref_1][ref_5]。 ### 三、KV缓存对并行计算的优化体现 KV缓存优化的是**单步内**自注意力计算的**数据准备**和**计算量**，从而让硬件（如GPU）能更高效地执行并行计算： 1. **减少冗余计算，聚焦并行核心**：GPU擅长对大规模数据进行并行处理。如果没有缓存，每一步都需要为所有历史token重新计算K、V投影，这部分计算是重复且冗余的，占用了大量的并行计算资源。KV缓存消除了这部分冗余，使得GPU的并行计算能力可以更集中地应用于**注意力分数的计算（矩阵乘法）和加权求和**这一核心且无法避免的环节[ref_3][ref_6]。 2. **提升内存访问效率**：缓存K、V向量后，这些数据可以常驻在GPU的高速显存中。在计算注意力时，直接读取缓存数据比重新通过线性层计算并加载数据具有更高的内存带宽利用率和更低的延迟，这有利于提升并行计算单元的吞吐量[ref_4][ref_6]。 3. **支持高效的批处理**：在实际推理服务中，通常会同时处理多个用户请求（批处理）。每个请求都有自己的KV缓存。由于每一步的计算量（主要是对新token的Q与缓存的K进行矩阵乘）变得相对固定且更轻量，这使得批处理中不同长度序列的计算更容易达到平衡，GPU的并行计算资源利用率更高[ref_4]。 ### 四、KV缓存的挑战与高级优化 尽管KV缓存带来了巨大的加速，但它也引入了新的挑战，主要是**显存占用随序列长度线性增长**。每个token在每个注意力头、每个模型层都需要缓存其K和V向量。对于长序列生成，缓存可能消耗大量显存[ref_1][ref_4]。 针对此问题，业界提出了多种优化技术，进一步提升了并行计算和资源利用的效率： | 优化技术 | 原理 | 对并行计算的影响 | | :--- | :--- | :--- | | **PagedAttention（如vLLM）** | 将KV缓存组织成固定大小的“块”，类似操作系统内存分页。不同序列的缓存块可以非连续存储，减少内存碎片，允许更灵活、更高吞吐量的批处理[ref_4]。 | 极大地提高了显存利用率，使得GPU能够同时并行处理更多请求的生成任务，提升整体吞吐量。 | | **量化（Quantization）** | 将KV缓存中的浮点数（如FP16）转换为低精度格式（如INT8、INT4）。显著减少缓存的内存占用[ref_4]。 | 降低数据存储和传输开销，使更多数据能同时加载到计算核心附近，提升并行计算效率。 | | **CPU Offloading** | 将不活跃的、较早生成的token的KV缓存从GPU显存转移到主机内存（CPU RAM）[ref_4]。 | 节省宝贵的GPU显存，用于服务更多并发请求或生成长度更长的序列，从系统层面提升并行处理能力。 | | **动态KV缓存** | 根据策略（如重要性评分）选择性丢弃或压缩历史缓存中的部分token信息[ref_1]。 | 在保证生成质量的同时，控制缓存大小，维持单步注意力计算的数据量在一个可高效并行处理的范围内。 | 综上所述，**KV缓存通过避免重复计算Key和Value向量，将Transformer推理时单步自注意力计算的核心部分优化为线性复杂度**。这不仅大幅减少了计算量，更重要的是，它让每一步的计算变得更加规整和轻量，从而允许GPU等硬件更充分、更高效地发挥其大规模并行计算的优势，是Transformer模型实现高效推理的关键技术之一[ref_1][ref_3][ref_5]。