前言
本文依托论文《Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving》来讲解kimi的后端服务架构Mooncake。按照自己的思路来梳理论文中的一些关键信息。
背景
服务端面临的问题
随着大模型技术越来越强，很多应用都是以Maas（Model as a Service)的方式对外提供服务，服务端的能力受模型的能力约束。对于C端应用来说，期望的服务端优化目标一般时最大化有效吞吐量，以保证大量用户的正常体验。这种优化进行的同时一般也要满足SLO（服务等级目标）。在大模型应用的背景下，这里的目标通常是首次生成令牌的时间（TTFT）和令牌之间的时间（TBT） 。一般优化的目标期望在这些点之间进行trade off。
常见的解决方案
对于最大化吞吐量的目标而言，目前的优化方式一般是：
- 尽可能多地重复使用 KVCache，以减少所需的计算资源（比如之前文章中提到的PageAttention,RAGCache等）。
- 最大化每个批次中的令牌数量，以提高模型浮点运算利用率（MFU）。
虽然这两种方式一般可以提高吞吐量，但在SLO上可能会有所下降。在分布式推理的场景下，从远程复用KVCache会增加TTFT，而大的batch则会增加TBT。如果期望在最大化有效吞吐量的情况下同时保证SLO等在预期的约束下，那么就需要有更好的方式来对系统进行设计。
Mooncake架构
Mooncake是Kimi为了解决上述问题而提出的服务后端架构。它是一个以KVCache为中心的解聚合架构。优化思路：将处理流程分为prefill阶段和decode阶段。不同阶段的资源使用和优化目标不同。prefill阶段主要进行KVCache的计算，主要优化TTFT； decode阶段主要进行后续token的生成，主要优化TBT。 它的架构整体如下：
在这这个架构下，大模型的处理被划分为了prefill和decode两个阶段。对于每个阶段，分别使用Scheduler对各自的实例进行调度。在这个框架下，一个请求的处理流程如下：
- 在请求来临时，KVCache-centric Conductor会分配一对Prefill和Decode实例用于后续该请求的处理。
- 请求被分发到指定的prefill实例上进行全量的KVCache的计算。
- 将计算完成的KVCache传递到decode实例。
- 将请求加入decode batch并在decode实例上进行进行，生成最终结果。
接下来看一下这个流程中一些更具体的细节。
KVCache调度
KVCache的Block被存储在CPU中的pool里面。当需要在GPU中使用到KVCache或者需要将KVCache传送到远程节点时，使用一个基于RDMA的组件Messenger来完成。这里的hash是基于之前块的hash结果和当前块的内容，这样做的目的是为了防止重复。
Prefill调度
在prefill阶段，需要尽可能的复用之前已经生成的KVCache。这就带来一个思路：选择prefill实例的时候，可以选择有更多可复用KVCache的实例。但是完全依赖这个策略可能会引起负载的不均衡，进而导致SLO不达预期。因此在进行prefill实例选取的时候，除了考虑可复用的KVCache的大小，还要考虑实例的负载情况。在prefill阶段，可能需要进行chunking，在生成KVCache的过程中，不同chunks会并行执行，执行完成之后会会流式的传递到decode实例中。
在这个过程中，如果没有cache的token数量超过了一个阈值，那么prefix阶段将请求内容划分成多个chunk并且流水线执行。chunk能够提高decode的计算密度，增强MFU。
Decode调度
在解码阶段，请求会被调度到decode实例上去进行decode。此时当前请求会被动态加入到当前的某个batch中去进行后续计算。decode实例会在GPU加载已经完全传输到本地CPU DRAM中的全量KVCache信息, 并进行后续的计算。
Construct全局调度
全局调度器Construct来完成整体零散组件的调度。Construct的职责是：基于KVCache分布和负载来调度请求。同时也会进行一些KVCache block的复制和交换。
过载保护
Early Rejection
如果请求超过了服务的能力，那么服务会处于过载状态，进而导致服务的SLO下降。在这里最简单的策略是直接丢弃。最简单的策略是Early Rejection:Conductor会基于prefill pool和 decode pool较大的那个负载来决定是否丢弃当前请求。这种方式通过拒绝请求的方式减少了无用计算，增强了负载均衡能力。但直接这样应用会带来一个负载波动问题。这种负载波动问题的根源在于预测解码负载与实际执行之间的时间差。示例如下：
基于预测的Early Rejection
为了解决这个问题，Mooncacke中使用了一种基于预测的Early Rejection的方法。这种方法下，会去预测prefill下一阶段decode的负载，并基于这个预测来决定是否拒绝当前请求。这里的核心是预测模块。在Mooncake中，使用的是系统级的预测：系统级预测并不试图预测单个请求的完成时间，而是在指定时间后估计实例的总批次计数或总令牌时间（TBT）状态。这种类型的预测是持续进行的，并且需要较低的精度，因此更适合过载场景。
总结
本文从面临的问题出发，逐步深入了解mooncake的框架内部组成以及组件之间处理请求的workflow。其中一些原理涉及分布式推理，这部分内容暂时还没有深入了解，可能理解上会有一些偏差，感兴趣的同学可以了解一下。目前计划中暂不深入分布式推理，如果后续深入这一块，再回过头来进行补充。
