Sarathi-Serve：chunked-prefills

最新推荐文章于 2025-04-01 09:47:31 发布

smartcat2010

最新推荐文章于 2025-04-01 09:47:31 发布

阅读量562

点赞数 3

分类专栏： LLM推理优化文章标签：大模型 GPU

本文链接：https://blog.youkuaiyun.com/smartcat2010/article/details/144952828

版权

LLM推理优化专栏收录该内容

39 篇文章

订阅专栏

将prefiling阶段的句子，和decoding阶段的句子，混合到一个batch里，为什么能增大throughput？

左图，Prefill阶段，因为是Compute-bound的，计算单元已经打满了，再增大batch也不会增大吞吐量了。

右图，Decode阶段，因为是Memory-bound的（实际是memory-latency-bound)，增大batch几乎没有增大延迟，而计算单元远远没打满，因此可以随batch增大而线性增长。

将处于这2个阶段的句子，混合到同一个batch里，好处：

prefill搭上decode的便车，能用上decode阶段被浪费的算力。
decode搭上prefill的便车，合并数据的读取次数，做到1次读取，大家共享。

Chunked-prefills

将prefill阶段的句子，拆分成多个block，每个block包含一部分tokens的计算。

目的：

1. 单卡执行时，降低decode阶段的token的延迟。（如果和整个长句子的prefill阶段样本捆绑，则decode阶段的句子本token延迟显著增大；如果前者切分的短些，则后者延迟可显著减少）。

2. pipeline并行时，可减少气泡。

pp并行的气泡问题：

上图，AB是同一个batch，CD是同一个batch；p下标是prefill阶段，d下标是decode阶段。

ABCD这4个样本的prompt长度各不相同，因此prefill阶段的耗时各不相同。

解决方法：

A被拆成了2个chunk；B被拆成了3个Chunk；C被拆成了2个；D被拆成了2个；

优先调度decode阶段的token，在显存&计算余量的缝隙里，考虑加入多少个token的prefill阶段。如上图， $A_{d1}$ 优先确保，然后才让 $C_{p1}$ 的一些tokens和其共享slot。详细解释：

decode-maximal batching原则：即优先往batch中添加需要做decode的序列，直到添加不动为止（即我们预留给decode的KV cache空间已经不足了，无法存放新的KV cache了）。然后我们再根据这个batch中剩余的tokens预算，对需要做prefill的序列做chunk切割，把对应的prefill tokens添加进batch中。

注意：前面的chunk，计算量小；后面的chunk，计算量大；

解决：prefill阶段，前面的chunk，给的token多些；后面的chunk，给的token少些。