[GH-ISSUE #4918] [RFC]: nvlink_intra support for VMM (cuMemCreate) buffers, including PyTorch expandable segments #1250

New issue

Open

opened 2026-10-29 21:24:25 +00:00 by bsakaguchi · 2 comments

bsakaguchi commented

2026-10-29 21:24:25 +00:00

Owner

Originally created by @whn09 on GitHub (Oct 25, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4918

Changes proposed

Motivation

nvlink_intra shares registered GPU buffers with cudaIpcGetMemHandle. That call fails for memory from the CUDA virtual memory management API (cuMemCreate + cuMemMap), so such buffers cannot be registered at all. The common case is PyTorch with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, which serving stacks often enable (#2511 reported it and was closed without a fix for this transport; #3538 fixed the RDMA dma-buf path only).

Proposed design

Detection and registration.

cuMemRetainAllocationHandle succeeds only for VMM memory, so it decides between the existing cudaIpc path and the new one.
For VMM memory, cuMemGetAddressRange returns only the physical allocation mapped at an address, not the whole reservation. Registration therefore walks the range chunk by chunk and requires every byte to be mapped.
Each chunk must have CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR in requestedHandleTypes; otherwise registration fails with a message that says what to change.

Publishing. One BufferDesc covers the chunks. Its shm_name carries, behind a magic and version, the chunk layout and the address of a per-process fd server.

No metadata schema change.
Old readers see a payload that is not a cudaIpcMemHandle_t and fail with the existing "Mismatched NVLink data transfer method" error instead of crashing.

Passing fds. An exported POSIX fd is a per-process object, so it travels over a Unix domain socket with SCM_RIGHTS.

SOCK_SEQPACKET, abstract by default (@mooncake-nvlink-<pid>-<nonce>), or a socket file in MC_NVLINK_FD_DIR for peers that share a directory but not a network namespace.
The server accepts the same uid or root (SO_PEERCRED).
It exports fresh fds per request and closes them after sending, so registered buffers hold no fds. The importer closes them after import.
pidfd_getfd (PyTorch's approach) is not used: with Yama ptrace_scope=1 it only works for descendants.

Import. On first use, the initiator fetches the fds, runs cuMemImportFromShareableHandle, maps the chunks back to back into a new reservation, and grants access to the visible devices (falling back to the current device). Teardown unmaps per chunk, frees the reservation and releases the handles.

Configuration.

MC_NVLINK_VMM=0 restores today's behaviour.
MC_NVLINK_FD_DIR selects socket files instead of an abstract socket.
Legacy buffers are unchanged.

Out of scope

FABRIC handles (they need IMEX). On hosts with fabric access PyTorch picks FABRIC by default; the error message points to expandable_segments_handle_type:posix_fd.
Inter-node transfers (MNNVL nvlink, RDMA, EFA).
Tracking allocator growth. Only memory that is mapped at registration time is shared.

Validation (HGX B200, driver 595.91, CUDA 13.2, PyTorch 2.13)

New unit tests:
- fd channel: no CUDA, runs on every Linux build.
- Two-process GPU test: three chunks of different sizes, writes across both boundaries, rejection of non-shareable memory and of partly mapped ranges, no fd leak over 20 cycles, exporter release while the peer still has the mapping.
PyTorch expandable segments at 1 and 4 GiB, read and write, data checked on both sides.
Bandwidth, VMM vs. cudaMalloc, interleaved n=3:
- 268 MiB: 733.8 vs 743.8 GB/s.
- 112 MiB: 682.5 vs 689.9 GB/s.
- A batch of 900 x 114 KiB is 10% slower. The same gap appears in a single process without Mooncake (0.826 vs 0.723 ms), so it comes from copying into VMM memory, not from the import.
Separate PID namespaces (both peers PID 1): cudaIpc fails as today, VMM works.
Separate network namespaces: an abstract socket fails with an explicit message; MC_NVLINK_FD_DIR works.

Before submitting a new issue...

Make sure you already searched for relevant issues and read the documentation

Originally created by @whn09 on GitHub (Oct 25, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4918 ### Changes proposed ## Motivation `nvlink_intra` shares registered GPU buffers with `cudaIpcGetMemHandle`. That call fails for memory from the CUDA virtual memory management API (`cuMemCreate` + `cuMemMap`), so such buffers cannot be registered at all. The common case is PyTorch with `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`, which serving stacks often enable (#2511 reported it and was closed without a fix for this transport; #3538 fixed the RDMA dma-buf path only). ## Proposed design **Detection and registration.** - `cuMemRetainAllocationHandle` succeeds only for VMM memory, so it decides between the existing cudaIpc path and the new one. - For VMM memory, `cuMemGetAddressRange` returns only the physical allocation mapped at an address, not the whole reservation. Registration therefore walks the range chunk by chunk and requires every byte to be mapped. - Each chunk must have `CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR` in `requestedHandleTypes`; otherwise registration fails with a message that says what to change. **Publishing.** One `BufferDesc` covers the chunks. Its `shm_name` carries, behind a magic and version, the chunk layout and the address of a per-process fd server. - No metadata schema change. - Old readers see a payload that is not a `cudaIpcMemHandle_t` and fail with the existing "Mismatched NVLink data transfer method" error instead of crashing. **Passing fds.** An exported POSIX fd is a per-process object, so it travels over a Unix domain socket with `SCM_RIGHTS`. - `SOCK_SEQPACKET`, abstract by default (`@mooncake-nvlink-<pid>-<nonce>`), or a socket file in `MC_NVLINK_FD_DIR` for peers that share a directory but not a network namespace. - The server accepts the same uid or root (`SO_PEERCRED`). - It exports fresh fds per request and closes them after sending, so registered buffers hold no fds. The importer closes them after import. - `pidfd_getfd` (PyTorch's approach) is not used: with Yama `ptrace_scope=1` it only works for descendants. **Import.** On first use, the initiator fetches the fds, runs `cuMemImportFromShareableHandle`, maps the chunks back to back into a new reservation, and grants access to the visible devices (falling back to the current device). Teardown unmaps per chunk, frees the reservation and releases the handles. **Configuration.** - `MC_NVLINK_VMM=0` restores today's behaviour. - `MC_NVLINK_FD_DIR` selects socket files instead of an abstract socket. - Legacy buffers are unchanged. ## Out of scope - FABRIC handles (they need IMEX). On hosts with fabric access PyTorch picks FABRIC by default; the error message points to `expandable_segments_handle_type:posix_fd`. - Inter-node transfers (MNNVL `nvlink`, RDMA, EFA). - Tracking allocator growth. Only memory that is mapped at registration time is shared. ## Validation (HGX B200, driver 595.91, CUDA 13.2, PyTorch 2.13) - New unit tests: - fd channel: no CUDA, runs on every Linux build. - Two-process GPU test: three chunks of different sizes, writes across both boundaries, rejection of non-shareable memory and of partly mapped ranges, no fd leak over 20 cycles, exporter release while the peer still has the mapping. - PyTorch expandable segments at 1 and 4 GiB, read and write, data checked on both sides. - Bandwidth, VMM vs. cudaMalloc, interleaved n=3: - 268 MiB: 733.8 vs 743.8 GB/s. - 112 MiB: 682.5 vs 689.9 GB/s. - A batch of 900 x 114 KiB is 10% slower. The same gap appears in a single process without Mooncake (0.826 vs 0.723 ms), so it comes from copying into VMM memory, not from the import. - Separate PID namespaces (both peers PID 1): cudaIpc fails as today, VMM works. - Separate network namespaces: an abstract socket fails with an explicit message; `MC_NVLINK_FD_DIR` works. ### Before submitting a new issue... - [X] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

bsakaguchi added the

RFC

label

2026-10-29 21:24:25 +00:00

bsakaguchi commented

2026-10-29 21:24:26 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Oct 25, 2026):

Thanks for opening this issue, @whn09!

Field	Value
Issue	#4918
GitHub user ID	`964527`
Reporter	@whn09

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.

@github-actions[bot] commented on GitHub (Oct 25, 2026): Thanks for opening this issue, @whn09! | Field | Value | |-------|-------| | **Issue** | #4918 | | **GitHub user ID** | `964527` | | **Reporter** | @whn09 | 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](https://kvcache-ai.github.io/Mooncake/) · [Contributing guide](https://github.com/kvcache-ai/Mooncake/blob/main/CONTRIBUTING.md) > This message was posted automatically by the issue bot.

bsakaguchi commented

2026-10-29 21:24:27 +00:00

Author

Owner

@whn09 commented on GitHub (Oct 25, 2026):

Implementation: #4919

@whn09 commented on GitHub (Oct 25, 2026): Implementation: #4919

bsakaguchi referenced this issue

2026-10-29 21:32:50 +00:00

[PR #1250] [CLOSED] [TE] Refactor TransferEngine into Impl + forwarding wrapper #2106