Mooncake NDS Integration

导言

手里已经有一套面向 NPU 的 NDS send/receive 接口，下一步是把它接入 Mooncake。乍看之下，这只是把 GDS 的 cuFile* 调用替换为 NDS API；但沿源码真正走一遍后，会遇到两个容易误判的事实：Mooncake 同时保留了新旧两代 GDS 路径，而 Mooncake Store 的磁盘副本读写目前又绕开了它们。

因此，接入点不应先选旧 NVMeoFTransport，也不能只新增一个 transport 就宣称 Store 已经获得 NPU 直读 SSD。更稳妥的路径是：先把 NDS 实现为 TENT 的 NdsTransport，复用其选择、批次、状态与回退机制；再单独改造 Store 的文件副本路径。

本文以 Mooncake main 的提交 468fbf63 为代码截面。文中的 NDS 指内部提供的 NPU Direct Storage 接口，不把它假定为某个公开标准或 NVIDIA 产品。

接口假设

目前没有 NDS SDK 的真实头文件与签名。下文暂时假设 send 表示 NPU HBM → 文件/SSD，receive 表示 文件/SSD → NPU HBM，并假设 SDK 至少能够表达设备地址、文件或句柄、文件偏移、长度和异步完成对象。正式编码前必须用真实接口校正这个映射。

先找到真正的 GDS 路径

Mooncake 中能搜到两套与 GDS 相关的实现。它们都调用 cuFile，却不处于同一代架构，也不具有同样的接入价值。

旧路径：NVMeoFTransport

经典 Transfer Engine 中的 NVMeoFTransport 通过 USE_NVMEOF 构建，协议名是 nvmeof。它管理 cuFile buffer、文件 handle、Batch ID 和异步事件，测试也能直接调用 allocateBatchID → submitTransfer → getTransferStatus 完成读写。

问题在于，它没有真正接入通用 MultiTransport 批次。当前 submitTransferTask 会直接返回 NotImplemented；对应测试绕过通用调度层，直接持有 Transport* 调用专用接口，参见 nvmeof_transport_test.cpp。

这条路径仍可用来理解 cuFile 的注册、文件 handle 与事件映射，但不适合作为新 NDS 接入的主模板。照着它复制，容易得到一个只能被专用测试调用、无法被当前运行时自动选择的孤立后端。

新路径：TENT GdsTransport

TENT 是 Mooncake 新一代 Transfer Engine。构建时设置 -DUSE_TENT=ON，既可以使用原生 API，也可以通过 MC_USE_TENT=1 让经典 TE 接口委托给 TENT，见 TENT C++ API。

TENT 的 GDS 路径更完整：

1. 构建发现：启用 CUDA 且找到 cufile 库和头文件时定义 USE_GDS。
2. 运行时装载：配置 transports/gds/enable 后，TransportLoader 创建 GdsTransport。
3. 统一文件段：file://path 被解析为 FileSegmentDesc，transport 再按 Segment ID 取回路径。
4. 策略选择：默认文件策略按 GDS → IOURING 排序；GDS 不可用时可以回退到 CPU staging。
5. 批次生命周期：allocateSubBatch 复用昂贵的 cuFile batch handle，submitTransferTasks 把大请求切成最多 16 MiB 的 slice，getTransferStatus 轮询并聚合每个 slice 的完成状态。
6. 失败隔离：取消是 best effort；只要底层仍可能引用参数或用户 buffer，失败批次就不能立即复用，而是进入 quarantine，等所有 I/O 真正终态后再回收。
7. 内存注册：CUDA buffer 通过 cuFileBufRegister/cuFileBufDeregister 加入或移出 GDS 能力集合。

核心提交与状态处理集中在 gds_transport.cpp。这里最值得复用的不是 cuFileBatchIOSubmit 这一个调用，而是它周围那一整套 注册、切片、提交、轮询、取消、终态确认和资源回收契约。

NDS 应该接在哪里

下图把当前 GDS、目标 NDS 和 Store 的实际磁盘路径放在同一张图里。上半部分是 NDS 第一阶段应该进入的 TENT 文件数据面；下半部分说明为什么新增 NdsTransport 后，Store 仍不会自动走 NDS。

图：蓝色是当前 GDS，紫色是目标 NDS 与第二阶段接线，绿色是 Store 目前绕过 TENT 的文件读写路径。

由此可以把目标拆成两个互不冒充的里程碑：

1. TENT 可用：普通 Transfer Engine 请求能够在 NPU buffer 与 file:// Segment 之间走 NDS，并具备回退、状态与错误处理。
2. Store 可用：LOCAL_DISK 副本的读写显式改走 TENT/NDS，且对象的分配、可见性、校验与故障恢复语义仍成立。

第一项是 transport 接入，第二项是存储后端重构。它们应分两个 PR 或至少两个可独立回滚的提交完成。

第一步：钉死 NDS 契约

不要一开始就在 NdsTransport 里散落 SDK 调用。先加一层很薄的 NdsBackend，把 Mooncake 需要的语义与内部 SDK 的具体名字隔开。

enum class NdsDirection { kReadFromFile, kWriteToFile };

struct NdsIo {
    NdsDirection direction;
    void* npu_ptr;
    uint64_t file_offset;
    size_t length;
};

class NdsBackend {
public:
    Status initialize(const NdsConfig& config);
    Status registerBuffer(void* addr, size_t length);
    Status unregisterBuffer(void* addr);
    Status openFile(const FileBufferDesc& file, NdsFileHandle& handle);
    Status submit(const NdsFileHandle& file,
                  std::span<const NdsIo> ios,
                  NdsRequestHandle& request);
    Status poll(NdsRequestHandle request, NdsCompletion& completion);
    Status cancel(NdsRequestHandle request);
};

在适配层中再完成方向映射：

Request::WRITE  → NDS send    → NPU HBM 写入文件
Request::READ   → NDS receive → 文件读入 NPU HBM

编码前至少确认以下问题，否则 transport 的生命周期无法定型：

• send/receive 接受路径、文件描述符，还是 SDK 自己的文件 handle；
• NPU buffer 是否必须预注册，注册粒度和注销约束是什么；
• 文件 offset、设备地址和长度分别要求怎样的对齐；
• 提交返回的是 request、event、队列槽位，还是同步状态；
• 一次请求是否支持多个 slice，最大 batch depth 是多少；
• cancel 是否存在，已入队请求能否保证不再访问用户 buffer；
• SDK 是否线程安全，完成队列由谁推进，进程退出时如何 drain。

适配层的价值

NdsBackend 不负责 Mooncake 策略，只负责把 SDK 变成稳定、可模拟的最小接口。这样单元测试可以用 fake backend 注入 partial completion、超时和失败，而不必依赖真实 NPU 与 SSD。

第二步：让运行时认识 NPU

这是当前代码里最容易漏掉、也最先应该修的地方。

TENT 已经在 policy schema 中接受 local_memory: "npu"，但 AscendPlatform::getMemoryType 返回的仍是 MTYPE_CUDA。TransportSelector 又把 MTYPE_CUDA 转成字符串 cuda，因此写出 local_memory: "npu" 并不会真的命中 Ascend NPU buffer。

应先完成三项修改：

1. 在 MemoryType 中新增 MTYPE_NPU，并让 Ascend platform 返回它。
2. 在 matchesMemoryPattern 和设备内存判断中加入 MTYPE_NPU。
3. 把文件能力从含混的 gpu_to_file 拆出 npu_to_file，避免 GDS 和 NDS 因共享一个布尔位而被错误互选。

最小改动是继续把 NPU 当作 gpu_to_file，再靠 policy 把 GDS/NDS 分开；但这会把正确性压在配置顺序上。长期可维护的实现应在类型和 capability 两层都区分 CUDA 与 NPU。

同时要修复回退 transport。当前 IOUringTransport 只有看到 MTYPE_CUDA 才分配 CPU staging buffer。加入 MTYPE_NPU 后，应把这类判断统一成 isDeviceMemory(type)，否则 NDS 不可用时会把 NPU 地址误当成普通 CPU 指针交给 io_uring。

第三步：加入 NDS transport

扩展公开枚举

在固定提交 468fbf63 中，TransportType 依次到 MPCOMM，随后是 sentinel。为了不改变已有 transport 的数值，应把 NDS 追加在 kNumTransportTypes 之前，而不是插进 GDS 附近。

当前截面下需要同步更新：

• tent/include/tent/common/types.h：TransportType::NDS、transportTypeName、parseTransportType；
• tent/include/tent/transfer_engine.h：C API 宏 TRANSPORT_NDS，在该提交上对应数值 13；
• tent/src/python/pybind.cpp：暴露 TransportType.NDS，补充必要的数值一致性断言；
• 所有以 kSupportedTransportTypes 为长度的数组会随 sentinel 自动扩容，但仍需跑覆盖测试。

这里的 13 只对固定提交成立。如果上游同时新增 transport，应以合并后的枚举顺序重新生成并核对，而不是永久硬编码这篇文章里的数字。

实现生命周期

建议新增：

tent/include/tent/transport/nds/nds_transport.h
tent/src/transport/nds/nds_transport.cpp
tent/src/transport/nds/CMakeLists.txt

NdsTransport 至少实现以下接口：

TENT 接口	NDS 责任
install/uninstall	初始化 SDK、读取 queue depth、停止接收新任务并 drain 完成队列
addMemoryBuffer/removeMemoryBuffer	注册或注销 NPU HBM，并把 NDS 写入 buffer 的 transports
allocateSubBatch/freeSubBatch	分配请求槽位；保存 SDK request、slice 范围和稳定的参数存储
submitTransferTasks	解析 File Segment，将 READ/WRITE 映射为 receive/send，按 SDK 上限和对齐切片
getTransferStatus	轮询 completion，聚合字节数，保留第一个终态错误
cancelTransferTask	best effort 取消；未确认终态前不得释放 SDK 仍可能访问的 buffer 或参数
capabilities	声明 npu_to_file=true，不要伪装成 CUDA GDS

GDS 固定使用 16 MiB slice 是 cuFile 当前实现的选择，不要原样复制成 NDS 常数。NDS 应从 SDK 限制、NPU 页粒度、NVMe 队列深度和真实 KV block 分布确定切片大小，并允许通过配置调节。

错误路径也应复制语义，而不是复制代码：某个 slice 失败后可以请求取消兄弟 slice，但 task 只有在所有底层 I/O 都不再触碰用户内存时才能公布终态并释放 batch。否则上层可能在 DMA 尚未停止时复用 KV buffer。

接入构建与装载

构建层建议增加：

-DUSE_NDS=ON
-DNDS_ROOT=/path/to/nds

CMake 应在 NDS_ROOT 下寻找头文件和库，找到后定义 USE_NDS、链接 SDK，并把 tent_xport_nds 加入 TENT transport 集合。运行时在 TransportLoader 中按以下配置创建实例：

{
  "transports": {
    "nds": {
      "enable": true,
      "io_batch_depth": 32
    }
  }
}

io_batch_depth 只是配置形态示例，默认值必须根据 NDS SDK 的真实队列约束确定，不能因为 GDS 当前默认 32 就直接照搬。

修正选择与回退

当前默认文件策略没有 memory filter，候选顺序固定为 GDS → IOURING。接入 NDS 后，不应简单改成 NDS → GDS → IOURING，因为那会让正确性依赖 capability 检查是否足够严密。

更清晰的策略是按本地内存类型分三条规则：

{
  "policy": [
    {
      "name": "npu_file",
      "segment_type": "file",
      "local_memory": "npu",
      "transports": ["nds", "io_uring"]
    },
    {
      "name": "cuda_file",
      "segment_type": "file",
      "local_memory": "cuda",
      "transports": ["gds", "io_uring"]
    },
    {
      "name": "cpu_file",
      "segment_type": "file",
      "local_memory": "cpu",
      "transports": ["io_uring"]
    }
  ]
}

预期结果必须是：NPU 优先 NDS、CUDA 优先 GDS、CPU 走 io_uring；NDS 或 GDS 缺失时，设备内存先经过受控的 CPU staging 再落盘。

复用文件段，必要时再扩元数据

如果 NDS 接受 path 或普通 fd，可以直接复用 file://path 与现有 FileBufferDesc{path,length,offset}，不用创造 nds://。这能让 GDS、NDS 和 io_uring 共享同一个文件 Segment，只由 selector 决定执行后端。

如果 NDS 还需要 namespace、设备队列、存储 endpoint 或专用文件句柄，就应为 FileBufferDesc 增加类似 Memory Segment 的 transport_attrs，按 transport 保存序列化属性。不要把 NDS 私有字段直接塞进通用结构顶层，否则每加入一个存储后端都会扩大公共 schema。

第四步：单独接通 Store

新增 NdsTransport 后，Transfer Engine 的文件请求已经可用，但 Mooncake Store 的 LOCAL_DISK 副本仍不会自动经过它。

当前读路径是：

submitFileReadOperation
    → FilereadWorkerPool
    → StorageBackend::LoadObject
    → StorageFile::vector_read / preadv

源码证据分别见 submitFileReadOperation 与 StorageBackend::LoadObject。这条链路直接在 Store worker 中执行文件 I/O，没有创建 file:// Segment，也没有提交 TENT batch。

先改读路径

读比写更适合作为 Store 集成的第一刀：文件、对象 offset 和目标 NPU slice 都已经确定，不需要先改变对象可见性。

可以增加一个 TransferEngineFileOperationState：

1. 根据 disk replica 的文件路径创建或缓存 file://path Segment；
2. 为每个目标 NPU slice 构造 Request::READ；
3. target_offset 指向对象值区在文件中的真实偏移；
4. 提交 TENT batch，并把轮询结果转换为现有 TransferFuture；
5. NDS 不可用时按 policy 回退到 io_uring staging；
6. 保留现有 file worker 作为可配置回退，直到 NDS 路径通过故障注入。

对于 OffsetAllocatorStorageBackend，不能直接从 record 起点把整个记录读进 NPU。它的格式是 header + key + padding + value，只有 value 区按 4 KiB 对齐。安全做法是继续在 CPU 侧读取并校验 header/key，再用 NDS receive 只搬 value 区。

写路径需要 reserve/commit

写路径不能只把 vector_write 换成 send。当前 BatchOffload 在同一个函数中完成空间分配、header/key/CRC 构造、磁盘写入和内存索引发布；它没有暴露给外部 DMA writer 使用的 reserve/commit 边界。

建议把它拆成三段：

prepare_record
    → 分配 extent，确定 record/value offset、seq、header、key 与 padding

NDS send
    → 只把 NPU 中的 value 写到 4 KiB 对齐的 value offset

commit_record
    → 确认所有 I/O 完成后更新 shard map、FIFO index、计数并通知 Master

任何失败都必须释放尚未发布的 allocation，且不能让 Master 看见一个未完成的 LOCAL_DISK 副本。当前 record layout 已经明确为未来 GDS/DMA writer 保留了 4 KiB 对齐，并用 seq 防止恢复时接受 checkpoint 之后的 torn write，见 storage_backend.h。

CRC 是第二个边界。当前默认 CRC-32C 覆盖 header prefix、key 和 value；如果 value 从未经过 CPU，就无法沿用现有 CPU 计算方式。可选方案只有三类：

• SDK 或 NPU 侧同时计算 CRC，再把结果写回 header；
• 为 direct write 清除 kFlagHasCrc，接受恢复时少一层 torn/stale 检测；
• 额外把 value 拉回 CPU 计算 CRC，但这会抵消 direct storage 的主要意义。

在没有设备侧校验能力时，第二种与现有格式兼容，但必须把可靠性降级写进配置、指标与故障测试，不能静默发生。

修改顺序

为了让每一步都能独立验证，建议按以下顺序提交：

1. SDK 契约测试：实现 fake NdsBackend，固定 send/receive、offset、alignment、completion 与 cancel 语义。
2. NPU 类型修复：增加 MTYPE_NPU、npu_to_file 和通用 isDeviceMemory；先让 selector 测试全部通过。
3. 公共枚举与绑定：追加 TransportType::NDS，更新 C API、Python、名称解析和数组边界测试。
4. TENT transport：实现注册、文件 context、sub-batch、切片、提交、轮询、取消与 quarantine。
5. 构建和配置：接入 USE_NDS/NDS_ROOT、loader、默认配置与自定义 policy。
6. TENT 端到端：用临时文件验证 NPU→文件、文件→NPU、非对齐尾部、partial failure、取消与回退。
7. Store 只读接入：新增 TENT file operation，先覆盖普通 file replica，再覆盖 OffsetAllocator 的 value offset。
8. Store 写入重构：引入 prepare/send/commit，处理 CRC、失败回滚、淘汰与恢复。
9. 性能验收：在真实 KV block 分布上比较 NDS、io_uring staging 与现有 Store worker，而不是只测单个大块峰值带宽。

对应的测试矩阵至少应覆盖：

层次	必测行为
Selector	NPU→NDS、CUDA→GDS、CPU→io_uring、NDS 缺失时安全回退
Buffer	注册/注销、重复注册、越界、未注册内存、进程退出清理
File	path/fd 生命周期、offset 与 length 对齐、文件截断、权限错误
Batch	多 slice、乱序完成、部分失败、取消后 drain、状态字节数单调
Store read	header/key 校验、value 直读 NPU、淘汰期间 extent pin
Store write	完成前不可见、失败释放 extent、重启恢复、CRC 开关语义
性能	小块时延、聚合吞吐、CPU 占用、队列深度、注册成本和 P99

最终判断

现在可以确认的是：Mooncake 当前最完整的 GPU 直存储接入点是 TENT GdsTransport，而不是经典 NVMeoFTransport；它提供的真正模板是 transport 生命周期与运行时选择。也可以确认，Store 的文件副本当前由自己的 worker 和 StorageBackend 读写，因此仅新增 NDS transport 不足以打通 NPU KV Cache 到 SSD 的端到端路径。

证据还不能支持的是 NDS 的真实方向、对齐、取消、注册与完成语义。没有 SDK 头文件和最小示例，就不应把伪接口直接写成生产代码。

所以最实际的下一步不是一次性改 Store，而是先拿真实 NDS API 完成 NdsBackend 契约表和 fake test。只要这层钉死，TENT 接入就有清晰模板；待普通 file:// 读写稳定后，再把 Store 的 read 与 write 分开推进，系统边界会清楚得多。

参考文献

• Mooncake TENT C++ API
• Mooncake TENT Transport Selector
• Mooncake TENT GDS Transport
• Mooncake legacy NVMeoF Transport
• Mooncake Store transfer task
• Mooncake Store storage backend
• NVIDIA GPUDirect Storage API Reference
• NVIDIA GPUDirect Storage Troubleshooting Guide