第二章：推理加速核心：预填充（Prefill）与解码（Decode）的深度解析与实现

在上一章中，我们了解了 LLM 推理的基本流程。本章将深入探讨推理过程中的两个核心阶段：预填充（Prefill）和解码（Decode），它们是影响推理性能的关键因素。理解这两个阶段的原理和优化方法对于提升 LLM 推理效率至关重要。

2.1 预填充（Prefill）阶段：原理与优化

预填充（Prefill）阶段发生在模型接收到完整输入 Prompt 之后，但在开始生成第一个输出 token 之前。这个阶段的主要任务是处理输入的 Prompt，计算出所有输入 token 的上下文表示，并初始化后续解码阶段所需的关键数据结构——KV 缓存（Key-Value Cache）。

计算流程的数学原理与图示说明：

假设我们有一个批次大小为 batch_size 的输入 Prompt，每个 Prompt 的长度为 seq_len，模型的隐藏层维度为 hidden_dim。输入 Embedding 层的输出形状为 (batch_size, seq_len, hidden_dim)。

对于 Transformer 模型的每一层，自注意力机制是核心。对于输入 $H_{in}$（形状为 (batch_size, seq_len, hidden_dim)），我们通过线性变换得到 Query (Q)、Key (K) 和 Value (V) 矩阵。这些权重矩阵 $W_Q,W_K,W_V$ 的形状为 (hidden_dim, num_heads * head_dim)，其中 num_heads 是注意力头的数量，head_dim 是每个注意力头的维度 (hidden_dim = num_heads * head_dim)。

因此，Q、K、V 的形状分别为 (batch_size, seq_len, num_heads * head_dim)。通常会将 num_heads 维度分离出来，得到形状为 (batch_size, num_heads, seq_len, head_dim) 的 Q、K、V。

接下来，计算注意力分数。对于每个注意力头，我们将 Query 和 Key 的转置相乘：

$AttentionScore=Q\cdot K^T$

这里的矩阵乘法发生在形状为 (batch_size, num_heads, seq_len, head_dim) 的 Q 和形状为 (batch_size, num_heads, head_dim, seq_len) 的 $K^T$ 之间，得到形状为 (batch_size, num_heads, seq_len, seq_len) 的注意力分数矩阵。

然后，对注意力分数进行缩放（除以 $\sqrt{head\_dim}$）和 Softmax 归一化，得到注意力权重矩阵，形状仍然是 (batch_size, num_heads, seq_len, seq_len)。

最后，将注意力权重与 Value 矩阵相乘，得到自注意力层的输出：

$AttentionOutput=AttentionWeight\cdot V$

这里的矩阵乘法发生在形状为 (batch_size, num_heads, seq_len, seq_len) 的注意力权重和形状为 (batch_size, num_heads, seq_len, head_dim) 的 V 之间，得到形状为 (batch_size, num_heads, seq_len, head_dim) 的输出。这个输出会被重新组合成形状为 (batch_size, seq_len, hidden_dim) 并传递给下一层。

计算复杂度分析：

自注意力机制在预填充阶段的时间复杂度主要由注意力分数矩阵的计算 ($Q\cdot K^T$) 和加权求和 ($AttentionWeight\cdot V$) 决定。这两个矩阵乘法的复杂度都是 $O(batch\_size\cdot num\_heads\cdot seq\_le{n}^2\cdot head\_dim)$。由于 $hidden\_dim=num\_heads\cdot head\_dim$，所以整体复杂度可以表示为 $O(batch\_size\cdot seq\_le{n}^2\cdot hidden\_dim)$。对于长序列的输入，这个平方项会使得计算量显著增加，成为预填充阶段的性能瓶颈之一。在后续章节中，我们将探讨一些优化的注意力机制，例如其复杂度可以降低到 $O(seq\_len\cdot hidden\_dim\cdot log(seq\_len))$ 甚至 $O(seq\_len\cdot hidden\_dim)$。

KV 缓存的初始化过程与代码实现：

在预填充阶段，对于 Transformer 模型的每一层，模型都会计算得到 Key (K) 和 Value (V) 矩阵。这些 K 和 V 矩阵会被缓存起来，这就是所谓的 KV 缓存。对于每个输入 token 和每一层，都会存储对应的 Key 和 Value 向量。KV 缓存的形状通常是 (batch_size, num_heads, seq_len, head_dim)。

在 vLLM 中，当您调用 llm.generate(prompt) 时，预填充阶段会自动进行 KV 缓存的初始化。vLLM 会在处理输入 Prompt 的过程中，计算每一层的 Key 和 Value 向量，并将它们存储在 GPU 内存中以供后续解码步骤使用。

from vllm import LLM

model_name = "Qwen/Qwen2.5-7B"
llm = LLM(model=model_name)

prompt = "The capital of France is "
# 当调用 generate 时，vLLM 会对 prompt 进行预填充，并初始化 KV 缓存
outputs = llm.generate(prompt, max_tokens=5)

for output in outputs:
    generated_text = output.outputs[0].text
    print(f"Generated text: {generated_text}")

在上述代码中，当 llm.generate(prompt, max_tokens=5) 被调用时，vLLM 会首先对 “The capital of France is " 这个 Prompt 进行预填充。在此过程中，对于 Prompt 中的每个 token（“The”, " capital”, " of", " France", " is"），vLLM 会在模型的每一层计算其对应的 Key 和 Value 向量，并将这些向量存储到 KV 缓存中。KV 缓存的结构可以想象成一个多维数组，它按层、按注意力头、按 token 位置存储了 Key 和 Value 信息。

预填充阶段的并行性与性能优化：

预填充阶段的一个关键优势在于其高度的并行性。由于整个输入 Prompt 在开始时是已知的，模型可以同时计算所有 token 在每一层的表示。这使得预填充阶段能够充分利用 GPU 的并行计算能力，显著提高处理速度。

为了进一步优化预填充阶段的性能，主流的技术包括：

• 高效注意力机制（Efficient Attention Mechanisms）：
  • FlashAttention： 这是一种通过重新组织注意力计算过程，减少 GPU 内存读写次数，从而显著加速计算的高效注意力机制。它尤其在处理长序列时表现出色，可以将自注意力的复杂度从 $O(seq\_le{n}^2)$ 降低到更接近 $O(seq\_len)$。vLLM 等现代推理框架通常会集成或利用这些高效的注意力实现，以提升预填充速度。
• 批处理（Batching）：
  • 在实际应用中，通常会同时处理多个独立的推理请求。预填充阶段可以将这些请求的 Prompt 组成一个批次进行处理。通过批处理，可以更有效地利用 GPU 的计算资源，提高整体的吞吐量。vLLM 本身就支持高效的批处理，能够同时处理多个 Prompt 的预填充过程。

2.2 解码（Decode）阶段：原理与实现

解码（Decode）阶段在预填充阶段完成之后开始。在这个阶段，模型以自回归的方式逐个生成输出 token。每生成一个 token，该 token 就会被添加到已生成的序列中，并作为下一步生成的输入。

自回归生成机制详解与可视化：

解码阶段从预填充阶段处理的输入 Prompt 的最后一个 token 开始，目标是生成后续的输出序列。假设预填充阶段处理了 $n$ 个输入 token，解码阶段的目标是生成接下来的 $m$ 个输出 token $(y_1,y_2,...,y_m)$。

在每一步 $t$（从 1 到 $m$），模型会基于已经生成的序列 $(x_1,...,x_n,y_1,...,y_{t-1})$ 来预测下一个 token $y_t$。对于 Decoder-only 模型，在解码的每一步，通常只将上一步生成的 token 作为当前 Transformer 层的输入（除了第一步，输入是预填充的最后一个 token）。然而，模型内部的自注意力机制仍然可以访问包括原始 Prompt 和所有已生成 token 在内的完整序列的信息，这是通过 KV 缓存实现的。

KV 缓存的查询和更新机制（代码示例）：

KV 缓存在解码阶段是加速的关键。在预填充阶段，我们已经为输入 Prompt 中的所有 token 计算了 Key 和 Value 向量并存储在 KV 缓存中。在解码的每一步，假设模型生成了一个新的 token $y_t$。为了预测下一个 token $y_{t+1}$，模型需要计算 $y_t$ 的 Key 和 Value 向量。然后，这个新的 Key 和 Value 向量会被追加到 KV 缓存中，扩展缓存的长度。

当模型在某一步需要计算自注意力时，对于当前要预测的 token $y_{t+1}$（需要计算其 Query 向量），它会与 KV 缓存中所有历史的 Key 向量（包括来自原始 Prompt 和之前已生成的 token）进行比较，计算注意力权重。然后，使用这些权重对 KV 缓存中对应的 Value 向量进行加权求和，得到上下文信息。

以下代码示例展示了使用 vLLM 进行解码的过程。在 llm.generate 的调用中，KV 缓存的查询和更新是自动处理的：

from vllm import LLM

model_name = "Qwen/Qwen2.5-7B"
llm = LLM(model=model_name)

prompt = "The weather today is "
# 指定生成最多 20 个 token
outputs = llm.generate(prompt, max_tokens=20)

for output in outputs:
    generated_text = output.outputs[0].text
    print(f"Generated text: {generated_text}")

在上述代码中，vLLM 在预填充 "The weather today is " 之后，KV 缓存中已经包含了这 5 个 token 的 Key 和 Value 信息。在接下来的解码过程中，每生成一个新的 token（例如 “sunny”），vLLM 会计算 “sunny” 的 Key 和 Value 向量，并将它们添加到 KV 缓存的末尾。当模型预测下一个 token时，其 Query 向量会与 KV 缓存中所有 6 个 Key 向量进行注意力计算。这个过程会重复进行，直到生成 20 个新的 token 或达到结束条件。

单步解码的实现与性能分析：

在解码阶段的每一步，模型主要进行以下操作：

1. 接收上一步生成的 token 的 Embedding。
2. 计算该 token 在所有 Transformer 层的 Query、Key 和 Value 向量。
3. 将当前生成 token 的 Key 和 Value 向量更新到 KV 缓存中。
4. 在自注意力计算中，当前 token 的 Query 向量会与 KV 缓存中所有历史 token 的 Key 向量进行比较，计算注意力权重。Value 向量会根据这些权重进行加权求和，得到上下文向量。
5. 模型最后一层的输出会经过线性层和 Softmax 函数，得到下一个 token 的概率分布。
6. 根据解码策略（例如采样）从概率分布中选择下一个 token。

单步解码的计算成本主要在于自注意力机制的计算，其复杂度与当前 KV 缓存的长度（等于原始 Prompt 长度加上已生成的 token 数量）成正比。KV 缓存的关键作用在于，它避免了在每一步都重新计算原始 Prompt 的 Key 和 Value 向量。然而，随着生成序列的长度增加，KV 缓存的大小也会增长，可能导致内存带宽成为瓶颈。

2.3 预填充与解码的协同工作与性能瓶颈

数据流分析与性能瓶颈识别：

预填充和解码阶段共同完成了 LLM 的推理过程。预填充为解码准备了初始的 KV 缓存，而解码阶段则迭代地利用和更新这个缓存来生成最终的输出。

• 数据流： 输入 Prompt -> Tokenization -> Embedding -> 预填充 (生成 KV 缓存) -> 解码 (利用并更新 KV 缓存，逐个生成 token) -> Detokenization -> 输出文本。

• 性能瓶颈识别：

  • 预填充阶段： 对于极长的输入 Prompt，自注意力机制的计算量仍然很大，可能成为计算瓶颈。同时，加载模型权重和初始 KV 缓存的内存开销也需要考虑。
  • 解码阶段： 解码的串行自回归特性是主要的瓶颈。虽然 KV 缓存减少了重复计算，但每一步仍然需要进行注意力计算，并且随着生成序列的增长，KV 缓存的大小也会增加，可能导致内存带宽瓶颈。此外，生成长序列会显著增加总的推理时间。

实操：对比分析不同 Prompt 长度下预填充和解码的耗时：

import time
from vllm import LLM
from transformers import AutoTokenizer

model_name = "Qwen/Qwen2.5-7B"
llm = LLM(model=model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
output_tokens = 100  # 固定生成 100 个 token
prompt_lengths = [10, 50, 100, 200]

print("Analyzing inference time with varying prompt lengths:")
results = {}
for prompt_len in prompt_lengths:
    prompt = "This is a test prompt. " * prompt_len
    input_ids = tokenizer.encode(prompt)
    prompt_token_count = len(input_ids)
    start_time = time.time()
    outputs = llm.generate(prompt, max_tokens=output_tokens)
    end_time = time.time()
    total_time = end_time - start_time
    results[prompt_token_count] = total_time
    print(f"Prompt length (tokens): {prompt_token_count}, Total inference time: {total_time:.4f} seconds")

print("\nResults:")
for length, time in results.items():
    print(f"Prompt Length: {length}, Inference Time: {time:.4f} seconds")

# 分析：观察不同 Prompt 长度下总推理时间的变化。
# 当 Prompt 长度增加时，预填充阶段的计算量会增加，导致总推理时间变长。
# 而解码阶段生成的 token 数量固定，其耗时相对稳定。
# 可以尝试使用 matplotlib 等库将这些结果绘制成折线图，横轴是 Prompt 长度，纵轴是推理时间，
# 这样可以更直观地看到 Prompt 长度对推理性能的影响。
# 随着 Prompt 长度的增加，总的推理时间应该会呈现上升趋势，这主要是由于预填充阶段的计算量增加所致。

补充：动态 Batching 的初步介绍：原理与在预填充和解码阶段的应用

动态 Batching 是一种提高 LLM 推理吞吐量的关键技术，尤其是在高并发的场景下。

• 原理： 动态 Batching 允许推理系统在运行时动态地将多个独立的推理请求组合成一个批次进行处理。这个批次的组成可以根据请求的当前处理阶段（预填充或解码）、序列长度以及其他策略进行调整，以最大化 GPU 的利用率。

• 在预填充阶段的应用： 当多个新的推理请求到达时，可以将它们的 Prompt 组成一个批次，一起进行预填充计算。这可以显著提高 GPU 的并行计算效率。

• 在解码阶段的应用： 动态 Batching 在解码阶段更为复杂但也更为关键。例如，vLLM 采用了名为 PagedAttention 的技术，它允许在解码过程中更灵活地管理 KV 缓存。不同请求的 KV 缓存可以以非连续的方式存储在 GPU 内存中，从而避免了因批处理中不同长度的序列而产生的内存浪费。当有新的请求加入或已有请求完成时，系统可以动态地分配和释放内存页，实现更高效的批处理和更高的吞吐量。

通过动态 Batching，推理系统可以在保证较低延迟的同时，显著提高单位时间内处理的请求数量。我们将在后续章节中更深入地探讨这些高级优化技术，包括 PagedAttention 的具体实现原理。

内容同步在我的微信公众号 ：智语Bot