第 3 章 V1:Buffer 初始化与通信域建立
本章回答:V1 的 Buffer 在第一次 dispatch 前要建立什么状态?主路径是 Python 汇集 peer 信息,C++ 分配本地资源,再同步 IPC 与 NVSHMEM 通信域。
3.1 Python 先确定 rank 与交换句柄
Buffer.__init__ 接受 PyTorch ProcessGroup 或 MPI communicator,并将 rank、group size 与缓冲区参数交给 C++ runtime(V1 deep_ep/buffer.py:32-84)。之后每个 rank 收集本地 device ID 和 CUDA IPC handle;跨节点或 low-latency 模式还会设置 NVSHMEM IBGDA 环境参数,选出唯一 ID 的根 rank,再调用 runtime.sync(...) 完成初始化(同文件 :86-125)。
Python 先创建 C++ runtime,再交换设备号和 IPC handle:
self.num_nvl_bytes = num_nvl_bytes
self.num_rdma_bytes = num_rdma_bytes
self.low_latency_mode = low_latency_mode
self.explicitly_destroy = explicitly_destroy
self.runtime = deep_ep_cpp.Buffer(self.rank, self.group_size, num_nvl_bytes, num_rdma_bytes, low_latency_mode, explicitly_destroy)
# Synchronize device IDs
local_device_id = self.runtime.get_local_device_id()
device_ids = all_gather_object(local_device_id)
# Synchronize IPC handles
local_ipc_handle = self.runtime.get_local_ipc_handle()
ipc_handles = all_gather_object(local_ipc_handle)跨节点或低延迟分支设置 IBGDA 并交换根 unique ID,最后同步 C++ runtime:
if self.runtime.get_num_rdma_ranks() > 1 or low_latency_mode:
assert num_qps_per_rank > 0
os.environ['NVSHMEM_DISABLE_P2P'] = '0' if allow_nvlink_for_low_latency_mode else '1'
os.environ['NVSHMEM_IB_ENABLE_IBGDA'] = '1'
os.environ['NVSHMEM_IBGDA_NUM_RC_PER_PE'] = f'{num_qps_per_rank}'
if (low_latency_mode and self.rank == 0) or (not low_latency_mode and self.runtime.get_rdma_rank() == 0):
root_unique_id = self.runtime.get_local_nvshmem_unique_id()
nvshmem_unique_ids = all_gather_object(root_unique_id)
root_unique_id = nvshmem_unique_ids[0 if low_latency_mode else self.runtime.get_root_rdma_rank(True)]
self.runtime.sync(device_ids, ipc_handles, root_unique_id)
assert self.runtime.is_available()这个初始化顺序意味着 ProcessGroup 并不直接承载 GPU 数据面。它主要用于 rank 编排与对象交换;同节点的数据共享由 CUDA IPC 句柄建立,跨节点数据面则由 NVSHMEM 初始化。
3.2 C++ runtime 的资源所有权
C++ 构造函数根据 rank 计算 rdma_rank = rank / 8 与 nvl_rank = rank % 8,将一组最多 8 个本地 rank 视为 NVLink domain;它申请 local NVLink buffer、barrier 元数据、32 MiB workspace 和映射到 GPU 的 host counters(V1 csrc/deep_ep.cpp:26-80)。get_local_ipc_handle 导出本地分配的内存句柄,Python 收集后由 sync 在 peer 间打开。
rank 拆分和本地 IPC allocation 的核心逻辑如下:
// Get ranks
CUDA_CHECK(cudaGetDevice(&device_id));
rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS;
num_rdma_ranks = std::max(1, num_ranks / NUM_MAX_NVL_PEERS), num_nvl_ranks = std::min(num_ranks, NUM_MAX_NVL_PEERS);
if (num_nvl_bytes > 0) {
// Local IPC: alloc local memory and set local IPC handles
CUDA_CHECK(cudaMalloc(&buffer_ptrs[nvl_rank], num_nvl_bytes + barrier_signal_bytes + buffer_ptr_bytes + barrier_signal_ptr_bytes));
CUDA_CHECK(cudaIpcGetMemHandle(&ipc_handles[nvl_rank], buffer_ptrs[nvl_rank]));
buffer_ptrs_gpu = reinterpret_cast<void**>(static_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes + barrier_signal_bytes);
barrier_signal_ptrs[nvl_rank] = reinterpret_cast<int*>(static_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes);
barrier_signal_ptrs_gpu = reinterpret_cast<int**>(static_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes + barrier_signal_bytes + buffer_ptr_bytes);
// No need to synchronize, will do a full device sync during `sync`
CUDA_CHECK(cudaMemsetAsync(barrier_signal_ptrs[nvl_rank], 0, barrier_signal_bytes, comm_stream));
}常量定义给出了 rank 分组与 workspace 大小:
#define NUM_MAX_NVL_PEERS 8
#define NUM_MAX_RDMA_PEERS 20
#define NUM_WORKSPACE_BYTES (32 * 1024 * 1024)
#define NUM_MAX_LOCAL_EXPERTS 1024
#define NUM_BUFFER_ALIGNMENT_BYTES 128资源由 C++ 对象持有:Python 的 buffer 通过 runtime 字段持有扩展对象,destroy() 最终进入 C++ Buffer::destroy。销毁前先设备同步,并在释放 IPC 或 NVSHMEM 资源前执行相应 barrier;若设置 explicitly_destroy=True,调用方必须主动调用 destroy(),否则析构会警告潜在泄漏(csrc/deep_ep.cpp:84-90,143-183)。
析构策略和资源释放顺序如下:
Buffer::~Buffer() noexcept(false) {
if (not explicitly_destroy) {
destroy();
} else if (not destroyed) {
printf("WARNING: destroy() was not called before DeepEP buffer destruction, which can leak resources.\n");
fflush(stdout);
}
}
void Buffer::destroy() {
EP_HOST_ASSERT(not destroyed);
CUDA_CHECK(cudaDeviceSynchronize());
if (num_nvl_bytes > 0) {
intranode::barrier(barrier_signal_ptrs_gpu, nvl_rank, num_nvl_ranks, comm_stream);
CUDA_CHECK(cudaDeviceSynchronize());
if (is_available()) {
for (int i = 0; i < num_nvl_ranks; ++ i) if (i != nvl_rank)
CUDA_CHECK(cudaIpcCloseMemHandle(buffer_ptrs[i]));
}
CUDA_CHECK(cudaFree(buffer_ptrs[nvl_rank]));
}
#ifndef DISABLE_NVSHMEM
if (is_available() and num_rdma_bytes > 0) {
CUDA_CHECK(cudaDeviceSynchronize());
internode::barrier();
internode::free(rdma_buffer_ptr);
internode::finalize();
}
#endif
CUDA_CHECK(cudaFree(workspace));
CUDA_CHECK(cudaFreeHost(const_cast<int*>(moe_recv_counter)));
CUDA_CHECK(cudaFreeHost(const_cast<int*>(moe_recv_expert_counter)));
destroyed = true;
available = false;
}sequenceDiagram
participant P as Python Buffer
participant C as C++ Buffer
participant G as 同节点 rank
participant N as NVSHMEM ranks
P->>C: 创建本地 Buffer
C->>C: 分配 NVLink buffer / workspace / counters
P->>G: all_gather device ID 与 CUDA IPC handle
P->>N: 选 root unique ID 并 all_gather
P->>C: sync(peer 信息, root ID)
C->>G: 打开 peer IPC 内存
C->>N: 初始化 NVSHMEM 与 RDMA storage
Note over P,N: dispatch / combine 可用3.3 初始化约束来自拓扑假设
V1 C++ 构造中写死 NUM_MAX_NVL_PEERS == 8,并要求超过本地 NVLink peer 数时 rank 数能按该值整除。high-throughput internode 路径假定节点内经 NVLink、相同 GPU index 的节点间经 RDMA;低延迟模式另有 IBGDA QP 数和容量约束(deep_ep/buffer.py:47-52)。这说明部署拓扑是算法参数的一部分,不是运行环境的无关细节。
3.4 失败与释放边界
Python 初始化中途失败时,C++ 对象可能已经分配 CUDA 或 NVSHMEM 资源;V1 文档特别提醒析构清理可能在 Python 异常处理中挂住,因此提供显式销毁选项。销毁本身是集体操作:一边的 rank 离开、另一边仍继续使用缓冲区,会使 barrier 无法完成。源码直接显示了同步和资源释放顺序;它没有提供可由单 rank 安全恢复的故障协议。
本章收束
- V1 先通过 host collectives 交换进程信息,再建立 CUDA IPC / NVSHMEM 数据面。
- C++
Buffer持有 workspace、通信缓冲区与计数器,并负责集体清理。 - 节点内 8-rank 分组、NVLink 可见性、RDMA 可达性共同限定可用拓扑。
- 下一章从这些已建好的资源出发,追踪一次 V1 dispatch 到 combine。