graph TB subgraph Mooncake集群 M["Master Service\n(元数据 + 空间管理)"] MD["Metadata Service\n(可选,或P2P握手)"] S1["Store Service 1\n(贡献内存到全局池)"] S2["Store Service 2"] SN["Store Service N"] end
subgraph SGLang节点 SGL1["SGLang Server\n(Mooncake Client)"] SGL2["SGLang Server\n(Mooncake Client)"] end
M <-->|"分配/释放空间\n驱逐策略"| S1 M <-->|分配/释放空间| S2 M <-->|分配/释放空间| SN MD <-->|连接握手元数据| SGL1 MD <-->|连接握手元数据| SGL2
subgraph "L2 — CPU DRAM(快,中等容量)" CPU["Host KV Cache Pool\n(MHATokenToKVPoolHost)\n默认 2× GPU VRAM 大小"] end
subgraph "L3 — 分布式 Mooncake(慢但近乎无限)" MC["MooncakeStore\n(RDMA 零拷贝,跨节点共享)"] end
REQ["新请求"] -->|"前缀命中查询\n(HiRadixCache)"| GPU GPU -->|L1 miss| CPU CPU -->|L2 miss,prefetch| MC GPU -->|"evict (write-through/write-back)"| CPU CPU -->|"evict (backup)"| MC MC -->|"load back"| CPU CPU -->|"load back"| GPU
HiCacheController 负责调度以上所有层间的异步数据迁移:
write-through: 写入 L1 同时写入 L2/L3
write-back: eviction 触发时写入 L2/L3
prefetch: 请求到来前提前从 L3 加载到 L2
6. MooncakeStore 核心工作流
6.1 初始化流程
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
flowchart TD A["MooncakeStore.__init__()"] --> B["_import_mooncake_store()"] B --> C["MooncakeDistributedStore()"] C --> D["_load_config()\n优先级: extra_config > JSON文件 > 环境变量"] D --> E{standalone_storage?} E -->|Yes| F["store.setup_dummy()\n通过本地 RPC 连接独立 Store Service"] E -->|No| G{"共享 TransferEngine?"} G -->|Yes, 配置相同| H["复用已有 mooncake_transfer_engine"] G -->|No| I["创建新 transfer_engine"] H --> J["store.setup()"] I --> J F --> K["warmup()\n写入+读取4KB数据验证连通性"] J --> K K --> L["register_mem_pool_host()\n注册 HostKVCache buffer 到 RDMA"] L --> M["Ready"]
6.2 batch_set(KV 写入 L3)完整流程
1 2 3 4 5 6 7 8 9 10 11 12 13 14
flowchart TD A["batch_set_v1(keys, host_indices)"] --> B["_batch_preprocess()\n生成 key_strs, buffer_ptrs, buffer_sizes"] B --> C{is_mla_backend?} C -->|Yes| D["_get_mla_buffer_meta()\n每个 key 对应一个 KV 合并分量"] C -->|No| E{should_split_heads?} E -->|Yes| F["_get_mha_split_heads_buffer_meta()\n按 split_factor 展开为多个 rank 的 key"] E -->|No| G["_get_mha_buffer_meta()\n每 key → k + v 两个分量"] D --> H["_batch_exist(key_strs)\n批量检查哪些 key 已存在"] F --> H G --> H H --> I["过滤出未存在的 key\n跳过已有数据(幂等写入)"] I --> J["_put_batch_zero_copy_impl()\nbatch_put_from(ptrs, sizes)\nRDMA 零拷贝写入 Mooncake"] J --> K["_batch_postprocess()\n解析每对 k/v 的成功状态"] K --> L["返回 List[bool]"]
6.3 batch_get(KV 从 L3 读取)完整流程
1 2 3 4 5 6 7 8 9
flowchart TD A["batch_get_v1(keys, host_indices)"] --> B["_batch_preprocess()"] B --> C["_get_batch_zero_copy_impl()\nbatch_get_into(ptrs, sizes)\nRDMA 零拷贝读入 HostKVCache"] C --> D["_batch_postprocess()\n检查返回字节数 > 0 表示成功"] D --> E["返回 List[bool]"]
B2["batch_exists(keys)"] --> F["构造 query_keys(加 _k/_v 后缀)"] F --> G["_batch_exist(query_keys)\nbatch_is_exist()"] G --> H["返回连续命中页数(整数)"]
sequenceDiagram participant User as 用户请求 participant Enc as MMEncoder (rank 0) participant Local as 本地 hash_to_metadata participant MC as Mooncake 集群 participant ViT as ViT 模型
User->>Enc: 含图/音的多模态请求
Note over Enc: Step 1: 批量检查缓存命中 Enc->>Local: batch_is_exist(image_hashes) alt 本地命中 Local-->>Enc: hit = True (无需网络) else 本地未命中 Enc->>MC: batch_is_exist(missing_hashes) [async] MC-->>Enc: 全局命中情况 end Note over Enc: Rank 0 广播命中 mask 到所有 TP ranks
Note over Enc: Step 2: 所有 TP ranks 并行计算 Cache Miss Enc->>ViT: _encode_missing(missing_indices) ViT-->>Enc: new_embedding_slices
sequenceDiagram participant Router as Router participant P1 as Prefill Worker 1\n(HiCache enabled) participant P2 as Prefill Worker 2\n(HiCache enabled) participant MC as Mooncake 集群 participant D as Decode Worker
Note over P1: 请求 A,新前缀 [sys_prompt + user_turn_1] Router->>P1: 请求 A P1->>P1: HiCache miss → 全量计算 KV P1->>MC: batch_set_v1(keys, host_indices)\n[write-through / write-back]
Note over P2: 请求 B,相同前缀 [sys_prompt + user_turn_1] Router->>P2: 请求 B P2->>P2: L1/L2 miss P2->>MC: batch_exists(keys) → 命中! P2->>MC: batch_get_v1(keys, host_indices)\n[RDMA 零拷贝加载到 CPU DRAM] MC-->>P2: KV pages → CPU DRAM (L2) P2->>P2: transfer L2→L1 (GPU) P2->>D: KV Transfer (RDMA) D->>D: Decode 阶段
subgraph "Mooncake 集群" MCM["Master Service"] MCDM["Metadata Service"] MCP1["Store Service 1"] MCP2["Store Service 2"] end
LLM <-->|"读写 KV"| L1 HRC -->|"evict/load"| CC CC <-->|"GPU↔CPU 传输"| L1 CC <-->|"L2↔L3 传输"| L2 CC <-->|"batch_set/batch_get\n(RDMA 零拷贝)"| MS MS <-->|"RDMA"| MCP1 MS <-->|"RDMA"| MCP2 MCM -->|"管理"| MCP1 MCM -->|"管理"| MCP2 MCDM <-->|"握手"| MS
ENC -->|"图像/音频"| EC EC -->|"batch_get/put"| MES MES <-->|"RDMA"| MCP1 EC -->|"Embedding 向量"| LLM
10.2 mooncake_store 目录的两大核心价值
用途
类
场景
核心收益
KV Cache L3 后端
MooncakeStore
PD 解耦 / 长上下文 / 高并发
突破单节点内存限制,跨实例 KV 共享,RDMA 高带宽
Embedding 全局缓存
MooncakeEmbeddingStore + EmbeddingCacheController
EPD 多模态
消除重复 ViT 计算,跨节点 Embedding 共享,后台异步 I/O
10.3 设计亮点
零拷贝优先:通过 MooncakeHostTensorAllocator 将 CPU DRAM 直接纳入 RDMA 管理,所有数据传输均通过指针完成,无额外内存拷贝。