大模型——推理优化—

大模型——推理优化——KV Cache

在本文中，我们将详细介绍KV Cache，这是一种大模型推理加速的方法。
正如其名称所示，该方法通过缓存Attention中的K和V来实现推理优化。

一、大模型推理的冗余计算

我们先简单观察一下基于Decoder架构的大模型的生成过程
用户输入“中国的首都”，模型续写得到的输出为“是北京”，模型的生成过程如下：

将“中国的首都”输入模型，得到每个token的注意力表示（绿色部分）。使用“首都”的注意力表示，预测得到下一个token为“是”（实际还需要将该注意力表示映射成概率分布logits，为了方便叙述，我们忽略该步骤）。
将“是”拼接到原来的输入，得到“中国的首都是”，将其输入模型，得到注意力表示，使用“是”的注意力表示，预测得到下一个token为“北”。
将“北”拼接到原来的输入，依此类推，预测得到“京”，最终得到“中国的首都是北京”

存在的问题：
在每一步生成中，仅使用输入序列中的最后一个token的注意力表示，即可预测出下一个token。但模型还是并行计算了所有token的注意力表示，其中产生了大量冗余的计算（包含qkv映射，attention计算等），并且输入的长度越长，产生的冗余计算量越大。

例如：

在第一步中，我们仅需使用“首都”的注意力表示，即可预测得到“是”，但模型仍然会并行计算出“中国”，“的”这两个token的注意力表示。
在第二步中，我们仅需使用“是”的注意力表示，即可预测得到“北”，但模型仍然会并行计算“中国”，“的”，“首都”这三个token的注意力表示。

二、Self-Attention过程解析

2.1 公式解析

假设输入序列长度为 $n$ ，第 $j$ 个token对于整个输入序列的注意力表示如下公式：

$b^{j} = \sum_{i=1}^{n}softmax(q^{j} \cdot k^{i})v^{i}$

第 $j$ 个token对于整个输入序列的注意力表示的计算步骤大致如下：

向量映射：将输入序列中的每个token的词向量分别映射为 $q,k,v$ 三个向量。
注意力计算：使用 $q^{j}$ 分别与每个 $k$ 进行点乘，得到第 $j$ 个token对每个token的注意力分数。
注意力分数归一化：对注意力分数进行softmax，得到注意力权重。
加权求和：注意力权重与对应的向量 $v$ 加权求和，最终得到第 $j$ 个token的注意力表示。

2.2 过程实例

下面将以图像的方式帮助大家更形象地理解Self Attention。

假设：

$a = a^{1}a^{2}a^{3}a^{4}$
$a^{1}$ 对于整个输入序列 $a$ 的注意力值是 $b^{1}$

根据上面的Self-Attention公式得出：

$b^{1} = \sum_{i=1}^{4}softmax(q^{1} \cdot k^{i})v^{i}$

继续观察 $a^{2}$ 对于整个输入序列 $a$ 的注意力 $b^{2}$ 表示，即：
$b^{2} = \sum_{i=1}^{4}softmax(q^{2} \cdot k^{i})v^{i}$

三、KV Cache

3.1 原理

在推理阶段，当输入长度为 $n$ ，我们仅需使用即可预测出下一个token，但模型却会并行计算出，这部分会产生大量的冗余计算。
而实际上 $b^{n}$ 可直接通过公式 $b^{n} = \sum_{i=1}^{n}softmax(q^{n} \cdot k^{i})v^{i}$ 算出，即 $b^{n}$ 的计算只与 $q^{n}$ 、所有 $k$ 和 $v$ 有关
KV Cache的本质是以空间换时间，它将历史输入的token的 $k$ 和 $v$ 缓存下来，避免每步生成都重新计算历史token的 $k$ 和 $v$ 以及注意力表示 $b^{1}...b^{n-1}$ ，而是直接通过 $b^{n} = \sum_{i=1}^{n}softmax(q^{n} \cdot k^{i})v^{i}$ 的方式计算得到 $b^{n}$ ，然后预测下一个token。

3.2 KV cache过程

第一步生成时，缓存 $K,V$ 均为空，输入为“中国的首都”，模型将按照常规方式并行计算：

并行计算得到每个token对应的 $k,v$ ，以及注意力表示 $b^{1},b^{2},b^{3}$ 。
使用 $b^{3}$ 预测下一个token，得到“是”。
更新缓存，令 $K=[k^{1},k^{2},k^{3}],V=[v^{1},v^{2},v^{3}]$ 。

第二步生成时，计算流程如下：

仅将“是”输入模型，对其词向量进行映射，得到 $q^{4},k^{4},v^{4}$ 。
更新缓存，令 $K=[k^{1},k^{2},k^{3},k^{4}],V=[v^{1},v^{2},v^{3},v^{4}]$ 。
计算 $b^{4} = \sum_{i=1}^{4}softmax(q^{4} \cdot k^{i})v^{i}$ ，预测下一个token，得到“北”

第三步生成时，计算流程如下：

仅将“北”输入模型，对其词向量进行映射，得到 $q^{5},k^{5},v^{5}$ 。
更新缓存，令 $K=[k^{1},k^{2},k^{3},k^{4},k^{5}],V=[v^{1},v^{2},v^{3},v^{4},v^{5}]$ 。
计算 $b^{5} = \sum_{i=1}^{5}softmax(q^{5} \cdot k^{i})v^{i}$ ，预测下一个token，得到“京”

上述生成流程中，只有在第一步生成时，模型需要计算所有token的 $k,v$ ，并且缓存下来。
此后的每一步，仅需计算当前token的 $q^{n},k^{n},v^{n}$ ，更新缓存 $K,V$ ，然后使用 $q^{n},K,V$ 即可算出当前token的注意力表示，最后用来预测一下个token。

3.3 代码修改

这里参考gpt2里面的代码实现

https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py

query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)

query = self._split_heads(query, self.num_heads, self.head_dim) # 当前token对应的query
key = self._split_heads(key, self.num_heads, self.head_dim) # 当前token对应的key
value = self._split_heads(value, self.num_heads, self.head_dim) # 当前token对应的value

if layer_past is not None:
past_key, past_value = layer_past # KV Cache
key = torch.cat((past_key, key), dim=-2) # 将当前token的key与历史的K拼接
value = torch.cat((past_value, value), dim=-2) # 将当前token的value与历史的V拼接

if use_cache is True:
present = (key, value)
else:
present = None

# 使用当前token的query与K和V计算注意力表示
if self.reorder_and_upcast_attn:
attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask)
else:
attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)