LoongServe论文解读：prefill/decode分离、弹性并行、零KV Cache迁移

LoongServe 论文解读：prefill/decode 分离、弹性并行、零 KV Cache 迁移

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

论文提出了一种支持弹性分配的推理框架，通过引入弹性序列并行（Elastic Sequence Parallelism，简称 ESP）机制，动态地将 request 的 prefill 和 decode 阶段分配到 instance group 上。每个 group 可以根据负载的需求变化动态地 scale up 或者 scale down，并且没有 KV Cache 的迁移开销。

问题背景

Transformer LLM 推理过程分为两个阶段：prefill 和 decode。

• prefill 阶段：将用户输入的 prompts 生成 q、k、v，存入 KV Cache（为 decode 阶段缓存）。这一步计算并行好，是计算密集型 compute bound
• decode 阶段：由最新产生的 tokens 生成 q、k、v，计算它与之前所有 tokens 的 attention，这一步需要从 KV Cache 中读取前面所有 token 的 key、value，因此是内存密集型 memory bound。

在 long context 背景下，prefill 和 decode 阶段对计算和显存的需求非常不平衡，decode 阶段对 KV Cache 的容量需求是动态增长的、不能事先预测和分配（decode 阶段直到输出 end of sentence token 才停止）。如果把 prefill 和 decode 阶段混合部署，会造成 memory 等资源浪费，影响推理请求的吞吐量和延迟。因此，一种方法是将 prefill 和 decode 阶段分离，将 prefill 和 decode 部署在不同的 GPU 组。但是，prefill 和 decode 分离会带来新的问题：

1. 如何细粒度地调度请求
2. 请求从 prefill 到 decode 阶段可能需要迁移，迁移开销大
3. GPU 组是静态划分的，无法跨组协同导致 memory 碎片（比如遇到 memory 需求很大的超长 request，哪怕所有 GPU 组的 memory 余量加起来可以满足，也无法服务）

相关工作

vLLM 提出了 Paged Attention 算法，将 attention 算法产生的连续的 key value 向量按照 block 进行组织和管理，以减少显存碎片。vLLM 还借鉴操作系统当中的虚拟内存和分页思想优化 Transformer 模型推理中产生的 KeyValue Cache，大大提高了显存当中 KV Cache 的利用效率。但 vLLM 是基于单机层面上设计，只能将 GPU 中的 block swap 到单机的内存当中。

SplitWise、DistServe 和 TetriInfer 将 prefill 和 decode 分离到不同的 GPU 组中以避免干扰，但它们的静态并行性和分区策略并不灵活，无法处理动态工作负载。

Infinite-LLM 针对 long context 场景提出分布式 DistAttention，将 KV Cache 分割成 rblock，一个 node 可以借用别的 node 上的空闲显存。但它没有做 prefill/decode 分离，并且仍然需要周期性的 KV Cache 迁移来维持局部性，并且没有考虑不同请求之间或不同阶段之间的弹性资源需求。

设计思路

论文提出一种新的并行策略：弹性序列并行（Elastic Sequence Parallelism, ESP），并构建了一个分布式的大型语言模型（LLM）服务系统——LoongServe，以充分释放 ESP 的潜力。

LoongServe 由一组弹性实例和一个 global manager 组成。这些弹性实例可以动态地将自己组织成一组不相交的 ESP（Elastic Sequence Parallelism）组，每个 ESP 组并行处理请求批次。它们还支持高效的 scale up 和 scale down，并且不需要 KV Cache 的迁移。这样，弹性实例的 GPU 内存实际上形成了一个统一的分布式 KV Cache 缓存池，可以灵活地以 token 为粒度存储请求的键值张量，减少 GPU 内存碎片。Global manager 以 iteration 为粒度动态地调度请求、调整 ESP 组和管理分布式 KV Cache 缓存池，

Sequence Parallelism

为了 ESP 组内并行地处理请求，LoongServe 采用 Sequence Parallelism。在 attention 阶段，将 tokens 分到不同的实例，每个实例：

1. 计算自己分到的 tokens 的 q、k，计算 local attention。
2. 将自己的 k、v 传给下一个邻居实例。
3. 从上一个邻居实例接收到 k、v，更新 attention，将接收到的 k、v 传递给下一个邻居，直到接收到所有的 k。

下面看 LoongServe 是如何进行 ESP 组的 scale up 和 scale down 的。

Scale down

ESP 组内的计算需求下降时，可以进行 scale down，降低 DoP（degree of parallelism）。此时，需要将 KV Cache 集中到少数的实例。由于 sequence parallelism 机制，k、v 会在组间传递，实例只需要选择性地保存一部分 k、v，而不需要额外进行 KV Cache 迁移。

如下图所示，要将实例 1~3 组成的 ESP 缩成实例 1~2 时，若将 kv1~4 保存在实例 1，kv5~6 保存在实例 2，实例 1 只需要在组间传递 kv 的时候选择性地保存 kv2，实例 2 只需要选择性地保存 kv6。

Scale up

ESP 组内的计算或者显存需求上升时，可以进行 scale up，提高 DoP。此时，需要使新加入的实例可以参与并行计算。LoongServe 将 sequence parallelism 扩展到 decode 阶段，采用 multi-master distributed decoding 机制。

每一个 master 实例负责 batch 当中的某一个 request。对于一个 decode 阶段的请求来说，master 和组内其他的实例上都存储一部分 kv，master 将新 token 的 q 传给其他实例，其他实例在本地计算 attention，然后将结果汇总到 master。

除了 attention 层以外的层（例如 FFN）都在 master 本地计算，其他实例不参与。这样就达到了分布式 KV Cache 缓存池的效果，只要 master 上有空余的 memory 可以分配给自己负责的 request，就可以继续迭代。

调度算法

算法分为 dispatching, elastic instance allocation, batching 和 elastic scaling plan generation 四个部分。

• dispatching：负责从待处理的 requests 中选择要处理的请求集合 Rp
• elastic instance allocation：决定将 Rp 分配给哪些实例集合 Ep
• batching：根据 Rp 和 Ep 确定每个请求的并行度（DoP）
• elastic scaling plan generation：以 iteration 为粒度动态地规划 ESP 组的 scale up 和 scale down。

具体算法见论文第五章。

测试