Skip to content

UCT/CUDA_IPC: Support VMM with multiple memory allocations #11283

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tomerg-nvidia wants to merge 1 commit into
base: master
Choose a base branch
from
Open

UCT/CUDA_IPC: Support VMM with multiple memory allocations #11283

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 19 additions & 5 deletions src/uct/cuda/cuda_ipc/cuda_ipc.inl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,22 +83,36 @@ uct_cuda_ipc_check_and_pop_ctx(int is_ctx_pushed)
}

static UCS_F_ALWAYS_INLINE ucs_status_t uct_cuda_ipc_get_remote_address(
uct_cuda_ipc_extended_rkey_t *rkey, uint64_t raddr, CUdevice cu_dev,
uct_cuda_ipc_unpacked_rkey_t *rkey, uint64_t raddr, CUdevice cu_dev,
void **laddr_p, void **base_addr_p)
{
ucs_status_t status;
ptrdiff_t offset;
void *mapped_addr;

status = uct_cuda_ipc_map_memhandle(rkey, cu_dev, &mapped_addr,
#if HAVE_CUDA_FABRIC
if (rkey->super.super.ph.handle_type ==
UCT_CUDA_IPC_KEY_HANDLE_TYPE_VMM_MULTI) {
offset = raddr - rkey->super.super.d_bptr;
ucs_assertv(offset <= rkey->super.super.b_len,
"offset:%ld b_len:%lu", offset, rkey->super.super.b_len);
*laddr_p = (void*)(rkey->contig_va + offset);
if (base_addr_p != NULL) {
*base_addr_p = NULL;
}
return UCS_OK;
}
#endif

status = uct_cuda_ipc_map_memhandle(&rkey->super, cu_dev, &mapped_addr,
UCS_LOG_LEVEL_ERROR);
if (ucs_unlikely(status != UCS_OK)) {
return status;
}

offset = UCS_PTR_BYTE_DIFF(rkey->super.d_bptr, raddr);
ucs_assertv(offset <= rkey->super.b_len, "offset:%ld b_len:%lu", offset,
rkey->super.b_len);
offset = UCS_PTR_BYTE_DIFF(rkey->super.super.d_bptr, raddr);
ucs_assertv(offset <= rkey->super.super.b_len, "offset:%ld b_len:%lu",
offset, rkey->super.super.b_len);
*laddr_p = UCS_PTR_BYTE_OFFSET(mapped_addr, offset);
if (base_addr_p != NULL) {
*base_addr_p = mapped_addr;
Expand Down
8 changes: 0 additions & 8 deletions src/uct/cuda/cuda_ipc/cuda_ipc_cache.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -297,14 +297,6 @@ uct_cuda_ipc_open_memhandle_legacy(CUipcMemHandle memh, CUdevice cu_dev,
}

#if HAVE_CUDA_FABRIC
static void
uct_cuda_ipc_init_access_desc(CUmemAccessDesc *access_desc, CUdevice cu_dev)
{
access_desc->location.type = CU_MEM_LOCATION_TYPE_DEVICE;
access_desc->flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
access_desc->location.id = cu_dev;
}

static ucs_status_t
uct_cuda_ipc_open_memhandle_vmm(const uct_cuda_ipc_rkey_t *key, CUdevice cu_dev,
CUdeviceptr *mapped_addr,
Expand Down
5 changes: 2 additions & 3 deletions src/uct/cuda/cuda_ipc/cuda_ipc_ep.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,9 +120,8 @@ uct_cuda_ipc_post_cuda_async_copy(uct_ep_h tl_ep, uint64_t remote_addr,
return status;
}

status = uct_cuda_ipc_get_remote_address(&key->super, remote_addr,
cuda_device, &mapped_rem_addr,
&mapped_addr);
status = uct_cuda_ipc_get_remote_address(key, remote_addr, cuda_device,
&mapped_rem_addr, &mapped_addr);
if (ucs_unlikely(status != UCS_OK)) {
goto out;
}
Expand Down
4 changes: 4 additions & 0 deletions src/uct/cuda/cuda_ipc/cuda_ipc_iface.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -311,6 +311,10 @@ static void uct_cuda_ipc_complete_event(uct_iface_h tl_iface,
uct_cuda_ipc_event_desc_t);
ucs_status_t status;

if (cuda_ipc_event->mapped_addr == NULL) {
Copy link
Copy Markdown
Contributor

@tvegas1 tvegas1 Mar 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we integrate VMM_MULTI in the cuda_ipc cache, releasing some refcnt here? This could enable LRU as implemented in #11245, to avoid VA exhaustion on some platforms?

return; /* VMM_MULTI persistent mapping: cleanup at rkey_release */
}

status = uct_cuda_ipc_unmap_memhandle(cuda_ipc_event->pid,
cuda_ipc_event->pid_ns,
cuda_ipc_event->d_bptr,
Expand Down
Loading