Prefill v.s. Decode 流程
Prefill 阶段: LLM 推理的初始化过程, Prefill 完成的是对用户输入的完整 prompt 处理, 生成出第一个输出 token 并且构建 KV Cache. 这个阶段一次性处理所有的输入 token, 完成从 encoding 到 隐藏层状态的全量计算; 简而言之: 初始化 KV Cache
Decode 阶段: focus 逐个 token 的生成过程, 完成 Prefill 后, 基于已经生成的 KV Cache, 以自回归的方式逐个预测 token, 并在这个过程中增量更新 KV Cache
Prefill v.s. Decode 计算流程详细对比
1.Prefill 执行的是全序列并行计算, Decode 执行的是串行计算
2.Prefill 处理所有输入 token 的全局注意力, Decode 仅关注当前 token 与前序 Cache 的局部注意力
3.Prefill v.s. Decode
Prefill 的阶段流程
# 伪代码
prompt_embeddings = tokenize_and_embed("Explain quantum mechanics")
hidden_states = transformer_layers(prompt_embeddings) # 全量计算
kv_cache = generate_kv_cache(hidden_states) # 构建初始Cache
first_token = lm_head(hidden_states[:, -1, :]) # 生成首token
Decode 的阶段流程
current_token = first_token # 拿到 current_token
for _ in range(max_length):
token_embedding = embed(current_token)
hidden_state = transformer_step(token_embedding, kv_cache) # 增量计算
next_token = lm_head(hidden_state)
update_kv_cache(kv_cache, hidden_state) # Cache扩展
4.Prefill v.s. Decode 计算复杂度对比
假设 $n$ 为序列长度, $d$ 为隐藏层长度:
(i). Prefill 阶段的计算复杂度为 $O(n^2d)$, 也就是说计算复杂度和输入长度的平方成正比
(ii). Decode 阶段的计算复杂度为 $O(nd)$, 每生成一个 token 都需要执行固定量的计算, 计算复杂度因为受到自回归特性的影响, 总计算量和输出长度呈现线性关系
Prefill v.s. Decode 性能指标解析
Prefill 阶段的性能指标是 (Time to First Token, TTFT) 也就是接受请求到产出首个 Token 的时间
Decode 阶段的性能指标是 (Time Per Output Token, TPOT) 也就是生成单个 token 的平均时间
影响 Prefill 阶段 TTFT 的因素为
1.输入长度与计算复杂度
2.批处理大小和并行度
3.硬件的计算吞吐量
影响 Decode 阶段 TPOT 的因素为
1.KV Cache 的访问效率
2.自回归依赖导致的串行化
3.显存带宽利用率
Prefill v.s. Decode 优化策略
Prefill 阶段优化:
1.动态批处理: 合并多个请求的 prefill 计算, 利用矩阵乘法的并行性提升吞吐, 假设有 1 个 512 tokens 的 prefill 请求, 现在这种请求有 4 个的话, 可以合并为 2048 x d 的矩阵运算
2.选择性缓存: 对 prompt 中的固定模板的 KV Cache 进行预计算, 比如 system prompt 这个固定的部分
3.混合精度计算: 在 attention 计算中使用 FP8 格式, 配合 Hopper 架构的 Transformer Engine, 可将 Prefill 时间减少 40 %
Decode 阶段优化:
1.连续缓存管理: 采用 PagedAttetion 等技术, 将 KV Cache 划分为固定大小的内存页, 比如 256 tokens/页, 减少内存碎片
2.推测解码 Speculative Decoding: 使用小模型生成候选序列, 大模型仅进行验证
3.内存带宽优化: 通过 KV Cache 量化 (FT16 -> INT8) 和内存布局优化 (Channel-Last 格式), 使得显存带宽利用率提升 60%
转载请注明来源 goldandrabbit.github.io