Mooncake：将 P / D 分离进行到底

千姿百态

众所周知，ML sys 博大精深。四个月前我就在朋友圈见到了 Mooncake 这篇工作，当时觉着和自己差着十万八千里。不曾想命运弄人，四个月后还是扎扎实实拜读了一番，深受启发。
PS：Mooncake 一作是我的本科队友，我曾多次和他在贵系知名的软件工程和人工神经网络课程上组队。如今他在国内意气风发，直击业界痛点，而我在美国水深火热  ...
PD 分离的几层境界

言归正传，摘录方佳瑞学长的观点，PD 分离有几个层面：
Level -1：PD 在单卡上融合，这和 PD 分离南辕北辙。
Level 0：PD 分开计算，但是放在一张卡上调度。这是推理框架的 default 实现，目前的 SGLang v 0.3 也采用了这种方案。
Level 1：分离 P D 在同构设备的同构网络中。比如 P D 都在同一个 A800 节点内，不需要在集群层面进行改造。
Level 2：分离 P D 到同构设备的不同网络中。譬如 P D 分别占据一个完整的 A800 节点，不需要显著改造集群，但是需要建立集群间的高效通讯，譬如 RDMA。
Level 3：彻底分离 P D 到异构设备的异构网络中。譬如 P 占据一个完整的 MI 300 集群，而 D 占据另一个完整的 H20 集群，需要对集群进行显著的改造，并且需要节点间的高效通讯。
两天前我做了个简单的讨论：
这里进一步扩展下结论：

• 异构设备未必能够降低成本：虽然理论上可以采用更便宜的卡，但是 networking 成本不可忽略，而在一个大规模 data center 中部署多种机型也有不可忽略的 fragmentation 代价。
  
• 做 PD 分离对于业务性质有特殊的要求，譬如处理的 context window 和 shared prefix 等等。理论上越高等级的 PD 分离可以降低成本，实际上还需要严谨的评估。
  

Mooncake：将 P / D 分离进行到底

论述了这么多，接着参考业界已有的成熟方案看看 P / D 分离进行到底后的效果。值得再次强调的是，P / D 分离进行到底对于 kimi chat 这样的应用应该是非常 adorable 的，毕竟 kimi 的长文本能力是其突出卖点。不过就我个人而言，我对长文本的需求基本来自 LLM 读文档，而今天我发觉没有什么成熟的应用不具有这个功能。
Introduction

• Mooncake 相比于开源框架而言，能利用的资源更多，但是需要考虑的限制也更复杂。就我所知，当前的 open source serving engine 很难考虑到如何 fully utilize DRAM SSD RDMA 这样的硬件条件。然而，open source engine 可能也不太用考虑 SLO（service level objective 也即用户优先等级）的问题？
  
• Mooncake 将计算操作和硬件资源都进行了彻底的解构。对于后者而言，所有的硬件都被解耦而后重新建立了对应的资源池。每种资源池（比如 CPU 池、DRAM 池、SSD 池）各自有独立的优化目标，以期望整体达到最大优化。
  
• P D 彻底分节点进行导致设计者必须考虑 KV cache 在 P D 节点之间传输的问题。如下图所示，KV Cache 需要考虑 P D 之间所有的传输需求。
  

4. KV Cache 传输的调度几乎是 Mooncake 的核心。直观来说，有两个优化方案：首先是尽可能复用 KV cache；第二是尽可能加大单个 batch 内的 token 数目，用以最大化 MFU（Model FLOPs Utilization）。然而，复用 KV cache 涉及到了 KV cache 远距离传输的问题，也可能这些 KV cache 需要从更底层的存储设备读取上来，这会加大 TTFT（time to first token，产生第一个 token 的时间）；高强度传输 KV cache 也可能导致网络卡顿；而更大的 batch size 往往会导致 TPOT （time per output token，解码每个 token 的时间，这在原文中称为 time between tokens）增大。因此，这些直观的【通量最大】优化可能降低延迟，违背了 SLO。
5. 基于这些复杂的 trade off，参考上图，Mooncake 设计了 global scheduler（conductor）。对于每个 request，Conductor 需要为其选择一组用于 prefiil 和 decode 的设备，然后进行调度。首先迁移尽可能多的 KV cache 到 prefill 设备上；接着在 prefill 设备上通过 chunked and layer-wise prefill 的方法连续地 stream prefill 所得的 KV cache 到 decode 设备上；最后在 decode 设备上加载 KV cache，将此 request 加入到 continuous batching 中，完成 decode。
6. 这样的调度听上去非常自然，但是策略的限制非常复杂。如前所述，在 prefill 阶段，四处传输 KV cache 可能会增大 TTFT。为此，conductor 也需要预测某块 KV cache 可以复用的可能，对高强度使用的 KV cache block 进行交换以及复制。最火热的 block，比如 system prompt 的 block 几乎每个设备都该本地缓存一份；而冷门的 block，比如某个用户上传的邪门文档，可能就该被早日擦除。基于这些考虑，大部分的 DRAM 都会被用于维持 global KV cache pool，这样又给 scheduler 能利用的 DRAM 上了压力。
7. 与此相对，decode 阶段的优化目标显著不同。但是，如前文所述，尽可能增大 batch 内的 tokens 数目会增大 TPOT。
8. 更加糟糕的是，工业级部署的 serving engine 时常要面对极其不均匀的服务压力，在 overload 和 far from enough 之间反复横跳。这里 Mooncake 实现了一个【拒绝策略】（老实说我觉得这是现实的需求，但是我的服务要是经常被某个产品拒绝，我可能就不用这个产品了）
9. 为了实现这样的拒绝策略，设计者希望能够预测下一个阶段的服务器负载程度并且拒绝【某些】可能服务不过来了的请求。这一策略接下来讨论。
10. 采用 chunked pipeline parallelism（CPP）的方式将单个 request 拆分为多个 chunk，这些 chunks 可以按照顺序在不同的 prefill nodes 上完成计算。这有助于减小长请求的 TTFT（【老实说】我不太理解如何是减小，在我看来这会让序列的 TTFT 更加稳定可控）。相比于传统的 sequence parallelism（SP）而言，CPP 减轻了网络压力且 simplifies the reliance on frequent elastic scaling.【这句我没读懂  】
Preliminary

• prefill 阶段是 compute intensive 的，除非 request 很短。考虑到 attention 的计算是平方复杂度的而 MLP 是线性复杂度的，故而 prefill 的计算时间增长相对序列长度是超线性的。prefill 主要决定了 TTFT。
  

2. decode 阶段是 memory intensive 的，因此增大 decode 的计算时间相对 batch 内的 token 数量是亚线性的。（简单来说 decode 阶段增大 batch size 的话，大多时候 GPU 都有相对充足的计算能力来应对线性增大的 tokens）。decode 主要决定了 TPOT。
3. decode 阶段常见的优化是 continuous batching：在每一轮 decode 之前，decode 调度器检测所有请求的状态，将加入新到的请求，而将已经完成的请求清除出去。
4. 在遵循 SLO 的前提下，如何最大化 overall throughput？这一指标在其他工作中又被称为 goodput。
5. 在过载的条件下，如果一个请求完成了 prefill 之后才被 decode 拒绝，那么 prefill 的资源仍旧被浪费了。所以，拒绝某个请求应该尽早。
Overview of Mooncake Architecture

• Mooncake 将 GPU 集群内的 CPU DRAM SSD RDMA 资源分别建立了资源池。
  

2. 在 GPU 和 GPU 之间转移 KV cache 由单独的 GPU direct RDMA 设备完成，被称为 messager。原文中基于哈希的前缀存储能够为上层用户提供 context caching API。
3. 如果某个请求中未被 cache 的 tokens 数目超过了一个特定的 threshold（prefill_chunk）。这个请求将会被拆为多个 chunks，然后流水执行。通常而言，prefill_chunk 的大小大于 1k。
4. 在 prefill 阶段进行预估而没有被拒绝的请求可能在 decode 阶段出于 SLO 的缘故而被拒绝，这也是一个显著的 staleness 问题。
Real-world Request Trace

• 作者发布了一个真实情况下的请求数据集。为了保留前缀性，这些请求都采样自某段时间内某些 session。作者给出了 timestamp，input_length，output_length 和 hash_ids。
  

2. 这一 trace 没有保留任何实际的 tokens 或者 text。【说实话，我还没搞懂没有 tokens 的话，这能咋用  】
Prefill Pool Implementation

• 考虑了 chunked prefill 后，是否还需要做 PD 分离（分卡）是值得讨论的。chunked prefil 将 request tokens 分为多个 chunks，加入 continuous batching 中。这样处理有两个明显的好处：在不做 PD 分离的情况下，所有节点被平等对待，调度器不用这么麻烦；每个 chunk 可以捎带 decode 请求，增大了 decode 的 MFU（其实这里原文叙述的和我理解【有出入】，因为 chunked prefill 原文讲的是将 decode 请求塞到 prefill chunk 里面，但是 Mooncake 原文叙述的是将 chunked prefill 内联到 decode batch 中，我认为这表达的是一个意思？）
  

2. 即便如此，Mooncake 仍然选择将 PD 分离到底，并且规定只有当请求较短不用做 chunk 且符合 SLO 时，才将这样的请求内联到 decode 请求中【这是说 decode 节点也要负责一部分短序列的 prefill 么？】这是出于两个考虑：长文本的 prefill 需要特殊的 cross-node parallelism；节约 VRAM  。
3. 对于长序列而言，很有可能输入序列的长度会是输出序列长度的数十倍，这对 TTFT 带来了很大的挑战。将长序列拆分到单个节点上的多个显卡是 desirable 的，然而拆分到多个节点上并不现实。这需要两个昂贵的 RDMA-based all-reduce operations per layer，严重影响了 MFU。
4. 考虑最新的 sequence parallelism（SP）：SP 在不同节点之间分割单个 request 并实现了加速。SP 方法基于 Ring Attention 或者 Striped Attention 实现，利用 attention 的可结合性并且每一 layer 要求一次 cross-node 通讯。整体上减少了网络开销并且提高了 MFU。即便如此，跨节点 SP 的 MFU 仍旧不如单一节点的 TP。
5. 出于以上考虑，Mooncake 采用了 chunked pipeline parallelism（CPP）方法来应对长序列的 prefill 请求。CP 和训练时的 PP 类似，仅仅会在每个 pipeline 的边界处需要跨节点的通讯，这样通讯时间可以被其他通讯时间 overlap 掉。这样带来了更好的 MFU 和更低的网络通讯。此外，CPP 可以很自然的处理长短序列，短序列不会由于等待长序列的 prefill 而带来太长的 overhead。训练框架中的 pp 已经被广泛探索了，但是推理框架的 pp 最近才逐渐出现。
6. prefill 是逐层进行的，因此可以利用 computation time overlap 掉 transfer and dumping KV cache 的时间。如此重叠带来了可见的好处——prefill scheduler 可以忽略可用的 VRAM 大小，只要能够容得下当前的序列。【我不理解  】譬如在结构图中所示的那样，prefill scheduler 只考虑 KVCache 的分布和可用的 DRAM 大小即可。不过，此时剩下的 VRAM 大小有待进一步利用。例如，OpenAI 最近开放了 Batch API，允许用户以半价来得到异步的请求（TTFT 可能高达 24h  ），这非常适合一些不需要即时响应的任务（譬如合成数据  ）。这种请求甚至也没有 TPOT 要求（或者说要求相当的宽裕），因此如果有足够的 VRAM 存储相应的 KVCache，似乎也可以将这些 decode 请求内联到 prefill stage 里，以提高 MFU。
KV Cache-centric Scheduling

• 大多数的 LLM serving systems 依靠每个 engine 上当前分配的请求数来评估负载情况，以此设计 load-balance 算法。然而 Mooncake 进行了彻底的 P D 分离，prefill engine 的选择不仅需要考虑当前 engine 的负载，还要考虑 prefill cache 的命中率和这些命中 KVCache block 的分布。如前所述，一味将 prefill 请求发往具有最长前缀的 engine 可以减少计算成本，但是 transfer and load 这些 KV cache 可能会增大 TTFT 或者违背 SLO。简单来说，这时候可能调度到一些看上去前缀匹配更差的节点反而有利。为了解决这些复杂问题，Mooncake 需要设计 cache-aware prefill / decode scheduling 算法——同时考虑到 KV cache 的匹配程度以及 engine 上的负载情况。
  

2. 基于此，conductor （global scheduler）需要做一些估计。具体而言，conductor 需要估计一个 prefill 请求的执行时间。因此，Mooncake 训练并且部署了一个估计模型 estimater，在每个 prefill engine 上，estimater 根据 request 的长度、匹配到的 cache 长度与 engine 当前的负载程度来得到一个大概的 prefill wait time。最后，conductor 将这个 prefill 请求指定到最短的预估 TTFT engine 上，同时更新这个 engine 上的 cache 和 queue time。【我猜测这个实际过程更复杂，毕竟还要考虑到负载很低的 engine 可以迁移其他 engine 的 KV cache 来做 prefill，可能还更快？】如果 SLO 没法被满足，conductor 会在这一层就直接 raise HTTP 429 Too Many Requests 到更高层的 layer。
3. 这个估计模型的估测准确么？事实上准确，毕竟 moonshot 数据够多  此外，大量的计算是可以并行的，譬如 TTFT 最后的累加，所以预测模型的运行时间几乎可以忽略。
4. 更复杂的估计事实上是 KV cache 传输的开销，如我所料。这同时取决于传输 block 的大小和当前的网络情况。
5. 在 Mooncake 中，每个 prifll engine 都有自身存储的 prefix cache。这些 cache 的频率千差万别，譬如 system prompt 每人都有，而用户上传的长文本文件大概率只会在一个 session 中利用。如何分享一些高频的 block 来提高 hit rate 并避免频繁读取的开销，这是很值得考虑的问题。
6. 一个虚空打靶（straw-man）的回答是，我们可以再搞一个预测模型来预测每个 block 未来的使用情况。然而这次和预测 prefill tiem 不同，一个 block 未来的使用情况简直差异大的不行。显然，我自己可能一个 session 里面相同的两句话，一会儿隔了 1mins，一会儿隔个好几天   明显 prefill time 会相比用户接着使用这个 session 的情况好预测的多。做不到精准的预测，Mooncake 还是开发了一个 heuristic-based automated hot-spot migration scheme（这名字太长了...）来均衡负载情况。
7. 如前文所述，确实会有一些 request，其 prefix cache 最长的 engine 负载过大。这种情况下，conductor 会迁移那份最长的 prefix cache 到其他 engine 上去，如果这个迁移时间比起预估的 prefill time 更短。当然，如果从 prefix 最优的 engine 到 load 更低的 engine 的迁移时间过长，也可以选择放弃这次 prefix cache 迁移，直接在那个负载更低的 engine 上做 prefill。（这里用了手动设定的一个 threshold 来衡量是否值得迁移，作者将这个 threshold 的设计留给了后人）
Overload-oriented Scheduling

• 前文充分论述了在每个请求都要被执行时的 scheduling 方法，然而实际上 kimi chat 服务过载时，一些请求会被 reject（again，产品经理不怕被挂路灯么  ）。当容忍的服务数量超过了特定限度后，会有一些幸运儿被 reject 或者 deferred。这是 Mooncake 率先提出的场景...
  

2. 在服务器超载时拒绝服务，就得首先定义什么是超载，也即什么是负载。和不做 PD 分离的 serving 系统不同，你很难界定 Mooncake 这种将 PD 分离进行到底的系统的负载情况。基于此，作者认定【符合 SLO 的 overall goodput】为系统的负载，而又分别定义了 load of TTFT 以及 load of TPOT。进而基于这两个负载指标决定是否接受一个 new request 的 prefill 请求以及是否执行完成 prefill 后的 request 的 decode 操作。
3. 可以想见，这种拒绝请求策略存在明显的 staleness 的问题，也即如果一个 request 在 prefill 前被判断能够执行，结果执行完了 prefill 后，先前预设给他的 decode engine 这时候超载了，这个 request 还是会被 reject 掉，导致 prefill 阶段的计算被浪费。为此，需要尽可能在 prefill 阶段同时预估能否在【这个阶段】完成 prefill 并且在【下个阶段】完成 decode，作者将其称为 early rejection。
4. 如果完全不考虑这个 staleness，会出现 prefill 和 decode 出现交错达峰的情况，比如下图。


5. 如何预测下个阶段的 decode engine 负载情况呢？我们设想有两个 level 的可行方法。首先是 request level，细粒度地预测每个 request 可能的生成 tokens 数量。不过，这个在技术上还是非常难。（同样也可以假想下，这个预测确实 variance 太大了。如此看来，只有 prefill 耗时是比较好预测的，其他的比如某个 KV cache block 在未来时间的使用情况以及某个 request 生成序列的长度，都挺难预测的。）
6. 在更高一层的 system level 上，不必预测每个 request 生成的 tokens 的长度，反而可以预测下个阶段的 decode batch size 或者下个阶段的 TPOT。这种预测的粒度更低，更适合超载情况。Mooncake 选择了一个【我看不懂的】system-level prediction。
评估

• 作者分别采用了 arXiv submission、L-Eval、合成数据和他们自己收集的 real data 来做 benchmark。对于 real data，按照他们真实的 arrive time 来发送请求，而其他数据集采用泊松分布作为请求发送时间。
  

Related Work

• DistServe 尝试优化了资源分配与并行策略达到更好的 GPU goodput。
  

2. TetriInfer 整合了 chunked prefill 与 PD 分离方案，并对 P D 两个阶段都用了基于预测的 scheduling 算法。
3. AttentionStore 和 Mooncake 有异曲同工之妙，然而 Mooncake 进一步发展了 long-context 情况下的处理方案。
4. 与其他线性变换操作不同，decode 节点的 attention 操作的 arithmetic intensivity 仅仅与 attention head 数量 / K-V heads 数量成正比。增加 decode batch size 无法显著增大 attention 操作的计算强度，并且 attention 操作比起线性操作更加 memory-bounded。因此，将 attention 操作与其他线性操作分离可能是一大优化方向。在这个想法上，MLA 直接抬高了 attention 操作的算数强度，前景广阔。
总结

Mooncake 将 PD 分离进行到底，在工业级别应用上考虑到了多种我在学校很难考虑的硬件资源，可以说是 PD 分离时代的典范。希望我接下来做的工作不是去破解工业界秘密  
顺带，Mooncake 作者团队准备开源 distributed kvcache pool 的部分代码，并且和开源社区合作，推动 kvcache 池化层的 protocol，敬请期待！

原文地址：https://zhuanlan.zhihu.com/p/1711346141