SIGN IN SIGN UP

UCT/CUDA_COPY/MD: Preserve registration extent for multi-handle VMM (#11461)

* UCT/CUDA_COPY/MD: Preserve registration extent for multi-handle VMM

For multi-handle CUDA VMM allocations (multiple cuMemCreate handles
mapped contiguously into a single virtual address range via cuMemMap),
cuMemGetAddressRange() returns only the bounds of the cuMem handle
containing the base pointer, not the full mapped extent.

uct_cuda_copy_md_sync_memops_get_address_range() unconditionally used
that potentially-shorter alloc_length to set mem_info->alloc_length
when UCX_CUDA_COPY_REG_WHOLE_ALLOC was on or auto, silently shrinking
the caller-requested registration extent. The truncated length flowed
into cuMemGetHandleForAddressRange() and ibv_reg_dmabuf_mr(), producing
a memory region that only covered the first cuMem handle's physical
pages. Any subsequent RDMA operation that crossed a chunk boundary
failed with IBV_WC_LOC_PROT_ERR; downstream this surfaced as
UCS_ERR_CONNECTION_RESET via the SRQ-attached QP being flushed.

Preserve the caller's requested extent when the underlying VMM
allocation is smaller than the requested registration. CUDA exports a
dmabuf spanning the full multi-handle VA range when asked with the
full size; the whole-allocation expansion was only valid for
single-handle allocations.

* AUTHORS: Add Nicolas 'Pixel' Noble

* UCT/CUDA_COPY/MD: Test multi-handle VMM registration extent

Add a gtest for the multi-handle VMM case fixed by the previous
commit. Allocate three cuMem handles contiguously into one VA
range via cuMemAddressReserve + cuMemCreate + cuMemMap, query an
offset address with the size of one chunk, and assert that
uct_md_mem_query preserves the caller's requested extent
(alloc_length >= chunk_size * num_chunks) rather than shrinking
to cuMemGetAddressRange's per-handle bounds.

Suggested by tvegas1 in the PR review thread.

* UCT/CUDA_COPY/MD: Fix VMM test to exercise the multi-handle guard

The previous test queried one chunk-size at a chunk-aligned offset,
which makes cuMemGetAddressRange return exactly that chunk and the
new is_vmm guard never fires. Those assertions only held under the
withdrawn cuPointerGetAttribute(RANGE_SIZE) approach.

Query the full multi-handle range instead so the guard fires:
caller asks for chunk_size * num_chunks, cuMemGetAddressRange
returns one chunk, alloc_length < length, the guard preserves the
caller's extent. Drop the base_address assertion since the guard
returns early before mem_info->base_address is updated.

* UCT/CUDA_COPY/MD: Trim VMM comments per review

Shorten the multi-handle VMM source comment to note only that the
physical allocation may be smaller than the mapped virtual range, and
reduce the test comment without naming the internal function, per
review feedback.

* GTEST/CUDA: Reuse cuda_vmm_mem_buffer in the cuda_copy MD test

The multi-handle VMM allocation in test_cuda_ipc_md.cc duplicated what
cuda_vmm_mem_buffer already does. Teach that class to map several physical
allocations into one virtual reservation and drop the local helpers.

- init() takes a chunk count and maps that many allocations of the
  granularity-aligned chunk size into the reserved range
- error unwinding moved into cleanup(), which is also the destructor path
  and only unmaps the chunks that were actually mapped
- vmm_multi_handle_range queries the buffer extent through size()

Signed-off-by: Nicolas 'Pixel' Noble <nicolas@nobis-crew.org>

---------

Signed-off-by: Nicolas 'Pixel' Noble <nicolas@nobis-crew.org>
Co-authored-by: Mikhail Brinskiy <brminich@users.noreply.github.com>
N
Nicolas Noble committed
6d75cec7cc3a252257f0d05d1c2674c8f9eddff8
Parent: 8f00ee1
Committed by GitHub <noreply@github.com> on 8/14/2026, 8:10:32 AM