Repository navigation
support gpu direct rdma - #3144
randomkang wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
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.
| clock_gettime(CLOCK_MONOTONIC, &end); | ||
| double time_us = (end.tv_sec - start.tv_sec) * 1e6 + (end.tv_nsec - start.tv_nsec) / 1e3; |
There was a problem hiding this comment.
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.
| 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; |
There was a problem hiding this comment.
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.
| 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; |
| struct timespec start, end; | ||
| clock_gettime(CLOCK_MONOTONIC, &start); |
There was a problem hiding this comment.
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.
| struct timespec start, end; | ||
| clock_gettime(CLOCK_MONOTONIC, &start); |
There was a problem hiding this comment.
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.
|
|
||
| char header_buf[12]; | ||
| const size_t n = source->copy_to(header_buf, sizeof(header_buf)); | ||
| size_t n = 0; |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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?
|
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. |
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. |
041bef0 to
9e3b2af
Compare
|
@yanglimingcn @zchuango all comments have been resolved, please review again. |
|
Turn multiple commits into one |
ok |
|
LGTM |
| return (void *)d_data; | ||
| } | ||
|
|
||
| void* get_cpu_mem(int gpu_id, int64_t cpu_mem_size) { |
There was a problem hiding this comment.
是的。这里申请的是pinned memory,使用pinned memory和显存进行交换,会比普通的memory快。在使用cudaMallocHost分配pinned memory的时候,建议先执行cudaSetDevice,这样分配的pined memory就会与对应的gpu id关联。
|
|
||
| std::atomic<BlockPoolAllocators*> BlockPoolAllocators::instance_{nullptr}; | ||
|
|
||
| BlockPoolAllocators* BlockPoolAllocators::singleton() { |
There was a problem hiding this comment.
你看brpc自带的singleton.h是否能满足你的需求。
There was a problem hiding this comment.
我看了下,可以用brpc自带的Singleton,已改
|
|
||
| void* IOBuf::get_first_data_ptr() { | ||
| if (_ref_num() == 0) { | ||
| return 0; |
|
@yanglimingcn @zchuango Are there any other issues with this pr? |
There was a problem hiding this comment.
🟡 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.
| "//bazel/config:brpc_with_gdr": [ | ||
| "-lcuda", | ||
| "-lcudart", | ||
| ], |
There was a problem hiding this comment.
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.
| ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket, | ||
| bool /*read_eof*/, const void*) { |
| #include "butil/raw_pack.h" // RawPacker RawUnpacker | ||
| #include "butil/memory/scope_guard.h" | ||
| #include "butil/raw_pack.h" // RawPacker RawUnpacker |
There was a problem hiding this comment.
butil/raw_pack.h is duplicated.
There was a problem hiding this comment.
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 currentif/else if,GlobalGdrInitializeOrDie()won’t run when both--use_rdma=true… · New Ifnew(std::nothrow)fails,_rdma_epstays null but_rdma_stateis 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 UsingLOG(FATAL)in the protocol parser will terminate the whole process on unexpected… · New This uses type-punning and potentially unaligneduint32_tloads fromheader_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 defaultrdma_prepared_qp_cntfrom 1024 to 256, which is a… · New The allocation fast-path always takes two timestamps and computesduration, even in non-DEBUG… · New The initialcudaStreamSynchronize()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
🧠 Review effort: Lite
Give feedback about Copilot approvals in this survey to enter a drawing for a $150 gift card.
| bool is_gpu_memory = msg->payload.is_gpu_memory(); | ||
| if (!is_gpu_memory) { | ||
| LOG(FATAL) << "message is not on gpu!!!"; | ||
| } |
There was a problem hiding this comment.
We expect data is on gpu here, so the fatal is necessary and can help to find errors.
| 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])); |
|
please resolve merge conflict |
ok |
wwbmmm
left a comment
There was a problem hiding this comment.
[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.
| 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 |
There was a problem hiding this comment.
[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.
| // 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) { |
There was a problem hiding this comment.
[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).
| _rdma_state = RDMA_OFF; | ||
| _socket = socket; | ||
| _default_connect = options.app_connect; | ||
| return; |
There was a problem hiding this comment.
[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.
|
Is there other problem for this pr? @wwbmmm @chenBright |
There was a problem hiding this comment.
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
When--use_gdr=trueand the default--use_rdma=true, the first branch runs and GDR is never… · New Ifnew(std::nothrow)fails,_rdma_epstays null but_rdma_stateis still set to… With the currentif/else if,GlobalGdrInitializeOrDie()won’t run when both--use_rdma=true… 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 configurationprefetch_d2h_sizeis hard-coded to 512 and implicitly tied to the CPU block pool block size (also… · New This forcesMLX5_SCATTER_TO_CQE=0for all processes built withBRPC_WITH_GDR, even when… · New The initialcudaStreamSynchronize()before enqueueing copies is redundant if you already… The allocation fast-path always takes two timestamps and computesduration, even in non-DEBUG… This PR changes the defaultrdma_prepared_qp_cntfrom 1024 to 256, which is a… Same issue as the parser:LOG(FATAL)here can crash the server if the message is unexpectedly… This uses type-punning and potentially unaligneduint32_tloads fromheader_buf/string literal,… UsingLOG(FATAL)in the protocol parser will terminate the whole process on unexpected… Select the configured CUDA device in each GPU operation
And 3 more that still need to be addressed.
Review effort: Lite
| #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 |
There was a problem hiding this comment.
If user do not want to use gdr, then should not build brpc with gdr, and this environment variable setting will not take effect.
|
All Copilot issues have been resolved or commented. @chenBright @wwbmmm |
|
please resolve merge conflicts |
wwbmmm
left a comment
There was a problem hiding this comment.
[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(); |
There was a problem hiding this comment.
[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();
}
#endifAlternatively drop both explicit calls and rely on TransportFactory::ContextInitOrDie, as server.cpp already does.
There was a problem hiding this comment.
when gdr is open, rdma must be open.Flag validation has been added.
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>



What problem does this PR solve?
Issue Number: resolve #3102
Problem Summary:
What is changed and the side effects?
Changed:
Side effects:
Performance effects:
Breaking backward compatibility:
Check List: