Chunked Prefill:别让长 prompt 阻塞整个推理引擎
前面几篇文章里,我们围绕 KV cache 做了几层优化。
Block Manager 让 KV cache 变成可分配、可释放、可复用的 blocks。PagedAttention 让 attention 可以通过 block table 读取不连续的 KV cache。Prefix caching 进一步让多个请求复用相同前缀的 KV cache,避免重复 prefill。
这些优化解决了很多问题。
但还有一种场景仍然很麻烦:一个请求真的带着很长的新 prompt,并且这个 prompt 没法命中 prefix cache。
例如用户上传了一篇很长的文档,让模型总结。
或者 RAG 系统检索出一大段上下文,而这段上下文之前没有被缓存过。
又或者代码助手把很长的文件内容拼进 prompt。
这类请求必须做长 prefill。
长 prefill 的问题不只是“它自己慢”。
更重要的是,它会影响其他正在 decode 的请求。
如果一个长 prompt 一次性进入 prefill,占用了大量 GPU 时间,那么其他正在流式输出的请求可能会卡住。用户看到的效果就是:模型前面还在持续输出,突然停顿了一段时间。
Chunked prefill 要解决的正是这个问题。
它的核心思路是:不要让一个长 prefill 一次性吃完整个调度预算,而是把长 prompt 切成多个 chunk,分多轮处理。这样调度器就可以在长 prefill 的多个 chunk 之间插入 decode,让已经开始输出的请求继续流畅生成。
这篇文章我们要讨论:为什么长 prefill 会阻塞 decode,chunked prefill 如何改变调度模型,它和 KV cache、Block Manager、PagedAttention 有什么关系,以及第一版 mini vLLM 应该如何实现一个清晰可控的版本。
长 prefill 为什么会造成卡顿
先回到 prefill 和 decode 的区别。
prefill 处理完整 prompt。
如果 prompt 有 8192 个 token,那么模型要一次性处理这 8192 个 token,并为它们写入 KV cache。
decode 每一轮通常只处理每个 running Sequence 的一个 token。
假设系统里有 128 个 running 请求,每个请求每轮 decode 一个 token,那么这一轮 decode 的输入 token 数大概是 128。
现在来了一个新请求,prompt 长度是 8192。
如果调度器把这个请求完整 prefill 放进下一轮,那么它的 prefill token 数远远大于这一轮 decode token 数。
这会导致一个现象:长 prefill 可能主导这一轮的 GPU 时间。
从用户体验角度看,正在流式输出的 128 个请求都要等这个长 prefill 执行完,才能继续得到下一个 token。
这就是流式输出卡顿。
问题不在于长请求不应该被处理。长请求当然也要处理。
问题在于,它不应该一次性占用太多连续时间,让其他请求完全没有机会推进。
这和操作系统调度很像。
如果一个长任务一直霸占 CPU,交互式任务就会卡顿。
所以操作系统会做时间片。
Chunked prefill 在 LLM 推理里扮演的角色有点类似:把一个大 prefill 拆成多个小片段,让调度器可以在片段之间重新做决策。
Chunked prefill 改变了什么
没有 chunked prefill 时,一个 waiting 请求只有两种状态:
还没 prefill
prefill 完成,进入 running
prefill 是一次性动作。
有了 chunked prefill 后,prefill 不再是一次性动作,而是一个可以分多轮推进的过程。
一个长 prompt 会被拆成多个 chunk。
例如 prompt 长度是 8192,chunk size 是 1024。
那么 prefill 会变成 8 个 chunk:
chunk 0: token 0 到 1023
chunk 1: token 1024 到 2047
chunk 2: token 2048 到 3071
...
chunk 7: token 7168 到 8191
每一轮调度器只处理其中一部分 chunk。
处理完一个 chunk 后,请求还不能进入 decode,因为完整 prompt 还没处理完,也还不能采样第一个输出 token。
只有最后一个 chunk 处理完成后,模型才有 prompt 最后位置的 logits,可以采样第一个 generated token。
所以 Sequence 的状态也要变得更细。
以前是:
WAITING -> RUNNING -> FINISHED
现在可以理解成:
WAITING -> PREFILLING -> RUNNING -> FINISHED
PREFILLING 表示这个请求正在分块 prefill,但还没完成完整 prompt。
这个状态非常重要。
因为它告诉 Scheduler:这个请求已经占用了一部分 KV cache,也已经计算了一部分 prompt,但还没有首 token 输出。
它既不是普通 waiting,也不是 running。
它是一个正在进行中的 prefill 任务。
Chunked prefill 和 KV cache 的关系
Chunked prefill 能成立,是因为 KV cache 可以逐步写入。
处理 chunk 0 时,模型为 token 0 到 1023 写入 KV cache。
处理 chunk 1 时,模型需要让 token 1024 到 2047 attend 到前面 token 0 到 1023 的 KV cache,并把自己的 KV 写入后续位置。
也就是说,chunked prefill 的每个 chunk 都依赖之前 chunk 的 KV cache。
这和 decode 有点像,但又不完全一样。
decode 每轮输入一个 token,读取历史 KV,写入当前 token KV。
chunked prefill 每轮输入一段 suffix tokens,读取之前 chunks 的 KV,同时这一段内部还要做 causal attention。
所以它可以看成介于 prefill 和 decode 之间的形态:
普通 prefill:
一次性输入完整 prompt
decode:
每次输入一个 token
chunked prefill:
每次输入 prompt 的一段连续 tokens
处理某个 chunk 时,模型需要看到:
之前 chunks 的 KV cache
当前 chunk 内部的 causal attention
当前 chunk 里的第一个 token 可以看之前所有 chunks。
当前 chunk 里的第二个 token 可以看之前所有 chunks 加当前 chunk 第一个 token。
以此类推。
这意味着 ModelRunner 需要支持一种新的执行路径。
它不是普通完整 prefill,也不是单 token decode,而是:
prefill_chunk(sequence, chunk_token_ids, start_pos)
其中 start_pos 表示这个 chunk 在完整 prompt 中的起始位置。
这个 start_pos 非常关键,因为 position ids 必须接在已经缓存的 prefix 后面。
如果 position 错了,模型就会认为这段 token 出现在错误位置,输出自然会错。
Sequence 需要维护 prefill 进度
引入 chunked prefill 后,Sequence 需要维护更多状态。
以前一个 waiting Sequence 只需要保存完整 prompt tokens。
现在它还需要知道 prompt 已经 prefill 到哪里。
可以引入一个字段:
prefilled_token_count: int
它表示 prompt 中已经完成 prefill,并且 KV cache 已经写入的 token 数量。
初始时:
prefilled_token_count = 0
处理第一个 chunk 后:
prefilled_token_count = chunk_size
直到:
prefilled_token_count = prompt_len
这时完整 prompt prefill 完成,可以根据最后一个 prompt token 的 logits 采样第一个 generated token,并把 Sequence 状态改成 running。
这里要注意 prefilled_token_count 和前面文章里提到的 cached_token_count 的关系。
在 prompt prefill 阶段,它们通常一致,表示已经写入 KV cache 的 prompt tokens 数量。
当 prefill 完成并采样出第一个 generated token 后,generated token 还没有写入 KV cache。它会作为下一轮 decode 的输入。
所以状态语义仍然需要保持清楚:
prefilled_token_count:
prompt 中已经 prefill 的 token 数
cached_token_count:
已经写入 KV cache 的 token 数
generated_token_ids:
已经采样出来、要返回给用户的 token
对于普通请求,prefill 完成后:
prefilled_token_count = prompt_len
cached_token_count = prompt_len
generated_token_ids = [first_generated_token]
第一轮 decode 后:
cached_token_count = prompt_len + 1
generated_token_ids = [first_generated_token, second_generated_token]
Chunked prefill 没有改变这个语义,只是让 prompt 的 cache 写入过程变成多步。
Scheduler 如何调度 prefill chunk
有了 chunked prefill 后,Scheduler 不再把 waiting 请求的 prefill 成本看作完整 prompt 长度。
它会为每个 waiting 或 prefilling Sequence 选择一个 chunk。
比如配置:
max_prefill_chunk_size = 1024
一个 prompt 长度为 8192 的请求,第一次被调度时最多处理 1024 个 token。
处理完后,它仍然留在 prefilling 队列。
下一轮再处理下一个 1024 tokens。
直到最后一个 chunk。
调度时,Scheduler 需要决定每一轮 token budget 如何在 decode 和 prefill chunks 之间分配。
一种简单策略是:
先为 running decode 预留一部分 budget
剩余 budget 用来调度 prefill chunks
每个 prefill chunk 不超过 max_prefill_chunk_size
这样可以避免长 prefill 完全挤占 decode。
例如一轮总 token budget 是 4096。
可以要求 decode 至少有机会运行,然后剩余部分再给 prefill。
running decode: 128 tokens
prefill chunk: 最多 3968 tokens
如果系统里 running 请求很多,也可以限制 decode 最多占多少 budget,避免 waiting 请求饥饿。
所以 chunked prefill 本质上不是一个孤立优化,而是让调度器有了更细的控制粒度。
以前,长 prefill 是一个不可拆的大任务。
现在,它变成多个可调度的小任务。
一旦任务可拆,调度器就可以在吞吐和延迟之间做更细致的权衡。
Chunked prefill 对首 token 延迟的影响
Chunked prefill 能改善 decode 卡顿,但它可能影响长 prompt 请求自己的首 token 延迟。
为什么?
因为没有 chunked prefill 时,如果一个长 prompt 一旦开始 prefill,它会连续执行到完成,然后立刻产生第一个 token。
有了 chunked prefill 后,这个长 prompt 可能被拆成多个 chunk,中间穿插其他请求的 decode 或 prefill。它自己的完整 prefill 完成时间可能变长。
也就是说,它牺牲了一部分长请求自己的连续执行速度,换来整个系统更平滑的服务体验。
这就是系统调度里的典型取舍。
如果你只看单个长请求,chunked prefill 可能让它更慢。
但如果你看整个服务,它可以降低其他请求的 decode 卡顿,改善整体尾部延迟和流式输出稳定性。
所以 chunked prefill 不应该被理解成“让长 prompt 更快”。
更准确地说,它是:
让长 prefill 不再长时间独占推理引擎
它改善的是系统层的公平性和交互体验。
当然,在某些 workload 下,chunked prefill 也可能提升吞吐,因为它让 decode batch 更稳定,减少长 prefill 带来的调度抖动。但这不是唯一目标。
Chunk size 如何选择
chunk size 是 chunked prefill 的核心参数。
如果 chunk size 太大,长 prefill 仍然会造成明显卡顿。
如果 chunk size 太小,prefill 被切得太碎,调度次数增加,kernel 启动和中间状态管理开销也会增加。
所以 chunk size 是一个取舍。
大 chunk 的特点是:
prefill 效率更高
调度开销更小
但可能阻塞 decode 更久
小 chunk 的特点是:
更容易和 decode 交错
流式输出更平滑
但 prefill 效率可能下降
调度和状态管理开销更高
在教学版 mini vLLM 里,可以先让 chunk size 等于一个固定配置,比如 512 或 1024。
后续可以做 benchmark 比较不同 chunk size 下的表现。
重点关注几个指标:
长 prompt 请求的首 token 延迟
普通 running 请求的每 token 延迟
整体吞吐
P90 / P99 延迟
GPU 利用率
不同 workload 下,最佳 chunk size 可能不同。
如果请求普遍是短 prompt,chunked prefill 影响不大。
如果长 prompt 很多,chunk size 的选择会明显影响体验。
这也是为什么生产系统中调度参数通常需要结合实际流量调优。
和 Prefix Caching 的关系
Prefix caching 和 chunked prefill 都围绕 prefill 优化,但它们解决的问题不同。
Prefix caching 解决的是重复前缀。
如果一个 prompt 的前缀已经算过,就直接复用 KV cache,减少需要计算的 tokens。
Chunked prefill 解决的是长 prefill 阻塞。
如果一个 prompt 仍然很长,就把它拆成多个 chunk,避免一次性占用太久。
它们可以同时工作。
一个请求进来后,系统可以先查 prefix cache。
如果命中前缀,那么真正需要 prefill 的只是 suffix。
如果 suffix 仍然很长,再对 suffix 做 chunked prefill。
例如完整 prompt 长度是 12000。
Prefix cache 命中了前 8000。
剩余 suffix 是 4000。
如果 chunk size 是 1000,那么只需要处理 4 个 chunks。
这条链路可以理解成:
完整 prompt
先做 prefix cache match
得到未命中的 suffix
对 suffix 做 chunked prefill
最后进入 decode
这种组合很自然。
Prefix caching 减少 prefill 总量。
Chunked prefill 控制 prefill 粒度。
一个解决“算不算”的问题。
一个解决“怎么调度”的问题。
和 Block Manager 的关系
Chunked prefill 也会影响 Block Manager。
普通 prefill 时,系统可以一次性为完整 prompt 分配所有 blocks。
chunked prefill 有两种分配策略。
第一种是一次性预分配完整 prompt 所需 blocks。
请求开始 prefill 前,就把整个 prompt 的 KV cache 空间都分配好。
这样实现简单,后续每个 chunk 只需要按位置写入。
缺点是,如果这个请求在 prefill 过程中被取消,或者长时间等待,它会提前占用大量 blocks。
第二种是按 chunk 分配 blocks。
每处理一个 chunk,只为这个 chunk 需要写入的新 tokens 分配 blocks。
这样更节省资源,也更符合按需原则。
缺点是每一轮调度前都要检查和分配 blocks,实现更复杂。
对于 mini vLLM,建议采用按 chunk 分配。
因为这更符合 Block Manager 的设计思想:需要多少 cache,就分配多少 block。
同时也能避免一个长 prompt 一进来就占用大量 KV cache,导致其他请求无法运行。
但按 chunk 分配要求 Scheduler 在调度每个 chunk 前检查 block budget。
例如某个 chunk 有 1024 tokens,block size 是 16,那么最多需要 64 个 blocks。
如果当前 free blocks 不够,就不能调度这个 chunk。
这再次说明,Scheduler 必须同时考虑:
token budget
sequence budget
KV cache block budget
调度器越往后演进,就越不像一个普通队列,而更像一个资源分配器。
chunked prefill 中的输出问题
在普通 prefill 中,prefill 完成后会采样第一个 generated token,用户可以看到首 token。
但 chunked prefill 的中间 chunk 不会产生用户可见输出。
因为 prompt 还没处理完,模型还不能根据完整 prompt 采样第一个 token。
例如 prompt 被拆成 4 个 chunks。
chunk 0 完成:不输出
chunk 1 完成:不输出
chunk 2 完成:不输出
chunk 3 完成:采样第一个 token,输出
所以 Sequence 在 PREFILLING 状态下通常不会产生 RequestOutput,最多只更新内部进度。
这也意味着,chunked prefill 可能让一个请求在系统内部执行了多轮,但用户仍然没有看到输出。
因此我们需要区分两种进度。
内部进度:
prefilled_token_count 增加
KV cache 持续写入
用户可见进度:
首 token 是否产生
是否开始 streaming
这对监控也有影响。
一个请求可能已经被执行了一部分 prefill,但首 token latency 仍然在增长。
如果只看用户输出,会以为请求一直没开始。
如果看内部状态,会知道它正在 prefill。
真实系统里,这类状态对于排查长 prompt 延迟很重要。
chunked prefill 的状态机
现在我们可以重新定义 Sequence 状态。
简单版状态机可以是:
WAITING:
请求刚进入,还没有开始 prefill
PREFILLING:
prompt 已经 prefill 一部分,但还没完成
RUNNING:
完整 prompt prefill 完成,已经产生首 token,正在 decode
FINISHED:
请求结束,资源释放
状态转换是:
WAITING -> PREFILLING
处理第一个 prefill chunk,但 prompt 未完成
WAITING -> RUNNING
prompt 很短,一个 chunk 就完成,并产生首 token
PREFILLING -> PREFILLING
继续处理下一个 chunk,但还没完成 prompt
PREFILLING -> RUNNING
最后一个 chunk 完成,采样首 token
RUNNING -> FINISHED
decode 达到停止条件
WAITING / PREFILLING / RUNNING -> FINISHED
请求被取消或出错
这个状态机比原来多了一个 PREFILLING,但它让系统语义更准确。
尤其是取消请求时,PREFILLING 请求也要释放已经分配的 blocks。
不能只处理 RUNNING 请求的释放。
否则长 prompt 在 prefill 中被取消时,会造成 KV cache 泄漏。
ModelRunner 需要支持 prefill_chunk
为了实现 chunked prefill,ModelRunner 最好有一个独立入口:
run_prefill_chunk(sequence, chunk_token_ids, start_pos)
它要做几件事。
首先,构造 chunk 的 input_ids。
然后,构造 position_ids,从 start_pos 开始。
接着,拿到当前 Sequence 已经存在的 block table,让当前 chunk 可以 attend 到之前的 KV cache。
然后,为当前 chunk 分配写入 slots,把 chunk tokens 的 K 和 V 写入对应 blocks。
最后,如果这是最后一个 chunk,需要返回最后一个 prompt token 的 logits,用来采样第一个 generated token。
如果不是最后一个 chunk,可以不返回 logits,或者返回但 Engine 不使用。
从实现上看,最后一个 chunk 和中间 chunk 的区别在于:只有最后一个 chunk 需要触发 sampling。
这可以让 Engine 逻辑更明确。
如果 chunk 后 prompt 未完成:
更新 prefilled_token_count
不采样
不输出
如果 chunk 后 prompt 完成:
用最后位置 logits 采样首 token
状态转 RUNNING
产生 RequestOutput
这样 chunked prefill 就能自然接入现有 Engine.step。
调度策略的第一版设计
第一版 chunked prefill 调度可以保持简单。
每一轮先调度一批 running sequences 做 decode。
然后用剩余 token budget 调度 prefill chunks。
对于每个 waiting 或 prefilling Sequence,计算它本轮最多能处理多少 prompt tokens:
remaining_prompt_tokens = prompt_len - prefilled_token_count
chunk_len = min(remaining_prompt_tokens, max_prefill_chunk_size, remaining_token_budget)
如果 chunk_len 大于 0,并且 Block Manager 有足够 blocks,就调度这个 chunk。
这里有一个策略问题:waiting 和 prefilling 谁优先?
如果优先 waiting,系统会不断启动新请求,但旧的 prefilling 请求可能很久才完成,首 token 延迟变差。
如果优先 prefilling,已经开始 prefill 的请求能更快完成首 token,但新请求可能等待更久。
第一版可以优先 prefilling,再处理 waiting。
原因是 prefilling 请求已经占用了一部分 KV cache。如果长期不继续推进,就会占着资源却没有用户可见输出。让它尽快完成 prefill,进入 running 或 finished,通常更合理。
但这也不是绝对最优。
后续可以根据等待时间和优先级做更复杂策略。
第一版的目标仍然是:语义清晰,链路跑通。
chunked prefill 的正确性验证
验证 chunked prefill 的核心方法,仍然是和普通完整 prefill 对齐。
给定同一个 prompt,用两种方式运行:
方式 A:
一次性完整 prefill
方式 B:
分多个 chunk prefill
在 greedy decoding 下,两者产生的 token 应该一致。
测试时要覆盖这些情况。
prompt 长度小于 chunk size。
prompt 长度等于 chunk size。
prompt 长度刚好跨过 chunk size。
prompt 长度跨多个 chunks。
chunk size 不是 block size 的整数倍。
chunk 边界和 block 边界不对齐。
最后一种特别重要。
例如 block size 是 16,chunk size 是 1000。
1000 不是 16 的整数倍。
这意味着某个 chunk 可能结束在 block 中间,下一个 chunk 会继续写同一个 block。
Block Manager 必须正确处理这种情况。
如果写入 slot 计算依赖 chunk 边界,而不是依赖全局 token index,就很容易出错。
正确做法是永远使用完整 prompt 中的 token index 来计算 block 和 offset。
也就是:
global_token_index = prefilled_token_count + local_index_in_chunk
logical_block_id = global_token_index // block_size
offset = global_token_index % block_size
只要使用全局 token index,chunk 边界和 block 边界是否对齐都不影响正确性。
这就是 chunked prefill 实现中最重要的细节之一。
小结
这一篇我们引入了 chunked prefill。
它解决的问题不是让单个长 prompt 一定更快,而是避免长 prefill 一次性阻塞整个推理引擎。
没有 chunked prefill 时,一个长 prompt 可能占用大量连续 GPU 时间,导致正在 running 的请求 decode 卡顿。引入 chunked prefill 后,长 prompt 被拆成多个 chunk,Scheduler 可以在 chunk 之间插入 decode,让流式输出更加稳定。
Chunked prefill 会让请求状态机多出一个 PREFILLING 状态。这个状态表示 prompt 已经部分写入 KV cache,但还没有完成完整 prefill,也还没有产生首 token。
Sequence 需要维护 prefilled_token_count。Block Manager 需要按 chunk 分配和写入 KV cache slots。ModelRunner 需要支持 prefill_chunk,并保证 position ids、block table、slot mapping 都基于完整上下文位置。Scheduler 则需要在 token budget、sequence budget 和 block budget 之间做更细粒度的分配。
它和 prefix caching 可以自然结合。
Prefix caching 先减少需要计算的 prompt tokens。
Chunked prefill 再把剩余未命中的 suffix 拆成多个 chunk 调度。
这两个优化一个解决“少算”,一个解决“分开算”。
到这里,mini vLLM 已经具备了一个较完整的在线推理引擎结构:动态调度、KV cache block 管理、PagedAttention、prefix cache、chunked prefill、流式输出和采样控制。
下一篇文章,我们会进入 benchmark。
到现在为止,我们已经做了很多优化,但不能只凭感觉说它更好。我们需要设计一套实验,分别测首 token 延迟、每 token 延迟、吞吐、显存利用率和尾部延迟,观察 naive generate、KV cache、continuous batching、PagedAttention、prefix caching、chunked prefill 分别带来了什么变化。