Skip to content

UCT/CUDA_COPY: Fix incorrect mem_flags detection #11635

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

Merged
brminich merged 10 commits into from
Jul 10, 2026
+128 −25
Merged

UCT/CUDA_COPY: Fix incorrect mem_flags detection #11635

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
a5d2ba6
TEST/GTEST/UCT: Test cuda_copy mem_query without context
tomerg-nvidia Jul 8, 2026
4614172
UCT/CUDA_COPY: Push context before detecting mem_flags
tomerg-nvidia Jul 8, 2026
65d1073
TEST/GTEST: format fix
tomerg-nvidia Jul 8, 2026
be6bdd2
UCT/CUDA_COPY: Track pushed context
tomerg-nvidia Jul 9, 2026
ef7c5ac
UCT/CUDA_COPY: Use existing context if push failed for dmabuf export
tomerg-nvidia Jul 9, 2026
4f88c06
UCT/CUDA_COPY: Mark host-located vmm as registerable
tomerg-nvidia Jul 9, 2026
568aa5b
TEST/GTEST: Test host-based vmm registerability
tomerg-nvidia Jul 10, 2026
93c23d0
TEST/GTEST: Remove bad skip
tomerg-nvidia Jul 10, 2026
9ca826f
TEST/GTEST: Add a workaround for memtype cache pollution
tomerg-nvidia Jul 10, 2026
43a60f7
UCT/CUDA/CUDA_COPY: Remove shared context
tomerg-nvidia Jul 10, 2026
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
51 changes: 40 additions & 11 deletions src/uct/cuda/cuda_copy/cuda_copy_md.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -482,7 +482,8 @@ uct_cuda_copy_mem_release_fabric(uct_cuda_copy_alloc_handle_t *alloc_handle)

static int uct_cuda_copy_detect_vmm(const void *address,
ucs_memory_type_t *vmm_mem_type,
CUdevice *cuda_device)
CUdevice *cuda_device,
int *is_host_located)
{
#ifdef HAVE_CUMEMRETAINALLOCATIONHANDLE
CUmemGenericAllocationHandle alloc_handle;
Expand All @@ -497,8 +498,9 @@ static int uct_cuda_copy_detect_vmm(const void *address,
return 0;
}

*vmm_mem_type = UCS_MEMORY_TYPE_UNKNOWN;
*cuda_device = CU_DEVICE_INVALID;
*vmm_mem_type = UCS_MEMORY_TYPE_UNKNOWN;
*cuda_device = CU_DEVICE_INVALID;
*is_host_located = 0;

status = UCT_CUDADRV_FUNC_LOG_DEBUG(
cuMemGetAllocationPropertiesFromHandle(&prop, alloc_handle));
Expand All @@ -512,7 +514,8 @@ static int uct_cuda_copy_detect_vmm(const void *address,
(prop.location.type == CU_MEM_LOCATION_TYPE_HOST_NUMA) ||
(prop.location.type == CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT)) {
/* TODO: Marking as CUDA to allow cuda_ipc access vmm for now */
*vmm_mem_type = UCS_MEMORY_TYPE_CUDA;
*vmm_mem_type = UCS_MEMORY_TYPE_CUDA;
*is_host_located = 1;
} else
#endif
if (prop.location.type == CU_MEM_LOCATION_TYPE_DEVICE) {
Expand Down Expand Up @@ -650,7 +653,7 @@ static ucs_status_t
uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md,
const void *address, size_t length,
ucs_memory_info_t *mem_info,
int *is_async_managed)
int *is_async_managed, int *is_host_located)
{
#define UCT_CUDA_MEM_QUERY_NUM_ATTRS 4
CUmemorytype cuda_mem_type = CU_MEMORYTYPE_HOST;
Expand All @@ -665,8 +668,10 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md,
ucs_status_t status;

*is_async_managed = 0;
*is_host_located = 0;

is_vmm = uct_cuda_copy_detect_vmm(address, &mem_info->type, &cuda_device);
is_vmm = uct_cuda_copy_detect_vmm(address, &mem_info->type, &cuda_device,
is_host_located);
if (is_vmm) {
if (mem_info->type == UCS_MEMORY_TYPE_UNKNOWN) {
return UCS_ERR_INVALID_ADDR;
Expand Down Expand Up @@ -843,7 +848,7 @@ uct_cuda_copy_md_dmabuf_t uct_cuda_copy_md_get_dmabuf(const void *address,
static uint8_t
uct_cuda_copy_md_detect_mem_flags(uct_cuda_copy_md_t *md,
const ucs_memory_info_t *mem_info,
int is_async_managed,
int is_async_managed, int is_host_located,
const uct_cuda_copy_md_dmabuf_t *dmabuf)
{
int close_dmabuf = 0;
Expand All @@ -853,6 +858,11 @@ uct_cuda_copy_md_detect_mem_flags(uct_cuda_copy_md_t *md,
return 0;
}

/* Host-located CUDA VMM is registerable even if dmabuf export fails. */
if (is_host_located) {
return UCS_MEM_FLAG_REGISTRABLE;
}

if (mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) {
return UCS_MEM_FLAG_REGISTRABLE;
}
Expand Down Expand Up @@ -895,8 +905,11 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
.fd = UCT_DMABUF_FD_INVALID,
.offset = 0
};
int dmabuf_queried = 0;
int is_async_managed = 0;
int dmabuf_queried = 0;
int is_async_managed = 0;
int is_host_located = 0;
CUdevice cur_cuda_device = CU_DEVICE_INVALID;
CUdevice avail_cuda_device = CU_DEVICE_INVALID;
ucs_memory_info_t cached_mem_info;
ucs_memory_info_t addr_mem_info;
ucs_status_t cache_status;
Expand All @@ -918,7 +931,8 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
&cached_mem_info);
status = uct_cuda_copy_md_query_attributes(md, address, length,
&addr_mem_info,
&is_async_managed);
&is_async_managed,
&is_host_located);
if (status != UCS_OK) {
return status;
}
Expand Down Expand Up @@ -949,6 +963,17 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
mem_attr->alloc_length = addr_mem_info.alloc_length;
}

/* dmabuf export and mem_flags detection require a current context on the
* memory's device; reuse the current one if suitable, otherwise push it */
if (addr_mem_info.sys_dev != UCS_SYS_DEVICE_ID_UNKNOWN) {
if (uct_cuda_ctx_primary_push_avail(0, addr_mem_info.sys_dev,
&cur_cuda_device, &avail_cuda_device,
UCS_LOG_LEVEL_DEBUG) != UCS_OK) {
/* best-effort: continue with whatever context is current */
avail_cuda_device = cur_cuda_device;
}
}

if ((mem_attr->field_mask & UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_FD) ||
(mem_attr->field_mask & UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_OFFSET)) {
dmabuf = uct_cuda_copy_md_get_dmabuf(addr_mem_info.base_address,
Expand All @@ -967,7 +992,7 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,

if (address != NULL) {
addr_mem_info.mem_flags = uct_cuda_copy_md_detect_mem_flags(
md, &addr_mem_info, is_async_managed,
md, &addr_mem_info, is_async_managed, is_host_located,
dmabuf_queried ? &dmabuf : NULL);
ucs_memtype_cache_update(addr_mem_info.base_address,
addr_mem_info.alloc_length, addr_mem_info.type,
Expand All @@ -989,6 +1014,10 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
ucs_close_fd(&dmabuf.fd);
}

if (cur_cuda_device != avail_cuda_device) {
uct_cuda_ctx_primary_pop_and_release(avail_cuda_device);
}

return UCS_OK;
}

Expand Down
102 changes: 88 additions & 14 deletions test/gtest/uct/cuda/test_switch_cuda_device.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/**
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025. ALL RIGHTS RESERVED.
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025-2026. ALL RIGHTS RESERVED.
*
* See file LICENSE for terms.
*/
Expand Down Expand Up @@ -61,7 +61,8 @@ class cuda_vmm_mem_buffer {
void *ptr() const;

protected:
void init(size_t size, unsigned handle_type);
void init(size_t size, unsigned handle_type,
CUmemLocationType location_type = CU_MEM_LOCATION_TYPE_DEVICE);

private:
size_t m_size = 0;
Expand All @@ -87,20 +88,23 @@ void *cuda_vmm_mem_buffer::ptr() const
return (void*)m_ptr;
}

void cuda_vmm_mem_buffer::init(size_t size, unsigned handle_type)
void cuda_vmm_mem_buffer::init(size_t size, unsigned handle_type,
CUmemLocationType location_type)
{
size_t granularity = 0;
CUmemAllocationProp prop = {};
CUmemAccessDesc access_desc = {};
size_t granularity = 0;
CUmemAllocationProp prop = {};
CUmemAccessDesc access_desc[2] = {};
unsigned num_access = 1;
bool host_located = (location_type != CU_MEM_LOCATION_TYPE_DEVICE);
CUdevice device;
if (cuCtxGetDevice(&device) != CUDA_SUCCESS) {
UCS_TEST_ABORT("failed to get the device handle for the current "
"context");
}

prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
prop.location.id = device;
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = location_type;
prop.location.id = host_located ? 0 : device;
if (handle_type != 0) {
prop.requestedHandleTypes = (CUmemAllocationHandleType)handle_type;
}
Expand All @@ -123,10 +127,16 @@ void cuda_vmm_mem_buffer::init(size_t size, unsigned handle_type)
goto err_address_free;
}

access_desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
access_desc.location.id = device;
access_desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
if (cuMemSetAccess(m_ptr, m_size, &access_desc, 1) != CUDA_SUCCESS) {
access_desc[0].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
access_desc[0].location.id = device;
access_desc[0].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
if (host_located) {
access_desc[1].location.type = location_type;
access_desc[1].location.id = 0;
access_desc[1].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
num_access = 2;
}
if (cuMemSetAccess(m_ptr, m_size, access_desc, num_access) != CUDA_SUCCESS) {
goto err_mem_unmap;
}

Expand All @@ -139,7 +149,7 @@ void cuda_vmm_mem_buffer::init(size_t size, unsigned handle_type)
err_mem_release:
cuMemRelease(m_alloc_handle);
err:
UCS_TEST_SKIP_R("failed to allocate CUDA fabric memory");
UCS_TEST_SKIP_R("failed to allocate CUDA VMM memory");
}

#if HAVE_CUDA_FABRIC
Expand All @@ -155,6 +165,16 @@ cuda_fabric_mem_buffer::cuda_fabric_mem_buffer(size_t size,
}
#endif

#if CUDA_VERSION >= 12020
class cuda_host_vmm_mem_buffer : public cuda_vmm_mem_buffer {
public:
cuda_host_vmm_mem_buffer(size_t size, ucs_memory_type_t mem_type)
{
init(size, 0, CU_MEM_LOCATION_TYPE_HOST_NUMA);
}
};
#endif

UCS_TEST_P(test_switch_cuda_device, detect_mem_type_cuda)
{
detect_mem_type<mem_buffer>(UCS_MEMORY_TYPE_CUDA);
Expand Down Expand Up @@ -432,6 +452,22 @@ class test_mem_alloc_device : public test_switch_cuda_device {
#endif
}

void query_registrable_no_current_context(void *address, size_t size)
{
uct_md_mem_attr_v2_t mem_attr = {};
ucs_status_t query_status = UCS_ERR_NO_ELEM;

std::thread([&]() {
mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE |
UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS;
query_status = uct_md_mem_query_v2(md(), address, size, &mem_attr);
}).join();

ASSERT_UCS_OK(query_status);
EXPECT_EQ(UCS_MEMORY_TYPE_CUDA, mem_attr.mem_type);
EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE);
}

private:
std::vector<ucs_sys_device_t> m_sys_dev;

Expand Down Expand Up @@ -518,6 +554,44 @@ UCS_TEST_P(test_mem_alloc_device, no_current_context_cuda_registrable,
EXPECT_UCS_OK(free_status);
}

UCS_TEST_P(test_mem_alloc_device, no_current_context_user_mem_registrable,
"CUDA_COPY_ASYNC_MEM_TYPE=cuda")
{
constexpr size_t size = 4 * UCS_MBYTE;
CUdeviceptr dptr = 0;

ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc(&dptr, size));

query_registrable_no_current_context(reinterpret_cast<void*>(dptr), size);

EXPECT_EQ(CUDA_SUCCESS, cuMemFree(dptr));
}

UCS_TEST_P(test_mem_alloc_device, no_current_context_vmm_mem_registrable,
"CUDA_COPY_ASYNC_MEM_TYPE=cuda")
{
constexpr size_t size = 4 * UCS_MBYTE;
cuda_vmm_mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA);

query_registrable_no_current_context(buffer.ptr(), size);
}

#if CUDA_VERSION >= 12020
/* Host-located VMM is not dmabuf-exportable but is registrable by IB.
*
* TODO: REG_WHOLE_ALLOC=off is needed because of an existing issue with
* base_address being set to 0, causing cache pollution with this type of memory
*/
UCS_TEST_P(test_mem_alloc_device, host_vmm_mem_registrable,
"CUDA_COPY_REG_WHOLE_ALLOC=off")
{
constexpr size_t size = 4 * UCS_MBYTE;
cuda_host_vmm_mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA);

query_registrable_no_current_context(buffer.ptr(), size);
}
#endif

_UCT_MD_INSTANTIATE_TEST_CASE(test_mem_alloc_device, cuda_cpy);

class test_p2p_no_current_cuda_ctx : public uct_p2p_rma_test {
Expand Down