Skip to content

UCT/CUDA_COPY: Fix incorrect mem_flags detection - v1.22.x #11645

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 3 commits into from
Jul 10, 2026
+189 −25
Merged

UCT/CUDA_COPY: Fix incorrect mem_flags detection - v1.22.x #11645

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9c6d3e8
UCT/CUDA: Cache allocation flags with active context - v1.22.x
gleon99 Jul 9, 2026
83f5244
Merge branch 'v1.22.x' into port/11621-v1.22
gleon99 Jul 9, 2026
cf14e9c
UCT/CUDA_COPY: Fix incorrect mem_flags detection
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
68 changes: 57 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 @@ -352,6 +352,7 @@ uct_cuda_copy_mem_alloc(uct_md_h uct_md, size_t *length_p, void **address_p,
uct_mem_h *memh_p)
{
uct_cuda_copy_md_t *md = ucs_derived_of(uct_md, uct_cuda_copy_md_t);
uct_md_mem_attr_v2_t mem_attr;
ucs_status_t status;
uct_cuda_copy_alloc_handle_t *alloc_handle;
ucs_log_level_t log_level;
Expand Down Expand Up @@ -426,6 +427,22 @@ uct_cuda_copy_mem_alloc(uct_md_h uct_md, size_t *length_p, void **address_p,
allocated:
uct_cuda_copy_sync_memops(alloc_handle->ptr, alloc_handle->is_vmm);

/* Cache memory flags as part of uct_cuda_copy_md_mem_query() before
* restoring the CUDA context.
*/
mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE |
UCT_MD_MEM_ATTR_V2_FIELD_SYS_DEV |
UCT_MD_MEM_ATTR_V2_FIELD_BASE_ADDRESS |
UCT_MD_MEM_ATTR_V2_FIELD_ALLOC_LENGTH |
UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS;

status = uct_cuda_copy_md_mem_query(uct_md, (void*)alloc_handle->ptr,
alloc_handle->length, &mem_attr);
if (status != UCS_OK) {
(void)uct_md_mem_free(uct_md, alloc_handle);
goto out;
}

*memh_p = alloc_handle;
*address_p = (void*)alloc_handle->ptr;
*length_p = alloc_handle->length;
Expand Down Expand Up @@ -465,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 @@ -480,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 @@ -495,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 @@ -633,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 @@ -648,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 @@ -826,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 @@ -836,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 @@ -878,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 addr_mem_info;
ucs_status_t status;

Expand All @@ -897,7 +927,8 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
if (address != NULL) {
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 All @@ -921,6 +952,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 @@ -939,7 +981,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 @@ -961,6 +1003,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
146 changes: 132 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 All @@ -8,6 +8,7 @@
#include <uct/test_p2p_rma.h>

extern "C" {
#include <ucs/memory/memtype_cache.h>
#include <ucs/sys/ptr_arith.h>
#include <uct/base/uct_md.h>
}
Expand Down Expand Up @@ -60,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 @@ -86,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 @@ -122,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 @@ -138,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 @@ -154,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 @@ -431,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 @@ -474,6 +511,87 @@ UCS_TEST_P(test_mem_alloc_device, same_device_cuda_fabric_implicit,
test_same_device_alloc(UCS_MEMORY_TYPE_CUDA, false);
}

UCS_TEST_P(test_mem_alloc_device, no_current_context_cuda_registrable,
"CUDA_COPY_ASYNC_MEM_TYPE=cuda")
{
ucs_memory_info_t mem_info = {};
ucs_status_t lookup_status = UCS_ERR_NO_ELEM;
ucs_status_t alloc_status = UCS_ERR_NO_MEMORY;
CUresult ctx_status = CUDA_ERROR_UNKNOWN;
CUcontext cuda_ctx = nullptr;

/* UCP initializes the cache before allocating rendezvous fragments. */
ucs_memory_info_t init_mem_info = {};
(void)ucs_memtype_cache_lookup(&init_mem_info, sizeof(init_mem_info),
&init_mem_info);

std::thread([&]() {
ctx_status = cuCtxGetCurrent(&cuda_ctx);
if ((ctx_status != CUDA_SUCCESS) || (cuda_ctx != nullptr)) {
return;
}

alloc_status = allocate(UCS_MEMORY_TYPE_CUDA);
if (alloc_status != UCS_OK) {
return;
}

lookup_status = ucs_memtype_cache_lookup(mem.address, mem.length,
&mem_info);
}).join();

EXPECT_EQ(CUDA_SUCCESS, ctx_status);
EXPECT_EQ(nullptr, cuda_ctx);
ASSERT_UCS_OK(alloc_status);

ucs_status_t free_status = uct_mem_free(&mem);

EXPECT_UCS_OK(lookup_status);
if (lookup_status == UCS_OK) {
EXPECT_EQ(UCS_MEMORY_TYPE_CUDA, mem_info.type);
EXPECT_TRUE(mem_info.mem_flags & UCS_MEM_FLAG_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