Skip to content

Add ability to do RDMA without nvidia-peermem #704

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
stmatengss merged 1 commit into from
Aug 9, 2025
+92 −17
Merged

Add ability to do RDMA without nvidia-peermem #704

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions mooncake-common/common.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
9 changes: 8 additions & 1 deletion mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -160,7 +167,7 @@ class RdmaContext {
ibv_gid gid_;

RWSpinlock memory_regions_lock_;
std::vector<ibv_mr *> memory_region_list_;
std::vector<struct MemoryRegionMeta> memory_region_list_;
std::vector<RdmaCq> cq_list_;

std::shared_ptr<EndpointStore> endpoint_store_;
Expand Down
92 changes: 78 additions & 14 deletions mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,10 @@

#include "transport/rdma_transport/rdma_context.h"

#ifdef USE_CUDA
#include <cuda.h>

@stmatengss stmatengss Aug 1, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

By default, the CUDA library is not required. It should be included using #ifdef...#endif.

@misterwilliam misterwilliam Aug 1, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

#endif

#include <fcntl.h>
#include <sys/epoll.h>

Expand Down Expand Up @@ -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);

@staryxchen staryxchen Aug 3, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When deregistering MR, do we consider whether the memory is VRAM or DRAM?

@misterwilliam misterwilliam Aug 4, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As far as I know it doesn't matter. I've also tested with pytorch and my current implementation works.

# NVIDA cuda
>>> gt = torch.tensor([[3., -2.], [2., -2.]], device=torch.device("cuda", 0))
>>> engine.register_memory(gt.data_ptr(), gt.nbytes)
0
>>> engine.unregister_memory(gt.data_ptr())
0
# CPU host memory
>>> t = torch.tensor([[2., -2.], [2., -2.]])
>>> engine.register_memory(t.data_ptr(), t.nbytes)
0
>>> engine.unregister_memory(t.data_ptr())
0

if (ret) {
PLOG(ERROR) << "Failed to unregister memory region";
}
Expand Down Expand Up @@ -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;
}

Expand All @@ -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;
}
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -503,4 +567,4 @@ int RdmaContext::submitPostSend(
const std::vector<Transport::Slice *> &slice_list) {
return worker_pool_->submitPostSend(slice_list);
}
} // namespace mooncake
} // namespace mooncake
Loading