团子云技术 Lite 1.048596
Go back

Mooncake TE 阅读手记-02-Buddy Allocator 与两层池化

团团虾导读：registerLocalMemory 涉及 RDMA MR 注册（内核态操作）+ etcd 同步，不能频繁调用。Mooncake 的解法是在 Segment 之上构建两层内存池——Python 层的 Buddy Allocator（16 个 size class，2GB 大块）和 Store 层的 Offset Allocator。这篇追踪了 doBuddyAllocate 的递归分裂逻辑和 freeManagedBuffer 的只放回 free_list 不释放设计。

[Blog 1] Buffer, Segment, Slice — Transfer Engine 的三层数据抽象

阅读版本: Mooncake v0.3.10.post2-104-geaf724ab (commit eaf724ab, 2026-05-20)

理解 Transfer Engine 的架构，首先要理清它在不同层次上如何描述”内存”和”数据传输”。TE 定义了三个核心抽象层次：Buffer（可操作的内存块）、Segment（可发现的数据空间） 和 Slice（传输的最小执行单元）。三者是严格的嵌套关系。

1. Buffer — 已注册的真实内存块

在 TE 的上下文里，Buffer 不是一个笼统的缓冲区概念，而是特指 BufferDesc 描述的一段已经注册到 RDMA 网卡（或其它 Transport）的连续物理内存。

数据结构（mooncake-transfer-engine/include/transfer_metadata.h:52-65）：

struct BufferDesc {
    std::string name;
    uint64_t addr;              // 虚拟地址
    uint64_t length;            // 大小
#ifdef ENABLE_MULTI_PROTOCOL
    std::string protocol;       // 多协议模式下的传输协议 (cxl/tcp/rdma)
#endif
    std::vector<uint32_t> lkey; // RDMA 本地密钥 (每 HCA context 一个)
    std::vector<uint32_t> rkey; // RDMA 远程密钥 (每 HCA context 一个)
    std::string shm_name;       // NVLink/HIP 共享内存名
    uint64_t offset;            // CXL 偏移量
    std::vector<std::string> tseg;      // UB/URMA 段名
    std::vector<uint32_t> l_seg_index;  // UB/URMA 段索引
};

关键点：每个 BufferDesc 背后都对应一次真实的硬件内存注册。以 RDMA 为例，注册时会调用 ibv_reg_mr() 将内存页 pin 住并获取 lkey/rkey。这是一次昂贵的内核态操作，也是为什么不能频繁注册/注销的根本原因。

注意 protocol 字段位于 #ifdef ENABLE_MULTI_PROTOCOL 条件编译块内，仅在多协议模式下存在。

Transport 层的通用 Buffer 入口（mooncake-transfer-engine/include/transport/transport.h:375-378）：

struct BufferEntry {
    void *addr;
    size_t length;
};

这是批量注册/注销使用的精简版描述符，只含地址和长度，location 等元数据由外部传入。

2. Segment — 可发现的命名数据空间

Segment 是注册内存的”目录”。一个进程把自己的内存注册到 TE 后，这些 Buffer 被组织在一个有名字的 Segment 下面，远端节点通过 openSegment("name") 发现并获取访问所需的所有信息。

SegmentDesc 结构（mooncake-transfer-engine/include/transfer_metadata.h:88-108）：

struct SegmentDesc {
    std::string name;           // 全局唯一名字
    std::string protocol;       // rdma / tcp / cxl / ...
    std::vector<DeviceDesc> devices;   // HCA/NIC 设备信息
    Topology topology;                  // 拓扑信息
    std::vector<BufferDesc> buffers;    // 已注册的 Buffer 列表
    std::vector<NVMeoFBufferDesc> nvmeof_buffers; // NVMe-oF 文件
    RankInfoDesc rank_info;     // 节点标识信息
    // ...
};

Segment 在传输流程中的角色：

+--------------------------------------------------+
|  1. 元数据层 (Segment)                             |
|     发布/发现 Buffer 地址、rkey、拓扑、设备信息       |
|     "我的数据在哪？你的数据在哪？"                    |
+--------------------------------------------------+
|  2. 控制层 (RPC / Handshake)                       |
|     交换 RDMA QP 信息、建立连接                     |
|     "我们怎么连上？"                                |
+--------------------------------------------------+
|  3. 数据层 (Transport: RDMA / TCP / NVLINK / ...)   |
|     实际的数据搬移                                  |
|     "开始传数据"                                   |
+--------------------------------------------------+

工作流程：

  1. 注册 — registerLocalMemory(addr, length, "cpu:0", true) 创建 local segment 并将 Buffer 发布到 etcd/HTTP/Redis
  2. 发现 — openSegment("peer_segment_name") 从元数据存储拉取远端的 SegmentDesc，获取 addr、rkey、GID/LID 等所有必要信息
  3. 传输 — 构造 TransferRequest {READ, local_buf, segment_id, offset, length} 发起 RDMA 读/写

Segment 缓存（transfer_metadata.h:220-222）：

// 本地内存缓存
std::unordered_map<uint64_t, std::shared_ptr<SegmentDesc>> segment_id_to_desc_map_;
std::unordered_map<std::string, uint64_t> segment_name_to_id_map_;

缓存通过 syncSegmentCache() 按需同步，可通过 force_update=true 强制从 metadata 存储拉取最新数据，或由 metadata 变更通知触发刷新。并非基于 TTL 的自动过期机制。

3. Slice — 传输执行的最小粒度

Slice 是 Transport 层提交到硬件的原子操作，用户不可见。一个 TransferRequest 可能被拆成多个 Slice，每个 Slice 对应一个硬件级别的 RDMA WR/RDMA RD 操作。

Slice 结构（mooncake-transfer-engine/include/transport/transport.h:104-165）：

struct Slice {
    enum SliceStatus { PENDING, POSTED, SUCCESS, TIMEOUT, FAILED };

    void *source_addr;
    size_t length;
    TransferRequest::OpCode opcode;
    SegmentID target_id;
    SliceStatus status;
    TransferTask *task;  // 所属的传输任务

    union {
        struct { uint64_t dest_addr; uint32_t lkey; uint32_t rkey; ... } rdma;
        struct {
            void *dest_addr;
            void *cuda_stream;  // cudaStream_t, 用于异步 NVLink 传输
        } local;
        struct { uint64_t dest_addr; } tcp;   // TCP
        struct { uint64_t offset; int cufile_desc; ... } nvmeof;  // NVMe-oF
        struct { void *dest_addr; } cxl;      // CXL
        // ...
    };
};

关键设计细节：

4. 三者的嵌套关系

Segment  -- 一个进程暴露给集群的数据空间（逻辑概念）
  |
  +-- Buffer[0]  8GB, addr=0x7f..., rkey=..., location="cpu:0"
  +-- Buffer[1] 16GB, addr=0x7f..., rkey=..., location="numa:1"
  +-- Buffer[2]  4GB, addr=0x7f..., rkey=..., location="cuda:0"
         |
         | 一次 TransferRequest READ(target_id=seg, offset=5GB, length=3GB)
         | 被 Transport 拆分为:
         v
  Slice[0]  -- Buffer[1], offset=0GB,  length=1MB    (一次 RDMA READ WR)
  Slice[1]  -- Buffer[1], offset=1MB, length=1MB    (下一次 RDMA READ WR)
  Slice[2]  -- Buffer[2], offset=0GB, length=1MB    (跨 Buffer 边界)
  ...
  Slice[N]  -- Buffer[2], offset=..., length=1MB
         |
         | 所有 Slice 完成 ->
         v
  TransferTask.success_slice_count == slice_count -> Batch 完成

一句话总结：

Segment = 合同（"我的数据在这，地址和密钥的目录"）
Buffer  = 仓库里的货架（真实可用的已注册连续内存）
Slice   = 一次搬运动作（实际提交到网卡/总线的原子读写操作）

代码验证

Buffer 注册流程（mooncake-transfer-engine/src/transfer_engine_impl.cpp:511-534）：

int TransferEngineImpl::registerLocalMemory(void* addr, size_t length,
                                            const std::string& location,
                                            bool remote_accessible,
                                            bool update_metadata) {
    // 1. 检查地址重叠
    if (checkOverlap(addr, length)) { ... return ERR_ADDRESS_OVERLAPPED; }
    // 2. 遍历所有已安装的 Transport，逐一注册
    for (auto transport : multi_transports_->listTransports()) {
        int ret = transport->registerLocalMemory(addr, length, location,
                                                  remote_accessible, update_metadata);
        if (ret < 0) return ret;
    }
    // 3. 记录到本地 overlap 检查集合
    std::unique_lock<std::shared_mutex> lock(mutex_);
    insertMemoryRegionLocked({addr, length, location, remote_accessible});
    return 0;
}

Buffer 追加到 Segment（mooncake-transfer-engine/src/transfer_metadata.cpp:1059-1071）：

int TransferMetadata::addLocalMemoryBuffer(const BufferDesc &buffer_desc,
                                           bool update_metadata) {
    {
        RWSpinlock::WriteGuard guard(segment_lock_);
        auto new_segment_desc = std::make_shared<SegmentDesc>();
        auto &segment_desc = segment_id_to_desc_map_[LOCAL_SEGMENT_ID];
        *new_segment_desc = *segment_desc;  // COW 拷贝
        segment_desc = new_segment_desc;
        segment_desc->buffers.push_back(buffer_desc);  // 追加新 Buffer
    }
    if (update_metadata) return updateLocalSegmentDesc();  // 发布到 etcd
    return 0;
}

关键点：采用 Copy-on-Write 模式 — 拷贝旧 SegmentDesc，在新副本上追加 Buffer，然后原子替换指针。多次 registerLocalMemory 都在同一个 Segment 的 buffers 数组里追加条目，而不创建新 Segment。

ThreadLocalSliceCache（transport.h:242-281）验证了 Slice 有 thread-local 对象池，容量为 4096。

Share this post on:
Share this post via WhatsApp Share this post on Facebook Share this post on X Share this post via Telegram Share this post on Pinterest Share this post via email
Previous Post
Mooncake TE 阅读手记-03-TENT 与经典 TE 架构演进
Next Post
Mooncake TE 阅读手记-01-Buffer Segment Slice 三层抽象