1. 什么是 PD 分离

在 KV Cache^[1] 中介绍到在 LLM 的推理阶段，其实用到的是最终 Embedding 矩阵的最后一行，而影响最后一行的最关键的计算是 Attention 的计算，在 Attention 中与其相关的是新增 Token 的 $Q$ 向量，完整 Input 的 $K$ 矩阵以及完整 Input 的 $V$ 矩阵，因此，我们可以将该新增 Token 之前所有 Token 对应的 $K,V$ 向量缓存起来，这便是 KV Cache 的含义，生成的过程可由下图表示：

这样每次生成新的 Token 时，可以按照上述的方案，但是 LLM 在生成第一个 Token 之前，需要处理一长串的输入，我们称这一长串的输入为 Prompt。在这里，还是以参考文献[1]中的例子，假设 Prompt 为“大模型的推理”，完整的序列为“大模型的推理过程也被称为自回归过程”。为了方便描述 LLM 的推理过程，我们简化模型，我们取上述图中的 N 为 1，同时，Token 为每一个中文的词，实际的过程中会根据不同的分词方法而不同。

在生成词的过程中，需要分为两步来看，首先，我们看首个 Token 的生成，首个 Token 的计算过程如下图所示：

如上，我们发现，Prompt 这部分的文本时可以直接都输入，构造出输入的 Embedding 矩阵，假设向量的维度为 6，那么 Embedding 矩阵的维度是 $6\times 6$，分别与 $W_Q,W_K,W_V$ 三个矩阵运算后便得到 Prompt 的 $Q,K,V$ 矩阵：

$$\begin{matrix} Q = E\times W_Q \\ K = E\times W_K \\ V = E\times W_V \end{matrix}$$

其中，$W_Q,W_K,W_V$ 的矩阵维度为 $6\times 8$，最终，$Q,K,V$ 的维度也是 $6\times 8$。根据 Masked Multi-Head Attention 的公式：

$$\mathrm{Maksed\_Attention}(Q,K,V)=\mathrm{softmax}(\frac{QK^\top}{\sqrt{d_k}}+M)V$$

其中，$M$ 为掩码矩阵。最终 Attention 输出一个 $6\times 8$ 的矩阵，再经过 FFN 后取 Embedding 矩阵的最后一行预测下一个词。

其次，我们看剩余的 Token 的生成，如参考文献[1] 中所描述的，根据 KV Cache，我们知道在生成首个 Token “过”之后，我们需要先缓存之前所有 Token 的 K 和 V 向量，包括“大”，“模”，“型”，“的”，“推”，“理”。此时，将新生成的 Token “过”，计算其 $Q,K,V$，便可以预测解析来的词，这一步与第一步有明显的不同，同时，在 KV Cache 中再缓存好 Token “过”的 K 和 V 向量。

我们将上面两步分别称为 Prefill 和 Decoding 两个阶段。由于 KV Cache 的存在，Decoding 阶段的计算量大幅降低，这两个阶段也因此呈现出截然不同的计算特点：

• Prefill 阶段：以计算密集型（compute-bound）为主。该阶段需要对整个输入序列执行完整的 Transformer 计算，涉及大量矩阵乘法和注意力操作，计算复杂度较高。
• Decoding 阶段：以内存密集型（memory-bound）为主。由于 KV Cache 的加入，Decoding 阶段只需对新增的 token 与缓存内容进行增量计算，计算量显著降低，但频繁的内存访问和更新使其更加依赖内存带宽。

以上的过程便称为 PD 分离。

2. Prefill 阶段

通过上述的分析，不难发现，在 Prefill 阶段，其主要任务是批量处理完整的输入序列，模型需要从第一层一直计算到最后一层，逐层计算每个 token 与其他所有 token 的注意力关系，最终将过程中生成的每个 Token 的 K 向量和 V 向量存储到 KV Cache 中。

在此阶段，算力要优先于内存的带宽，同时，计算复杂度也会随着序列的长度呈现出平方的增长，因为 LLM 的核心计算量来自于注意力机制的计算，其计算复杂度为 $O\left(n^2d\right)$，其中 $n$ 为序列长度，$d$ 为模型的维度。因此在 Prefill 阶段，重点需要优化长文本的计算。

针对 Prefill 阶段的优化具有重要的意义，通过上面的分析，我们发现，Prefill 阶段的结果是首个 Token，因此，对于 Prefill 阶段的优化，能够降低 LLM 的首 Token 延迟，也就是 TTFT（Time to First Token），当前已有的一些方法包括：

• Prefix Cache：对常用的前缀序列提前缓存其计算结果，避免重复计算，提高效率。
• Chunked Prefill：将长序列分块处理，优化内存占用和并行计算效率。

3. Decode 阶段

在 Decode 阶段，其主要任务是以自回归的方式逐个生成 Token，需要用到的是前一步生成的 Token 以及 Prefill 阶段缓存好的 K 向量和 V 向量。这个阶段，计算不再是瓶颈，每次计算新的 Token 时，注意力机制的计算复杂度为 $O\left(nd\right)$，计算量与序列长度时线性关系，然而，在此阶段，需要大量访问 KV Cache，以取得对应 Token 的 K 向量和 V 向量，因此对于内存带宽的要求要高于算力。

Decode 阶段的优化，主要有

• Flash Decoding：通过高效的内存布局和缓存管理策略，减少内存访问延迟并提高显存利用率。

4. 总结

KV Cache 的提出，将 LLM 的推理过程分成了两个阶段，分别是 Prefill 阶段和 Decode 阶段，其中，Prefill 阶段处理的是 LLM 的 Prompt，用于产出首个 Token 以及存储 Prompt 中所有 Token 的 K 向量和 V 向量，此阶段算力要优先于内存的带宽；Decode 阶段则是根据新生成的 Token 以及历史所有 Token 的 K 向量和 V 向量，用于生成新的 Token，这个阶段，算力相比较于对 KV Cache 的访问，对内存带宽的要求要高于算力。

参考文献

[1] LLM 推理优化技术：KV Cache