2月26日，Mooncake的论文获得「计算机存储顶会 FAST 2025」Best Paper，这也是国内连续第三年拿到FAST Best Paper。同时，Mooncake 团队宣布和 vLLM 团队已经合作制定了一个多阶段路线图。这次整合将为 vLLM 引入 P/D（Prefill/Decode）分解和全局 KVCache 设计。

Mooncake：

GitHub - kvcache-ai/Mooncake: Mooncake is the serving platform for Kimi, a leading LLM service provigithub.com/kvcache-ai/Mooncake

大模型训练/推理高性能优化

原始论文发表于 2024年，论文地址：Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving

分离式部署的动机是什么？

大模型推理分成 prefill 和 decode 两个阶段，prefill 阶段是计算密集型，decode 阶段是访存密集型且受限于显存，两者的资源利用率不一致。decode 阶段在显存满足的情况下应给与更多的 tokens，来提高资源利用率，分离式部署就可以很好的满足这点需求。

通过分离式部署，prefill 和 decode 可以采取对各自有利的分布式配置、可以采取对各自有利的 GPU 架构，来利用有限的 GPU 资源最大化吞吐。

但是分离式部署，需要将 prefill 的 kv cache 传输到 decode，引入了通信开销；常用的解决办法有如下几点：

prefill 阶段按 chunk 或者 layer 来传递，计算完一个就传递一个
kv cache 量化，采样 int 4 或者 int8，在精度达标的情况下，减少了通信量

分离式部署一般需要做什么？

针对 prefill instance，尽可能提高 kv cache 命中率，来减少重复计算
prefill 按 chunk/layer 计算，【计算】和【传递 kv cache 给 decode 】流式并行
decode 加载 kv cache，并且将请求添加到 continuous batching 进行处理

chunked pipeline parallelism

针对每个请求，都会被划分为 chunks，同一请求的不同 chunk 在不同的节点上面进行计算，这样多个请求就可以并行起来，从而减少 TTFT。

prefill 阶段的优化目标

prefill 阶段需要尽可能提高 kv cache 命中率，减少重复计算。这就会涉及到针对常用的 kv blocks 需要被拷贝到 prefill 节点，不常用的 kv blocks 需要释放掉。在常见的分离式部署框架中，一般由 KV Router 模块来管理。

decode 阶段的优化目标

decode 阶段需要集合尽可能多的 tokens 到一个 batch 来进行推理，从而提高 MFU（Model Flops Utilization），这就会涉及到显存的问题，可能显存会不够用，同时还不能过多的影响 TBT（Time Between Tokens）。TBT 是用户能比较直接感受到的一个指标，TBT 过高，用户体验会较差。

Mooncake 是为了解决什么？

Mooncake 主要是遇到了 GPU 资源紧张，在实际应用场景，有限的资源情况下，如果处理过载调度（overload-oriented scheduling）的问题。

Mooncake 提出了哪些方案？

Mooncake 提出了一种 early rejection 机制，通过预测将来的负载来提前拒绝某些请求，从而减少计算资源浪费。因为实际高峰场景中会遇到某个请求完成 prefill 计算后，decoding 阶段并没有资源槽来做，还不如提前结束反馈给客户端，避免资源的进一步拥堵。

(1). KVCache Pool

Mooncake 的 KVCache 缓存池是在 CPU 上的。上面的流程图解释的比较清晰，先看命中了哪些 prefix cache，然后 prefill instance 进行加载，再计算 incremental cache blocks，然后将 prefix cache blocks + incremental cache blocks 传输到 decode 节点上。

(2). Inference Workflow

上图是一个完整推理的流程图，可以发现 mooncake 使用了 prefill computation 和 kv cache load and store 并行来减少传输 overhead。由于 kv cache 主要以 CPU 作为中介存储点，decode 阶段也是计算和加载异步执行，边加载边计算。

kv cache 放在 CPU 上因为线上并发比较高，GPU 显存有限，而 CPU 存储空间大。实际场景 CPU 内存也是不够用的，还需要用到硬盘，涉及到 kv cache offload 相关的进一步调度。

(3). KVCache-centric Scheduling

Mooncake 关于 KVCache 调度逻辑见上图。

针对 prefill instance 的选择，会考虑 prefix cache 命中长度和 KVCache blocks 重复使用的分布情况。新进来一个请求时，input tokens 会被切分成几个 blocks，针对每个 block 计算一个 hash key，然后将这些 block keys 与每个 prefill instance cache keys 去匹配，来识别 prefix match length（prefix_len）。

调度器会根据每个请求的长度和 prefix_len 来估计 prefill 需要执行的时间（使用离线测试集制作的预测模型），然后加上预计需要的等待时间（队列里面所有请求的 prefill 时长之和），就是这个请求的预估 TTFT。

由于大模型结构的一致性，prefill 时间还是比较好预估的，难预估的是 transfer 时间，因为 transfer 是异步的。

总结

Mooncake 这篇论文核心是为了解决 GPU 资源有限且请求文本较长的场景，提出的 early-rejection 方案可以参考一下。具体实际应用场景下的大模型推理分布式部署，还需要结合业务情况来具体优化，但是常见的优化手段，如 kv cache 量化、P 到 D 按 layer 传输、prefix cache 命中率提升这些都是必须的。另外包括异构场景下，P 和 D 不同的分布式配置，会引入什么问题，又会有哪些优化，论文并没有涉及。

发布于 2025-05-12 21:46