原理概述

在AI大模型推理场景中，系统需在有限硬件资源下实现低延迟、高吞吐的文本生成。传统单阶段处理方式会导致显存占用随输入长度指数级增长，计算资源无法有效复用。工程实践中通过将推理过程拆分为Prefill（预填充）和Decode（解码）两个阶段，实现计算与存储资源的解耦优化。这种阶段划分已成为主流技术框架的核心设计，本文将系统解析其底层运行机制与优化原理。

背景问题

当处理长文本输入或连续对话场景时，传统单阶段推理面临三大挑战：

1. 显存爆炸：自注意力机制需存储所有token的键值对（KV Cache），输入长度每增加1倍，显存占用增长4倍（平方关系）
2. 计算冗余：每次生成新token都需重新计算全部历史token的注意力权重
3. 延迟累积：解码阶段无法并行处理，每个token生成需等待前序计算完成

某行业常见技术方案通过阶段拆分，将输入处理与输出生成解耦，使系统能够针对不同阶段特性实施专项优化。

核心概念

理解该原理需掌握三个基础概念：

1. KV Cache：存储自注意力计算中间结果的缓存结构，包含Key和Value矩阵
2. 注意力掩码：控制模型可见范围的矩阵，解码阶段使用因果掩码确保自回归特性
3. 并行维度：批处理（Batch）与序列长度（Sequence Length）的维度划分策略

系统组成

优化系统包含四大核心模块：

1. 输入处理器：负责文本分词、嵌入编码和初始KV Cache生成
2. 计算调度器：动态分配Prefill/Decode阶段的计算资源
3. 缓存管理器：维护KV Cache的生命周期与存储策略
4. 输出生成器：实现token采样策略与生成结果后处理

工作流程

以”我是谁？”的推理过程为例，完整流程分为六个步骤：

1. 输入分词
   将原始文本拆解为token序列：["我", "是", "谁", "？"]

2. Prefill阶段

   # 伪代码示例
   def prefill(tokens):
       kv_cache = []
       for token in tokens:
           # 并行计算所有token的注意力
           attention_output = multi_head_attention(token, kv_cache)
           # 更新KV Cache（包含所有历史token）
           kv_cache.append(compute_kv(token, attention_output))
       return kv_cache

   该阶段完成三件事：

   • 生成初始嵌入向量
   • 计算所有token间的注意力权重
   • 构建完整的KV Cache矩阵

3. 缓存优化
   实施两项关键优化：

   • 分块存储：将KV Cache按注意力头（Head）维度拆分
   • 压缩算法：采用量化技术将FP32数据压缩至INT8

4. Decode阶段

   def decode(kv_cache, max_length):
       output_tokens = []
       for _ in range(max_length):
           # 仅计算最后一个token的注意力
           logits = forward_pass(kv_cache[-1:])
           # 采样新token
           new_token = sample(logits)
           # 更新KV Cache
           kv_cache.append(compute_kv(new_token))
           output_tokens.append(new_token)
       return output_tokens

   该阶段特点：

   • 每次仅处理1个新token
   • 复用Prefill阶段的KV Cache
   • 支持温度采样等生成策略

5. 并行控制
   实现两种并行模式：

   • 批并行：同时处理多个独立请求
   • 流水线并行：将不同层分配到不同设备

6. 结果后处理
   包括：

   • 特殊token过滤
   • 重复词抑制
   • 语法校验

关键机制

显存优化机制

1. KV Cache分页：将大矩阵拆分为多个固定大小页面，按需加载
2. 梯度检查点：在Decode阶段选择性释放中间计算结果
3. 显存池化：建立跨请求的显存共享池

计算优化机制

1. 注意力模式切换：
   • Prefill阶段使用全注意力
   • Decode阶段切换为滑动窗口注意力
2. CUDA核融合：将多个小操作合并为单个CUDA核调用
3. 张量并行：将矩阵运算拆分到多个GPU

延迟隐藏机制

1. 异步KV Cache更新：在生成当前token时预计算下一个token的KV对
2. 推测解码：并行生成多个候选序列，通过验证器筛选最优结果
3. 流式输出：边计算边返回结果，减少用户感知延迟

示例说明

以处理1024 token输入、生成256 token输出的场景为例：

阶段	计算量	显存占用	并行策略
Prefill	O(n²)	O(n)	批处理并行
Decode	O(n)	O(1)	流水线并行

优化效果：

• 显存占用降低60%（从48GB→19GB）
• 首token延迟从320ms降至120ms
• 持续生成吞吐量提升3倍

技术优势与限制

优势

1. 资源解耦：计算资源与显存资源可独立扩展
2. 弹性适配：支持从手机到数据中心的多样化部署
3. 动态调整：可根据负载情况动态分配资源

限制

1. 冷启动延迟：Prefill阶段存在不可压缩的最小延迟
2. 上下文限制：Decode阶段仍受KV Cache总大小限制
3. 工程复杂度：需维护两套不同的计算逻辑

常见误区

1. 误解阶段划分目的：并非单纯为了性能优化，更重要的是解决显存爆炸问题
2. 混淆并行维度：批并行与流水线并行解决不同层面的问题
3. 忽视缓存一致性：Decode阶段修改KV Cache需同步到所有副本

总结

Prefill与Decode的阶段划分本质是计算-存储解耦设计，通过将输入处理与输出生成分离，实现了：

1. 显存占用的线性增长控制
2. 计算资源的按需分配
3. 生成质量的可控调节

这种设计已成为现代AI推理引擎的基础架构，在百度智能云等平台的模型服务中广泛应用。理解其原理有助于开发者更好地进行性能调优、资源规划和故障排查，特别是在处理长文本、实时交互等复杂场景时具有重要指导意义。