brpc brpc支持GPU Direct RDMA

Is your feature request related to a problem?

Describe the solution you'd like 在GPU上训练时使用brpc的RDMA传输，期望使用GDR能力，直接将RPC请求写入到GPU的buffer中，再通过cudamemcp或者gdrcopy从GPU上拷贝出协议部分进行解析，以减少CPU和GPU之间的数据传输。

Describe alternatives you've considered 进行的尝试：将block_pool.cpp中的BlockPool替换成了通过cuda分配的显存，但在使用IOBuf时，会在分配出来的内存上进行placement new，

而这导致CPU直接访问显存，产生core dumped

因此，有以下三种想法： 1.修改block_pool.cpp，新增一个ObjectPool用于分配Block的实体，将Block的data指向BlockPool -> 即 Block的实体部分与数据部分分离在此基础上，修改IOBuf的create block接口，objmem_allocate分配一个block出来，blockmem_allocate分配一段内存用于传输RPC请求

2.增加RDMA Write/Read接口；原本brpc的Send/Recv用于传输控制数据（例如，对端GPU地址），再使用RDMA Write直接将训练数据写到对端GPU地址上

3.使用append_with_user_data，将attachment从显存上分配，发送时使用RDMA的scatter-gatther特性，使其request请求写入到主机内存，attachment写入到显存上。但是此接口也是会在分配出来的内存上进行placement new，因此也存在core dumped的问题，也需要修改IOBuf的接口

Additional context/screenshots 想请教这三个中是否有具有可行性的方案

Sep 22 '25 14:09 sunce4t

void* ExtendBlockPoolByUser(void* region_base, size_t region_size, int block_type); 这个函数支持将用户申请的内存，注册到rdma来使用，能否满足需求呢？

Sep 23 '25 08:09 yanglimingcn

void* ExtendBlockPoolByUser(void* region_base, size_t region_size, int block_type); 这个函数支持将用户申请的内存，注册到rdma来使用，能否满足需求呢？

在我的尝试中，我使用了

 bool InitBlockPool(RegisterCallback cb) {
     if (!cb) {
         errno = EINVAL;
@@ -322,7 +373,20 @@ bool InitBlockPool(RegisterCallback cb) {
     g_tls_info_mutex = new butil::Mutex;

     if (FLAGS_rdma_memory_pool_user_specified_memory) {
-        return true;
+        if(FLAGS_rdma_memory_gpu_memory) {
+
+          void * gpu_mem = InitGPUMemoryForRDMA();
+          if(gpu_mem == NULL) {
+              LOG(ERROR) << "Could not init GPU for RDMA";
+              return false;
+          }
+          LOG(INFO) << "Init GPU Memory: " << gpu_mem;
+          if(ExtendBlockPoolByUser(gpu_mem, FLAGS_rdma_memory_pool_initial_size_mb, BLOCK_DEFAULT) != NULL) {
+              return true;
+          }
+          else return false;
+       }
+       return true;
     }

ExtendBlockPoolByUser注册GPU上的内存，这可以注册成功，但是正如我上面所说的，IOBuf在创建block的时候，会在这块内存上进行placement new构造block对象，也就是说直接在GPU内存上进行构造了，造成core dumped

我目前正在尝试第一个想法，下面是示例，模仿原有的block pool，新增AllocObj（还有其余接口），用于分配一块内存来构造block对象

struct ObjIdleNode {
    void* start;
    size_t len;
    ObjIdleNode* next;
};
static int g_obj_region_num = 0;
static Region g_obj_regions[RDMA_MEMORY_POOL_MAX_REGIONS];
static __thread ObjIdleNode* tls_obj_idle_list = NULL;
static __thread size_t tls_obj_idle_num = 0;
static __thread bool tls_obj_inited = false;
static butil::Mutex* g_tls_obj_info_mutex = NULL;
static size_t g_tls_obj_info_cnt = 0;
static size_t* g_tls_obj_info[1024];
static void* AllocObjFrom(int obj_type) {
    void* ptr = NULL;
    if (tls_obj_idle_list != NULL){
        CHECK(tls_obj_idle_num > 0);
        ObjIdleNode* n = tls_obj_idle_list;
        ptr = n->start;
        if (n->len > g_obj_block_size[obj_type]) {
            n->start = (char *)n->start + g_obj_block_size[obj_type];
            n->len -= g_obj_block_size[obj_type];
            return ptr;
        } else if(n->len == g_obj_block_size[obj_type]) {
            tls_obj_idle_list = n->next;
            tls_obj_idle_num--;
            butil::return_object<ObjIdleNode>(n);
            return ptr;
        } else if(n->len < g_obj_block_size[obj_type]) {
            // TODO
        }
        return ptr;
    }

    uint64_t index = butil::fast_rand() % g_buckets;
    BAIDU_SCOPED_LOCK(*g_obj_info->lock[obj_type][index]);
    ObjIdleNode* node = g_obj_info->idle_list[obj_type][index];
    if (!node) {
        BAIDU_SCOPED_LOCK(g_obj_info->extend_lock);
        node = g_obj_info->idle_list[obj_type][index];
        if (!node) {
            // There is no block left, extend a new region
            if (!ExtendObjPoolByUser(FLAGS_rdma_memory_pool_increase_size_mb,
                                 block_type)) {
                LOG_EVERY_SECOND(ERROR) << "Fail to extend new region. "
                                        << "You can set the size of memory pool larger. "
                                        << "Refer to the help message of these flags: "
                                        << "rdma_memory_pool_initial_size_mb, "
                                        << "rdma_memory_pool_increase_size_mb, "
                                        << "rdma_memory_pool_max_regions.";
                return NULL;
            }
            node = g_obj_info->idle_list[obj_type][index];
        }
    }
    if (node) {
        ptr = node->start;
        if (node->len > g) {
            node->start = (char *)node->start + g_obj_block_size[obj_type];
            node->len -= size;

            g_obj_info->idle_size[obj_type][index] -= g_obj_block_size[obj_type];
        } else if(node->len == g_obj_block_size[obj_type]) {
            g_obj_info->idle_list[obj_type][index] = node->next;
            butil::return_object<ObjIdleNode>(node);
            g_obj_info->idle_size[block_type][index] -= g_obj_block_size[obj_type];
        } else if(node->len < g_obj_block_size[obj_type]) {
            // TODO
        }
    } else {
        return NULL;
    }

    // Move more blocks from global list to tls list
    if (obj_type == OBJ_BLOCK_DEFAULT) {
        node = g_obj_info->idle_list[0][index];
        tls_obj_idle_list = node;
        ObjIdleNode* last_node = NULL;
        while (node) {
            if (tls_obj_idle_num > (uint32_t)FLAGS_rdma_memory_pool_tls_cache_num / 2
                    || node->len > g_obj_block_size[0]) {
                break;
            }
            tls_obj_idle_num++;
            last_node = node;
            node = node->next;
        }
        if (tls_obj_idle_num == 0) {
            tls_obj_idle_list = NULL;
        } else {
            g_obj_info->idle_list[OBJ_BLOCK_DEFAULT][index] = node;
        }
        if (last_node) {
            last_node->next = NULL;
        }
    }

    // if (locked) {
    //     g_dump_mutex->unlock();
    // }
    return ptr;
}

Sep 23 '25 08:09 sunce4t

这是GPU内存的限制吗？在GPU内存上直接初始化，和在GPU内存直接修改这块内存的数据，应该是一回事。

Sep 23 '25 08:09 yanglimingcn

这是GPU内存的限制吗？在GPU内存上直接初始化，和在GPU内存直接修改这块内存的数据，应该是一回事。

是的，我们用的是cuMemAlloc这个API，这并不为CPU和GPU之间提供统一地址空间，因此直接访问就报错了；也可以使用cudaMallocManaged统一内存管理，或者cudaHostAlloc去进行映射。但都是有一定开销的。

Sep 23 '25 08:09 sunce4t

这个内存开销是在初始化分配时候有开销，在平时使用过程也会有开销吗？地址动态映射的开销是吗？

Sep 23 '25 08:09 yanglimingcn

这个内存开销是在初始化分配时候有开销，在平时使用过程也会有开销吗？地址动态映射的开销是吗？

以cudaMallocManaged为例，主要是CPU和GPU访问时，如果找不到对应页，就会产生缺页中断，在CPU和GPU之间进行页面迁移另外，我们还担忧，如果使用cudaMallocManaged，页面被迁移到Host内，此时如果使用RDMA要发送此页面的数据的话，可能产生地址转换问题

Sep 23 '25 08:09 sunce4t

@sunce4t 请教一下，训练使用GPU Direct RDMA的时候，需要精细化管理显存，让每个RPC的数据写入指定的显存吗？还是直接从池子里取一块显存写入？

Sep 23 '25 09:09 chenBright

@sunce4t 请教一下，训练使用GPU Direct RDMA的时候，需要精细化管理显存，让每个RPC的数据写入指定的显存吗？还是直接从池子里取一块显存写入？

hi，最好的方式肯定是能够精细化管理。目前我们想先把这个功能搞起来，哪怕是池子的方式也行，再一步步改进。

Sep 23 '25 14:09 sunce4t

@sunce4t 请教一下，训练使用GPU Direct RDMA的时候，需要精细化管理显存，让每个RPC的数据写入指定的显存吗？还是直接从池子里取一块显存写入？

@yanglimingcn @chenBright hi，如果我们attachment里存放显存地址，在CutFromIOBufList中解析这个attachment，这是否修改难度较小

Sep 24 '25 14:09 sunce4t

@sunce4t 请教一下，训练使用GPU Direct RDMA的时候，需要精细化管理显存，让每个RPC的数据写入指定的显存吗？还是直接从池子里取一块显存写入？

@yanglimingcn @chenBright hi，如果我们attachment里存放显存地址，在CutFromIOBufList中解析这个attachment，这是否修改难度较小

CutFromIOBufList是在发送消息的时候使用的，解析这块在ProcessRpcRequest这里面。你是打算attachment都直接使用显存的内存吗？

Sep 25 '25 03:09 yanglimingcn

@sunce4t 请教一下，训练使用GPU Direct RDMA的时候，需要精细化管理显存，让每个RPC的数据写入指定的显存吗？还是直接从池子里取一块显存写入？

@yanglimingcn @chenBright hi，如果我们attachment里存放显存地址，在CutFromIOBufList中解析这个attachment，这是否修改难度较小

CutFromIOBufList是在发送消息的时候使用的，解析这块在ProcessRpcRequest这里面。你是打算attachment都直接使用显存的内存吗？

是的，下面是一个例子：

void * device_ptr = InitGPUMemoryForRDMA(attachment_size);
              mr = ibv_reg_mr(brpc::rdma::GetRdmaPd(), device_ptr, attachment_size,
                    IBV_ACCESS_LOCAL_WRITE |
                    IBV_ACCESS_REMOTE_ATOMIC |
                    IBV_ACCESS_REMOTE_READ |
                    IBV_ACCESS_REMOTE_WRITE);
            _attachment.append_user_data_with_meta(device_ptr, attachment_size, [](void*){}, mr->lkey);

在发送方，我们将attachment内的数据替换成显存中的数据，接收方postrecv时也可以类似这样去替换成显存现在问题是收到消息进入ProcessNewMessage后，这里需要加一些从显存拷贝协议头部+meta+body 到内存的逻辑我看了ProcessNewMessage这个函数，里面有一些显式copy的逻辑，也有使用cutn的零拷贝的逻辑，这些地方感觉都需要修改。

Sep 25 '25 09:09 sunce4t

attachment，在接收端，你希望直接用显存来接收吗？

Sep 25 '25 09:09 yanglimingcn

attachment，在接收端，你希望直接用显存来接收吗？

是的，attachment其实就是训练过程所需要的一些数据，因此我们希望直接从一台机器的显存传输到另一台机器的显存上。

Sep 25 '25 09:09 sunce4t

attachment，在接收端，你希望直接用显存来接收吗？

是的，attachment其实就是训练过程所需要的一些数据，因此我们希望直接从一台机器的显存传输到另一台机器的显存上。

现在brpc协议，meta部分的长度不固定，感觉不太好搞，这种我觉得自定义一个协议会比较好，比如： 1、头部长度在消息里面，可以直接截取消息头。 2、剩下的是attachment可以直接申请显存的内存，接收消息。

Sep 26 '25 07:09 yanglimingcn

attachment，在接收端，你希望直接用显存来接收吗？

是的，attachment其实就是训练过程所需要的一些数据，因此我们希望直接从一台机器的显存传输到另一台机器的显存上。

现在brpc协议，meta部分的长度不固定，感觉不太好搞，这种我觉得自定义一个协议会比较好，比如： 1、头部长度在消息里面，可以直接截取消息头。 2、剩下的是attachment可以直接申请显存的内存，接收消息。

好的，我们尝试一下。

Sep 28 '25 07:09 sunce4t

1.修改block_pool.cpp，新增一个ObjectPool用于分配Block的实体，将Block的data指向BlockPool -> 即 Block的实体部分与数据部分分离在此基础上，修改IOBuf的create block接口，objmem_allocate分配一个block出来，blockmem_allocate分配一段内存用于传输RPC请求

2.增加RDMA Write/Read接口；原本brpc的Send/Recv用于传输控制数据（例如，对端GPU地址），再使用RDMA Write直接将训练数据写到对端GPU地址上

3.使用append_with_user_data，将attachment从显存上分配，发送时使用RDMA的scatter-gatther特性，使其request请求写入到主机内存，attachment写入到显存上。但是此接口也是会在分配出来的内存上进行placement new，因此也存在core dumped的问题，也需要修改IOBuf的接口

@sunce4t 我觉得这三个方案里面只有第2个方案是可行的，iobuf的设计本来就是用来管理内存中的数据，而不是GPU显存的数据。你既然用了GDR那么数据包就没必要通过brpc来发送了，只要用brpc来传输控制信息就可以。当然你可以在brpc上面再封装一层，把brpc传输控制信息+GDR传输显存数据封装成一个接口方便业务使用。

Oct 12 '25 06:10 wwbmmm

1.修改block_pool.cpp，新增一个ObjectPool用于分配Block的实体，将Block的data指向BlockPool -> 即 Block的实体部分与数据部分分离在此基础上，修改IOBuf的create block接口，objmem_allocate分配一个block出来，blockmem_allocate分配一段内存用于传输RPC请求 2.增加RDMA Write/Read接口；原本brpc的Send/Recv用于传输控制数据（例如，对端GPU地址），再使用RDMA Write直接将训练数据写到对端GPU地址上 3.使用append_with_user_data，将attachment从显存上分配，发送时使用RDMA的scatter-gatther特性，使其request请求写入到主机内存，attachment写入到显存上。但是此接口也是会在分配出来的内存上进行placement new，因此也存在core dumped的问题，也需要修改IOBuf的接口

@sunce4t 我觉得这三个方案里面只有第2个方案是可行的，iobuf的设计本来就是用来管理内存中的数据，而不是GPU显存的数据。你既然用了GDR那么数据包就没必要通过brpc来发送了，只要用brpc来传输控制信息就可以。当然你可以在brpc上面再封装一层，把brpc传输控制信息+GDR传输显存数据封装成一个接口方便业务使用。

hi，我们暂时没有使用第二个方案，而是在第三个方案基础上进行尝试了

Oct 13 '25 08:10 sunce4t

1.修改block_pool.cpp，新增一个ObjectPool用于分配Block的实体，将Block的data指向BlockPool -> 即 Block的实体部分与数据部分分离在此基础上，修改IOBuf的create block接口，objmem_allocate分配一个block出来，blockmem_allocate分配一段内存用于传输RPC请求 2.增加RDMA Write/Read接口；原本brpc的Send/Recv用于传输控制数据（例如，对端GPU地址），再使用RDMA Write直接将训练数据写到对端GPU地址上 3.使用append_with_user_data，将attachment从显存上分配，发送时使用RDMA的scatter-gatther特性，使其request请求写入到主机内存，attachment写入到显存上。但是此接口也是会在分配出来的内存上进行placement new，因此也存在core dumped的问题，也需要修改IOBuf的接口

@sunce4t 我觉得这三个方案里面只有第2个方案是可行的，iobuf的设计本来就是用来管理内存中的数据，而不是GPU显存的数据。你既然用了GDR那么数据包就没必要通过brpc来发送了，只要用brpc来传输控制信息就可以。当然你可以在brpc上面再封装一层，把brpc传输控制信息+GDR传输显存数据封装成一个接口方便业务使用。

hi，社区考虑加一个使用GDR的样例吗 @wwbmmm @yanglimingcn @chenBright

Oct 15 '25 05:10 sunce4t

社区考虑加一个使用GDR的样例吗

@sunce4t 欢迎贡献GDR样例代码。

Oct 22 '25 02:10 chenBright

请教下这个需要区分Nv的卡型或者架构么？需要gdrcopy的库支持么？

Nov 09 '25 10:11 Huixxi

请教下这个需要区分Nv的卡型或者架构么？需要gdrcopy的库支持么？

GDR本身没有依赖卡型和架构；cuda版本和ofed版本都需要特定版本及以上（通常是满足的，不必特定去安装）；gdrcopy我们目前有使用，加速D2H拷贝，也可以不用，直接使用cudaMemCpy；目前我们这边正在整理代码，先提供一版不使用gdrcopy的；

需要注意的就是：在H系列的卡之前，用GDR的话，不能保证RDMA写入立刻对GPU可见，需要一次额外的访问（https://github.com/NVIDIA/nccl/issues/1702）

Nov 10 '25 04:11 sunce4t