Skip to content

support gpu direct rdma - #3144

Open
randomkang wants to merge 1 commit into
apache:masterfrom
randomkang:gdr
Open

randomkang wants to merge 1 commit into
apache:masterfrom
randomkang:gdr

Conversation

@randomkang

@randomkang randomkang commented Nov 9, 2025 •

Copy link
Copy Markdown
Contributor

What problem does this PR solve?

Issue Number: resolve #3102

Problem Summary:

What is changed and the side effects?

Changed:

  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, we will prefetch 512B to memory

Side effects:

  • Performance effects:

  • Breaking backward compatibility:


Check List:

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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.

Comment thread src/butil/iobuf.cpp Outdated
Comment thread src/brpc/rdma/rdma_endpoint.cpp Outdated
Comment thread src/brpc/rdma/rdma_endpoint.cpp Outdated
Comment thread src/brpc/rdma/rdma_helper.cpp Outdated
Comment thread src/butil/gpu/gpu_block_pool.h Outdated
Comment thread src/butil/iobuf.cpp Outdated
Comment on lines +763 to +764
clock_gettime(CLOCK_MONOTONIC, &end);
double time_us = (end.tv_sec - start.tv_sec) * 1e6 + (end.tv_nsec - start.tv_nsec) / 1e3;

Copilot AI Nov 9, 2025

Copy link

Choose a reason for hiding this comment

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

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 uses AI. Check for mistakes.
Comment thread src/butil/iobuf.cpp Outdated
Comment on lines +1418 to +1422
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 AI Nov 9, 2025

Copy link

Choose a reason for hiding this comment

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

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 uses AI. Check for mistakes.
Comment thread src/butil/gpu/gpu_block_pool.cpp Outdated
Comment thread src/butil/iobuf.cpp Outdated
Comment on lines +1394 to +1395
struct timespec start, end;
clock_gettime(CLOCK_MONOTONIC, &start);

Copilot AI Nov 9, 2025

Copy link

Choose a reason for hiding this comment

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

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 uses AI. Check for mistakes.
Comment thread src/butil/iobuf.cpp Outdated
Comment on lines +744 to +745
struct timespec start, end;
clock_gettime(CLOCK_MONOTONIC, &start);

Copilot AI Nov 9, 2025

Copy link

Choose a reason for hiding this comment

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

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 uses AI. Check for mistakes.
Comment thread src/brpc/policy/baidu_rpc_protocol.cpp Outdated

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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.

@randomkang randomkang Nov 14, 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.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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.

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.

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?

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

Copy link
Copy Markdown
Contributor

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

Copy link
Copy Markdown
Contributor Author

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

Copy link
Copy Markdown
Contributor Author

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

Comment thread docs/cn/gdr.md Outdated
Comment thread docs/cn/gdr.md Outdated
Comment thread src/brpc/acceptor.h Outdated
Comment thread src/butil/gpu/gpu_block_pool.cpp Outdated
@yanglimingcn

Copy link
Copy Markdown
Contributor

Turn multiple commits into one

@randomkang

Copy link
Copy Markdown
Contributor Author

Turn multiple commits into one

ok

@yanglimingcn

Copy link
Copy Markdown
Contributor

LGTM

return (void *)d_data;
}

void* get_cpu_mem(int gpu_id, int64_t cpu_mem_size) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

这个入参,应该是gpu_id吗

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.

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

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

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

BlockPoolAllocators* BlockPoolAllocators::singleton() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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

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.

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

Comment thread src/butil/iobuf.cpp Outdated

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

这里返回nullptr

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.

OK

@randomkang

Copy link
Copy Markdown
Contributor Author

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

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

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Comment thread BUILD.bazel
Comment on lines +112 to +115
"//bazel/config:brpc_with_gdr": [
"-lcuda",
"-lcudart",
],

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.

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.

Comment thread src/brpc/gdr_transport.cpp
Comment thread src/brpc/policy/baidu_rpc_protocol.cpp Outdated
Comment thread src/brpc/policy/baidu_rpc_protocol_gpu.cpp Outdated
Comment thread src/brpc/rdma/rdma_helper.cpp
Comment thread example/rdma_performance/client.cpp Outdated
Comment thread src/brpc/rdma/rdma_endpoint.cpp Outdated
Comment thread src/butil/gpu/gpu_block_pool.cpp Outdated
Comment thread src/brpc/acceptor.h Outdated
Comment on lines +73 to +74
ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket,
bool /*read_eof*/, const void*) {
Comment on lines +29 to +31
#include "butil/raw_pack.h" // RawPacker RawUnpacker
#include "butil/memory/scope_guard.h"
#include "butil/raw_pack.h" // RawPacker RawUnpacker

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

butil/raw_pack.h is duplicated.

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.

ok

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

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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
1 resolved since last review

🧠 Review effort: Lite


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

Comment thread example/rdma_performance/client.cpp Outdated
Comment thread src/brpc/rdma_transport.cpp
Comment thread src/brpc/policy/baidu_rpc_protocol_gpu.cpp Outdated
Comment thread src/brpc/policy/baidu_rpc_protocol_gpu.cpp
Comment on lines +171 to +174
bool is_gpu_memory = msg->payload.is_gpu_memory();
if (!is_gpu_memory) {
LOG(FATAL) << "message is not on gpu!!!";
}

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.

We expect data is on gpu here, so the fatal is necessary and can help to find errors.

Comment thread src/brpc/rdma/rdma_endpoint.cpp Outdated
Comment thread src/butil/gpu/gpu_block_pool.cpp
Comment on lines +447 to +457
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

wwbmmm commented Oct 8, 2026

Copy link
Copy Markdown
Contributor

please resolve merge conflict

@randomkang

Copy link
Copy Markdown
Contributor Author

please resolve merge conflict

ok

@wwbmmm wwbmmm left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[replied by brpc-oncall robot] The latest push fixes most earlier findings; three issues remain on the newly changed lines: the reverted socket-mode comment, MLX5_SCATTER_TO_CQE being set after the RDMA context is already open, and an early return that leaves the transport partially initialized.

Comment thread src/brpc/acceptor.h Outdated
std::shared_ptr<SocketSSLContext> _ssl_ctx;

// Choose to use a certain socket: 0 TCP, 1 RDMA
// Choose to use a certain socket: 0 TCP, 1 RDMA, 2 UBRING, 3 GDR

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[replied by brpc-oncall robot] The merge-conflict resolution reverted this comment to 2 UBRING, 3 GDR, but the enum is SOCKET_MODE_URMA = 3, SOCKET_MODE_GDR = 4 (see src/brpc/socket_mode.h). Please restore 0 TCP, 1 RDMA, 2 UBRING, 3 URMA, 4 GDR so it stays in sync with the enum.

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.

ok

Comment thread src/brpc/rdma/rdma_helper.cpp Outdated
// is enabled, small messages are inlined directly into host memory, which
// breaks this requirement and causes operations to fail. Force it to 0
// regardless of the user-provided value.
if (setenv("MLX5_SCATTER_TO_CQE", "0", 1) != 0) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[replied by brpc-oncall robot] MLX5_SCATTER_TO_CQE is set here, i.e. after the RDMA context has already been created: TransportFactory::ContextInitOrDie calls RdmaTransport::ContextInitOrDie() (which calls GlobalRdmaInitializeOrDie() → ibv_open_device) before GdrContextInitOrDie(), and the example also calls GlobalRdmaInitializeOrDie() first. libmlx5 reads the MLX5_* environment variables when the device context is created, so setting it after ibv_open_device has no effect. Please move it before the RDMA global initialization (or drop it and keep the value as an externally documented prerequisite).

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.

ok

Comment thread src/brpc/rdma_transport.cpp Outdated
_rdma_state = RDMA_OFF;
_socket = socket;
_default_connect = options.app_connect;
return;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[replied by brpc-oncall robot] This early return happens before _on_edge_trigger is set and before _tcp_transport is created/initialized, while _rdma_state is set to RDMA_OFF. RdmaTransport::CutFromIOBuf/CutFromIOBufList/WaitEpollOut dereference _tcp_transport unconditionally, so any later call on this transport would crash. Even though the socket was just SetFailed, please leave the transport in a fully initialized state (e.g. fall through to the common tail of DoInit) rather than returning in the middle.

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.

ok

@randomkang

Copy link
Copy Markdown
Contributor Author

Is there other problem for this pr? @wwbmmm @chenBright

@chenBright
chenBright requested a balanced review from Copilot October 20, 2026 16:18

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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

23 open issues

And 3 more that still need to be addressed.

Review effort: Lite

Comment thread example/rdma_performance/client.cpp
Comment thread src/brpc/policy/baidu_rpc_protocol_gpu.cpp Outdated
Comment on lines +482 to +494
#if BRPC_WITH_GDR
// GPUDirect RDMA requires all received data to be written into the GPU
// memory buffer specified at post-receive time. When MLX5_SCATTER_TO_CQE
// is enabled, small messages are inlined directly into host memory, which
// breaks this requirement and causes operations to fail. libmlx5 reads
// MLX5_* environment variables when the device context is created (i.e.
// inside ibv_open_device below), so it must be forced to 0 before any
// device is opened, regardless of the user-provided value.
if (setenv("MLX5_SCATTER_TO_CQE", "0", 1) != 0) {
PLOG(ERROR) << "Fail to set MLX5_SCATTER_TO_CQE to 0";
ExitWithError();
}
#endif // if BRPC_WITH_GDR

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.

If user do not want to use gdr, then should not build brpc with gdr, and this environment variable setting will not take effect.

@randomkang

Copy link
Copy Markdown
Contributor Author

All Copilot issues have been resolved or commented. @chenBright @wwbmmm

@wwbmmm

wwbmmm commented Oct 25, 2026

Copy link
Copy Markdown
Contributor

please resolve merge conflicts

@wwbmmm wwbmmm left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[replied by brpc-oncall robot] Most previous findings are now correctly addressed (prefetch size derived from the CPU pool block size, bool-returning GPU buffer fills with error propagation, cut/append failure checks, DEBUG-only timing, GetGPUIndex in the example). Two items remain: the GDR init ordering in the example client (below) and the still-unaddressed IOBuf::copy_from_gpu() device-memory validation; the Bazel GDR build also still relies on ambient C_INCLUDE_PATH/LD_LIBRARY_PATH rather than declaring the CUDA dependency.

}
#ifdef BRPC_WITH_GDR
if (FLAGS_use_gdr) {
brpc::rdma::GlobalGdrInitializeOrDie();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[replied by brpc-oncall robot] With the else if turned into a plain if, --use_rdma=false --use_gdr=true now calls GlobalGdrInitializeOrDie() while RDMA is still uninitialized. GlobalGdrInitializeOrDieImpl passes GetRdmaPd() to InitGPUBlockPool, but g_pd is only set by GlobalRdmaInitializeOrDieImpl, so it is still nullptr: ibv_reg_mr(nullptr, ...) fails and BlockPoolAllocator hits its LOG(FATAL). Please make GDR imply the RDMA initialization here, e.g. initialize RDMA when either flag is set:

if (FLAGS_use_rdma || FLAGS_use_gdr) {
    brpc::rdma::GlobalRdmaInitializeOrDie();
}
#ifdef BRPC_WITH_GDR
if (FLAGS_use_gdr) {
    brpc::rdma::GlobalGdrInitializeOrDie();
}
#endif

Alternatively drop both explicit calls and rely on TransportFactory::ContextInitOrDie, as server.cpp already does.

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.

when gdr is open, rdma must be open.Flag validation has been added.

@randomkang

Copy link
Copy Markdown
Contributor Author

please resolve merge conflicts

Done

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

brpc支持GPU Direct RDMA

6 participants