MemServe: Flexible Mem Pool for Building Disaggregated LLM Serving with Caching

发表时间: 2024-06 · arXiv:2406.17565 (preprint)

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

作者/机构：Cunchen Hu, Heyang Huang, Junhao Hu, Jiang Xu, Xusheng Chen, Tao Xie, Chenxi Wang, Sa Wang, Yungang Bao, Ninghui Sun, Yizhou Shan / Huawei Cloud, UCAS, ICT, CAS, Peking University

速读

一句话结论
提出了一种名为 MemServe 的统一大语言模型推理服务系统，通过引入弹性内存池 MemPool，首次将上下文缓存与分离式推理结合，大幅降低了作业完成时间和首 token 延迟。

要解决什么问题
当前的大语言模型推理系统无法同时兼容两类主流的 KV Cache 优化技术：跨请求的上下文缓存（Context Caching）和请求内的分离式推理（Disaggregated Inference）。原有的系统架构假设 KV Cache 只是绑定在单一实例、单一请求上的中间数据。当引入分离式推理（将 Prefill 和 Decode 阶段拆分到不同物理实例）时，请求被物理打散，现有的上下文缓存机制完全失效，因为 Decode 实例缺乏将生成的 KV Cache 传回 Prefill 实例以供未来复用的机制。此外，现有的请求调度器通常基于负载或会话 ID 进行分发，缺乏全局视角的缓存感知能力，导致跨会话的 KV Cache 复用率极低，严重限制了多轮对话和长文本场景下的吞吐上限。

怎么做的
核心思路是剥离出独立的弹性内存池 MemPool，将分布式的内存管理、索引和网络传输统一抽象，从而在底层打通分离式推理与上下文缓存。该设计主要由三个关键部件构成。首先是 MemPool，它接管了集群内所有的 CPU DRAM 和 GPU HBM，并提供内存分配、索引和分布式传输三类 API。其索引层采用基数树将 Prompt Token 映射到历史 KV Cache，并记录数据所在的物理实例。其次是支持缓存的分离式推理架构，通过调用 MemPool 的传输并插入接口，Decode 实例在请求结束后能将新生成的 KV Cache 传回并插入到 Prefill 实例的索引中，完成闭环，实现全量上下文缓存。为了解决 PagedAttention 机制下 KV Cache 内存不连续导致底层网络库传输开销过大的问题，作者引入了块聚合技术，将原本分散的内存块聚合成大块，使网络 API 调用次数减少了 $2 * L$ 倍（$L$ 为模型层数）。最后是全局调度器，它维护了一棵全局 Prompt 树。当新请求到达时，调度器会基于算子级代价模型 $exec(x,y)$（$x$ 为 Prompt 长度，$y$ 为缓存命中率）来预估执行时间。调度器首先通过以下公式选择排队延迟与预估执行时间之和最小的实例 $p$：

$$ \operatorname { a r g m i n } _ { p } \sum _ { \mathrm { c u r r e n t } } e x e c ( x ^ { \prime } , y _ { p } ^ { \prime } ) + e x e c ( x , y _ { p } ) $$


如果发现另一个实例 $p^\prime$ 拥有更高的缓存命中率，调度器会通过以下条件判断是把缺失的 KV Cache 从 $p^\prime$ 传给 $p$，还是直接在 $p$ 上重新计算：

$$ \mathrm { t r a n s f e r } ( y _ { p } , y _ { p } ^ { \prime } ) \leq e x e c ( x , y _ { p } ) - e x e c ( x , y _ { p } ^ { \prime } ) $$
通过这套机制，系统在保持分离式推理高并发优势的同时，最大化了全局 KV Cache 的复用率。

效果如何
实验在一台搭载 8 张 H800-80GB GPU 的服务器上进行，测试模型为 Llama2-13B（张量并行度设为 2）。对比基线为原生的 vLLM（代表传统的 Prefill 与 Decode 同机部署路线）。作者测试了纯同机部署、同机部署加缓存、纯分离式推理、分离式推理加缓存四种设置。在真实对话数据集 ShareGPT 上，纯分离式推理相比原生 vLLM 将平均作业完成时间优化了 30%；在此基础上叠加 MemPool 提供的上下文缓存后，平均作业完成时间进一步优化了 17%，P99 作业完成时间进一步优化了 29%，平均首 token 延迟降低了 58%。在长文本问答数据集 LooGLE（具有长 Prompt 和高比例共享前缀）上，分离式推理结合上下文缓存展现出更强的优势，相比纯分离式推理，平均作业完成时间进一步优化了 26.9%，平均首 token 延迟优化了 56.2%。该方案也存在一定的局限性与代价。首先，全局调度器的 Prompt 树仅在响应返回时更新，无法实时感知底层实例的缓存驱逐事件，可能导致路由信息滞后（作者通过设置 TTL 缓解）。其次，目前的分布式传输强依赖单机内的 NVLink 和 NCCL 点对点通信，尚未实现基于 RDMA 的跨节点传输，且为了解决 NCCL 缺乏目的地址指定功能带来的乱序问题，当前实现只能在每个通信器中单线程运行。

主要贡献

大型语言模型（LLM）服务已经从无状态系统演变为有状态系统，利用了上下文缓存（Context Caching）和分离式推理（Disaggregated Inference）等技术。这些优化技术延长了KV Cache的生命周期和作用域，因此需要一种全新的架构方法。本文提出了MemServe，这是一个将请求间（inter-request）和请求内（intra-request）优化结合在一个统一系统中的框架。

核心问题与研究目标：
现有的LLM服务系统无法同时应用请求间和请求内的依赖利用优化。当前的上下文缓存（请求间）方法在设计时未考虑请求内拆分的场景，导致分离式推理（请求内）无法从上下文缓存中受益，因为它缺乏将KV Cache从Decode实例传回Prefill实例以供未来重用的机制。同样，序列并行（Sequence Parallelism）将KV Cache分布在多个实例中，也缺乏保存和重用它的机制。此外，当前的系统缺乏自上而下的整体设计，通常基于负载或会话ID进行调度，无法最大化跨会话的KV Cache重用。本文的研究目标是设计一个统一的架构，从根本上解决分布式实例间KV Cache的管理和传输问题。

创新点：
1. 引入MemPool（弹性内存池）：这是一个管理集群所有内存（包括CPU DRAM和GPU HBM）的基础组件，提供了一套丰富的API，用于管理分布式内存和KV Cache。
2. 首次结合上下文缓存与分离式推理：利用MemPool API，MemServe不仅在标准的Prefill-Decode共置（PD-colocated）实例上实现了上下文缓存，还在分离式推理架构中成功引入了上下文缓存机制。
3. 全局调度器与局部性感知策略：为了最大化KV Cache的重用，MemServe采用了一个基于全局提示词树（Global Prompt Trees）的局部性感知路由策略。
4. 内存布局与网络传输协同优化：针对现有AI网络协议栈中离散内存布局和点对点网络原语带来的开销，MemServe提出了使用大页（Huge Pages）聚合内存块的优化方案。

背景知识

生成式LLM推理分为两个阶段。LLM推理涉及根据输入提示词生成一系列输出Token。该过程包含两个阶段：Prefill（预填充）和Decode（解码）。在Prefill阶段，模型处理提示词以生成KV Cache，即自注意力机制中生成的键值对。在Decode阶段，模型迭代地使用KV Cache来生成Token。KV Cache的大小随着生成Token数量的增加而线性增长。

请求间优化利用请求之间的依赖关系。上下文缓存是该类别中唯一已知的技术。为了构建上下文缓存，模型存储并重用自注意力机制的KV Cache，以避免在相似或重复请求中进行冗余计算。这在多个请求共享公共前缀（如系统提示词）的场景中非常有用。实现这一目标需要两个基本机制：首先，需要一个索引来发现请求之间的依赖关系并找到保存的KV Cache（如表2所示）；其次，需要修改推理引擎和注意力算子以重用历史KV Cache（参考【40，Efficiently programming large language models using sglang+2023+arXiv】）。

请求内优化利用单个请求内部的依赖关系。两个著名的例子是分离式推理（参考【23，Splitwise: Efficient generative llm inference using phase splitting+2023+arXiv】等）和序列并行（参考【17，Infinite-llm: Efficient llm service for long context with distattention and distributed kvcache+2024+arXiv】）。通常，将Prefill与Decode分离可以减少这两个阶段之间的干扰，并允许它们在异构硬件上独立扩展。然而，这会将单个请求分解为两个子请求，并需要将KV Cache从Prefill严格传输到Decode。序列并行也是如此，分布式实例需要以严格的方式交换自注意力的输出。总体而言，请求内优化需要高效的机制在实例之间传输KV Cache。

方法细节

MemServe被设计为一个支持请求间和请求内优化的大规模LLM服务系统。它包含三个主要组件：全局调度器、多种类型的推理实例以及弹性内存池（MemPool）。MemPool提供了一组用于内存分配、索引管理和分布式传输的API。MemServe利用MemPool API在常规和分离式推理架构之上构建上下文缓存。全局调度器将用户的推理请求转发到正确的推理实例，它使用基于新型分布式提示词树的局部性感知策略，以最大化KV Cache的重用。
图1. MemServe架构

MemPool管理推理集群中的所有内存，并提供三类API。MemPool运行在每个推理实例内，统一管理包括CPU DRAM和GPU HBM在内的所有本地内存，并提供分布式内存池API。这些API分为内存块（Memory Block）、索引（Index）和分布式传输（Distributed Transfer）三类。推理引擎可以使用 alloc_mem 和 free_mem 等内存块API分配固定大小的内存块来存储KV Cache；使用 insert、match、delete 等索引API将活跃KV Cache转换为历史KV Cache并建立Token到地址的映射；使用 transfer 和 transfer_with_insert 等分布式API在实例间交换KV Cache。

MemPool通过索引层将提示词Token映射到历史KV Cache。在LLM服务领域存在三种索引方法：Token ID、Session ID和Document ID。Token ID因其通用性被采用，能够处理任何共享提示词前缀的情况。为了实现该索引，MemPool采用了SGLang提出的基数树（Radix Tree），并进行了两个关键扩展：第一，由于MemPool同时管理GPU HBM和CPU DRAM，基数树能够引用系统中任何位置的数据；第二，为支持全局调度器，增加了一个字段以指示哪个推理实例持有该数据。为了最小化数据重塑开销，MemPool在将活跃KV Cache插入历史KV Cache时保持原始内存布局，其索引粒度与推理引擎的配置保持一致（例如，vLLM中节点指向16个Token的KV Cache块）。

MemPool提供分布式API以屏蔽底层的异构性进行数据交换。数据传输工作流被分解为三个步骤：分配（Allocation）、传输（Transmission）和插入（Insertion）。首先，发送方实例向接收方发起请求，接收方调用 alloc_mem 在本地分配HBM或DRAM，并将地址列表及并行配置返回给发送方。接着，发送方通过最快的可用路径将KV Cache传输给接收方。最后，如果调用的是 transfer_with_insert，接收方会在本地调用 insert 将新传输的Token和KV Cache插入到本地索引中。transfer_with_insert 能够避免额外建立映射的网络往返，特别适用于将历史KV Cache从Decode实例传回Prefill实例。
图2. MemPool传输API

数据传输步骤必须处理三种异构性：并行度、内存和网络。为了处理非对称并行度，发送方首先检查KV Cache在张量并行或流水线并行维度上的划分方式，随后划分本地缓存并调用相应的网络原语。如果历史KV Cache已被交换到DRAM中，就会出现内存不对称。MemPool始终尝试使用数据拷贝最少的最快链路传输数据。如果硬件支持（如NVLINK连接的HBM和DRAM），处理内存不对称仅需一次内存拷贝；但在常规GPU服务器上，数据路径中不可避免地会产生额外的内存拷贝。
图3. MemPool支持的用例

MemPool以单实例为粒度处理故障和动态缩放。MemServe的集群管理（CM）模块负责维护集群配置、添加或移除实例以及监控健康状况。由于内存块和分布式传输API可以修改远程实例的状态，当某个实例发生故障时，其他实例的飞行中请求将会超时。CM通过定期心跳检测故障，并向所有运行中的实例广播更新后的集群信息。收到通知后，每个实例会释放由故障实例分配的任何内存块，以防止内存泄漏。

MemServe通过四个设计里程碑在分离式推理中实现上下文缓存。上下文缓存和分离式推理由于缺乏KV Cache管理机制而无法共存。MemServe利用MemPool API逐步构建完整的解决方案：
1. PD-Basic：通过修改推理引擎，Prefill实例调用 transfer API 将生成的活跃KV Cache连同元数据（请求ID、采样参数等）传输给Decode实例。
2. PD-Caching-1：在Prefill实例上启用缓存。Prefill实例调用 insert 将活跃KV Cache转为历史KV Cache。这适用于共享长前缀的负载，但在多轮对话中会导致Prefill实例反复向Decode实例发送相同的活跃KV Cache，浪费带宽。
3. PD-Caching-2：在Decode实例上启用缓存以减少重复数据移动。Prefill实例改用 transfer_with_insert，使Decode实例接收后直接将其插入本地索引。请求结束后，Decode实例调用 insert 保存Decode阶段产生的KV Cache。这减少了数据移动，但Prefill实例依然缺乏Decode阶段的历史缓存。
4. PD-Caching-3：实现完整的上下文缓存。在请求结束后，Decode实例调用 transfer_with_insert 将Decode阶段产生的KV Cache回传给Prefill实例。这样，Prefill实例的缓存随对话轮数线性增长。
图4. 使用MemPool API在分离式推理中增强上下文缓存

针对离散内存布局和网络原语不足，MemServe提出了内存块聚合优化。在将活跃KV Cache从Prefill传到Decode时，存在按层传输（by-layer）和按请求传输（by-request）两种方式。由于PagedAttention引入的分页机制将KV Cache存储在固定大小的细粒度内存块中（如每层2个块），而现有的AI网络栈（如NCCL）缺乏高效的点对点、Gather和Scatter原语，导致网络API的调用次数等于离散内存块的数量。这使得在高负载下，频繁的网络调用带来巨大开销。MemServe提出通过聚合小KV块来减少碎片，类似于使用大页（huge pages）。具体而言，将每层的2个块聚合成一个大小为 $2 * L$ 的大块。这仅适用于按请求传输（by-request-agg），能将网络API调用次数减少 $2 * L$ 次，显著提升网络性能。
图5. 分离式推理的网络与内存优化

MemServe设计了一个算子级的成本模型来预测执行时间并指导调度。成本模型 $exec(x, y)$ 预测长度为 $x$ 且缓存命中率为 $y$ 的提示词的Prefill执行时间。该模型用于：1) 全局调度器通过公式 $\operatorname { a r g m i n } _ { \hat { p } } \sum _ { \mathrm { c u r r e n t } x ^ { \prime } \mathrm { o n } \hat { p } } exec ( x ^ { \prime } , y _ { \hat { p } } ^ { \prime } ) + exec ( x , y _ { \hat { p } } )$（第一项为队列延迟）进行局部性感知和负载均衡调度；2) 通过条件 $\mathrm { transfer } ( y _ { p } , y _ { p } ^ { \prime } ) \leq exec ( x , y _ { p } ) - exec ( x , y _ { p } ^ { \prime } )$ 决定是从缓存命中率更高的实例传输KV Cache，还是直接重新计算。

成本模型通过拟合算子级别的执行时间来实现更高的可扩展性。相比于架构级模型，算子级模型在改变张量并行（TP）或流水线并行（PP）配置时无需重新校准。例如，对于计算密集型算子，定义 $op ( x , y ) = ( \eta - 1 ) \cdot T _ { fullwave } + T _ { lastwave }$，其中 $\eta = \lceil B _ { total } / SMs _ { num } \rceil$；对于内存密集型算子，由于多级缓存难以确定吞吐极限，直接拟合延迟与读写操作数量的关系，如Prefix Attention定义为 $op _ { attention } ( x , y ) = ax ^ { 2 } y + bx ^ { 2 } + cx + d$；对于常数算子（如归一化、激活），使用简单的线性关系建模。

全局调度器（GS）利用全局提示词树实现局部性感知调度。GS维护三种类型的全局提示词树（针对Prefill-only、Decode-only和PD-colocated实例）。每棵树在基数树节点中增加了一个指向存储KV Cache实例的字段。当请求到达时，GS运行分词器将提示词转为Token IDs，并发查询所有类型的树。策略模块根据查询结果和当前负载，选择具有最长公共前缀（即最大历史KV Cache）的实例。如果存在其他实例持有该实例缺失的额外KV Cache，策略引擎会输出这些实例及对应的Token IDs。最后，GS将请求发送给选定的实例，并在响应返回时更新全局提示词树。为了处理底层实例的本地驱逐事件导致的状态不同步，全局提示词树配置了生存时间（TTL）。
图6. 全局调度器架构

实验环境

图7. 工作负载统计
图7. 工作负载统计

实验结果

端到端应用性能（ShareGPT, LooGLE, ReAct）
* 实验内容：比较四种设置：PD（原生vLLM）、PD-CC（PD+缓存）、1P1D（分离式推理）、1P1D-CC（分离式推理+完整缓存）。
* 实验结果与分析：
* 在ShareGPT上，与PD相比，分离式推理（1P2D）将平均和P99的作业完成时间（JCT）分别提高了30%和42%。引入上下文缓存（1P2D-CC）进一步将平均和P99 JCT提高了17%和29%，并将首字延迟（TTFT）提高了58%和45%。
* 在LooGLE上，分离式推理将平均和P99 JCT提高了10.3%和10.8%。上下文缓存进一步将平均和P99 JCT提高了26.9%和22.5%，TTFT提高了56.2%和45.2%。
* 在ReAct上，分离式使得平均和P99 JCT提高了40.8%和53.1%。缓存进一步使JCT提升26.7%和21.4%，TTFT提升78.5%和84.9%。
* 结论：分离式推理+上下文缓存（1P1D-CC）在所有工作负载中均表现最佳，显著提升了JCT和TTFT（详见Fig 8）。
图8. 端到端评估

MemPool API 性能微基准测试

缓存机制对比测试
* 实验内容：对比原生vLLM的哈希索引与MemPool的基数树索引在Prefill阶段的开销。
* 实验结果与分析：随着提示词长度增加，原生vLLM的哈希前缀机制产生了巨大的开销，而MemPool的基数树索引开销极小。
* 结论：使用MemPool进行上下文缓存引入的开销微乎其微（详见Fig 10）。
图10. 缓存研究

内存聚合与传输机制测试
* 实验内容：对比原始离散内存布局和提议的聚合内存布局（Agg_Block）在NCCL传输2048-token KV Cache时的性能；对比By-Layer、By-Req和By-Req-Agg的延迟。
* 实验结果与分析：聚合方法大幅优于原始内存布局。对于By-Req-Agg，它在高负载下显著优于By-Layer和原生的By-Req。
* 结论：内存聚合有效减少了网络API调用，大幅提升了传输性能（详见Fig 11, Fig 12）。
图11. 网络与内存布局优化研究
图12. 对比By-Layer、By-Req和By-Req-Agg

上下文缓存成本模型与调度策略测试
* 实验内容：测试不同提示词长度、批处理大小、内存位置对TTFT的影响；对比算子级与架构级成本模型的准确性；对比不同全局调度策略。
* 实验结果与分析：缓存命中率越大，TTFT改善越明显。算子级成本模型在TP=2时的预测误差远低于架构级模型，展现出更好的可扩展性。基于提示词树的调度策略相比于会话内调度，使P99 TTFT提高了59%。
* 结论：算子级成本模型更准确且易扩展；全局提示词树调度能最大化KV Cache重用（详见Fig 13, Fig 14, Fig 15）。
图13. 上下文缓存成本模型
图14. 成本模型准确性
图15. 全局调度器策略

补充细节

MemPool底层的网络通信实现采用了NCCL和Socket。MemPool包含暴露API的Python库和执行数据传输的C++核心。在HBM之间，MemPool使用NCCL的点对点 send 和 recv API进行传输；如果涉及DRAM，则使用Socket API。由于NCCL是为张量和流水线并行设计的集合通信库，其 send 和 recv API只指定源地址而不指定目标地址，难以在多线程下保证收发顺序。因此，MemPool每个NCCL通信器只使用单线程以保证顺序。同时，由于NCCL缺乏Gather/Scatter原语，MemPool必须多次调用收发对以跨越异构并行实例传输数据。

MemServe通过深度修改vLLM实现了分离式推理与上下文缓存。为了对接MemPool，实现代码替换了vLLM原有的缓存引擎和基于哈希的前缀缓存机制。为了实现内存块聚合（Block Aggregation），MemServe修改了几个关键的CUDA内核，例如 paged_attention、swap_blocks 以及 reshape_and_cache。

相关工作对比：
* 分离式推理：Splitwise、TetriServe、DistServe和Dejavu等工作通过临时修改推理引擎来实现分离式推理。而MemServe采用不同方法，先抽象出独立的MemPool组件，再将分离式推理作为MemPool的一个用例来实现。
* 上下文缓存：Pensieve、SGLang、Prompt Cache等均在PD共置设置下构建缓存。MemServe是首个利用MemPool API在分离式推理之上逐步构建上下文缓存的系统。

结论

本文提出了MemServe，这是一个旨在通过统一请求间（Inter-request）和请求内（Intra-request）优化来提升LLM服务效率的新型系统。MemServe的核心是一个管理分布式实例间KV Cache的分布式MemPool。通过利用MemPool API，MemServe成功构建了上下文缓存、分离式推理以及两者的组合方案。端到端的实验结果表明，MemServe能够大幅度改善作业完成时间（JCT）、首字延迟（TTFT）以及每个输出Token的时间（TPOT）。未来的工作将探索MemPool支持的更多优化组合方案。

参考文献引用汇总