《Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving》阅读笔记

最近昇腾提供的大 EP PD 分离推理解决方案非常火，很多开发者都开始使用了。正好这两天也看了一篇 PD 分离的经典论文，就是 Kimi 采用的 PD 分离架构：Mooncake。

背景

传统的大模型推理方式存在一个问题，就是一个 batch 内的所有请求输出长度很可能不同，导致有的输入已经完成推理了，但是必须要等最后一个请求完成推理后，这个 batch 推理才算完成，才会进行下一个 batch 的计算，这就造成了计算资源浪费。为了解决这个问题，就有了 continus batching 特性，当一个 batch 内的某个请求完成推理后，就立即把下一个请求提上来补上，这样就能使得宽度为 batch-size 的“管道”中持续有请求在推理。但是这样又引入了新的问题：新请求的全量计算和旧请求的增量计算放在一起计算的时候，由于全量计算的输入 token 远大于增量推理的输入 token（一般为 1），所以在组 batch 的时候会形成空泡，造成计算资源浪费。

于是，PD 分离被提出来了。P 指的是 prefill，即首 token 推理，D 指的是 decode，即增量推理；分离指的是 prefill 计算和 decode 计算放在不同的计算资源上进行，而不是在相同的芯片上，P 算完了 D 接着算。干湿分离，最大的好处就是让 prefill 计算和 decode 计算互不干扰，解决 P 和 D 的根本矛盾：P 是计算密集型的行为，而 D 是存储密集型的行为。当然，PD 分离也会带来一个问题，就是 P 产生的 kvcache 需要通过网络传输到 D 节点，这个过程需要跨卡甚至跨机！

MoonCake 架构

上图是 MoonCake 的架构图，非常形象，一盒月饼。上面 2 块是 prefill 节点，下面 2 块是 decode 节点。当然，正式系统不一定是 2/2 的配置，还可能是其他数量的配置，不同数量配比对应的性能也不同。每块“月饼”的灰色部分是 GPU 和显存，黄色部分是 CPU 和内存、硬盘存储，“月饼”之间通过 RDMA（Remote Direct Memory Access）传输。左侧是调度系统，主要负责给新来的请求选择合适的 P 节点和 D 节点。右侧是 P 节点和 D 节点的优化目标。

对于 P 节点来说，目标是最大化 cache reuse，P 节点会使用 prefix cache 特性，简单来说就是用新请求的输入去匹配旧请求的输入，如果能匹配上，就用之前请求算出来的 kvcache，避免重复计算。如果 cache reuse 越高，就越能减少计算量、提升效率。当然，P 节点的限制包含 TTFT（time to first token）、最低 MFU（Model FLOPs Utilization）以及内存限制。对于 D 节点来说，目标是优化吞吐量，也就是每秒能处理的请求个数。限制是 TBT（time between tokens）和显存限制。

需要注意的是，整个架构的 kvcache transfer engine 是一个 distributed kvcache pool，也就是各个节点的内存池信息是互通的，有助于 kvcache 的共享传输。

当然，对于我们在背景中提到的网络传输问题，论文也给出了答案，认为只要这个公式满足，kvcache 的传输成本就不会成为阻碍：B/G > 2ds/[gqa×(apd +bd^2)]。其中 B 是 kvcache 的传输带宽，G 是芯片算力，d 代表模型的隐藏层维度，a、b、s 是常量，p 是序列长度，gqa 是 q 头数量除以 kv 头数量。根据这个公式，d 越大，所需带宽越小，并且论文通过实验证明，只有 B 达到 100 Gbps，就有正收益。

调度算法

上图是 PD 分离推理系统处理推理请求的调度算法。

算法输入：P 节点实例池、D 节点实例池、请求 R、cache block size（在这里不用关注）

算法输出：选择的 P 实例和 D 实例。

算法思路是先选择 P 节点，把所有 P 节点遍历一遍，预估选择每个 P 节点计算的总耗时，耗时包括 3 个部分：kvcache 传输时间、任务排队时间、做计算的时间。选择耗时最短的 P 节点作为输出，再根据负载情况选择 D 节点。

算法详细步骤如下：

step1：计算请求 R 的 prefix hash 值，然后根据每个 P 节点的 prefix hash 池信息，找到一个前缀匹配最优的节点，记为“best instance”，对应的 prefix 匹配长度记为”best len"；

step2：遍历每个 P 节点，假设在该节点的 prefix_len 和 best_len 满足这个关系：best_len/prefix_len>kvcache_balancing_threshold，说明这个节点的 prefix_len 太短了，重新做 prefix 计算的成本会比从最佳节点拷贝 kvcache 过来的成本高，所以在这种情况下，需要估算 T_transfer；如果 prefix_len 长度还可以，不满足 best_len/prefix_len>kvcache_balancing_threshold，那么就不需要从最佳节点拷贝 kvcache 过来，直接在这个节点进行计算即可，这种情况 T_transfer=0；

step3：使用经验模型估算每个 P 节点的排队时间和 prefill 计算时间；

step4：计算每个节点的总耗时，在遍历过程中，如果发现更短的总耗时，就更新 TTFT 和实例；

step5：完成 P 节点遍历，选择 TTFT 最短的 P 实例，根据负载情况选择 D 节点。

step6：如果 TTFT 或者 TBT 不满足 SLO，拒绝请求；

step7：进行 kvcache 拷贝传输。

整个算法过程非常简单，但论文中也说了，真正难的、在工程上有挑战的是经验模型 EstimateKVCacheTransferTime、EstimatePreffllQueueTime 和 EstimatePreffllExecutionTime 的构建。

实验结果

来看一下论文的几组实验结论：

1，和 vllm 框架相比，PD 分离框架的 TPOT 更短；

2，和 vllm 框架相比，PD 分离框架的 TTFT 更短、prefix cache 命中率更高；

3，P 和 D 节点数量配比会影响 PD 系统的吞吐。

详细的内容可以参考论文。