KV Cache的原理与实现

爱听歌的周童鞋

已于 2025-02-19 21:14:46 修改

阅读量1.2k

点赞数 20

分类专栏： LLM 文章标签： LLM 推理框架 KV Cache

于 2024-12-31 21:15:26 首次发布

本博客上原创文章未经本人许可，不得用于商业用途。转载请注明出处，否则保留追究法律责任的权利

本文链接：https://blog.csdn.net/qq_40672115/article/details/144857492

版权

LLM 专栏收录该内容

20 篇文章

订阅专栏

前言

KuiperInfer 作者推出的《自制大模型推理框架课程》，链接。记录下个人学习笔记，仅供自己参考😄

video：自制大模型推理框架-KVCache动手实现-原理篇

repo：https://github.com/zjhellofss/KuiperLLama

refer：大模型推理加速：看图学KV Cache

本次课程我们来学习第十五次课程—KV Cache 的原理与实现

课程大纲可以看下面的思维导图

在这里插入图片描述

0. 简述

KV Cache 已成为当前大型模型推理框架的标准配置，其主要功能是通过牺牲一定的空间复杂度来换取时间效率的提升，从而加快大型模型的推理预测速度。但是它只能用于 Decoder 架构的模型如 GPT、LLama 等，这是因为 Decoder 有 Causal Mask，在推理的时候前面已经生成的字符不需要与后面的字符产生 attention，从而使得前面已经计算的 K 和 V 可以缓存起来。

下面我们就来一起看下 KV Cache 具体是怎么做的

Note：关于 Transformer 和注意力机制大家感兴趣的可以看看李宏毅老师或者 3Blue1Brown 的视频讲解

1. Attention的计算

我们先看下不使用 KV Cache 的推理过程，假设模型最终生成了 “遥遥领先” 四个字

当模型生成第一个 “遥” 字时，input=“<s>”，“<s>” 是 起始字符，此时 attention 的计算如下：

在这里插入图片描述

Note：为了看上去方便，我们 暂时忽略 scale 项 $\sqrt{d}$

第 1 步计算的时候

当步长等于 1 的时候，我们有一个维度为 1xdim 的输入 Token，1 是输入 token 的个数，通过查询矩阵 $W_q$ 和键矩阵 $W_k$ 将输入序列（input token）映射得到 $Q$ 和 $K$ ，而 $QK^T$ 表示了所有可能 键—查询对 之间点积的网格，接着对每行做 softmax，此时可以将每一行看作权重，表示查询和键的相关度，我们称这个网格为 注意力模式 即 Attention Pattern

根据注意力模式就能让模型推断出每个词与哪些词有关，然后要做的就是去更新嵌入向量，此时会用到值矩阵 $W_v$ ，通过矩阵 $W_v$ 将输入序列映射得到 $V$ ，与前面的 Attention Pattern 相乘就得到最终的注意力输出

在步长等于 1 的时候 $Q$ 的维度是 1xdim， $K$ 的维度同样为 1xdim，而 $QK^T$ 得到的是一个 1x1 的权重，随后利用 $V$ 对 1xdim 的矩阵进行加权即可，用公式表示为：
$\color{gold}Att_1\color{black}(Q,K,V)=\mathrm{softmax}(\color{gold}Q_1\color{black}K_1^T)\overrightarrow{V_1}=\mathrm{softmaxed}(\color{gold}Q_1\color{black}K_1^T)\overrightarrow{V_1}$
Note：softmaxed 表示已经按行进行了 softmax

当模型生成第二个 “遥” 字时，input=“<s>遥”，此时 attention 的计算如下：

在这里插入图片描述

第 2 步计算的时候

在第二步的计算过程中，我们处理了两个输入 Token，维度是 2xdim。因此，当这些 Token 转换为嵌入向量与 $W_q$ 和 $W_k$ 矩阵相乘进行映射时，我们得到的 $Q$ 和 $K$ 的维度是 2xdim。但是第 2 步在 $Q$ 矩阵与 $K^T$ 矩阵相乘时，我们必须注意一个关键点： $Q$ 矩阵的第一行与 $K^T$ 矩阵的第 2 列是不会相乘的，这一机制是 Transformer Decoder 结构种的 Casual Mask 机制。

该机制的目的在于，在计算注意力（Attention）的过程种，将这些 Token 从注意力机制种屏蔽掉，确保模型在预测时仅关注过去和当前的 token，不能让后面的词影响前面的词，从而使得模型基于每个时间步骤可用的信息进行预测。也就是说需要强制将 $QK^T$ 矩阵左上方位置的值都设为负无穷，这样应用 softmax 之后，它们就都会变为 0。对于 $Q$ 矩阵的第一行而言， $K$ 矩阵的第二列代表的就是未来的数据

用公式表示第二步计算的过程如下：
$\begin{aligned} \color{gold}Att_{step2}\color{black}(Q,K,V) & =\mathrm{\color{black}softmax}( \begin{bmatrix} \color{gold}Q_1\color{black}K_1^T & \color{black}-\infty \\ \color{red}Q_2\color{black}K_1^T & \color{red}Q_2\color{black}K_2^T \end{bmatrix}) \begin{bmatrix} \color{black}\overrightarrow{V_1} \\ \color{black}\overrightarrow{V_2} \end{bmatrix} \\ & =( \begin{bmatrix} \color{black}\mathrm{softmaxed}(\color{gold}Q_1\color{black}K_1^T) & \color{black}\mathrm{softmaxed}(-\infty) \\ \color{black}\mathrm{softmaxed}(\color{red}Q_2\color{black}K_1^T) & \color{black}\mathrm{softmaxed}(\color{red}Q_2\color{black}K_2^T) \end{bmatrix}) \begin{bmatrix} \color{black}\overrightarrow{V_1} \\ \color{black}\overrightarrow{V_2} \end{bmatrix} \\ & =( \begin{bmatrix} \color{black}\mathrm{softmaxed}(\color{gold}Q_1\color{black}K_1^T) & \color{black}0 \\ \color{black}\mathrm{softmaxed}(\color{red}Q_2\color{black}K_1^T) & \color{black}\mathrm{softmaxed}(\color{red}Q_2\color{black}K_2^T) \end{bmatrix}) \begin{bmatrix} \color{black}\overrightarrow{V_1} \\ \color{black}\overrightarrow{V_2} \end{bmatrix} \\ & =( \begin{bmatrix} \color{black}\mathrm{softmaxed}(\color{gold}Q_1\color{black}K_1^T)\times\overrightarrow{V_1}+0\times\overrightarrow{V_2} \\ \color{black}\mathrm{softmaxed}(\color{red}Q_2\color{black}K_1^T)\times\overrightarrow{V_1}+\mathrm{softmaxed}(\color{red}Q_2\color{black}K_2^T)\times\overrightarrow{V_2} \end{bmatrix}) \\ & =( \begin{bmatrix} \color{black}\mathrm{softmaxed}(\color{gold}Q_1\color{black}K_1^T)\times\overrightarrow{V_1} \\ \color{black}\mathrm{softmaxed}(\color{red}Q_2\color{black}K_1^T)\times\overrightarrow{V_1}+\mathrm{softmaxed}(\color{red}Q_2\color{black}K_2^T)\times\overrightarrow{V_2} \end{bmatrix}) \end{aligned}$
根据上面的推理，第二步的注意力 Attention 的计算公式为：
$\begin{aligned} & \color{gold}Att_{1}\color{black}(Q,K,V)=\mathrm{softmaxed}(\color{gold}Q_{1}\color{black}K_{1}^{T})\overrightarrow{V_{1}} \\ & \color{red}Att_{2}\color{black}(Q,K,V)=\mathrm{softmaxed}(\color{red}Q_{2}\color{black}K_{1}^{T})\overrightarrow{V_{1}}+\mathrm{softmaxed}(\color{red}Q_{2}\color{black}K_{2}^{T})\overrightarrow{V_{2}} \end{aligned}$
我们可以发现，由于 $Q_1K_2^T$ 这个值会被 mask 掉，因此：

$Q_1$ 在第二步参与的计算与第一步是一样的，而且第二步生成的 $V_1$ 也仅仅依赖于 $Q_1$ ，与 $Q_2$ 毫无关系
$V_2$ 的计算也仅仅依赖于 $Q_2$ ，与 $Q_1$ 毫无关系

当模型生成第三个 “领” 字时，input=“<s>遥遥”，此时 attention 的计算如下：

在这里插入图片描述

第 3 步计算的时候

第三步计算同理，详细的推导和第二步类似，其计算公式为：
$\begin{aligned} & \color{gold}Att_{1}\color{black}(Q,K,V)=\mathrm{softmaxed}(\color{gold}Q_{1}\color{black}K_{1}^{T})\overrightarrow{V_{1}} \\ & \color{red}Att_{2}\color{black}(Q,K,V)=\mathrm{softmaxed}(\color{red}Q_{2}\color{black}K_{1}^{T})\overrightarrow{V_{1}}+\mathrm{softmaxed}(\color{red}Q_{2}\color{black}K_{2}^{T})\overrightarrow{V_{2}} \\ & \color{purple}Att_3\color{black}(Q,K,V)=\mathrm{softmaxed}(\color{purple}Q_{3}\color{black}K_{1}^{T})\overrightarrow{V_{1}}+\mathrm{softmaxed}(\color{purple}Q_{3}\color{black}K_{2}^{T})\overrightarrow{V_{2}}+\mathrm{softmaxed}(\color{purple}Q_{3}\color{black}K_{3}^{T})\overrightarrow{V_{3}} \end{aligned}$
同样的 $Att_k$ 只与 $Q_k$ 有关

当模型生成第四个 “先” 字时，input=“<s>遥遥领”，此时 attention 的计算如下：

在这里插入图片描述

第 4 步计算的时候

第四步和之前类似，其计算公式为：
$\begin{aligned} & \color{gold}Att_{1}\color{black}(Q,K,V)=\mathrm{softmaxed}(\color{gold}Q_{1}\color{black}K_{1}^{T})\overrightarrow{V_{1}} \\ & \color{red}Att_{2}\color{black}(Q,K,V)=\mathrm{softmaxed}(\color{red}Q_{2}\color{black}K_{1}^{T})\overrightarrow{V_{1}}+\mathrm{softmaxed}(\color{red}Q_{2}\color{black}K_{2}^{T})\overrightarrow{V_{2}} \\ & \color{purple}Att_3\color{black}(Q,K,V)=\mathrm{softmaxed}(\color{purple}Q_{3}\color{black}K_{1}^{T})\overrightarrow{V_{1}}+\mathrm{softmaxed}(\color{purple}Q_{3}\color{black}K_{2}^{T})\overrightarrow{V_{2}}+\mathrm{softmaxed}(\color{purple}Q_{3}\color{black}K_{3}^{T})\overrightarrow{V_{3}} \\ & \color{darkred}Att_4\color{black}(Q,K,V)=\mathrm{softmaxed}(\color{darkred}Q_{4}\color{black}K_{1}^{T})\overrightarrow{V_{1}}+\mathrm{softmaxed}(\color{darkred}Q_{4}\color{black}K_{2}^{T})\overrightarrow{V_{2}}+\mathrm{softmaxed}(\color{darkred}Q_{4}\color{black}K_{3}^{T})\overrightarrow{V_{3}}+\mathrm{softmaxed}(\color{darkred}Q_{4}\color{black}K_{4}^{T})\overrightarrow{V_{4}} \end{aligned}$
从图中以及上述公式中可以看出，我们在第四步中积攒了大量的重复计算，这些信息明明在前三步中已经被计算过了，却还是要在第四步中做重复的计算。

根据上面的图和公式我们可以得到以下结论：

当前计算方式存在大量冗余计算
$Att_k$ 只与 $Q_k$ 有关
推理第 $x_k$ 个字符的时候只需要输入字符 $x_{k-1}$ 即可

我们每一步其实只需要根据 $Q_k$ 计算 $Att_k$ 就可以，之前已经计算的 Attention 完全不需要重新计算。但是 $K$ 和 $V$ 是全程参与计算的，所以这里我们需要把每一步的 $K$ 和 $V$ 缓存起来，所以说叫 KV Cache 好像有点不太对，因为 KV 本来就需要全程计算，可能叫增量 KV 计算会跟更好理解。

2. 简化Attention的计算

具体的我们可以将原本 kxdim 维度的 Query 矩阵拆分为两部分：

第一部分是包含第 0 行至 k-1 行的 Query1 矩阵，其维度是 (k-1)xdim
第二部分是仅包含第 k 行的 Query2 矩阵，维度是 1xdim，在执行第 k 次自回归计算时，我们只需要关注 Query 矩阵的第 k 行即 Query2 矩阵和 Key 矩阵第 0 到 k 列进行矩阵相乘就可以了

此时上述 Attention 的计算流程简化后如下图所示：

在这里插入图片描述

3. KV Cache

3.1 K-Cache

KV Cache 顾名思义就是缓存一部分 K 矩阵和 V 矩阵，就像我们前面所说的 K 矩阵同样可以分为 Key1 矩阵和 Key2 矩阵，Key1 矩阵是在前 K-1 步骤中得到的，Key2 是在当前步骤中得到的。

我们从前面的图中也可以看出下一步的 Key 矩阵中的前 k-1 列和上一步中 Key 矩阵的前 k 列是相同的，例如在第 4 步中 Key 矩阵的第 1、第 2 和第 3 列与第 3 步中的 Key 矩阵是相同的。

为什么会相同呢，这是因为在第 3 步中 Key 矩阵是这么计算得到的：
$[\mathrm{input \ token \ 1},\mathrm{input \ token \ 2},\mathrm{input \ token \ 3}]\times W_k$
将每个 token 通过 embedding 转换为嵌入向量之后与查询矩阵 $W_k$ 相乘即可得到相应的查询向量

而在第 4 中同理有：
$[\mathrm{input \ token \ 1},\mathrm{input \ token \ 2},\mathrm{input \ token \ 3},\mathbf {input \ token \ 4 \ 当前步骤的新词}]\times W_k$
因此我们可以发现当我们的 token 转换为嵌入向量与查询矩阵相乘产生查询向量时，其实相同 token 产生的查询向量是一样的，也就是说第 4 步中得到的 Key 矩阵的前三列和第 3 步中的 Key 矩阵是相同的

既然是相同的，我们为什么不在计算 Key 矩阵的时候保存前 K-1 列，在第 k 步的时候只计算第 k 列，即当前的 input token x Wk 矩阵，随后再拼接成一个完整的 Key 矩阵，那这种记录 Key 矩阵前 k-1 列的方法就是所谓的 K-Cache。

3.2 V-Cache

对于 V 矩阵同样的，V 矩阵也是由输入的多个 token 的嵌入向量和值矩阵 $W_v$ 相乘得到的，在第 k 步和第 k-1 步得到的 V 矩阵有一部分是重叠的

所以我们需要将前一步计算出来的 Value1 矩阵放在一块 Value-Cache 区域中，等到当前步的时候我们只需要将当前的输入 token 和 Wv 矩阵进行相乘得到 Value2 矩阵，再将它们拼接起来就可以得到完整的 Value 矩阵并开始注意力的计算。

4. 为KV-Cache申请空间

从前面的分析中我们可以得知，为了给执行了 k 步计算的 Transformer 块准备 KV Cache 的空间，我们所需申请的空间大小应为 kxdim，如果整个模型包含了 N 个这样的 Transformer 块，那么总共需要申请的显存空间将是 Nxkxdim

在 KuiperLLama 这个 repo 中为 KV Cache 申请空间的代码主要是在 kuiper/source/model/llama3.cpp#L469，如下所示：

// kv cache
tensor::Tensor key_cache(base::DataType::kDataTypeFp32, config_->layer_num_, config_->seq_len_,
                          config_->kv_dim_, true, alloc);
tensor::Tensor value_cache(base::DataType::kDataTypeFp32, config_->layer_num_, config_->seq_len_,
                            config_->kv_dim_, true, alloc);

由于在每一步自回归预测中，我们都会增加一个输入单词，因此单词的总数量最多不会超过最大序列长度 max_seq_len，因此所需申请的显存空间大小可以表示为 max_seq_lenxNxdim，所以我们在 LLama2Model::init_mem() 方法中申请了这块内存空间用于后续使用

5. 对KV-Cache空间的拆分

那我们在每个步长中对于 KV-Cache 空间怎么来进行拆分呢？在 KuiperLLama 中对应的代码在 kuiper/source/model/model.cpp#L22，如下所示：

std::pair<tensor::Tensor, tensor::Tensor> Model::slice_kv_cache(int32_t layer_idx,
                                                                int32_t token_pos) const {
  int32_t layer_offset = layer_idx * config_->seq_len_ * config_->kv_dim_;
  int32_t cache_offset = layer_offset + token_pos * config_->kv_dim_;

  float* key_cache_ptr =
      const_cast<float*>(get_buffer(ModelBufferType::kKeyCache).ptr<float>(cache_offset));
  float* val_cache_ptr =
      const_cast<float*>(get_buffer(ModelBufferType::kValueCache).ptr<float>(cache_offset));

  tensor::Tensor key(base::DataType::kDataTypeFp32, config_->kv_dim_, false, nullptr,
                     key_cache_ptr);
  tensor::Tensor val(base::DataType::kDataTypeFp32, config_->kv_dim_, false, nullptr,
                     val_cache_ptr);
  key.set_device_type(device_type_);
  val.set_device_type(device_type_);
  return {key, val};
}

简单来说，我们首先需要索引到第几个 transformer 块即 layer_offser，接着找到 token pos 步的存放位置，最终得到的就是第 i 个 transformer 块中第 k 步的 KV-Cache 空间即 cache_offset，可以简单的理解为 kv_cache[layer_index,token_pos,:]

索引完成之后，我们将这个索引位置对应的区域封装为两个 Tensor，并将这两个 Tensor 指定特定类型后返回，以便用于存储缓存数据。

总结

这节课我们学习了 KV-Cache 的原理，Transformer 的 Attention 计算过程可以简化，只需要计算当前 token 产生的 attention 分数即可，通过 KV-Cache 将 K 和 V 矩阵的值进行缓存减少重复计算。

OK，以上就是第 15 次课的全部内容了，下节我们将去学习 KV Cache 的实现和自注意力计算中的 QKV，敬请期待😄

这是博主 2024 年度的最后一篇文章了，感谢大家长期以来的关注，我们 2025 年见，最后提前祝大家新年快乐🎉🎉🎉