近期，国产大模型公司「月之暗面」联合清华大学团队发布了一篇名为《Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving》的论文，详细介绍了Mooncake——由月之暗面创建的超人气智能助手 Kimi的底层推理平台架构的细节。
论文如下（滑动查看）。由于是纯英文硬核技术论文，就不占用太长篇幅了。感兴趣的朋友可以在公众号后台回复「月之暗面」获取原论文PDF下载。
本文聚焦于该论文的解读，分为两部分：
Part 1：月之暗面工程副总裁、AI Infra负责人许欣然对该论文的“碎碎念”；
Part 2：论文作者之一，清华大学助理教授章明星对该论文的解读：
（1）在月之暗面做月饼，Kimi 以 KVCache 为中心的分离式推理架构
（2）Kimi “泼天的流量”怎么接，分离架构下基于预测的调度策略
（以上内容来自知乎。）
背景：Mooncake是什么？
Mooncake 采用以 KVCache 为中心的解耦架构，将预填充集群与解码集群分离，并充分利用 GPU 集群中未充分利用的 CPU、DRAM 和 SSD 资源，实现 KVCache 的解耦缓存。
Mooncake 的核心是其以 KVCache 为中心的调度程序，它在最大化整体有效吞吐量和满足与延迟相关的服务级别目标 (SLO) 要求之间取得平衡。与假设所有请求都会得到处理的传统研究不同，Mooncake 面临着高度超载场景带来的挑战。为了缓解这些问题，我们开发了一种基于预测的早期拒绝策略。实验表明，Mooncake 在长上下文场景中表现出色。与基线方法相比，Mooncake 在某些模拟场景中可以实现高达 525% 的吞吐量提升，同时遵守 SLO。在实际工作负载下，Mooncake 的创新架构使Kimi能够处理 75% 以上的请求。
Part 1 许欣然：关于Mooncake的碎碎念
这周把 Kimi 底层用的推理系统方案写了个简单的 Tech Report：https://github.com/kvcache-ai/Mooncake?tab=readme-ov-file。
我们坚信一个系统方案只有实际在场景中验证过的，才有分享出来的必要性。本论文与很多 Prefill/Decoding 分离的论文不同的是，这套方案已经在大规模集群上进行几个月的验证并证明了方案的有效性。这也是为什么一个 POC 写在所有业内论文之前的系统，直到今天才发布出来跟大家见面。
目前这套系统承载了 Kimi 线上80% 以上的流量，效果很好也为产品带来了更多的设计空间。
系统是需要跟随着应用快速变化的，同时也需要“硬件厂商”和“云厂商”早点接受新的理念才能跟上浪潮。发出这篇论文，主要是希望给各方提供一些信心，提供一些推理规模足够大场景下的必然优化思路。
趁这个机会，希望给各家硬件厂商和云厂商一些“暴论”：
Mooncake 这类的分离策略会是一个长期来看的趋势
现在、立刻、马上真能省很多钱（毕竟不能公开规模和每日请求的 pattern，如果你说省不了那你都对）
KVCache 的容量会长期保持高位，因此围绕着 KVCache 来优化是非常有必要的
"private memory per request” 是整个推理系统优化中的关键瓶颈（否则 groq 万岁），会有很多努力来降低 KVCache 的大小，但同时会有更多动力来增大
分离之后，允许整个系统往 “算力/$” 和 “带宽/$” 的两个方向独立发展，对硬件优化是更友好的
AR 的模型架构惯性在短期内难以颠覆，因此总可以认定 decoding 的成本总会跟 bandwidth 成本有非常强的正相关性质
“带宽/$” 低一个数量级的硬件方案已经在肉眼可见的范围内了
纵观历史，“算力/$” & “带宽/$” 同时最优的芯片似乎还没出现过，集群上必然要拆分成异构的两个部分
Mooncake 方案和 MLA、各种 KVCache 的压缩方案 都是正交的，KVCache 变小了意味着 Mooncake 的方案收益会更明显
的确，有很多方向可能会让 Mooncake 的架构变成没必要的方案
given hardware lottery，新架构演进会是一个相当缓慢的过程（不 AR、用 RAG 做 Attention 等等方案），不能因噎废食
包括我们自己也在投入很多资源在 break 现有框架上，因此有理由相信在可见的未来推理方案还会变动
由于目前海量的推理压力，所以软件系统做为一个迭代速度极快的方案，就应该一代模型一次跟进
预测这个状态至少会持续2~3年，因此集群层面现在已经值得做拆分了
芯片层面值得做为一个重要的设计考量，在芯片的 IO 能力上要多预留一些能力
最后，再次非常感谢团队里几位同学 &@ZHANG Mingxing老师的付出~
最后打个广告
Moonshot AI 的 AI Infra 团队招聘正式员工 / 实习生，base 地点 北京/上海
我们希望你：代码能力强，对分布式系统或CUDA有较强的经验。能从第一性原理出发，独立发现系统关键瓶颈并做出设计决策。
欢迎投递简历到 xuxinran@moonshot.cn
Part 2
（1）在月之暗面做月饼，Kimi 以 KVCache 为中心的分离式推理架构
Disclaimer：和论文不同本文夹带过量私货个人观点，不代表本人单位更不代表 Moonshot，纯技术讨论。
TL;DR
Mooncake 是由月之暗面创建的超人气智能助手 Kimi的底层推理平台。
本系列则是对应技术报告的插科打诨浓缩版介绍。和强调 inclusion 四平八稳啥都讲一点的论文格式不同，这里更多的是想讨论一些当前还未形成共识的一些 design choice （私货警告）。特别的本篇主要讨论 Mooncake 的分离式架构，讨论点包括但不限于 TBT or TPOT，Prefill 节点应不应当独立存在以及独立的话如何多节点并行处理，KVCache cache 全局调度的原理和调度策略，Decode 还能不能进一步分离等等。
总体架构
Mooncake 的架构是非常典型的分离式架构，将单个同构 GPU 集群的资源打散并重新组织成三个可以独立弹性伸缩的资源池。其中 Prefill Pool 处理用户输入，主要对 Time To First Token (TTFT) 负责。同时因为 Prefill 相对计算密集，这一部分也承担着抬高整体资源利用率的任务。Prefill 处理完之后对应的 KVCache 会被送到 Decode Pool 进行 autoregression 式的流式输出。虽然我们希望尽可能攒大的 batch 以提升 MFU，但这一部分主要需要对 Time Between Tokens (TBT) 负责。
Mooncake 架构
这里我们使用 TBT 而非另外一些工作中常用的 Time Per Output Token (TPOT) 的原因在于虽然其实本来这两者应该等价，但实际计算 TPOT 的时候经常会被简单地等价成计算 average TPOT。然后就直接用了吞吐的倒数或者说总生成时间除以总生成 token 数。TBT 的定义则更加明确的指定为两个 token 生成间的延迟，因此基于 TBT 上限设计 Service Level Objective (SLO) 能更好的反映流式交互时的用户体验。直观的理解的话大家可以想象如果 GPT-4o 的语音回复是说一句停顿一段时间再说一句的说话方式的话，即便 average TPOT 能小一些体验还是不够好。
在 Prefill/Decode 之外我们还利用每台 HGX 机器上组成了一个 KVCache Pool 来进行全局的 Prefix Cache。相比 vLLM 当前单机的 Prefix Cache 通过全局的调度能够大幅度提升复用率从而提升总吞吐。由此带来了一系列如何调度，分配，复制 KVCache 的问题，而且必须和 Prefill/Decode 的调度 co-design，因此我们把 Mooncake 的架构称之为以 KVCache 为中心。
在此基础上接下来我们分别讨论 Prefill/KVCache/Decode Pool 的关键 design choice。
Prefill Pool
论文这一章我原本起了个“Prefill: To Separate or Not. Is it a Question?” 的名字。主要原因在于由于 Chunked Prefill 的提出实际上 Prefill 和 Decode 节点的边界已经模糊了起来。不过除非是特别短的 prompt （可以直接一次性加入到 decode 的 continues batch 里提升 MFU 但不违反 TBT SLO 的那种）不然我们还是倾向于使用独立的 Prefill 节点。原因主要是两方面。
首先我们定义了一个 VRAM occupation cost 的概念，即 KVCache 大小乘以其在显存中留驻的时间。使用 chunked prefill 的话由于每个 chunk 还需要和其他 decode request batch 在一起处理因此是显著增加了对应 KVCache 在显存中留驻的时间的，因此等效的降低了宝贵显存空间的利用率。当然显存空出来了有没有地方用就是另外一个问题了。这里我们在论文 4.2 节中主要讨论了和未来 batch job 结合的可能性。
另外更重要的一个原因则主要是在于长文本场景下我们需要在 Prefill 集群中使用特殊的多节点分布式划分方法来压低 TTFT。这里我们讨论了相比现在被讨论的更多的 SP 策略直接使用 Chunked PP 策略在调度和减少通讯量方面的优越性，前者大幅度简化架构，后者则能够省出宝贵的传输带宽来搬运 KVCache。具体策略其实就是 TeraPipe 的推理简化版本，因为只有 forward 没有 backword 的话不需要啥动态规划来平衡每个 Chunk 的计算量，只要给每个 Chunk 设置一个最小的值就行。
KVCache Pool
都 2024.6 了单节点的 Prefix Cache 已经不是什么新鲜东西了。这里主要讨论一下全局调度能够成立的底层原因：对于每 X byte 的 KVCache，其生成所需的算力正比于 X*hd 再乘以一个较大的常数，这里 hd 对应模型的 hidden dimension。因此只要每卡算力比如 A100 的 220TFLOPS 和每卡通讯带宽比如 CX7 的 100Gbps 的比值小于 hd 乘以这个常数，那么从远端传输 KVCache 相比原地重算不仅仅减少了计算量还减少了 TTFT。又省钱用户体验还更好了自然是何乐而不为。
但是上述公式中具体的 hd，常数是和模型结构强相关的所以不是说一定可以，只能说原则上比较大的模型应该都行。另外虽然网卡有 100Gbps，但是如果因为比如 KVCache 下沉到慢速存储上去了，网络拥挤了之类的原因跑不到这么高那么自然也就不一定成立。所以还是需要以 KVCache 为中心针对不同请求的 TTFT SLO 分别去调度。
在论文的 5.2 节中我们主要是提供了一种简单的基于 heuristic 的热点识别和复制方法。倒不是说有多么精妙，但重点是简单好实现，不用特别去预测未来的 KVCache 使用模式，而且在实际的使用中效果不错的方案。未来会尝试更多的调度算法不过这个可以作为一个 strong baseline 来使用。
顺带这里感慨一句，目前标准的 8 卡 HGX 服务器每台会配 9~10 张 100Gbps 乃至 200Gbps 的 RDMA 网卡，至少 2TB 的内存和一大堆的高速 SSD。假设一千台的话就是近万张网卡互连起来的数 PB 共享内存池。做 Disaggregated Memory 挺久了第一次见这样的富裕仗。可惜作为一个 bandwidth bound 场景在架构上能设计的东西没那么多特殊的奇淫技巧可用，更多还是工程实现。
Decode Pool
Decode Pool 本身业界的实践是最多的。特别是开源有极快速地在迭代的 vLLM，该有的不是已经有了也是在 roadmap 里。这里主要是夹带一下 future work 中讨论的增加面向大容量大带宽设计的高性价比设备，形成异构集群的私货。从去年开始就断断续续和很多人讨论这个话题，虽然听闻了很多的风声，不过实际正用上的设备目测还得有个两年。而且硬件毕竟迭代慢，出来了算法变了（比如 MLA）或者被老黄的大力出奇迹创飞了就很尴尬。
不过实际上我们几个月前就挂出但没多少人看的另一篇论文也提到，其实不需要等新的硬件，只要组合基于 GDDR 的设备其实就能将 Decode 部分进一步拆分成 attention 算子和 Linear 算子两个池子从而进一步提升整体性能。不过目前云上的 GDDR 设备一般不会和旗舰显卡集群放在一起，而且为了节省成本只会配很差的 CPU 和网卡，所以想要实际部署这一套方案还有很多挑战。而且当前设备毕竟不是专门为这个场景设计的，attention offload 策略性价比的提升的潜力没有办法充分发挥出来，为此进一步将架构复杂化而且增大一定的 TBT 不一定很合适。所以归根结底还是得看看有没有硬件厂商愿意赌一把了。
Part 2
（2）Kimi “泼天的流量”怎么接，分离架构下基于预测的调度策略
Disclaimer：和论文不同本文夹带过量私货和个人观点，不代表本人单位更不代表 Moonshot，纯技术讨论。
TL;DR
正如 @许欣然 所说，本文公开的目的之一就是推动硬件厂商和云厂商向前篇提到的分离式，乃至未来异构分离的方向演化。作为我个人已知范围内承载了最大规模在线流量的分离式推理系统，Mooncake 也积攒了一些独特的实践经验。下篇这里主要是想要讨论其中非常关键的一项：Kimi 在承接急速增长的需求时，面向过载（Overload-oriented）场景的调度策略。
Kimi 也会累
和一般按照峰值容量规划，调度主要解决如何尽可能利用闲置资源的传统工作不同，Kimi 自从破圈之后面临的是完全不同的，如何在不断扩容还天天过载的情况下尽可能保障用户体验的问题。虽然是一个非常凡尔赛的甜蜜烦恼，同时 infra 也已经在最大程度避免，但在高峰期不可避免地还是需要偶尔向大家喊累。
这里主要是我们发现由于当前的负载情况已经无法按照 SLO 提供响应，所以与其提供一个完全无保障的长时间排队（违反 TTFT SLO）或者一字一顿的慢速回答（违反 TBT SLO）这样极大伤害体验的服务，不如先暂停一小部分服务。至于为什么是喊累而不是显示 429 Too Many Requests，当然是因为 Kimi 是大家的智能*伙伴*，伙伴累了要休息一下希望大家可以体谅。
分离式架构在过载场景下的挑战
也正是在这样的环境下正式上线了 Mooncake。虽然 Mooncake 大幅提升了 Kimi 能够承载的总吞吐，但 Prefill 和 Decode 分离的架构在过载调度上也的确引入了一些新的挑战。这里最重要的问题就是 Prefill 和 Decode 是分别对 TTFT 和 TBT 负责的，而且有一个关键的时间差。所以在坏的情况下 Prefill 前调度器觉得可以同时满足 TTFT 和 TBT，但是 Prefill 后由于 Decode 集群的过载 TBT 无法满足了。这样就进入了一个是违反 SLO 还是浪费已经花费了的 Prefill 资源的两难问题。
为了解决上述问题，一个自然的，也是被 Mooncake 采用的解决方案就是同时综合 Prefill 和 Decode 两个集群的情况，然后以两者中更高负载更危险的那个集群为去判定是否可以接受服务。由于会因为在 Prefill 集群有空闲的情况下由于未来 Decode 集群的负载问题提前拒绝一些服务，这样的策略被我们称之为 Early Reject。
看起来很直接，但实际部署之后我们观测到集群负载出现了奇怪的颠簸现象。可以看到 Prefill 和 Decode 集群的负载就和跷跷板一样一边上去一边下来，然后交替。
仔细分析之后发现的原因下图有一个很直观的展示。由于 Prefill 和 Decode 集群负载之间的时间差，如果简单的参考当前 Decode 集群负载去拒绝请求的话会导致 Decode 集群负载被消化的时候 Prefill 没有及时跟上，由此产生了跷跷板效应。
基于预测的调度
为了解决上述的问题，我们进一步设计和上线了基于预测的调度策略。原理也很直观，能够预测未来（特别是新加入的请求完成 Prefill 阶段的时刻）Decode 集群负载的话自然就可以打消时间差的问题平滑整个系统的负载。
具体来说，对于未来的某个时刻 t，首先我们将现有的 Prefill 集群中 t 时已完成的请求加入 Decode 集群，这一点可以通过预测 Prefill 时间和排队时间来实现。然后，我们假设每个请求的 decode 时间均为 t_d，将 Decode 集群中 t 时已结束（即decode 时间超过 t_d）的请求移除。这一步我们也在尝试通过预测每个请求实际会输出的 token 数做一个更精确的预测。最后，我们利用事先拟合的模型根据 Decode 集群中 t 时仍在处理的请求预测 TBT，作为 Decode 集群的预测负载并根据这个数据决定是否需要 Early Reject。
未来：调度会越来越复杂，也会越来越重要
从目前的趋势来看，未来 LLM Serving 的负载会愈发的复杂和多元化。原本同等地位的请求会在厂商和用户的驱动下基于多种多样的诉求进一步分化。比如有些批处理请求可以完全不考虑 TTFT/TBT 仅仅需要保障 Job Completion Time (JCT) SLO，有些高级/付费用户无疑需要保证更加可靠和更高的 SLO。另外外部基于 Agent 的任务也不仅仅看单个 request，而是看整个 Workflow 的综合 SLO，加上 Code Executor 等外部环境的交互后也会进一步复杂化整个系统。
在这些领域上我们也是刚刚起步，探索了一些，但是需要研究的更多，也希望借此机会和大家更多的探讨这一部分未来可能的优化空间。
| |
热门跟贴