【cuda】deepep 学习 cudaHostGetDevicePointer cudaHostAllocMapped

https://blog.csdn.net/KIDGIN7439/article/details/146131893?spm1001.2014.3001.5502notify_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 实现中流程通常是这样的计算布局: 在 GPU 上计算所有元数据需要发送/接收多少 token 等。GPU - CPU 同步: 将计算出的“需要接收的 token 总数”从 GPU 拷贝到 CPU。CPU 分配显存: CPU 根据收到的数量调用torch.empty或cudaMalloc在 GPU 上为接收缓冲区分配正确大小的显存。开始通信: 启动dispatchKernel将数据发送到伙伴 GPU并接收来自伙伴 GPU 的数据到新分配的缓冲区。这个流程的问题在于第 2 步和第 3 步是阻塞的。CPU 必须等待 GPU 完成元数据计算然后才能进行显存分配。这期间 GPU 可能处于空闲状态无法与后续的通信操作重叠造成了“气泡”bubble。notify_dispatch的重叠Overlap机制deep_ep通过精巧的设计将这个串行过程变成了部分并行的流程。关键要素:映射的锁页内存 (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。notify_dispatchKernel 的双重任务: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 的“地址地图”都计算完毕准备好进行数据传输。CPU 端的并行操作:在 C 主机端代码 (internode_dispatch函数) 中一旦启动了notify_dispatchKernelCPU不会等待 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 T5T1: CPU 启动notify_dispatchKernel。T1 - T2: GPU 快速完成尺寸交换和总数计算并写入 Mapped Memory。T2: CPU 在轮询中检测到 Size立即开始分配显存。T2 - T4:这就是重叠发生的地方CPU 正在执行耗时的torch::empty显存分配并非瞬时完成。同时GPU 并没有停下它在继续执行notify_dispatchKernel 的任务 B计算详细的前缀和矩阵并进行硬件屏障同步。T4:dispatchKernel 可以被启动了因为它需要的所有输入包括 CPU 分配好的recv_x缓冲区和 GPU 计算好的前缀和矩阵都已经准备就绪。总结这个机制的精髓在于它巧妙地将一个大的notify_dispatch任务分解为“快速获取尺寸”和“慢速计算布局”两个子任务并利用Mapped Pinned Memory这一零拷贝的 CPU-GPU 通信桥梁实现了提前释放 CPU: CPU 不再需要等待整个布局计算完成只需要等待最先计算出的尺寸信息。隐藏显存分配延迟: CPU 分配显存的操作与 GPU 后续的布局计算和同步操作并行执行从而将显存分配的开销“隐藏”在了 GPU 的计算时间之下。通过这种方式deep_ep有效地减少了 CPU-GPU 同步点压缩了通信流水线中的“气泡”从而提升了整体的 All-to-All 通信性能。好的我们来结合代码深入且详细地解释映射的锁页内存 (Mapped Pinned Memory)在deep_ep中是如何被用来实现高效 CPU-GPU 同步的。为什么需要这种特殊的内存在dispatch过程中GPU 完成了复杂的 token 接收和排序后需要将一个关键信息通知给 CPU“我总共收到了多少个 token”CPU 需要这个信息来知道后续应该处理多少数据。为下一次combine操作分配正确大小的输出张量。在 Python 端将接收到的 token 数量返回给用户。传统的做法是GPU 将接收数量写入一个普通的 GPU 显存变量。CPU 调用cudaMemcpy(dst_cpu, src_gpu, ...)将这个值从 GPU 显存拷贝回 CPU 内存。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_castint*(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_castint*(moe_recv_expert_counter),0));for(inti0;iNUM_MAX_LOCAL_EXPERTS;i)moe_recv_expert_counter[i]-1;步骤详解cudaMallocHost(moe_recv_counter, ..., cudaHostAllocMapped):cudaMallocHost: 请求 CUDA 驱动分配锁页内存 (Pinned Memory)。普通的主机内存由malloc或new分配是可分页的其物理地址可能随时被操作系统改变。而锁页内存的物理地址是固定的这使得 GPU 的 DMA 引擎可以直接访问它。cudaHostAllocMapped: 这是一个关键标志。它告诉驱动这块锁页内存不仅要被锁定还要被映射到 GPU 的地址空间。结果: 驱动在系统 RAM 中分配了一块内存并将moe_recv_counter指向这块内存的CPU 虚拟地址。cudaHostGetDevicePointer(moe_recv_counter_mapped, ...):作用: 询问 CUDA 驱动“你刚刚为我创建的那块映射内存它在当前 GPU 设备上的虚拟地址是什么”结果: 驱动将这块内存的GPU 虚拟地址写入moe_recv_counter_mapped这个指针变量。*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 | ------------ ------------在dispatchKernel 中的使用现在我们来看intranode_dispatch函数是如何使用这些指针的// in Buffer::intranode_dispatch(...)// ... (当不是 cached_mode 时)// 1. CPU 将这块内存重置为 -1表示新一轮通信开始*moe_recv_counter-1;for(inti0;inum_local_experts;i)moe_recv_expert_counter[i]-1;// 2. 将 GPU 视角的指针 moe_recv_counter_mapped 传递给 notify_dispatch Kernelintranode::notify_dispatch(...,moe_recv_counter_mapped,// GPU-side pointer...,moe_recv_expert_counter_mapped,// GPU-side pointer...);在notify_dispatchCUDA 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 主机端的等待dispatchKernel 启动后是异步的CPU 不会等待它完成。但 CPU 需要知道结果于是它进入了一个忙等待busy-wait循环// in Buffer::intranode_dispatch(...)// 3. CPU 进入等待循环autostart_timestd::chrono::high_resolution_clock::now();while(true){// 4. CPU 通过自己的指针 moe_recv_counter 读取内存值num_recv_tokensstatic_castint(*moe_recv_counter);// 读取每个专家的计数值boolready(num_recv_tokens0);for(inti0;inum_local_expertsandready;i)readymoe_recv_expert_counter[i]0;// 5. 检查是否收到了有效值 (不再是-1)if(ready)break;// GPU 已写入完毕跳出循环// ... (超时检查)}// 6. CPU 现在知道了接收到的 token 数量num_recv_tokens_per_expert_liststd::vectorint(moe_recv_expert_counter,moe_recv_expert_counternum_local_experts);这个循环是整个机制的精髓CPU 不断地、重复地通过自己的指针moe_recv_counter去轮询 (polling)那块共享内存的值。一开始它读到的都是 -1。一旦 GPU 上的 Kernel 完成计算并写入了新的值例如 1024CPU 在下一次轮询时就会读到 1024。它发现值不再是 -1就知道 GPU 已经完成了它的任务于是跳出循环。总结映射的锁页内存提供了一种低延迟、无显式拷贝的 CPU-GPU 通信机制。创建:cudaMallocHostcudaHostAllocMappedcudaHostGetDevicePointer创建了一块物理上在系统 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_castint*(moe_recv_rdma_counter),0));// 3. CPU 初始化这块内存*moe_recv_rdma_counter-1;这里发生了什么cudaMallocHost(..., cudaHostAllocMapped): 这不仅仅是分配了一块普通的 CPU 内存。Pinned (锁页): 它告诉操作系统这块物理内存不允许被交换到磁盘上。这对于 DMA (直接内存访问) 传输是必需的因为 GPU 需要一个稳定不变的物理地址。Mapped (映射): 这是关键它告诉 CUDA 驱动这块物理内存不仅要对 CPU 可见还要映射到 GPU 的虚拟地址空间中。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 会等待这个更新。这个过程是这样的GPU 更新: 在某个notify_dispatchKernel 内部会有类似这样的代码// (在 GPU Kernel 中) // ... 经过计算得到了 rdma_recv_count ... *moe_recv_rdma_counter_mapped rdma_recv_count;GPU 通过它的入口 (_mapped指针) 将计算结果写入了那块共享的物理内存。CPU 检查: 在 C 的internode_dispatch函数中有一个while循环// (在 CPU 代码中)while(true){// CPU 通过自己的入口读取这块内存num_rdma_recv_tokensstatic_castint(*moe_recv_rdma_counter);if(num_rdma_recv_tokens0){// 检查值是否已从 -1 更新break;}// ... (超时检查) ...}CPU 在这个循环中反复地读取*moe_recv_rdma_counter的值。由于 CPU 和 GPU 操作的是同一块物理内存一旦 GPU 将值更新了例如从 -1 更新为 128CPU 下一次读取*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 异步通信和同步的关键技术。