[TransferEngine] Add NCCL host and device transport backend#2852
[TransferEngine] Add NCCL host and device transport backend#2852akhillanger wants to merge 11 commits into
Conversation
![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.
|
++ @UNIDY2002 @staryxchen from the RFC #2630 |
github-actions
Bot
added
the
documentation
|
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.
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.
There was a problem hiding this comment.
I kept this compile probe because the regular library target does not instantiate the header-only device helpers. The CUDA example has two separate opt-in gates—BUILD_EXAMPLES=ON and BUILD_NCCL_DEVICE_TRANSPORT_EXAMPLE=ON, with the latter disabled by default—so it is not compiled by a normal USE_NCCL_DEVICE=ON build. The probe provides CI coverage that instantiates the NCCL and IBGDA APIs from the same consumer template without requiring runtime GPU or NIC hardware. I added comments explaining this distinction.
Disclosure: This response was generated with help of AI.
staryxchen
left a comment
staryxchen
left a comment
There was a problem hiding this comment.
Thanks for contribution!
I have a question: The PR description says "A TE read maps to ncclGet", but the actual READ implementation uses requestReversePut — a synchronous handshake to the peer, which then performs the reverse PUT via handleReversePut → putAndWait. But there's no ncclGet call in the code.
More importantly, this puts a synchronous RPC on the data path. In the RFC #2630 discussion, a cross-node RPC on the data path was called out as one of the core issues with the POC, which motivated the control-plane/data-plane separation principle. Additionally, handleReversePut runs cudaEventSynchronize on the handshake daemon thread, blocking all other peers' handshakes on that thread.
|
@staryxchen thanks for the review: You are correct. There was a discrepancy in PR description. Also, I removed the reversed-PUT implementation. Currently NCCL does not have a host-side When a submitted task's opcode is READ, |
UNIDY2002
left a comment
UNIDY2002
left a comment
There was a problem hiding this comment.
Thank you for your contribution.
My main concern is that this PR is not only adding an NCCL transport. It also introduces generic framework changes: protocol-scoped metadata removal, protocol-dispatched handshake handlers, default-handler replacement, and new local metadata transaction behavior.
Those changes seem motivated by a broader multi-transport coexistence model. That model is not part of the current TE host-transport contract, and I do not think we should expand TE in that direction in this PR. In the current patch, the protocol-scoped remove APIs are only used by NCCL, while existing transports still use the old address-only path, which makes the framework change look partial and NCCL-specific rather than a coherent transport-layer design.
Please keep this PR focused on the NCCL backend within the existing TE model. If the goal is to support richer multi-transport coexistence, protocol-dispatched handshakes, or per-transport registration/removal semantics, that work should be developed under the TENT framework instead. I would be happy to review and help align that direction there, but I do not think these framework changes should be introduced through this NCCL TE transport PR.
alogfans
left a comment
alogfans
left a comment
There was a problem hiding this comment.
The PR should be refactored to minimize the change of existing modules.
|
Thank you for the comments @UNIDY2002 and @alogfans. I will work on addressing your concerns and get back. |
UNIDY2002
left a comment
UNIDY2002
left a comment
There was a problem hiding this comment.
The code is much cleaner now.
LGTM on the device transport part.
@staryxchen @alogfans When you have a chance, could you please take a look at the host transport part as well? Thanks!
| cudaSetDevice(local_device); | ||
| cudaEvent_t event = nullptr; | ||
| cudaError_t cuda_result = | ||
| cudaEventCreateWithFlags(&event, cudaEventDisableTiming); |
There was a problem hiding this comment.
Each WRITE slice allocates a cudaEvent_t into slice->nccl.event, and it is only destroyed inside poll() on success or error. But TransferTask::~TransferTask() (transport.h) deallocates slices via the cache without any transport-side cleanup, so a batch destroyed before every POSTED slice has been polled to completion leaks the live CUDA events.
NCCL is the first TE transport that owns a per-slice CUDA handle requiring explicit destruction, so the existing slice lifecycle does not cover this. Could you add a cleanup path — e.g. a transport hook invoked on slice deallocation, or track in-flight events in Impl and destroy them in ~Impl()?
There was a problem hiding this comment.
Addressed in 16d1a19. I added a one-shot cleanup callback to each slice, invoked by ThreadLocalSliceCache::deallocate() before the slice is deleted or cached. NCCL installs a callback that destroys any remaining cudaEvent_t on the correct CUDA device.
Submission-error and polling-completion paths now use the same idempotent helper. The event handle and callback are cleared before cleanup or cache reuse, preventing both double destruction and accidental inheritance by another transport. I also added tests covering cleanup through TransferTask::~TransferTask() and callback clearing across slice-cache reuse.
| return true; | ||
| } | ||
|
|
||
| std::shared_ptr<NcclSession> getOrCreateSession( |
There was a problem hiding this comment.
Once a session enters sessions_, it is only removed in ~Impl. If initialize() fails (failed_ = true), waitReady returns nullptr for every subsequent transfer to that peer, and the session cannot be recreated because buffers_frozen_ is already latched — the only recovery is destroying the whole transport.
For an experimental backend this may be acceptable, but could you at least document it in "Device constraints", or evict failed sessions so a later transfer can retry bootstrap?
There was a problem hiding this comment.
Addressed in 16d1a19 using the documentation option. The public NCCL host contract and build constraints now state that the first valid WRITE freezes the registered-buffer catalog, a terminal session failure remains cached for that endpoint/device pair, and subsequent transfers do not retry bootstrap. Recovery requires recreating the NCCL-only TransferEngine on both peers and registering the buffers again.
| return attributes.device; | ||
| } | ||
|
|
||
| std::string makeSessionKey(const std::string& local_name, int local_device, |
There was a problem hiding this comment.
makeSessionKey orders by local_name < peer_name, and local_rank is assigned the same way (local_server_name_ < peer_name ? 0 : 1). When local_name == peer_name (same-node cross-device), both sides take the else branch with swapped arguments, producing different keys; and both compute rank = 1, so neither creates the unique ID — bootstrap deadlocks.
Same-node transfers should use NVLink, so this is likely out of scope, but the latent bug is there. If you want it correct, compare the (name, device) tuple instead of name alone:
if (std::tie(local_name, local_device) < std::tie(peer_name, peer_device)) {
...
}There was a problem hiding this comment.
Addressed in 16d1a19. Session canonicalization now compares the complete (server name, CUDA device) tuple, and the same comparator is used for makeSessionKey() and both outbound and inbound rank-election paths.
At the ordering layer, equal server names on different devices now produce the same canonical key in both directions, complementary ranks, and exactly one rank-0 unique-ID creator. Same-engine transfers remain intentionally rejected and should use the intra-node NVLink/P2P path; this fixes the latent ordering assumption without claiming same-engine NCCL support.
|
@akhillanger Hi! Would you help resolve the conflicts with the main branch, so that the CI could be triggered? Thanks. |
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.
Make buffer catalog freezing and metadata removal atomic, drain queued puts when event recording fails, and strictly validate NCCL handshake payloads.
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.
Apply the repository-required clang-format 20 style to all C++ and CUDA files changed by this PR.
16d1a19 to
5c8362a
Compare
Kernels such as DeepEP V2 require rail-team connectivity. Is there any plan or interest to support this feature in the future? @akhillanger |
|
Good point. DeepEP V2's hybrid path does require rail connectivity: it checks railedGinType, requests NCCL_GIN_CONNECTION_RAIL, and uses ncclTeamTagRail for the inter-node phase. DeepEP also retains a direct mode using full world connectivity, so this requirement applies specifically to its hierarchical hybrid path rather than every V2 kernel. This PR implements only full world-team GIN connectivity. We can support rail connectivity as a follow-up so this initial backend remains focused. |
staryxchen
left a comment
staryxchen
left a comment
There was a problem hiding this comment.
LGTM, only a few nits
| @@ -0,0 +1,338 @@ | |||
| // Copyright 2024 KVCache.AI | |||
There was a problem hiding this comment.
| // Copyright 2024 KVCache.AI | |
| // Copyright 2026 KVCache.AI |
| @@ -0,0 +1,302 @@ | |||
| // Copyright 2024 KVCache.AI | |||
There was a problem hiding this comment.
| // Copyright 2024 KVCache.AI | |
| // Copyright 2026 KVCache.AI |
| @@ -0,0 +1,545 @@ | |||
| // Copyright 2024 KVCache.AI | |||
There was a problem hiding this comment.
| // Copyright 2024 KVCache.AI | |
| // Copyright 2026 KVCache.AI |
| @@ -0,0 +1,80 @@ | |||
| // Copyright 2024 KVCache.AI | |||
There was a problem hiding this comment.
| // Copyright 2024 KVCache.AI | |
| // Copyright 2026 KVCache.AI |
|
@staryxchen fixed the copyright issues you pointed. Thanks. |
UNIDY2002
left a comment
UNIDY2002
left a comment
There was a problem hiding this comment.
LGTM. Let's wait for the CI results.
6 participants
Summary
This PR adds two opt-in experimental NCCL backends to Mooncake Transfer Engine:
Transportinterface.Host RMA transport
installTransport("nccl")installs a WRITE-only transport in aTransferEngine(false)instance. The transport uses the existing P2P metadata handshake to exchange an opaque NCCL bootstrap payload and registered CUDA-buffer metadata.Each peer/device session creates a two-rank NCCL communicator, collectively registers corresponding buffers as symmetric NCCL windows, and submits transfers with
ncclPutSignal. CUDA events report asynchronous completion through the normal Transfer Engine status API.The host backend follows the classic Transfer Engine single-transport model:
Status::NotSupportedTransport; NCCL 2.30 has no public host-side Get operation.Device transport
getOrCreateNcclTransport()provides host-side setup for an NCCL device communicator. Applications exchange the generated NCCL unique ID through their existing control plane, initialize every rank collectively, and allocate/register symmetric buffers.A compact
NcclDeviceContextis passed by value to CUDA kernels. The device helpers provide:NCCL-specific declarations live in
transport/device/nccl_device_transport.h; the generic device transport API is unchanged.Build
Both backends require CUDA and NCCL 2.30.4 or newer and are disabled by default.
cmake -S . -B build \ -DUSE_CUDA=ON \ -DUSE_NCCL_HOST=ON \ -DUSE_NCCL_DEVICE=ON cmake --build build -jThe options can be enabled independently:
USE_NCCL_HOSTrequires the publicnccl.hhost API.USE_NCCL_DEVICEadditionally requires the experimentalnccl_device.hAPI.NCCL_ROOTselects a nonstandard NCCL installation.Optional manual examples are enabled with:
Device constraints
initialize()binds the transport to the current CUDA device; that device must remain current for later host calls, shutdown, and destruction.libnccl. Device code must be rebuilt or re-JITed after an NCCL upgrade.Examples and tests
nccl_host_transport_exampleuses two GPUs to validate concurrent and bidirectional WRITE transfers and explicit READ rejection.nccl_device_transport_exampleuses one process per rank to validate LSA or GIN put/signal completion from a device kernel.transfer_metadata_testcovers ordered NCCL GPU-buffer metadata and opaque handshake-payload round trips.device_backend_api_compatibility_testinstantiates one templated put operation for NCCL and IBGDA contexts.Validation
ENABLE_MULTI_PROTOCOLconfigurations.