如何利用Kimi解读Kimi的KVCache技术细节

最近Kimi公布了一篇Mooncake: Kimi's KVCache-centric Architecture for LLM Serving的文章，详细介绍了Kimi背后的推理架构，因此笔者想到用Kimi解读Kimi，梳理相关技术要点如下，供大家参考：

文章 "Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving" 介绍了 Mooncake 这一为大型语言模型(LLM)服务而设计的服务平台。以下是从文章中提取的技术要点、解决的问题以及提升效果的总结：

解决的问题

高负载挑战：Mooncake 面临由于高度超载场景带来的挑战，这与传统研究不同，后者假设所有请求都会得到处理134。
早期拒绝策略：为了缓解高负载问题，开发了基于预测的早期拒绝政策13。

Mooncake 架构的核心特点：

KVCache-中心化架构：采用了解耦的预填充（prefill）和解码（decoding）集群，以KVCache为中心进行调度和优化。
资源的充分利用：利用GPU集群中未充分利用的CPU、DRAM和SSD资源，通过分布式KVCache实现高效的缓存机制。
KVCache中心化调度程序：开发了一个调度程序，用以平衡整体吞吐量的最大化和满足延迟相关的服务级别目标（SLOs）。

架构设计动机与目标：

多样化工作负载：应对LLM服务中的多样化工作负载，包括不同的输入/输出长度、请求到达频率和分布。
优化问题：作为模型即服务（MaaS）提供商，Kimi旨在在满足不同SLOs的约束下，最大化整体有效吞吐量。

架构实施细节：

预填充池的实现：介绍了如何实现一个独立的预填充节点池，并通过分块管道并行性（Chunked Pipeline Parallelism, CPP）机制来处理长上下文输入。
KVCache中心化请求调度算法：详细描述了如何进行基于KVCache的请求调度，包括自动化热点迁移方案，以在不精确预测未来KVCache使用情况下复制热门KVCache块。
过载导向调度：讨论了在过载情况下如何决定接受或拒绝传入请求的策略，包括基于预测的早期拒绝政策。

实验结果：

性能提升：在长上下文场景中，与基线方法相比，Mooncake在某些模拟场景中实现了高达525%的吞吐量增长，同时遵守SLOs。
实际工作负载下的表现：在真实工作负载下，Mooncake的架构使Kimi能够处理多75%的请求。

未来工作方向：

从目前的趋势来看，未来 LLM Serving 的负载会愈发的复杂和多元化。原本同等地位的请求会在厂商和用户的驱动下基于多种多样的诉求进一步分化。比如有些批处理请求可以完全不考虑 TTFT/TBT 仅仅需要保障 Job Completion Time (JCT) SLO，有些高级/付费用户无疑需要保证更加可靠和更高的 SLO。另外外部基于 Agent 的任务也不仅仅看单个 request，而是看整个 Workflow 的综合 SLO，加上 Code Executor 等外部环境的交互后也会进一步复杂化整个系统。

异构加速器的潜力：探索使用异构加速器，以及新技术如存内计算（Process-in-Memory）或混合键合（Hybrid Bonding）技术。
KVCache压缩算法：研究减少KVCache大小的算法，以提高缓存效率和吞吐量。
调度策略：开发考虑不同请求优先级和SLOs的高级调度政策。

Mooncake架构概览

Mooncake的架构革新在于其KVCache中心化的设计理念。这种设计允许系统将预填充(prefill)和解码(decoding)任务分配到不同的集群，从而实现资源的最优分配和任务的高效处理。预填充集群负责生成模型的初始状态，而解码集群则利用这些状态来处理具体的请求。

架构图示：

实验表明，Mooncake 在长上下文场景中表现出色。与基线方法相比，Mooncake 在某些模拟场景中可以实现高达 525% 的吞吐量提升，同时遵守 SLO。在实际工作负载下，Mooncake 的创新架构使Kimi能够处理 75% 以上的请求。

从结构上看，Mooncake由全局调度器（Conductor）、Prefill节点集群、Decoding节点集群和分布式KVCache池几部分组成，另外还有RDMA通信组件（Messenger）。

其中全局调度器是用户请求到达系统后的第一站，它负责接收请求并根据KV缓存分布和负载情况，将请求调度到Prefill和Decoding节点。

调度器在调度时需要综合考虑KV缓存的复用长度、负载均衡等因素，实现KV缓存复用的最大化。

具体到Mooncake，它采用了一种启发式的自动热点迁移策略，可以在不需要精确预测未来访问的情况下自动复制热点KV缓存块。

同时，这种动态复制热点KV缓存块的方式，也是实现均衡负载的一种重要途径。

实验结果表明，与随机调度和负载均衡调度相比，Mooncake的调度策略可以显著降低TTFT（Time To First Token，首个Token延迟），提高系统性能。

Mooncake的另一个创新之处在于对GPU集群中未充分利用的资源的再利用。通过智能地分配任务到CPU、DRAM和SSD，Mooncake承载了 Kimi 线上80% 以上的流量，不仅提升了硬件的使用效率，还增强了系统的处理能力。

Mooncake的调度程序是整个架构的核心。它负责在保证服务级别目标(SLOs)的同时，最大化整体有效吞吐量。调度程序采用了先进的算法来预测请求的负载，并动态调整资源分配。

高负载挑战与解决方案

Mooncake 采用的解决方案就是同时综合 Prefill 和 Decode 两个集群的情况，然后以两者中更高负载更危险的那个集群为去判定是否可以接受服务。由于会因为在 Prefill 集群有空闲的情况下由于未来 Decode 集群的负载问题提前拒绝一些服务，这样的策略被我们称之为 Early Reject。

看起来很直接，但实际部署之后我们观测到集群负载出现了奇怪的颠簸现象。可以看到 Prefill 和 Decode 集群的负载就和跷跷板一样一边上去一边下来，然后交替。

在高负载情况下，传统的服务架构往往会遇到性能瓶颈。Mooncake通过实时监测系统负载，并评估当前的资源使用情况，来识别潜在的高负载场景。

为了应对高负载带来的挑战，Mooncake开发了一种基于预测的早期拒绝策略。这种策略通过预测模型来评估即将到来的请求，并在系统达到临界负载之前拒绝一些请求，从而避免系统过载。

性能提升的实证分析

Mooncake的性能提升在模拟场景中得到了验证。通过在不同的负载条件下测试系统，Mooncake显示出了其在处理大规模请求时的卓越能力。

吞吐量图表：

在实际工作负载下，通过优化资源分配和调度策略，Mooncake使得Kimi能够处理更多的请求，同时保持了服务的稳定性和响应速度。

请求处理能力：

长上下文场景的优势

在人工智能领域，尤其是自然语言处理（NLP）中，长上下文处理能力是衡量一个语言模型性能的关键指标之一。长上下文意味着模型需要理解和记忆大量的信息，以便生成准确和连贯的输出。这对于诸如文档摘要、机器翻译、问答系统等应用至关重要。

应用场景示例：

文档摘要：自动生成长篇文章或报告的摘要。
机器翻译：处理包含复杂句式和专业术语的长文本。
问答系统：在对话中理解并引用之前的交流内容。

Mooncake的架构特别适合处理长上下文场景，这得益于其高效的资源管理和调度策略。通过分离预填充和解码任务，Mooncake能够为每个请求分配足够的上下文信息，同时保持高性能。

由于 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。

技术实现细节

KVCache的实现机制

整个Mooncake系统设计的核心，是围绕着KV缓存展开的。（KV缓存用于存储键-值对（Key-Value Pairs），主要优势在于可以简单高效地访问和检索数据，在大模型当中可以提高推理速度并减少计算资源消耗。）之所以这样做，是因为团队预计KV缓存的容量会长期保持高位，因此围绕KV缓存进行优化十分必要。

KVCache是Mooncake架构中的一个关键组件，它负责存储和管理请求的上下文信息。KVCache的设计允许快速读写操作，这对于处理长上下文请求至关重要。

在 Prefill/Decode 之外我们还利用每台 HGX 机器上组成了一个 KVCache Pool 来进行全局的 Prefix Cache。相比 vLLM 当前单机的 Prefix Cache 通过全局的调度能够大幅度提升复用率从而提升总吞吐。由此带来了一系列如何调度，分配，复制 KVCache 的问题，而且必须和 Prefill/Decode 的调度 co-design，因此我们把 Mooncake 的架构称之为以 KVCache 为中心。

在此基础上接下来我们分别讨论 Prefill/KVCache/Decode Pool 的关键 design choice。

在这个架构中，每个组件都有其特定的职责，通过高度模块化的设计，Mooncake能够灵活地扩展和维护。Mooncake 的核心是其以 KVCache 为中心的调度程序，它在最大化整体有效吞吐量和满足与延迟相关的服务级别目标 (SLO) 要求之间取得平衡。与假设所有请求都会得到处理的传统研究不同，Mooncake 面临着高度超载场景带来的挑战。为了缓解这些问题，我们开发了一种基于预测的早期拒绝策略。

完成调度之后，任务会分别交由Prefill和Decoding节点进行运算。Prefill节点接收到调度器转发过来的请求后，会从KV缓存池中读取缓存，执行预计算并生成新的KV缓存。对于长上下文请求，Mooncake还会分块流水并行的方式，使用多个节点并行处理来降低延迟。而Decoding节点除了接收调度器发来的请求外，还会收到Prefill阶段生成的KV缓存，节点会对这些缓存执行解码并生成最终结果。

大容量、高性能的KV缓存存储由缓存池提供；RDMA通信组件则凭借其高带宽、低延迟的优势，负责在不同节点之间的KV缓存传输。除了采取以KV缓存为中心的工作流程外，Mooncake还有另一个重要特点——分离式的架构。采取分离式架构的重要因素之一，是在于Prefill和Decoding两个阶段的计算特性差异很大。具体来说，它们分别要对TTFT和TBT（Time Between Tokens，Token间延迟）负责。这就导致了两者在计算复杂度、内存访问方式、并行粒度和对延迟的敏感度上都存在差异，所以，月之暗面团队对GPU集群也进行了相应的拆分，以便将它们分别部署在不同节点集群上，实现资源隔离和专门优化。

另外，Mooncake中的KV缓存池也是分布式的，同时充分利用了GPU集群中空闲的CPU、DRAM和SSD资源，实现了大容量、高带宽的KV缓存存储和传输，同时也减少了闲置资源的浪费。