Sarathi-Serve 是一个高效大型语言模型（LLM）推理调度器，旨在解决LLM推理中吞吐量和延迟之间的权衡问题。通过引入“分块预填充（chunked-prefills）”和“无停顿调度（stall-free scheduling）”技术，Sarathi-Serve能够在保持低延迟的同时显著提高推理吞吐量。

论文：(OSDI 2024)Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve

代码：microsoft/sarathi-serve: A low-latency & high-throughput serving engine for LLMs

该文章参考自：https://zhuanlan.zhihu.com/p/12679786211

背景

LLM服务特性

LLM推理分为两个阶段：

Prefill（预填充）：处理输入提示并生成首个输出令牌，计算密集但延迟高。
Decode（解码）：逐个生成后续令牌，延迟低但计算利用率低。

当前的LLM推理调度器大致可以分为两类，即根据它们在批处理请求时如何调度预填充和解码阶段，分为预填充优先和解码优先：

传统的请求级批处理系统，如FasterTransformer，采用解码优先的调度策略。

只有批处理中的所有请求都完成了它们的解码阶段后，该批处理才算完成，即只要有一个或多个请求正在进行解码，就不会调度新的预填充。
现有调度策略（如vLLM、Orca）引入迭代级调度，使用连续批处理技术。

每当GPU内存可用时，优先调度预填充阶段的请求。预填充优先的调度器具有更好的吞吐量，因为这样允许后续的解码以高批次大小运行。然而，优先处理预填充会干扰正在进行的解码，导致高延迟（生成停滞）。

总的来说，现有LLM服务存在以下问题：

吞吐量与延迟的冲突：优先处理Prefill（提高吞吐量）会导致解码阶段的高延迟（生成停滞），而优先处理Decode（降低延迟）则会牺牲吞吐量。
管道并行中的气泡：混合Prefill和Decode的批次因计算时间差异导致GPU资源浪费。
长提示处理效率低：长输入提示的Prefill阶段耗时过长，加剧延迟波动。

存在的问题与挑战

预填充和解码阶段的成本分析。

LLM推理的两个阶段——预填充和解码——表现出截然相反的行为，其中批量处理可以极大地提高解码阶段的吞吐量，但对预填充吞吐量几乎没有影响；序列长度极大地影响预填充的时间，但批量大小几乎不影响解码的延迟。

解码批处理在内存受限的状态下运行，导致计算未得到充分利用。这意味着可以在解码批处理的同时处理更多的令牌，而不会显著增加其延迟。
优化线性操作对于提高大型语言模型（LLM）的推理效率至关重要。

从图中我们可以看到，预填充时间随着序列长度增加而近乎二次增长，且线性操作占用了大部分的运行时间成本。虽然注意力成本随着序列长度的增加而呈二次增长，但在高序列长度下，线性操作仍然贡献了超过80%的总时间。
预填充和解码阶段的计算特性是不同的。

在上图（左）中，用（计算量/访存量）衡量两个阶段的计算强度，Decode处于访存受限区域，Prifill处于计算受限区域，最理想的是平衡点是操作的算术强度与设备的FLOPS-to-Bandwidth比率相匹配。

右图显示了LLaMA2-70B中线性层计算在一次迭代中的总执行时间随token数量的变化。在开始时，当批次处于受内存限制的状态时，执行时间仅略有增加，但随后随着批次变为受计算限制的状态，执行时间呈线性增长。

吞吐量与延迟的权衡。

vLLM和Orca都是prefill优先，vLLM无法混合pd阶段的batch，Orca通过线性层批处理做到了pd混合batch，但由于prefill阶段的Attention算子耗时长，还是会拖累TBT。FastTransformer是decoding优先，无法混合pd阶段的batch，stall了prefill阶段，导致TTFT大。

当今最先进的系统使用预填充优先级调度，批越大，吞吐量越高，延迟越高，要根据所需的SLO在吞吐量和延迟之间进行权衡。
流水线气泡。

推理过程中存在的三种类型的气泡：
1. 由于连续两个微批次中prefill令牌的数量不同而产生；
2. 由于prefill和decode阶段的计算时间不同而产生；
3. 由于微批次之间decode计算时间的差异而产生，因为注意力成本取决于累积的上下文长度（KV缓存的大小），并在不同请求之间变化。
总的来说，是由于微批之间的计算不均匀导致，本质上主要还是因为无法完美耦合prefill和decode两个计算、调度特性的不同的阶段。

解决方案

分块预填充

允许在多个迭代中以小块形式计算prefill。基于前面的挑战可以发现，适度序列长度的prefill请求可以有效地使GPU计算能力达到饱和，那么利用这一机制形成具有适当token数量的批次，将序列prefill的一部分加入到decode批中来，以充分利用计算潜力，同时不违反TBT（总批处理时间）服务级别目标。

无停顿批处理

Sarathi-Serve调度器是一个迭代级别的调度器，它利用chunked prefill和预填充与解码的合并来提高吞吐量，同时最小化延迟，作者称这种方法为无停顿批处理。

Sarathi-Serve首先根据用户指定的服务级别目标（SLO）计算每批可执行的最大令牌数预算。（详见下一小节）

算法流程：

在每个调度迭代中，首先将所有正在decode阶段的请求放入批次中（第6-8行）
对于未完成预填充的请求，分块加入（第9-12行）
只有在所有正在运行的请求都被容纳后，才接受新请求（第13-20行）

注意：在向批次添加预填充请求时，要计算在该批次的剩余token预算中可容纳的最大块大小（第11、15行）

通过限制每次迭代的计算负载，无停顿批处理确保解码不会因为并行的预填充块而经历生成停顿。

确定token预算

Token预算的确定需平衡两个相互制约的因素：

延迟目标（TBT SLO）：较小的Token预算（分块更细）可降低单次迭代延迟，但可能导致：
1. GPU利用率下降，频繁分块增加调度开销；
2. KV缓存重复访问：每个分块需访问之前所有分块的KV缓存，导致显存读取次数增加（例如，分块数为N时，首个分块的KV缓存被加载N−1次）。
分块预填充开销：
1. Tile-Quantization效应：GPU矩阵乘法（Matmul）的硬件优化要求分块大小与GPU的Tile尺寸对齐（如256）。若分块大小不匹配（如257），计算时间可能增加32%。
2. 管道并行气泡：较大的分块导致批次间运行时间差异大，产生GPU闲置（气泡）；过小的分块则因算术强度低和固定开销（如内核启动）降低效率。

优化策略：

通过分析不同分块大小的性能（如预填充时间、解码延迟），找到满足TBT SLO的最大Token预算。
根据GPU的Tile尺寸调整分块大小，避免计算浪费。

文中没有具体说明优化的策略，而是说使用 Vidur（LLM推理性能模拟器）进行场景化分析，结合模型、硬件和并行策略（如TP/PP）动态优化Token预算。

算法实现

基于 vLLM 进行扩展和优化。通过使用 FlashAttention v2 和 FlashInfer 的内核，Sarathi-Serve 支持分块预填充。

评估

实验部分评估了Sarathi-Serve在不同模型和硬件配置下的性能，包括Mistral-7B、Yi-34B、LLaMA2-70B和Falcon-180B模型，以及单个A100 GPU、两个A100 GPU（TP2并行）、8个A40 GPU（TP4和PP2并行）等硬件配置。实验使用了两个数据集：openchat_sharegpt4和arxiv_summarization，分别代表多轮对话和科学文献摘要生成的任务。

关键结论

吞吐量提升：Sarathi-Serve在Mistral-7B模型上实现了2.6倍的服务容量提升，在Yi-34B模型上实现了3.7倍的提升（与vLLM相比）。在Falcon-180B模型上，使用流水线并行时，Sarathi-Serve提供了高达5.6倍的端到端服务容量提升。
延迟优化：在Yi-34B模型上，与不使用分块预填充的混合批次相比，Sarathi-Serve的延迟增加仅为25%，而使用完整预填充的混合批次延迟增加了28.3倍。
流水线并行优化：Sarathi-Serve通过创建计算需求均匀的混合批次，减少了流水线并行中的气泡，从而提高了GPU利用率，使得在普通以太网连接的多节点部署中也能高效运行。