Prefill and Decode

Prefill 与 Decode 两阶段及 TTFT/ITL 指标示意图

1. 只要三条结论,为什么还得等?

还有十分钟开会。你把一份几十页的项目文档发给 AI:“帮我挑出三条最值得提醒老板的风险,简短一点。”

屏幕安静了一会儿。你忍不住想:只要三条结论,怎么还没开始?

过了一会儿,回答开始出现,三条结论很快就写完了。

换一天,你只发了一句"帮我写一篇两千字的演讲稿",它倒是很快开口,却要写上好一阵。

“多久开始回答"和"之后写得多快”,背后有两段不同的计算:

prefill 处理已经给定的输入,decode 根据已有内容继续生成。

先看图里的第一个输出。

请求发出到收到它的时间,叫 TTFT(time to first token);之后相邻输出 token 的间隔,叫 ITL(inter-token latency)。

两者都能影响你觉得 AI 快不快。

理解 prefill 和 decode,能帮助我们把问题问得更具体。

输入经过 prefill,保存 prompt 的 KV,再选出第一个输出,后续输出接着出现。这里只概览两个阶段;TTFT 还可能包含排队、输入处理和传输。箭头与播放时长不代表实测耗时。

2. Prefill:要先处理几十页输入

那份文档里,前面写着"周五上线",后面写着"支付接口还没联调"。

你想得到的风险判断,需要利用这些输入信息。

Prefill 就发生在这里:模型把 prompt 的 token 转成 representation,逐层执行 attention 和前馈网络,同时保存各层的 K/V,供后续生成使用。

要求回答简短,并不会自动省掉对长文档的输入处理。

不过,处理文档无须像生成回答那样,一个 token 接着一个 token 等。

把例子缩到四个输入 token:它们已经全在 prompt 里,内容都已确定。

进入某一层后,可以把四个位置的 representation 排成矩阵,批量计算 Q、K、V。

你可能会问:

第二个位置要用到第一个位置的信息,怎么还能并行?

看下面两行。

输入 h₁ 产生 Q₁、K₁、V₁;输入 h₂ 产生 Q₂、K₂、V₂。

第一个位置读取自己的 K/V,第二个位置用 Q₂ 读取前两个位置的 K/V。

关键在那条斜线:第二个位置需要 K₁/V₁,它们从这一层的输入就能算出来,无须等待第一个位置的 attention 输出 O₁。 因此,所需 Q/K/V 准备好后,两行 attention 可以批量、并行计算。

同一层、一个 attention head 的两位置示意。斜线标出 K₁/V₁ 对第二个位置的贡献;O₁、O₂ 是各位置的 attention 输出,O₂ 的计算不依赖 O₁。省略归一化、位置处理、多头合并等细节。

"并行"不要求 GPU 在同一个瞬间完成所有操作。

Causal mask 仍限定每个位置只能读取自己和前文的 K/V;读取范围有先后边界,计算输出无需按位置排队。

层与层之间则仍有依赖:这一层经过 attention、残差连接和前馈网络等操作得到的输出,成为下一层的输入,再计算新的 Q/K/V。

回到文档,"支付接口还没联调"对应的后方位置可以参考前面的上线安排;前方位置无法读取后来才出现的内容。

走完最后一层,最后一个输入位置的 representation 用于预测第一个新 token。

经过采样或其他选择规则,回答终于开始出现。

留意下面的变化:第一个新 token 刚出现时,cache 里仍只有输入位置的 KV。

新 token 自己的 KV,要等它被送回模型才会产生。

四个输入位置的 causal attention 与首 token 边界。x1–x4 是示意 token,y1 是第一个新 token;图中 KV 行代表各层各自保存的状态。

3. Decode:刚写出来的内容,为什么又要送回去?

假设回答刚生成一个 token “领导”,接下来生成冒号。

"领导"会变成下一步的输入:模型处理它,读取前文的 KV,计算并保存它在各层的新 KV,再预测"冒号"可以成为下一步的输入,如此继续,直到停止。

图里把这两个新 token 简写成 y1、y2。

沿着外侧回路看:上一步选出的输出,就是下一步送入模型的输入。

用 y1 指代示意 token “领导”,后续 y2、y3 同样是示意符号。动画展开两个连续步骤:新 token 送回模型,两层各自追加 KV,旧条目留在原位,再选出下一个 token。

不妨数一下:如果原来有四个输入 token,选出 y1 时,cache 仍有四个位置;把 y1 送回模型处理后,才增加到五个,并得到预测 y2 所需的输出。

那份几十页的文档呢?

它的有效 KV 留在原处。

Causal attention 下,追加后文不会反过来改变前文的 representation,因此无须每写一个新 token,就把整份文档从头再算一遍。

新 query 仍要读取相关历史 K/V,这部分工作还在。

为什么不把整段回答一起算出来?

因为 prompt 已经全部给定,回答后面的 token 还没确定。

常规自回归 decode 要先选出当前 token,后一步才拿得到自己的新输入。

4. 读长报告,和写长演讲稿,慢在同一个地方吗?

现在回到开头的两个任务:

A:“读完这份长报告,只给我三条结论。” 输入多、输出少,要重点关注 prefill。

B:“请写一篇两千字的演讲稿。” 输入短、输出长,反复进行的 decode 可能占据大部分等待时间。

左侧是长报告换三条结论,右侧是短提示换长演讲稿。图形只对比输入与输出的多少,不代表实测 token 数、耗时或速度。

两个阶段使用同一个 decoder-only Transformer、同一套参数,硬件面对的工作却不同。

Prefill 一次带来许多已知位置。

同一份权重可以参与多行输入的矩阵计算,通常更有机会复用已经读入的权重。

单请求的普通 decode,每步只带来一行新输入,却仍要执行模型、读取权重和相关 KV。

于是就有一个常见倾向:

长输入的 prefill 往往更受计算能力影响,小 batch 的 decode 往往更受内存带宽影响。

后者的瓶颈常出在把数据送到计算单元的速度上。

这条经验要结合任务看。

增大 batch、改变模型或 kernel,瓶颈都可能变化;full attention 下,上下文越长,新 query 要读取的历史 KV 通常也越多。

开头快,不保证一整篇都写得快。

所以下次看到一个 tokens/s 数字,可以多问一句:

它是在读多长的输入、写多长的回答、同时服务多少人?

5. 别人的长文档,会让你的回答卡住吗?

你的演讲稿正写到一半,同事又提交了一整本产品手册。

两条请求如果共享计算资源,长 prefill 占用过长的连续执行时间,就可能让你的后续输出出现停顿。

怎么办?

一个办法是把长文档的输入计算拆成小段,在调度迭代中给正在生成的回答安排计算机会。

这叫 chunked prefill,实现中也可以把 prefill chunk 与 decode 工作放进同一批处理。

这相当于把工作分几次安排。

上下文保持连续:处理后面的 chunk 时,仍会读取前面已经保留的 KV。

另一个办法是分配不同的资源:一组 worker 专门处理输入,另一组接着生成回答。

这叫 prefill/decode disaggregation,两边可以分别调整容量和调度策略。

A 的输出继续增长,B1–B3 分块处理并保留各自的前缀状态。随后展示另一种安排:把 B 的 KV 和继续生成所需状态交给 decode 资源池。A、B 的请求状态始终独立,两种方法可以组合,动作时长不代表实际性能。

这时必须完成一次交接。

后一组要拿到文档对应的 KV,以及继续生成所需的状态,例如 token 位置、已经选出的新 token,才能从原来的地方接着算。

两边的模型参数、位置处理和 cache 格式需要兼容,通常也都要具备执行同一模型的能力。

分工可以减少阶段间干扰,但交接有传输、转换和协调成本。

请求很短时,来回传数据甚至可能比省下的时间还多。

还有更直接的问题:

这份文档刚才已经处理过,再问一次呢?

比如先问"总结风险",再问"列出截止日期"。

如果请求开头的实际 token prefix 完全一致、执行条件兼容,且这部分 KV 仍被保留,就可以通过 prefix caching 复用它。

新的问题仍需要 prefill,新的回答仍需要 decode。

三个办法各管一件事:prefix caching 省去可复用的输入计算;chunked prefill 把剩余输入计算分段安排;disaggregation 把两个阶段放到不同资源上。

6. Spark 处理输入,Mac 接着写,行不行?

沿着这条思路,问题就很自然了:能不能让一台机器处理长文档,再让另一台接着生成回答?

EXO 展示过这样一个组合:DGX Spark 执行 prefill,通过网络传递 KV,再由 Mac Studio M3 Ultra 执行 decode,延续同一条请求。

交接也可以边算边进行:前面某层的 KV 准备好,先发出去,同时继续计算后面的层。

这样,部分传输与计算可以重叠;末尾尚未传完的数据、格式转换,仍然要花时间。

第三方演示中的阶段分工,以及可实现的逐层 KV 传输重叠。

回到开会前那份文档:你只要三条结论,输入却有几十页。

Prefill 要处理这些已知内容;decode 利用保存的状态,把回答逐步生成出来。

下次觉得 AI 慢,可以先分清:是迟迟没有第一个输出,还是回答开始后写得慢?

再把输入长度、输出长度和并发量放进去看,优化的问题就具体多了。