[TransferEngine] Add NCCL host and device transport backend#2852
[TransferEngine] Add NCCL host and device transport backend#2852akhillanger wants to merge 7 commits into
Conversation
Hide NCCL communicators and windows behind opaque Mooncake contexts, add lane-defined LSA/GIN helpers and capability routing, and provide collective-safe buffer allocation. Align NCCL 2.30.4 CMake discovery, CXI guards, lifecycle semantics, and run-scoped example bootstrap.
akhillanger
requested review from
alogfans,
chestnut-Q and
doujiang24
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces NCCL-based transport backends to the Mooncake transfer engine, adding support for both device-initiated (USE_NCCL_DEVICE) and host-submitted RMA (USE_NCCL_HOST) communication paths, along with corresponding examples and handshake protocol integration. The code review identified several critical issues in the new host transport implementation: a state desynchronization risk in unregisterMemory if metadata updates fail, a potential GPU memory safety issue if cudaEventRecord fails after a successful transfer submission, and multiple potential crashes or exceptions due to unvalidated JSON parsing during handshake and buffer catalog decoding.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| int unregisterMemory(void* addr, bool update_metadata) { | ||
| { | ||
| std::lock_guard<std::mutex> sessions_lock(sessions_mutex_); | ||
| if (!sessions_.empty() || !bootstrap_ids_.empty()) { | ||
| LOG(ERROR) << "[Host NCCL] cannot unregister memory while " | ||
| "NCCL sessions own collective windows"; | ||
| return ERR_INVALID_ARGUMENT; | ||
| } | ||
| } | ||
|
|
||
| bool found = false; | ||
| { | ||
| std::lock_guard<std::mutex> lock(buffers_mutex_); | ||
| auto it = std::find_if( | ||
| local_buffers_.begin(), local_buffers_.end(), | ||
| [addr](const BufferInfo& buffer) { | ||
| return buffer.addr == reinterpret_cast<uint64_t>(addr); | ||
| }); | ||
| if (it != local_buffers_.end()) { | ||
| local_buffers_.erase(it); | ||
| found = true; | ||
| } | ||
| } | ||
| if (!found) return ERR_ADDRESS_NOT_REGISTERED; | ||
| return metadata_->removeLocalMemoryBuffer(addr, update_metadata); | ||
| } |
There was a problem hiding this comment.
In unregisterMemory, the buffer is erased from local_buffers_ before calling metadata_->removeLocalMemoryBuffer. If removeLocalMemoryBuffer fails, the local state becomes desynchronized from the metadata store. We should only erase the buffer after a successful metadata update, or roll back the erasure on failure.
int unregisterMemory(void* addr, bool update_metadata) {
{
std::lock_guard<std::mutex> sessions_lock(sessions_mutex_);
if (!sessions_.empty() || !bootstrap_ids_.empty()) {
LOG(ERROR) << "[Host NCCL] cannot unregister memory while "
"NCCL sessions own collective windows";
return ERR_INVALID_ARGUMENT;
}
}
bool found = false;
BufferInfo removed_info;
{
std::lock_guard<std::mutex> lock(buffers_mutex_);
auto it = std::find_if(
local_buffers_.begin(), local_buffers_.end(),
[addr](const BufferInfo& buffer) {
return buffer.addr == reinterpret_cast<uint64_t>(addr);
});
if (it != local_buffers_.end()) {
removed_info = *it;
local_buffers_.erase(it);
found = true;
}
}
if (!found) return ERR_ADDRESS_NOT_REGISTERED;
int result = metadata_->removeLocalMemoryBuffer(addr, update_metadata);
if (result != 0) {
std::lock_guard<std::mutex> lock(buffers_mutex_);
local_buffers_.push_back(removed_info);
}
return result;There was a problem hiding this comment.
Addressed in 9f3cdfb. unregisterMemoryLocked now performs the protocol-scoped metadata removal before erasing the entry from local_buffers_, while the caller holds buffers_mutex_. The metadata side stages a cloned descriptor, publishes it first when requested, and commits the local metadata cache only after publication succeeds, so a lookup or publication failure leaves both catalogs unchanged. Batch unregister also prevalidates every address and uses one bulk transaction, making it all-or-nothing. I added coverage for failed atomic bulk removal, successful bulk removal, and same-address protocol-scoped removal.
Disclosure: This response was AI-generated.
| ncclResult_t result = ncclPutSignal( | ||
| source, length, ncclUint8, owner_rank, window->handle, | ||
| dest_addr - destination.addr, 0, 0, 0, comm_, stream_); | ||
| if (result == ncclSuccess && event) { | ||
| cuda_result = cudaEventRecord(event, stream_); | ||
| } | ||
| if (saved_device >= 0) cudaSetDevice(saved_device); |
There was a problem hiding this comment.
If ncclPutSignal succeeds but cudaEventRecord fails, the transfer is already submitted to the GPU stream, but the function returns an error. The caller might immediately free or overwrite the source buffer, leading to silent data corruption or use-after-free on the GPU. We should synchronously wait for the stream to finish if cudaEventRecord fails to ensure safety.
ncclResult_t result = ncclPutSignal(
source, length, ncclUint8, owner_rank, window->handle,
dest_addr - destination.addr, 0, 0, 0, comm_, stream_);
if (result == ncclSuccess && event) {
cuda_result = cudaEventRecord(event, stream_);
if (cuda_result != cudaSuccess) {
cudaStreamSynchronize(stream_);
}
}
if (saved_device >= 0) cudaSetDevice(saved_device);There was a problem hiding this comment.
Addressed in 9f3cdfb. After a successful ncclPutSignal, a cudaEventRecord failure now synchronizes stream_ before returning the error. This drains the already-queued put before the caller can reuse or free its source buffer. The original event-record error is preserved; if the fallback synchronization also fails, that error is appended and the session is marked failed. The normal WRITE and reverse-PUT READ paths passed the two-GPU and forced-network GB200 tests.
Disclosure: This response was AI-generated.
| int onHandshake(const HandShakeDesc& peer_desc, | ||
| HandShakeDesc& local_desc) { | ||
| local_desc.protocol = kHandshakeProtocol; | ||
| Json::Value request; | ||
| std::string error; | ||
| if (!decodeJson(peer_desc.payload, &request, &error)) { | ||
| local_desc.reply_msg = "Invalid NCCL handshake payload: " + error; | ||
| return 0; | ||
| } | ||
|
|
||
| const std::string op = request["op"].asString(); | ||
| Json::Value response; |
There was a problem hiding this comment.
Accessing request["op"].asString() without validating that request is an object and contains a string "op" can cause a crash or throw a Json::LogicError if the peer sends a malformed payload. We should validate the JSON structure before accessing its fields.
int onHandshake(const HandShakeDesc& peer_desc,
HandShakeDesc& local_desc) {
local_desc.protocol = kHandshakeProtocol;
Json::Value request;
std::string error;
if (!decodeJson(peer_desc.payload, &request, &error)) {
local_desc.reply_msg = "Invalid NCCL handshake payload: " + error;
return 0;
}
if (!request.isObject() || !request.isMember("op") || !request["op"].isString()) {
local_desc.reply_msg = "Missing or invalid 'op' field in handshake request";
return 0;
}
const std::string op = request["op"].asString();There was a problem hiding this comment.
Addressed in 9f3cdfb. decodeJson now uses strict parsing and rejects any non-object root. onHandshake then requires op to exist and be a string before calling asString(), and the dispatch is protected by a Json::Exception boundary so malformed peer input is returned as a handshake rejection rather than escaping the daemon callback. The outer transfer-handshake decoder now also rejects wrong-typed present fields.
Disclosure: This response was AI-generated.
| buffers->clear(); | ||
| for (const auto& item : value) { | ||
| BufferInfo buffer; | ||
| buffer.addr = item["addr"].asUInt64(); | ||
| buffer.length = item["length"].asUInt64(); | ||
| buffer.device_id = item["device_id"].asInt(); | ||
| if (!buffer.addr || !buffer.length || | ||
| buffer.device_id != expected_device) { | ||
| if (error) *error = "Invalid NCCL buffer catalog entry"; | ||
| return false; | ||
| } | ||
| buffers->push_back(buffer); | ||
| } |
There was a problem hiding this comment.
Calling asUInt64() and asInt() on item fields without verifying their existence and types can cause crashes or exceptions if the peer sends a malformed buffer catalog. We should validate each item's structure.
buffers->clear();
for (const auto& item : value) {
if (!item.isObject() || !item.isMember("addr") || !item["addr"].isUInt64() ||
!item.isMember("length") || !item["length"].isUInt64() ||
!item.isMember("device_id") || !item["device_id"].isInt()) {
if (error) *error = "Malformed buffer catalog entry";
return false;
}
BufferInfo buffer;
buffer.addr = item["addr"].asUInt64();
buffer.length = item["length"].asUInt64();
buffer.device_id = item["device_id"].asInt();
if (!buffer.addr || !buffer.length ||
buffer.device_id != expected_device) {
if (error) *error = "Invalid NCCL buffer catalog entry";
return false;
}
buffers->push_back(buffer);
}There was a problem hiding this comment.
Addressed in 9f3cdfb. decodeBuffers now validates that the catalog is an array and that every entry is an object with UInt64 address/length fields and an integer device ID before conversion. It also rejects zero values, address-range overflow, the wrong device, empty catalogs, and duplicate or overlapping ranges. Decoding uses a temporary vector and only replaces the output after the entire catalog validates, so a malformed later entry cannot leave a partial result.
Disclosure: This response was AI-generated.
Make buffer catalog freezing and metadata removal atomic, drain queued puts when event recording fails, and strictly validate NCCL handshake payloads.
|
++ @UNIDY2002 @staryxchen from the RFC #2630 |
Make kernel operations intentionally unchecked, shrink the by-value context, add compile-only NCCL/IBGDA traits coverage, and document the exact NCCL device-code version contract.
akhillanger
requested review from
ShangmingCai,
stmatengss and
ykwd
as code owners
github-actions
Bot
added
the
documentation
Apply the repository-required clang-format 20 style to all C++ and CUDA files changed by this PR.
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
UNIDY2002
left a comment
UNIDY2002
left a comment
There was a problem hiding this comment.
Thank you for your contribution! I have a few comments that should be addressed before we move on.
| option(USE_NCCL_DEVICE "option for enabling the NCCL DeviceTransport backend" OFF) | ||
| option(USE_NCCL_HOST "option for enabling the NCCL host RMA transport" OFF) |
There was a problem hiding this comment.
We have two options here, but I only see USE_NCCL_DEVICE in the build.md doc.
There was a problem hiding this comment.
This seems to be a compile-time probe. I'm wondering whether it's necessary, since the regular build process should already reveal whether these features are supported.
| ASSERT_EQ(re, 0); | ||
| } | ||
|
|
||
| TEST_F(TransferMetadataTest, BulkMemoryRemovalIsAtomic) { |
There was a problem hiding this comment.
Why do we need this case? It seems irrelevant to NCCL transport.
| if (!root.isObject()) return ERR_INVALID_ARGUMENT; | ||
| for (const char *field : {"protocol", "payload", "local_nic_path", | ||
| "local_gid", "peer_nic_path", "reply_msg"}) { | ||
| if (root.isMember(field) && !root[field].isString()) { | ||
| return ERR_INVALID_ARGUMENT; | ||
| } | ||
| } | ||
| if (root.isMember("local_lid") && !root["local_lid"].isUInt()) { | ||
| return ERR_INVALID_ARGUMENT; | ||
| } | ||
| if (root.isMember("qp_num")) { | ||
| if (!root["qp_num"].isArray()) return ERR_INVALID_ARGUMENT; | ||
| for (const auto &qp : root["qp_num"]) { | ||
| if (!qp.isUInt()) return ERR_INVALID_ARGUMENT; | ||
| } | ||
| } | ||
| #ifdef USE_BAREX | ||
| if (root.isMember("barex_port") && !root["barex_port"].isUInt()) { | ||
| return ERR_INVALID_ARGUMENT; | ||
| } | ||
| #endif | ||
| #ifdef USE_EFA | ||
| if (root.isMember("efa_addr") && !root["efa_addr"].isString()) { | ||
| return ERR_INVALID_ARGUMENT; | ||
| } | ||
| #endif | ||
| #ifdef USE_CXI | ||
| if (root.isMember("cxi_addr") && !root["cxi_addr"].isString()) { | ||
| return ERR_INVALID_ARGUMENT; | ||
| } | ||
| #endif | ||
| #ifdef USE_UB | ||
| if (root.isMember("jetty_num")) { | ||
| if (!root["jetty_num"].isArray()) return ERR_INVALID_ARGUMENT; | ||
| for (const auto &jetty : root["jetty_num"]) { | ||
| if (!jetty.isUInt()) return ERR_INVALID_ARGUMENT; | ||
| } | ||
| } | ||
| #endif |
There was a problem hiding this comment.
I'm afraid these checks are too strict, because sometimes we may launch TE of different versions.
Setting fallback values when the new fields (i.e. protocol and payload) do not exist should be more acceptable, just like the handling of local_lid and local_gid below.
| } | ||
|
|
||
| int TransferMetadata::removeSegmentDesc(const std::string &segment_name) { | ||
| std::lock_guard<std::mutex> txn_guard(local_segment_txn_mutex_); |
There was a problem hiding this comment.
Will introducing a mutex here have side effects? I'm not sure. cc: @alogfans @staryxchen
There was a problem hiding this comment.
removeSegmentDesc is less frequently used so I think it doesn't have side effects.
| int TransferMetadata::addLocalMemoryBuffer(const BufferDesc &buffer_desc, | ||
| bool update_metadata) { | ||
| std::lock_guard<std::mutex> txn_guard(local_segment_txn_mutex_); | ||
| std::shared_ptr<SegmentDesc> updated; | ||
| { | ||
| RWSpinlock::WriteGuard guard(segment_lock_); | ||
| auto new_segment_desc = std::make_shared<SegmentDesc>(); | ||
| auto &segment_desc = segment_id_to_desc_map_[LOCAL_SEGMENT_ID]; | ||
| *new_segment_desc = *segment_desc; | ||
| segment_desc = new_segment_desc; | ||
| segment_desc->buffers.push_back(buffer_desc); | ||
| auto it = segment_id_to_desc_map_.find(LOCAL_SEGMENT_ID); | ||
| if (it == segment_id_to_desc_map_.end() || !it->second) { | ||
| LOG(ERROR) << "Local segment descriptor not found"; | ||
| return ERR_METADATA; | ||
| } | ||
| updated = std::make_shared<SegmentDesc>(*it->second); | ||
| updated->buffers.push_back(buffer_desc); | ||
| it->second = updated; | ||
| } | ||
| if (update_metadata) { | ||
| return updateSegmentDesc(updated->name, *updated); | ||
| } | ||
| if (update_metadata) return updateLocalSegmentDesc(); | ||
| return 0; | ||
| } |
There was a problem hiding this comment.
Why do we have to change the impl of addLocalMemoryBuffer, removeLocalMemoryBuffer, startHandshakeDaemon, etc? They seems irrelevant to this PR.
| if (metadata_conn_string == P2PHANDSHAKE) { | ||
| rpc_binding_method = "P2P handshake"; | ||
| desc.rpc_port = findAvailableTcpPort(desc.sockfd); | ||
| if (desc.rpc_port == 0) { | ||
| LOG(ERROR) << "P2P: No valid port found for local TCP service."; | ||
| return -1; | ||
| // Preserve the legacy dynamic-port behavior for callers that pass | ||
| // only a host name. The parser supplies the default handshake | ||
| // port for those callers, but that port was not explicitly | ||
| // requested and cannot be shared by multiple local processes. | ||
| // Explicit non-zero ports are honored so independently launched | ||
| // peers can advertise predictable endpoints. | ||
| if (!hasExplicitPort(local_server_name) || port == 0) { | ||
| desc.rpc_port = findAvailableTcpPort(desc.sockfd); | ||
| if (desc.rpc_port == 0) { | ||
| LOG(ERROR) | ||
| << "P2P: No valid port found for local TCP service."; | ||
| return -1; | ||
| } |
There was a problem hiding this comment.
Would you help me review these lines of codes? @alogfans
| @@ -1261,14 +1449,55 @@ int TransferMetadata::getRpcMetaEntry(const std::string &server_name, | |||
|
|
|||
| int TransferMetadata::startHandshakeDaemon( | |||
There was a problem hiding this comment.
Does startHandshakeDaemon call only once in normal execution? If not, how to update on_receive_handshake correctly? If necessary, it's better to seperate this feature to updateHandshakeHandler().
| } | ||
| { | ||
| std::lock_guard<std::mutex> lock(handshake_handler_mutex_); | ||
| handshake_daemon_started_ = true; |
There was a problem hiding this comment.
Why need this here? I think it's better to keep it in startHandshakeDaemon () only.
5 participants
Summary
Add an optional NCCL backend to Mooncake Transfer Engine (TE), supporting both:
ncclPutandncclGetThe backend is opt-in and does not change the existing P2P, RDMA, IBGDA, CXI, TENT, or EP paths.
RFC: #2630 (This PR adds NCCL backend for both host and device APIs).
Motivation
NCCL now provides host and device APIs for one-sided communication over registered windows. Integrating these APIs
into TE allows applications to use NCCL-managed transport selection while retaining Mooncake's resource ownership,
registration, and submission model.
This also provides a common NCCL foundation for future Mooncake-EP integration without requiring EP kernels to manage NCCL communicators or windows directly.
Implementation
ncclPut.ncclGet.Semantics
NCCL communicator creation and window registration are collective operations. All participating ranks must invoke them
in the same order.
Host submissions are asynchronous and stream ordered. A TE write maps to
ncclPut, while a TE read maps toncclGet.Completion is reported only after the corresponding CUDA event completes.
The lifecycle remains externally serialized, consistent with the existing TE transport implementations.
Compatibility
Validation
Out of scope