跳转至

Mooncake Codebase Architecture

导言

上一篇文章从 FAST'25 论文出发,解释了 P/D 解耦、分布式 KV Cache 与调度机制。论文读懂以后,打开代码仓却很容易再次迷路:Connector、Mooncake Store、Transfer Engine、TENT 看起来像四个并列组件,实际却跨越 vLLM 与 Mooncake 两个仓库,并分别承担框架适配、对象管理、字节搬运和新传输内核

本文固定在 Mooncake 6a00c353 与 vLLM 5bbc58c0,从仓库结构、请求流、时序和关键类关系重新组织这些概念,最后给出一套可重复的源码走读与开发 SOP。

一句话结论:Connector 把推理框架语义翻译成 KV 操作,Store 管对象与生命周期,Transfer Engine / TENT 搬运字节;读代码时必须把控制面与数据面分开追。

先纠正命名错觉

论文中的组件名描述的是系统职责,仓库目录描述的是可构建模块,两者不是一一对应关系。尤其要先接受三个事实:

  1. Connector 的主实现属于 vLLM。当前 vLLM 内置 MooncakeConnectorMooncakeStoreConnector;Mooncake 仓内的 mooncake_connector_v1.py 是兼容旧版 vLLM 的 out-of-tree 实现,代码也明确提示 vLLM 0.13.0 之后应使用内置 Connector(mooncake_connector_v1.py:507-509)。
  2. Mooncake Store 不负责发明传输协议。它在 Transfer Engine 之上增加 key、replica、lease、checksum、淘汰和失败收敛等对象语义。
  3. TENT 不是另一个上层系统。它是 Transfer Engine NEXT,在同一个 mooncake::TransferEngine 外层接口后面与 Classic TE 二选一。

版本边界

本文解释的是上述两个固定提交。Mooncake 仍在快速演进,目录、Python 包装与 vLLM Connector 接口会继续变化。阅读最新 main 时,应先重新固定提交号,再验证本文给出的入口是否仍然存在。

仓库不是一个 Engine

Mooncake 代码仓分层

图:Mooncake 与 vLLM 的代码分层。蓝色表示框架适配和控制流,橙色表示数据搬运主路径。

从上往下看,代码可以分为五层:

层次 主要目录或类 负责什么 不负责什么
推理框架适配 vLLM MooncakeConnector* 接收 scheduler / worker hooks,处理 token 命中、block 分配与 load/save 不定义 RDMA、TCP 等传输细节
Python/C++ 绑定 mooncake-integration/python/mooncake/mooncake-wheel/ 暴露 mooncake.enginemooncake.store 不承担核心调度策略
Store 语义层 mooncake-store/ 管理对象、replica、lease、checksum、淘汰与失败收敛 Master 不转发 KV 数据
字节搬运层 mooncake-transfer-engine/ 注册内存、打开远端 segment、提交批量读写、聚合状态 不理解 token、layer 或 KV block
扩展能力 mooncake-ep/mooncake-pg/mooncake-reshard/mooncake-p2p-store/ EP 通信、进程组、权重重分片或轻量 P2P Store 不是 Connector → Store → TE 主链的必经模块

CMakeLists.txt 默认打开 WITH_TEWITH_STORE,而 WITH_EPWITH_P2P_STORE 等能力按需开启(CMakeLists.txt:19-30)。这说明仓库更像一个共享基础设施单仓,不是只能整体启动的单体服务。

mooncake-common/ 则是横跨这些模块的公共底座,放置元数据/RPC、构建选项与通用定义;mooncake-integration/ 通过 pybind 把 C++ TransferEngine 和 Store 暴露给 Python(mooncake-integration/CMakeLists.txt)。因此从 Python 入口向下走读时,不应在包装层停住,而要继续跨到对应的 C++ facade。

Connector 是边界适配器

Connector 的核心价值不是“搬数据”,而是把 vLLM 的请求生命周期翻译成下游能理解的操作。当前有两条语义不同的路径。

Connector 选择与请求流

P/D 直接传输

vLLM 的内置 MooncakeConnector仍按 scheduler / worker 分工:

  • Scheduler 侧决定请求可复用多少 token、需要分配哪些目标 block,并生成 Connector metadata。
  • Worker 侧注册 KV cache 内存区域,加载 mooncake.engine.TransferEngine,通过 bootstrap / ZMQ 控制通道交换 endpoint 与 block 信息。
  • Transfer Engine根据远端 segment 与偏移,把 Producer KV 直接写入 Consumer 已分配的 KV 区域。

这条路径是点到点的 P→D 交接。控制通道负责发现和握手,KV 数据不经过 Store Master,也不要求先形成一个可跨请求查询的 Store 对象。

MooncakeConnector 的 P 到 D 直传时序

图:概念时序聚焦代码职责,不等同于某个 vLLM 版本的全部异步回调顺序。粗橙线才是 KV 数据。

共享外部 KV 池

vLLM 的 MooncakeStoreConnector同样拆成 scheduler / worker,但语义换成了共享缓存:

  1. Scheduler 查询外部 Store 命中,生成请求的 LoadSpecscheduler.py:51-150)。
  2. Worker 初始化 MooncakeDistributedStoreworker.py:1330-1408)。
  3. 加载时调用批量 get,把命中对象写入本地 KV buffers;保存时调用批量 put,把新 KV 注册为可复用对象。
  4. Store Client 先向 Master 获取 replica 元数据,再通过 Transfer Engine 与持有数据的远端 Client 传输。

两者的选择标准不是“哪个更新”,而是目标语义:

问题 MooncakeConnector MooncakeStoreConnector
主要目标 一次 P→D 交接 跨请求、跨实例复用 KV
控制信息 endpoint、目标 block、请求状态 key、replica、lease、命中信息
数据路径 Producer → Consumer Store Client → Store Client
Master 不需要 管元数据,不搬数据
适合的走读起点 Connector Worker Store Scheduler / Worker

Store 管对象,TE 搬字节

Store 最关键的设计边界是:Master 在控制面,Client 在数据面。官方设计文档也把 Store 概括为 Master Service 与 Client 两个核心组件,并明确数据在 Client 之间传输而绕过 Master(mooncake-store.md)。

Get 为例,公开入口先查询对象的位置和 lease,再进入带 QueryResult 的读取:

auto query_result = Query(object_key);
if (!query_result) {
    return tl::unexpected(query_result.error());
}
return Get(object_key, query_result.value(), slices);

源码:client_service.cpp:1139-1145。其中 Query 实际调用 master_client_.GetReplicaList,并把 replica、lease TTL 与 checksum 组装成 QueryResultclient_service.cpp:1211-1222)。随后真正的 Get 选择可用副本、执行 TransferRead,并校验 checksum 与 lease。

Put 则是一个三阶段协议:

  1. PutStart:向 Master 申请副本位置并取得写入计划。
  2. TransferWrite:Client 按 memory、disk/DFS 等副本类型传输切片。
  3. PutEnd / PutRevoke:根据成功副本数提交或撤销元数据。

这也是 Store 比裸 Transfer Engine 多出的价值:字节已经写到远端,不代表一个对象已经满足复制策略并可安全对外可见。相关收敛逻辑位于 Client::Put

Store 到 Transfer Engine 的桥是 TransferSubmitter。它把对象切片变成 TransferRequest,申请 batch ID,提交给稳定的 engine_ facade,再用 TransferFuture 聚合完成状态:

BatchID batch_id = engine_.allocateBatchID(batch_size);
Status s = engine_.submitTransfer(batch_id, requests);

源码:transfer_task.cpp:1239-1271。到这里,key、replica 和 lease 已经被压缩成 source + target_id + target_offset + lengthTransfer Engine 从此只看字节范围。

Mooncake 关键类关系

TENT 如何接入

TENT 是 Transfer Engine NEXT。它引入声明式 Request、动态 TransportSelector、切片级调度、运行时故障收敛和可插拔 transport,但仍隐藏在原有 mooncake::TransferEngine facade 后面(TENT overview)。

这里有一个容易漏掉的双重开关

  1. 构建时使用 -DUSE_TENT=ON,让 mooncake-transfer-engine/tent/ 进入构建;默认值是 OFF(common.cmake:174)。
  2. 运行时设置 MC_USE_TENT=1;否则外层 facade 仍创建 Classic TransferEngineImpl
if (getenv("MC_USE_TENT") || getenv("MC_USE_TEV1")) {
    use_tent_ = true;
}

源码:transfer_engine.cpp:393-409MC_USE_TEV1 是兼容别名;真正的分派发生在 initallocateBatchIDsubmitTransfer 与状态查询等 facade 方法中。

当走 TENT 分支时,facade 会把 Classic TransferRequest 逐字段转换为 tent::Request,再调用 impl_tent_->submitTransfer;状态查询也反向翻译为旧的 TransferStatustransfer_engine.cpp:621-665transfer_engine.cpp:750-778)。因此上层 Store 与 Connector 不必同时重写,但编译进 TENT 不等于运行时已经使用 TENT

进一步走读

如果要继续追 TENT 的文件传输和请求对象换手,可分别阅读 Mooncake TENT GDSMooncake TENT Request Path。这两篇聚焦 TENT 内部,本文只解释它在整个仓库中的位置。

论文组件与开源边界

论文中的 Conductor 是全局调度与准入组件,但在本文固定的 Mooncake 根提交中,情况需要谨慎描述:仓库包含 conductor-architecture-design.md,文档设想了 mooncake-conductor/ 的 Go 目录结构;然而该提交的根目录树并没有可构建的 mooncake-conductor/ 源码目录。

所以更准确的说法是:

  • 论文 Conductor解释生产系统应承担哪些调度职责。
  • 当前开源 Connector落实推理框架内部的 scheduler / worker 生命周期适配。
  • Store Master只管理对象与副本元数据,不能被直接等同为论文 Conductor。

这三个控制面彼此有关,但不应根据名字或一张论文图强行合并成同一个类。

典型源码 SOP

面对这种跨仓、跨语言、带编译开关的系统,最有效的 SOP 不是从根目录顺序读文件,而是围绕一个数据对象和一个场景穿刺。

走读一个请求

  1. 固定 revision。同时记录 Mooncake 与集成框架提交号,避免一边读新版 vLLM Connector、一边引用旧版 Mooncake Python 包装。
  2. 写清对象和场景。例如“一个请求的第 12 层 KV,从 Prefill GPU 传到 Decode GPU”,或“一个 prefix key 从 Store 命中后写入本地 KV blocks”。
  3. 选择 Connector 路径。一次 P→D 交接从 MooncakeConnectorWorker 开始;共享缓存命中从 MooncakeStoreConnectorScheduler/Worker 开始。
  4. 分开画两条线。控制面记录 key、endpoint、replica、lease 与状态;数据面只记录 source、target、offset、length 与 transport。
  5. 跨过绑定层。看到 mooncake.enginemooncake.store 后,立即找到 mooncake-integration/ 的 pybind 定义和对应 C++ 类。
  6. 追到稳定 facade。Store 路径应经过 Client → TransferSubmitter → TransferEngine;再根据构建与环境变量判断进入 Classic TE 还是 TENT。
  7. 最后找测试。先运行离改动最近的单元测试,再扩大到 Store/TE 集成测试;不要先尝试全仓硬件矩阵。

建议维护一张最小走读表:

步骤 入口 输入对象 输出对象 控制/数据 下一跳
1 Connector Scheduler request + tokens connector metadata 控制 Connector Worker
2 Store Client object key + slices replica plan 控制 MasterClient
3 TransferSubmitter slices + replica TransferRequest[] 数据 TransferEngine
4 TE facade batch + requests Classic/TENT request 数据 transport

修改和验证代码

一个面向贡献的常规流程可以压缩为:

  1. 限定模块。PR 标题使用 [Store][TransferEngine][TENT][Integration] 等前缀;超过 500 行且不含测试的架构变化先写 RFC issue(CONTRIBUTING.md:17-40)。
  2. 补最近的测试。Store 语义至少覆盖元数据成功/失败和传输成功/失败;transport 修改至少覆盖 submit、poll、部分失败和资源释放。
  3. 窄构建。先只打开需要的模块和 transport,保留 BUILD_UNIT_TESTS=ON
  4. 窄测试。先运行目标目录 CTest/pytest,再参考 GitHub Actions 扩展配置。
  5. 只格式化改动文件。提交前运行基于 origin/main...HEAD 的 pre-commit,避免把全仓历史格式噪声带进功能 PR(CONTRIBUTING.md:67-85)。

基础 TE + Store 的示例构建骨架是:

git checkout 6a00c3534d87b47561e78e07cdc5f3277bef25de
cmake -S . -B build -G Ninja \
    -DWITH_TE=ON -DWITH_STORE=ON \
    -DWITH_STORE_RUST=OFF \
    -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=OFF \
    -DBUILD_BENCHMARK=OFF -DUSE_TCP=ON
cmake --build build -j
ctest --test-dir build --output-on-failure

TENT 的构建与运行时选择应同时出现:

cmake -S . -B build-tent -G Ninja \
    -DWITH_TE=ON -DWITH_STORE=OFF \
    -DWITH_STORE_RUST=OFF \
    -DUSE_TENT=ON -DUSE_TCP=ON \
    -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=OFF \
    -DBUILD_BENCHMARK=OFF
cmake --build build-tent -j
MC_USE_TENT=1 ctest \
    --test-dir build-tent/mooncake-transfer-engine/tent/tests \
    --output-on-failure

官方 CI 的 TENT lane 同样使用 -DUSE_TENT=ON 构建,并从 mooncake-transfer-engine/tent/tests 运行 CTest(.github/workflows/ci.yml:627-680)。但这些命令是否能在本机直接通过,还取决于依赖、元数据服务、RDMA/GPU/SSD 等环境;它们是验证入口,不是脱离环境的成功承诺。

git fetch origin main
pre-commit run --files $(git diff --name-only \
    --diff-filter=ACMR origin/main...HEAD)

不要迷信统一脚本

在固定提交 6a00c353scripts/ 中没有 run_ci_test.sh。遇到文档或旧笔记提及它时,应以当前 .github/workflows/ci.yml、模块 CMakeLists.txt 和实际测试目录为准,而不是假设仓库存在一个永远稳定的全量入口。

总结

从代码实现看,Mooncake 的主链可以压缩成两种形式:

Direct P→D:
vLLM MooncakeConnector
    → mooncake.engine.TransferEngine
    → Classic TE / TENT
    → remote KV memory

Shared Store:
vLLM MooncakeStoreConnector
    → mooncake.store.MooncakeDistributedStore
    → mooncake::Client
    → MasterClient(元数据) + TransferSubmitter(数据)
    → mooncake::TransferEngine
    → Classic TE / TENT

真正稳定的理解方式不是背目录,而是始终追问三件事:谁把请求语义翻译成 KV 操作,谁决定对象是否有效,谁实际搬运字节。一旦把 Connector、Store 和 Transfer Engine 放回这三个职责中,论文架构、vLLM 集成与 Mooncake C++ 内核就能对齐;TENT 也不再是一个突然出现的“新 Engine”,而是稳定 facade 后面的下一代搬运实现。

参考资料

评论