From d68b4e80d117a7a174ef8d91d2cfb16185ecabb9 Mon Sep 17 00:00:00 2001 From: William Chang Date: Fri, 1 Aug 2025 19:32:15 +0000 Subject: [PATCH] Add ability to do RDMA without nvidia-peermem Migrate from ibv_reg_mr() to ibv_reg_dmabuf_mr(). This allows us to not need nvidia-peermem for RDMA because ibv_reg_dmabuf_mr() does not require nvidia-peermem on nvidia GPUs while ibv_reg_mr() does require nvidia-peermem. Need to introduce a new struct (MrMeta) to track starting address of MR because ibv_reg_dmabuf_mr() is iova based and doesn't set mr->address to starting address, so we need to track it ourselves. --- mooncake-common/common.cmake | 8 +- .../transport/rdma_transport/rdma_context.h | 9 +- .../transport/rdma_transport/rdma_context.cpp | 92 ++++++++++++++++--- 3 files changed, 92 insertions(+), 17 deletions(-) diff --git a/mooncake-common/common.cmake b/mooncake-common/common.cmake index 13cfd2dc1c..8c2d82ddf9 100644 --- a/mooncake-common/common.cmake +++ b/mooncake-common/common.cmake @@ -54,7 +54,7 @@ add_compile_definitions(GLOG_USE_GLOG_EXPORT) option(BUILD_EXAMPLES "Build examples" ON) -option(BUILD_UNIT_TESTS "Build uint tests" ON) +option(BUILD_UNIT_TESTS "Build unit tests" ON) option(USE_CUDA "option for enabling gpu features" OFF) option(USE_NVMEOF "option for using NVMe over Fabric" OFF) option(USE_TCP "option for using TCP transport" ON) @@ -68,7 +68,7 @@ option(USE_HTTP "option for enable http as metadata server" ON) option(WITH_RUST_EXAMPLE "build the Rust interface and sample code for the transfer engine" OFF) option(WITH_METRICS "enable metrics and metrics reporting thread" ON) option(USE_3FS "option for using 3FS storage backend" OFF) - +option(WITH_NVIDIA_PEERMEM "disable to support RDMA without nvidia-peermem. If WITH_NVIDIA_PEERMEM=OFF then USE_CUDA=ON is required." ON) option(USE_LRU_MASTER "option for using LRU in master service" OFF) set(LRU_MAX_CAPACITY 1000) @@ -145,6 +145,10 @@ if(USE_3FS) message(STATUS "3FS storage backend is enabled") endif() +if(WITH_NVIDIA_PEERMEM) + add_compile_definitions(WITH_NVIDIA_PEERMEM) +endif() + set(GFLAGS_USE_TARGET_NAMESPACE "true") find_package(yaml-cpp REQUIRED) find_package(gflags REQUIRED) diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h index d068592c9a..2837396739 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h @@ -45,6 +45,13 @@ struct RdmaCq { volatile int outstanding; }; +struct MemoryRegionMeta { + // mr->addr is not set to starting address for iova based mr. Therefore we + // track it ourselves. + void *addr; + struct ibv_mr *mr; +}; + // RdmaContext represents the set of resources controlled by each local NIC, // including Memory Region, CQ, EndPoint (QPs), etc. class RdmaContext { @@ -160,7 +167,7 @@ class RdmaContext { ibv_gid gid_; RWSpinlock memory_regions_lock_; - std::vector memory_region_list_; + std::vector memory_region_list_; std::vector cq_list_; std::shared_ptr endpoint_store_; diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp index fbf5560375..bf2ac367c1 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp @@ -14,6 +14,10 @@ #include "transport/rdma_transport/rdma_context.h" +#ifdef USE_CUDA +#include +#endif + #include #include @@ -148,7 +152,7 @@ int RdmaContext::deconstruct() { endpoint_store_->destroyQPs(); for (auto &entry : memory_region_list_) { - int ret = ibv_dereg_mr(entry); + int ret = ibv_dereg_mr(entry.mr); if (ret) { PLOG(ERROR) << "Failed to unregister memory region"; } @@ -200,14 +204,51 @@ int RdmaContext::registerMemoryRegion(void *addr, size_t length, int access) { << "shrink it to " << globalConfig().max_mr_size; length = (size_t)globalConfig().max_mr_size; } - ibv_mr *mr = ibv_reg_mr(pd_, addr, length, access); - if (!mr) { + + MemoryRegionMeta mrMeta; +#if !defined(WITH_NVIDIA_PEERMEM) && defined(USE_CUDA) + // Implement register memory in a way that does not assume the presence of + // nvidia-peermem. If memory is on CPU call ibv_reg_mr() as usual. If memory + // is on GPU then use ibv_reg_dmabuf_mr() instead which does not require + // nvidia-peermem. + CUmemorytype memType; + CUresult result = cuPointerGetAttribute( + &memType, CU_POINTER_ATTRIBUTE_MEMORY_TYPE, (CUdeviceptr)addr); + + // Register memory depending on whether memory is on host or GPU. + if (result != CUDA_SUCCESS || memType == CU_MEMORYTYPE_HOST) { + mrMeta.addr = addr; + mrMeta.mr = ibv_reg_mr(pd_, addr, length, access); + } else if (memType == CU_MEMORYTYPE_DEVICE) { + size_t allocSize; + cuPointerGetAttribute(&allocSize, CU_POINTER_ATTRIBUTE_RANGE_SIZE, + (CUdeviceptr)addr); + int dmabuf_fd; + result = cuMemGetHandleForAddressRange( + &dmabuf_fd, (CUdeviceptr)addr, allocSize, + CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD, 0); + if (result != CUDA_SUCCESS) { + const char *errStr; + cuGetErrorString(result, &errStr); + LOG(ERROR) << "Failed to retrieve dmabuf for " << (uintptr_t)addr + << " cuda error=" << errStr; + return ERR_CONTEXT; + } + mrMeta.addr = addr; + mrMeta.mr = ibv_reg_dmabuf_mr(pd_, 0 /* offset */, length, + (uintptr_t)addr, dmabuf_fd, access); + } +#else + mrMeta.addr = addr; + mrMeta.mr = ibv_reg_mr(pd_, addr, length, access); +#endif + if (!mrMeta.mr) { PLOG(ERROR) << "Failed to register memory " << addr; return ERR_CONTEXT; } RWSpinlock::WriteGuard guard(memory_regions_lock_); - memory_region_list_.push_back(mr); + memory_region_list_.push_back(mrMeta); return 0; } @@ -218,9 +259,9 @@ int RdmaContext::unregisterMemoryRegion(void *addr) { has_removed = false; for (auto iter = memory_region_list_.begin(); iter != memory_region_list_.end(); ++iter) { - if ((*iter)->addr <= addr && - addr < (char *)((*iter)->addr) + (*iter)->length) { - if (ibv_dereg_mr(*iter)) { + if (iter->addr <= addr && + addr < (char *)(iter->addr) + iter->mr->length) { + if (ibv_dereg_mr(iter->mr)) { LOG(ERROR) << "Failed to unregister memory " << addr; return ERR_CONTEXT; } @@ -237,9 +278,9 @@ uint32_t RdmaContext::rkey(void *addr) { RWSpinlock::ReadGuard guard(memory_regions_lock_); for (auto iter = memory_region_list_.begin(); iter != memory_region_list_.end(); ++iter) - if ((*iter)->addr <= addr && - addr < (char *)((*iter)->addr) + (*iter)->length) - return (*iter)->rkey; + if (iter->addr <= addr && + addr < (char *)(iter->addr) + iter->mr->length) + return iter->mr->rkey; LOG(ERROR) << "Address " << addr << " rkey not found for " << deviceName(); return 0; @@ -249,9 +290,9 @@ uint32_t RdmaContext::lkey(void *addr) { RWSpinlock::ReadGuard guard(memory_regions_lock_); for (auto iter = memory_region_list_.begin(); iter != memory_region_list_.end(); ++iter) - if ((*iter)->addr <= addr && - addr < (char *)((*iter)->addr) + (*iter)->length) - return (*iter)->lkey; + if (iter->addr <= addr && + addr < (char *)(iter->addr) + iter->mr->length) + return iter->mr->lkey; LOG(ERROR) << "Address " << addr << " lkey not found for " << deviceName(); return 0; @@ -404,6 +445,29 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port, return ERR_CONTEXT; } +#if !defined(WITH_NVIDIA_PEERMEM) && defined(USE_CUDA) + // Verify dmabuf support which is required if not using nvidia-peermem. + // Assume device index matches. + CUdevice cuDevice; + CUresult result = cuDeviceGet(&cuDevice, i); + if (result != CUDA_SUCCESS) { + LOG(ERROR) << "Failed to query CUDA device"; + return ERR_CONTEXT; + } + int dmaBufSupported; + result = cuDeviceGetAttribute( + &dmaBufSupported, CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED, cuDevice); + if (result != CUDA_SUCCESS) { + LOG(ERROR) << "Failed to query CUDA device attributes"; + return ERR_CONTEXT; + } + if (!dmaBufSupported) { + LOG(ERROR) << "DMA BUF supported required for GPU RDMA without " + "nvidia-peermem"; + return ERR_CONTEXT; + } +#endif + ibv_port_attr port_attr; ret = ibv_query_port(context, port, &port_attr); if (ret) { @@ -503,4 +567,4 @@ int RdmaContext::submitPostSend( const std::vector &slice_list) { return worker_pool_->submitPostSend(slice_list); } -} // namespace mooncake \ No newline at end of file +} // namespace mooncake