4. 自注意力机制推理：单步计算、KV Cache的引入与原理

好，咱们今天聊点硬核的。自注意力机制，说白了就是Transformer能“看懂”上下文的关键。但推理的时候，如果每一步都从头算一遍，那效率就太低了。我当年第一次跑大模型推理时，看着那慢吞吞的速度，心里就一个想法——这玩意儿肯定能优化。

嗯，今天我们就来拆解一下：单步计算是怎么做的？KV Cache又是怎么救场的？

4.1 单步自注意力计算：从零开始

先看最基础的情况。假设我们有一个序列，长度是 n，每个位置对应一个向量。自注意力的目标就是让每个位置“看到”所有其他位置。

计算过程分三步：

  1. 生成Q、K、V：每个输入向量 x_i 乘以三个不同的权重矩阵，得到 Q_i、K_i、V_i。
  2. 计算注意力分数：score = Q * K^T / sqrt(d_k)，然后做 softmax。
  3. 加权求和：用 softmax 后的权重去加权 V，得到输出。

公式长这样：

Attention(Q, K, V) = softmax(Q * K^T / sqrt(d_k)) * V

这里 d_k 是每个头的维度。为什么要除以 sqrt(d_k)？我刚开始学的时候也纳闷。后来在项目中调参发现，如果不除，softmax 后的分布会特别“尖”，梯度容易消失。说白了，就是让数值更稳定。

4.2 推理时的痛点：重复计算

训练时，我们一次性把整个序列扔进去，并行计算，爽得很。

但推理不一样。推理是自回归的——每次只生成一个 token，然后把这个 token 拼到序列末尾，再算下一步。

你想想看，如果每次生成新 token 时，都把前面所有 token 的 Q、K、V 重新算一遍，那计算量就是 O(n^2) 的。序列越长，越慢。我见过一个项目，序列长度到 2048 时，推理速度直接掉了两个数量级。

注意： 千万别在推理时重复计算前面 token 的 K 和 V。这是新手最容易踩的坑。我曾经在优化一个对话模型时，发现 GPU 利用率极低，一查代码，好家伙，每次生成都重新算了一遍完整的注意力矩阵。

4.3 KV Cache：用空间换时间

那怎么优化呢？答案就是 KV Cache。

核心思想很简单：既然每一步生成的 K 和 V 只依赖于当前 token，而且前面的 token 不会再变了，那我们就把它缓存起来。

具体做法：

  • 维护两个缓存数组：K_cache 和 V_cache。
  • 每生成一个新 token，只计算它的 K_new 和 V_new。
  • 把 K_new 拼接到 K_cache 末尾，V_new 同理。
  • 计算注意力时，直接用完整的 K_cache 和 V_cache，而不是重新算。

这样，每一步的计算量就从 O(n^2) 降到了 O(n)。注意，这里的 n 是当前序列长度。虽然还是线性增长，但相比平方级，已经是天壤之别了。

关键点： KV Cache 只缓存 K 和 V，不缓存 Q。因为 Q 每次都要跟所有 K 做点积，所以 Q 必须实时计算。

4.4 流程图：KV Cache 的工作流程

下面这张图，是我用 SVG 画的，展示了带 KV Cache 的自注意力推理流程。你看一眼就明白了。

输入序列 Token 1 Token 2 Token 3 Token 4 ← 当前生成 计算 K/V K1, V1 K2, V2 K3, V3 K4, V4 KV Cache [K1,V1] [K2,V2] [K3,V3] [K4,V4] 注意力计算 Q4 * [K1..K4]^T → softmax → * V 输出 4 输入 Token 已缓存 K/V 新计算 K/V KV Cache 存储

你看，每次只计算新 token 的 K 和 V，然后拼到缓存里。注意力计算时，Q 只跟缓存里的所有 K 做点积。这样，每一步的计算量就只跟当前序列长度有关，而不是平方级。

4.5 实际实现中的细节

在实际代码里，KV Cache 通常是一个张量，形状是 [batch_size, num_heads, max_seq_len, head_dim]。但注意，我们不会一开始就分配最大长度，而是动态增长的。

我习惯用 torch.cat 来拼接，但更好的做法是预分配一块连续内存，然后通过索引写入。这样可以避免频繁的内存分配和拷贝。

下面是一个简化的伪代码：

# 初始化 KV Cache
K_cache = torch.zeros(batch_size, num_heads, max_len, head_dim)
V_cache = torch.zeros(batch_size, num_heads, max_len, head_dim)
step = 0

def generate_next_token(x):
    # x: 当前 token 的 embedding
    Q = linear_q(x)
    K = linear_k(x)
    V = linear_v(x)
    
    # 存入缓存
    K_cache[:, :, step:step+1, :] = K
    V_cache[:, :, step:step+1, :] = V
    
    # 取有效部分
    K_eff = K_cache[:, :, :step+1, :]
    V_eff = V_cache[:, :, :step+1, :]
    
    # 计算注意力
    scores = torch.matmul(Q, K_eff.transpose(-2, -1)) / sqrt(d_k)
    attn = softmax(scores, dim=-1)
    out = torch.matmul(attn, V_eff)
    
    step += 1
    return out
小技巧： 在推理时，可以把 Q、K、V 的线性层合并成一个矩阵乘法，减少 kernel launch 的开销。我试过，在 batch size 为 1 时，能提升 10%-15% 的吞吐。

4.6 避坑指南

最后，分享几个我踩过的坑：

  • 缓存长度管理：别忘了在每一步更新 step 计数器。我曾经忘记更新，结果每一步都只看到第一个 token，生成的内容完全不对。
  • 内存泄漏：如果序列特别长，KV Cache 会占用大量显存。我建议设置一个最大长度，超过后做截断或压缩。
  • 多 batch 推理：不同序列的长度可能不同，KV Cache 需要做 padding 或使用变长支持。我一般用 attention_mask 来处理。

嗯，KV Cache 的原理其实不复杂，但实现细节决定了推理效率。你想想看，如果没有它，大模型推理根本跑不起来。所以，理解它、用好它，是每个做推理优化的工程师的必修课。

公众号:蓝海资料掘金营，微信deep3321

← 上一章 📖 返回目录 下一章 →