CUDA IPC共享内存-以DeepEP为例

前言

• deepep当中，RDMA到nvlink域的转发之前一直以为是先从RDMA recv buffer到nvlink的本地buffer，然后再从nvlink本地buffer到其他GPU的nvlink buffer
• 实际上，deepep是直接从RDMA recv buffer到其他GPU的nvlink buffer
• 这就涉及到cuda ipc共享内存的概念

相关代码

if (num_nvl_bytes > 0) {
    // Local IPC: alloc local memory and set local IPC handle
    // 在这个GPU上分配内存，并把分配的内存的地址存在buffer_ptrs[nvl_rank]中
    CUDA_CHECK(cudaMalloc(&buffer_ptrs[nvl_rank], num_nvl_bytes + fifo_bytes + buffer_ptr_bytes + task_ptr_bytes));
    // 获取该内存块的IPC句柄（相当于获取"内存钥匙"），输出：IPC句柄，输入：内存地址
    CUDA_CHECK(cudaIpcGetMemHandle(&ipc_handles[nvl_rank], buffer_ptrs[nvl_rank]));
    buffer_ptrs_gpu = reinterpret_cast<void**>(reinterpret_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes + fifo_bytes);

    // Set task fifo
    EP_HOST_ASSERT(NUM_MAX_FIFO_SLOTS % num_nvl_ranks == 0);
    task_fifo_ptrs[nvl_rank] = reinterpret_cast<int*>(reinterpret_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes);
    task_fifo_ptrs_gpu = reinterpret_cast<int**>(reinterpret_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes + fifo_bytes + buffer_ptr_bytes);

    // No need to synchronize, will do a full device sync during `sync`
    CUDA_CHECK(cudaMemsetAsync(task_fifo_ptrs[nvl_rank], 0, fifo_bytes, comm_stream));
}


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() == CUDA_IPC_HANDLE_SIZE);
    // 如果是远程GPU
    if (offset + i != rank) {
        // 复制他的钥匙
        std::memcpy(ipc_handles[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE);
        // 用钥匙打开他的门，输入：IPC句柄，输出：把获得权限的远程内存地址存在buffer_ptrs[i]中
        CUDA_CHECK(cudaIpcOpenMemHandle(&buffer_ptrs[i], ipc_handles[i], cudaIpcMemLazyEnablePeerAccess));
        // 设置任务队列指针
        task_fifo_ptrs[i] = reinterpret_cast<int*>(reinterpret_cast<uint8_t*>(buffer_ptrs[i]) + num_nvl_bytes);
    } else {
        // 如果是自己的GPU，检查自己的钥匙是否正确
        EP_HOST_ASSERT(std::memcmp(ipc_handles[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE) == 0);
    }
}

总结

• 所以实际上，buffer_ptrs数组中有num_nvl_ranks个元素，每个元素都是一个GPU的内存地址
• 其中buffer_ptrs[本地GPU的rank]是本地分配的
• 其他buffer_ptrs[远程GPU的rank]是通过cuda ipc打开的远程内存地址
• 这样就实现了直接从RDMA recv buffer到远程GPU的nvlink buffer的零拷贝传输