Mooncake：Kimi 的开源服务平台，由 KVCache 驱动的解耦架构

Mooncake 是Kimi的服务平台，Kimi 是由Moonshot AI提供的领先LLM服务。https://github.com/kvcache-ai/Mooncake 这个代码库存放了 Mooncake 的技术报告，未来还将用于开源相关的数据。

Mooncake 是月之暗面公司提供的领先大语言模型（LLM）服务平台Kimi的服务平台。它采用了一种以KVCache为中心的解耦架构，将预填充（prefill）和解码集群分离开来。它还利用了GPU集群中未充分利用的CPU、DRAM和SSD资源，实现了一个解耦的KVCache缓存。Mooncake的核心是其以KVCache为中心的调度器，该调度器在最大化整体有效吞吐量的同时，满足与延迟相关的服务级别目标（SLOs）。与传统研究假设所有请求都会被处理不同，Mooncake在高度过载的场景下面临挑战。

为了缓解这些挑战，我们开发了一种基于预测的早期拒绝策略。实验表明，在长上下文场景中表现出色。与基线方法相比，在某些模拟场景中，Mooncake的吞吐量可以提高多达525%，同时遵守SLOs。在实际工作负载下，Mooncake的创新架构使Kimi能够处理多75%的请求。

开发 Mooncake 的动机

随着大语言模型（LLM）在各种场景中的迅速普及，LLM服务的工作负载变得极为多样化。这些工作负载在输入/输出长度、到达频率和分布上各不相同，最重要的是，它们对服务级别目标（SLOs）的要求也不同。作为一种模型即服务（MaaS）提供商，Kimi的主要目标之一是解决一个具有多种复杂约束的优化问题。优化目标是最大化整体有效吞吐量，这直接影响收入，而约束条件反映了不同层次的SLOs，这些SLOs通常涉及满足与延迟相关的要求，主要是首个令牌时间（TTFT）和令牌间时间（TBT）。

要实现这一目标，前提是充分利用GPU集群中可用的各种资源。具体来说，虽然GPU服务器目前作为高度集成的节点提供（如DGX/HGX超级计算机），但有必要将其解耦并重组为几个解耦的资源池，每个资源池针对不同但协作的目标进行优化。例如，许多研究人员建议将预填充服务器与解码服务器分开，因为这两个LLM服务阶段具有非常不同的计算特性，其中KVCache在请求从预填充服务器移动到解码服务器时发生转换。

基于这一想法，我们发现 KVCache 的调度是 LLM 服务调度的核心。为了提高整体吞吐量，通常有两种一般方法：

1. 尽可能多地重用KVCache以减少所需的计算资源；

2. 最大化每个批次中的令牌数量以提高模型浮点运算利用率（MFU）。然而，从远程位置重用KVCache会延长TTFT，而大批量会导致更大的TBT。因此，同时利用这两种面向吞吐量的优化可能会导致违反与延迟相关的SLOs。

根据上述指南，提出了一种以 KVCache 为中心的解耦设计，用于调度和优化。图1展示了我们当前的以KVCache为中心的LLM服务解耦架构，名为Mooncake。对于每个请求，全局调度器（Conductor）需要选择一对预填充和解码实例，并按以下步骤调度请求：

1. 将尽可能多的可重用KVCache转移到选定的预填充实例；

2. 在块/层中完成预填充阶段，并持续将输出KVCache流式传输到相应的解码实例；

3. 加载KVCache并将请求添加到解码实例的连续批处理过程中以生成请求输出。

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

相比之下，解码阶段有不同的优化目标和限制。目标是在解码批处理中聚合尽可能多的令牌以提高MFU。然而，这一目标不仅受TBT SLO的限制，还受到VRAM中可容纳的聚合KVCache总大小的限制。

更重要的是，现有的LLM服务研究假设资源充足，集中于提高资源利用率。相比之下，目前GPU/加速器供应有限，许多MaaS提供商面临严重的过载问题，尤其是在高峰期。在这种情况下进行调度提出了现有研究未探讨的独特挑战。例如，我们需要预测未来的负载，并在预填充阶段后如果没有可用的解码插槽时提前拒绝某些请求，以节省计算资源。然而，这种早期拒绝策略的直接实现会意外地导致过载波动。这使我们不得不预测特定查询的生成长度并进行短期的整体负载预测以实现更好的拒绝策略。同时有必要对不同请求的优先级进行分类以实现基于优先级的调度。在本文中，我们总结了这些问题为面向过载的调度，并展示了我们的初步研究结果。

Mooncake 的设计和结果

首先概述了Mooncake的架构，包括其主要组件和处理请求的典型工作流程。然后，我们描述了其实现过程中所做的主要设计选择，特别是那些当前研究中未涵盖的部分。

我们讨论了如何实现一个独立的预填充节点池，该节点池能够无缝处理上下文长度的动态分布。我们采用了一种分块流水线并行（CPP）机制，将单个请求的处理扩展到多个节点，这对于减少长上下文输入的TTFT是必要的。与传统的基于序列并行（SP）的解决方案相比，CPP减少了网络消耗，并简化了对频繁弹性扩展的依赖。这一机制进一步辅以层级预填充，使KVCache的流传输能够重叠延迟。

接下来，我们详细介绍了以KVCache为中心的请求调度算法，该算法平衡了实例负载和以TTFT和TBT SLOs衡量的用户体验。这包括一种基于启发式的自动热点迁移方案，在不需要精确预测未来KVCache使用情况的情况下复制热点KVCache块。实验结果表明，我们的缓存感知调度可以显著降低实际场景中的TTFT。在使用公共数据集、模拟数据和实际工作负载的端到端实验中，Mooncake在长上下文场景中表现出色。与基线方法相比，Mooncake在满足SLOs的情况下，吞吐量最多可提高525%。在实际工作负载下，Mooncake使Kimi能够处理多75%的请求。

最后，与假设所有请求都将被处理的现有LLM服务研究不同，由于Kimi的用户请求迅速增长，Mooncake始终面临过载。因此，Mooncake的调度涉及根据系统负载确定是否接受或拒绝传入请求。在§6中，我们讨论了我们实施的独特早期拒绝策略，该策略在过载场景中减少了浪费的计算资源。我们进一步探讨了由直接早期拒绝引起的负载波动问题，以及如何通过预测未来负载来缓解这一问题。

Mooncake 目前是服务Kimi的主要平台，已成功应对了指数级的工作负载增长，证明了其在扩展到大型和高度过载的工作负载方面的有效性。然而，还有许多问题需要探索，这些未来方向也包括在本文中。

初步研究和问题定义

现代大语言模型（LLM）基于Transformer架构，利用注意力机制和多层感知器（MLP）来处理输入。流行的基于Transformer的模型，如GPT和LLaMA，采用仅解码结构。每个推理请求在逻辑上分为两个阶段：预填充阶段和解码阶段。

在预填充阶段，所有输入的令牌并行处理。此阶段生成第一个输出令牌，同时存储已计算键和值的中间结果，称为KVCache。然后，解码阶段使用此KVCache自回归地生成新令牌，将计算中的新键和值添加到KVCache中。预填充阶段能够同时处理输入令牌，这通常使其计算密集，短请求除外。由于注意力网络的计算复杂度随输入长度呈二次增长，而MLP的复杂度呈线性增长，预填充阶段的计算时间通常随输入长度超线性增长，如图2左侧所示。

相反，由于自回归生成的限制，解码阶段每批次一次只能处理一个令牌。这使其受到内存限制，并导致计算时间随批次大小次线性增加，如图2右侧所示。解码阶段广泛使用的一种优化是连续批处理。在每次迭代之前，调度器检查所有请求的状态，将新到达的请求添加到批次的预填充阶段，同时移除已完成的请求。

由于预填充和解码阶段的不同特性，MaaS提供商设定了不同的指标来衡量其相应的服务级别目标（SLOs）。具体来说，预填充阶段主要关注从请求到达到生成第一个令牌之间的延迟，即首个令牌时间（TTFT）。另一方面，解码阶段关注的是同一请求的连续令牌生成之间的延迟，即令牌间时间（TBT）。

作为MaaS提供商，确保满足服务协议定义的SLO指标是至关重要的。例如，TTFTP 90 = 4×这样的指标表示90%的推理请求的TTFT不会超过在相同条件下单个请求运行时间的四倍。在本文的端到端实验中，我们设定TTFTP 90 = 10×和TBTP 90 = 5×。在实际部署中，我们设定了固定的TTFT和TBT SLOs。如果监控检测到未满足的SLOs，我们要么增加推理资源，要么拒绝一些传入的请求。

然而，由于目前GPU供应紧张，弹性扩展推理集群通常不可行。因此，决定拒绝哪些请求成为面向过载调度的核心问题。我们的主要目标是在遵守SLOs的同时最大化整体吞吐量，其他研究中将这种概念称为“有效吞吐量”（goodput）。我们的方法不同之处在于，只有完全完成执行的请求才计入有效吞吐量的衡量标准。否则，所有先前消耗/生成的令牌都不计数，相应的资源被浪费。换句话说，如果某请求无法在SLO下完成其全部执行，则应尽早拒绝该请求。实现这一目标不仅涉及优化预填充和解码阶段的架构，还需要开发预测短期未来负载的能力。

Mooncake 的解耦架构概述

如图1所示，Mooncake采用了一种解耦架构，不仅将预填充节点与解码节点分开，还将GPU集群的CPU、DRAM、SSD和RDMA资源组合起来，实现了一个解耦的KVCache。这个解耦缓存利用未充分利用的资源，提供充足的缓存容量和传输带宽，从而实现高效的近GPU前缀缓存，而无需额外成本。

图3展示了KVCache块的存储和传输逻辑。在CPU内存中，KVCache以分页块的形式存储。根据请求模式，它可以使用如LRU（最近最少使用）、LFU（最不常用）或基于请求特性的缓存淘汰算法。这些KVCache块在CPU和GPU之间的传输由一个名为Messenger的独立（GPUDirect）基于RDMA的组件处理。这种架构还使我们能够向外部用户提供上下文缓存API，以提高KVCache的重用率。

为了调度所有这些解耦的组件，Mooncake在其核心实现了一个名为Conductor的全局调度器。Conductor负责根据当前的KVCache分布和工作负载分配请求。如果对未来推理有利，它还会复制或交换某些KVCache块。具体而言，图4展示了一个请求的典型工作流程。一旦标记完成，Conductor选择一对预填充节点和一个解码节点，并开始包括四个步骤的工作流程：

1. KVCache重用：选定的预填充节点（组）接收一个请求，其中包括原始输入、可重用前缀缓存的块ID和分配给请求的完整缓存块ID。它根据前缀缓存块ID从远程CPU内存加载前缀缓存到GPU内存中，以启动请求。如果不存在前缀缓存，则跳过此步骤。此选择平衡了三个目标：尽可能多地重用KVCache、平衡不同预填充节点的工作负载，并保证TTFT SLO。这导致了一个以KVCache为中心的调度，进一步讨论见§5。

2. 增量预填充：预填充节点（组）使用前缀缓存完成预填充阶段，并将新生成的增量KVCache存储回CPU内存。如果未缓存的输入令牌数量超过一定阈值（prefill_chunk），则预填充阶段被分成多个块并以流水线方式执行。选择此阈值以充分利用相应GPU的计算能力，通常超过1000个令牌。使用分块但仍解耦的预填充节点的原因见§4.1。

3. KVCache传输：上述的Messenger服务部署在每个节点中，以管理和传输这些缓存。每个Messenger作为其各自推理实例中的独立进程运行，接收信号以促进高速、跨机器的KVCache传输。此步骤异步执行，并与上述增量预填充步骤重叠，流式传输由每个模型层生成的KVCache到目标解码节点的CPU内存中，以减少等待时间。图4：推理实例的工作流程。（*）对于预填充实例，KVCache层的加载和存储操作逐层进行，并与预填充计算并行，以减轻传输开销）。（†）对于解码实例，异步加载与GPU解码同时进行，以防止GPU空闲时间。

4. 解码：在解码节点的CPU DRAM中接收到所有KVCache后，请求以连续批处理方式加入下一个批次。Conductor根据解码节点的当前负载预先选择解码节点，以确保其不违反TBT SLO。然而，本地调度器会再次检查该SLO，因为预填充阶段后的预期负载可能已发生变化。这种双重检查可能导致请求被拒绝，在这种情况下，相应的预填充成本将被浪费。

https://arxiv.org/pdf/2407.00079