[neverhook]

Summary

This PR provides the minimum Transfer Engine prerequisite for RFC #2914.

It adds CUDA HOST_NUMA Fabric allocation and closes the exact-range NVLink
registration lifecycle. It intentionally does not add the Store EGM pool or
the complete Provider/Consumer data path.

Stack

This is [1/3] of the RFC implementation stack:

1. [1/3] [TransferEngine] Add HOST_NUMA Fabric allocation core #2966 — HOST_NUMA/Fabric Transfer Engine prerequisite (this PR).
2. [2/3] [Store] Add EGM-backed Store pool over NVLink #3335 — Provider-owned EGM Store pool.
3. [3/3] [Store] Avoid staging for remote GPU reads #3431 — direct remote-memory reads into GPU destinations.

The fork branches form this exact commit chain. GitHub still displays each
community PR against main because a cross-fork PR cannot use another fork
branch as its base.

Core changes

HOST_NUMA placement

• Adds NvlinkHostNumaAllocation for Fabric-capable CUDA HOST_NUMA memory.
• Validates the IPC switch and every visible GPU's VMM, Fabric, and HOST_NUMA
  VMM capabilities independently from NUMA discovery.
• Auto discovery sorts and deduplicates GPU-local NUMA IDs and validates each
  discovered node's granularity.
• Explicit placement validates only the selected node's granularity, so an
  unrelated node cannot block the override.
• Grants access to all visible GPUs in one batched cuMemSetAccess call.

Retryable exact-range registration lifecycle

• Inserts a handle-free registration placeholder before retaining the CUDA
  allocation handle. A duplicate or map-allocation failure therefore cannot
  leave an untracked retained handle.
• If cuMemRetainAllocationHandle() fails, returns ERR_MEMORY, removes the
  placeholder, and does not publish metadata.
• Export failure, publication rollback, and public unregister share the same
  lock-held cleanup helper. A failed registration can be retried at the same
  address: pending cleanup finishes before a new handle is acquired, while a
  successful registration still rejects duplicates.
• Registration pins protect the allocation's VA throughout registration,
  rollback, and deferred/batch cleanup. Release() refuses to unmap a pinned
  range. A range being released remains unavailable for registration until
  address-free succeeds.
• Batch deletion retains handles and pins until metadata publication succeeds;
  partial cleanup retries skip already completed work.
• Metadata cleanup is recorded before external publication completes, so a
  later retry does not repeat a completed delete. ERR_ADDRESS_NOT_REGISTERED
  remains idempotent.
• A failed cuMemRelease() preserves the registration record and handle
  ownership; a later unregister retries only the remaining release stage.
• Allocation-owned release follows the same rule: ownership is cleared only
  after a successful release, and a failed call can be retried.
• Every rollback checks the release result instead of silently discarding it.
• Create() requires an empty output. An error may still return a partial
  owner; the caller must retain it and retry Release(). Destruction remains
  best-effort, not a replacement for successful explicit cleanup.

CUDA documents that these APIs may report errors from earlier asynchronous
work and does not explicitly guarantee failure atomicity. This PR deliberately
treats a non-successful release as retaining ownership so cleanup remains
retryable. That policy and its accepted ambiguity are documented in the code.
See the CUDA Driver API virtual-memory documentation.

Lifetime and tests

• Makes OwnedRangeRegistry a process-lifetime heap singleton to avoid static
  destruction-order hazards.
• Adds hardware-free scripted coverage for capability/placement failures,
  retain/export/publication failures, metadata retry, idempotent unregister,
  release retry, concurrent register/release, VA reuse, and partial batch cleanup.

CI integration

• Runs the two fake-driver CTest entries in the existing unit-tests build
  with USE_MNNVL=ON, reusing the CUDA driver runtime setup.
• Removes the separate NVLink build and its path filter.
• Labels the real GPU entry nvlink_gpu and excludes only that entry from the
  CPU runner. Hardware and fake-driver cases use separate GTEST_FILTER
  environments, so they do not run twice.
• Adds an optional label exclusion to the shared CTest action. Nightly keeps
  its existing test selection when the input is omitted.

Scope boundary

This PR does not change Store behavior, Consumer mapping/cache, wire schema,
TransferMetadata behavior, RDMA behavior, or Store configuration. Store EGM pooling,
the full Provider/Consumer path, and GB200/NVL72 integration validation belong
to the follow-up Store layer. GB200 validation is therefore not a merge gate
for #2966.

Review size

• 1 logical commit.
• 12 changed files.
• +2,570/-2 lines.

Validation

This revision consolidates the CI build and test selection; allocation,
registration, transfer, and cleanup behavior are unchanged.

• clang-format 20 check on PR-changed C/C++ lines.
• git diff --check and all applicable touched-file pre-commit hooks.
• CUDA 12.8 ARM64 focused build.
• nvlink_vmm_unit (2 tests).
• CUDA-off Transfer Engine build.
• CI shell-script regressions (12 tests, including 5 new CTest-action tests).
• CTest JSON selection: only the real GPU entry is excluded; both
  fake-driver entries remain selected and ran successfully.
• New-head complete unit-tests sweep and GitHub Actions CI Gate.

Local CUDA-on/off regression tests passed on the composed stack (#3431).
GitHub Actions passed the complete CTest sweep (226 entries) and CI Gate
for this revision. Both fake-driver entries ran successfully; the hardware-only
entry was excluded on the CPU runner.

Repository-wide pre-commit run --all-files is not a merge gate.

Related RFC: #2914

Module

• Transfer Engine (mooncake-transfer-engine)
• CI/CD

Type of Change

• New feature
• Bug fix

Checklist

• I have performed a self-review of my own code
• I have formatted changed C++ lines using clang-format 20.1.8
• I have added tests to prove my changes are effective
• For changes >500 LOC: RFC [RFC]: EGM-backed DRAM Store tier over Multi-Node NVLink #2914 exists

AI Assistance Disclosure

• AI tools were used

Codex assisted with the rebase, focused lifecycle fixes and tests, scoped
validation, and PR preparation. The human submitter remains responsible for
reviewing, understanding, and defending every changed line.

[gemini-code-assist]

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.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[he-yufeng]

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.

[neverhook]

@he-yufeng Thanks for the detailed lifecycle review. Addressed both points in 5e06423.

1. Cleanup-pending observability

• The CreateWithDriverApi failure now includes CleanupPendingOwnerCount in the returned status.
• Failed cleanup retries emit a throttled ERROR log with the retained base, length, location type/id, and CUDA error.
• Retryable staged ownership for cuMemUnmap and cuMemAddressFree remains unchanged.

2. cuMemRelease at-most-once ownership

• I did not rely on failure being atomic. Every relevant lifecycle path now consumes its handle-owned marker before invoking cuMemRelease.
• A non-success result is treated as ambiguous and is logged, but the same opaque handle is never retried. This covers factory post-map release, owner Release, consumer import cleanup, cached mapping cleanup, registration failure, and unregistration retry.
• Fake-driver tests model an error that has already consumed the reference and assert exactly one release call.

Also addressed the two batching comments: provider allocation and consumer lazy import now build complete access descriptor arrays before acquiring CUDA ownership, then issue one cuMemSetAccess call. Unrelated whole-header formatting churn was reverted.

Validation for head 5e06423: clang-format 20 passed, and both ARM64 Python 3.10/3.12 jobs passed Build project plus Build and gate NVLink VMM unit tests in https://github.com/neverhook/Mooncake/actions/runs/29802040245. The broader manually triggered matrix is still finishing, but the two NVLink gates are complete.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[neverhook]

Current version summary (3c850601):

• Separates visible-GPU Fabric/VMM capability checks from NUMA discovery: explicit placement validates only the selected node's granularity, while all visible GPUs still receive one batched cuMemSetAccess grant.
• Closes exact HOST_NUMA registration cleanup: metadata deletion is not repeated, publication failures roll back immediately, ERR_ADDRESS_NOT_REGISTERED is idempotent, and failed retained-handle release remains retryable.
• Makes OwnedRangeRegistry process-lifetime and adds scripted regression coverage for placement, rollback, publication, metadata, and handle-release failure paths.
• Validation is complete: focused CUDA-on tests pass 2/2, CUDA-off builds 113/113 targets, and exact-head GitHub Actions pass 28/28 checks including CI Gate.

Scope remains the minimum Transfer Engine prerequisite for #2914. Store EGM pool integration, the complete Provider/Consumer path, and GB200/NVL72 end-to-end validation stay in the follow-up Store-only PR and are not merge gates for #2966.

[neverhook]

Stack rebase update:

• Rebased onto the current main; the [1/3] [TransferEngine] Add HOST_NUMA Fabric allocation core #2966 patch is range-diff equivalent.
• Preserved one logical commit and the existing 8-file scope.
• Local validation: clang-format 20, git diff --check, applicable touched-file hooks, CUDA-on nvlink_vmm_unit 2/2, and CUDA-off Transfer Engine build.
• [2/3] [Store] Add EGM-backed Store pool over NVLink #3335 and [3/3] [Store] Avoid staging for remote GPU reads #3431 were replayed on top in dependency order.

Exact-head GitHub Actions are now the remaining automated gate. GB200/NVL72 integration remains a later Store-stack gate, not a gate for #2966.

[neverhook]

Stack maintenance update:

• Rebased the complete [1/3] -> [2/3] -> [3/3] stack onto current main; this layer remains patch-equivalent by range-diff.
• The composed head passed diff/format/pre-commit checks, CUDA-on focused tests (nvlink_vmm_unit 2/2, egm_store_pool_unit 1/1, direct-read 3/3), and CUDA-off Store/Transfer Engine/Master/Python validation.
• Fresh exact-head GitHub Actions have been triggered and are currently queued.

GB200 was not rerun because this was a patch-equivalent baseline rebase with no functional-path change.

[neverhook]

Review and stack update:

• Documented the callee-side removeLocalMemoryBuffer() contract without changing its signature or behavior.
• Added a dedicated NVLink path output and adapted the focused step to current main's single Python 3.12 build-flags job.
• Rebased the full [1/3] -> [2/3] -> [3/3] stack onto current main. Local CUDA-on/off focused builds and lifecycle tests pass.
• Replied to and resolved both new review threads. Fresh exact-head Actions are queued.

[neverhook]

Rebased the stack on current main. The HOST_NUMA and registration patch is unchanged. Focused tests pass; the fresh CI run is still finishing.

[neverhook]

Rebased on current main, including the SSD-test fix in #3870. The HOST_NUMA patch is unchanged. Fresh CI has been triggered for this head.

[neverhook]

The stack is rebased on current main. I kept the new CI layout and wired the focused NVLink tests into its path filter; the allocation and registration code did not need changes.

[staryxchen]

PR#3978 is merged, does this overlap with your changes?

[neverhook]

PR#3978 is merged, does this overlap with your changes?

both add HOST_NUMA/EGM support, but #3978 targets TENT’s MnnvlTransport; #2966 targets the non-TENT NvlinkTransport used by #3335. #3978 does not implement the Store pool or Store direct-read changes.

[staryxchen]

Overall the approach looks right. Left a few inline comments, please take a look.

[Aionw]

The new nvlink path filter in .github/workflows/ci.yml is wide enough that the focused NVLink step runs on most PRs, which I don't think is the intent.

Replaying the filter over the last 150 first-parent merges on main (head e699484):

rule	matches
nvlink filter as written	78 / 150 (52%)
merges actually touching nvlink_transport/	1 / 150

The list (mooncake-transfer-engine/**, mooncake-common/**, cmake/**, extern/**, CMakeLists.txt, dependencies.sh, .github/actions/**, ci.yml) is almost character-for-character the existing tent filter, and the step lives in build-flags, which is already gated on src == 'true'. So the only thing the new output filters out is PRs that touch src but not the Transfer Engine; anything touching TE pays for it.

What it pays is a full-repo configure + build (-S . -B build-nvlink-vmm-tests, USE_CUDA=ON -DUSE_MNNVL=ON) added to a job that already runs its own Transfer Engine build with the opposite flags (USE_CUDA=OFF -DUSE_MNNVL=OFF), so sccache reuse between the two is limited.

Suggestion: drop the dedicated build rather than tune the filter. The unit-tests job already configures USE_CUDA=ON + BUILD_UNIT_TESTS=ON, and its Configure CUDA driver runtime step already resolves the driver library and creates libcuda.so.1. Adding -DUSE_MNNVL=ON there builds nvlink_transport/ (two .cpp files, no new external deps — nvlink-allocator is already built under USE_CUDA), and both fake-driver tests then run in the existing ctest sweep. The nvlink filter output and the 37-line step in build-flags can both go away, and the tests end up where the rest of the unit tests already run.

One prerequisite: add_test(NAME nvlink_transport_test ...) registers the whole binary, and the pre-existing NvlinkTransportTest.WriteAndRead needs a real GPU (cudaSetDevice/cudaMalloc, exit(EXIT_FAILURE) on error). It would need a label (e.g. nvlink_gpu) plus --label-exclude on the main ctest run; the nvlink_vmm_unit cases run unchanged.

If you would rather keep the dedicated step, narrowing the filter alone also fixes the trigger rate: mooncake-transfer-engine/src/transport/nvlink_transport/**, mooncake-transfer-engine/include/transport/nvlink_transport/**, mooncake-transfer-engine/tests/nvlink_*, mooncake-transfer-engine/tests/CMakeLists.txt, mooncake-common/common.cmake and .github/workflows/ci.yml replays as 10/150 (~7%) over the same window.

[neverhook]

Thanks @Aionw, I went with the shared build. The fake-driver cases now run in unit-tests; the separate NVLink build and filter are gone. Only the real-GPU entry is excluded on CPU runners. The full CTest sweep and CI Gate passed across the stack.

[staryxchen]

Looks good from TE side

[Aionw]

LGTM for ci