【LLM】大模型推理加速 KV-Cache

模型推理过程

在了解KVCache之前，我们需要知道Transformer类大模型的推理过程。
对于LLM进行一次前向传播也就是生成一个token的过程可以被分解成以下步骤：

• 文本$T_{input}$经过Tokenizer后得到n个Token，即 { T 1 , T 2 , . . . , T n } \{T_1,T_2,...,T_n\} {T1​,T2​,...,Tn​}
• Token经过嵌入层后得到Token Embedding { x 1 0 , x i 0 , . . . , x n 0 } \{x_1^0,x_i^0,...,x_n^0\} {x10​,xi0​,...,xn0​}，其中0代表第0层，$x$是D维向量
• Token Embedding 经过L层变换之后，得到 { x 1 L , x i L , . . . , x n L } \{x_1^L,x_i^L,...,x_n^L\} {x1L​,xiL​,...,xnL​}
• generation阶段，将最后一层的最后一个Token Embedding即$x_n^L$取出，与lm_head中的vocabulary Embedding { e 1 , e 2 , . . . , e V } \{e_1,e_2,...,e_V\} {e1​,e2​,...,eV​}进行运算，得到概率 { p 1 , p 2 , . . . , p V } \{p_1,p_2,...,p_V\} {p1​,p2​,...,pV​}，最后从概率采样中选择一个产生新的$T_{n+1}$

流程如图所示

如果我们一直用$T_{[1,i]}$去生成$T_{i+1}$也是可行的，但很容易发现一个问题，Token Embedding中$[1,i]$这部分的值是不需要重新计算的。

我们发现每次计算$T_{i+1}$时，只需要额外计算$x_i^{\mathrm{0...}L}$的值就可以了。

KV Cache原理

当我们有了 { x 1 l , x 2 l , . . . , x n l } \{x_1^l,x_2^l,...,x_n^l\} {x1l​,x2l​,...,xnl​}的数据，再加上新来的$x_{n+1}^l$计算$x_{n+1}^{l+1}$的计算量只是一次attention query。

我们只关注$x_{n+1}^{l+1}$：
$$s_i=(W_Q{x}_{n+1}^l{)}^T(W_K{x}_i^l),1\le i\le n+1$$

$$y_h=\sum _{i=1}^{n+1}\frac{s_i}{{\sum }_{j=1}^{n+1}{e}_j}{W}_V{x}_i^l,1\le i\le n+1$$

$$x_{n+1}^{l+1}=Conca{t}_1^H{y}_h$$

其中$W_{Q/K/V}$是QKV的投影矩阵，H是多头自注意力中head的个数。中间省略了dot product、layernorm、feed forward等步骤。

过程如图片所示。图片来源

在每次计算过程中得到的中间值$W_K{x}_i^l$和$W_V{x}_i^l$，将它们保存下来，就得到了K Cache和V Cache。

KV Cache的存储

KV Cache的总大小是2nHD，其中n是token数量，H是head数量，D是$x_i^l$的维度。
目前有三种解决方案：

• 分配一个最大容量的缓冲区，但需要提前知道最大的token数量。
• 动态分配缓冲区大小，类似vector的方式。
• 将数据拆散，按最小单元格存储，用一份元数据记录每一块数据的位置。

最后一种方式也就是现在常用的PageAttention，也是vllm的核心技术。