一、 推理过程:Prefill+Decode
大型语言模型(LLM)的推理过程通常分为两个核心阶段:Prefill(预填充)阶段和Decode(解码)阶段。这两个阶段在计算特性、资源需求和优化策略上存在显著差异,共同构成了LLM从输入处理到文本生成的完整流程。以下是详细解析:
1. Prefill(预填充)阶段
目标:处理输入的完整提示(prompt),为后续生成准备上下文和中间状态。
核心任务:
- Token化与嵌入:将输入文本转换为token序列,并通过嵌入层映射为高维向量(形状为
[batch_size, seq_len, hidden_dim]
)。 - 全量注意力计算:对输入的所有token执行并行自注意力计算,生成每个token与上下文的关联性(注意力矩阵)。输入是完整的上下文,因此模型会执行一次完整的 masked self-attention.
- KV Cache初始化:缓存每一层Transformer的Key(K)和Value(V)向量,避免后续重复计算(会在 decode 阶段被反复使用)。
计算特性:
- 高度并行:输入token已知,可一次性计算所有位置的注意力(矩阵-矩阵运算)。
- 计算密集型:复杂度为O(n²)(n为输入长度),GPU算力利用率高,但长提示处理延迟显著。
- 显存占用:KV Cache的存储需求随输入长度平方增长,可能成为瓶颈。
优化技术:
- FlashAttention:优化注意力计算的内存访问模式,减少显存开销。
- 分块处理:将长输入分段处理,降低单次计算压力。
2. Decode(解码)阶段
目标:基于Prefill阶段的KV Cache,自回归生成输出token,直至达到终止条件(如EOS标记或最大长度)。
核心任务:
- 增量生成:每次仅处理最新生成的1个token,通过查询(Q)与历史KV Cache计算注意力。
- KV Cache更新:将当前token的K/V向量追加到缓存中,供下一步使用。
- 输出选择:通过解码策略(如贪心、采样或集束搜索)从概率分布中选择下一个token。
计算特性:
- 严格串行:每步依赖前序所有token,无法并行(矩阵-向量运算)。
- 内存带宽限制:需频繁加载历史KV Cache,计算量小但内存访问压力大。
- 显存增长:KV Cache随生成序列线性扩展,长文本生成可能耗尽显存。
优化技术:
- KV Cache压缩:分页存储(PagedAttention)、低精度量化(INT8)。
- 推测解码:用小模型预生成草稿,大模型并行验证,加速2-3倍。
- 动态批处理:合并多个请求的生成步骤,提高GPU利用率。
3. 阶段对比与协同
特性 | Prefill阶段 | Decode阶段 |
---|---|---|
输入形式 | 完整prompt(已知所有token) | 逐token生成(每次1个) |
并行性 | 高度并行 | 严格串行 |
计算瓶颈 | 算力限制(Compute-bound) | 内存带宽限制(Memory-bound),每生成一个 token,都需要访问所有历史的 KV Cache(多层、多头) |
主要优化目标 | 减少长提示处理延迟 | 提高生成吞吐量、降低显存占用 |
协同关系:Prefill为Decode提供初始KV Cache,两者通过缓存复用实现高效推理。例如,Prefill生成的KV Cache使Decode阶段避免重复计算prompt的注意力,将总复杂度从O(n³)降至O(n²)。
4. 其他可选阶段
部分框架会扩展以下阶段:
- 后处理(Post-Processing):对生成结果进行过滤、重排序或格式化。
- 缓存管理:动态释放或压缩历史KV Cache,适应长对话场景。
总结
LLM推理的两阶段划分源于输入与输出的计算特性差异:Prefill充分利用并行性处理已知输入,Decode则通过串行和缓存复用实现高效生成。这一设计平衡了计算效率与资源消耗,是LLM落地应用的核心优化基础。
二、Prefill阶段解析
LLM(大语言模型)在 Prefill(预填充)阶段 主要完成对输入提示(prompt)的预处理和上下文理解,为后续的生成(Decode)阶段奠定基础。以下是该阶段的核心任务和实现细节:
1. 输入处理与Token化
- 文本转Token序列:将用户输入的文本(如“请解释Transformer的原理”)通过分词器(Tokenizer)转换为模型可理解的Token ID序列。
- 嵌入映射:每个Token通过嵌入层(Embedding Layer)转换为高维向量(如768维或1024维),形成形状为
[batch_size, seq_len, hidden_dim]
的输入矩阵。
2. 全量注意力计算与KV Cache生成
- 并行计算:Prefill阶段的核心优势是能够并行处理所有输入Token。模型通过自注意力机制(Self-Attention)计算每个Token与之前所有Token的关联性(包括自身),生成注意力分数矩阵。
- 数学过程:对每个位置的Token,计算其Query(Q)与所有Key(K)的点积,经Softmax归一化后加权求和Value(V)向量。
- KV Cache初始化:每一层Transformer计算的Key和Value向量被缓存,形成KV Cache。这是后续Decode阶段复用中间结果、避免重复计算的关键。