[TransferEngine] Add NVLink VMM Fabric memory lifecycle#2966
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 NVLink Virtual Memory Management (VMM) allocation support and Fabric memory transport capabilities to the Mooncake transfer engine, refactoring memory lifecycle management in NvlinkTransport and adding comprehensive tests. The review feedback highlights critical safety and performance improvements, including addressing potential null pointer dereferences when retrieving the local segment descriptor in TransferMetadata, and optimizing driver overhead by batching cuMemSetAccess calls rather than executing them in a loop for each device.
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.
| RWSpinlock::WriteGuard guard(segment_lock_); | ||
| auto new_segment_desc = std::make_shared<SegmentDesc>(); | ||
| auto &segment_desc = segment_id_to_desc_map_[LOCAL_SEGMENT_ID]; | ||
| auto& segment_desc = segment_id_to_desc_map_[LOCAL_SEGMENT_ID]; | ||
| old_segment_desc = segment_desc; | ||
| new_segment_desc = std::make_shared<SegmentDesc>(); | ||
| *new_segment_desc = *segment_desc; |
There was a problem hiding this comment.
There is a risk of a null pointer dereference if segment_id_to_desc_map_[LOCAL_SEGMENT_ID] is not initialized or is null. We should add a defensive null check before dereferencing segment_desc.
RWSpinlock::WriteGuard guard(segment_lock_);
auto& segment_desc = segment_id_to_desc_map_[LOCAL_SEGMENT_ID];
if (!segment_desc) {
return ERR_METADATA;
}
old_segment_desc = segment_desc;
new_segment_desc = std::make_shared<SegmentDesc>();
*new_segment_desc = *segment_desc;| RWSpinlock::WriteGuard guard(segment_lock_); | ||
| auto new_segment_desc = std::make_shared<SegmentDesc>(); | ||
| auto &segment_desc = segment_id_to_desc_map_[LOCAL_SEGMENT_ID]; | ||
| auto& segment_desc = segment_id_to_desc_map_[LOCAL_SEGMENT_ID]; | ||
| old_segment_desc = segment_desc; | ||
| new_segment_desc = std::make_shared<SegmentDesc>(); | ||
| *new_segment_desc = *segment_desc; |
There was a problem hiding this comment.
There is a risk of a null pointer dereference if segment_id_to_desc_map_[LOCAL_SEGMENT_ID] is not initialized or is null. We should add a defensive null check before dereferencing segment_desc.
RWSpinlock::WriteGuard guard(segment_lock_);
auto& segment_desc = segment_id_to_desc_map_[LOCAL_SEGMENT_ID];
if (!segment_desc) {
return ERR_METADATA;
}
old_segment_desc = segment_desc;
new_segment_desc = std::make_shared<SegmentDesc>();
*new_segment_desc = *segment_desc;| int device_count = 0; | ||
| result = fabric_driver_api_.device_get_count(&device_count); | ||
| if (result != CUDA_SUCCESS || device_count <= 0) { | ||
| LOG(ERROR) << "NvlinkTransport: cuDeviceGetCount " | ||
| "failed during lazy import: " | ||
| << result; | ||
| return ERR_CONTEXT; | ||
| } | ||
|
|
||
| for (int ordinal = 0; ordinal < device_count; ++ordinal) { | ||
| CUdevice device; | ||
| result = | ||
| fabric_driver_api_.device_get(&device, ordinal); | ||
| if (result != CUDA_SUCCESS) { | ||
| LOG(ERROR) << "NvlinkTransport: cuDeviceGet failed " | ||
| "during lazy import: " | ||
| << result; | ||
| return ERR_CONTEXT; | ||
| } | ||
| CUmemAccessDesc access = {}; | ||
| access.location.type = CU_MEM_LOCATION_TYPE_DEVICE; | ||
| access.location.id = device; | ||
| access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; | ||
| result = fabric_driver_api_.mem_set_access( | ||
| attempt.address(), entry.length, &access, 1); | ||
| if (result != CUDA_SUCCESS) { | ||
| LOG(ERROR) | ||
| << "NvlinkTransport: cuMemSetAccess failed for " | ||
| "visible device " | ||
| << ordinal << ": " << result; | ||
| return ERR_MEMORY; | ||
| } | ||
| } |
There was a problem hiding this comment.
Calling cuMemSetAccess individually for each device in a loop can introduce significant driver overhead and performance bottlenecks, especially on multi-GPU systems. Batching all access configurations into a single cuMemSetAccess call is much more efficient.
int device_count = 0;
result = fabric_driver_api_.device_get_count(&device_count);
if (result != CUDA_SUCCESS || device_count <= 0) {
LOG(ERROR) << "NvlinkTransport: cuDeviceGetCount "
"failed during lazy import: "
<< result;
return ERR_CONTEXT;
}
std::vector<CUmemAccessDesc> access_descs;
access_descs.reserve(device_count);
for (int ordinal = 0; ordinal < device_count; ++ordinal) {
CUdevice device;
result =
fabric_driver_api_.device_get(&device, ordinal);
if (result != CUDA_SUCCESS) {
LOG(ERROR) << "NvlinkTransport: cuDeviceGet failed "
"during lazy import: "
<< result;
return ERR_CONTEXT;
}
CUmemAccessDesc access = {};
access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
access.location.id = device;
access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
access_descs.push_back(access);
}
result = fabric_driver_api_.mem_set_access(
attempt.address(), entry.length, access_descs.data(), access_descs.size());
if (result != CUDA_SUCCESS) {
LOG(ERROR)
<< "NvlinkTransport: cuMemSetAccess failed: " << result;
return ERR_MEMORY;
}| auto grant_access = [&](CUmemLocationType location_type, | ||
| int location_id) -> Status { | ||
| CUmemAccessDesc access = {}; | ||
| access.location.type = location_type; | ||
| access.location.id = location_id; | ||
| access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; | ||
| CUresult access_result = api.mem_set_access(ptr, length, &access, 1); | ||
| if (access_result != CUDA_SUCCESS) | ||
| return cudaDriverFailure("cuMemSetAccess", access_result); | ||
| return Status::OK(); | ||
| }; | ||
|
|
||
| if (options.location_type == LocationType::HOST_NUMA) { | ||
| status = | ||
| grant_access(CU_MEM_LOCATION_TYPE_HOST_NUMA, options.location_id); | ||
| if (!status.ok()) { | ||
| return rollback(status); | ||
| } | ||
| } | ||
|
|
||
| int device_count = 0; | ||
| result = api.device_get_count(&device_count); | ||
| if (result != CUDA_SUCCESS) { | ||
| status = cudaDriverFailure("cuDeviceGetCount", result); | ||
| return rollback(status); | ||
| } | ||
| if (device_count <= 0) { | ||
| return rollback(Status::NotSupportedTransport( | ||
| "VMM allocation requires at least one visible CUDA device")); | ||
| } | ||
|
|
||
| for (int ordinal = 0; ordinal < device_count; ++ordinal) { | ||
| CUdevice device; | ||
| result = api.device_get(&device, ordinal); | ||
| if (result != CUDA_SUCCESS) { | ||
| status = cudaDriverFailure("cuDeviceGet", result); | ||
| return rollback(status); | ||
| } | ||
| status = grant_access(CU_MEM_LOCATION_TYPE_DEVICE, device); | ||
| if (!status.ok()) { | ||
| return rollback(status); | ||
| } | ||
| } |
There was a problem hiding this comment.
Calling cuMemSetAccess individually for each device in a loop can introduce significant driver overhead and performance bottlenecks, especially on multi-GPU systems. Batching all access configurations into a single cuMemSetAccess call is much more efficient.
std::vector<CUmemAccessDesc> access_descs;
if (options.location_type == LocationType::HOST_NUMA) {
CUmemAccessDesc access = {};
access.location.type = CU_MEM_LOCATION_TYPE_HOST_NUMA;
access.location.id = options.location_id;
access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
access_descs.push_back(access);
}
int device_count = 0;
result = api.device_get_count(&device_count);
if (result != CUDA_SUCCESS) {
status = cudaDriverFailure("cuDeviceGetCount", result);
return rollback(status);
}
if (device_count <= 0) {
return rollback(Status::NotSupportedTransport(
"VMM allocation requires at least one visible CUDA device"));
}
for (int ordinal = 0; ordinal < device_count; ++ordinal) {
CUdevice device;
result = api.device_get(&device, ordinal);
if (result != CUDA_SUCCESS) {
status = cudaDriverFailure("cuDeviceGet", result);
return rollback(status);
}
CUmemAccessDesc access = {};
access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
access.location.id = device;
access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
access_descs.push_back(access);
}
if (!access_descs.empty()) {
CUresult access_result = api.mem_set_access(ptr, length, access_descs.data(), access_descs.size());
if (access_result != CUDA_SUCCESS) {
status = cudaDriverFailure("cuMemSetAccess", access_result);
return rollback(status);
}
}
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
4aef80a to
ae700b3
Compare
he-yufeng
left a comment
he-yufeng
left a comment
There was a problem hiding this comment.
Thanks for this — the VMM lifecycle handling here is careful work. A few things I want to call out and two points worth a closer look before merge. I focused the review on the new NvlinkVmmAllocation lifecycle in nvlink_vmm_allocation.cpp.
What's done well. The reverse-order Release() with per-stage commit and retained ownership on failure is the right shape for CUDA VMM teardown, and pre-allocating the intrusive CleanupPendingVmmOwnerNode before acquiring any CUDA resource (so recording a rollback can never itself fail under memory pressure) is a nice touch. Moving the owned-range provenance removal to the very start of Release(), before the first unmap, so a stuck teardown stops satisfying IsExactOwnedRange rather than staying eligible for a fresh remote registration, is exactly the safe direction.
1. A single stuck teardown can silently wedge the whole allocation path (availability). CreateWithDriverApi opens with:
if (!RetryCleanupPendingOwners()) {
return Status::Memory("a previous VMM factory rollback is still pending cleanup");
}If any quarantined owner keeps failing Release() (driver persistently rejecting a teardown stage), every subsequent Create fails from here on, process-wide, with no upper bound and no escape valve. The retention itself is defensible, but the failure is invisible to operators: CleanupPendingOwnerCount() exists but isn't wired to any metric or log-once, so the symptom presents as "all Fabric VMM allocations suddenly fail" with the root cause buried. Could we at least surface the pending count (a gauge, or a throttled LOG(ERROR) naming the retained range) so this degraded state is observable? Whether a bounded-retry-then-give-up policy is acceptable is a product call, but the observability gap seems worth closing regardless.
2. Rollback re-releases a handle whose first cuMemRelease returned an error (worth confirming). In the success path the handle is released once and the flags cleared:
result = api.mem_release(handle);
if (result != CUDA_SUCCESS) {
status = cudaDriverFailure("cuMemRelease", result);
return rollback(status); // -> owner->Release() -> mem_release(handle) again
}
owner->handle_owned_ = false;When that mem_release fails, handle_owned_ is still true, so the rollback path's Release() calls cuMemRelease on the same handle a second time. This is only safe if a failing cuMemRelease is guaranteed to be atomic (no state change on error), so the retry acts on a still-valid handle rather than double-releasing one that was partially torn down. I believe that's the driver's contract, but given this is the one spot where the same handle can be released twice, it'd be good to either confirm that assumption explicitly or clear handle_owned_ before entering rollback here and lean on the quarantine path if the handle truly leaked. Same reasoning applies to the mid-Create failures that reach rollback while mapped_/address_reserved_ are set — those are fine because each stage's flag gates its own teardown, but the handle stage is the asymmetric one.
Neither of these blocks the design, which I think is sound. #2 is a small correctness clarification; #1 is the one I'd most like addressed since a wedged allocation path is hard to diagnose in production. Happy to look again once you've had a chance to weigh in.
3 participants
Description
Adds the Transfer Engine primitives required by RFC #2914 to own, publish,
import, and clean up CUDA VMM allocations over the existing NVLink transport.
This supersedes the withdrawn #2934 with a narrower review boundary.
NvlinkVmmAllocationfor DEVICE and CUDA HOST_NUMA allocationlocations, including Fabric capability checks, alignment, exact range
provenance, and reverse-order ownership cleanup.
handle. Ambiguous commit-before-error outcomes are compensated inside
NvlinkTransport; no genericTransferMetadatatransaction model is added.ownership quarantined until cleanup succeeds or process exit.
terminal sentinel. The shared
TransferTask/completion model is unchanged.requiring an exact transport set is intentionally left to the Store PR.
Related RFC: #2914
Explicitly excluded from this PR
TransferMetadatapublication transactions or new sharedTransferTaskstate.finalizeTransferResult, repeated-result observation fields, and periodicevent-driven polling state.
*_for_testingcallbacks.Transfer Engine contract lands.
Module
mooncake-transfer-engine)Type of Change
How Has This Been Tested?
Local checks:
Results:
nvlink_vmm_unittests pass inboth Python 3.10 and 3.12 ARM64 jobs.
including all four Ubuntu/Python wheel integration combinations.
The ARM64 job builds only:
and requires CTest to discover at least one
nvlink_vmm_unittest beforerunning the label.
The earlier development branch used by #2934 completed a dual-node GB200
Provider/Consumer data-plane run, but that evidence is not claimed as a
current-revision result after this scope reduction.
Checklist
pre-commit run --all-filesand all hooks pass.AI Assistance Disclosure
Codex assisted with replaying the latest
main, removing unrelated changes,implementing the narrowed lifecycle semantics, and preparing tests and this
draft PR. The human submitter remains responsible for reviewing,
understanding, and defending every changed line.