[njhill]

What?

rocm_ipc shares VMM memory by exporting its allocation as a dma-buf file descriptor. Peers import it, map it and grant access to it, as cuda_ipc does for CUDA VMM memory. When VMM memory can't be shared, it is registered without an IPC handle, and peers report the key as unreachable. Adds rocm_ipc/test_rocm_ipc_vmm.share.

Why?

Device memory mapped with the HIP virtual memory API (hipMemCreate + hipMemMap, hsa_amd_vmem_* in ROCr) can't be used with UCX today, while the equivalent CUDA VMM memory (cuMemCreate + cuMemMap) can. Frameworks that grow GPU buffers in place reserve an address range and map memory into it on demand. vLLM's extensible KV cache (vllm-project/vllm#56492) does this, and its NIXL KV connector registers that memory with UCX. On ROCm that fails at registration, then crashes in rocm_copy, and rocm_ipc can't share the memory at all.

hsa_amd_ipc_memory_create doesn't support VMM memory, so rocm_ipc failed to register it. ucp_mem_map registers on every MD that supports the memory type, so this failed the whole registration whenever rocm_ipc was enabled.

This is one of four PRs that make VMM memory work on ROCm:

• UCT/ROCM: Treat HSA_EXT_POINTER_TYPE_HSA_VMEM as ROCm memory #12137 – memory type detection
• UCT/ROCM/COPY: Do not access VMM device memory with the CPU #12138 – no CPU access to VMM memory in rocm_copy
• UCM/ROCM: Hook the HSA virtual memory map and unmap calls #12139 – UCM hooks for hsa_amd_vmem_map/unmap
• UCT/ROCM/IPC: Share VMM memory through dma-buf file descriptors #12140 – sharing VMM memory through rocm_ipc (this PR)

This PR includes the commit of #12137 ("UCT/ROCM: Treat HSA_EXT_POINTER_TYPE_HSA_VMEM as ROCm memory"), which it depends on. Only the last commit is new here.

How?

• Export: the exporter exports the allocation mapped at the registered address with hsa_amd_vmem_export_shareable_handle. The packed key carries a handle type, a server id and an fd id.
• Fd transfer: pidfd_getfd needs ptrace access to the exporter, which Yama's default ptrace_scope=1 denies for unrelated processes. Instead, the exporter passes the fd with SCM_RIGHTS over an abstract unix socket, served from the UCS async thread. Only processes of the same user are served, checked with SO_PEERCRED (rocm_ipc_fd.c). Registrations of the same allocation share one export, so peers map it once.
• Import: the importer imports the fd, maps it at a new address and grants all GPUs access. The mapping is cached like an IPC attachment.
• Fallback: some ROCm runtimes can't share VMM allocations. They report the queried address instead of the mapping base in hsa_amd_pointer_info, and crash on imported allocations. A probe on a small local mapping detects this once. On such runtimes, and when an allocation can't be exported, VMM memory is registered without an IPC handle, so UCP selects another transport.
• Local VMM buffer: the local agent copies it, and hsa_amd_agents_allow_access isn't called on it, since that applies only to pool allocations.
• Compatibility: the packed rkey grows, so peers on different versions don't interoperate over rocm_ipc.

Real sharing needs these ROCr fixes: ROCm/rocm-systems#13990, ROCm/rocm-systems#13991, ROCm/rocm-systems#13992 and ROCm/rocm-systems#11449, which is merged.

Testing

The new test registers, packs and unpacks VMM memory in a second mapping. Without this change it fails at mem_reg ("Address not valid"). With it:

• stock ROCr: the key is reported unreachable, and the trace shows the fallback;
• patched ROCr: put_zcopy and get_zcopy through rocm_ipc move data correctly, and the trace shows the fd export, the import mapping and cache hits.

Tested on MI350X (gfx950) with the vLLM ROCm CI image (ROCm 10.0 TheRock wheels, ROCr 1.21.0). The existing rocm_ipc tests pass.

End to end: vLLM's NIXL prefill/decode accuracy test ran with VMM KV caches on both sides. With patched ROCr, the KV reads use zero-copy rocm_ipc, with steady-state transfers of 14–16 ms average and 23–24 ms P90. With stock hipMalloc caches over the staging path they were 20–26 ms average and 35–42 ms P90. Accuracy matches the non-VMM baseline. With stock ROCr the test passes over the fallback.

🤖 Generated with Claude Code

[svc-nvidia-pr-review]

🤖 Starting review — findings will be posted here when done.

[edgargabriel]

@njhill thank you for your work. As a heads up, we are having similar PRs also in the works at AMD, which we plan to file shortly. There are really three things that we need to fix :

• hipMallocAsync being reported as VMM memory starting from a recent rocm release
• support for hipVMM functions and sharing through POSIX FD mechanism
• support for hipVMM functions and sharing through fabric handles for MI455/Helios setups.

Our approach is slightly different though, we do not want to intermerge everything in a single TL (rocm_ipc), we plan to introduce a new TL (rocm_fabric) and new memory types UCS_MEMORY_ROCM_VMM

[njhill]

@njhill thank you for your work. As a heads up, we are having similar PRs also in the works at AMD, which we plan to file shortly. There are really three things that we need to fix :

• hipMallocAsync being reported as VMM memory starting from a recent rocm release
• support for hipVMM functions and sharing through POSIX FD mechanism
• support for hipVMM functions and sharing through fabric handles for MI455/Helios setups.

Our approach is slightly different though, we do not want to intermerge everything in a single TL (rocm_ipc), we plan to introduce a new TL (rocm_fabric) and new memory types UCS_MEMORY_ROCM_VMM

Thanks @edgargabriel! I'm curious do you have any eta for a UCX release containing these fixes?

[ChuanLi1101]

@edgargabriel @mshanthagit Routing this for review as the ROCm CODEOWNERs. This is part of the VMM-on-ROCm series (#12137 to #12140) that vLLM extensible KV cache (vllm-project/vllm#56492) needs for NIXL registration on ROCm. Edgar, your note above says AMD plans a rocm_fabric TL and UCS_MEMORY_ROCM_VMM instead of extending rocm_ipc. Could you say whether you want to review this rocm_ipc dma-buf path, or whether the AMD PR should supersede it? Any ETA for a release with these fixes would also help.