kvcache-ai / kvcache-ai/Mooncake
[RFC]: Make NVLink/MNNVL remote mappings restart-safe with explicit allocation identity
- Dominant language
- C++
- Stars
- 6.6k
- Forks
- 1.2k
- Avg merge
- 3d 5h
- Merged PRs (30d)
- 312
Description
## 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](https://github.com/kvcache-ai/Mooncake/blob/98ff4e4787e99265d25938139551841350ca5f4e/mooncake-transfer-engine/include/transport/nvlink_transport/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](https://github.com/kvcache-ai/Mooncake/blob/98ff4e4787e99265d25938139551841350ca5f4e/mooncake-transfer-engine/src/transport/nvlink_transport/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](https://github.com/kvcache-ai/Mooncake/blob/98ff4e4787e99265d25938139551841350ca5f4e/mooncake-transfer-engine/src/transfer_metadata.cpp#L961-L1059).
- `closeSegment()` is a no-op and cannot release per-segment mappings: [transfer_engine_impl.cpp#L528-L570](https://github.com/kvcache-ai/Mooncake/blob/98ff4e4787e99265d25938139551841350ca5f4e/mooncake-transfer-engine/src/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](https://github.com/kvcache-ai/Mooncake/blob/98ff4e4787e99265d25938139551841350ca5f4e/mooncake-transfer-engine/tent/include/tent/runtime/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](https://github.com/kvcache-ai/Mooncake/blob/98ff4e4787e99265d25938139551841350ca5f4e/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h#L35-L40), [segment_manager.h#L71-L137](https://github.com/kvcache-ai/Mooncake/blob/98ff4e4787e99265d25938139551841350ca5f4e/mooncake-transfer-engine/tent/include/tent/runtime/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](https://github.com/kvcache-ai/Mooncake/blob/98ff4e4787e99265d25938139551841350ca5f4e/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h#L112-L124), [mnnvl_transport.cpp#L575-L665](https://github.com/kvcache-ai/Mooncake/blob/98ff4e4787e99265d25938139551841350ca5f4e/mooncake-transfer-engine/tent/src/transport/mnnvl/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](https://github.com/kvcache-ai/Mooncake/blob/98ff4e4787e99265d25938139551841350ca5f4e/mooncake-transfer-engine/tent/src/transport/mnnvl/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](https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html):
- 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](https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__DEVICE.html) 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](https://github.com/kvcache-ai/Mooncake/blob/98ff4e4787e99265d25938139551841350ca5f4e/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp#L413-L468), [worker_pool.cpp#L472-L505](https://github.com/kvcache-ai/Mooncake/blob/98ff4e4787e99265d25938139551841350ca5f4e/mooncake-transfer-engine/src/transport/rdma_transport/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.
Contributor guide
Research direction
Start with NvlinkTransport in nvlink_transport.h and nvlink_transport.cpp, then trace SegmentDesc refresh in transfer_metadata.cpp and closeSegment() in transfer_engine_impl.cpp. Compare this with BufferDesc, SegmentManager, and MnnvlTransport in the TENT files, including its thread-local cache and uninstall() path. Done means capability-aware lookup, safe invalidation and resource retirement, refresh-then-fail-closed behavior, and coverage for legacy and TENT mappings.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- distributed-systems
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 32/100