1. Mooncake 简介
Kimi与清华大学联合参与的项目论文《Mooncake: Trading More Storage for Less Computation-A KVCache-centric Architecture for Serving LLM Chatbot》获得计算机存储领域顶会 FAST 2025「最佳论文」奖。
Mooncake 是 Kimi 的底层推理服务平台,该平台采用一种以键值缓存(KVCache)为中心的分离架构,不仅将预填充和解码集群分开,还高效利用了 GPU 集群中未充分利用的 CPU、DRAM、SSD 和 NIC 资源,建立了一个分离的 KVCache。Mooncake 的核心是其以 KVCache 为中心的全局缓存和调度器,旨在在严格的延迟相关服务级别目标(SLOs)下最大化吞吐量。
实验表明,Mooncake 在处理长上下文输入的场景中表现出色。在使用真实数据进行的测试中,与基线方法相比,Mooncake在符合SLOs的情况下,将有效请求处理能力提高了59%~498%。目前,Mooncake已在数千个节点上运行,每天处理超过1000亿个token。在实际部署中,Mooncake的创新架构使Kimi 在NVIDIA A800和H800集群上分别比以前的系统多处理115%和107%的请求。
2. Background
- LLM推理分为两个阶段:Prefill与Decoding
- LLM都是基于Transformer的自回归模型。它根据之前的tokens不断地生成下一个token,此类模型的推理可以分为prefill和decode两个阶段:
- 预填充阶段(Prefilling):在 LLM 推理过程中,预填充阶段是指模型在收到输入序列后,生成第一个输出token之前的计算过程
- 解码阶段(Decoding):解码阶段是指在预填充阶段生成第一个token后,模型通过迭代生成后续token的过程,直到生成结束标志或达到最大序列长度
- LLM都是基于Transformer的自回归模型。它根据之前的tokens不断地生成下一个token,此类模型的推理可以分为prefill和decode两个阶段:
2.1 预填充(Prefill)
预填充阶段是模型开始生成响应之前的初始阶段,主要用于处理输入Prompt,主要包括:
- 上下文设置:为模型提供必要的背景信息或初始文本,以便为响应设定场景
- 数据预处理:格式化和结构化输入数据,进行分词
- 参数初始化:设置文本生成所需的模型参数和配置
计算流程:
假设输入Prompt为:"The quick brown fox jumps over the lazy dog."
假设模型的词汇表将其拆分为8个token(简化表示):["The", "quick", "brown", "fox", "jumps", "over", "the", "lazy"]
计算过程:
- Token Embedding:输入文本通过分词器转换为Token ID,并查询嵌入矩阵获得词向量
- Transformer计算:所有Token的嵌入并行通过模型的多层Transformer结构进行前向传播,计算出最终的隐藏状态(hidden states)
- Logits计算:最后一层的输出经过线性变换(Linear Layer),得到用于预测下一个Token的Logits分布
由于所有Token可以同时计算,预填充阶段的计算是高度并行化的,因此在批处理(batching)中效率较高。
优化方式:
- 缓存管理:使用KV Cache存储计算的Key/Value,以减少重复计算
- FlashAttention优化注意力计算:减少访存开销,加速Transformer计算
- Fused Kernels(融合算子):将多步计算合并,提高计算效率
- 分布式推理优化:
- 流水线并行(Pipeline Parallelism):拆分Transformer层,提高多卡利用率
- 张量并行(Tensor Parallelism):拆分矩阵运算,减少单卡计算负担
2.2 解码(Decode)
解码阶段解码是模型根据给定的预填上下文生成文本的阶段,用于逐步生成新Token,它包括:
- 标记生成:模型根据给定的上下文预测序列中的下一个token
- 采样策略:决定在生成过程中如何选择token的技术,如贪婪搜索、束搜索或温度采样
- 后处理:将生成的token转换回可读文本,进一步改进输出以确保连贯性和相关性
与预填充不同,解码是一个自回归(autoregressive)过程,每次只能计算一个新Token,因此计算是顺序的,无法完全并行。
计算流程:
假设模型已经计算完前面的输入Prompt "The quick brown fox jumps over the lazy",现在需要生成下一个Token:
- 获取最新Token的隐藏状态:取最后一个Token "lazy",查找嵌入,计算其隐藏状态
- 计算注意力机制:只需计算当前Token对之前所有Token的注意力(使用KV Cache减少重复计算)
- 生成Logits并采样:通过Softmax获取下一个Token的概率分布,根据不同采样策略(Greedy, Top-K, Top-P)选择下一个Token,如 "dog"
- 重复上述步骤,直到达到终止条件(最大长度、终止Token等)
优化方式:
- KV Cache复用:解码阶段只需计算新Token的Key/Value,并追加到缓存,而无需重新计算之前Token的注意力
- 并行化策略:
- Speculative Decoding(推测解码):提前预测多个Token并并行计算,减少延迟
- Multi-Step Decoding(多步解码):部分模型可尝试并行预测多个Token后再修正
- 分布式优化:
- 流水线并行:多个GPU流水线处理不同批次的解码请求
- Tensor并行:将每个Token的计算分布到多个设备上加速运算
2.3 KV Cache
KV缓存主要用于存储 Transformer 解码器中的 Key(键)和 Value(值),减少重复计算,加速自回归推理。在解码时,当前时间步的 Query(Q)只需要与之前存储的 K-V 进行计算,避免了重复计算,提高推理效率。
常见的 KV 缓存存储格式:
kv_cache = {
"key": torch.zeros(batch_size, num_heads, max_seq_len, head_dim),
"value": torch.zeros(batch_size, num_heads, max_seq_len, head_dim)
}
每次解码时,会将新计算出的 key 和 value 追加到 kv_cache。
2.4 推理性能指标
-
两个重要指标:Time To First Token(TTFT),Time-Between Tokens(TBT)(论文中使用这两个指标)
上图可知,Prefilling阶段对输入批次进行处理,同时生成Decode阶段的第一个token,在Decode阶段,每次利用历史的KV Cache以自回归的形式,每一轮都生成一个token,在这两个阶段中,有两个重要指标。第一个是在Prefill阶段生成第一个token的时间,TTFT,第二个是Decode阶段两个token之间生成的间隔时间,TBT
其他指标:
- TTFT (首Tokens时延,Time to First Token):用户发起请求到模型返回第一个 Token 所需的时间
- E2E Latency (端到端时延,End-to-End Latency):从请求发送到完整接收所有输出 Token 的总耗时
- TPOT(Time Per Output Token):模型生成单个 Token 的平均耗时
- ITL(Token 延时,Inter-Token Latency):生成每个 Token 的实际间隔时间
- 吞吐(Throughput):单位时间内系统处理的 Token 总数(如 Tokens/秒)
2.5 推理请求SLO
SLO(Service Level Objective,服务级别目标)是论文中提到的一个概念,【Mooncake的核心是其以KVCache为中心的调度器,旨在平衡最大化整体有效吞吐量同时满足与延迟相关的服务级别目标(SLO)】,主要是TTFT和TBT
在大语言模型服务器中,多个用户可能会同时发送请求,服务器需要根据不同用户的请求类型动态分配 GPU 资源。由于请求的特性各不相同,资源需求也有所差异。
例如,一些用户发送的是快问快答请求,期望服务器能迅速响应,此类请求对 TTFT(Time to First Token) 和 TBT(Total Block Time) 要求较高,需要更短的首字节延迟和整体处理时间。而对于批处理请求,如输入一篇文档并进行全文翻译,由于处理内容较长,整体耗时较多,但对 TTFT 和 TBT 的要求相对宽松。
当请求量激增时,服务器的计算资源可能无法满足所有请求。在调度过程中,如果系统判断某些请求无法在合理时间内完成,或者资源已接近上限,可能会优先满足高优先级任务,并拒绝部分预计无法按 SLO(服务级别目标)满足的请求。
现有方案的目标:满足SLO的前提下,最大化整体吞吐
- 以空间换效率:KV Cache
现有方案仅使用GPU缓存,DRAM容量有限,理论缓存命中率为50%(有的高有的低,概率大概50%):
Mooncake 就是希望尽可能扩大存储KV Cache的资源,容纳更多的KV Cache,减少对 GPU 计算资源的占用。(KV Cache 通常存储在VRAM中,而 GPU 显存资源有限,存储过多的 KV Cache 会影响其他计算任务的执行)
Mooncake 通过扩展 KV Cache 的存储能力,将其部分或全部存放在其他存储介质(如 CPU 内存)上,以减少对 GPU 显存的依赖。这种优化策略不仅能提升推理任务的并发能力,还能在相同 GPU 资源下支持更多的用户请求,提高整体系统吞吐量。
3. Mooncake 论文
3.1 MoonCake Architecture
Mooncake 是一个围绕 KVCache进行调度和优化的 LLM推理架构,旨在提高推理效率,同时最大程度地复用 KVCache,减少计算和内存开销。
Mooncake 采用 KVCache 分离架构,并通过 全局调度(Conductor)+ 局部调度(实例调度器) 结合的方式优化推理性能:
1. KVCache-centric Conductor(全局调度器):
- Cache-aware Prefill Scheduler:在预填充阶段,调度器需要选择一个合适的预填充实例,并尽可能复用已有的 KVCache
- KVCache Balance Scheduler:管理 KVCache 的平衡,确保高频访问的数据能快速提供,避免热点数据访问瓶颈
- Load-balance Decoding Scheduler:在解码阶段,调度器需要选择合适的解码实例,并优化批处理(batching)策略,提高吞吐量
2. Mooncake 将计算和存储资源划分为不同的池,每个池负责特定任务:
- Prefill Pool(预填充池):负责预填充阶段的计算,主要在 GPU 进行,并利用 Paged KVCache(分页 KVCache)
- KVCache Pool(KV 缓存池):由 CPU/DRAM/SSD 组成的分布式 KVCache 资源池,支持跨节点的数据存取
- Decoding Pool(解码池):负责解码阶段的计算,主要在 GPU 进行,并利用 Paged KVCache 进行缓存管理
3. 调度流程
Mooncake 采用 全局调度 + 局部调度 的方式进行推理计算:
- 预填充实例选择(Prefill Scheduling)
- 最大化 KVCache 复用,减少重复计算,提高预填充阶段的效率
- 由于 KVCache 分布在不同层级存储(DRAM/SSD),如果从较慢的存储(SSD)中读取数据,会导致 TTFT延迟,影响性能
- 解决方案:
- 提前预测 KVCache 访问需求,预先加载热点数据,减少传输等待时间
- 热点数据多副本存储,避免多个任务同时请求同一块数据时发生网络拥塞
- 冷数据淘汰,减少无用缓存占据高性能存储资源
- 预填充阶段计算(Chunked Prefill)
- 计算任务被拆分成 多个 Chunk(数据块) 或 Layer(层) 进行处理,提高并行度
- 计算过程中,KVCache 需要不断传输到解码实例,以便后续解码阶段可以无缝衔接
- 解码实例选择(Decoding Scheduling)
- 解码实例需要 加载 KVCache 并加入批处理队列(Continuous Batching)
- 目标是 最大化吞吐量(Throughput),保证 TBT 生成速度不会受限
4. 优化目标
- Prefill Stage
- 最大化缓存复用,减少重复计算,提升推理效率
- 约束:
- TTFT SLO:必须保证首个 token 生成的延迟在可接受范围内
- 最小化 MFU:尽量减少不必要的内存占用
- KVCache < DRAM:缓存大小不能超过 DRAM 限制,否则会溢出到 SSD,影响访问速度
- Decoding Stage
- 最大化吞吐量,提高 Token 生成速率
- 约束:
- TBT SLO:必须保证生成 token 速度满足要求,不影响用户体验
- KVCache < VRAM:解码时,缓存大小不能超过 GPU VRAM 限制,否则会溢出到更慢的存储,降低性能
3.2 Mooncake’s Disaggregated Archtecture
分离式架构:
- Prefill与Decode计算资源分开
- KVCache和计算分离开
上图展示了 KVCache 在 CPU 内存中的存储和传输逻辑:
- KVCache 结构
- KVCache 以分页块的形式存储在 CPU 内存中,每个块的哈希值由其自身哈希和前缀哈希共同确定
- 黄色块(Prefix Cache Blocks)表示已缓存的前缀
- 粉色块(Incremental Cache Blocks)表示增量缓存部分
- 灰色块(Unallocated Cache Blocks)表示未分配的缓存区域
- 缓存匹配 & 增量计算
- KVCache 通过哈希匹配实现高效的前缀重用,图中 a 到 e 成功匹配 ✅,f 发生不匹配 ❌,需要计算新的 KVCache
- f 之后的 KV 需要增量计算,并存储到缓存池中
- KVCache 传输与调度
- Messenger 组件 负责 KVCache 在 CPU 和 GPU 之间的传输,使用 GPUDirect RDMA 加速数据移动
- Prefill Instance 负责预填充 KVCache
- Decoding Instance 负责加载 KVCache 用于推理
- Conductor 调度器 负责全局管理 KVCache,调度 KVCache 复用、复制或置换,提高推理性能
Prefill 与 Decode 两个阶段解耦
在同一块 GPU 上同时执行 Prefill 和 Decode,优先级策略的选择会直接影响推理性能,需要在 TBT 和 TTFT之间进行权衡。如果 Prefill 阶段优先,即在新请求进入时优先完成预处理,则新请求的 TTFT 会降低,使用户能更快看到模型的首个输出,但是会占用较多 GPU 资源,从而影响 Decode 阶段 的执行,导致已在运行的生成任务的 TBT 增大。
反之,如果 Decode 阶段优先,即让已处理的请求更快完成 Token 生成,则可以缩短 TBT,提高推理吞吐量,但这可能会延迟新请求的 Prefill 处理,导致 TTFT 增加,使用户感知到的首 Token 响应时间变长。
上图描述了一个请求的完整推理流程,包括 KVCache 复用、增量预填充、KVCache 传输以及解码:
- KVCache Reuse
- Conductor 选择 一个合适的 prefill 节点(或一个 prefill 组)处理请求
- 该 prefill 节点接收请求,其中包含:原始输入(raw input)、可复用的前缀缓存(prefix cache block IDs)、完整请求所需的缓存块(full cache block IDs)
- 该 prefill 节点会根据 前缀缓存的 block ID,将 CPU 内存中的 KVCache 加载到 GPU,作为计算的初始状态
- 如果没有可用的 KVCache 前缀缓存,则跳过此步骤
- Incremental Prefill
- prefill 节点使用 复用的 KVCache 计算 未缓存的部分(增量计算)
- 计算完成后,将新的 KVCache 存回 CPU 内存
- 若未缓存的 token 数量超过某个阈值(prefill_chunk,通常大于 1000 tokens),则执行流水线处理:将 prefill 任务分块(chunked execution),分批处理 KVCache 计算(采用 分块(chunked)计算,预填充计算能与 KVCache 传输并行执行,提高吞吐)
- KVCache Transfer
- Messenger 组件 管理和传输 KVCache
- 每个 Messenger 运行在 独立的推理进程 中,异步管理跨机器 KVCache 传输(使用 GPUDirect RDMA)
- 传输过程与 prefill 计算重叠执行,计算增量 KVCache 同时 传输已计算好的 KVCache,避免等待
- Decoding
- 当所有 KVCache 传输至 解码节点的 CPU 内存(DRAM) 后,请求加入 连续批处理(continuous batching) 队列
- Conductor 预先选择解码节点,基于当前负载保证 不违反 TBT(Time-Between-Tokens)SLO
- 二次负载检查:本地调度器(local scheduler)会重新检查当前负载,如果预期负载发生变化,可能会拒绝请求,如果请求被拒绝,则 前面的 prefill 计算成本会被浪费
从图上可以看到:
- Prefill Instance:在 Prefill 阶段,模型需要处理新的输入,并生成对应的 KV Cache
- (s1) KVCache Reuse:如果输入包含已处理的前缀(Prefix),则可以复用之前的 KV Cache,减少计算开销
- (s2) Incremental Prefill:对于新输入的部分,增量计算新的 KV Cache(Incremental KVCache)
- KV Cache 存储策略:
- GPU 端:存储 Prefix KVCache 和 Incremental KVCache
- CPU 端:同样存储 Prefix KVCache 和 Incremental KVCache
- Layer-wise Load and Store(∗):KV Cache 采用 按层(Layer-wise)加载和存储,并与 Prefill 计算 并行执行,以减少数据传输的开销(降低 Prefill 过程中的等待时间)【按层加载和存储 KV Cache,减少数据传输开销,与计算并行执行】
- Decoding Instance:模型基于 KV Cache 生成 Token
- Full KVCache 加载:
- GPU 端和 CPU 端都需要存储完整的 KV Cache(Full KVCache)
- (s3) KVCache Transfer:Prefill 阶段完成后,KV Cache 需要传输到 Decoding Instance
- (s4) Decoding:GPU 执行解码计算
- Async Load(†):解码阶段采用 异步加载(Async Load),即在 GPU 解码计算的同时,将 KV Cache 逐步加载到显存,防止 GPU 由于等待数据而进入空闲状态(提高 GPU 利用率)【异步加载 KV Cache,防止 GPU 空闲,提高推理吞吐量】
- Full KVCache 加载:
3.3 块流水线并行CPP
Attention计算过程:第一个过程是根据Query和Key计算权重系数,第二个过程根据权重系数对Value进行加权求和。而第一个过程又可以细分为两个阶段:第一个阶段根据Query和Key计算两者的相似性或者相关性;第二个阶段对第一阶段的原始分值进行归一化处理;这样,可以将Attention的计算过程抽象为如图所示的三个阶段:
序列并行:将一个序列分为多份,分别发往不同的GPU,每个GPU分别计算Q、K、V(上图所示三个阶段),最后把所有计算结果传输到一个GPU进行计算。如下图,假设将一个序列分为a、b、c三个子序列,b和c完成计算后,传输到GPU1(线上方表示计算、下方表示传输)
作者认为这种方式仅适合长请求,例如GPU1可以容纳a、b、c组成的这个序列,就没有必要采用序列并行分配到不同的GPU,反而增加传输通信开销。
块流水线并行:将a这个子序列,分为1、2、3三个阶段(如上图的三个阶段)的操作,在GPU1中进行阶段1的操作,完成后传输到GPU2,在GPU2进行计算第二个阶段,然后传输到GPU3进行第三个阶段的计算,以此类推。长请求和短请求都适用,粒度更小,传输不会阻塞。
核心思想:
- 节点分组:在集群中,每 X 个节点被分为一个流水线预填充节点组(并行化预填充计算)
- 输入拆分:每个请求的输入被拆成多个块(chunk),每个块大小不超过 prefill_chunk(同一请求的不同块可以由不同节点同时处理,并行度更高)
CPP的好处:
- 减少跨节点通信,提升计算利用率(MFU)
- 只在每个流水线阶段的边界进行通信,而不是频繁同步,通信可以与计算重叠,减少了通信开销
- 减少与KVCache 传输竞争的网络资源争用,优化推理时的带宽使用
- 适应不同上下文长度
- 短上下文的预填充不会有额外的开销
- 长上下文的预填充能够避免频繁调整节点分区,减少调度开销
3.4 逐层预填充
优化KVCache的VRAM占用,提高计算效率,并利用计算和KVCache传输的重叠来减少推理延迟
核心思路:
1. VRAM 是宝贵资源,目标是最小化 KVCache 占用
- KVCache 的占用成本 = S(KVCache 大小)× T(存活时间)
- 如果请求被分块处理,且与解码请求交错执行,则 T 会增加,导致更大的 KVCache 占用成本
2. 逐层预填充:计算 + KVCache 传输重叠
- 预填充是逐层进行的,可以异步加载和存储 KVCache,让计算与数据传输并行执行,进一步降低VRAM占用
- 在 Mooncake 系统 中:
- KVCache 的 加载(load)和存储(store)是异步执行的
- 计算某层注意力前,等待该层的 KVCache 加载完成,并触发下一层的 KVCache 加载
- 计算完成后,启动该层 KVCache 的存储
- 所有层计算完后,再等待所有 KVCache 存储操作完成
好处:
- 由于 KVCache 传输和计算并行,执行时间主要受 KVCache 加载时间或标准预填充时间决定,实验表明,这种方法能有效降低长上下文推理的延迟:
- 预填充调度只需要关注 KVCache 分布和 DRAM 可用空间,而不是 VRAM 大小
3.5 KV Cache 中心调度
核心目标是最大化缓存命中率,优化负载均衡,减少请求的预填充时间,从而降低推理延迟(TTFT)
预填充全局调度
传统的 LLM 负载均衡策略一般只考虑请求数量,而 Mooncake 的调度策略额外考虑了:
- 前缀缓存命中长度(pref ix_len):如果请求的前缀已经在 KVCache 里,直接复用缓存可以减少计算开销
- KVCache 复用分布:不同实例的 KVCache 可能存储了不同的前缀内容,合理分配可提高缓存命中率,减少重复计算
缓存感知的全局调度算法:
- 请求被拆分为多个块,每个块计算一个哈希键,连接前后块的哈希键,形成唯一标识
- 查找预填充实例的缓存,计算前缀匹配长度(pref ix_len)
- 估算执行时间(基于请求长度和 pref ix_len)
- vLLM 也有类似的 KVCache 复用逻辑,但仅限于本地缓存,而 Mooncake 是全局调度
- 预测模型基于离线测试数据,误差较小
- 选择 TTFT 最短的实例进行分配,并更新该实例的缓存和队列时间
- 如果无法满足 SLO(服务质量目标),直接返回 HTTP 429 Too Many Requests
缓存负载均衡
问题:缓存访问频率不均
- system prompt 几乎每个请求都会访问(高频缓存)
- 用户长文本缓存 可能只被单个用户访问(低频缓存)
目标:
- 高缓存命中率:让高频 KVCache 尽量被复用
- 低负载:让请求均匀分布,避免部分节点过载
启发式热点迁移方案:
1. 避免将请求强制分配到最长前缀匹配的实例
2. Conductor 发现更优的 KVCache 位置后,主动迁移缓存
- 让新的实例从原持有者获取 KVCache,并存储到本地
- 只在 额外 prefill 时间 < KVCache 传输时间 时执行,避免无效复制
3. 动态计算 vs 直接缓存匹配
- 如果远程匹配长度不大于本地匹配的一个阈值,优先直接计算输入 token 而不是传输 KVCache
- 这样可以减少 prefill 时间,同时促进 KVCache 自动复制到更多机器上,形成更优的缓存分布
小结
1. Mooncake 的 KVCache 调度不只是负载均衡,而是缓存感知的全局调度
- 计算请求匹配的 KVCache 位置,选择 TTFT 最短的实例
- 结合缓存分布、实例负载、网络传输进行最优调度
- 采用 HTTP 429 机制防止超载
2. 缓存负载均衡使用了动态的热点迁移策略
- 主动迁移:将 KVCache 复制到更优位置,提高缓存复用率
- 如果 KVCache 传输成本高,直接计算会更优
3. 提高缓存命中率,减少重复计算,让系统适应动态工作负载,避免节点过载,降低 TTFT,提升推理吞吐量
4. 总的来说,Mooncake 的 Conductor 通过缓存感知调度 + 动态 KVCache 迁移,实现了高效的长上下文推理优化
3.6 拓扑感知路径选择
- 增加热门KV Cache副本数量来缓解拥塞,但传输可能会受到互联的PCle交换机带宽限制
- 每个节点都有一个拓扑矩阵,矩阵里面包含首选、备选网卡,实现更高效的数据传输
3.7 高负载场景下的调度策略
1. 过载场景下的调度:
- 传统LLM服务通常假设所有请求都可以处理,但现实中资源增长往往赶不上请求增长,导致过载问题
- Mooncake 采用SLO满足度(TTFT和TBT的SLO约束)来定义负载,而非简单的请求数量比率
2. 提前拒绝策略:
- 由于解码阶段的计算资源有限,预填充完成的请求可能在解码阶段被拒绝,导致计算资源浪费
- 解决方案是在预填充阶段开始之前评估解码实例的负载,如果超过阈值则提前拒绝请求
3. 负载波动问题:
- 早期拒绝虽然减少了无效计算,但会导致预填充和解码实例之间的负载波动,降低资源利用率
- 主要原因是解码负载的预测滞后,导致预填充和解码之间的负载相位错位
4. 基于预测的早期拒绝:
- 通过预测解码负载来调整提前拒绝策略,减少波动
- 请求级别预测:估计单个请求的输出长度,精度要求高,成本较高
- 系统级别预测:估算整个系统的平均解码负载,更适用于超载场景
3.8 性能
4. Mooncake 项目
Mooncake 通过 Transfer Engine 进行高效的数据传输,并在此基础上实现了 P2P Store(用于分布式 checkpoint 共享)和 Mooncake Store(用于 vLLM 的 KVCache 池化)。此外,通过对 vLLM 的修改,使其能够更高效地执行 预填充-解码解耦,从而优化大规模推理任务的吞吐量和效率。
4.1 Transfer Engine
Mooncake 的底层是 Transfer Engine,Transfer Engine 是一个高性能数据传输框架,提供统一接口来管理 DRAM、VRAM 和 NVMe 之间的数据传输,同时屏蔽底层硬件细节。其核心功能是提供高速、可靠且灵活的数据传输,支持多种传输协议,包括:
- TCP 传输
- RDMA 传输(包括 RoCE、IB、eRDMA)
- NVIDIA GPUDirect RDMA
- NVMe-over-Fabric (NVMe-of)
- CXL 传输(尚在开发中)
相比于 PyTorch 分布式通信库 gloo 和传统的 TCP 传输,Mooncake Transfer Engine 具备更低的 I/O 延迟,提升了数据交换效率。
性能:
- 在 4×200 Gbps RoCE 网络上,Transfer Engine 提供 87 GB/s 带宽
- 在 8×400 Gbps RoCE 网络上,Transfer Engine 提供 190 GB/s 带宽
- 约 2.4x~4.6x 的加速效果,相比于 TCP
4.2 P2P Store
基于 Transfer Engine,Mooncake 提供了 P2P Store 库,用于 在集群节点之间共享临时对象(例如 checkpoint 文件)。P2P Store 通过分布式存储方式,避免了单个机器的带宽瓶颈,提高了大规模集群的训练和推理效率。
性能:满载利用硬件带宽,例如:在 25Gbps NIC 上,get replica 吞吐量达 3.1 GB/s
4.3 vLLM 集成
vLLM 社区正在开发 分布式 Prefill/Decode 解耦,现有实现依赖 nccl 和 gloo 进行数据传输,难以高效实现跨机器解耦。
Mooncake 还对 vLLM 进行了改动,使 Transfer Engine 与其深度集成,主要用于 预填充-解码(prefill-decode)解耦。借助 RDMA 设备,这种解耦方式更加高效,减少了推理过程中的通信延迟。
Mooncake 方案:
- Transfer Engine 替代 nccl 和 gloo:通过 RDMA 设备加速 KVCache 传输,提供更简单的接口和更高效的 RDMA 资源利用率
- xPyD Prefill/Decode 解耦:通过 Mooncake Store 进行优化,提高推理吞吐量
4.4 Mooncake Store
基于 Transfer Engine,Mooncake 还提供了 Mooncake Store,它支持 分布式的 KVCache 池化,用于 vLLM 的 xPyD(扩展并行解码)解耦,进一步优化推理过程中缓存的存取效率。
Mooncake Store 是 专为 LLM 推理优化的分布式 KVCache 存储引擎,提供 Put、Get 和 Remove 等对象级 API,支持 vLLM 的 KVCache 池化和xPyD(扩展并行解码)解耦。
4.5 小结
组件 | 作用 | 关键优化点 |
Transfer Engine | 高效数据传输框架 | RDMA/NVMe-of/TCP,NUMA 感知路径选择,错误恢复 |
P2P Store | 分布式 checkpoint 共享 | 纯客户端架构,避免单点带宽瓶颈 |
Mooncake Store | LLM KVCache 存储 | 适用于 vLLM xPyD,支持高效对象存储 |
vLLM 集成 | 预填充-解码解耦 | 替换 nccl/gloo,实现跨机器 KVCache 共享 |
Performance
Backend/Setting | Output Token Throughput (tok/s) | Total Token Throughput (tok/s) | Mean TTFT (ms) | Median TTFT (ms) | P99 TTFT (ms) |
Transfer Engine (RDMA) | 12.06 | 2042.74 | 1056.76 | 635.00 | 4006.59 |
TCP | 12.05 | 2041.13 | 1414.05 | 766.23 | 6035.36 |
5. Quick Start
5.1 环境准备
- RDMA 驱动:需要安装 Mellanox OFED(适用于 IB/RoCE 网络)
- 系统:Linux(x86_64)
- 编译工具:
-
gcc
、g++
9.4+ -
cmake
3.16+ -
make
-
- 推荐依赖
- CUDA 12.1+
- Go 1.20+
- Rust Toolchain
- hiredis
- curl
5.2 安装步骤
# 获取源码
git clone https://github.com/kvcache-ai/Mooncake.git
cd Mooncake
# 安装依赖
bash dependencies.sh
# 编译
mkdir build && cd build
cmake .. # 可加 -D 选项启用特定功能,如 -DUSE_CUDA
make -j$(nproc) # 使用全部 CPU 核心编译
style="display: none !important;">