问题背景

在衡量大模型推理服务性能时，通常关注吞吐和延迟这两个指标。现有的系统往往采用批处理的方式来增加系统吞吐，但这样会带来TBT（time-between-tokens）延迟的增加。现有系统都无法同时满足推理任务的高吞吐量和低延迟。 

现有方案的问题

1. Prefill 和 Decode 阶段存在较大差异

1）批处理极大地提高了Decode阶段的吞吐量，但对Prefill阶段的吞吐量影响很少。将Prefill和Decode阶段内可以分为线性层（Linear）、注意力层（Attention）和其他。线性层占运行时间成本的绝大部分。即使在序列长度较高的情况下，线性层仍占总时间的 80% 以上。因此，优化线性层对改进大模型推理非常重要。

2）Decode阶段的计算资源利用率较低。由于线性层是优化大模型推理的关键，针对线性层中的矩阵乘法操作进行分析，可以得出Decode阶段的计算强度非常小。并且通过实验可以发现，在Decode阶段，可以同时处理更多的token，而不会显著增加其延迟。

2. 现有调度策略不适合大模型在线推理服务

这张图比较了现有推理系统的不同调度策略，请求 A 和 B 在开始时处于解码阶段，经过一次迭代后，请求 C 和 D 进入系统。

根据调度的优先策略不同，目前的调度策略可以分为Prefill优先和Decode优先两类。

Prefill优先：vLLM和Orca都是使用基于 FCFS的 迭代级别（Iteration-level）批处理，并优先执行Prefill阶段。其中vLLM只支持仅Prefill和仅Decode的组合执行，而Orca支持Prefill和Decode的混合执行。这两个推理系统都是通过最大化decode阶段的batch size来提高系统整体吞吐。但是优先执行Prefill阶段会抢占Decode阶段的执行，造成请求A和请求B的decode阶段的延迟响应，造成Decode阶段的生成停顿（Generation Stall），表现为延迟突增（Latency Spike）。

Decode优先：请求级别（Request-level）的批处理系统（如FasterTransformer）采用。这种策略优化了TBT这一延迟指标，但严重限制了吞吐量。因为即使批次中的某些请求提前完成，也会减小batch size继续执行，直到最后一个请求完成。

而理想的调度策略是将Prefill和Decode阶段进行交错与混合执行，同时考虑到吞吐和延迟间的权衡。

3. 流水线气泡严重影响推理性能

现有的方案认为在并发任务推理场景下，微批处理（Micro-batching）可以消除流水线并行中的气泡。但是实验表明，即使采用迭代级调度，在使用流水线并行（Pipeline Parallelism）的分布式LLM推理服务中，流水线气泡也会浪费大量 GPU 周期。尽管微批处理可以缓解这一问题，但由于 LLM 推理中预填充和解码长度的差异，标准的微批处理调度仍会导致气泡，浪费 GPU 周期并降低系统整体吞吐量。

Ocra 中两阶段的流水线因为批次执行时间不均匀，仍然存在流水线气泡。根据气泡产生的原因可以分为三类：

·      由于两个连续微批次中的 Prefill 阶段 token 数量不同而产生的气泡

·      由于 Prefill 和 Decode 阶段的计算时间不同而产生的气泡

·      由于微批次之间的 Decode 阶段计算时间不同而产生的气泡

当Prefill阶段的时间更长、频率更高，随着输入提示词（Prompt）长度和batch size的增加，这一问题会更加严重。理想的情况是创建执行时长相对统一的批次，来最小化气泡。

设计方案

为了解决之前的挑战，实现高吞吐和可预测的尾时延，Sarathi-Serve 系统使用了两个关键技术：分块预填充（Chunked-prefill）和无停顿批处理（Stall-free Batching）。

1. 分块预填充 Chunked Prefill

为了增加整体吞吐，最直接的方法就是将计算受限的Prefill阶段和内存受限的Decode阶段混合执行。然而，在实际使用场景中，输入的提示词往往包括数千个token。把耗时长的Prefill阶段与Decode混合执行，将导致TBT延迟的增加。

这个问题的根本原因在于Prefill阶段执行时长明显大于Decode阶段。分块预填充技术通过将Prefill阶段拆分为多个小块，每次迭代只执行一个Prefill块，从而缩短每次混合执行中的Prefill部分的时间，避免延长Decode阶段。而且就算是把Prefill拆分了，其计算需求也足够大，能通过混合执行，减少Decode阶段的GPU计算资源浪费。

2. 无停顿批处理 Stall-free Batching

Sarathi-Serve的调度策略针对减少推理请求停顿进行了优化。不同于vLLM和Orca，Sarathi-Serve利用Decode迭代中的计算密度空隙来执行Prefill chunk，在不延迟系统中的Decode请求执行的情况下最大化吞吐。

该调度策略按顺序填充当前批次的解码任务、部分完成的预填充任务和新请求。在将Prefill请求添加到批次时，可以计算在剩余token预算内能容纳的最大块大小。只有在所有正在运行的请求都被纳入到下一迭代批次后，才接受新请求。只要确保每个批次处理的token数不超过设置的上限，通过合理设置批次的token执行上限，最终可以实现无停顿批处理，避免了延迟峰值，同时保持高效的计算资源利用。

3. Token 数上限的确定

确定token上限需要在TBT的SLO要求和chunked-prefill的开销之间进行权衡。Token上限越小，则可以在一定范围内尽可能减少TBT，降低混合执行中Prefill操作对Decode操作的影响。但是这样会导致过度分块，使GPU计算资源利用率下降（tile-quantization效应）以及反复访问KV Cache。除此之外，token上限还会影响流水线气泡和系统吞吐量。

于是本文使用了Vidur分析器和模拟器来确定在特定部署场景下最大化系统容量的token上限。

系统实现

Sarathi-Serve 推理系统基于 vLLM 框架进行实现，添加了对 FlashAttention v2 和 FlashInfer 的分块预填充的支持。并且使用 NCCL 进行流水线并行和张量并行的通信。