LLM推理优化——KV Cache篇（百倍提速）

LLM推理优化——KV Cache篇（百倍提速）

注意：KV Cache本质上是空间换时间的技术。与计算机组成原理中的cache不同，它不涉及访存优化。

不知道大家在用LLM的时候，有没有注意到一个问题：我们在输入我们的问题之后，需要等待一段漫长的时间才能看到第一个字符的响应，而等待第一个之后，后续的响应却非常之快，这就是使用了KV Cache加速后的带来的优势。

LLM推理过程与自注意力机制

LLM的推理过程和训练过程是有所区别的，LLM在训练过程中使用因果掩码并行化训练，无需像RNN一样等待之前的结果运算结束。但是他的推理过程却类似RNN，并且有大量的重复计算。详细解释可以参照上一篇Blog：LLM的训练与推断。我们已经知道LLM推断时实际上是一种自回归模型$f_{\theta }$，假设在$t$步及其之前的输入和推断表示为$x_{1:t}$，我们可以用公式表示它：

$$x_{1:t+1}=f_{\theta }(x_{1:t})\text{\hspace{1em}(1)}$$

从公式1我们可以看到，即使$x_{1:t}$已经是计算过的，但是在计算$x_{t+1}$时却还要重复计算$x_{1:t}$,当序列愈来愈长，这个计算代价也越来越不可接受，尤其绝大多数计算都是在做无用功（重复计算）。

基于此，我们就有一个朴素的优化思考:我们能否去除掉这些重复计算？那答案是当然可以，不然就没有这篇博客了🤣。

那么如何去除呢？如果有动态规划、搜索基础的小伙伴可能会比较眼熟，类似于常用的剪枝操作。将后续仍然需要使用的计算结果保存起来，后续使用时不必再次计算。

那么来看一下具体是如何实施的，从LLM的训练与推断中我们可以知道，需要重复计算的矛盾点来自于自注意力机制，首先回顾一下它：
$$attn(Q,K,V)=softmax\left(\frac{Q^{T}K}{\sqrt{d_k}}\right)V$$
我们可以用行列向量的形式表示第t步计算$x_t$时的$(Q^{T}K{)}_t$
$$\begin{array}{rl}(Q^{T}K{)}_t& =\left(\begin{array}{cccc}{q}_1^T{k}_1& q_1^T{k}_2& \dots & q_1^T{k}_t\\ q_2^T{k}_1& q_2^T{k}_2& \dots & q_2^T{k}_t\\ ⋮& ⋮& \ddots & ⋮\\ q_t^T{k}_1& q_t^T{k}_2& \dots & q_t^T{k}_t\end{array}\right)\\ (Q^{T}K{)}_{t+1}& =\left(\begin{array}{ccccc}{q}_1^T{k}_1& q_1^T{k}_2& \dots & q_1^T{k}_t& q_1^T{k}_{t+1}\\ q_2^T{k}_1& q_2^T{k}_2& \dots & q_2^T{k}_t& q_2^T{k}_{t+1}\\ ⋮& ⋮& & ⋮& ⋮\\ q_t^T{k}_1& q_t^T{k}_2& \dots & q_t^T{k}_t& q_t^T{k}_{t+1}\\ q_{t+1}^T{k}_1& q_{t+1}^T{k}_2& \dots & q_{t+1}^T{k}_t& q_{t+1}^T{k}_{t+1}\end{array}\right)\end{array}$$

从式子中可以看到，在计算$(Q^{T}K{)}_{t+1}$时，前t行t列已经被计算过，我们只需要计算最后一行并且保存即可。由$q_{t+1}^T{k}_1,q_{t+1}^T{k}_2,\dots ,q_{t+1}^T{k}_t,q_{t+1}^T{k}_{t+1}$可知，我们还需要知道之前每次计算出的$k$值，所以在保存每行值之外，也应当保存$k_1,k_2,\dots ,k_{t+1}$的值。

这里可能会有疑惑，最后一列用了之前的$q_1,q_2,\dots ,q_t$,为什么不保存$q$的值。实际上，$Q^{T}K$矩阵会经过一个掩码操作，他最后得到是一个下三角矩阵，我们不必再次计算，直接补0即可。

保存$V$值亦是同理。

保存$K,V$是不是KV Cache名称的由来呢？