AB
AiBoss
Wiki

什么是分块预填充(Chunked Prefill)?一种把长输入切块与解码请求混批的推理调度技术

分块预填充(Chunked Prefill)是一种大语言模型推理调度技术:把一次很长的预填充(prefill)拆成若干较小的块,与正在进行的解码(decode)请求放进同一批次处理。vLLM 官方文档将其列为实验性特性,NVIDIA TensorRT-LLM 也提供同名特性。它主要用来改善 token 间延迟(ITL)与 GPU 利用率,代价是首 token 延迟(TTFT)与吞吐之间的取舍需要调参。

分块预填充(Chunked Prefill)是一种面向大语言模型推理服务的调度技术:它把一次较长的预填充(prefill,即对输入提示词全部 token 计算 KV 缓存并生成首个输出 token 的阶段)切分成若干较小的块(chunk),再把这些块与正在进行的解码(decode,即逐 token 生成后续输出的阶段)请求编排进同一个批次中执行。它要解决的问题是:预填充与解码对 GPU 资源的利用方式不同,若把二者严格分开调度,要么让解码请求排队等待长预填充,要么让 GPU 在只跑解码时算力闲置。

为什么重要

按 vLLM 官方文档的描述,在默认调度策略下,vLLM 的调度器优先处理预填充,并且不会把预填充与解码放进同一个批次。该文档称,这一策略优化了首 token 延迟(time to first token,TTFT),但会造成较慢的 token 间延迟(inter token latency,ITL)以及低效的 GPU 利用率。

NVIDIA 开发者博客从计算特征角度解释了这一矛盾:预填充阶段要处理全部输入 token 以计算 KV 缓存,计算密集,能充分利用 GPU 的大规模并行算力;解码阶段每次只处理新生成的一个 token,计算强度低得多。TensorRT-LLM 支持在途批处理(in-flight batching),让多个请求并行处理、预填充与解码阶段同时进行;但该博客指出,相比分阶段批处理(后者会暂停解码请求去处理预填充请求)虽然更高效,这种做法仍可能引入延迟,因为解码阶段要等预填充请求完成才能继续。

换言之,长输入请求会形成瓶颈:它占据大量计算资源,期间解码请求得不到及时推进,用户看到的是输出「卡顿」;而若为保护解码而限制预填充,GPU 算力又得不到充分利用。分块预填充正是针对这一矛盾提出的调度方案。

工作机制

以 vLLM 官方文档的说明为主线,分块预填充的核心做法可以拆成以下几点:

  1. 改变优先级:启用分块预填充后,调度策略从「预填充优先」改为「解码优先」。调度器会先把所有待处理的解码请求排入批次,然后再调度预填充。
  2. 按 token 预算调度预填充:在还有可用 token 预算(由 max_num_batched_tokens 控制)时,调度器才安排待处理的预填充请求。
  3. 切块:如果最后一个待处理的预填充请求无法完整放入 max_num_batched_tokens,就把它切分成块,只放入能容纳的部分,剩余部分留待后续迭代。
  4. 混批执行:由于块变小了,计算密集的预填充请求与访存密集的解码请求可以落在同一个批次里,从而提升 GPU 利用率。

vLLM 文档把该策略的收益归纳为两点:一是解码请求被优先处理,改善了 ITL 与生成解码过程;二是把计算密集(预填充)与访存密集(解码)的请求放进同一批次,有助于获得更好的 GPU 利用率。

NVIDIA 开发者博客则从内存角度描述了 TensorRT-LLM 的实现思路:由于内存占用取决于每次迭代处理的 token 数量,使用预填充块可以把内存消耗与传入请求的上下文长度解耦,系统因而能在不增加内存需求的前提下处理更长的上下文和更高的并发。该博客还介绍了动态块大小(dynamic chunk sizing):它依据 GPU 利用率指标自动确定激活缓冲区大小,免去开发者手动配置最大输入序列长度。

两种实现都强调块大小是一个需要权衡的参数。vLLM 文档称,较小的批大小因预填充对解码的打断更少而获得更好的 ITL,较大的批大小则因能把更多预填充放入批次而获得更好的 TTFT;若 max_num_batched_tokens 等于 max_model_len,效果几乎等同于默认调度策略(区别在于仍然优先处理解码)。NVIDIA 博客同样指出,较大的块大小会减少处理预填充序列所需的迭代次数、降低 TTFT,但同时会增加时间开销,因此需要根据期望的用户体验与可用 GPU 资源做取舍。

典型例子

vLLM 官方文档给出了具体的启用方式:在构造 LLM 对象时传入 enable_chunked_prefill=True,例如以 meta-llama/Llama-2-7b-hf 为模型,并通过 max_num_batched_tokens 调节性能。文档说明,分块预填充下 max_num_batched_tokens 的默认值为 512,并称在初始基准测试中该默认值在 A100 上具有最好的 ITL;文档同时提醒,默认批大小(512)是针对 ITL 优化的,吞吐可能低于默认调度器,若以吞吐为目标建议把 max_num_batched_tokens 设为大于 2048。文档将该特性标注为实验性(experimental)。

NVIDIA 开发者博客介绍的是 NVIDIA TensorRT-LLM 中的同名特性,文章署名作者为 Amr Elmeleegy、Nick Comly 与 Sharan Chetlur,发布日期为 2024 年 11 月 15 日。该文称分块预填充可提高 GPU 利用率并简化开发者的部署体验,并提到它建立在上一篇讨论 TensorRT-LLM 高级 KV 缓存优化特性的文章之上,那篇文章称相关特性在需要系统预填充的用例中带来最高 5 倍的性能提升。该文还提到,读者可查阅 TensorRT-LLM 的 GitHub 文档开始使用分块预填充。

边界与常见误解

它不是无代价的加速开关。vLLM 文档明确指出,默认批大小 512 是针对 ITL 优化的,吞吐可能低于默认调度器;也就是说,启用分块预填充是在 TTFT、ITL 与吞吐之间重新分配预算,而非全面提升。NVIDIA 博客也强调块大小存在用户体验与系统吞吐之间的权衡。

它不等于「预填充消失了」。分块只是把一次长预填充拆成多次迭代完成,总计算量并未减少;改变的是这些计算与解码请求在时间上的交错方式。

不要把某一家的默认值当成通用结论。512 这一默认值出自 vLLM 文档所述的初始基准测试(A100 环境),是针对 ITL 的选择;不同硬件、不同负载下的最优值需要自行测量。NVIDIA 博客所述的动态块大小是 TensorRT-LLM 的实现机制,并非所有推理框架都具备。

成熟度需按来源区分。vLLM 文档将该特性称为实验性特性;NVIDIA 博客则将其作为 TensorRT-LLM 的正式特性介绍。两者对同一概念的定位表述并不相同,引用时应注明出处。

它不解决所有长上下文问题。该技术缓解的是调度层面的资源争用与内存占用随上下文增长的问题,KV 缓存本身仍需存储;NVIDIA 博客把分块预填充与 KV 缓存优化并列为不同层面的手段。

参考资料