LoongServe: Efficiently Serving Long-Context Large Language Models with Elastic Sequence Parallelism

发表时间: 2024-04 · arXiv:2404.09526 (SOSP 2024)

原文: https://arxiv.org/abs/2404.09526

作者/机构：
Bingyang Wu, Shengyu Liu, Yinmin Zhong, Xuanzhe Liu, Xin Jin (School of Computer Science Peking University)
Peng Sun (Shanghai AI Lab)

速读

一句话结论 针对长上下文大语言模型推理中资源需求剧烈波动的问题，提出了一种弹性序列并行（Elastic Sequence Parallelism, ESP）机制并构建了 LoongServe 系统，通过在迭代粒度上实时动态调整并行度，在长文本数据集上实现了最高 5.81 倍的吞吐量提升。

要解决什么问题 随着大语言模型支持的上下文窗口越来越长（如 1M tokens），推理系统面临极端的资源需求波动。这种波动体现在两个维度：一是不同请求的输入长度差异巨大，导致显存和算力需求相差悬殊；二是同一个请求在 Prefill（预填充，计算密集型）和 Decoding（解码，访存密集型）两个阶段的资源特征完全不同。现有的系统卡在静态并行策略上。如果采用固定的张量并行或序列并行，无法同时兼顾两个阶段的效率；如果采用将 Prefill 和 Decoding 分离到不同 GPU 组的路线（如 DistServe），在阶段切换时需要跨节点迁移海量的 KV Cache（例如 7B 模型处理 1M 长度单请求会产生 488GB 的 KV Cache），导致极高的通信延迟。此外，现有系统通常要求单个请求的 KV Cache 必须连续存放在同一个实例中（局部性约束），这使得即使集群总显存有空余，也无法拼凑起来服务长文本请求，造成严重的显存碎片化，甚至导致长文本请求直接因为显存分配失败而被拒绝。

怎么做的 核心思路是引入弹性序列并行（ESP），打破静态配置，在每次迭代（Iteration）级别动态决定每个请求批次的并行度（DoP），并实现跨实例的细粒度 KV Cache 统一管理。为了让弹性伸缩不产生致命的通信开销，LoongServe 设计了两个关键机制。第一是用于 Prefill 阶段结束后的主动缩容（Proactive Scale-down）。传统做法是算完再搬运数据，而 LoongServe 利用序列并行在计算注意力时本就需要环形通信传递 KV 数据的特性，在 Prefill 计算的过程中，直接让目标缩容组的实例把路过的 KV 张量按需存入自己的显存池。这样 Prefill 算完时，数据已经就位，实现了零额外开销的缩容。第二是用于 Decoding 阶段的多主节点分布式解码（Multi-master Distributed Decoding）扩容机制。当解码阶段显存不足或算力遇到瓶颈时，系统直接加入新实例而不迁移已有的 KV Cache。不同实例作为不同请求的“主节点”各自保存新生成的 KV 张量，彻底打破了局部性约束，并在交换 Query 张量时将通信与本地注意力计算重叠。为了在几十毫秒的调度窗口内找出最优解，系统采用了一个四步多项式时间调度算法，其中利用解析模型结合历史数据来预估执行时间：

$$ T_p(R) = \alpha_p + \beta_p \sum_{r \in R} r.\text{input\_len} + \gamma_p \sum_{r \in R} r.\text{input\_len}^2 $$

调度器会基于该公式，利用动态规划算法将长度相近的请求打包，并为每个批次分配合适的实例和并行度，从而在算力效率和显存占用之间取得最佳平衡。

效果如何 实验在 8 张 A800 80GB GPU 组成的单机以及 16 卡双机环境下进行，模型选用 LWM-1M-Text（基于 Llama-2-7B 架构，支持 1M 上下文），测试数据混合了 ShareGPT、L-Eval 和 LV-Eval 等真实长文本数据集。对比基线涵盖了三种主流路线：代表静态张量并行的 vLLM、代表长文本分块预填充（Chunked Prefill）路线的 DeepSpeed-MII 和 LightLLM w/ SplitFuse，以及代表 Prefill-Decoding 分离路线的 DistServe。在严格的延迟服务等级协议（SLO）下，LoongServe 的吞吐量比分块预填充路线最高提升了 3.85 倍，比分离路线最高提升了 5.81 倍。特别是在超长文本（如 LV-Eval）测试中，DistServe 因为静态显存切分导致单阶段显存不足直接触发 OOM，而 LoongServe 凭借全局 KV Cache 池顺利完成推理。在双机多节点测试中，LoongServe 依然保持了极高的扩展性。作者也指出了方法的局限性：在 Decoding 阶段如果 Batch Size 较小，扩容引入的通信和同步开销会相对明显（约 10% 的额外延迟），但全局调度器会动态评估收益，避免在小 Batch 时触发无效扩容。

主要贡献

随着大型语言模型（LLMs）的上下文窗口迅速增加（如Anthropic的Claude-3、Google的Gemini-1.5以及UC Berkeley的LWM均支持1M上下文窗口），不同请求之间以及同一请求的不同阶段（预填充阶段和解码阶段）在资源使用上产生了巨大的差异。受限于静态并行策略，现有的LLM服务系统无法有效利用底层资源来服务不同阶段的可变长度请求。具体而言，服务长上下文LLMs在GPU计算和GPU内存方面带来了重大挑战。例如，当使用LWM模型服务一个输入长度为1M token的单个请求时，仅键值缓存（KV cache）的GPU内存消耗就可达488GB，远超当前最先进GPU的内存容量；同时，注意力机制的计算复杂度与输入序列长度呈平方关系，使得长序列处理的计算密集度极高。

为了解决这一问题，本文提出了一种新的并行范式——弹性序列并行（Elastic Sequence Parallelism, ESP），以弹性地适应不同请求和阶段之间的差异。基于ESP，本文设计并构建了LoongServe，这是一个LLM推理服务系统，其核心贡献包括：
1. 实时弹性调整并行度：通过在实时中弹性调整并行度（Degree of Parallelism, DoP）来提高计算效率。
2. 优化通信效率：通过减少键值缓存（KV cache）的迁移开销，并将部分解码通信与计算重叠，从而提高通信效率。
3. 消除GPU内存碎片：通过在实例之间管理单token粒度的键值缓存，消除组隔离带来的GPU内存碎片，提高GPU内存效率。
4. 全面的实验验证：在多种真实的混合数据集下进行评估，结果表明，与分块预填充（chunked prefill）相比，LoongServe的吞吐量提高了高达 $3.85 \times$；与预填充-解码分离（prefill-decoding disaggregation）架构相比，吞吐量提高了高达 $5.81 \times$。

背景知识与设计动机

LLM推理过程：大多数流行的LLMs采用Transformer架构【53, Attention is all you need, 2017, NIPS】。模型通常由堆叠的Transformer层组成，每层包含一个注意力层（Attention）和一个前馈神经网络层（FFN）。为了避免冗余计算，LLM服务系统会缓存token的中间状态（即KV cache），并将其重用于未来的token生成。这种优化将整个生成过程分为两个阶段：预填充阶段（prefill phase）和解码阶段（decoding phase）。预填充阶段在单次迭代中处理所有输入token以构建KV cache并生成第一个输出token，因此是计算密集型的；而解码阶段只需为新生成的输出token计算KV cache，计算量相对较轻。
不同阶段不同长度请求的可扩展性

现有LLM服务系统的局限：为了加速LLM推理，现有解决方案采用了高效的GPU内核实现（如Flash Attention【15, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022, NIPS】和Flash Decoding【14, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023, arXiv】），并利用模型并行（如张量并行【46, Megatron-LM: Training Multi-billion Parameter Language Models using Model Parallelism, 2019, arXiv】）来跨GPU划分模型参数。然而，模型并行的程度必须在启动系统之前静态确定。为了减轻长上下文的影响，分块预填充（chunked prefill）【20, DeepSpeed-FastGen, 2024, arXiv】【41, Splitwise, 2023, arXiv】将长上下文分割成块并与解码阶段逐块处理，但这仍然会导致两个阶段之间的干扰【61, DistServe: Disaggregating Prefill and Decoding..., 2024, OSDI】。为了避免干扰，预填充-解码分离架构【61】将两个阶段分离到不同的GPU组中，但这会导致由于频繁迁移产生的高通信开销、由于资源不匹配导致的GPU计算低效，以及由于组间隔离导致的GPU内存碎片。

序列并行及其局限性：为了加速长上下文LLM的训练，许多系统提出了序列并行【12, Striped Attention: Faster Ring Attention for Causal Transformers, 2023, arXiv】。在本文中，作者将Striped Attention扩展到了LLM服务场景。输入序列在处理前被排列、分成多个段，然后分发到不同的实例。除了注意力层外，各实例在没有通信的情况下执行。在注意力层，每个实例同时使用其局部查询 $q_i$ 和键值 $kv_j$ 张量计算注意力，并将键值张量发送给其相邻实例。序列并行与MHA、MQA和GQA等流行注意力机制兼容，其计算复杂度与张量并行相同，且消耗较少的GPU内存用于缓冲区激活。然而，它不能直接应用于LLM服务场景，因为它只支持预填充阶段，且在训练场景中每阶段的并行度是固定的，而服务系统需要处理高度动态的推理流量。
预填充阶段的序列并行
固定序列并行与张量并行的比较

动机与挑战：尽管现有的LLM服务系统支持长上下文LLMs，但其静态特性与LLM服务中的动态性之间存在固有的不匹配。首先，随着上下文窗口的增加，不同输入长度请求在计算和GPU内存消耗上的资源需求差异巨大（例如服务1M上下文模型时内存差异可达1,000,000倍）。其次，即使对于单个请求，预填充和解码阶段的资源需求也大相径庭。预填充阶段是计算密集型的，增加DoP可以显著加速；而解码阶段增加DoP可能因额外通信开销导致性能提升微乎其微。此外，现有解决方案【8, 20, 21, 41, 61】由于局部性约束（请求的整个或大部分KV cache必须驻留在单个实例中）导致了严重的GPU碎片问题。为了从根本上解决这些问题，本文提出了弹性序列并行（ESP），通过增加对解码阶段的支持、灵活管理跨阶段的KV cache，并动态分配请求的输入token到各个实例，允许在不重新划分LLM参数的情况下调整DoP。然而，释放ESP的潜力面临两大挑战：一是弹性扩展的通信开销过高可能会抵消灵活资源分配带来的好处；二是动态负载形成了一个复杂的调度空间，需要在几十毫秒的迭代粒度内做出高效的调度决策。
基于组的策略导致的KV cache碎片化

方法细节与系统设计

3 LoongServe Overview

系统架构概览：为了解决上述挑战，作者设计并构建了一个分布式LLM服务系统LoongServe，以充分释放ESP的潜力。LoongServe由一组弹性实例（elastic instances）和一个全局管理器（global manager）组成。这些弹性实例可以动态地将自身组织成一组不相交的ESP组，以不同可配置的并行度（DoP）并行处理请求批次。它们还支持高效的弹性向上和向下扩展，且没有额外开销，以适应请求不断变化的资源需求。LoongServe弹性实例的GPU内存还构成了一个统一的分布式键值缓存池（unified distributed key-value cache pool），可用于在弹性实例之间以单token为粒度灵活存储请求的键值张量，从而减少GPU内存碎片。LoongServe全局管理器负责利用全局视图管理请求、弹性实例和统一的分布式键值缓存池。
LoongServe系统概览

动态调度与执行流程：在每次迭代中，基于缩放信息库（Scaling Information Base, SIB）的分析结果，全局管理器会动态调整现有批次的DoP、弹性实例的分组策略、新到达请求的批处理和分发策略，以及键值缓存的放置策略，以实时提高吞吐量并降低请求延迟。在此过程中，伴随全局管理器的LoongServe分发器（dispatcher）按照全局管理器的要求，将新到达的请求分发到一组特定的弹性实例。同时，给定全局管理器生成的DoP和缩放计划，弹性控制器（elasticity controller）命令弹性实例更新其配置，以形成相应的ESP组并并行处理请求。LLM推理生成的键值缓存被放置在缩放计划指定的位置。全局管理器监控请求的进度、弹性实例的资源使用情况以及键值缓存池，以更新其决策。

4 LoongServe Elastic Instances

弹性实例与生命周期管理：弹性实例是LoongServe中最小的独立执行单元。每个实例维护一个模型权重的副本，并在等量的GPU上采用统一的模型并行策略。在每次迭代之前，全局管理器动态地将弹性实例分配到多个并行组中，每个组处理特定批次的计算，并具有不同的并行度（DoP）。并行组中的实例数量影响其相应的DoP。弹性实例还可能接收来自全局管理器的潜在弹性缩放计划，这需要重新分配实例的并行组。如果是这样，在完成当前迭代的计算后，实例需要实现弹性缩放，以确保后续计算可以按照新定义的DoP进行，这意味着必要的键值张量已经存在于每个实例的KV Cache池中。

请求生命周期与缩放时机：请求的生命周期包含了预填充和解码阶段，其中由于预填充阶段的计算复杂度远高于解码阶段，因此在预填充阶段之后总是需要缩小批次的并行度（scale down）。随着生成的token越来越多，解码阶段的计算复杂度增加，并行组中的GPU内存可能会被键值缓存填满，这导致需要扩大并行组的并行度（scale up）。全局管理器还可以选择性地缩小解码批次的规模，为预填充批次留出更多资源。为了减轻缩放机制的开销，作者设计了一套零开销的ESP机制，允许在预填充阶段向下弹性缩放和在解码阶段向上缩放，而无需任何额外的键值张量迁移。
请求的生命周期

4.1 Elastic Scale-down

向下缩放的必要性：在批次完成预填充阶段后，其解码阶段的计算需求显著下降。在这种情况下，通常有利于缩小其并行组的规模，以释放资源给其他批次。具体而言，对于一个DoP为 $d$ 的并行组 $R$，弹性向下缩放机制需要将该组缩小为一个新的DoP为 $d^\prime$ 的并行组 $R^\prime$，其中 $d^\prime < d$。主要的挑战是确保并行组 $R$ 中请求的所有键值张量都能高效地转移到新的并行组 $R^\prime$ 中。

现有解决方案的缺陷：现有的实践【21, 41, 61】使用反应式迁移（reactive migration），即在预填充阶段之后将键值张量从并行组 $R$ 迁移到新的并行组 $R^\prime$。然而，这种迁移开销是不可忽视的，并且随着请求序列长度的增加呈线性增长。此外，反应式迁移还存在GPU碎片问题。它要求 $R$ 中每个实例至少有 $O(blsh/d)$ 的未使用GPU内存用于键值张量，以容纳新生成的键值张量，然后才能进行迁移（其中 $b, l, s, h$ 分别为批大小、层数、序列长度和隐藏维度）。这可能导致即使总可用插槽足够，单个实例也会出现内存溢出（OOM）错误。

主动迁移（Proactive Migration）机制：为了消除迁移开销，作者提出了一种新的无额外通信开销的向下缩放机制，称为主动迁移。核心观察是，在预填充阶段，序列并行本质上在并行组中循环键值张量。作者没有在预填充阶段之后进行反应式迁移，而是在预填充阶段选择性地将键值张量保留在新并行组 $R^\prime$ 实例的KV缓存池中，从而实现零开销的弹性向下缩放。

主动迁移的工作原理：以包含三个实例的并行组 $R$（DoP=3）为例，如果指令要求将其缩小为仅包含前两个实例的新并行组 $R^\prime$，并且前四个token存储在实例1中，其余token存储在实例2中。在这种情况下，当从相邻实例接收键值张量时，除了计算注意力和将其发送到下一个实例外，前两个实例还会根据指令要求选择性地将键值张量保存到其键值缓存池中。在预填充阶段计算结束后，请求的键值张量就已经存在于 $R^\prime$ 的KV缓存池中了。

主动迁移的优势：与反应式迁移相比，主动迁移不产生额外的迁移开销，因为它重用了现有的通信结果。此外，主动迁移消除了反应式迁移施加的内存限制，并允许根据每个实例的内存可用性制定任何token级别的KV Cache分配计划，而不会导致计算负载不平衡。它还重用了序列并行中现有的缓冲空间，避免了额外的GPU内存分配。因为计算是逐层进行的，该缓冲区只需要临时存储单层的键值张量，其大小即 $O(bsh/d)$，甚至可能小于纯模型并行所需的 $O(bsh)$。
预填充阶段的弹性向下缩放

4.2 Elastic Scale-up

向上缩放的必要性：由于解码阶段的计算模式较轻量，解码批次通常以较小的DoP执行以提高整体效率。然而，随着解码的进行，生成的键值张量的大小可能会超过并行组中KV缓存池的容量，这就需要添加新的实例来扩展容量。此外，如果批大小足够大，解码阶段的计算可能会受到计算限制（compute-bound）。在这些情况下，使用更多的GPU可以减少延迟，但也需要高效的向上缩放操作。在扩大并行组时，主要的挑战是确保新添加的实例能够高效地参与正在进行的计算，而不会产生额外的开销。

现有解决方案的缺陷：之前的工作如张量并行【46】和FlashDecoding【14】仅支持在单个实例内的多个GPU上进行分布式解码计算。当实例的资源（如GPU内存）不足时，它们必须将批次中的部分请求迁移到另一个实例，然后在不同的实例中并行处理它们。然而，将请求的整个键值缓存迁移到另一个实例会产生巨大的迁移开销，甚至高于解码步骤本身。此外，它要求请求的整个或大部分键值张量必须存储在单个实例中，导致内存碎片问题。

单主分布式解码（Single-master distributed decoding）：作者首先将序列并行扩展到解码阶段，称为单主分布式解码，允许多个实例参与单个批次的解码计算。在并行组中，指定一个实例为主实例（master instance），负责驱动计算过程。在每个注意力层，主实例首先计算查询和键值张量，将键值张量保存到其本地KVCache池中，并将查询张量发送给并行组中的其他实例。所有实例并行执行局部注意力计算，并将结果发送回主实例。之后，主实例继续计算其他局部层（如FFN层），并开始执行下一个注意力层。只要单个实例有足够的内存来存储新生成的键值张量，就可以通过将其指定为主实例，在多个实例之间处理整个批次。当扩大并行组时，全局管理器只需将新实例添加到并行组中，并指示它们执行，而无需迁移现有的键值张量。

多主分布式解码（Multi-master distributed decoding）：然而，单主方法在批大小很大时存在局限性。在内存管理方面，它要求主实例中未使用的插槽足以存储下一次迭代中生成的键值张量，这会导致内存碎片问题。在计算方面，由于FFN等局部层都在主实例上执行，当解码阶段变得受计算限制时，其性能受限于主实例中的计算资源。为了解决这些问题，作者进一步将其扩展为多主分布式解码。在一个并行组中包含多个主实例。不同的主实例负责批次中的不同请求。它们将相应的键值张量保存到本地KV缓存池中，以打破内存碎片问题，并在多个主实例之间并行化局部层计算，以提高计算效率。当主实例交换查询张量时，它们之间的通信可以进一步与其主导请求的局部注意力计算重叠。
解码阶段的弹性向上缩放

5 LoongServe Global Manager

调度空间的复杂性：借助高效的弹性缩放机制来改变DoP，全局管理器负责利用它们在弹性实例之间高效地调度请求和键值张量。系统中可能存在挂起队列 $P = \{r_1, r_2, ..., r_{n_P}\}$ 中的新到达请求，以及处于解码阶段的请求批次集合 $B = \{B_1, B_2, ..., B_{n_B}\}$，其中每个解码批次 $B_i$ 与现有的并行组 $G_i$ 相关联。弹性实例要么由于上一次迭代中的向下缩放操作而处于空闲状态，要么正在执行解码批次。在每次迭代中，全局管理器需要从 $P$ 中分发一些请求在当前迭代中执行预填充阶段，为它们分配弹性实例，决定批处理策略，并更改现有解码批次的状态。所有这些方面，包括批处理策略、每个批次的DoP以及键值张量的放置，都是可配置的，并且在不同迭代中可能不同，形成了一个与请求和弹性实例数量呈指数关系的调度空间。
灵活调度空间的图示

四步调度算法：全局管理器面临的主要挑战是在几十毫秒的严格延迟限制内，从这个复杂的调度空间实时生成高效的调度计划。为了解决这些问题，作者提出了一种可扩展的四步调度算法。核心思想是将这个调度问题解耦为四个子问题：分发（dispatching）、弹性实例分配（elastic instance allocation）、批处理（batching）和弹性缩放计划生成（elastic scaling plan generation）。在每个步骤中，LoongServe全局管理器在多项式时间内为相应方面生成高效的计划，然后将这些计划组合成最终的调度计划。

5.1 Dispatching

分发策略与内存约束：分发步骤是从挂起队列 $P$ 中选择一个请求子集 $R_p$，在当前迭代中执行预填充阶段。在此步骤中，全局管理器考虑GPU计算和GPU内存的资源可用性。与之前的工作【28, 58】一样，全局管理器按照先到先得（FCFS）的顺序扫描 $P$。在GPU内存方面，键值张量的内存消耗是首要考虑因素。除非有足够的未使用插槽，否则全局管理器不会将请求添加到 $R_p$ 中。此外，由于驱逐长序列请求会产生巨大的重新计算开销，全局管理器会通过考虑基于当前状态和用户提供的请求最大序列长度的最大未来键值缓存消耗来避免未来的驱逐和重新计算【54, Fast Distributed Inference Serving..., 2023, arXiv】。如果可能触发驱逐，则不添加该请求。

计算约束与临界点评估：在GPU计算方面，首要考虑的是分发 $R_p$ 的效率及其对处于解码阶段的现有请求 $B$ 的影响。全局管理器使用分析模型和来自SIB（§5.5）的分析结果，以被占用实例 $E_p$（初始化为所有空闲实例的子集）上的迭代时间 $T(R_p, E_p)$ 来衡量它们。首先，存在一个临界点，即 $R_p$ 从内存受限转变为计算受限。在该点之前，向 $R_p$ 添加更多请求可提高GPU计算效率；在该点之后，仅延长执行时间而效率提升微乎其微。全局管理器通过分析预填充批次受内存限制的迭代时间的上限来估计此点，当 $R_p$ 的迭代时间超过此临界点时，停止向 $R_p$ 添加请求。

抢占成本与收益计算：其次，如果新请求利用了被 $B$ 占用的实例中的一些未使用键值缓存插槽，向 $R_p$ 添加更多请求可能会干扰现有的某些批次 $B_p \subseteq B$。为简化此问题，保守地考虑最坏情况，即 $R_p$ 可能抢占 $B_p$。在这种情况下，对于每个 $B_{p,i}$，其并行组 $G_{p,i}$ 中实例的未使用键值插槽可用于向 $R_p$ 添加额外的新请求子集 $R_{p,i}^\prime$。全局管理器分析执行 $R_{p,i}^\prime$ 的性能收益和抢占 $B_{p,i}$ 的成本，以决定是否将 $R_{p,i}^\prime$ 添加到 $R_p$ 并扩展 $E_p$。性能成本（即抢占对输出token延迟的影响）公式如下：

$$Cost = \sum_{r \in B_{p,i}} \frac{T(R_p \cup R_{p,i}^\prime, E_p \cup G_{p,i})}{r.\mathrm{output\_len}}$$


其中输出长度是现有输出token的数量，这类似于之前工作中基于输出长度的等待时间估计【36, 42, 55】。

性能收益公式与复杂度：至于性能收益，其被公式化为对 $R_{p,i}^\prime$ 的输入token延迟的影响：

$$Gain = \sum_{r \in R_{p,i}^\prime} \frac{(\mathrm{AvgLat_d} - \min(B_{p,i}.\mathrm{exec\_time}))^+}{r.\mathrm{input\_len}}$$


在这个方程中，$\mathrm{AvgLat_d}$ 是解码阶段已完成请求的平均执行时间，$\min(B_{p,i}.\mathrm{exec\_time})$ 是解码阶段请求的已执行时间。它们之间的减法估计了在最坏情况下 $R_{p,i}^\prime$ 必须等待解码批次多长时间。如果收益大于成本，全局管理器将 $R_{p,i}^\prime$ 添加到 $R_p$ 并将 $G_{p,i}$ 添加到 $E_p$。此步骤的复杂度为 $O(n_B)$。

5.2 Elastic Instance Allocation

实例分配与预抢占处理：在生成要执行的确切 $R_p$ 后，全局管理器需要在此步骤中为它们决定实际的弹性实例分配。在此步骤中，主要关注的是在最大化GPU计算效率的同时，减轻预填充阶段（即 $R_p$）和解码阶段（即 $B$）之间的干扰。全局管理器首先将空闲实例分配给 $R_p$。如果空闲实例的未使用键值缓存插槽不足，$R_p$ 可以抢占几个具有最多未使用键值缓存插槽的实例，以获得足够的插槽。为了避免抢占，全局管理器会尽可能将抢占实例中现有的键值张量迁移到其他活动实例。因此，$R_p$ 可以获得具有足够键值缓存插槽的弹性实例 $E_p$。

进一步优化实例分配：然而，通过向计算密集的预填充阶段分配更多弹性实例，仍然有可能实现更好的性能。为此，全局管理器重复考虑是否将具有最少已用键值缓存插槽的弹性实例 $e_{min}$ 分配给 $R_p$。约束条件是使用 $e_{min}$ 的批次可以将其键值张量迁移到解码阶段的其他实例。在这种情况下，输入token延迟的减少（收益）如下：

$$Gain = \sum_{r \in R_p} \frac{T(R_p, E_{\hat{p}}) - T(R_p, E_{\hat{p}} \cup e_{min})}{r.\mathrm{input\_len}}$$

迁移成本与复杂度：但这也会产生迁移开销（成本），公式如下：

$$Cost = \sum_{r \in R_p} \frac{V(e_{min})}{\mathrm{avg\_bandwidth} \cdot r.\mathrm{input\_len}}$$


在这个方程中，$V(e_{min})$ 是 $e_{min}$ 中现有键值张量的体积，$\mathrm{avg\_bandwidth}$ 是 $e_{min}$ 和目标实例之间的平均带宽。目标实例始终是具有最多未使用键值缓存插槽的实例。全局管理器重复将 $e_{min}$ 分配给 $R_p$，直到收益小于成本。此步骤的复杂度为 $O(m)$，其中 $m$ 是弹性实例的数量。

5.3 Batching

动态规划批处理问题公式化：在决定了 $R_p$ 和 $E_p$ 之后，此步骤优化 $R_p$ 在 $E_p$ 上的批处理策略，以进一步最小化预填充阶段的延迟。在此步骤中，主要关注的是为具有不同序列长度的请求分配不同的DoP。作者将此批处理问题公式化为动态规划（DP）问题。优化目标是最小化 $R_p$ 在 $E_p$ 上的输入延迟。核心见解是具有相似序列长度的请求具有相似的特征，应该被批处理在一起。因此，全局管理器首先根据请求的序列长度按降序对请求进行排序。分配的弹性实例也根据它们的位置和未使用键值缓存插槽的数量按升序排序。设 $f[i][k]$ 为使用前 $k$ 个弹性实例时，前 $i$ 个请求的最小输入延迟。DP方程可以公式化为：

$$f[i][k] = \min_{\stackrel{0 < j \leq i, 0 < l \leq k,}{D[j, i] \leq V[l, k]}} (f[j][l] + T(R[j, i], E[l, k]))$$

前缀和与回溯生成计划：方程中的 $D[j, i]$ 是从 $j$ 到 $i$ 的请求的token数，$V[l, k]$ 是从 $l$ 到 $k$ 的弹性实例的未使用键值缓存插槽数。区间内这些值的总和可以通过预先维护一个前缀和数组来计算。$T(R[j, i], E[l, k])$ 是在使用从 $l$ 到 $k$ 的弹性实例时，从 $j$ 到 $i$ 的请求的输入延迟总和。所有请求的最小输入延迟总和可以在多项式时间内找到。在更新 $f[i][k]$ 时，全局管理器记录请求的最后拆分点 $split_{req}[i][k]$（即 $j$）和弹性实例的最后拆分点 $split_{ins}[i][k]$（即 $l$）。全局管理器使用它们进行回溯，以生成批处理计划和每个批次的相应DoP。

四边形不等式优化：如果天真地为所有 $0 < i \leq n$ 和 $0 < j \leq m$ 更新 $f[i][k]$，此DP算法的时间复杂度为 $O(|R_p|^2 \cdot |E_p|^2)$。然而，作者注意到 $split_{req}[i][k]$ 和 $split_{ins}[i][k]$ 具有单调性属性。因此，通过使用四边形不等式属性【57, Efficient dynamic programming using quadrangle inequalities, 1980, STOC】，此问题可以优化为 $O((|R_p| + |E_p|)^2)$，在实践中足够高效。

5.4 Elastic Scaling Plan Generation

主动向下缩放策略：如前所述，全局管理器还需要生成用于主动向下和向上缩放的弹性缩放计划。对于主动向下缩放，核心见解是解码阶段的扩展性较差。如图2所示，在大多数情况下，解码阶段请求的最小最佳DoP是相似的，并且小于其他情况。因此，在启动时将模型并行的程度设置为最小最佳DoP。在运行时，全局管理器只需将DoP向下缩放至请求的键值张量能够适应相应弹性实例的最小DoP。这对解码阶段的大多数请求来说是最佳的，即使对于具有更大最佳DoP的解码请求，这也是接近最佳的，因为将更多弹性实例留给持续时间更长的计算密集型预填充阶段更为有利。

向上缩放策略：对于向上缩放，当GPU计算或GPU内存不足时，全局管理器进行向上缩放。GPU计算不足是指解码阶段受到计算限制。因为FFN层首先成为计算瓶颈，且其复杂度与批大小相关，全局管理器使用预先分析的批大小阈值来检测它。只要多主解码能够减少内存碎片或执行时间，就会使用它。每个主节点生成的新键值张量数量被设置为尽可能均匀。

5.5 Optimizations

基于SIB的分析模型：不同场景下预填充阶段的迭代时间指导着决策的制定。由于包含不同输入长度和不同DoP的请求组合海量，无法仅通过预先配置来覆盖所有情况。因此，作者提出了一种分析模型来估计它们，公式如下：

$$\mathrm{T}_{\mathcal{P}}(R) = \alpha_{\mathcal{P}} + \beta_{\mathcal{P}} \cdot \sum_{r \in R} r.\mathrm{input\_len} + \gamma_{\mathcal{P}} \cdot \sum_{r \in R} r.\mathrm{input\_len}^2$$


在这个方程中，$\alpha_p$、$\beta_p$ 和 $\gamma_p$ 分别是捕获恒定开销、线性计算（如FFN层）和二次计算（如注意力层）的系数。它们是基于少量分析结果通过最小二乘法训练得出的。对于不同的并行策略，训练不同的系数。

系统实现细节与相关工作

实现细节：LoongServe包含约1.5万行基于C++、CUDA、Python和Triton【50】的代码，并重用了vLLM【28】和LightLLM【48】的一些组件。前端API类似于OpenAI API。全局管理器主要用Python实现，但核心逻辑（如批处理算法）用C++实现以加速循环函数。使用Ray【38】与弹性实例通信，并优化RPC参数以减少序列化开销。弹性实例在单token粒度上使用PagedAttention管理KV缓存池。作者调整了StripedAttention的块大小以跳过短序列中的冗余计算，并为解码阶段实现了带有额外参数的自定义Flash-Decoding。实例间通信基于NCCL，张量并行和序列并行使用不同的NCCL通信器。
相关工作对比：与现有的LLM服务系统（如SARATHI【8】、DeepSpeed-FastGen【20】、SplitWise【41】、DistServe【61】、Infinite-LLM【33】）相比，LoongServe提出了无额外开销的弹性缩放机制和ESP，以无局部性约束的方式处理动态工作负载。与训练中的序列并行和弹性训练相比，LoongServe专门针对服务场景，支持解码阶段、动态DoP和高效的KV缓存管理，并且不损害原始LLM的准确性。

实验设置

  • 模型：LWM-1M-Text模型（1百万token上下文窗口，架构同Llama-2-7B）。
  • 硬件配置：单节点评估使用一台配备8张NVIDIA A800 80GB GPU、128个CPU、2048 GB主机内存和4个200 Gbps InfiniBand网卡的服务器（GPU间NVLink带宽为400 GB/s）。多节点评估使用两台相同的服务器（共16张GPU）。
  • 软件配置：PyTorch 2.0.0, CUDA 12.2, OpenAI Triton 2.1.0, HuggingFace tokenizers 0.15.2。
  • 数据集：请求到达模式由泊松过程生成。输入/输出长度从真实数据集采样：ShareGPT（4-2.3K tokens）、L-Eval（2.7K-210.5K tokens）、LV-Eval（15.1K-497.3K tokens）以及上述数据集等概率混合的Mixed数据集。
  • 基线系统：vLLM（TP=8）、DeepSpeed-MII（动态SplitFuse，TP=8，仅在ShareGPT评估因长序列OOM）、LightLLM w/ SplitFuse（TP=8）、DistServe（预填充和解码各分配4个GPU，DoP=4）。LoongServe设置TP=2，ESP=4。
  • 评估指标：标准化每token延迟、标准化输入延迟、标准化输出延迟。设定延迟服务等级目标（SLO）为推理延迟的 $25 \times$，比较此SLO下的最大吞吐量。

实验结果

1. 端到端性能（End-to-End Performance）
* 实验内容：在四个真实数据集上比较LoongServe与四个基线系统的吞吐量和延迟。
* 实验结果：如图10所示，LoongServe的输出延迟显著优于其他基线。与vLLM相比，LoongServe在总吞吐量和输入吞吐量上提升了高达 $4.64 \times$ 和 $4.00 \times$。与DeepSpeed MII和LightLLM w/ SplitFuse相比，LoongServe的吞吐量分别提升了高达 $3.85 \times$ 和 $3.37 \times$。与DistServe（在长序列数据集上因4卡内存不足导致OOM）相比，LoongServe提升了高达 $5.81 \times$ 和 $3.58 \times$。
* 分析结论：由于LoongServe使用不同的弹性实例执行不同阶段，解码阶段得到了很好的保护；统一的分布式KV缓存池使其能有效处理长上下文请求，避免了OOM问题。
不同LLM服务系统在真实工作负载下的平均延迟

2. 多节点性能（Multi-Node Performance）
* 实验内容：在16-GPU集群上使用Mixed数据集评估多节点扩展性。LoongServe设置ESP=8。
* 实验结果：如图11所示，LoongServe在所有指标上均取得最佳性能。与vLLM相比，总吞吐量和输入吞吐量分别提升了高达 $1.86 \times$ 和 $1.72 \times$；与LightLLM w/ SplitFuse相比，分别提升了 $3.37 \times$ 和 $3.11 \times$，同时显著降低了输出延迟。
* 分析结论：LoongServe在多节点设置下扩展良好，能够通过为不同阶段的请求选择合适的DoP来避免短请求的不必要通信开销并扩大长请求的并行度。
多节点性能

3. 弹性序列并行（ESP）的消融实验
* 实验内容：在不同Zipf分布参数的序列长度下，比较LoongServe（动态ESP）与传统张量并行（无ESP, TP=8）、静态混合并行（无ESP, TP=2, SP=4）和复制并行（无ESP, TP=2 x 4）的P90 goodput。
* 实验结果：如图12所示，LoongServe在不同序列分布下将P90 goodput分别提升了 $2.33 \times$、$1.98 \times$ 和 $1.53 \times$。
* 分析结论：仅引入静态序列并行不足以应对动态推理工作负载。LoongServe的四步调度算法动态调整策略，证明了ESP对大多数请求都是有益的。
不同序列长度分布下的P90 goodput

4. 弹性向上缩放（Elastic Scale-up）的消融实验
* 实验内容：在ShareGPT数据集上比较开启与关闭弹性向上缩放的LoongServe性能，并统计缩放触发频率。
* 实验结果：如图13a所示，带有弹性向上缩放的LoongServe的P90 goodput比不带该功能的版本高出 $2.87 \times$。图13b显示，全局管理器平均每10秒触发7.12次弹性向上缩放操作。
* 分析结论：对于输入短而输出长的请求，随着输出长度的增加需要频繁向上缩放，证明了弹性向上缩放处理动态工作负载的必要性。
弹性向上缩放的消融实验

5. 弹性缩放机制的开销（Scaling Overhead）
* 实验内容：在不同批大小和输入长度下，测量带有和不带有缩放操作的批次前向传递时间。
* 实验结果：如图14所示，向下缩放仅引入了极小的开销（小于2%）。对于向上缩放，在大批大小下，将计算分布到更多实例显著减少了计算时间，使每次迭代延迟改善了 $2 \times$；在小批大小下，由于通信和同步，向上缩放开销较高，但仍可接受（小于10%）。
* 分析结论：LoongServe的缩放机制开销极低，且系统能够动态选择最佳策略。
弹性缩放机制的开销

6. 分析模型的准确性（Accuracy of Analytical Model）
* 实验内容：评估LoongServe分析模型在不同并行策略下的准确性。
* 实验结果：如图15所示，分析模型对具有不同序列长度和并行策略的请求批次实现了高准确度（偏差小于10%）。
* 分析结论：该模型足以可靠地指导LoongServe全局管理器的调度决策。
LoongServe分析模型的准确性

结论

为了在动态工作负载下服务长上下文LLMs，本文提出了弹性序列并行（ESP）并构建了LoongServe系统。该系统提供了一套无额外开销的弹性缩放机制以及针对ESP的可扩展调度算法。在多种真实数据集上的评估表明，与现有解决方案相比，LoongServe同时显著提高了预填充阶段和解码阶段的性能。这为未来在更复杂、更长上下文的生成式AI应用中实现高效的资源利用和低延迟服务奠定了坚实的基础。