kimi chat大模型的200万长度无损上下文可能是如何做到的？

这周，清华和Moonshot发了一个技术报告https://github.com/kvcache-ai/Mooncake/blob/main/Mooncake-v1.pdf，介绍Kimi背后的LLM服务系统Mooncake，它采用分离式设计，将Prefill和Decode两阶段解耦，构建了一个全局KVCache Pool，实现以Cache为中心的调度。

Moonshot作为MaaS头部厂商，以其过硬的技术产品实力和明星的团队阵容闻名于世。和其他大模型公司不一样，他们很少发技术报告或对外做技术分享。这次Mooncake技术报告，让大家得对其技术得以管中窥豹。

论文的通信作者为Moonshot的许欣然与清华大学计算机系的章明星，两位均是重量级大咖，也分别在知乎宣传了Mooncake的工作。许欣然在AISys领域深耕多年，曾执掌MegEngine，工程经验丰富，如今在Moonshot担任工程副总裁一职。章明星过去研究领域是分布式系统、图数据库，其研究成果在2017年成为大陆作者单位在OSDI上的首篇一作论文，还曾经是ACM-ICPC World Finals，如今也开始关注大模型系统领域，大模型正汇聚着各行各业最顶尖人才，共襄盛举。

ZHANG Mingxing：Mooncake (1): 在月之暗面做月饼，Kimi 以 KVCache 为中心的分离式推理架构

关于 Mooncake 的碎碎念

Mooncake分离式架构动机是Prefill和Decode阶段性质不同，Prefill是计算密集，受限算力带宽用不满，Decode是访存密集性，受限带宽算力用不满。所以用同一种硬件部署两阶段往往顾此失彼，不是最有性价比。因此，最近很多工作对二者进行拆分，和Mooncake最相似的是今年5月份发布的微软和华盛顿大学的工作Splitwise https://arxiv.org/pdf/2311.18677，它里面列出了Prefill和Decode不同的七个Insights值得大家仔细读一下。因为Mooncake开发也需要一段时间，它和Splitwise应该是不谋而合的同期工作。

拆分Prefill/Decode之后，LLM推理系统就更像一个分布式内存系统+流处理系统，这就是传统计算机系统研究者最擅长的领域。某大佬和我讲的sys三板斧，batch， cache，调度都可以招呼上。比如，Decode可以进一步拆成Attention和非Attention算子分离调度，也是章明星团队近期的一个工作叫Attention Offloadinghttps://arxiv.org/abs/2405.01814。

高屋建瓴的分析二位作者的知乎已经讲得很透彻了，本文对技术报告具体系统设计做了一些注释，一方面用费曼学习法帮助自己理解，另一方面也给大家一些参考，方便读者读Mooncake文章时交叉验证。学习过程也和

@手抓饼熊

有很多交流讨论，看英文效率比较低的话，推荐大家看

@手抓饼熊

的最新中文翻译，英文论文有可能也是中文->GPT->英文得来的，再用GPT翻译回去会没准更加原汁原味。

Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving

3 Overview of Mooncake’s Disaggregated Archtecture

如下面Figure 1所示，Mooncake采用了分离式架构。这里分离有两层含义：

一是，将Prefill与Decode计算资源分开，这与前人工作无异，如splitwise和distserve等；Prefill阶段优化目标是利用request间存在共同前缀的机会，尽可能复用KVCache，同时满足TTFT（Time To First Token） SLO，最大化MFU（论文中似乎笔误成minimum）和KVCache小于CPU内存限制。Decode优化目标为最大化吞吐，满足TBT（Time between tokens ，Decode阶段两个Token之间的时间）SLO和KVCache小于GPU显存限制。

二是，将KVCache和计算分离开，它将GPU集群的CPU、DRAM、SSD和RDMA资源分组组成Distributed KVCache Pool，KVCache也是分块以Paged方式管理，KVCache Blocks如何在Pool中调度，请求和复用KVCache乃本文精髓。

我们跟随论文的安排，跟Figure 4先走马观花过一遍Mooncake处理一个request的流程。一个request到达（tokenized之后的），调度程序会选择一对Prefill Instance和Decode Instance，并启动包含四个步骤的工作流程：

s1）KVCache Reuse：Prefill Instance将request分成token block，考虑request之间存在共同前缀，需要尽可能将token block调度到Prefix KVCache最长的节点处理，来resuse KVCache。为此作者提出一种以KVCache为中心的调度，将在章节5 KVCache-centric Scheduling中详细介绍。

s2）Incremental Prefill：使用Prefix Cache，Prefill Instance只需要计算Prefix Cache没覆盖的部分。长序列计算需要多卡并行，用TeraPipehttps://proceedings.mlr.press/v139/li21y.html方式做流水并行。将在章节5 Implementation of the Prefill Pool中详细介绍。

s3）KVCache Transfer：和Splitwise一样，分离是设计需要将KVCache从Prefill Instance搬运到Decode Instance。Mooncake通过异步传输，与上述Incremental Prefill步骤重叠，将每个模型层生成的KVCache流式传输到目标Decode Instance的CPU内存，以减少等待时间。

s4）Decoding：在Decoding Instance的CPU DRAM中接收到所有KVCache后，请求Continous Batching处理。

工作流中，s3, s4平平无奇，前人设计如Splitwise、DistServe和vLLM等已经覆盖。s1和s2比较精彩，本文接下来的章节4和章节5来详细介绍。

4 Implementation of the Prefill Pool

章明星说原本章叫“Prefill: To Separate or Not. Is it a Question?” 。这说明分离式设计并不是共识。

Prefill和Decode的计算性质不同，前者吃计算，后者吃带宽。这会带来很多次生影响，比如Batching方式就不同，Prefill不需要batching，Decode需要大Batching。处理Prefill和Decode有融合派和分离派两大流派。

融合派：将Prefill放到Decode step间隙或者和某个Decode step一起推理。2022年OSDI的LLM Continous Batching 开山之作Orca将Prefill和Decode Step融合在一个Batching Step做forward，Orca时代还没有Paged Attention，还需要Selective Batching来将Attention解Batching。2023年vLLM做Batching时，prefill和decoding则是独立forward的，一个Batching step要么处理decoding要么处理prefill。prefill直接调用xformers处理计算密集的prefill attn计算；decoding手写CUDA PA处理访存密集的attn计算。后来，以Sarathihttps://arxiv.org/pdf/2401.08671和FastGenhttps://arxiv.org/pdf/2403.02310为代表，将prefill序列拆成chunk，chunk prefilling可以插入到decode phase中，甚至和decode step融合。比如，flash attention的flash_attn_with_kvcache函数支持q部分有kvcache（decode）部分没有（prefill）。章明星也提到，对于短的prefill chunk和decode step融合和纯做decode step延迟可能相同，这相当于prefill白嫖decode没用满的算力。对这段Continous Batching发展历史感兴趣可以读：大模型推理核心技术之Continuous Batching和我的WXG往事。融合派的缺点是，Prefill和Decode在相同设备上，二者并行度被绑定。如果prompt长度有限，prefill阶段占比很小基本不到10%，所以忍一忍也无所谓。不过，对于长序列当Prefill比例升高，其Prefill并行度不够灵活的缺陷就暴露出来。

分离派：考虑Prefill/Decode性质差异，人们开始尝试把Prefill和Decode放到不同的设备来分别处理，比如，Splitwisehttps://arxiv.org/abs/2311.18677和DistServehttps://arxiv.org/pdf/2401.09670，Mooncake也是对他们的继承和发展。分离派可以给Prefill和Decode设置不同的并行度，二者有更大的灵活性。许欣然提到分离派允许整个系统往 “算力/$” 和 “带宽/$” 的两个方向独立发展，对硬件优化也是更友好的。在机房里，分离派可以用不同GPU混部来降本，比如H800做Prefill，H20做Decode，二者用RDMA互联。分离派遇到的最大挑战是如何在不同设备之间传输KVCache，集群需要高速网络来互联，而网络成本不容小觑。所以，分离派的硬件往往需求很高端，目前得是是A/H卡，硬件成本高且无法Scale也为人诟病。

这一节将如何并行处理长序列的Prefill阶段。

分离派一个优势就是Prefill并行度灵活，为了降低长序列TTFT，可以分配多卡甚至多机GPU并行处理，比如长序列我们就多分配一些GPU，短序列少点。如何做并行？长序列Prefill的batch size是1，没法用数据并行。张量并行性能不满足，它通信量很大无法扩展出节点，另一方面GQA的head number也限制了它并行度。那是否可以用序列并行（SP）？Ulysses通信量远低于TP，Ring可以和计算重叠，这里也感恩Mooncake引用了我们的最近的工作USPhttps://zhuanlan.zhihu.com/p/698031151。但是SP推理需要在每个卡上replicate模型参数的，对大模型不利，如果用ZeRO方式shard参数，通信量都增加了很多；而且，SP每层都要通信，占用宝贵的网络带宽，网络带宽还得留着给KVCache送Decode Instance用呢。

Mooncake为Prefill阶段设计了Chunked Pipeline Parallelism (CPP) ，其实就是TeraPipehttps://arxiv.org/abs/2102.07988。TeraPipe正是vLLM核心作者Zhuohan Li的2021年的工作，当时是用在流水线训练里，这里只用它的forward过程即可。TeraPipe将模型沿着Layer切分成N个stage，放在不同设备上，将输入沿着序列维度切分成若干chunk。这样就可以在一个推理任务里，不同chunk流水在不同stage上的计算就可以流水起来起来。如下图所示，切分方式看左图，流水方式看右图一个Forward过程即可。

TeraPipe用于训练时不能均分序列，一个token要和之前所有token做计算（和Attention Causal Mask一体两面），位置越靠后计算量越大，均分序列的话Attention计算会负载不均衡，所以越靠后sequence chunk越小才好，所以TeraPipe论文用动态规划来找最佳划分策略。在Mooncake中，章明星说TeraPipe只有 forward 没有 backword 的话不需要啥动态规划来平衡每个 Chunk 的计算量，只要给每个 Chunk 设置一个最小的值就行。这点我是有些疑惑的，我觉得forward也需要和训练类似的负载均衡策略，如下图上半部分。

TeraPipe做推理好处 1）它仅在每个stage的边界处需要跨节点通信，这可以很容易地与计算重叠。这导致更好的MFU和与KVCache传输的网络资源争用更少。2）短和长上下文均适用。原文给了解释原因是bringing no significant overhead for short context prefill and avoiding frequent dynamic adjustment of node partitioning。我理解是：layer分布在不同设备不需要改变，长短上下文都TeraPipe，长文相当于用上了多卡资源，尽管有些气泡，相当于并行加速了，可以满足TTFT。短文本来也不用并行TTFT单GPU也可以满足，因为TeraPipe通信少，所以和一个设备做Prefill时间一样，因此和单个设备做Prefill比没有明显开销。

我觉得这里还有一些问题可以深入讨论一下，第一，是流水并行的气泡问题，可以放一些Prefill阶段不同GPU的扩展性。第二，TeraPipe可以和SP组成混合并行，更容易去扩展到多机。第三，TeraPipe方式切分参数会导致Prefill并行度没法变化，切分成8个stage就必须一直做PP=8的并行了，因此，不能弹性改变Prefill的计算资源。当然，Mooncake可能在集群里放置很多Prefill Instance，每个Instance的并行度不同，然后在Instance之间做request-level的load balance。论文目前虽然没细说，章明星后面会开文章详细介绍Prefill并行这块。

这里安利一下我们团队的DiT扩散模型的并行推理工作PipeFusionhttps://github.com/PipeFusion/PipeFusion也用了TeraPipe式的token-level切分的流水并行，特殊之处它利用了上一个diffusion step的stale KV，而且diffusion model不需要causal mask。

4.2 Layer-wise Prefill

这一节比较晦涩，我理解是在Prefill阶段对KVCache做CPU Offloading（或者传输到Decode Instance？）。这些操作都是异步的，而且是Layer-wise的，也就是一层Prefill一个Layer的KVCache Block算出来，就立刻transfer（发给decode Instance）或dump（cpu offload）来最小化GPU显存占用。

为何着急忙慌赶人家KVCache走呢？我理解是因为Prefill机器比Decode少，因为Prefill负载占LLM推理的比例低，但是Prefill产生的KVCache和Decode消耗KVCache一样多，所以Decode那边为了把硬件榨干，需要让KVCache刚好用满GPU显存，那Prefill显存肯定不够，必须Offload。这也是Mooncake论文Figure 1中之所以写成“Prefill阶段KVCache < DRAM”，“Decode阶段KVCache < VRAM”的原因。

Mooncake论文Figure 5试图证明Layer-wise有效果。这个图画的草率了，全文没有提到Serialized是什么意思。我理解他是想说Splitwise论文中Fig. 11（本文4.1节前面的图），也就是KVCache从Prefill Instance通过Layer-wise方式传输给Decode Instance，这个可以和Prefill计算重叠起来，甚至和Decode第一个step部分计算重叠起来。Splitwise用如下图方法来缓解KVCache传输的延迟，采用异步流式来和Prefill计算重叠，我觉得Mooncake也是类似。

5 KVCache-centric Scheduling

这一章开始讲调度了。

KVCache Pool用Paged方式管理KVCache（本章简称Cache）。如下图所示，黄色是已经有的Prefix Cache Blocks，其他request算了，本request可以复用。粉色是本request自己算的Incremental Cache Blocks。这Pool也会用一些Cache策略来新陈代谢，比如LRU、LFU之类的。

Prefill节点接收到一个request，它根据prefix cache ID将前prefix cache从远程CPU内存加载到本地GPU内存，以启动request的prefill计算。如果不存在prefix cache，则需要自己Prefill计算了。这种选择平衡了三个目标：尽可能多地重用KVCache（三板斧中Caching）、平衡不同预填充节点的工作负载，并保证TTFT SLO（三板斧中Scheduling）。

5.1 Prefill Global Scheduling

本节介绍如何给Prefill计算做reuse kvcache blocks和load balance。在Mooncake它将一个request切分成token blocks处理，类似FastGen和Sarathin中的Chunk。路由这些token blocks到哪台机器，要考虑因素很多。MoonShot用户很多，request之间有重叠部分，可以reuse共同前缀，prefix cache hit length越长计算越少。我猜测这个就是Kimi Context Cachinghttps://mp.weixin.qq.com/s/anZhObqWPLWZrNQCFv64Ag能力的来源。但是，都路由到prefix cache hit length最长的机器，机器之间会负载不均衡，为了load balance，还需要考虑distribution of reusable KVCache blocks。本节也callback了KVCache为中心的宗旨，是让request主动去找KVCache。

先说怎么找到max prefix cache来尽可能reuse kvcache blocks。

一个token block经过prefill计算产生一个KVCache block，而且一个token block prefill计算需要所有Prefix token blocks的Prefix KVCache blocks。这些KVCache blocks都是在Distributed KVCache Pool中，不一定在本地内存，怎么快速检索这些KVCache blocks呢？需要建立一个token block -> KVCache block的映射关系，根据映射关系去检索prefix token blocks的KVCache blocks。这个映射关系是一个Hash Table。

为了快速找到一个token block最大前缀max prefix token blocks，Hash设计有讲究。每个Block的Hash Key是基于前一个Block的Hash Key计算得到，图中B=Hash(A+b)。如果两个Block Hash Key相同，不仅该token block的KVCache block，之前所有prefix KVCache也都可以复用。检索代价从多项式降低到线性。这个技巧很巧妙，来自vllm。

再说怎么做load balance。

借助上面算出来的max prefix cache length信息，注意每个机器都不一样。可以用request长度+此机器的prefix长度估计每个机器执行该request时间，再加上等待时间，作为最终TTFT。将request分配给具有最短TTFT的机器，并相应更新该机器的缓存和队列时间。如果无法实现SLO，直接返回HTTP 429 Too Many Requests。

5.2 Cache Load Balancing

本节介绍如何KVCache负载根据使用频率做再平衡。

在Mooncake集群中，每个Prefill机器管理其自己本地的prefix caches。这些caches的使用频率差异很大。例如，系统提示几乎被每个请求访问，而存储来自本地长文档内容的cache可能只被一个用户使用。我们希望不同机器cache使用频率相近，因此要调整cache的位置。

因为很难预测一个cache未来使用频率。Mooncake提出了一种基于启发式的自动热点迁移方案。上一节所属，request可能并不总是被定向到具有最长prefix cache length的Prefill机器上。在这种情况下，如果估计的prefill时间短于cache传输时间，因为要排队等待，cache位置和request转发到另一个机器。该机器主动检索KVCache并将其存储在本地。另外，如果远程机器的最佳prefix cache length不长，还可以重计算。

这两种策略合在一起，不仅减少了prefill时间，还促进了热点缓存的自动复制，使其能够更广泛地分布在多台机器上。