KV Cache 节省的计算次数推演

我们知道 LLM 语言模型的基础架构源自 Transformer。而 LLM 在推理时基本只使用 Decoder 模块。Decoder 的核心部分可简单分为 2 块：1. Attention 2. MLP。我们的 K,V Cache 发生在 Attention 阶段。

我们做如下假设：

1. 输入 $x$ 只有 1 条（Batch=1），维度为 $T \times C$，$T$ 表示输入的 token 个数，$C$ 表示每个 token 用长度为 $C$ 的向量表示。
2. 权重 $Q, K, V$ 矩阵的维度为 $head\_size \times C$，$head\_size$ 表示多头 Attention 中每头的维度，一般 $head\_size < C$。

为论证简单，我们以计算 $K$ 为例：

$$K = x W^T$$

其中：

• $x \in \mathbb{R}^{T \times C}$
• $W \in \mathbb{R}^{head\_size \times C}$，实际计算中使用 $W^T \in \mathbb{R}^{C \times head\_size}$
• 根据矩阵乘法，$K \in \mathbb{R}^{T \times head\_size}$

这一步本质是对输入 $x$ 进行一次线性投影，简单理解在空间中找到另一个投影点，维度从 $C$维 变成 $head\_size$ 维。

为方便说明，先从 $T=1$ 开始。

首次计算：

$$K = x W^T$$

具体表示为：

$$[x_{11}, x_{12}, \dots, x_{1C}] \cdot \begin{bmatrix} w_{11} & w_{12} & \cdots & w_{1\,head\_size} \\ w_{21} & w_{22} & \cdots & w_{2\,head\_size} \\ \vdots & \vdots & \ddots & \vdots \\ w_{C1} & w_{C2} & \cdots & w_{C\,head\_size} \end{bmatrix} = [k_{11}, k_{12}, \dots, k_{1\,head\_size}]$$

其中：

$$k_{11} = \sum_{i=1}^{C} x_{1i} w_{i1} \\ k_{12} = \sum_{i=1}^{C} x_{1i} w_{i2} \\ \vdots \\ k_{1\,head\_size} = \sum_{i=1}^{C} x_{1i} w_{i\,head\_size}$$

如果我们将一次内积乘法算作 1 次计算（此处将整个向量内积抽象为一次计算，便于复杂度分析），那么 $T=1$ 时共做了 $head\_size$ 次计算。

当 $T=2$ 时：

$$\begin{bmatrix} x_{11} & x_{12} & \cdots & x_{1C} \\ x_{21} & x_{22} & \cdots & x_{2C} \end{bmatrix} \cdot \begin{bmatrix} w_{11} & w_{12} & \cdots & w_{1\,head\_size} \\ w_{21} & w_{22} & \cdots & w_{2\,head\_size} \\ \vdots & \vdots & \ddots & \vdots \\ w_{C1} & w_{C2} & \cdots & w_{C\,head\_size} \end{bmatrix} = \begin{bmatrix} k_{11} & k_{12} & \cdots & k_{1\,head\_size} \\ k_{21} & k_{22} & \cdots & k_{2\,head\_size} \end{bmatrix}$$

此时总计算次数为 $2 \times head\_size$ 次。

对比不同输入长度 $T$ 下的计算次数：

所以，当从第 1 个 token 一直推理到第 $n$ 个 token 时，仅以 $K$ 的计算为例——

• 关 KV Cache 的累计计算次数 = $(1+2+3+\dots+n) \times head\_size = \frac{n(n+1)}{2} \times head\_size$
• 开 KV Cache 的累计计算次数 = $n \times head\_size$

在线性投影阶段，计算开销从 $O(n^2)$ 降低到 $O(n)$。对于 $V$ 的计算，情况完全一样，本文以 $K$ 为代表。

总结：KV Cache 推理在线性投影阶段计算开销对比