极市平台 | 月之暗面kimi底层推理系统方案揭秘

本文来源公众号“极市平台”，仅用于学术分享，侵权删，干货满满。

原文链接：月之暗面kimi底层推理系统方案揭秘

极市导读

本文介绍了月之暗面科技有限公司开发的Kimi AI的底层推理系统方案Mooncake，以及在过载情况下的调度策略，包括早期拒绝和基于预测的早期拒绝策略，以提高资源利用率和系统稳定性。 

太长不看版（作者大佬自己的在知乎碎碎念）：本论文与很多 Prefill/Decoding 分离的论文不同的是，这套方案已经在大规模集群上进行几个月的验证并证明了方案的有效性。目前这套系统承载了 Kimi 线上80% 以上的流量， 效果很好也为产品带来了更多的设计空间。这也是为什么一个 POC 写在所有业内论文之前的系统，直到今天才发布出来跟大家见面。系统是需要跟随着应用快速变化的，同时也需要 硬件厂商 和 云厂商 早点接受新的理念才能跟上浪潮。

系统是需要跟随着应用快速变化的，同时也需要 硬件厂商 和 云厂商 早点接受新的理念才能跟上浪潮。发出这篇论文，主要是希望给各方提供一些信心，提供一些推理规模足够大场景下的必然优化思路。

趁这个机会，希望给各家硬件厂商和云厂商一些“暴论”

1.Mooncake 这类的存算分离策略会是一个长期趋势。

• 现在、立刻、马上真能省很多钱（毕竟不能公开规模和每日请求的 pattern，如果你说省不了那你都对）。

• KVCache 的容量会长期保持高位，因此围绕着 KVCache 来优化是非常有必要的。"private memory per request” 是整个推理系统优化中的关键瓶颈（否则 groq 万岁），会有很多努力来降低 KVCache 的大小，但同时会有更多动力来增大。

• 分离之后，允许整个系统往 “算力/”和“带宽” 的两个方向独立发展，对硬件优化是更友好的。AR 的模型架构惯性在短期内难以颠覆，因此总可以认定 decoding 的成本总会跟 bandwidth 成本有非常强的正相关性质。“带宽/低一个数量级的硬件方案已经在肉眼可见的范围内了。纵观历史，“算力/” & “带宽/$” 同时最优的芯片似乎还没出现过，集群上必然要拆分成异构的两个部分。

*2.Mooncake 方案和 MLA、各种 KVCache 的压缩方案 都是正交的，KVCache 变小了意味着 Mooncake 的方案收益会更明显。

*3.的确，有很多方向可能会让 Mooncake 的架构变成没必要的方案。

• given hardware lottery，新架构演进会是一个相当缓慢的过程（不 AR、用 RAG 做 Attention 等等方案），不能因噎废食。

• 包括我们自己也在投入很多资源在 break 现有框架上，因此有理由相信在可见的未来推理方案还会变动。

• 由于目前海量的推理压力，所以软件系统做为一个迭代速度极快的方案，就应该一代模型一次跟进。

• 预测这个状态至少会持续2~3年，因此集群层面现在已经值得做拆分了。

• 芯片层面值得做为一个重要的设计考量，在芯片的 IO 能力上要多预留一些能力。

Github地址：https://github.com/kvcache-ai/Mooncake/tree/main

原文（powered by kimi/互相赋能就让它自己宣传自己吧）

1 引言

1.1 动机：开发 Mooncake 的原因

随着大型语言模型（LLMs）在各种场景中的迅速采用，对LLM服务的工作负载已经变得显著多样化。这些工作负载在输入/输出长度、到达频率和分布上有所不同，最重要的是，它们需要不同种类的服务级别目标（SLOs）。作为一个模型即服务（MaaS）提供商，Kimi [5] 的一个主要目标是解决一个具有多个复杂约束的优化问题。优化目标是在满足不同级别的SLOs的约束下，最大化整体有效吞吐量，这直接影响收入。

为了实现这一目标，一个先决条件是充分利用GPU集群中可用的各种资源。具体来说，尽管当前GPU服务器以高度集成的节点形式提供（例如，DGX/HGX超级计算机 [6]），但有必要将它们解耦并重新结构化为几个分离的资源池，每个池针对不同但协作的目标进行优化。例如，许多研究人员 [7, 8, 9] 已经建议将预填充服务器(prefill servers)与解码服务器(decoding servers)分开，因为LLM服务的这两个阶段具有非常不同的计算特性，在请求从预填充服务器转移到解码服务器时，KVCache会发生变化。

基于这个想法，我们发现KVCache的调度是LLM服务调度的核心。为了提高整体吞吐量，通常有两种一般方法：1）尽可能多地重用KVCache以减少所需的计算资源；2）最大化每个批次中的token数量以提高模型FLOPs利用率（MFU）。然而，从远程位置重用KVCache会延长首次token的时间（TTFT），而较大的批次大小会导致更大的token间时间（TBT）。因此，这两种面向吞吐量的优化的使用可能会导致违反与延迟相关的SLOs。

根据上述指导方针，我们提出了一个以KVCache为中心的分离设计，用于调度和优化。图1展示了我们当前的以KVCache为中心的分离架构，名为Mooncake。对于每个请求，全局调度器（Conductor）需要选择一对预填充和解码实例，并按以下步骤调度请求：1）尽可能多地将可重用的KVCache转移到选定的预填充实例；2）以块/层的方式完成预填充阶段，并将输出的KVCache连续地流式传输到相应的解码实例；3）在解码实例中加载KVCache，并将请求添加到连续批处理过程中以生成请求输出。

尽管这个过程看起来简单，但由于许多限制，选择策略却相当复杂。在预填充阶段，主要目标是尽可能多地重用KVCache以避免冗余计算。然而，等待存储在低级存储上的KVCache可能会导致违反TTFT SLO。此外，对KVCache服务器的高需求可能会导致网络拥塞，延长等待时间。因此，Conductor还负责预测KVCache块的未来使用情况，并相应地执行调度操作，例如交换和复制。最热门的块应该复制到多个节点以避免获取拥塞，而最冷的块应该被交换出去以减少保留成本。预填充调度还受到预填充节点上DRAM空间可用性的限制，特别是当大量内存被保留用于全局KVCache池时。

相比之下，解码阶段有不同的优化目标和约束条件。目标是将尽可能多的token聚合到一个解码批次中以提高MFU。然而，这一目标不仅受到TBT SLO的限制，还受到可以包含在VRAM中的聚合KVCache总大小的限制。

更重要的是，现有的LLM服务研究假设资源充足，并专注于提高资源利用率。相比之下，当前的GPU/加速器供应有限，许多MaaS提供商面临严重的过载问题，尤其是在高峰时段。在这种情况下的调度提出