support gpu direct rdma by randomkang · Pull Request #3144 · apache/brpc

randomkang · 2025-11-09T02:51:10Z

What problem does this PR solve?

Issue Number: resolve #3102

Problem Summary:

What is changed and the side effects?

Changed:

recv all data on gpu first
the gpu block is alloced from a gpu block pool
brpc header, meta and body will be copied from gpu to cpu to process.
To decrease the d2h counts, we will prefetch 512B to memory

Side effects:

Performance effects:
Breaking backward compatibility:

Check List:

Please make sure your changes are compilable.
When providing us with a new feature, it is best to add related tests.
Please follow Contributor Covenant Code of Conduct.

Copilot

Pull Request Overview

This PR adds GPU Direct RDMA (GDR) support to the BRPC framework, enabling efficient data transfer between GPU memory and RDMA-capable network devices. The implementation includes GPU memory pool management, CUDA stream pooling for asynchronous memory operations, and protocol-level changes to handle GPU-resident data.

Added GPU memory block pool allocator with configurable block sizes
Implemented CUDA stream pooling for optimized device-to-host and device-to-device transfers
Integrated GDR support into IOBuf, RDMA endpoint, and RPC protocol parsing

Reviewed Changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
src/butil/iobuf.h	Added GPU-specific copy and cut methods with conditional compilation guards
src/butil/iobuf.cpp	Implemented `cutn_from_gpu` and `copy_from_gpu` methods for GPU memory operations
src/butil/gpu/gpu_block_pool.h	New header defining GPU memory pool allocator, stream pool, and region management
src/butil/gpu/gpu_block_pool.cpp	Implementation of GPU memory allocation, CUDA stream management, and memory copy operations
src/brpc/rdma/rdma_helper.h	Added function declaration for GPU index retrieval
src/brpc/rdma/rdma_helper.cpp	Added GPU index configuration and initialization logic for GDR block pool
src/brpc/rdma/rdma_endpoint.h	Added remote receive window tracking and GDR-specific receive posting method
src/brpc/rdma/rdma_endpoint.cpp	Modified flow control, completion handling, and receive buffer posting for GDR support
src/brpc/policy/baidu_rpc_protocol.cpp	Enhanced message parsing to detect and handle GPU memory with optimized prefetch
bazel/config/BUILD.bazel	Added build configuration for GDR feature flag
BUILD.bazel	Added CUDA dependencies and GPU block pool source to build

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2025-11-09T08:37:24Z

+    clock_gettime(CLOCK_MONOTONIC, &end);
+    double time_us = (end.tv_sec - start.tv_sec) * 1e6 + (end.tv_nsec - start.tv_nsec) / 1e3;


The variables start, end, and time_us are defined and computed but the result is only used in a commented-out log statement (lines 766-767). Consider removing this timing code if it's not being used, or uncomment the logging if it's needed for debugging.

Copilot · 2025-11-09T08:37:25Z

+    double time_us = (end.tv_sec - start.tv_sec) * 1e6 + (end.tv_nsec - start.tv_nsec) / 1e3;
+    size_t copied_bytes = n - m;
+
+    // LOG(INFO) << "GDRCopy: " << copied_bytes << " bytes, "
+    //           << time_us << " us" << ", to_gpu " << to_gpu;


The variables time_us and copied_bytes are computed but only used in commented-out log statements (lines 1421-1422). Consider removing this unused code or uncommenting the logging if it's needed for debugging.

Suggested change

double time_us = (end.tv_sec - start.tv_sec) * 1e6 + (end.tv_nsec - start.tv_nsec) / 1e3;

size_t copied_bytes = n - m;

// LOG(INFO) << "GDRCopy: " << copied_bytes << " bytes, "

// << time_us << " us" << ", to_gpu " << to_gpu;

Copilot · 2025-11-09T08:37:25Z

+    struct timespec start, end;
+    clock_gettime(CLOCK_MONOTONIC, &start);


The timing variables start and end are declared and used but the computed time_us value is only for commented-out logging. If timing is not actively needed, consider removing this instrumentation code.

Copilot · 2025-11-09T08:37:26Z

+    struct timespec start, end;
+    clock_gettime(CLOCK_MONOTONIC, &start);


The timing variables start and end are declared and used but the computed time_us value is only for commented-out logging. If timing is not actively needed, consider removing this instrumentation code.

yanglimingcn · 2025-11-10T01:25:08Z

+
    char header_buf[12];
-    const size_t n = source->copy_to(header_buf, sizeof(header_buf));
+    size_t n = 0;


I think it would be better to put the baidu_rpc_protocol.cpp file containing GPU functionality into a separate file called baidu_rpc_with_gpu_protocol.cpp, or to define a new protocol. Currently, the macro definitions generate too many if-else statements.

Maybe put GPU functionality into a separate file or define a new protocol is two heavy, we reorg the code to decrease the macro definitions. Is it ok? @yanglimingcn

I think it's better to separate new protocol file for GDR feature， Besides, I have a idea we can also define a new protocol meta that includes both request meta and response meta carrying the GDR attribute value., the corresponding meta should also contain the GPU address information that both sides need to register with IB. Doing the GDR operations inside processRequest and ProcessRpcResponse method step will probably be more reasonable.

If the server or client wants to know the GPU memory region of the other side, we can send this information with a separate rpc transport. But i do not see the the benefits of this approach. Can you explain more?

yanglimingcn · 2025-12-02T07:31:08Z

I believe the long-term solution is to have users register memory with RDMA, allowing users to customize this memory based on their data organization methods.
https://zhuanlan.zhihu.com/p/376989325 This link contains some relevant details.

randomkang · 2025-12-14T12:53:32Z

I believe the long-term solution is to have users register memory with RDMA, allowing users to customize this memory based on their data organization methods. https://zhuanlan.zhihu.com/p/376989325 This link contains some relevant details.

In this pr, brpc will recv data into fragmented gpu blocks and user must call IOBuf::copy_from_gpu to copy these gpu blocks into a continous hbm in order to use it. The d2d copy is time-consuming and not necessary.

In feature, we can let the user assign the gpu destination directly with interface like "rdma_memory_pool_user_specified_memory" and recv the data into it with the opcode like IBV_WR_RDMA_READ/IBV_WR_RDMA_WRITE. Then we can skip the d2d copy.

Futhermore, we can parse the brpc protocol with gpu kernel and initiate RDMA communication with GPU. Then the control path and the data path both are on gpu like nccl gin.

randomkang · 2026-08-04T05:51:55Z

@yanglimingcn @zchuango all comments have been resolved, please review again.

yanglimingcn · 2026-08-25T11:46:03Z

Turn multiple commits into one

randomkang · 2026-08-25T12:59:31Z

Turn multiple commits into one

ok

yanglimingcn · 2026-08-27T11:25:48Z

LGTM

yanglimingcn · 2026-08-28T02:20:34Z

+    return (void *)d_data;
+}
+
+void* get_cpu_mem(int gpu_id, int64_t cpu_mem_size) {


这个入参，应该是gpu_id吗

是的。这里申请的是pinned memory，使用pinned memory和显存进行交换，会比普通的memory快。在使用cudaMallocHost分配pinned memory的时候，建议先执行cudaSetDevice，这样分配的pined memory就会与对应的gpu id关联。

yanglimingcn · 2026-08-28T02:22:48Z

+
+std::atomic<BlockPoolAllocators*> BlockPoolAllocators::instance_{nullptr};
+
+BlockPoolAllocators* BlockPoolAllocators::singleton() {


你看brpc自带的singleton.h是否能满足你的需求。

我看了下，可以用brpc自带的Singleton，已改

yanglimingcn · 2026-08-28T02:31:53Z


+void* IOBuf::get_first_data_ptr() {
+    if (_ref_num() == 0) {
+        return 0;


这里返回nullptr

randomkang · 2026-10-08T03:49:00Z

@yanglimingcn @zchuango Are there any other issues with this pr?

Copilot

🟡 Changes recommended

Critical build, handshake, parsing, allocation, and memory-pool correctness issues remain unresolved.

13 open findings

Bazel GDR build lacks CUDA include and dependency configuration · New
Dispatch GDR connections through the RDMA handshake · New
Stop response processing after GPU buffer fill failure · New
Prevent GPU parse errors from falling through to non-GDR parsers · New
Validate GDR block size before initializing the GPU pool · New
Propagate CUDA allocation failures instead of aborting · New
Synchronize GPU memory region lookup and publication · New
Use value-based BRPC_WITH_GDR guards in iobuf.cpp · New
Allocate GDR attachment on the configured GPU · New
Handle GPU IOBuf metadata allocation failure · New
Select the configured CUDA device in each GPU operation · New
Correct socket mode enum values in documentation · New
Add tests for GPU prefetch and malformed attachment handling · New

🧠 Review effort: Lite

Give feedback about Copilot approvals in this survey to enter a drawing for a $150 gift card.

randomkang · 2026-10-08T07:03:13Z

+    "//bazel/config:brpc_with_gdr": [
+        "-lcuda",
+        "-lcudart",
+    ],


We do not use any cuda kernel in this pr, so CUDA_HOME is not necessary. The cuda headers and libs can be found by C_INCLUDE_PATH and LD_LIBRARY_PATH.

+namespace brpc {
+
+void GdrTransport::Init(Socket *socket, const SocketOptions &options) {
+    DoInit(socket, options, true);


+#if BRPC_WITH_GDR
+        bool is_gpu_memory = msg->payload.is_gpu_memory();
+        if (is_gpu_memory) {
+            FillResBufGpu(&res_buf, msg.get(), meta, &res_buf_ptr, cntl);


+    do {
+        if (n >= 4) {
+            void* dummy = header_buf;
+            if (*(const uint32_t*)dummy != *(const uint32_t*)"PRPC") {
+                pe = PARSE_ERROR_TRY_OTHERS;
+                break;
+            }
+        } else {
+            if (memcmp(header_buf, "PRPC", n) != 0) {
+                pe = PARSE_ERROR_TRY_OTHERS;


+    if (!butil::gdr::InitGPUBlockPool(g_gpu_index, GetRdmaPd())) {
+        PLOG(ERROR) << "Fail to initialize RDMA GPU memory pool";
+        ExitWithError();
+    }


+                int gpu_id = 0;
+                cudaSetDevice(gpu_id);
+                cudaMalloc(&_addr, attachment_size);
+                auto pd = brpc::rdma::GetRdmaPd();
+                mr = ibv_reg_mr(pd, _addr, attachment_size,


+                _rbuf[_rq_received].append_user_data_gpu(device_ptr, g_gdr_recv_block_size, deleter, lkey);
+                _rbuf_data[_rq_received] = device_ptr;


+void GPUStreamPool::fast_d2d(std::vector<void*>& src_list,
+                                  std::vector<int64_t>& length_list,
+                                  void* dst) {
+#ifdef DEBUG
+    if (!verify_same_context()) {
+        LOG(FATAL) << "Context mismatch!";
+        return;
+    }
+#endif


+ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket,
+                            bool /*read_eof*/, const void*) {


chenBright · 2026-10-08T04:36:11Z

+#include "butil/raw_pack.h"                      // RawPacker RawUnpacker
+#include "butil/memory/scope_guard.h"
+#include "butil/raw_pack.h"                      // RawPacker RawUnpacker


butil/raw_pack.h is duplicated.

Copilot

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Copilot review overview

20 open findings

With the current if/else if, GlobalGdrInitializeOrDie() won’t run when both --use_rdma=true… · New
If new(std::nothrow) fails, _rdma_ep stays null but _rdma_state is still set to… · New
Use value-based BRPC_WITH_GDR guards in iobuf.cpp
Synchronize GPU memory region lookup and publication
Propagate CUDA allocation failures instead of aborting
Validate GDR block size before initializing the GPU pool
Prevent GPU parse errors from falling through to non-GDR parsers
Stop response processing after GPU buffer fill failure
Dispatch GDR connections through the RDMA handshake
Bazel GDR build lacks CUDA include and dependency configuration
Using LOG(FATAL) in the protocol parser will terminate the whole process on unexpected… · New
This uses type-punning and potentially unaligned uint32_t loads from header_buf/string literal,… · New
Same issue as the parser: LOG(FATAL) here can crash the server if the message is unexpectedly… · New
This PR changes the default rdma_prepared_qp_cnt from 1024 to 256, which is a… · New
The allocation fast-path always takes two timestamps and computes duration, even in non-DEBUG… · New
The initial cudaStreamSynchronize() before enqueueing copies is redundant if you already… · New
Select the configured CUDA device in each GPU operation
Handle GPU IOBuf metadata allocation failure
Allocate GDR attachment on the configured GPU
Add tests for GPU prefetch and malformed attachment handling

1 resolved since last review

Correct socket mode enum values in documentation

🧠 Review effort: Lite

Give feedback about Copilot approvals in this survey to enter a drawing for a $150 gift card.

    if (FLAGS_use_rdma) {
        brpc::rdma::GlobalRdmaInitializeOrDie();
    }
+#ifdef BRPC_WITH_GDR
+    else if (FLAGS_use_gdr) {
+        brpc::rdma::GlobalRdmaInitializeOrDie();


+void RdmaTransport::DoInit(Socket *socket, const SocketOptions &options, bool use_gdr) {
    CHECK(_rdma_ep == nullptr);
-    if (options.socket_mode == SOCKET_MODE_RDMA) {
-        _rdma_ep = new rdma::RdmaEndpoint(socket);
+    // gdr mode is a special mode of rdma mode.
+    // both rdma mode and gdr mode need init rdma::RdmaEndpoint.
+    if (options.socket_mode == SOCKET_MODE_RDMA ||
+            options.socket_mode == SOCKET_MODE_GDR) {
+        _rdma_ep = new(std::nothrow)rdma::RdmaEndpoint(socket, use_gdr);
+        if (!_rdma_ep) {
+            const int saved_errno = errno;
+            PLOG(ERROR) << "Fail to create RdmaEndpoint";
+            socket->SetFailed(
+                saved_errno, "Fail to create RdmaEndpoint: %s", berror(saved_errno));
+        }
        _rdma_state = RDMA_UNKNOWN;
    } else {
        _rdma_state = RDMA_OFF;


+    void* prefetch_d2h_data = NULL;
+    bool is_gpu_memory = source->is_gpu_memory();
+    if (!is_gpu_memory) {
+        LOG(FATAL) << "RpcMessage is not in gpu!!!";


+        if (n >= 4) {
+            void* dummy = header_buf;
+            if (*(const uint32_t*)dummy != *(const uint32_t*)"PRPC") {
+                pe = PARSE_ERROR_TRY_OTHERS;
+                break;
+            }
+        } else {


+    bool is_gpu_memory = msg->payload.is_gpu_memory();
+    if (!is_gpu_memory) {
+        LOG(FATAL) << "message is not on gpu!!!";
+    }


 DEFINE_int32(rdma_cqe_poll_once, 32, "The maximum of cqe number polled once.");
 DEFINE_int32(rdma_prepared_qp_size, 128, "SQ and RQ size for prepared QP.");
-DEFINE_int32(rdma_prepared_qp_cnt, 1024, "Initial count of prepared QP.");
+DEFINE_int32(rdma_prepared_qp_cnt, 256, "Initial count of prepared QP.");


+    auto startTime = std::chrono::high_resolution_clock::now();
+
+    std::lock_guard<std::mutex> lock(poolMutex);
+
+    if (!freeList) {
+        if (extendRegion() != 0) {
+            // Out of regions or registration failed; let the caller handle it.
+            errno = ENOMEM;
+            return nullptr;
+        }
+    }
+
+    BlockHeader* block = freeList;
+    freeList = freeList->next;
+
+    void* addr = block->addr;
+    get_bh_list()->Release(block);
+
+    totalAllocated++;
+    peakUsage = std::max(peakUsage, totalAllocated - totalDeallocated);
+
+    auto endTime = std::chrono::high_resolution_clock::now();
+    auto duration = std::chrono::duration_cast<std::chrono::nanoseconds>(endTime - startTime);
+
+#ifdef DEBUG
+    if (duration.count() > 1000) {
+        LOG(INFO) << "Slow allocation: " << duration.count() << " ns";
+    }
+#endif


+    std::lock_guard<std::mutex> stream_lock(d2d_locks_[stream_idx]);
+    CHECK_CUDA(cudaStreamSynchronize(d2d_streams_[stream_idx]));
+    for (int i = 0; i < segs; i++) {
+        if (length_list[i] == 0) {
+            continue;
+        }
+        CHECK_CUDA(cudaMemcpyAsync(static_cast<char*>(dst) + offset, src_list[i], length_list[i],
+                    cudaMemcpyDeviceToDevice, d2d_streams_[stream_idx]));
+        offset += length_list[i];
+    }
+    CHECK_CUDA(cudaStreamSynchronize(d2d_streams_[stream_idx]));


wwbmmm · 2026-10-08T12:47:02Z

please resolve merge conflict

randomkang · 2026-10-11T01:59:33Z

please resolve merge conflict

ok

1) recv all data on gpu first 2) the gpu block is alloced from a gpu block pool 3) brpc header, meta and body will be copied from gpu to cpu to process. 4) To decrease the d2h counts when processing header, meta and body, we will prefetch 512B to memory Co-authored-by: sunce4t <847480001@qq.com>

wwbmmm requested review from Copilot and yanglimingcn November 9, 2025 08:34

Copilot AI reviewed Nov 9, 2025

View reviewed changes

yanglimingcn reviewed Nov 10, 2025

View reviewed changes

Comment thread src/brpc/rdma/rdma_endpoint.cpp Outdated

sunce4t force-pushed the gdr branch from 6eedb3e to 1fb9370 Compare November 26, 2025 02:17

randomkang force-pushed the gdr branch 2 times, most recently from 041bef0 to 9e3b2af Compare June 2, 2026 16:41

randomkang force-pushed the gdr branch from 9e3b2af to 051d0b2 Compare August 3, 2026 17:11

yanglimingcn requested a lite review from Copilot August 25, 2026 10:00

Copilot started reviewing on behalf of yanglimingcn August 25, 2026 10:40 View session

Copilot stopped reviewing on behalf of yanglimingcn due to an error August 25, 2026 11:00
Copilot had to stop work due to a timeout.

yanglimingcn reviewed Aug 25, 2026

View reviewed changes

Comment thread docs/cn/gdr.md Outdated

yanglimingcn reviewed Aug 25, 2026

View reviewed changes

Comment thread docs/cn/gdr.md Outdated

yanglimingcn reviewed Aug 25, 2026

View reviewed changes

Comment thread src/brpc/acceptor.h Outdated

yanglimingcn reviewed Aug 25, 2026

View reviewed changes

Comment thread src/butil/gpu/gpu_block_pool.cpp Outdated

randomkang force-pushed the gdr branch from 2d74338 to bf328f5 Compare August 25, 2026 12:57

yanglimingcn reviewed Aug 28, 2026

View reviewed changes

randomkang force-pushed the gdr branch from bf328f5 to 11ee2f3 Compare August 28, 2026 08:50

randomkang force-pushed the gdr branch from 11ee2f3 to 28972ad Compare October 8, 2026 03:42

chenBright removed the request for review from Copilot October 8, 2026 04:39

chenBright requested review from zchuango and a lite review from Copilot and removed request for zchuango October 8, 2026 04:39

Copilot started reviewing on behalf of chenBright October 8, 2026 04:39 View session

Copilot AI reviewed Oct 8, 2026

View reviewed changes

chenBright reviewed Oct 8, 2026

View reviewed changes

wwbmmm requested a lite review from Copilot October 8, 2026 09:57

Copilot AI reviewed Oct 8, 2026

View reviewed changes

Copilot stopped reviewing on behalf of wwbmmm due to an error October 8, 2026 11:00
An unexpected error occurred. For more details, see the detailed logs in GitHub Actions.

randomkang force-pushed the gdr branch from 4ed2a6c to da3556c Compare October 11, 2026 01:56

randomkang force-pushed the gdr branch from f4041e5 to ed9c38c Compare October 11, 2026 15:16

		clock_gettime(CLOCK_MONOTONIC, &end);
		double time_us = (end.tv_sec - start.tv_sec) * 1e6 + (end.tv_nsec - start.tv_nsec) / 1e3;

		struct timespec start, end;
		clock_gettime(CLOCK_MONOTONIC, &start);


		std::atomic<BlockPoolAllocators*> BlockPoolAllocators::instance_{nullptr};

		BlockPoolAllocators* BlockPoolAllocators::singleton() {

		_rbuf[_rq_received].append_user_data_gpu(device_ptr, g_gdr_recv_block_size, deleter, lkey);
		_rbuf_data[_rq_received] = device_ptr;

		ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket,
		bool /read_eof/, const void*) {

Conversation

randomkang commented Nov 9, 2025 • edited by chenBright Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What problem does this PR solve?

What is changed and the side effects?

Check List:

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Nov 9, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 9, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Copilot AI Nov 9, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 9, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

randomkang Nov 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yanglimingcn commented Dec 2, 2025

Uh oh!

randomkang commented Dec 14, 2025

Uh oh!

randomkang commented Aug 4, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yanglimingcn commented Aug 25, 2026

Uh oh!

randomkang commented Aug 25, 2026

Uh oh!

yanglimingcn commented Aug 27, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

randomkang commented Oct 8, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

randomkang commented Nov 9, 2025 •

edited by chenBright

Loading

randomkang Nov 14, 2025 •

edited

Loading