Kimi的研究论文揭示了一种关键的推理架构，该架构承担了80%的流量处理任务。

编辑日期：2024年07月05日

近期，清华KVCache.ai团队与月之暗面的研究者们共同发表了一篇开创性的论文，首次揭示了备受瞩目的国产大模型Kimi的推理体系结构！Kimi作为人气爆棚的明星模型，一直以来都是关注的焦点，流量不断，甚至时常因高需求而超负荷运行。

随着论文的公开，人们关注的焦点——如何妥善应对这股庞大的流量——也揭晓了答案。

"Mooncake"，Kimi的推理架构，以其独特的分体式设计著称。在规划阶段，Mooncake就预见到可能面临的高流量挑战，并为此进行了专项开发。据测试，该系统在模拟环境中能提升525%的吞吐能力，而在实际应用中，它能有效处理超过75%的请求。据许欣然，月之暗面工程公司的副总裁，在知乎上的一篇文章透露，Kimi超过80%的流量都依赖于这个系统。Mooncake体系结构的关键在于其KV缓存设计，主要用于存储键值对，这种设计便于快速访问和检索数据，尤其在大型模型中，能加快推理速度并降低计算资源的消耗。由于团队预测KV缓存将长期维持高容量，因此对其优化成为了核心策略。

Mooncake架构包含四个核心组成部分：全局调度器（Conductor）、Prefill节点集群、Decoding节点集群以及分布式KVCache池，还配备有RDMA通信模块（Messenger）。全局调度器作为系统的第一道关卡，接收用户请求，并依据KV缓存的分布状态及负载状况，智能地将请求分发至Prefill和Decoding节点。

调度策略兼顾KV缓存的复用深度和负载均衡，旨在最大化缓存复用效率。Mooncake实施了一种基于启发式的热点KV缓存自动迁移策略，无需精确预判未来访问模式，即可动态复制热点缓存块。

这种动态复制机制对于维持负载平衡至关重要。实验证明，Mooncake的调度策略相比于随机调度和负载均衡调度，能显著减少TTFT（首次Token获取时间），显著提升系统性能。

调度完成后，任务将被分配给Prefill和Decoding节点分别进行计算。

当Prefill节点收到调度器传递的请求，它会从KV缓存库中检索数据，执行预计算操作，并创建新的KV缓存。

面对需要处理长上下文的请求，Mooncake采取分块流式并行的方法，利用多个节点的并行处理以减少延迟。

Decoding节点在接收到调度器的指令之余，也会获取Prefill阶段产生的KV缓存。随后，节点对这些缓存进行解码，最终生成结果。

Mooncake系统的一大特色是其内置的大规模、高效能KV缓存，这得益于专门的缓存池支持。在节点间的KV缓存数据传输中，RDMA通信模块凭借其卓越的高带宽和低延迟特性，确保了流畅进行。

Mooncake的另一显著特征是其采用的分体式架构设计。这一设计的出发点主要源于Prefill和Decoding两个处理阶段在计算性质上的显著差异。

具体而言，Prefill关注TTFT（Time To First Token，首次获取Token的时间），而Decoding阶段则聚焦于TBT（Time Between Tokens，Token间延迟）。因此，这两个阶段在计算难度、内存访问模式、并行处理的程度以及对延迟的容忍度上都有着本质的不同。

因此，月之暗面团队对GPU集群执行了细致的拆分策略，旨在将它们分散到各个独立的节点集群，以此达成资源的隔离与专项优化。

此外，Mooncake的KV缓存池采用分布式设计，巧妙地利用了GPU集群中未被占用的CPU、DRAM和SSD资源，构建出大容量、高吞吐量的KV缓存系统，有效减少了资源闲置的情况。

尽管Mooncake的分离式架构提高了效率，然而面对实际操作中的海量流量，系统仍面临重大挑战。为解决这一问题，作者提出了一种创新的应对方法。在高负荷情况下，调度的核心决策在于是否接纳新请求。利用Mooncake的分布式特性，可实施预填充阶段的早期拒绝策略，依据解码节点的负载状况，在早期就拒绝部分请求。系统以TTFT和TBT的SLO达标情况作为衡量负载的标准，其中TTFT的P90值需控制在空载处理时间的10倍以内，而TBT的P90值则限制在5倍。此早期拒绝策略有效降低了无谓的预填充计算，提升了资源使用率，不过它也带来了新问题——预填充和解码节点的负载波动，这可能导致资源利用率下降，进而影响系统性能。

这源于早期拒绝策略的决策延迟问题，如图所示：

为解决此问题，月之暗面团队提出了一种改良的早期拒绝策略，即预测性早期拒绝，以减小节点负载的波动。该策略的核心理念在于预测未来Decoding节点的负载状态，并依据预测结果来判断是否拒绝新请求。

预测可分别在请求级和系统级进行。尽管请求级预测复杂，需估算单个请求的执行时长，但系统级预测较为简便，只需预测整体系统负载。

Mooncake采取了一种简化的系统级预测方案，假设每个请求的处理时间遵循特定的分布规律，以此预测未来的系统负载状况。

实验证明，运用这种基于预测的早期拒绝策略，能显著减轻负载波动现象。

综合评估显示，Mooncake的架构与优化方案切实提升了推断服务的效能，尤其在处理复杂情境和长背景信息时表现出更强的优势。

在ArXiv Summarization及L-Eval数据集的测试中，Mooncake的吞吐量相较于基础模型vLLM分别实现了20%和40%的显著提升。

月饼模型在模拟数据集上的性能峰值可达到525%的吞吐量提升，而在实际数据集上，相较于vLLM，它能处理的请求量增加约75%。

在高负荷情况下，评估表明应用预测性早期拒绝策略后，被拒绝的请求从基准的4183个减至3589个，这表明系统处理请求的能力有所增强。

展望未来，清华大学计算机科学系的助理教授章明星认为，随着形势演变，大模型的运营将面临更错综复杂和多样的需求，调度任务将日益变得关键且复杂。

对于项目"月之暗面"的未来发展路径，许欣然提供了洞见——采取分布式策略将引导该系统在"计算能力/成本"和"带宽/成本"两个维度上独立演进，这将有利于实现对硬件的更高效优化。

相关资源： - 论文：https://arxiv.org/pdf/2407.00079 - GitHub仓库：https://github.com/kvcache-ai/Mooncake - 参考链接： [1]https://zhuanlan.zhihu.com/p/705910725 [2]https://zhuanlan.zhihu.com/p/706204757