【cuda】 deepep

cpp 复制代码

void Buffer::sync(const std::vector<int>& device_ids,
                  const std::vector<std::optional<pybind11::bytearray>>& all_gathered_handles,
                  const std::optional<pybind11::bytearray>& root_unique_id_opt) {
    EP_HOST_ASSERT(not is_available());

    // Sync IPC handles
    if (num_nvl_bytes > 0) {
        EP_HOST_ASSERT(num_ranks == device_ids.size());
        EP_HOST_ASSERT(device_ids.size() == all_gathered_handles.size());
        for (int i = 0, offset = rdma_rank * num_nvl_ranks; i < num_nvl_ranks; ++i) {
            EP_HOST_ASSERT(all_gathered_handles[offset + i].has_value());
            auto handle_str = std::string(all_gathered_handles[offset + i].value());
            EP_HOST_ASSERT(handle_str.size() == shared_memory::HANDLE_SIZE);
            if (offset + i != rank) {
                std::memcpy(&ipc_handles[i], handle_str.c_str(), shared_memory::HANDLE_SIZE);
                shared_memory_allocator.open_mem_handle(&buffer_ptrs[i], &ipc_handles[i]);
                barrier_signal_ptrs[i] = reinterpret_cast<int*>(static_cast<uint8_t*>(buffer_ptrs[i]) + num_nvl_bytes);
            } else {
                EP_HOST_ASSERT(std::memcmp(&ipc_handles[i], handle_str.c_str(), shared_memory::HANDLE_SIZE) == 0);
            }
        }

        // Copy all buffer and barrier signal pointers to GPU
        CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs, sizeof(void*) * NUM_MAX_NVL_PEERS, cudaMemcpyHostToDevice));
        CUDA_CHECK(cudaMemcpy(barrier_signal_ptrs_gpu, barrier_signal_ptrs, sizeof(int*) * NUM_MAX_NVL_PEERS, cudaMemcpyHostToDevice));
        CUDA_CHECK(cudaDeviceSynchronize());
    }

cudaMemcpy(...): 这两行代码将 CPU 上的 buffer_ptrs 和 barrier_signal_ptrs 数组，完整地拷贝到 GPU 上的 buffer_ptrs_gpu 和拷贝到哪个GPU上呢，到底是哪个GPU的buffer_ptrs_gpu呢

这个问题提得非常好！您触及了理解这段 C++ 代码与 torch.distributed 交互方式的一个关键点。

拷贝到哪个 GPU 上呢，到底是哪个 GPU 的 buffer_ptrs_gpu 呢？

答案是：拷贝到当前正在执行这段 C++ 代码的那个 GPU 进程所绑定的 GPU 上。

让我们来把整个流程串起来，您就会豁然开朗。

SPMD 模式下的执行流程

我们回到 SPMD (Single Program, Multiple Data) 的执行模式。当您运行 torch.multiprocessing.spawn(test_loop, ...) 时，系统会创建多个独立的 Python 进程。假设我们有 8 个 GPU（num_processes=8）。

进程与 GPU 绑定:
- 进程 0 会被分配并绑定到 GPU 0 (cuda:0)。
- 进程 1 会被分配并绑定到 GPU 1 (cuda:1)。
- ...
- 进程 7 会被分配并绑定到 GPU 7 (cuda:7)。
  这个绑定是通过 init_dist 函数内部的 torch.cuda.set_device(local_rank) 完成的。
独立的 Buffer 实例:
- 进程 0 在 GPU 0 上执行 test_loop，它会创建一个 Buffer 对象，我们称之为 buffer_instance_0。
- 进程 1 在 GPU 1 上执行 test_loop，它会创建另一个完全独立的 Buffer 对象，buffer_instance_1。
- ...以此类推，每个进程都有一个自己的 Buffer 实例。
独立的 C++ Buffer 对象:
- buffer_instance_0 内部的 self.runtime 指向一个 C++ deep_ep::Buffer 对象，我们称之为 cpp_buffer_0。这个 C++ 对象的所有 CUDA API 调用（如 cudaMalloc, cudaMemcpy）都会作用在当前设备 ，也就是 GPU 0 上。
- 同样，buffer_instance_1 内部的 cpp_buffer_1 的所有 CUDA API 调用都会作用在 GPU 1 上。

`cudaMemcpy` 的具体目标

现在我们来看这行代码，并分别代入两个进程的视角：

cpp 复制代码

CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs, sizeof(void*) * NUM_MAX_NVL_PEERS, cudaMemcpyHostToDevice));

在进程 0 (GPU 0) 的视角下：

buffer_ptrs_gpu : 这是 cpp_buffer_0 对象的一个成员变量。在 cpp_buffer_0 的构造函数中，通过 shared_memory_allocator.malloc 分配的内存块位于 GPU 0 的显存上。因此，cpp_buffer_0 的 buffer_ptrs_gpu 这个指针指向的是 GPU 0 显存中的一个地址。
buffer_ptrs : 这是 cpp_buffer_0 对象在 CPU 上的一个 void* 数组。经过 sync 函数的循环，它已经填满了指向 GPU 0、GPU 1、...、GPU 7 nvl_buffer 的指针。
cudaMemcpy 操作 :
- 源 : cpp_buffer_0 的 buffer_ptrs 数组（在 CPU 内存中）。
- 目标 : cpp_buffer_0 的 buffer_ptrs_gpu 指针所指向的地址（在 GPU 0 的显存中）。
- 结论 : 这行代码将 cpp_buffer_0 整理好的地址表，从 CPU 拷贝到了 GPU 0 的显存上。

在进程 1 (GPU 1) 的视角下：

buffer_ptrs_gpu : 这是 cpp_buffer_1 对象的成员变量，它指向 GPU 1 显存中的一个地址。
buffer_ptrs : 这是 cpp_buffer_1 对象在 CPU 上的 void* 数组。它同样填满了指向 GPU 0 到 GPU 7 nvl_buffer 的指针（注意，这些指针的值和 cpp_buffer_0 中的可能不同，因为每个 GPU 的虚拟地址空间是独立的）。
cudaMemcpy 操作 :
- 源 : cpp_buffer_1 的 buffer_ptrs 数组（在 CPU 内存中）。
- 目标 : cpp_buffer_1 的 buffer_ptrs_gpu 指针所指向的地址（在 GPU 1 的显存中）。
- 结论 : 这行代码将 cpp_buffer_1 整理好的地址表，从 CPU 拷贝到了 GPU 1 的显存上。

总结与图示

每个 rank 都在执行同样的代码，但操作的是自己绑定的 GPU 和属于自己的 C++ 对象实例。

复制代码

+------------------------------------------------+
|                   CPU Host Memory                |
|                                                |
|  +--------------------+  +--------------------+  |
|  |   cpp_buffer_0     |  |   cpp_buffer_1     |  | ...
|  | (属于进程 0)       |  | (属于进程 1)       |  |
|  |                    |  |                    |  |
|  | .buffer_ptrs[]     |  | .buffer_ptrs[]     |  |
|  | [ptr_to_gpu0,      |  | [ptr_to_gpu0,      |  |
|  |  ptr_to_gpu1, ...] |  |  ptr_to_gpu1, ...] |  |
|  +--------+-----------+  +--------+-----------+  |
|           |                      |               |
+-----------|----------------------|---------------+
            | cudaMemcpy           | cudaMemcpy
            | (在进程0中调用)        | (在进程1中调用)
            V                      V
+--------------------------+  +--------------------------+
|       GPU 0显存          |  |       GPU 1显存          |
|                          |  |                          |
|  +-------------------+   |  |  +-------------------+   |
|  | buffer_ptrs_gpu_0 |   |  |  | buffer_ptrs_gpu_1 |   |
|  | (cpp_buffer_0的)  |   |  |  | (cpp_buffer_1的)  |   |
|  +-------------------+   |  |  +-------------------+   |
|                          |  |                          |
+--------------------------+  +--------------------------+

所以，cudaMemcpy(...) 将地址表拷贝到了当前进程所绑定的那块 GPU 上，存放在属于当前进程的那个 Buffer 实例 的 buffer_ptrs_gpu 成员变量所指向的显存区域。

这个机制保证了每个 GPU 的 Kernel 都能在自己的显存中找到一张完整的、可以访问所有节点内伙伴的"地址地图"。

好的，我们来详细解释 for 循环中的 offset = rdma_rank * num_nvl_ranks 这个计算，以及它与 num_ranks == device_ids.size() 这个断言的关系。

这部分代码的核心目标 是：从一个全局的、扁平的句柄列表 all_gathered_handles 中，准确地定位并提取出当前物理节点内的所有句柄。

关键变量解释

让我们先明确几个变量的含义（以64卡，8节点的配置为例）：

rank: 当前 GPU 的全局 Rank ID (0-63)。
num_ranks: 全局总 Rank 数量 (64)。
rdma_rank: 当前 GPU 所在的节点 ID (0-7)。计算方式是 rank / 8。
nvl_rank: 当前 GPU 在其节点内的本地 ID (0-7)。计算方式是 rank % 8。
num_nvl_ranks: 一个节点内的 GPU 数量 (8)。
device_ids: 一个从 Python 端通过 all_gather_object 收集来的、包含所有 64 个 GPU 设备 ID 的列表。它的顺序与全局 Rank ID 一致。
all_gathered_handles: 同样是通过 all_gather_object 收集来的、包含所有 64 个 GPU 的 IPC 句柄的列表。它的顺序也与全局 Rank ID 一致。

断言：`num_ranks == device_ids.size()`

这个断言非常直观：

含义 : 确保我们从 Python 端收到的 device_ids 列表的大小，正好等于我们在 C++ 端知道的全局总 Rank 数量。
目的 : 这是一个健全性检查 (Sanity Check) 。如果这个断言失败，意味着 Python 前端和 C++ 后端关于"总共有多少个 GPU 参与通信"的认知不一致，这是一个严重的配置错误，程序必须立即终止。all_gathered_handles.size() 同样也应该等于 num_ranks。

循环与 `offset` 的计算：`offset = rdma_rank * num_nvl_ranks`

现在我们来关注 for 循环和 offset 的计算，这是整个逻辑的关键。

目标 : for 循环的意图是遍历当前节点内部的所有伙伴 GPU 。循环变量 i 代表节点内的本地 ID （从 0 到 7）。但是，all_gathered_handles 这个列表是全局的、按全局 Rank 排序 的。我们不能直接用 i 去索引它。

问题 : 如何从全局列表 all_gathered_handles[0...63] 中，找到属于当前节点的那一段？

解决方案 : 这就是 offset 的作用。offset 计算的是当前节点在全局列表中的起始位置。

offset = rdma_rank * num_nvl_ranks

让我们代入具体数值来理解：

假设当前 GPU 是 Rank 18:
- rank = 18
- rdma_rank = 18 / 8 = 2 (它在第 2 个节点上，节点ID从0开始)
- nvl_rank = 18 % 8 = 2
- num_nvl_ranks = 8
- offset = 2 * 8 = 16
假设当前 GPU 是 Rank 5:
- rank = 5
- rdma_rank = 5 / 8 = 0 (它在第 0 个节点上)
- nvl_rank = 5 % 8 = 5
- num_nvl_ranks = 8
- offset = 0 * 8 = 0

offset 的值告诉我们：

如果我在节点 2 上，那么属于我们这个节点的句柄，在全局 all_gathered_handles 列表中的起始索引是 16。
如果我在节点 0 上，那么起始索引就是 0。

`for` 循环的完整逻辑

现在，我们将 offset 和 for 循环结合起来看：

cpp 复制代码

// 假设当前 rank=18, rdma_rank=2, nvl_rank=2, offset=16

for (int i = 0; i < num_nvl_ranks; ++i) { // i 从 0 遍历到 7
    
    // global_rank_index = offset + i
    // ...
    auto handle_str = std::string(all_gathered_handles[offset + i].value());
    // ...
}

当 i = 0 (处理节点内第0张卡):
- 索引是 offset + i = 16 + 0 = 16。
- 代码会访问 all_gathered_handles[16]，这正好是全局 Rank 16 的句柄（即节点 2 上的第 0 张卡）。
当 i = 1 (处理节点内第1张卡):
- 索引是 offset + i = 16 + 1 = 17。
- 代码会访问 all_gathered_handles[17]，即全局 Rank 17 的句柄。
当 i = 2 (处理节点内第2张卡，也就是自己):
- 索引是 offset + i = 16 + 2 = 18。
- 代码会访问 all_gathered_handles[18]，即全局 Rank 18 (自己) 的句柄。
...
当 i = 7 (处理节点内第7张卡):
- 索引是 offset + i = 16 + 7 = 23。
- 代码会访问 all_gathered_handles[23]，即全局 Rank 23 的句柄。

通过 offset + i 这种方式，循环成功地将节点内的本地索引 i 转换为了全局句柄列表中的正确索引，从而遍历了从 Rank 16 到 Rank 23 的所有句柄------这正是节点 2 的所有成员。

总结

num_ranks == device_ids.size() 是一个前置的配置校验，确保 C++ 和 Python 对系统规模的认知是一致的。
offset = rdma_rank * num_nvl_ranks 是一个地址转换公式。它扮演了"基地址"的角色。
for 循环通过 offset + i 的方式，实现了从一个全局扁平的地址空间 （all_gathered_handles 数组）到一个两级层次化地址空间 （节点ID + 节点内ID）的正确索引，从而让每个 GPU 都能准确地找到并处理自己所在物理节点内所有伙伴 GPU 的信息。
复制代码
```
        shared_memory_allocator.open_mem_handle(&buffer_ptrs[i], &ipc_handles[i]);
        barrier_signal_ptrs[i] = reinterpret_cast<int*>(static_cast<uint8_t*>(buffer_ptrs[i]) + num_nvl_bytes); 这两句什么意思？
```

好的，我们来详细、深入地解释这两行代码，它们是建立节点内（Intra-node）P2P 通信和同步机制的关键步骤。

我们以 GPU 0 (nvl_rank=0) 的视角，分析它处理来自 GPU 1 (nvl_rank=1) 的句柄时，这两行代码具体做了什么。此时，循环变量 i 的值为 1。

第一行: `shared_memory_allocator.open_mem_handle(&buffer_ptrs[i], &ipc_handles[i]);`

中文直译 : "共享内存分配器，请使用 ipc_handles[1] 这个句柄，打开对应的内存，并把得到的地址存入 buffer_ptrs[1]。"

深入解释:

&ipc_handles[i] (即 &ipc_handles[1]): 这是一个输入参数。它指向我们之前从 Python 端收到的、属于 GPU 1 的 IPC 句柄。可以把它想象成 GPU 1 给 GPU 0 的一张"门禁卡"或者一个"地址簿条目"。
&buffer_ptrs[i] (即 &buffer_ptrs[1]) : 这是一个输出参数。我们将要把打开远程内存后得到的本地虚拟地址存放在这里。
open_mem_handle 函数内部:
- 这个函数会调用底层的 CUDA API，通常是 cudaIpcOpenMemHandle。
- cudaIpcOpenMemHandle 会执行一个非常重要的操作：它请求 CUDA 驱动和操作系统，在当前 GPU (GPU 0) 的虚拟地址空间中，创建一个新的内存映射。
- 这个新的内存映射，其物理后端并不是 GPU 0 自己的显存，而是通过 NVLink 连接指向 GPU 1 的物理显存中，由那个 IPC 句柄所代表的内存区域。
- 操作成功后，cudaIpcOpenMemHandle 会返回一个在 GPU 0 的地址空间中有效的虚拟地址指针。
结果:
- buffer_ptrs[1] 现在存储了这个新创建的虚拟地址指针。
- 关键效果 : 从这一刻起，当 GPU 0 上的任何 CUDA Kernel 代码试图通过 buffer_ptrs[1] 这个指针进行读写时，硬件（GPU 的内存控制器和 NVLink 硬件）会自动将这些请求路由到 GPU 1 的显存上。对 GPU 0 来说，它感觉就像在访问一块本地内存，但实际上数据流正在通过 NVLink 在两个 GPU 之间传输。

图示:

复制代码

  GPU 0 的虚拟地址空间                         GPU 1 的物理显存
+-------------------------+                   +--------------------+
|                         |                   |                    |
|                         |                   |   nvl_buffer_1     |
|   ...                   |                   |  (物理地址 0xABCD) |
|                         |                   |                    |
|   buffer_ptrs[1]        |                   +--------------------+
|  (虚拟地址 0x1234) ----(映射)---- NVLink ---->|
|                         |
|   ...                   |
|                         |
+-------------------------+

open_mem_handle 就是建立了从 0x1234 到 0xABCD 的这个映射关系。

第二行: `barrier_signal_ptrs[i] = reinterpret_cast<int>(static_cast<uint8_t>(buffer_ptrs[i]) + num_nvl_bytes);`

中文直译 : "将 buffer_ptrs[1] 这个通用指针，加上 num_nvl_bytes 字节的偏移量，然后将结果解释为一个 int 类型的指针，并存入 barrier_signal_ptrs[1]。"

深入解释:

背景 : 我们知道，每个 GPU 在最开始分配共享内存时，分配的是一块连续的大内存，其布局是：

复制代码

[ <--- nvl_buffer (num_nvl_bytes) ---> | <-- barrier_signals --> | <-- buffer_ptrs_gpu --> | ... ]
^                                       ^
|                                       |
(内存块起始地址)                       (偏移 num_nvl_bytes 处)

buffer_ptrs[i] (即 buffer_ptrs[1]) : 从第一行我们知道，这个指针指向 GPU 1 的那块大内存的起始地址。
static_cast<uint8_t*>(buffer_ptrs[i]) : C++ 不允许对 void* 类型的指针直接进行算术运算。为了进行字节级别的偏移计算，我们首先需要把它转换成一个字节指针，即 uint8_t*。
... + num_nvl_bytes : 这是指针的算术运算。它将指针向前移动 num_nvl_bytes 个字节。移动之后，这个新指针正好指向了 GPU 1 内存块中屏障信号量 (barrier_signals) 区域的起始位置。
reinterpret_cast<int*>(...) : 屏障信号量我们约定是 int 类型的。所以，我们将计算出的字节指针地址，重新解释（reinterpret_cast）为一个 int 类型的指针。
结果:
- barrier_signal_ptrs[1] 现在存储了一个可以直接访问 GPU 1 屏障信号量区域的指针。
- 关键效果 : 当 GPU 0 想要与 GPU 1 进行屏障同步时，它可以直接读写 barrier_signal_ptrs[1] 所指向的内存，来修改或检查 GPU 1 上的信号量值。

总结

这两行代码配合起来，实现了两个至关重要的目标：

建立数据通道 : 第一行代码通过 IPC 句柄，建立了从当前 GPU 到伙伴 GPU 数据缓冲区 (nvl_buffer) 的直接访问通道。
建立同步通道 : 第二行代码利用已经建立的数据通道的地址，通过简单的指针偏移，定位并获取了到伙伴 GPU 同步信号量区域 (barrier_signals) 的直接访问指针。

完成这两步之后，当前 GPU 就拥有了与节点内所有伙伴 GPU 进行数据交换 和状态同步所需的一切"硬件地址"，为后续所有复杂的节点内 All-to-All 通信和屏障操作铺平了道路。

SPMD 模式下的执行流程

cudaMemcpy 的具体目标

在进程 0 (GPU 0) 的视角下：

在进程 1 (GPU 1) 的视角下：

总结与图示

关键变量解释

断言：num_ranks == device_ids.size()

循环与 offset 的计算：offset = rdma_rank * num_nvl_ranks

for 循环的完整逻辑

总结

第一行: shared_memory_allocator.open_mem_handle(&buffer_ptrs[i], &ipc_handles[i]);

第二行: barrier_signal_ptrs[i] = reinterpret_cast<int*>(static_cast<uint8_t*>(buffer_ptrs[i]) + num_nvl_bytes);

总结

`cudaMemcpy` 的具体目标

断言：`num_ranks == device_ids.size()`

循环与 `offset` 的计算：`offset = rdma_rank * num_nvl_ranks`

`for` 循环的完整逻辑

第一行: `shared_memory_allocator.open_mem_handle(&buffer_ptrs[i], &ipc_handles[i]);`

第二行: `barrier_signal_ptrs[i] = reinterpret_cast<int>(static_cast<uint8_t>(buffer_ptrs[i]) + num_nvl_bytes);`