AI

Mooncake 统一内存池:AI Vibe Coding 与 Rust

Posted by w@hidva.com on March 15, 2026

最近我在做 Mooncake 接入 RL, 收益颇丰. 正好借这个机会, 总结下我们基于 Mooncake 所做的 local master / unify memory pool 设计. 一开始看上去, 这像是一个“把 DummyClient.get_buffer()/put_parts() 做快一点”的问题. 但真正把链路扒开来看就会发现, 问题远不止如此. 真正的问题是: 在一个 Pod 内, 数据平面被拆成了 global memory、local memory、hot cache、以及每个 dummy client 各自的 shm; 同一份数据在不同池子里来回复制, 生命周期和元数据也散落在不同地方. 只要这个结构不改, 拷贝次数, 容量规划, 驱逐策略, 以及后续扩展都会越来越别扭.

问题不只是少几次拷贝

在 RL, KVS 这些典型部署里, 一个 Pod 往往是 1 个 real client 带多个 dummy client. 此时 real client 这边会同时维护:

  1. global mem
  2. local mem
  3. local hot cache
  4. 与每个 dummy client 两两共享的 shm

DummyClient 的一次 get_buffer()/put_parts() 往往都夹着好几次本地中转:

DummyClient.get_buffer()
  -> RealClient
  -> local hot cache miss
  -> 从远端拉到 real client 侧 hot cache
  -> 再拷贝到 dummy 专属 shm
  -> dummy 基于 shm 构造 tensor/numpy

DummyClient.put_parts(key, data)
  -> data 经 rpc 拷贝到 real client 临时内存。
  -> 再从临时内存拷贝到 dummy 专属 shm
  -> 再次 dummy 专属 shm 传输到远端某个 global mem 中。

所以问题根本不是“这里多了一次 memcpy, 那里少了一次 memcpy”. 真正的问题是, 本地数据平面一开始就不是一个统一的资源视图. allocator 是散的, 副本是散的, 生命周期管理也是散的. 同一个 key 在 Pod 内可能同时躺在 hot cache、某个 dummy 的 shm、甚至别的局部 buffer 里. 这种模型下, 你很难得到一个真正统一的 zero-copy 访问路径.

另外还有一个问题也会越来越明显: Master 逐渐变成瓶颈. 当前架构里, 不管是 real client 还是 dummy client, 很多元数据操作、连接建立、资源协调都要和 Master 交互. 当 client 数量上去, 并发 get/put 上去之后, Master 既要处理心跳和状态同步, 又要响应大量元数据查询和分配请求. 这件事和多池模型其实属于同一类问题: 原本很多本地就可以解决的事情, 被拆成了很多需要额外协调的动作.

Local Master + 统一共享池 G

我们的做法, 是继续沿着 local master 这条路往下走, 并且进一步把 Pod 内的数据平面也彻底收拢起来. local master 解决的是控制面本地化, unify memory pool 解决的是数据面统一化: 把原先散落的多个内存池合并成统一共享池 G, 让本地副本、共享访问和 backing store 回到同一个资源视图里. 简单来说, 就是在 Pod/Node 内只保留一块统一共享池 G:

  1. RealClient 启动时创建 G
  2. DummyClient 启动时直接 attach 这块共享内存
  3. 本地对象缓存、共享访问区、以及反序列化后 tensor 的 backing store, 全都统一放进 G

这样原先分裂的 localmem、local hot cache、各个 dummy 专属 shm, 就都可以收敛成一个池子了. 看起来既然都共享同一块内存了, 那是不是直接把指针传来传去就结束了? 还差一点. 因为同一个 shm fd 在不同进程里 mmap 之后, 虚拟地址未必一样. 所以 RealClient 看到的 real_addr, 到 DummyClient 这边要先转成 offset, 再加到自己的 g_base 上. 这个细节不复杂, 但它把“共享一块物理内存”和“跨进程直接复用虚拟地址”这两件事分得很清楚. 统一共享池 G 之后, 有三件事情也必须一并统一:

  1. 统一 allocator. 所有本地 buffer 都从 G 里直接分配, 不再在多个池子之间各自为战.
  2. 统一元数据表. 用统一的 MetaMap / HandleTable 记录 key 到 buffer 的映射、大小、状态、引用信息.
  3. 统一生命周期管理. 通过 replica ref cnt 管理 memory replica 的生存期, 确保上层 tensor 还在引用时, 底层 buffer 不会被错误回收.

这个生命周期管理尤其关键. 因为 get_buffer() 最终反序列化出来的 tensor, 是直接引用 G 中那块内存的. 一旦底层 buffer 释放早了, 上层对象就直接悬空了. 这也是为什么 AcquireReplica / ReleaseReplica 这套引用计数机制在 unified pool 设计里不是可选项, 而是核心项.

另外 get_buffer(key) 这里还有个挺有意思的小细节. 同一个 dummy client 可能会多次获取同一个 key, 这时返回的底层 replica.ptr 很可能是同一个地址. 如果直接拿 ptr 当 active handle 的 key, 后一次会把前一次覆盖掉. 所以这里最终还是需要一个单调递增的 handle_id, 用来区分多次 get_buffer() 调用各自的生命周期.

Put/Get 链路怎么变短了

统一池子之后, 真正最直观的变化还是 put/get 链路本身.

Get: Pod 内只保留一份本地副本

原来的路径是:

dummy -> real -> hot cache -> dummy shm -> tensor

现在则变成:

  1. 先查 LocalMaster 的共享 MetaMap
  2. 如果 key 已经在本地, 直接 AcquireReplica
  3. 做一次地址翻译, 之后 split + 反序列化
  4. tensor 直接引用 G 中已有数据

如果 key 不在本地, 也不是像以前那样先拉到 hot cache, 再拷到别的池子里. 现在是:

  1. 先查远端 replica 信息
  2. 由 RealClient 执行一次 LocalCopy
  3. 通过 PutStart -> Get -> PutEnd 直接把远端数据写入共享池 G
  4. 更新本地元数据
  5. 后续所有 client 直接复用这一份本地副本

也就是说, 同一个 key 在一个 Pod 内只保留一份 memory replica. 这一点很重要. 因为它不只是让这一次 get 变快, 而是从结构上消除了“每多一个 dummy client, 就可能多一份局部副本”的趋势.

这里还有一个当时专门想过一轮的问题: LocalCopy 要不要传播 softpin. 我最后的结论是不要. 社区引入 softpin，如文档所示是用于 frequently accessed or important objects like system prompts；我个人理解是因为目前驱逐策略不是严格（或近似） LRU 下取舍的产物，实际上当前驱逐策略可以任务是某种随机的，也就是如果我们支持了严格/近似 LRU 策略，那么 system prompts 这类 frequently accessed objects 天然就会对驱逐免疫。在 RL/Omni 使用场景中，其更像是把 mooncake 作为一种数据中转站来使用，其使用模式一般是 put, get, remove；也即其需要的 softpin 机制是 put(key, softpin=True) 那么 mooncake 要确保在指定时间内这个 key 不能被驱逐；基于这个背景下，我们没有必要传播 softpin 了。

Put: 从“先写中转区”改成“原地序列化”

put 路径的变化也很直接:

  1. DummyClient 直接向 LocalMaster 申请 G 中一块连续空间
  2. msgspec.msgpack + enc_hook 直接把对象序列化到这块空间里
  3. 写完之后更新 MetaMap
  4. 其它 attach 到 G 的 client 立即可见

这里最核心的收益是: 去掉了先写私有 localmem, 再搬到别处去的那次中转. 从“先找个临时地方落一下”改成“直接写最终位置”. 这件事说起来平平无奇, 但一旦数据大起来、调用频繁起来, 差别就非常实在了.

一次还挺典型的 AI vibe coding

这次开发过程本身, 我觉得也挺适合拿来聊 AI vibe coding. 我现在越来越觉得, 那类“架构边界已经比较明确, 但链路长、机械改动多、需要反复对齐项目风格”的任务, 很适合让 AI 参与. 这次比较顺的原因, 大概有这么几条:

  1. 先给伪代码, 再让 AI 填细节. 比如 LocalCopy、acquire_buffer_dummy_local、release_buffer_dummy_local、local_get_buffer 这几条链路, 先把骨架画出来, AI 再去补实际实现. 这种方式比一句“帮我把 get_buffer 改成支持 local master”有效得多.
  2. 每一步都编译, 每一步都测. AI 很擅长铺代码, 但人最好不要等它一口气铺完整个功能再一起收尸. 一步一编译, 一步一测试, 反馈成本最低.
  3. 测试拓扑提前说清楚. 比如 real client 1/2, dummy client 11/12/21/22, 哪个 put, 哪些 get, 跨不跨 real client, 预期结果是什么. 这种带拓扑的测试描述非常适合 AI 执行.

我目前越来越倾向于把 AI 当一个不怕烦的实现同伴: 你给边界, 给伪代码, 给验收条件, 它负责把那些又长又机械的部分尽快铺出来.

为什么我觉得 Rust 更适合 AI vibe coding

最后说点题外话, 但其实也不算太题外. 我现在越来越觉得, Mooncake store 这种东西就应该用 Rust 写. Store 本身是 I/O 密集型系统, 天然适合 async; 而真要把 async、性能和工程可靠性放在一起看, Rust 依然是我最顺手的选择. 另外, Rust 也确实更适合 AI vibe coding.

第一个原因是, Rust 本身就更适合表达意图. enum + match、Result + ?、trait、pattern matching 这些语法设施, 很多时候能把“有哪些状态”、“错误怎么传播”、“谁拥有这块数据”写得非常直接. 这件事对人当然有帮助, 对 AI 也一样有帮助. 代码里显式表达出来的信息越多, AI 越不容易靠猜.

第二个原因更有意思: 人类最烦 Rust 的地方, 往往恰好是 AI 最不怕烦的地方. 比如 lifetime. 对人来说, lifetime 很容易让人写着写着就烦了, 尤其在 borrow 边界复杂、需要来回调整所有权的时候. 但对 AI 来说, 这类事情反而很适合. 它可以老老实实地根据编译器报错去缩短 borrow, 拆 owned/borrowed 结构, 调整函数签名, 反复改到通过为止. 很多时候, lifetime 更像是一种机械但繁琐的对齐过程, 这类体力活其实挺适合交给 AI.

第三个原因则更现实: Rust 编译通过之后, 可信度就是更高. 至少我个人会更信它. C++ 编译通过, 很多时候只是说明模板和语法终于拧过去了, 绝不意味着内存一定没问题. dangling pointer、use-after-free、double free、悬垂引用、迭代器失效, 这些东西完全可能在“编译成功”的前提下继续往线上跑. Rust 则不一样, borrow checker 和类型系统至少先把一大类最讨厌的坑堵掉了. 比如我们之前写 C++ coroutine 时, 还专门总结过几条使用规则:

提供的很多接口接受的是一个lambda参数，这个lambda有可能会在第一次co_await之后就析构了，也可能不析构，具体取决于接口的实现。常见会析构的的接口有：parallel_for_each、ec Invoke。这种情况建议的写法是把整个lambda抽成一个普通函数或者成员函数，然后parallel_for_each里直接return

coroutine变量生命周期一直是困扰我们的一个大问题，引用类型函数参数就是其中一种。值类型的参数在co_await时会被放到协程栈上，编译器就已经保证了生命周期。但是引用或指针类型的参数，编译器只是把一个指针放到了协程栈上，这个指针实际指向的对象可能就会析构了。

这类规则的麻烦不在于“做不到”, 而在于它太细、太隐式、太依赖经验. 人尚且容易忘, AI 就更容易在上下文不完整时踩坑. Rust 则倾向于把这些约束前置到类型系统、borrow checker 和编译器报错里. 这会让开发前期更拧巴一点, 但也让 AI 更容易沿着编译错误把事情一点点改对.

CATALOG