📘今日学习总结
KV Cache
作用
- 在 Transformer 模型中,KV Cache(Key-Value Cache,键值缓存) 指的是 自回归推理过程中缓存的 Key 和 Value,以避免重复计算。
- 它主要用于 自注意力机制(Self-Attention),在 解码(Decoder) 过程中缓存之前计算过的 Key(键)和 Value(值),从而加速生成。
- 在没有 KV Cache 的情况下,每次生成一个新的 token 时,模型都需要重新计算所有 token 的 Query(查询)、Key(键)、Value(值),导致计算冗余。
- 使用 KV Cache 后,只需要计算新 token 的 Query,并与已缓存的 KV 进行注意力计算,大幅提高推理速度,计算复杂度从 O(N²) 降到 O(N),显著加速推理。
原理
- 在标准 Transformer 自注意力(Self-Attention)中,每个 token 通过 QKV 机制 进行计算
\[Attention(Q, K, V) = \text{softmax}\left(\frac{QK^{T}}{\sqrt{d_{k}}}\right)V
\]
- 在训练阶段,所有 token 可以并行计算
QK^T,但在 自回归推理(Auto-Regressive Decoding)时,每个 token 只能依赖前面已经生成的 token,这就是 KV Cache 需要优化的地方。
\[Attention(Q_n, [K_1, K_2,..., K_{n-1}], [V_1, V_2, ..., V_{n-1}])
\]
KV Cache 的挑战与未来优化
显存占用
- KV Cache 需要存储所有已生成 token 的 Key 和 Value,对长序列推理时显存占用大,尤其是在多头注意力(Multi-Head Attention)中,每一层都需要缓存 K, V。
- 解决方案:
- Paged KV Cache:采用分块存储,减少 GPU 显存压力。
- FlashAttention:优化 GPU 访问 KV Cache 的方式,降低显存消耗。
动态 KV Cache
- 传统 KV Cache 需要线性存储 K, V,但在多轮对话、长文本生成中,需要删除无用缓存,避免显存爆炸。
- 解决方案:
- Sliding Window KV Cache:仅保留最近的 N 个 token 进行 KV 计算。
- 精细化 KV 复用:减少长序列存储需求,提高缓存利用率。