Mooncake论文阅读笔记

最近昇腾提供的大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系统的吞吐。

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

本文由博客一文多发平台 OpenWrite 发布！