From 9c6d3e8510487977791c3779c473e853a2c3d5e5 Mon Sep 17 00:00:00 2001 From: Leonid Genkin Date: Thu, 9 Jul 2026 12:23:14 +0300 Subject: [PATCH 1/2] UCT/CUDA: Cache allocation flags with active context - v1.22.x --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 17 +++++++ .../gtest/uct/cuda/test_switch_cuda_device.cc | 44 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index 80798e8a15a..de76992664c 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -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; @@ -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; diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index e5fa35275c1..c5a5ebf0643 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -8,6 +8,7 @@ #include extern "C" { +#include #include #include } @@ -474,6 +475,49 @@ 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); +} + _UCT_MD_INSTANTIATE_TEST_CASE(test_mem_alloc_device, cuda_cpy); class test_p2p_no_current_cuda_ctx : public uct_p2p_rma_test { From cf14e9cbf0a4a67855790f72a000df7a611cf24d Mon Sep 17 00:00:00 2001 From: Tomer Gilad Date: Fri, 10 Jul 2026 18:19:01 +0200 Subject: [PATCH 2/2] UCT/CUDA_COPY: Fix incorrect mem_flags detection --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 51 +++++++-- .../gtest/uct/cuda/test_switch_cuda_device.cc | 102 +++++++++++++++--- 2 files changed, 128 insertions(+), 25 deletions(-) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index de76992664c..d0141a61319 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -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; @@ -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)); @@ -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) { @@ -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; @@ -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; @@ -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; @@ -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; } @@ -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 addr_mem_info; ucs_status_t status; @@ -914,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; } @@ -938,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, @@ -956,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, @@ -978,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; } diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index c5a5ebf0643..1bf3a7cdfb8 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -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. */ @@ -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; @@ -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; } @@ -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; } @@ -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 @@ -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(UCS_MEMORY_TYPE_CUDA); @@ -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 m_sys_dev; @@ -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(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 {