[RFC]: Make NVLink/MNNVL remote mappings restart-safe with explicit allocation identity

[neverhook]

Summary

Mooncake's legacy NvlinkTransport and TENT MnnvlTransport cache imported CUDA mappings using a consumer-local segment ID and the provider virtual address. A cache hit returns the existing mapping without validating that it was created from the transport capability currently published in metadata.

This creates an ABA/lifecycle defect when a provider restarts or unregisters and re-registers memory while reusing the same endpoint/segment name and virtual address:

• CUDA VMM/Fabric path: an old imported mapping can keep the previous allocation alive, so a transfer may complete successfully against obsolete bytes.
• Legacy CUDA IPC path: the same address-derived cache bug exists, but CUDA specifies undefined behavior if the exporter frees the allocation before the importer closes its IPC mapping. The impact must therefore be described as failure, undefined access, or possible stale data—not as guaranteed preservation of the old allocation.

The required invariant is:

Once a consumer observes a newly published transport capability for a remote range, it must either use a mapping bound to that capability or fail clearly. It must never reuse an address-only mapping from an older capability.

Explicit capability identity is necessary for this invariant. It does not by itself detect provider restart or eliminate metadata staleness before a refresh; restart detection and refresh policy remain separate concerns.

Confirmed failure scenario

1. Provider P0 publishes segment E, remote VA A, and exported capability/handle H0.
2. A consumer imports H0, creates local mapping M0, and caches it under (target_id, A).
3. P0 exits, or unregisters the allocation.
4. Provider P1 reuses endpoint/segment E and VA A, but publishes H1 for a new allocation.
5. The consumer refreshes the segment descriptor and observes H1.
6. The transport still hits (target_id, A) and returns M0 without comparing H1 with the capability used to create M0.

For Fabric/VMM, CUDA allocation lifetime is extended until all mappings are unmapped and all handle references are released. The old mapping can therefore remain usable and expose obsolete bytes. For CUDA IPC, using an imported mapping after the exporter has freed the allocation is undefined behavior, but Mooncake still lacks a fail-closed mechanism preventing reuse of the old mapping.

Current code evidence

The links below are pinned to origin/main commit 98ff4e4787e99265d25938139551841350ca5f4e.

Legacy Transfer Engine

• NvlinkTransport stores remaps under (target_id, remote_base_addr); an entry contains only the local address and length: nvlink_transport.h#L70-L77.
• A cache hit returns the mapping before deserializing or validating the currently published handle. Both CUDA IPC and Fabric imports are inserted under the same address-derived key: nvlink_transport.cpp#L744-L846.
• Remote SegmentDesc objects can be replaced during refresh, but refresh does not invalidate transport mappings. A failed sync also leaves the prior descriptor cached: transfer_metadata.cpp#L961-L1059.
• closeSegment() is a no-op and cannot release per-segment mappings: transfer_engine_impl.cpp#L528-L570.

TENT

• BufferDesc carries address, length, transport handles, and attributes, but no allocation/registration capability identity: segment.h#L65-L86.
• SegmentDesc::machine_id is machine/boot identity, not a TE process or allocation incarnation; a process restart during the same host boot can retain the same value.
• TENT invalidates and refetches SegmentDesc snapshots through TTL and best-effort push, but this invalidation does not retire materialized MNNVL mappings: segment_manager.h#L35-L40, segment_manager.h#L71-L137.
• MnnvlTransport caches mappings by target_id and remote address. It also copies the mapping table into a thread-local cache; a thread-local hit bypasses metadata lookup entirely: mnnvl_transport.h#L112-L124, mnnvl_transport.cpp#L575-L665.
• uninstall() clears bookkeeping but has a TODO for closing opened entries, so mappings and associated CUDA/OS resources are not retired: mnnvl_transport.cpp#L131-L136.

Refreshing metadata alone is therefore insufficient: mapping lookup must validate the capability that created the mapping, and metadata invalidation must reach the mapping layer.

CUDA lifetime and resource-management issues

The current paths also have independent lifetime leaks that amplify the ABA problem.

According to the CUDA Driver API Virtual Memory Management documentation:

• an allocation is freed only after all mappings are unmapped and all handle/shareable-handle references are released;
• each successful cuMemImportFromShareableHandle() must be paired with cuMemRelease();
• every handle returned by cuMemRetainAllocationHandle() must be released with a corresponding cuMemRelease().

Current code paths do not consistently perform those pairs:

• Legacy and TENT provider registration retain an allocation handle for export but do not release that retained handle after export.
• Legacy and TENT consumers import a generic allocation handle and map it, but do not release the imported generic handle after a successful map.
• TENT uninstall() does not cuMemUnmap() or free the reserved VA.
• The POSIX-FD mode exports/imports descriptors without closing them in these paths.

For CUDA IPC, the CUDA Runtime API documentation states that freeing an exported allocation before the importing context calls cudaIpcCloseMemHandle() is undefined behavior. The IPC fix must therefore ensure stale mappings are closed and never reused; it must not rely on the old IPC allocation remaining valid.

Required identity model

The wire format should carry an explicit, opaque identity for the published transport capability. A single globally unique capability_id is sufficient for correctness; it can be implemented, for example, as:

• a random 128-bit value generated for each registration publication; or
• (provider_instance_id, registration_generation) where the provider instance changes at TE process startup and the generation changes for every new registration capability.

Separate allocation_id, provider identity, and generation fields may be useful for observability or future policy, but four independently meaningful fields are not required by the core cache-correctness invariant.

The capability identity must:

• never be derived only from virtual address;
• change after provider process restart;
• change after unregister/re-register, even if endpoint, VA, and length are reused;
• change when the published transport capability is replaced in a way that requires a new import;
• be serialized in backward-compatible metadata.

For older metadata without capability_id, a new consumer must not treat an address-only cache hit as valid across descriptor replacement. It should either compare the current transport handle and remap when it changes, or conservatively retire address-derived mappings on refresh.

Proposed implementation direction

1. Bind mappings to capability identity

Use a key equivalent to:

(target_id, remote_base_addr, capability_id)

A cache hit must validate the full identity and requested range. When the current descriptor carries a different capability, create a new mapping and atomically publish it instead of returning the address-only entry.

The raw Fabric/IPC handle may be fingerprinted as a transitional compatibility check, but an opaque CUDA handle's byte representation should not be treated as the long-term protocol identity.

2. Hold an immutable route and mapping lease through completion

Resolve a request to an immutable snapshot containing:

• target segment;
• remote range;
• capability identity;
• transport handle/attributes;
• a reference-counted mapping lease.

Hold the lease until the asynchronous CUDA operation completes. This prevents metadata refresh from creating a TOCTOU gap between route selection, submission, and retirement.

3. Retire mappings safely

Use one idempotent retirement primitive for close, uninstall, metadata invalidation, and error paths:

1. prevent new acquisitions of the old mapping;
2. wait for outstanding leases/CUDA completion events;
3. for VMM/Fabric, call cuMemUnmap() on the exact mapped range;
4. free the reserved VA with cuMemAddressFree();
5. for CUDA IPC, call cudaIpcCloseMemHandle();
6. release every retained/imported generic allocation handle with cuMemRelease();
7. close exported/imported POSIX file descriptors at the appropriate lifecycle boundary.

An imported VMM generic handle can normally be released after a successful cuMemMap(); the mapping itself remains until explicitly unmapped. Error paths must release partially acquired resources as well.

Thread-local caches must store generation-aware references to shared mapping objects, not detached copies of raw addresses that cannot observe invalidation.

4. Refresh once, then fail closed

On a capability/range mismatch:

1. force-refresh the TE segment descriptor once;
2. resolve the route again;
3. if the identity still cannot be matched, reject the transfer candidate.

A failed forced refresh must not authorize continued use of a mapping already known to conflict with current capability metadata.

5. Keep restart detection separate

Capability identity prevents ABA reuse after the new descriptor is observed. It does not itself notify consumers that a provider restarted.

Detection may continue to use registry replacement, TENT push invalidation, TTL, connection failure, explicit lifecycle events, or stronger provider liveness checks. Any requirement to reduce the pre-refresh stale-data window should be handled as a separate detection-policy change.

RDMA scope

The legacy RDMA path sends QP session, remote address, and rkey with each operation. Provider restart normally invalidates the old QP/MR; completion errors trigger endpoint deletion, forced metadata refresh, selection of the new rkey, and redispatch: worker_pool.cpp#L413-L468, worker_pool.cpp#L472-L505.

That gives RDMA practical failure-driven recovery, but it is not a formal global ABA proof because endpoint/address/rkey are not guaranteed lifetime-unique identities. RDMA behavior changes remain out of scope for the initial NVLink/MNNVL fix.

Scope

This RFC is limited to lifecycle and capability identity for the existing legacy NvlinkTransport and TENT MnnvlTransport implementations.

• RDMA behavior changes are out of scope for the initial patch.
• New memory kinds and new transport capabilities are out of scope.
• Restart detection latency is out of scope except where required to connect an already-observed metadata invalidation to mapping retirement.
• Store metadata should carry only references/routing information if needed; authoritative capability identity and raw transport handles remain TE-owned.

Suggested test matrix

Cover legacy TE and TENT where applicable:

• Fabric provider restart with the same endpoint/segment and remote VA, but a new capability/handle;
• same-process unregister/re-register of the same VA;
• metadata refresh followed by an address-cache hit;
• multiple TENT threads holding thread-local mapping-cache entries during invalidation;
• concurrent transfers while capability identity changes;
• in-flight-safe replacement and retirement;
• repeated import/invalidate/retire cycles with CUDA VA, allocation-handle, and FD leak checks;
• provider removal/tombstone followed by endpoint reuse;
• backward-compatible behavior when capability identity is absent;
• CUDA IPC tests proving the old IPC mapping is closed and not reused, without assuming exporter-free behavior is defined;
• RDMA regressions confirming existing QP/rkey recovery is unchanged.

The key Fabric assertion is:

After the consumer observes P1/H1, it must either access P1 through a mapping bound to H1's capability identity or fail clearly. It must never report success with P0's bytes.

The key IPC assertion is:

After the consumer observes P1/H1, it must never reuse the IPC mapping opened from H0; the old mapping must be closed or the operation must fail clearly.

Definition of done

• Backward-compatible TE metadata carries an explicit registration capability identity.
• Legacy NVLink and TENT MNNVL mapping caches include and validate that identity.
• Metadata invalidation reaches shared and thread-local mapping lookup state.
• Requests hold a generation-bound mapping lease through asynchronous completion.
• Retirement is in-flight-safe and releases mapped VA, IPC mappings, CUDA generic handles, and OS file descriptors.
• Provider and consumer retain/import calls have balanced release paths, including errors.
• Restart/re-registration tests cover endpoint and VA reuse and demonstrate fail-closed behavior after the new descriptor is observed.
• CUDA IPC behavior is described and tested without relying on undefined exporter-free semantics.
• RDMA remains behaviorally unchanged in the initial patch.

[github-actions]

Thanks for opening this issue, @neverhook!

Field	Value
Issue	#2832
GitHub user ID	41352674
Reporter	@neverhook

A maintainer will triage this when possible. To help us respond faster, please include:

• Mooncake version or commit SHA
• Environment (OS, CUDA/driver, RDMA stack if relevant)
• Steps to reproduce and expected vs. actual behavior

Useful links: Documentation · Contributing guide

This message was posted automatically by the issue bot.