【大模型】大模型推理中的Prefill、Decode与KV Cache详解

最新推荐文章于 2025-07-31 17:06:18 发布

烟锁池塘柳0

最新推荐文章于 2025-07-31 17:06:18 发布

阅读量925

点赞数 25

CC 4.0 BY-SA版权

分类专栏：机器学习、深度学习与大模型文章标签：深度学习人工智能机器学习

本文链接：https://blog.csdn.net/Zlyzjiabjw547479/article/details/149498273

机器学习、深度学习与大模型专栏收录该内容

37 篇文章

订阅专栏

大模型推理中的Prefill、Decode与KV Cache详解

在大模型的推理（Inference）过程中，Prefill、Decode和KV Cache是三个密不可分的概念，本文将解释它们如何协同工作，以及为何它们对于优化LLM推理速度至关重要。

文章目录

引言：大模型推理的性能瓶颈

众所周知，大型语言模型通常基于Transformer架构。Transformer模型的核心是自注意力（Self-Attention）机制，它允许模型在处理序列数据时，动态地评估不同部分（Token）之间的重要性。然而，这种机制在带来卓越性能的同时，也引入了巨大的计算复杂性。

在自回归（Autoregressive）的生成任务中，模型需要逐字（Token）生成内容。每生成一个新的Token，都需要将其与之前所有的Token进行一次注意力计算。随着序列长度的增加，计算量会呈二次方增长，导致推理延迟急剧上升，这成为了LLM服务化（Serving）的主要性能瓶颈。为了解决这一难题，KV Cache应运而生，并由此衍生出Prefill和Decode两个 distinct 的推理阶段。

1 Attention机制与KV Cache的诞生

1.1 Transformer 中的自注意力机制

Transformer 中的自注意力机制

在自注意力计算中，输入序列的每个Token会生成三个向量：查询（Query, Q）、键（Key, K）和值（Value, V）。这三个向量是通过将输入Token的嵌入（Embedding）与三个不同的权重矩阵相乘得到的。

Q (Query): 代表当前Token，用于去“查询”其他Token。
K (Key): 代表被查询的Token，用于与Q进行匹配度计算。
V (Value): 代表被查询Token的实际内容。

注意力得分是通过计算Q和所有K的点积得到的，然后通过Softmax函数进行归一化，最后将这些得分与对应的V相乘并加权求和，得到最终的输出。

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

1.2 性能瓶颈与KV Cache的提出

（本部分参考了 AIInfra/05Infer/02InferSpeedUp/01KVCache.pptx 当中的内容，
也可以参见这个视频介绍：【大模型推理】大模型推理 Prefill 和 Decoder 阶段详解）

在自回归生成中，假设我们已经生成了 i-1 个Token，现在要生成第 i 个Token。此时，模型需要计算第 i 个Token的Q向量，与前 i 个所有Token的K和V向量进行注意力计算：

没有KV Cache

不难发现，在这个过程中，前 i-1 个Token的K和V向量是已经被计算过的。如果每次生成新Token都重新计算这些历史K和V向量，将会造成巨大的计算浪费。

KV Cache 机制的核心思想就是：将过去所有Token的Key和Value向量缓存起来，避免重复计算

具体来说，在生成第 i 个Token时，我们只需要计算当前Token的Q, K, V向量。然后，将新生成的K和V向量与之前缓存的KV Cache进行拼接，形成一个包含了到当前位置为止所有Token的K和V集合。这样，在进行注意力计算时，我们只需要使用当前Token的Q向量和更新后的完整KV Cache即可：

有KV Cache

通过KV Cache，自注意力机制的计算复杂度从 $O(n^2)$ 降低到了 $O (n)$ ，其中 n 是序列长度，极大地提升了推理效率。

2 大模型推理的两个阶段：Prefill 与 Decode

引入KV Cache后，LLM的生成过程被自然地分为了两个阶段：Prefill阶段和Decode阶段。

2.1 Prefill（预填充）阶段

Prefill阶段发生在处理用户输入的提示（Prompt）时。 在这个阶段，模型会并行地处理输入Prompt中的所有Token，并为它们计算出相应的K和V向量，然后将这些向量填充到KV Cache中。

核心特点：

并行计算：由于Prompt是已知的，模型可以一次性地、并行地对所有Token进行计算。这使得该阶段的计算非常密集，可以充分利用GPU的并行计算能力。
计算密集型（Compute-Bound）：Prefill阶段涉及大量的矩阵乘法运算，其瓶颈主要在于GPU的计算能力。
生成第一个Token：此阶段的最终输出是生成回复的第一个Token，同时填充好初始的KV Cache。

举个例子，当用户输入“中国的首都是哪里？”时，模型会将这句话分词（Tokenize），然后并行处理“中国”、“的”、“首都”、“是”、“哪里”、“？”这几个Token，计算出它们各自的K和V向量，并将这些向量存入KV Cache。这个过程的耗时，直接影响了用户感知的“首字延迟”（Time to First Token, TTFT）。

2.2 Decode（解码）阶段

Decode阶段发生在逐个生成输出Token时。 当Prefill阶段完成后，模型进入自回归（Autoregressive）的生成循环。在每一步中，模型只处理上一步新生成的那个Token。

核心特点：

串行计算：每次只处理一个Token，其计算依赖于上一步的结果，因此是串行的、自回归的。
内存带宽密集型（Memory-Bandwidth-Bound）：在Decode的每一步，模型都需要从GPU内存中读取庞大的模型权重和完整的KV Cache。计算量本身不大（主要是一个向量和矩阵的乘法），但数据传输量巨大。因此，其瓶颈在于GPU显存的带宽。
利用并更新KV Cache：每生成一个新Token，都会计算出其K和V向量，并将其追加到KV Cache中，供下一个Token的生成使用。这个过程的耗时，影响了“字间延迟”（Time Between Tokens），即生成后续内容的速度。

继续上面的例子，在生成第一个Token“北”之后，Decode阶段开始。模型会计算“北”的K和V，并将其加入到之前由Prompt生成的KV Cache中。然后利用“北”的Q和更新后的KV Cache来生成下一个Token“京”。如此循环，直到生成完整的答案“北京。”。

3 Prefill、Decode与KV Cache的协同工作流程

下面我们通过一个流程图来总结这三者之间的协同关系：

4 优化与挑战

理解了Prefill、Decode和KV Cache的工作原理后，我们就能更好地理解当前LLM推理优化的方向。

KV Cache的内存占用：KV Cache虽然极大地提升了速度，但它也占用了大量的GPU显存。其大小与batch_size * sequence_length * num_layers * hidden_size成正比。对于长序列或大批量请求，KV Cache可能会耗尽显存。因此，KV Cache量化（Quantization）、稀疏化等技术被提出来减小其内存占用。
**Prefill与Decode的资源矛盾 **：Prefill是计算密集型，需要高计算能力的GPU；而Decode是内存带宽密集型，需要高显存带宽的GPU。将两者放在同一硬件上运行，可能会导致资源利用率不均。因此，业界提出 Prefill-Decode 分离（Disaggregation）的架构，使用不同硬件集群分别处理Prefill和Decode任务，以实现更优的资源配置和吞吐量。