Mooncake + PD 分离架构深度解析:从原理到数据流全链路拆解

一、背景:大模型推理的两个阶段

大模型生成一段回答,本质上分两步:

  • Prefill(预填充):读懂用户的输入,生成"理解笔记"——即 KV Cache。这一步计算量巨大,但可以并行处理,GPU 算力拉满。
  • Decode(解码):拿着 KV Cache,一个字一个字往外生成。每生成一个字都要回头查"笔记"。这一步无法并行,速度慢,但对算力需求不高,主要吃内存。

如果一台机器同时负责两个阶段:

审题(忙) → 写答案(慢) → 审题(忙) → 写答案(慢)
GPU满载     GPU空闲       GPU满载     GPU空闲

两个阶段对硬件的需求完全不同,放在一台机器上互相拖累。这就是 PD 分离(Prefill-Decode Disaggregation) 的动机:让 Prefill 节点和 Decode 节点各司其职。


二、Mooncake 是什么

Mooncake 是 Kimi(月之暗面)开源的分布式 KV Cache 基础设施,提供两个核心能力:

能力 说明
Transfer Engine 高效的 RDMA 点对点数据传输通道
Distributed KV Cache Store 集群级统一内存池,跨节点缓存 KV Cache 供复用

两个能力用在不同环节,配合使用。


三、整体部署架构

                  ┌─────────────────────┐
                  │  mooncake_connector  │
                  │  _proxy.py          │
                  │  (请求代理层)        │
                  │  :8000              │
                  └──────────┬──────────┘
                             │
            ┌────────────────┼────────────────┐
            ▼                                 ▼
   Prefill 节点群                       Decode 节点群
   :8100 (vLLM API)                    :8200 (vLLM API)
   :8998 (bootstrap)
            │                                 │
            └──── RDMA (Mooncake) ────────────┘
                  KV Cache 点对点传输

   ┌──────────────────────────────────┐
   │         etcd 集群                 │
   │   (KV Cache 元数据管理)           │
   │   chunk_hash → 所在节点           │
   └──────────────────────────────────┘

各组件职责

组件 角色 说明
Proxy 请求路由 接收用户请求,拆分成两个请求分发给 Prefill 和 Decode
Prefill 节点 计算 KV Cache 处理输入 prompt,生成 KV Cache
Decode 节点 生成回答 拿到 KV Cache 后逐 token 生成
Mooncake Transfer Engine 数据传输 各节点上的 RDMA 传输进程
etcd 元数据服务 记录每份 KV Cache chunk 存在哪个节点
LMCache 本地缓存管理 vLLM 内部的 KV Cache 管理,前缀缓存匹配

各节点存储配置

Prefill 节点 Decode 节点
GPU 显存(HBM) 正在计算的 KV Cache 正在生成的 KV Cache
CPU 内存(DRAM) KV Cache 缓存池(Mooncake 内存池的一部分) 显存不够时临时周转
SSD 磁盘 可选,冷数据保底层 不需要

四、Proxy 的请求拆分机制

Proxy 的核心逻辑:将一个用户请求拆成两个,分别发给 Prefill 和 Decode,通过相同的 transfer_id 配对。

发给 Prefill 的请求

req_data["kv_transfer_params"] = {
    "do_remote_decode": True,      # "你只算 KV,decode 别人做"
    "do_remote_prefill": False,
    "transfer_id": f"xfer-{request_id}",
}
req_data["stream"] = False
req_data["max_tokens"] = 1         # 只要 KV Cache,不需要生成回答

发给 Decode 的请求

req_data["kv_transfer_params"] = {
    "do_remote_decode": False,
    "do_remote_prefill": True,     # "KV Cache 从远端 Prefill 拉过来"
    "remote_bootstrap_addr": ...,  # Prefill 节点的地址
    "remote_engine_id": ...,       # Prefill 节点的引擎 ID
    "transfer_id": f"xfer-{request_id}",
}

时序

Proxy
  │
  ├── asyncio.create_task → Prefill(异步发出,fire-and-forget)
  │
  └── 立刻流式连接 Decode(不等 Prefill 返回)
  
Decode 内部会阻塞等待 Prefill 通过 RDMA 把 KV Cache 传过来,
收到完整数据后才开始生成。

Prefill 算完整个 prompt 后,一次性通过 RDMA 传给 Decode,不是边算边传。但因为 Prefill 本身只需几百毫秒,用户几乎感知不到这个等待。


五、调度策略

当前 Proxy 示例代码采用 Round-Robin 轮询,Prefill 和 Decode 各自独立轮询:

# Prefill 轮询
def prefiller_cycle(prefill_clients):
    while True:
        for prefill_client in prefill_clients:
            for i in range(prefill_client["dp_size"]):
                yield prefill_client, i

# Decode 轮询
app.state.decode_iterator = itertools.cycle(range(len(decode_clients)))

不看负载、不看缓存命中率、不看队列深度,无脑轮询。 这意味着:

  • 某个 Prefill 节点上有缓存,但请求可能被分配到另一个没有缓存的节点
  • 某个 Decode 节点已经满载,还是会继续往它派请求

这是示例代码的局限,生产环境需要实现 prefix-affinity routing(前缀亲和调度)。


六、KV Cache 缓存流转全链路

情况一:Prefill 本地命中

Prefill Node A 收到请求
    │
    ├─ 本地找到 KV Cache ✓
    ├─ 跳过计算
    └─ 点对点传给 Decode

最快路径,省掉了全部 Prefill 计算。

情况二:本地没有,远端有(部分命中)

Prefill Node A 收到请求(prompt 共 10 个 chunk)
    │
    ├─ 本地没有
    ├─ 查 etcd:chunk 1-7 在 Node B
    ├─ 从 Node B RDMA 拉取 chunk 1-7 到 Node A 本地
    ├─ Node A 自己计算缺失的 chunk 8-10
    ├─ 此时 Node A 本地有完整的 chunk 1-10
    ├─ 向 etcd 注册新计算的 chunk 8-10 元数据
    └─ Node A 把完整的 chunk 1-10 点对点传给 Decode

此时 Node B 上仍有 chunk 1-7,与 Node A 存在冗余。这是正常的:

  • Node A 挂了,Node B 的数据仍可用
  • 多个节点同时需要同一份数据时,可以分散带宽压力
  • LRU 淘汰机制会自动清理不再被访问的冗余数据

Decode 不参与这个查找过程。 它只知道"从 Prefill Node A 拿数据",不关心数据原来在哪。

情况三:全集群都没有

Prefill Node A 收到请求
    │
    ├─ 本地没有
    ├─ 查 etcd → 全集群都没有
    ├─ 从头计算整个 prompt 的 KV Cache
    ├─ 计算完 → 向 etcd 注册所有 chunk 的元数据
    └─ 点对点传给 Decode

最慢路径,需要完整 Prefill 计算。


七、Mooncake 统一内存池的本质

Mooncake 的统一内存池不是把数据集中存到某个地方,而是:

  • 数据留在计算它的那个节点上(CPU 内存 / SSD)
  • 通过 etcd 让全集群可见(注册元数据:chunk hash → 所在节点)
  • 需要时按需拉取(通过 Transfer Engine RDMA 传输)

类比:一个图书馆系统——每本书只放在某一个分馆,总馆有一个目录记录哪本书在哪,你要借书先查目录再去对应分馆取。不会在每个分馆都放一本。

KV Cache 的 Chunk 化存储

KV Cache 不以整个序列为单位存储,而是切分为固定大小的 chunk(如 256 tokens)

  • 每个 chunk 根据 token 内容计算唯一哈希
  • 相同 prompt 前缀的 chunk 在不同请求间自动复用
  • 粒度更细,复用率更高

数据不做副本

  • KV Cache 只存一份(在计算它的节点上)
  • 节点宕机 → 该节点的 KV Cache 丢失 → 下次请求触发重新 Prefill
  • 性能降级但结果正确(KV Cache 可从 prompt 重新计算)

多个节点碰巧持有相同 chunk(如情况二中 Node A 从 Node B 拉取后两边都有),属于自然冗余,由 LRU 自动管理,不是主动做副本。


八、Decode 节点上的 KV Cache 生命周期

Prefill 计算完成
    │
    │ RDMA 点对点传输
    ▼
Decode 节点接收 KV Cache → 加载到 GPU 显存
    │
    ├─ 逐 token 生成
    ├─ 每生成一个 token,KV Cache 增长一行
    ├─ 持续占用显存
    │
    ▼
生成结束(EOS / max_tokens)
    │
    └─ KV Cache 释放,显存回收

Decode 节点的 KV Cache 是"用完就扔"的工作副本:

Prefill 节点 Decode 节点
目的 缓存复用,供后续请求命中 本次生成的工作数据
生命周期 按 LRU 长期保留 生成结束即释放
大小 只有 prompt 部分的 KV prompt + 生成内容的 KV(更大)
回写 不需要,Prefill 侧已有 不需要,生成内容对其他请求无复用价值

生成过程中,同一份 prompt 的 KV Cache 在 Prefill 和 Decode 各有一份,属于正常的短暂冗余。


九、存储层级与淘汰机制

Prefill 节点的多级缓存

┌───────────┐
│ GPU HBM   │ ← 正在计算的 KV Cache
│ (~80GB)   │
└─────┬─────┘
      │
┌─────▼─────┐
│ CPU DRAM  │ ← KV Cache 缓存池(Mooncake 内存池主体)
│ (~512GB)  │
└─────┬─────┘
      │ DRAM 满了,LRU 淘汰
┌─────▼─────┐
│ NVMe SSD  │ ← 可选,冷数据保底层
│ (~2TB)    │
└───────────┘

何时数据会迁移到磁盘

只有 CPU 内存满了,才会触发 LRU 淘汰,将最久没被访问的 chunk 写入 SSD。不是主动迁移,是被动降级。

是否需要磁盘层

场景 建议
用户请求差异大,前缀重复率低 不需要,DRAM 够用,未命中直接重算
大量请求共享超长前缀(如 RAG 场景),且前缀种类多到 DRAM 放不下 需要,SSD 读取(1-5ms)比重算(100-500ms)快两个数量级

建议先只用 DRAM,上线后观察缓存命中率,命中率因内存不足频繁淘汰时再加磁盘层。


十、Chunk 大小设置建议

KV Cache 在 Mooncake 中按 chunk(固定 token 数的块)为单位存储和传输。chunk 大小直接影响缓存复用率和系统开销。

核心权衡

chunk 越小,复用率越高;chunk 越大,管理开销越低。

Prompt: "你是一个AI助手,请帮我分析这份财报..."
        ├── chunk 1 ──┤├── chunk 2 ──┤├── chunk 3 ──┤

chunk 小(如 64 tokens):
  切得细 → 两个 prompt 只要前面几十个 token 一样就能复用
  但 chunk 数量多 → etcd 元数据条目多,查询/管理开销大

chunk 大(如 512 tokens):
  切得粗 → 必须前 512 个 token 完全一致才能复用
  但 chunk 数量少 → 元数据精简,传输次数少

建议值

chunk 大小 适合场景 原因
256 tokens 通用推荐,多数场景的起点 复用率和管理开销的平衡点
64-128 tokens 短 prompt、前缀变化多(如多轮对话) 需要更细粒度的复用
512 tokens 超长系统提示词、RAG 文档前缀 前缀本身就很长且高度一致,大块传输更高效

影响因素详解

1. 复用率

系统提示词: 500 tokens(所有请求相同)
用户问题:   100 tokens(每个请求不同)

chunk = 256:
  chunk 1 (token 1-256)   ✓ 全部请求复用
  chunk 2 (token 257-500) ✓ 全部请求复用
  chunk 3 (token 501-600) ✗ 每个请求不同

chunk = 512:
  chunk 1 (token 1-512)   ✓ 复用
  chunk 2 (token 513-600) ✗ 不同

→ 这个场景下 512 和 256 复用效果差不多,512 更省元数据
多轮对话,前缀只有几十个 token 一致:

chunk = 256: 前 256 token 必须完全一样才命中 → 很难命中
chunk = 64:  前 64 token 一样就能命中 → 复用率高很多

2. 元数据开销

一个 10000 token 的 KV Cache:

chunk = 64  → 156 条元数据记录
chunk = 256 → 39 条
chunk = 512 → 20 条

chunk 越小,etcd 的存储和查询压力越大。

3. RDMA 传输效率

RDMA 每次传输有固定开销(建立连接、注册内存等)

chunk = 64:  传 39 个 chunk → 39 次 RDMA 操作
chunk = 256: 传 10 个 chunk → 10 次 RDMA 操作

大块传输更能发挥 RDMA 的带宽优势

4. 单个 chunk 的实际内存大小

chunk 的字节数取决于模型参数:

单个 token 的 KV Cache 大小 = 2 × num_layers × hidden_dim × 2(K和V) × dtype_bytes

以 Llama-70B (fp16) 为例:
≈ 2.5 MB / token

chunk = 256 tokens → 约 640 MB / chunk
chunk = 64  tokens → 约 160 MB / chunk

chunk 太大会导致单次 RDMA 传输数据量大、内存分配最小单位大容易浪费。

调优建议

初始部署
    │
    └─ 先用 256 tokens(通用起点)
         │
         ├─ 观察缓存命中率
         │
         ├─ 命中率低 + prompt 前缀短且多变
         │     → 调小到 128 或 64
         │
         ├─ 命中率高 + 前缀长且一致
         │     → 调大到 512
         │
         └─ etcd 查询延迟高 / 元数据量大
               → 调大

核心原则:chunk 大小要和业务中"前缀一致"的粒度匹配。 如果大部分请求共享的前缀长度是 2000 tokens,用 256 就很好;如果共享的只有 100 tokens,用 64 更合适。


十一、生产环境的优化方向

当前 Proxy 示例代码是 demo 级别,生产环境需要关注:

方面 现状 生产需要
调度策略 Round-Robin 无脑轮询 前缀亲和调度(prefix-affinity routing)
高可用 Proxy 单点 多实例 + 负载均衡
Prefill 失败处理 fire-and-forget,失败静默 重试 / 回退到其他节点
监控 缓存命中率、节点负载、传输延迟
超时控制 timeout=None 合理超时 + 熔断
Decode 容量 不感知 基于并发生成数和显存使用率调度

前缀亲和调度的核心思路:

用户请求到达 Proxy
    │
    ▼
计算 prompt 前缀的 hash
    │
    ▼
查表:这个 hash 的 KV Cache 在哪个节点?
    │
    ├── 命中 → 发给对应节点
    └── 未命中 → 按负载选最闲的节点

让请求去找缓存,而不是让缓存来找请求。

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐