团团虾导读:Mooncake Store 是 TE 的上层消费者,一套完整的分布式 KV Cache 对象存储引擎。这篇从 Store 如何使用 TE 的完整调用链出发——InitTransferEngine → 两类内存注册(Local Buffer vs Global Segment)→ Put/Get 流程 → TransferSubmitter 的三种传输策略 → SSD 分层存储,最后分析了远端 SSD 读时 TE 如何参与。理解 Store 也等于理解 TE 在真实业务中的用法。
[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 / ...) |
| 实际的数据搬移 |
| "开始传数据" |
+--------------------------------------------------+
工作流程:
- 注册 —
registerLocalMemory(addr, length, "cpu:0", true)创建 local segment 并将 Buffer 发布到 etcd/HTTP/Redis - 发现 —
openSegment("peer_segment_name")从元数据存储拉取远端的SegmentDesc,获取 addr、rkey、GID/LID 等所有必要信息 - 传输 — 构造
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
// ...
};
};
关键设计细节:
- Union 按 Transport 分派:不同传输协议有各自专用的目的地址/密钥字段,避免为不需要的字段浪费内存
- 原子完成回调:
markSuccess()/markFailed()使用__atomic_fetch_add更新 Task 计数 - Thread-Local 对象池:
ThreadLocalSliceCache(transport.h:242)缓存最多 4096 个 Slice,避免频繁 new/delete - 两种完成检测路径:
- 事件驱动模式(
USE_EVENT_DRIVEN_COMPLETION编译选项):Slice 完成时通过原子递增completed_slice_count检测完成,条件为prev_completed + 1 == task->slice_count - 轮询模式(传统模式):外部通过
getTransferStatus轮询,判断success_slice_count + failed_slice_count == slice_count
- 事件驱动模式(
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。