https://blog.csdn.net/KIDGIN7439/article/details/146131893?spm=1001.2014.3001.5502

notify_dispatch过程中会计算其他所有rank发送给当前rank多少token，写入到host的moe_recv_counter_mapped，还会计算其他所有rdma_rank发送给当前rank多少token，写入到host的moe_recv_rdma_counter_mapped，这里通过cpu端轮询这个值，轮询到之后就可以在notify_dispatch不结束的情况下开始分配显存，做到overlap。

详细解释

好的，您对 notify_dispatch 过程的理解非常精准！这确实是 deep_ep（尤其是在标准模式下）实现计算-通信重叠（overlap）的一个核心技巧。我们来详细拆解这个机制。

这个机制主要应用在非缓存模式 (non-cached mode) 的 internode_dispatch 函数中，因为此时每次通信的 token 数量都是动态变化的，需要实时计算。

背景：为什么需要这个机制？

在传统的 All-to-All 实现中，流程通常是这样的：

  1. 计算布局: 在 GPU 上计算所有元数据（需要发送/接收多少 token 等）。
  2. GPU -> CPU 同步: 将计算出的“需要接收的 token 总数”从 GPU 拷贝到 CPU。
  3. CPU 分配显存: CPU 根据收到的数量，调用 torch.empty 或 cudaMalloc 在 GPU 上为接收缓冲区分配正确大小的显存。
  4. 开始通信: 启动 dispatch Kernel，将数据发送到伙伴 GPU，并接收来自伙伴 GPU 的数据到新分配的缓冲区。

这个流程的问题在于，第 2 步和第 3 步是阻塞的。CPU 必须等待 GPU 完成元数据计算，然后才能进行显存分配。这期间 GPU 可能处于空闲状态，无法与后续的通信操作重叠，造成了“气泡”（bubble）。

notify_dispatch 的重叠（Overlap）机制

deep_ep 通过精巧的设计，将这个串行过程变成了部分并行的流程。

关键要素:

  1. 映射的锁页内存 (Mapped Pinned Memory):

    • 在 Buffer 构造函数中，通过 cudaMallocHost(..., cudaHostAllocMapped) 创建了 moe_recv_counter 和 moe_recv_rdma_counter。
    • 这种内存的特殊之处在于它同时拥有一个 CPU 地址和一个 GPU 地址。
    • CPU 可以直接访问 moe_recv_counter。
    • GPU Kernel 可以直接访问 moe_recv_counter_mapped。
    • 对这块内存的写入对 CPU 和 GPU 几乎是立即可见的（遵循一定的内存一致性模型），无需显式的 cudaMemcpy。

  2. notify_dispatch Kernel 的双重任务:
    internode::notify_dispatch 这个 C++ 函数（它会启动一个或多个 CUDA Kernel）被赋予了多个任务，并且这些任务在时间上是错开的：

    • 任务 A (早期执行):
      • 交换尺寸: GPU Kernel 会通过 NVLink 和 RDMA，与其他 GPU 交换“我准备发给你多少个 token”这类尺寸信息。
      • 计算接收总数: 每个 GPU 根据收到的信息，在 GPU 上累加出自己将要接收的 token 总数 (num_recv_tokens) 和将要通过 RDMA 接收的 token 总数 (num_rdma_recv_tokens)。
      • 写入 Mapped Memory: Kernel 立即将这两个总数通过 moe_recv_counter_mapped 和 moe_recv_rdma_counter_mapped 这两个 GPU 指针，写入到那块共享的锁页内存中。这是整个机制的触发点！
    • 任务 B (后期执行):
      • 计算详细布局: Kernel 继续在 GPU 上执行，计算更复杂的元数据，如前缀和矩阵 (channel_prefix_matrix 等)，这些是实际 dispatch 数据时所需的详细“地址地图”。这个过程比计算总数要耗时得多。
      • 硬件屏障同步: 与其他 GPU 进行硬件屏障同步，确保所有 GPU 的“地址地图”都计算完毕，准备好进行数据传输。

  3. CPU 端的并行操作:

    • 在 C++ 主机端代码 (internode_dispatch 函数) 中，一旦启动了 notify_dispatch Kernel，CPU 不会等待 Kernel 完成。
    • 它会立即进入一个 while(true) 循环，**忙等待 (spin-wait) ** *moe_recv_counter >= 0 和 *moe_recv_rdma_counter >= 0。
    • 由于 GPU Kernel 在任务 A 中很早就写入了这两个值，CPU 会很快检测到值的变化（不再是初始的 -1）。
    • 一旦轮询成功，CPU 就拿到了最终需要分配的接收缓冲区的大小 (num_recv_tokens)。
    • CPU 立刻调用 torch::empty(...) 为 recv_x, recv_topk_idx 等张量分配显存。

时间线图示

让我们用时间线来清晰地展示这个重叠过程：

传统串行流程:

CPU:  |----------------| 等待GPU | 分配显存 | 等待GPU | ...
      |                |         |          |         |
GPU:  | 计算布局 & Size |---------|----------| 通信... | ...
      ^                ^         ^
      |                |         |
      T1               T2        T3
      (T1->T2: 布局计算, T2->T3: CPU等待+分配, T3之后: 通信)

DeepEP 重叠流程:

CPU:  |----------| 轮询Size | 分配显存 |-----------------| 等待GPU | ...
      |          |          |          |                 |         |
GPU:  | 计算Size | 计算布局  | 屏障同步 | 通信...         |---------| ...
      ^          ^          ^          ^                 ^
      |          |          |          |                 |
      T1         T2         T3         T4                T5
  • T1: CPU 启动 notify_dispatch Kernel。
  • T1 -> T2: GPU 快速完成尺寸交换和总数计算，并写入 Mapped Memory。
  • T2: CPU 在轮询中检测到 Size，立即开始分配显存。
  • T2 -> T4: 这就是重叠发生的地方！
    • CPU 正在执行耗时的 torch::empty（显存分配并非瞬时完成）。
    • 同时，GPU 并没有停下，它在继续执行 notify_dispatch Kernel 的任务 B：计算详细的前缀和矩阵，并进行硬件屏障同步。
  • T4: dispatch Kernel 可以被启动了，因为它需要的所有输入（包括 CPU 分配好的 recv_x 缓冲区和 GPU 计算好的前缀和矩阵）都已经准备就绪。

总结

这个机制的精髓在于，它巧妙地将一个大的 notify_dispatch 任务分解为“快速获取尺寸”和“慢速计算布局”两个子任务，并利用 Mapped Pinned Memory 这一零拷贝的 CPU-GPU 通信桥梁，实现了：

  1. 提前释放 CPU: CPU 不再需要等待整个布局计算完成，只需要等待最先计算出的尺寸信息。
  2. 隐藏显存分配延迟: CPU 分配显存的操作与 GPU 后续的布局计算和同步操作并行执行，从而将显存分配的开销“隐藏”在了 GPU 的计算时间之下。

通过这种方式，deep_ep 有效地减少了 CPU-GPU 同步点，压缩了通信流水线中的“气泡”，从而提升了整体的 All-to-All 通信性能。

好的，我们来结合代码，深入且详细地解释映射的锁页内存 (Mapped Pinned Memory) 在 deep_ep 中是如何被用来实现高效 CPU-GPU 同步的。

为什么需要这种特殊的内存？

在 dispatch 过程中，GPU 完成了复杂的 token 接收和排序后，需要将一个关键信息通知给 CPU：“我总共收到了多少个 token？”

CPU 需要这个信息来：

  1. 知道后续应该处理多少数据。
  2. 为下一次 combine 操作分配正确大小的输出张量。
  3. 在 Python 端将接收到的 token 数量返回给用户。

传统的做法是：

  1. GPU 将接收数量写入一个普通的 GPU 显存变量。
  2. CPU 调用 cudaMemcpy(dst_cpu, src_gpu, ...) 将这个值从 GPU 显存拷贝回 CPU 内存。
  3. cudaMemcpy 是一个阻塞操作，CPU 会一直等待，直到数据拷贝完成。

这种方式虽然可行，但有两个缺点：

  • 延迟高: cudaMemcpy 本身有不可忽略的开销。
  • 同步开销大: 强制的阻塞同步会打断 CPU 的执行流。

映射的锁页内存 就是为了解决这个问题而生的。

deep_ep.cpp 中的实现

让我们来看 Buffer 构造函数中的相关代码：

// in Buffer::Buffer(...)

// MoE counter (for total received tokens)
// 1. 在主机端分配一块特殊的内存
CUDA_CHECK(cudaMallocHost(&moe_recv_counter, sizeof(int64_t), cudaHostAllocMapped));

// 2. 获取这块内存对应的 GPU 设备指针
CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_counter_mapped, const_cast<int*>(moe_recv_counter), 0));

// 3. 初始化这块内存
*moe_recv_counter = -1;

// MoE expert-level counter (for per-expert received tokens)
// ... (逻辑完全相同，只是分配了一个数组)
CUDA_CHECK(cudaMallocHost(&moe_recv_expert_counter, sizeof(int) * NUM_MAX_LOCAL_EXPERTS, cudaHostAllocMapped));
CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_expert_counter_mapped, const_cast<int*>(moe_recv_expert_counter), 0));
for (int i = 0; i < NUM_MAX_LOCAL_EXPERTS; ++i)
    moe_recv_expert_counter[i] = -1;
步骤详解：

  1. cudaMallocHost(&moe_recv_counter, ..., cudaHostAllocMapped):

    • cudaMallocHost: 请求 CUDA 驱动分配锁页内存 (Pinned Memory)。普通的主机内存（由 malloc 或 new 分配）是可分页的，其物理地址可能随时被操作系统改变。而锁页内存的物理地址是固定的，这使得 GPU 的 DMA 引擎可以直接访问它。
    • cudaHostAllocMapped: 这是一个关键标志。它告诉驱动，这块锁页内存不仅要被锁定，还要被映射到 GPU 的地址空间。
    • 结果: 驱动在系统 RAM 中分配了一块内存，并将 moe_recv_counter 指向这块内存的 CPU 虚拟地址。

  2. cudaHostGetDevicePointer(&moe_recv_counter_mapped, ...):

    • 作用: 询问 CUDA 驱动：“你刚刚为我创建的那块映射内存，它在当前 GPU 设备上的虚拟地址是什么？”
    • 结果: 驱动将这块内存的 GPU 虚拟地址 写入 moe_recv_counter_mapped 这个指针变量。

  3. *moe_recv_counter = -1:

    • 作用: CPU 通过它的地址 moe_recv_counter，将这块共享内存的初始值设为 -1。这是一种信令机制，表示 GPU 尚未写入有效值。

执行完这三步后，我们拥有了对同一块物理内存的两个“视角”：

+------------------------------------------------+
|                   System RAM                   |
|                                                |
|      +-----------------------------------+     |
|      |  物理内存块 (e.g., at addr 0xAAAA)  |     |
|      +-----------------------------------+     |
|         ^                           ^          |
|         |                           |          |
+---------|---------------------------|----------+
          | (CPU MMU 映射)            | (GPU MMU/IOMMU 映射)
          |                           |
    +-----|------+              +-----|------+
    | CPU 视角   |              | GPU 视角   |
    |------------|              |------------|
    | 指针:      |              | 指针:      |
    | moe_recv_counter |              | moe_recv_counter_mapped |
    +------------+              +------------+

在 dispatch Kernel 中的使用

现在，我们来看 intranode_dispatch 函数是如何使用这些指针的：

// in Buffer::intranode_dispatch(...)

// ... (当不是 cached_mode 时)

// 1. CPU 将这块内存重置为 -1，表示新一轮通信开始
*moe_recv_counter = -1;
for (int i = 0; i < num_local_experts; ++i)
    moe_recv_expert_counter[i] = -1;

// 2. 将 GPU 视角的指针 moe_recv_counter_mapped 传递给 notify_dispatch Kernel
intranode::notify_dispatch(...,
                           moe_recv_counter_mapped, // GPU-side pointer
                           ...,
                           moe_recv_expert_counter_mapped, // GPU-side pointer
                           ...
                          );

在 notify_dispatch CUDA Kernel 内部 (GPU 端):

  • Kernel 会进行复杂的计算，最终确定总共收到了 N 个 token。
  • Kernel 中一个特定的线程（通常是 threadIdx.x == 0 的线程）会执行一个原子写操作，例如： 
    // 伪代码，在 CUDA Kernel 中
if (threadIdx.x == 0) {
    // ... 计算出总接收数 total_received_count
    *moe_recv_counter_mapped = total_received_count;
    // ... 计算出每个专家的接收数
    for (int i = 0; i < num_local_experts; ++i) {
        moe_recv_expert_counter_mapped[i] = per_expert_count[i];
    }
}
  • GPU 通过它的指针 moe_recv_counter_mapped，直接向那块共享的物理内存写入了结果。

在 C++ 主机端的等待

dispatch Kernel 启动后是异步的，CPU 不会等待它完成。但 CPU 需要知道结果，于是它进入了一个忙等待（busy-wait）循环：

// in Buffer::intranode_dispatch(...)

// 3. CPU 进入等待循环
auto start_time = std::chrono::high_resolution_clock::now();
while (true) {
    // 4. CPU 通过自己的指针 moe_recv_counter 读取内存值
    num_recv_tokens = static_cast<int>(*moe_recv_counter);

    // 读取每个专家的计数值
    bool ready = (num_recv_tokens >= 0);
    for (int i = 0; i < num_local_experts and ready; ++i)
        ready &= moe_recv_expert_counter[i] >= 0;

    // 5. 检查是否收到了有效值 (不再是-1)
    if (ready)
        break; // GPU 已写入完毕，跳出循环

    // ... (超时检查)
}
// 6. CPU 现在知道了接收到的 token 数量
num_recv_tokens_per_expert_list = std::vector<int>(moe_recv_expert_counter, moe_recv_expert_counter + num_local_experts);

这个循环是整个机制的精髓：

  • CPU 不断地、重复地通过自己的指针 moe_recv_counter 去轮询 (polling) 那块共享内存的值。
  • 一开始，它读到的都是 -1。
  • 一旦 GPU 上的 Kernel 完成计算并写入了新的值（例如 1024），CPU 在下一次轮询时就会读到 1024。
  • 它发现值不再是 -1，就知道 GPU 已经完成了它的任务，于是跳出循环。

总结

映射的锁页内存提供了一种低延迟、无显式拷贝的 CPU-GPU 通信机制。

  • 创建: cudaMallocHost + cudaHostAllocMapped + cudaHostGetDevicePointer 创建了一块物理上在系统 RAM，但逻辑上 CPU 和 GPU 都能访问的“共享区域”。
  • 通信:
    • GPU -> CPU: GPU Kernel 使用其设备指针（..._mapped）向该区域写入结果。
    • CPU -> GPU: CPU 使用其主机指针轮询读取该区域，以检查 GPU 是否已完成任务。
  • 优势: 避免了 cudaMemcpy 的开销和阻塞，实现了非常快速的、基于轮询的异步状态通知。这对于需要 CPU 尽快知道 GPU 计算结果的场景（如动态确定下一步操作的缓冲区大小）至关重要。

cudaHostGetDevicePointer 请问CPU能访问mapped吗 ，比如直接检查 moe_recv_rdma_counter_mapped 这个是否不是-1 了

您好，这是一个非常棒的问题，直击了 Mapped Pinned Memory (锁页映射内存) 的核心工作原理！

请问 CPU 能访问 moe_recv_rdma_counter_mapped 吗？比如直接检查这个是否不是 -1 了？

答案是：不可以。CPU 不能直接访问 moe_recv_rdma_counter_mapped。

moe_recv_rdma_counter_mapped 是一个 GPU 上的设备指针 (device pointer)。CPU 代码如果直接尝试解引用这个指针，会导致段错误 (Segmentation Fault)，因为这个地址在 CPU 的内存地址空间中是无效的。

那么，CPU 是如何检查到 GPU 已经更新了值的呢？这就要靠与 _mapped 指针成对出现的那个原始的 CPU 指针了。

Mapped Pinned Memory 的工作机制

让我们回顾一下 moe_recv_rdma_counter 的创建过程：

// 1. 在 CPU 端分配一块特殊的内存：锁页内存 (Pinned Memory)
CUDA_CHECK(cudaMallocHost(&moe_recv_rdma_counter, sizeof(int), cudaHostAllocMapped));

// 2. 获取这块锁页内存在 GPU 上的 "别名" (设备指针)
CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_rdma_counter_mapped, const_cast<int*>(moe_recv_rdma_counter), 0));

// 3. CPU 初始化这块内存
*moe_recv_rdma_counter = -1;

这里发生了什么？

  1. cudaMallocHost(..., cudaHostAllocMapped): 这不仅仅是分配了一块普通的 CPU 内存。

    • Pinned (锁页): 它告诉操作系统，这块物理内存不允许被交换到磁盘上。这对于 DMA (直接内存访问) 传输是必需的，因为 GPU 需要一个稳定不变的物理地址。
    • Mapped (映射): 这是关键！它告诉 CUDA 驱动，这块物理内存不仅要对 CPU 可见，还要映射到 GPU 的虚拟地址空间中。

  2. cudaHostGetDevicePointer(...): 这个函数就是用来查询这个映射关系的。它会返回一个在 GPU 地址空间中有效的指针 (moe_recv_rdma_counter_mapped)，这个指针和 CPU 端的指针 (moe_recv_rdma_counter) 指向同一块物理内存。

两个指针，一块物理内存

现在，我们有了两个“入口”来访问同一块物理内存：

  • moe_recv_rdma_counter: 这是 CPU 的入口。它是一个标准的 C++ 指针，在 CPU 代码中可以安全地读写。
  • moe_recv_rdma_counter_mapped: 这是 GPU 的入口。它是一个 CUDA 设备指针，只能在 GPU Kernel 代码中安全地读写。

图示:

+-------------------------------------------------------------+
|                     同一块物理内存 (RAM)                      |
|                     (地址稳定，不会被换出)                      |
|                                                             |
|                           [ -1 ]                            |
|                                                             |
+-------------------------------------------------------------+
       ^                                    ^
       |                                    |
       |  (通过系统总线访问)                  |  (通过 PCIe/NVLink 访问)
       |                                    |
+------|------------------+      +-----------|-------------------+
|      CPU 代码空间        |      |       GPU Kernel 代码空间     |
|                          |      |                             |
|  int* moe_recv_rdma_counter; |      | int* moe_recv_rdma_counter_mapped; |
|  *moe_recv_rdma_counter;  |      | *moe_recv_rdma_counter_mapped; |
|  (安全，可读写)          |      | (安全，可读写)              |
|                          |      |                             |
+--------------------------+      +-----------------------------+

CPU 如何检查 GPU 的更新？

在 deep_ep 的代码中，GPU Kernel 会更新这块内存，而 CPU 会等待这个更新。这个过程是这样的：

  1. GPU 更新: 在某个 notify_dispatch Kernel 内部，会有类似这样的代码：

    // (在 GPU Kernel 中)
// ... 经过计算，得到了 rdma_recv_count ...
*moe_recv_rdma_counter_mapped = rdma_recv_count;

    GPU 通过它的入口 (_mapped 指针) 将计算结果写入了那块共享的物理内存。

  2. CPU 检查: 在 C++ 的 internode_dispatch 函数中，有一个 while 循环：

    // (在 CPU 代码中)
while (true) {
    // CPU 通过自己的入口读取这块内存
    num_rdma_recv_tokens = static_cast<int>(*moe_recv_rdma_counter);

    if (num_rdma_recv_tokens >= 0) { // 检查值是否已从 -1 更新
        break;
    }
    // ... (超时检查) ...
}

    CPU 在这个循环中，反复地读取 *moe_recv_rdma_counter 的值。由于 CPU 和 GPU 操作的是同一块物理内存，一旦 GPU 将值更新了（例如，从 -1 更新为 128），CPU 下一次读取 *moe_recv_rdma_counter 时，就会立即看到这个新值 128，然后循环结束。

总结

  • CPU 不能直接访问 moe_recv_rdma_counter_mapped，因为它是 GPU 设备指针。
  • CPU 可以直接访问 moe_recv_rdma_counter，因为它是 CPU 主机指针。
  • cudaHostGetDevicePointer 的魔力在于它建立了这两个指针到同一块物理内存的映射关系，使得 GPU 的写入可以被 CPU “零拷贝”地、高效地观察到。

所以，CPU 检查的是 *moe_recv_rdma_counter 是否不再是 -1，而不是 moe_recv_rdma_counter_mapped。这是实现高效 GPU-CPU 异步通信和同步的关键技术。

# 学习
Logo
天工造物局

更多推荐

SWE-RL进阶实战：如何将奖励函数集成到你自己的强化学习训练框架

SWE-RL是首个利用开源软件演进数据与规则奖励来扩展LLM推理能力的强化学习开源项目，其官方仓库中的核心奖励函数实现堪称"即插即用"的宝藏。本文将带你完整掌握SWE-RL奖励函数的集成方法，手把手教你如何把这套基于序列相似度的奖励机制迁移到自己的强化学习训练框架中，让大模型学会像人类工程师一样修改真实代码。## 为什么你需要SWE-RL的奖励函数？在传统的RLHF训练中，奖励信号往往来自

avatar 天工造物局

告别局域网：VisionProTeleop 外部网络模式实现跨地域机器人遥操作

想让机器人遥操作摆脱同一 WiFi 的限制？开源项目 VisionProTeleop 从 v2.5 版本起推出**外部网络模式**，用一组房间码就能让 Vision Pro 与远在实验室的机器人跨地域连接，实现在家远程操控机械臂的终极体验。本文将带你快速理解其原理，并手把手完成配置。[![VisionProTeleop 状态窗口展示跨地域遥操作连接信息](https://raw.gitcode

avatar 天工造物局
cover

（AAAI 2024）SCIC：基于情境因果影响的协作多智能体强化学习

avatar 天工造物局