From d72bcfb265ccb7c55e7850399b70025a313f2827 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 8 Sep 2026 10:26:41 +0000 Subject: [PATCH 01/19] UCT/CUDA_COPY: Skip dmabuf probe for async managed memory CUDA async managed allocations can report CU_POINTER_ATTRIBUTE_IS_MANAGED with no owning CUDA context. Mark this case as async-managed so mem_flags detection keeps it non-registerable and does not call cuMemGetHandleForAddressRange(). This avoids a CUDA driver crash during UCP memory detection when UCX_CUDA_COPY_DMABUF is enabled. Add a CUDA 13 regression test that queries mem_flags for async managed mempool memory with dmabuf enabled. --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 3 + .../gtest/uct/cuda/test_switch_cuda_device.cc | 63 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index 3ae04fa2363..d5881c40fac 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -730,6 +730,9 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, * provided address and length as base address and alloc length * respectively */ mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; + if (cuda_mem_ctx == NULL) { + *is_async_managed = 1; + } cu_err = cuMemRangeGetAttribute( (void*)&pref_loc, sizeof(pref_loc), diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index 12afb0ef78e..633e8bc3bd5 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -441,6 +441,69 @@ UCS_TEST_P(test_mem_alloc_device, no_current_context_user_mem_registrable, EXPECT_EQ(CUDA_SUCCESS, cuMemFree(dptr)); } +#if CUDART_VERSION >= 13000 +UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_not_registrable, + "CUDA_COPY_DMABUF=try") +{ + constexpr size_t size = 192; + cudaMemLocation location = {}; + cudaStream_t stream = nullptr; + cudaMemPool_t pool = nullptr; + void *buffer = nullptr; + int cuda_device; + int concurrent_managed_access; + uct_md_mem_attr_v2_t mem_attr = {}; + + ASSERT_EQ(cudaSuccess, cudaGetDevice(&cuda_device)); + ASSERT_EQ(cudaSuccess, + cudaDeviceGetAttribute(&concurrent_managed_access, + cudaDevAttrConcurrentManagedAccess, + cuda_device)); + if (!concurrent_managed_access) { + UCS_TEST_SKIP_R("device does not support concurrent managed access"); + } + + location.type = cudaMemLocationTypeDevice; + location.id = cuda_device; + ASSERT_EQ(cudaSuccess, + cudaMemGetDefaultMemPool(&pool, &location, + cudaMemAllocationTypeManaged)); + ASSERT_EQ(cudaSuccess, cudaStreamCreate(&stream)); + + cudaError_t cuda_status = cudaMallocFromPoolAsync(&buffer, size, pool, + stream); + if (cuda_status != cudaSuccess) { + EXPECT_EQ(cudaSuccess, cudaStreamDestroy(stream)); + ASSERT_EQ(cudaSuccess, cuda_status); + } + + cuda_status = cudaMemsetAsync(buffer, 0, size, stream); + if (cuda_status != cudaSuccess) { + EXPECT_EQ(cudaSuccess, cudaFreeAsync(buffer, stream)); + EXPECT_EQ(cudaSuccess, cudaStreamSynchronize(stream)); + EXPECT_EQ(cudaSuccess, cudaStreamDestroy(stream)); + ASSERT_EQ(cudaSuccess, cuda_status); + } + + cuda_status = cudaStreamSynchronize(stream); + if (cuda_status != cudaSuccess) { + EXPECT_EQ(cudaSuccess, cudaFreeAsync(buffer, stream)); + EXPECT_EQ(cudaSuccess, cudaStreamDestroy(stream)); + ASSERT_EQ(cudaSuccess, cuda_status); + } + + mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE | + UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS; + EXPECT_UCS_OK(uct_md_mem_query_v2(md(), buffer, size, &mem_attr)); + EXPECT_EQ(UCS_MEMORY_TYPE_CUDA_MANAGED, mem_attr.mem_type); + EXPECT_FALSE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); + + EXPECT_EQ(cudaSuccess, cudaFreeAsync(buffer, stream)); + EXPECT_EQ(cudaSuccess, cudaStreamSynchronize(stream)); + EXPECT_EQ(cudaSuccess, cudaStreamDestroy(stream)); +} +#endif + UCS_TEST_P(test_mem_alloc_device, no_current_context_vmm_mem_registrable, "CUDA_COPY_ASYNC_MEM_TYPE=cuda") { From 1d53937e0f0cd7d37582b3eb8d1912eb8664a7bd Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 8 Sep 2026 13:11:31 +0000 Subject: [PATCH 02/19] UCT/CUDA_COPY: Address async managed memory review Respect CUDA_COPY_ASYNC_MEM_TYPE when marking managed allocations without a CUDA context as async-managed. Add coverage for ordinary cudaMallocManaged memory remaining registrable and make the async managed mempool test skip when the device or driver does not support that allocation path. --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 4 +- .../gtest/uct/cuda/test_switch_cuda_device.cc | 38 ++++++++++++++++--- 2 files changed, 35 insertions(+), 7 deletions(-) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index d5881c40fac..053f5ea1034 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -730,7 +730,9 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, * provided address and length as base address and alloc length * respectively */ mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; - if (cuda_mem_ctx == NULL) { + if ((cuda_mem_ctx == NULL) && md->config.cuda_async_managed) { + /* Managed pool allocations are stream-ordered and do not have + * an owning CUDA context. */ *is_async_managed = 1; } diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index 633e8bc3bd5..1a6bcdb7be0 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -441,6 +441,28 @@ UCS_TEST_P(test_mem_alloc_device, no_current_context_user_mem_registrable, EXPECT_EQ(CUDA_SUCCESS, cuMemFree(dptr)); } +UCS_TEST_P(test_mem_alloc_device, managed_mem_registrable, + "CUDA_COPY_PREF_LOC=cpu") +{ + constexpr size_t size = 4 * UCS_MBYTE; + uct_md_mem_attr_v2_t mem_attr = {}; + void *buffer = nullptr; + cudaError_t cuda_status; + + cuda_status = cudaMallocManaged(&buffer, size); + if (cuda_status != cudaSuccess) { + UCS_TEST_SKIP_R("failed to allocate managed memory"); + } + + mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE | + UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS; + EXPECT_UCS_OK(uct_md_mem_query_v2(md(), buffer, size, &mem_attr)); + EXPECT_EQ(UCS_MEMORY_TYPE_CUDA_MANAGED, mem_attr.mem_type); + EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); + + EXPECT_EQ(cudaSuccess, cudaFree(buffer)); +} + #if CUDART_VERSION >= 13000 UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_not_registrable, "CUDA_COPY_DMABUF=try") @@ -453,6 +475,7 @@ UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_not_registrable, int cuda_device; int concurrent_managed_access; uct_md_mem_attr_v2_t mem_attr = {}; + cudaError_t cuda_status; ASSERT_EQ(cudaSuccess, cudaGetDevice(&cuda_device)); ASSERT_EQ(cudaSuccess, @@ -465,16 +488,18 @@ UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_not_registrable, location.type = cudaMemLocationTypeDevice; location.id = cuda_device; - ASSERT_EQ(cudaSuccess, - cudaMemGetDefaultMemPool(&pool, &location, - cudaMemAllocationTypeManaged)); + cuda_status = cudaMemGetDefaultMemPool(&pool, &location, + cudaMemAllocationTypeManaged); + if (cuda_status != cudaSuccess) { + UCS_TEST_SKIP_R("failed to get default managed memory pool"); + } + ASSERT_EQ(cudaSuccess, cudaStreamCreate(&stream)); - cudaError_t cuda_status = cudaMallocFromPoolAsync(&buffer, size, pool, - stream); + cuda_status = cudaMallocFromPoolAsync(&buffer, size, pool, stream); if (cuda_status != cudaSuccess) { EXPECT_EQ(cudaSuccess, cudaStreamDestroy(stream)); - ASSERT_EQ(cudaSuccess, cuda_status); + UCS_TEST_SKIP_R("failed to allocate async managed mempool memory"); } cuda_status = cudaMemsetAsync(buffer, 0, size, stream); @@ -488,6 +513,7 @@ UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_not_registrable, cuda_status = cudaStreamSynchronize(stream); if (cuda_status != cudaSuccess) { EXPECT_EQ(cudaSuccess, cudaFreeAsync(buffer, stream)); + EXPECT_EQ(cudaSuccess, cudaStreamSynchronize(stream)); EXPECT_EQ(cudaSuccess, cudaStreamDestroy(stream)); ASSERT_EQ(cudaSuccess, cuda_status); } From 458aeca5387f3c6adb60c8e12bc709d336a296cf Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 8 Sep 2026 14:52:45 +0000 Subject: [PATCH 03/19] GTEST/CUDA: Cover managed memory registration flags Add positive coverage for ordinary managed memory so CUDA managed allocations with an owning context keep reporting UCS_MEM_FLAG_REGISTRABLE. Cover both user cudaMallocManaged memory and memory allocated through uct_mem_alloc. --- .../gtest/uct/cuda/test_switch_cuda_device.cc | 39 +++++++++++++++---- 1 file changed, 31 insertions(+), 8 deletions(-) diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index 1a6bcdb7be0..20cbc0d099b 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -342,6 +342,24 @@ class test_mem_alloc_device : public test_switch_cuda_device { EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); } + void query_managed_registrable(void *address, size_t size) + { + CUpointer_attribute attr_type = CU_POINTER_ATTRIBUTE_CONTEXT; + uct_md_mem_attr_v2_t mem_attr = {}; + CUcontext cuda_mem_ctx = nullptr; + + ASSERT_EQ(CUDA_SUCCESS, + cuPointerGetAttribute(&cuda_mem_ctx, attr_type, + (CUdeviceptr)address)); + EXPECT_NE(nullptr, cuda_mem_ctx); + + mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE | + UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS; + EXPECT_UCS_OK(uct_md_mem_query_v2(md(), address, size, &mem_attr)); + EXPECT_EQ(UCS_MEMORY_TYPE_CUDA_MANAGED, mem_attr.mem_type); + EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); + } + private: std::vector m_sys_dev; @@ -441,12 +459,11 @@ UCS_TEST_P(test_mem_alloc_device, no_current_context_user_mem_registrable, EXPECT_EQ(CUDA_SUCCESS, cuMemFree(dptr)); } -UCS_TEST_P(test_mem_alloc_device, managed_mem_registrable, +UCS_TEST_P(test_mem_alloc_device, user_managed_mem_registrable, "CUDA_COPY_PREF_LOC=cpu") { constexpr size_t size = 4 * UCS_MBYTE; - uct_md_mem_attr_v2_t mem_attr = {}; - void *buffer = nullptr; + void *buffer = nullptr; cudaError_t cuda_status; cuda_status = cudaMallocManaged(&buffer, size); @@ -454,15 +471,21 @@ UCS_TEST_P(test_mem_alloc_device, managed_mem_registrable, UCS_TEST_SKIP_R("failed to allocate managed memory"); } - mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE | - UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS; - EXPECT_UCS_OK(uct_md_mem_query_v2(md(), buffer, size, &mem_attr)); - EXPECT_EQ(UCS_MEMORY_TYPE_CUDA_MANAGED, mem_attr.mem_type); - EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); + query_managed_registrable(buffer, size); EXPECT_EQ(cudaSuccess, cudaFree(buffer)); } +UCS_TEST_P(test_mem_alloc_device, uct_alloc_managed_mem_registrable, + "CUDA_COPY_PREF_LOC=cpu") +{ + ASSERT_UCS_OK(allocate(UCS_MEMORY_TYPE_CUDA_MANAGED)); + + query_managed_registrable(mem.address, mem.length); + + EXPECT_UCS_OK(uct_mem_free(&mem)); +} + #if CUDART_VERSION >= 13000 UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_not_registrable, "CUDA_COPY_DMABUF=try") From 31617b6ada38cd614c84db35132a80ce6a77cf3b Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 8 Sep 2026 16:10:13 +0000 Subject: [PATCH 04/19] UCT/CUDA_COPY: Detect managed memory pools Use the CUDA mempool pointer attribute to identify managed stream-ordered pool allocations. Such allocations are not registrable regardless of CUDA_COPY_ASYNC_MEM_TYPE, so skip dmabuf probing for them. Extend the regression test to cover CUDA_COPY_ASYNC_MEM_TYPE=cuda and use RAII for CUDA test allocations. --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 11 +- .../gtest/uct/cuda/test_switch_cuda_device.cc | 174 +++++++++++------- 2 files changed, 111 insertions(+), 74 deletions(-) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index 053f5ea1034..cf644f228e6 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -680,10 +680,11 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, ucs_memory_info_t *mem_info, int *is_async_managed, int *is_host_located) { -#define UCT_CUDA_MEM_QUERY_NUM_ATTRS 4 +#define UCT_CUDA_MEM_QUERY_NUM_ATTRS 5 CUmemorytype cuda_mem_type = CU_MEMORYTYPE_HOST; uint32_t is_managed = 0; CUcontext cuda_mem_ctx = NULL; + CUmemoryPool cuda_mempool = NULL; CUpointer_attribute attr_type[UCT_CUDA_MEM_QUERY_NUM_ATTRS]; void *attr_data[UCT_CUDA_MEM_QUERY_NUM_ATTRS]; CUdevice cuda_device; @@ -710,6 +711,8 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, attr_data[2] = &cuda_device; attr_type[3] = CU_POINTER_ATTRIBUTE_CONTEXT; attr_data[3] = &cuda_mem_ctx; + attr_type[4] = CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE; + attr_data[4] = &cuda_mempool; status = UCT_CUDADRV_FUNC_LOG_ERR( cuPointerGetAttributes(ucs_static_array_size(attr_data), @@ -730,9 +733,9 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, * provided address and length as base address and alloc length * respectively */ mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; - if ((cuda_mem_ctx == NULL) && md->config.cuda_async_managed) { - /* Managed pool allocations are stream-ordered and do not have - * an owning CUDA context. */ + if (cuda_mempool != NULL) { + /* Managed-pool allocations are stream-ordered and cannot be + * registered. */ *is_async_managed = 1; } diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index 20cbc0d099b..92fd36b91f1 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -19,6 +19,56 @@ extern "C" { #include +#if CUDART_VERSION >= 13000 +class cuda_async_managed_mem_buffer { +public: + cuda_async_managed_mem_buffer(cudaMemPool_t pool, size_t size) : + m_buffer(nullptr), m_stream(nullptr), m_status(cudaStreamCreate(&m_stream)) + { + if (m_status == cudaSuccess) { + m_status = cudaMallocFromPoolAsync(&m_buffer, size, pool, m_stream); + } + } + + ~cuda_async_managed_mem_buffer() + { + if (m_buffer != nullptr) { + cudaFreeAsync(m_buffer, m_stream); + } + if (m_stream != nullptr) { + cudaStreamSynchronize(m_stream); + cudaStreamDestroy(m_stream); + } + } + + void *ptr() const + { + return m_buffer; + } + + cudaError_t status() const + { + return m_status; + } + + cudaError_t memset(int value, size_t size) + { + return cudaMemsetAsync(m_buffer, value, size, m_stream); + } + + cudaError_t synchronize() + { + return cudaStreamSynchronize(m_stream); + } + +private: + void *m_buffer; + cudaStream_t m_stream; + cudaError_t m_status; +}; +#endif + + class test_switch_cuda_device : public test_md { protected: template void detect_mem_type(ucs_memory_type_t mem_type) const; @@ -360,6 +410,50 @@ class test_mem_alloc_device : public test_switch_cuda_device { EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); } +#if CUDART_VERSION >= 13000 + void test_async_managed_mem_pool_not_registrable() + { + constexpr size_t size = 192; + cudaMemLocation location = {}; + cudaMemPool_t pool = nullptr; + int cuda_device; + int concurrent_managed_access; + uct_md_mem_attr_v2_t mem_attr = {}; + cudaError_t cuda_status; + + ASSERT_EQ(cudaSuccess, cudaGetDevice(&cuda_device)); + ASSERT_EQ(cudaSuccess, + cudaDeviceGetAttribute(&concurrent_managed_access, + cudaDevAttrConcurrentManagedAccess, + cuda_device)); + if (!concurrent_managed_access) { + UCS_TEST_SKIP_R("device does not support concurrent managed access"); + } + + location.type = cudaMemLocationTypeDevice; + location.id = cuda_device; + cuda_status = cudaMemGetDefaultMemPool(&pool, &location, + cudaMemAllocationTypeManaged); + if (cuda_status != cudaSuccess) { + UCS_TEST_SKIP_R("failed to get default managed memory pool"); + } + + cuda_async_managed_mem_buffer buffer(pool, size); + if (buffer.status() != cudaSuccess) { + UCS_TEST_SKIP_R("failed to allocate async managed mempool memory"); + } + + ASSERT_EQ(cudaSuccess, buffer.memset(0, size)); + ASSERT_EQ(cudaSuccess, buffer.synchronize()); + + mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE | + UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS; + EXPECT_UCS_OK(uct_md_mem_query_v2(md(), buffer.ptr(), size, &mem_attr)); + EXPECT_EQ(UCS_MEMORY_TYPE_CUDA_MANAGED, mem_attr.mem_type); + EXPECT_FALSE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); + } +#endif + private: std::vector m_sys_dev; @@ -463,17 +557,9 @@ UCS_TEST_P(test_mem_alloc_device, user_managed_mem_registrable, "CUDA_COPY_PREF_LOC=cpu") { constexpr size_t size = 4 * UCS_MBYTE; - void *buffer = nullptr; - cudaError_t cuda_status; + mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA_MANAGED); - cuda_status = cudaMallocManaged(&buffer, size); - if (cuda_status != cudaSuccess) { - UCS_TEST_SKIP_R("failed to allocate managed memory"); - } - - query_managed_registrable(buffer, size); - - EXPECT_EQ(cudaSuccess, cudaFree(buffer)); + query_managed_registrable(buffer.ptr(), size); } UCS_TEST_P(test_mem_alloc_device, uct_alloc_managed_mem_registrable, @@ -490,66 +576,14 @@ UCS_TEST_P(test_mem_alloc_device, uct_alloc_managed_mem_registrable, UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_not_registrable, "CUDA_COPY_DMABUF=try") { - constexpr size_t size = 192; - cudaMemLocation location = {}; - cudaStream_t stream = nullptr; - cudaMemPool_t pool = nullptr; - void *buffer = nullptr; - int cuda_device; - int concurrent_managed_access; - uct_md_mem_attr_v2_t mem_attr = {}; - cudaError_t cuda_status; - - ASSERT_EQ(cudaSuccess, cudaGetDevice(&cuda_device)); - ASSERT_EQ(cudaSuccess, - cudaDeviceGetAttribute(&concurrent_managed_access, - cudaDevAttrConcurrentManagedAccess, - cuda_device)); - if (!concurrent_managed_access) { - UCS_TEST_SKIP_R("device does not support concurrent managed access"); - } - - location.type = cudaMemLocationTypeDevice; - location.id = cuda_device; - cuda_status = cudaMemGetDefaultMemPool(&pool, &location, - cudaMemAllocationTypeManaged); - if (cuda_status != cudaSuccess) { - UCS_TEST_SKIP_R("failed to get default managed memory pool"); - } - - ASSERT_EQ(cudaSuccess, cudaStreamCreate(&stream)); - - cuda_status = cudaMallocFromPoolAsync(&buffer, size, pool, stream); - if (cuda_status != cudaSuccess) { - EXPECT_EQ(cudaSuccess, cudaStreamDestroy(stream)); - UCS_TEST_SKIP_R("failed to allocate async managed mempool memory"); - } - - cuda_status = cudaMemsetAsync(buffer, 0, size, stream); - if (cuda_status != cudaSuccess) { - EXPECT_EQ(cudaSuccess, cudaFreeAsync(buffer, stream)); - EXPECT_EQ(cudaSuccess, cudaStreamSynchronize(stream)); - EXPECT_EQ(cudaSuccess, cudaStreamDestroy(stream)); - ASSERT_EQ(cudaSuccess, cuda_status); - } - - cuda_status = cudaStreamSynchronize(stream); - if (cuda_status != cudaSuccess) { - EXPECT_EQ(cudaSuccess, cudaFreeAsync(buffer, stream)); - EXPECT_EQ(cudaSuccess, cudaStreamSynchronize(stream)); - EXPECT_EQ(cudaSuccess, cudaStreamDestroy(stream)); - ASSERT_EQ(cudaSuccess, cuda_status); - } - - mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE | - UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS; - EXPECT_UCS_OK(uct_md_mem_query_v2(md(), buffer, size, &mem_attr)); - EXPECT_EQ(UCS_MEMORY_TYPE_CUDA_MANAGED, mem_attr.mem_type); - EXPECT_FALSE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); - - EXPECT_EQ(cudaSuccess, cudaFreeAsync(buffer, stream)); - EXPECT_EQ(cudaSuccess, cudaStreamSynchronize(stream)); - EXPECT_EQ(cudaSuccess, cudaStreamDestroy(stream)); + test_async_managed_mem_pool_not_registrable(); +} + +UCS_TEST_P(test_mem_alloc_device, + async_managed_mem_pool_not_registrable_cuda_type, + "CUDA_COPY_DMABUF=try", "CUDA_COPY_ASYNC_MEM_TYPE=cuda") +{ + test_async_managed_mem_pool_not_registrable(); } #endif From 0fb53b39c58703925a64339b40619952089f3d1a Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 8 Sep 2026 16:38:18 +0000 Subject: [PATCH 05/19] UCT/CUDA_COPY: Preserve managed pool device locality Report the GPU system device for managed pool allocations so the non-registrable cache entry does not trigger UCP memory detection on every operation. Guard mempool pointer attributes for CUDA versions before 11.2. --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 15 +++++++++++++++ test/gtest/uct/cuda/test_switch_cuda_device.cc | 2 ++ 2 files changed, 17 insertions(+) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index cf644f228e6..74656a94a43 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -680,11 +680,17 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, ucs_memory_info_t *mem_info, int *is_async_managed, int *is_host_located) { +#if CUDA_VERSION >= 11020 #define UCT_CUDA_MEM_QUERY_NUM_ATTRS 5 +#else +#define UCT_CUDA_MEM_QUERY_NUM_ATTRS 4 +#endif CUmemorytype cuda_mem_type = CU_MEMORYTYPE_HOST; uint32_t is_managed = 0; CUcontext cuda_mem_ctx = NULL; +#if CUDA_VERSION >= 11020 CUmemoryPool cuda_mempool = NULL; +#endif CUpointer_attribute attr_type[UCT_CUDA_MEM_QUERY_NUM_ATTRS]; void *attr_data[UCT_CUDA_MEM_QUERY_NUM_ATTRS]; CUdevice cuda_device; @@ -711,8 +717,10 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, attr_data[2] = &cuda_device; attr_type[3] = CU_POINTER_ATTRIBUTE_CONTEXT; attr_data[3] = &cuda_mem_ctx; +#if CUDA_VERSION >= 11020 attr_type[4] = CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE; attr_data[4] = &cuda_mempool; +#endif status = UCT_CUDADRV_FUNC_LOG_ERR( cuPointerGetAttributes(ucs_static_array_size(attr_data), @@ -733,11 +741,18 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, * provided address and length as base address and alloc length * respectively */ mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; +#if CUDA_VERSION >= 11020 if (cuda_mempool != NULL) { /* Managed-pool allocations are stream-ordered and cannot be * registered. */ *is_async_managed = 1; + mem_info->sys_dev = uct_cuda_get_sys_dev(cuda_device); + if (mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) { + return UCS_ERR_NO_DEVICE; + } + goto out_default_range; } +#endif cu_err = cuMemRangeGetAttribute( (void*)&pref_loc, sizeof(pref_loc), diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index 92fd36b91f1..778667bacca 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -447,9 +447,11 @@ class test_mem_alloc_device : public test_switch_cuda_device { ASSERT_EQ(cudaSuccess, buffer.synchronize()); 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_MEM_FLAGS; EXPECT_UCS_OK(uct_md_mem_query_v2(md(), buffer.ptr(), size, &mem_attr)); EXPECT_EQ(UCS_MEMORY_TYPE_CUDA_MANAGED, mem_attr.mem_type); + EXPECT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, mem_attr.sys_dev); EXPECT_FALSE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); } #endif From 1b27615b0f45f5392da6ad4fc3dda1cd0df7d095 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 8 Sep 2026 17:35:25 +0000 Subject: [PATCH 06/19] UCT/CUDA_COPY: Handle managed pool CPU devices Preserve the common CUDA CPU-device handling for managed pool allocations and cover GPU preferred location in the managed pool regression test. --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 10 +++++++--- test/gtest/uct/cuda/test_switch_cuda_device.cc | 6 ++++++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index 74656a94a43..0c18b63ba20 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -746,9 +746,13 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, /* Managed-pool allocations are stream-ordered and cannot be * registered. */ *is_async_managed = 1; - mem_info->sys_dev = uct_cuda_get_sys_dev(cuda_device); - if (mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) { - return UCS_ERR_NO_DEVICE; + if (cuda_device == CU_DEVICE_CPU) { + mem_info->sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + } else { + mem_info->sys_dev = uct_cuda_get_sys_dev(cuda_device); + if (mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) { + return UCS_ERR_NO_DEVICE; + } } goto out_default_range; } diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index 778667bacca..a7fd04b4657 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -581,6 +581,12 @@ UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_not_registrable, test_async_managed_mem_pool_not_registrable(); } +UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_gpu_pref_loc, + "CUDA_COPY_DMABUF=try", "CUDA_COPY_PREF_LOC=gpu") +{ + test_async_managed_mem_pool_not_registrable(); +} + UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_not_registrable_cuda_type, "CUDA_COPY_DMABUF=try", "CUDA_COPY_ASYNC_MEM_TYPE=cuda") From 12b311635337bbfb08b8df2def16761fa614440b Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 8 Sep 2026 19:15:49 +0000 Subject: [PATCH 07/19] UCT/CUDA_COPY: Keep managed-pool detection resilient --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index 0c18b63ba20..a9858e236f5 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -744,14 +744,17 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, #if CUDA_VERSION >= 11020 if (cuda_mempool != NULL) { /* Managed-pool allocations are stream-ordered and cannot be - * registered. */ + * registered. Use the pool device to avoid re-detection when + * CPU is the preferred location. */ *is_async_managed = 1; if (cuda_device == CU_DEVICE_CPU) { mem_info->sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; } else { mem_info->sys_dev = uct_cuda_get_sys_dev(cuda_device); if (mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) { - return UCS_ERR_NO_DEVICE; + ucs_diag("cu_device %d (for address %p...%p) " + "unrecognized", cuda_device, address, + UCS_PTR_BYTE_OFFSET(address, length)); } } goto out_default_range; From 07409aeaf6eac0333f44897ded87dffe53d10f5a Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 8 Sep 2026 20:12:46 +0000 Subject: [PATCH 08/19] UCT/CUDA_COPY: Use mempool attributes for detection --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 48 ++++++++----------- .../gtest/uct/cuda/test_switch_cuda_device.cc | 3 +- 2 files changed, 23 insertions(+), 28 deletions(-) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index a9858e236f5..74a863e9b97 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -747,28 +747,21 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, * registered. Use the pool device to avoid re-detection when * CPU is the preferred location. */ *is_async_managed = 1; - if (cuda_device == CU_DEVICE_CPU) { - mem_info->sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; - } else { - mem_info->sys_dev = uct_cuda_get_sys_dev(cuda_device); - if (mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) { - ucs_diag("cu_device %d (for address %p...%p) " - "unrecognized", cuda_device, address, - UCS_PTR_BYTE_OFFSET(address, length)); - } - } - goto out_default_range; - } + pref_loc = cuda_device; + } else #endif - cu_err = cuMemRangeGetAttribute( - (void*)&pref_loc, sizeof(pref_loc), - CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION, - (CUdeviceptr)address, length); - if ((cu_err != CUDA_SUCCESS) || (pref_loc == CU_DEVICE_INVALID)) { - pref_loc = (md->config.pref_loc == UCT_CUDA_PREF_LOC_CPU) ? - CU_DEVICE_CPU : - cuda_device; + { + cu_err = cuMemRangeGetAttribute( + (void*)&pref_loc, sizeof(pref_loc), + CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION, + (CUdeviceptr)address, length); + if ((cu_err != CUDA_SUCCESS) || + (pref_loc == CU_DEVICE_INVALID)) { + pref_loc = (md->config.pref_loc == UCT_CUDA_PREF_LOC_CPU) ? + CU_DEVICE_CPU : + cuda_device; + } } if (pref_loc == CU_DEVICE_CPU) { @@ -783,15 +776,16 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, } goto out_default_range; +#if CUDA_VERSION >= 11020 + } else if ((cuda_mempool != NULL) && md->config.cuda_async_managed) { +#else } else if ((cuda_mem_ctx == NULL) && md->config.cuda_async_managed) { - /* Currently virtual/stream-ordered CUDA allocations are typed as +#endif + /* Stream-ordered CUDA allocations are typed as * `UCS_MEMORY_TYPE_CUDA_MANAGED`. This may be changed using - * UCX_CUDA_COPY_ASYNC_MEM_TYPE env var. Ideally checking for - * `CU_POINTER_ATTRIBUTE_IS_LEGACY_CUDA_IPC_CAPABLE` would be better - * here, but due to a bug in the driver `cudaMalloc` also returns - * false in that case. Therefore, checking whether the allocation - * was not allocated in a context should also allows us to - * identify virtual/stream-ordered CUDA allocations. */ + * UCX_CUDA_COPY_ASYNC_MEM_TYPE env var. CUDA versions before 11.2 + * do not provide the memory-pool attribute, so use the missing + * context to identify these allocations. */ mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; *is_async_managed = 1; } else { diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index a7fd04b4657..af72c487eda 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -23,7 +23,8 @@ extern "C" { class cuda_async_managed_mem_buffer { public: cuda_async_managed_mem_buffer(cudaMemPool_t pool, size_t size) : - m_buffer(nullptr), m_stream(nullptr), m_status(cudaStreamCreate(&m_stream)) + m_buffer(nullptr), m_stream(nullptr), + m_status(cudaStreamCreate(&m_stream)) { if (m_status == cudaSuccess) { m_status = cudaMallocFromPoolAsync(&m_buffer, size, pool, m_stream); From cfc5a8e46aaca6e68410f7f936acfb62963ff21a Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 9 Sep 2026 08:46:26 +0000 Subject: [PATCH 09/19] UCT/CUDA_COPY: Detect managed memory pools by capability --- config/m4/cuda.m4 | 2 + src/uct/cuda/cuda_copy/cuda_copy_md.c | 48 ++++------ test/gtest/common/mem_buffer.cc | 64 ++++++++++++-- test/gtest/common/mem_buffer.h | 2 +- .../gtest/uct/cuda/test_switch_cuda_device.cc | 88 ++----------------- 5 files changed, 87 insertions(+), 117 deletions(-) diff --git a/config/m4/cuda.m4 b/config/m4/cuda.m4 index 553bd212d89..a510414b902 100644 --- a/config/m4/cuda.m4 +++ b/config/m4/cuda.m4 @@ -210,6 +210,8 @@ AS_IF([test "x$cuda_checked" != "xyes"], [AC_DEFINE([HAVE_CUMEMRETAINALLOCATIONHANDLE], [1], [Enable cuMemRetainAllocationHandle() usage])]) AC_CHECK_DECLS([CU_MEM_LOCATION_TYPE_HOST], + [], [], [[#include ]]) + AC_CHECK_DECLS([CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE], [], [], [[#include ]])]) # Check nvml header files diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index 74a863e9b97..d9607cf9818 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -680,19 +680,14 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, ucs_memory_info_t *mem_info, int *is_async_managed, int *is_host_located) { -#if CUDA_VERSION >= 11020 -#define UCT_CUDA_MEM_QUERY_NUM_ATTRS 5 -#else -#define UCT_CUDA_MEM_QUERY_NUM_ATTRS 4 -#endif +#define UCT_CUDA_MEM_QUERY_MAX_ATTRS 5 CUmemorytype cuda_mem_type = CU_MEMORYTYPE_HOST; uint32_t is_managed = 0; CUcontext cuda_mem_ctx = NULL; -#if CUDA_VERSION >= 11020 - CUmemoryPool cuda_mempool = NULL; -#endif - CUpointer_attribute attr_type[UCT_CUDA_MEM_QUERY_NUM_ATTRS]; - void *attr_data[UCT_CUDA_MEM_QUERY_NUM_ATTRS]; + void *cuda_mempool = NULL; + unsigned num_attrs = 4; + CUpointer_attribute attr_type[UCT_CUDA_MEM_QUERY_MAX_ATTRS]; + void *attr_data[UCT_CUDA_MEM_QUERY_MAX_ATTRS]; CUdevice cuda_device; int32_t pref_loc; int is_vmm; @@ -717,14 +712,14 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, attr_data[2] = &cuda_device; attr_type[3] = CU_POINTER_ATTRIBUTE_CONTEXT; attr_data[3] = &cuda_mem_ctx; -#if CUDA_VERSION >= 11020 - attr_type[4] = CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE; - attr_data[4] = &cuda_mempool; +#if HAVE_DECL_CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE + attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE; + attr_data[num_attrs] = &cuda_mempool; + num_attrs++; #endif status = UCT_CUDADRV_FUNC_LOG_ERR( - cuPointerGetAttributes(ucs_static_array_size(attr_data), - attr_type, attr_data, + cuPointerGetAttributes(num_attrs, attr_type, attr_data, (CUdeviceptr)address)); if (status != UCS_OK) { /* pointer not recognized */ @@ -741,17 +736,13 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, * provided address and length as base address and alloc length * respectively */ mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; -#if CUDA_VERSION >= 11020 if (cuda_mempool != NULL) { /* Managed-pool allocations are stream-ordered and cannot be * registered. Use the pool device to avoid re-detection when * CPU is the preferred location. */ *is_async_managed = 1; - pref_loc = cuda_device; - } else -#endif - - { + pref_loc = cuda_device; + } else { cu_err = cuMemRangeGetAttribute( (void*)&pref_loc, sizeof(pref_loc), CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION, @@ -776,16 +767,15 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, } goto out_default_range; -#if CUDA_VERSION >= 11020 - } else if ((cuda_mempool != NULL) && md->config.cuda_async_managed) { -#else } else if ((cuda_mem_ctx == NULL) && md->config.cuda_async_managed) { -#endif - /* Stream-ordered CUDA allocations are typed as + /* Currently virtual/stream-ordered CUDA allocations are typed as * `UCS_MEMORY_TYPE_CUDA_MANAGED`. This may be changed using - * UCX_CUDA_COPY_ASYNC_MEM_TYPE env var. CUDA versions before 11.2 - * do not provide the memory-pool attribute, so use the missing - * context to identify these allocations. */ + * UCX_CUDA_COPY_ASYNC_MEM_TYPE env var. Ideally checking for + * `CU_POINTER_ATTRIBUTE_IS_LEGACY_CUDA_IPC_CAPABLE` would be better + * here, but due to a bug in the driver `cudaMalloc` also returns + * false in that case. Therefore, checking whether the allocation + * was not allocated in a context should also allows us to + * identify virtual/stream-ordered CUDA allocations. */ mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; *is_async_managed = 1; } else { diff --git a/test/gtest/common/mem_buffer.cc b/test/gtest/common/mem_buffer.cc index c17f550da2b..10c2b37eab1 100644 --- a/test/gtest/common/mem_buffer.cc +++ b/test/gtest/common/mem_buffer.cc @@ -473,7 +473,29 @@ void *mem_buffer::allocate(size_t size, ucs_memory_type_t mem_type, bool async) } return ptr; case UCS_MEMORY_TYPE_CUDA_MANAGED: - CUDA_CALL(cudaMallocManaged(&ptr, size), ": size=" << size); + if (async) { +#if CUDART_VERSION >= 13000 + cudaMemLocation location = {}; + cudaMemPool_t pool; + int device; + + CUDA_CALL(cudaGetDevice(&device), ""); + location.type = cudaMemLocationTypeDevice; + location.id = device; + CUDA_CALL(cudaMemGetDefaultMemPool(&pool, &location, + cudaMemAllocationTypeManaged), + ""); + CUDA_CALL(cudaMallocFromPoolAsync(&ptr, size, pool, 0), + ": size=" << size); + CUDA_CALL(cudaStreamSynchronize(0), ""); +#else + UCS_TEST_ABORT("asynchronous allocation for " + + std::string(ucs_memory_type_names[mem_type]) + + " memory type is not supported"); +#endif + } else { + CUDA_CALL(cudaMallocManaged(&ptr, size), ": size=" << size); + } return ptr; #endif #if HAVE_ROCM @@ -563,7 +585,18 @@ void mem_buffer::release(void *ptr, ucs_memory_type_t mem_type, bool async) } break; case UCS_MEMORY_TYPE_CUDA_MANAGED: - CUDA_CALL(cudaFree(ptr), ": ptr=" << ptr); + if (async) { +#if CUDART_VERSION >= 13000 + CUDA_CALL(cudaStreamSynchronize(0), ""); + CUDA_CALL(cudaFreeAsync(ptr, 0), ": ptr=" << ptr); +#else + UCS_TEST_ABORT("asynchronous release for " + + std::string(ucs_memory_type_names[mem_type]) + + " memory type is not supported"); +#endif + } else { + CUDA_CALL(cudaFree(ptr), ": ptr=" << ptr); + } break; #endif #if HAVE_ROCM @@ -851,10 +884,31 @@ bool mem_buffer::is_async_supported(ucs_memory_type_t mem_type) if (!is_mem_type_supported(UCS_MEMORY_TYPE_CUDA)) { return false; } - return mem_type == UCS_MEMORY_TYPE_CUDA; -#else - return false; + + if (mem_type == UCS_MEMORY_TYPE_CUDA) { + return true; + } + +#if CUDART_VERSION >= 13000 + if (mem_type == UCS_MEMORY_TYPE_CUDA_MANAGED) { + cudaMemLocation location = {}; + cudaMemPool_t pool; + int device; + + if (cudaGetDevice(&device) != cudaSuccess) { + return false; + } + + location.type = cudaMemLocationTypeDevice; + location.id = device; + return cudaMemGetDefaultMemPool(&pool, &location, + cudaMemAllocationTypeManaged) == + cudaSuccess; + } #endif +#endif + + return false; } mem_buffer::mem_buffer(size_t size, ucs_memory_type_t mem_type) : diff --git a/test/gtest/common/mem_buffer.h b/test/gtest/common/mem_buffer.h index 2a7d6289ed0..2a00f1cc51e 100644 --- a/test/gtest/common/mem_buffer.h +++ b/test/gtest/common/mem_buffer.h @@ -124,7 +124,7 @@ class mem_buffer { /* Allocation mode. */ enum class alloc_mode { DEFAULT, /* Default allocation mode, using cudaMalloc */ - ASYNC /* Asynchronous allocation mode, using cudaMallocAsync */ + ASYNC /* Asynchronous CUDA allocation mode */ }; mem_buffer(size_t size, ucs_memory_type_t mem_type); diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index af72c487eda..c653c951dff 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -19,57 +19,6 @@ extern "C" { #include -#if CUDART_VERSION >= 13000 -class cuda_async_managed_mem_buffer { -public: - cuda_async_managed_mem_buffer(cudaMemPool_t pool, size_t size) : - m_buffer(nullptr), m_stream(nullptr), - m_status(cudaStreamCreate(&m_stream)) - { - if (m_status == cudaSuccess) { - m_status = cudaMallocFromPoolAsync(&m_buffer, size, pool, m_stream); - } - } - - ~cuda_async_managed_mem_buffer() - { - if (m_buffer != nullptr) { - cudaFreeAsync(m_buffer, m_stream); - } - if (m_stream != nullptr) { - cudaStreamSynchronize(m_stream); - cudaStreamDestroy(m_stream); - } - } - - void *ptr() const - { - return m_buffer; - } - - cudaError_t status() const - { - return m_status; - } - - cudaError_t memset(int value, size_t size) - { - return cudaMemsetAsync(m_buffer, value, size, m_stream); - } - - cudaError_t synchronize() - { - return cudaStreamSynchronize(m_stream); - } - -private: - void *m_buffer; - cudaStream_t m_stream; - cudaError_t m_status; -}; -#endif - - class test_switch_cuda_device : public test_md { protected: template void detect_mem_type(ucs_memory_type_t mem_type) const; @@ -411,41 +360,19 @@ class test_mem_alloc_device : public test_switch_cuda_device { EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); } -#if CUDART_VERSION >= 13000 void test_async_managed_mem_pool_not_registrable() { constexpr size_t size = 192; - cudaMemLocation location = {}; - cudaMemPool_t pool = nullptr; - int cuda_device; - int concurrent_managed_access; uct_md_mem_attr_v2_t mem_attr = {}; - cudaError_t cuda_status; - - ASSERT_EQ(cudaSuccess, cudaGetDevice(&cuda_device)); - ASSERT_EQ(cudaSuccess, - cudaDeviceGetAttribute(&concurrent_managed_access, - cudaDevAttrConcurrentManagedAccess, - cuda_device)); - if (!concurrent_managed_access) { - UCS_TEST_SKIP_R("device does not support concurrent managed access"); - } - location.type = cudaMemLocationTypeDevice; - location.id = cuda_device; - cuda_status = cudaMemGetDefaultMemPool(&pool, &location, - cudaMemAllocationTypeManaged); - if (cuda_status != cudaSuccess) { - UCS_TEST_SKIP_R("failed to get default managed memory pool"); + if (!mem_buffer::is_async_supported(UCS_MEMORY_TYPE_CUDA_MANAGED)) { + UCS_TEST_SKIP_R("asynchronous CUDA managed memory is not " + "supported"); } - cuda_async_managed_mem_buffer buffer(pool, size); - if (buffer.status() != cudaSuccess) { - UCS_TEST_SKIP_R("failed to allocate async managed mempool memory"); - } - - ASSERT_EQ(cudaSuccess, buffer.memset(0, size)); - ASSERT_EQ(cudaSuccess, buffer.synchronize()); + mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA_MANAGED, + mem_buffer::alloc_mode::ASYNC); + buffer.memset(0); mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE | UCT_MD_MEM_ATTR_V2_FIELD_SYS_DEV | @@ -455,7 +382,6 @@ class test_mem_alloc_device : public test_switch_cuda_device { EXPECT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, mem_attr.sys_dev); EXPECT_FALSE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); } -#endif private: std::vector m_sys_dev; @@ -575,7 +501,6 @@ UCS_TEST_P(test_mem_alloc_device, uct_alloc_managed_mem_registrable, EXPECT_UCS_OK(uct_mem_free(&mem)); } -#if CUDART_VERSION >= 13000 UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_not_registrable, "CUDA_COPY_DMABUF=try") { @@ -594,7 +519,6 @@ UCS_TEST_P(test_mem_alloc_device, { test_async_managed_mem_pool_not_registrable(); } -#endif UCS_TEST_P(test_mem_alloc_device, no_current_context_vmm_mem_registrable, "CUDA_COPY_ASYNC_MEM_TYPE=cuda") From ad87c78b1c1bdade9ab7b3f2f054d268c60fbd92 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 9 Sep 2026 10:22:58 +0000 Subject: [PATCH 10/19] GTEST/CUDA: Reuse managed default pool lookup --- test/gtest/common/mem_buffer.cc | 43 ++++++++++++++++++--------------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/test/gtest/common/mem_buffer.cc b/test/gtest/common/mem_buffer.cc index 10c2b37eab1..d4fd06910cc 100644 --- a/test/gtest/common/mem_buffer.cc +++ b/test/gtest/common/mem_buffer.cc @@ -52,6 +52,26 @@ #endif +#if HAVE_CUDA && CUDART_VERSION >= 13000 +static cudaError_t +mem_buffer_cuda_get_default_managed_pool(cudaMemPool_t *pool_p) +{ + cudaMemLocation location = {}; + int device; + cudaError_t cuda_status; + + cuda_status = cudaGetDevice(&device); + if (cuda_status != cudaSuccess) { + return cuda_status; + } + + location.type = cudaMemLocationTypeDevice; + location.id = device; + return cudaMemGetDefaultMemPool(pool_p, &location, + cudaMemAllocationTypeManaged); +} +#endif + #if HAVE_ROCM #include #include @@ -475,16 +495,9 @@ void *mem_buffer::allocate(size_t size, ucs_memory_type_t mem_type, bool async) case UCS_MEMORY_TYPE_CUDA_MANAGED: if (async) { #if CUDART_VERSION >= 13000 - cudaMemLocation location = {}; cudaMemPool_t pool; - int device; - - CUDA_CALL(cudaGetDevice(&device), ""); - location.type = cudaMemLocationTypeDevice; - location.id = device; - CUDA_CALL(cudaMemGetDefaultMemPool(&pool, &location, - cudaMemAllocationTypeManaged), - ""); + + CUDA_CALL(mem_buffer_cuda_get_default_managed_pool(&pool), ""); CUDA_CALL(cudaMallocFromPoolAsync(&ptr, size, pool, 0), ": size=" << size); CUDA_CALL(cudaStreamSynchronize(0), ""); @@ -891,19 +904,9 @@ bool mem_buffer::is_async_supported(ucs_memory_type_t mem_type) #if CUDART_VERSION >= 13000 if (mem_type == UCS_MEMORY_TYPE_CUDA_MANAGED) { - cudaMemLocation location = {}; cudaMemPool_t pool; - int device; - - if (cudaGetDevice(&device) != cudaSuccess) { - return false; - } - location.type = cudaMemLocationTypeDevice; - location.id = device; - return cudaMemGetDefaultMemPool(&pool, &location, - cudaMemAllocationTypeManaged) == - cudaSuccess; + return mem_buffer_cuda_get_default_managed_pool(&pool) == cudaSuccess; } #endif #endif From 13e872f4a19399a42fbab191d670dd399c404054 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 9 Sep 2026 11:34:32 +0000 Subject: [PATCH 11/19] UCT/CUDA_COPY: Track queried attribute count --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index d9607cf9818..8231e143c72 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -685,7 +685,7 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, uint32_t is_managed = 0; CUcontext cuda_mem_ctx = NULL; void *cuda_mempool = NULL; - unsigned num_attrs = 4; + unsigned num_attrs = 0; CUpointer_attribute attr_type[UCT_CUDA_MEM_QUERY_MAX_ATTRS]; void *attr_data[UCT_CUDA_MEM_QUERY_MAX_ATTRS]; CUdevice cuda_device; @@ -704,14 +704,14 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, return UCS_ERR_INVALID_ADDR; } } else { - attr_type[0] = CU_POINTER_ATTRIBUTE_MEMORY_TYPE; - attr_data[0] = &cuda_mem_type; - attr_type[1] = CU_POINTER_ATTRIBUTE_IS_MANAGED; - attr_data[1] = &is_managed; - attr_type[2] = CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL; - attr_data[2] = &cuda_device; - attr_type[3] = CU_POINTER_ATTRIBUTE_CONTEXT; - attr_data[3] = &cuda_mem_ctx; + attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_MEMORY_TYPE; + attr_data[num_attrs++] = &cuda_mem_type; + attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_IS_MANAGED; + attr_data[num_attrs++] = &is_managed; + attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL; + attr_data[num_attrs++] = &cuda_device; + attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_CONTEXT; + attr_data[num_attrs++] = &cuda_mem_ctx; #if HAVE_DECL_CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE; attr_data[num_attrs] = &cuda_mempool; From 2b732832cf2b984fa6774f239da85e59ca2e521b Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 9 Sep 2026 12:04:16 +0000 Subject: [PATCH 12/19] UCT/CUDA_COPY: Clarify managed pool detection --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 14 +++++++------- test/gtest/uct/cuda/test_switch_cuda_device.cc | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index 8231e143c72..4bbba6b153c 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -713,9 +713,8 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_CONTEXT; attr_data[num_attrs++] = &cuda_mem_ctx; #if HAVE_DECL_CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE - attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE; - attr_data[num_attrs] = &cuda_mempool; - num_attrs++; + attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE; + attr_data[num_attrs++] = &cuda_mempool; #endif status = UCT_CUDADRV_FUNC_LOG_ERR( @@ -738,8 +737,7 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; if (cuda_mempool != NULL) { /* Managed-pool allocations are stream-ordered and cannot be - * registered. Use the pool device to avoid re-detection when - * CPU is the preferred location. */ + * registered. Ignore PREF_LOC and use the pool device. */ *is_async_managed = 1; pref_loc = cuda_device; } else { @@ -774,8 +772,10 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, * `CU_POINTER_ATTRIBUTE_IS_LEGACY_CUDA_IPC_CAPABLE` would be better * here, but due to a bug in the driver `cudaMalloc` also returns * false in that case. Therefore, checking whether the allocation - * was not allocated in a context should also allows us to - * identify virtual/stream-ordered CUDA allocations. */ + * was not allocated in a context should also allow us to + * identify virtual/stream-ordered CUDA allocations. Keep this + * heuristic for non-managed allocations; managed-pool allocations + * are handled above. */ mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; *is_async_managed = 1; } else { diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index c653c951dff..f0ab3c641d8 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -362,7 +362,7 @@ class test_mem_alloc_device : public test_switch_cuda_device { void test_async_managed_mem_pool_not_registrable() { - constexpr size_t size = 192; + constexpr size_t size = 192; uct_md_mem_attr_v2_t mem_attr = {}; if (!mem_buffer::is_async_supported(UCS_MEMORY_TYPE_CUDA_MANAGED)) { From 3773519b94561fc42d0b35519ed588b522d608db Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 9 Sep 2026 12:35:37 +0000 Subject: [PATCH 13/19] GTEST/CUDA: Align managed pool test variables --- test/gtest/uct/cuda/test_switch_cuda_device.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index f0ab3c641d8..3a1ee80b04a 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -362,7 +362,7 @@ class test_mem_alloc_device : public test_switch_cuda_device { void test_async_managed_mem_pool_not_registrable() { - constexpr size_t size = 192; + constexpr size_t size = 192; uct_md_mem_attr_v2_t mem_attr = {}; if (!mem_buffer::is_async_supported(UCS_MEMORY_TYPE_CUDA_MANAGED)) { From 15e77817c63da5c3a80370ab952041b3b3ada721 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 9 Sep 2026 12:48:02 +0000 Subject: [PATCH 14/19] GTEST/CUDA: Clear managed pool probe errors --- test/gtest/common/mem_buffer.cc | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/test/gtest/common/mem_buffer.cc b/test/gtest/common/mem_buffer.cc index d4fd06910cc..bdfb30d38ec 100644 --- a/test/gtest/common/mem_buffer.cc +++ b/test/gtest/common/mem_buffer.cc @@ -67,8 +67,13 @@ mem_buffer_cuda_get_default_managed_pool(cudaMemPool_t *pool_p) location.type = cudaMemLocationTypeDevice; location.id = device; - return cudaMemGetDefaultMemPool(pool_p, &location, - cudaMemAllocationTypeManaged); + cuda_status = cudaMemGetDefaultMemPool(pool_p, &location, + cudaMemAllocationTypeManaged); + if (cuda_status != cudaSuccess) { + (void)cudaGetLastError(); /* do not leak the error to other tests */ + } + + return cuda_status; } #endif From 16a911b6a767d66f65cded12c9b6e60ec454ebe7 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 9 Sep 2026 13:07:39 +0000 Subject: [PATCH 15/19] GTEST/CUDA: Share async release path --- test/gtest/common/mem_buffer.cc | 15 +-------------- 1 file changed, 1 insertion(+), 14 deletions(-) diff --git a/test/gtest/common/mem_buffer.cc b/test/gtest/common/mem_buffer.cc index bdfb30d38ec..3ce927f5cab 100644 --- a/test/gtest/common/mem_buffer.cc +++ b/test/gtest/common/mem_buffer.cc @@ -589,6 +589,7 @@ void mem_buffer::release(void *ptr, ucs_memory_type_t mem_type, bool async) break; #if HAVE_CUDA case UCS_MEMORY_TYPE_CUDA: + case UCS_MEMORY_TYPE_CUDA_MANAGED: if (async) { #if CUDART_VERSION >= 11020 cudaStreamSynchronize(0); @@ -597,20 +598,6 @@ void mem_buffer::release(void *ptr, ucs_memory_type_t mem_type, bool async) UCS_TEST_ABORT("asynchronous release for " + std::string(ucs_memory_type_names[mem_type]) + " memory type is not supported"); -#endif - } else { - CUDA_CALL(cudaFree(ptr), ": ptr=" << ptr); - } - break; - case UCS_MEMORY_TYPE_CUDA_MANAGED: - if (async) { -#if CUDART_VERSION >= 13000 - CUDA_CALL(cudaStreamSynchronize(0), ""); - CUDA_CALL(cudaFreeAsync(ptr, 0), ": ptr=" << ptr); -#else - UCS_TEST_ABORT("asynchronous release for " + - std::string(ucs_memory_type_names[mem_type]) + - " memory type is not supported"); #endif } else { CUDA_CALL(cudaFree(ptr), ": ptr=" << ptr); From f473f14bf2d9bb61bcdf6d5b5dbfd55744424997 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 9 Sep 2026 16:40:19 +0000 Subject: [PATCH 16/19] UCT/CUDA_COPY: Register managed memory without dmabuf --- config/m4/cuda.m4 | 2 - src/uct/cuda/cuda_copy/cuda_copy_md.c | 64 ++++++++----------- .../gtest/uct/cuda/test_switch_cuda_device.cc | 16 ++--- 3 files changed, 34 insertions(+), 48 deletions(-) diff --git a/config/m4/cuda.m4 b/config/m4/cuda.m4 index a510414b902..553bd212d89 100644 --- a/config/m4/cuda.m4 +++ b/config/m4/cuda.m4 @@ -210,8 +210,6 @@ AS_IF([test "x$cuda_checked" != "xyes"], [AC_DEFINE([HAVE_CUMEMRETAINALLOCATIONHANDLE], [1], [Enable cuMemRetainAllocationHandle() usage])]) AC_CHECK_DECLS([CU_MEM_LOCATION_TYPE_HOST], - [], [], [[#include ]]) - AC_CHECK_DECLS([CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE], [], [], [[#include ]])]) # Check nvml header files diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index 4bbba6b153c..16093093e98 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -680,14 +680,12 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, ucs_memory_info_t *mem_info, int *is_async_managed, int *is_host_located) { -#define UCT_CUDA_MEM_QUERY_MAX_ATTRS 5 +#define UCT_CUDA_MEM_QUERY_NUM_ATTRS 4 CUmemorytype cuda_mem_type = CU_MEMORYTYPE_HOST; uint32_t is_managed = 0; CUcontext cuda_mem_ctx = NULL; - void *cuda_mempool = NULL; - unsigned num_attrs = 0; - CUpointer_attribute attr_type[UCT_CUDA_MEM_QUERY_MAX_ATTRS]; - void *attr_data[UCT_CUDA_MEM_QUERY_MAX_ATTRS]; + CUpointer_attribute attr_type[UCT_CUDA_MEM_QUERY_NUM_ATTRS]; + void *attr_data[UCT_CUDA_MEM_QUERY_NUM_ATTRS]; CUdevice cuda_device; int32_t pref_loc; int is_vmm; @@ -704,21 +702,18 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, return UCS_ERR_INVALID_ADDR; } } else { - attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_MEMORY_TYPE; - attr_data[num_attrs++] = &cuda_mem_type; - attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_IS_MANAGED; - attr_data[num_attrs++] = &is_managed; - attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL; - attr_data[num_attrs++] = &cuda_device; - attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_CONTEXT; - attr_data[num_attrs++] = &cuda_mem_ctx; -#if HAVE_DECL_CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE - attr_type[num_attrs] = CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE; - attr_data[num_attrs++] = &cuda_mempool; -#endif + attr_type[0] = CU_POINTER_ATTRIBUTE_MEMORY_TYPE; + attr_data[0] = &cuda_mem_type; + attr_type[1] = CU_POINTER_ATTRIBUTE_IS_MANAGED; + attr_data[1] = &is_managed; + attr_type[2] = CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL; + attr_data[2] = &cuda_device; + attr_type[3] = CU_POINTER_ATTRIBUTE_CONTEXT; + attr_data[3] = &cuda_mem_ctx; status = UCT_CUDADRV_FUNC_LOG_ERR( - cuPointerGetAttributes(num_attrs, attr_type, attr_data, + cuPointerGetAttributes(ucs_static_array_size(attr_data), + attr_type, attr_data, (CUdeviceptr)address)); if (status != UCS_OK) { /* pointer not recognized */ @@ -735,22 +730,14 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, * provided address and length as base address and alloc length * respectively */ mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; - if (cuda_mempool != NULL) { - /* Managed-pool allocations are stream-ordered and cannot be - * registered. Ignore PREF_LOC and use the pool device. */ - *is_async_managed = 1; - pref_loc = cuda_device; - } else { - cu_err = cuMemRangeGetAttribute( - (void*)&pref_loc, sizeof(pref_loc), - CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION, - (CUdeviceptr)address, length); - if ((cu_err != CUDA_SUCCESS) || - (pref_loc == CU_DEVICE_INVALID)) { - pref_loc = (md->config.pref_loc == UCT_CUDA_PREF_LOC_CPU) ? - CU_DEVICE_CPU : - cuda_device; - } + cu_err = cuMemRangeGetAttribute( + (void*)&pref_loc, sizeof(pref_loc), + CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION, + (CUdeviceptr)address, length); + if ((cu_err != CUDA_SUCCESS) || (pref_loc == CU_DEVICE_INVALID)) { + pref_loc = (md->config.pref_loc == UCT_CUDA_PREF_LOC_CPU) ? + CU_DEVICE_CPU : + cuda_device; } if (pref_loc == CU_DEVICE_CPU) { @@ -773,9 +760,7 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, * here, but due to a bug in the driver `cudaMalloc` also returns * false in that case. Therefore, checking whether the allocation * was not allocated in a context should also allow us to - * identify virtual/stream-ordered CUDA allocations. Keep this - * heuristic for non-managed allocations; managed-pool allocations - * are handled above. */ + * identify virtual/stream-ordered CUDA allocations. */ mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; *is_async_managed = 1; } else { @@ -897,6 +882,11 @@ uct_cuda_copy_md_detect_mem_flags(uct_cuda_copy_md_t *md, return 0; } + /* Managed memory is registered through ODP and is not dmabuf-exportable. */ + if (mem_info->type == UCS_MEMORY_TYPE_CUDA_MANAGED) { + return UCS_MEM_FLAG_REGISTRABLE; + } + /* Host-located CUDA VMM is registerable even if dmabuf export fails. */ if (is_host_located) { return UCS_MEM_FLAG_REGISTRABLE; diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index 3a1ee80b04a..92cdab24991 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -360,7 +360,7 @@ class test_mem_alloc_device : public test_switch_cuda_device { EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); } - void test_async_managed_mem_pool_not_registrable() + void test_async_managed_mem_pool_registrable() { constexpr size_t size = 192; uct_md_mem_attr_v2_t mem_attr = {}; @@ -375,12 +375,10 @@ class test_mem_alloc_device : public test_switch_cuda_device { buffer.memset(0); 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_MEM_FLAGS; EXPECT_UCS_OK(uct_md_mem_query_v2(md(), buffer.ptr(), size, &mem_attr)); EXPECT_EQ(UCS_MEMORY_TYPE_CUDA_MANAGED, mem_attr.mem_type); - EXPECT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, mem_attr.sys_dev); - EXPECT_FALSE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); + EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); } private: @@ -501,23 +499,23 @@ UCS_TEST_P(test_mem_alloc_device, uct_alloc_managed_mem_registrable, EXPECT_UCS_OK(uct_mem_free(&mem)); } -UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_not_registrable, +UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_registrable, "CUDA_COPY_DMABUF=try") { - test_async_managed_mem_pool_not_registrable(); + test_async_managed_mem_pool_registrable(); } UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_gpu_pref_loc, "CUDA_COPY_DMABUF=try", "CUDA_COPY_PREF_LOC=gpu") { - test_async_managed_mem_pool_not_registrable(); + test_async_managed_mem_pool_registrable(); } UCS_TEST_P(test_mem_alloc_device, - async_managed_mem_pool_not_registrable_cuda_type, + async_managed_mem_pool_registrable_cuda_type, "CUDA_COPY_DMABUF=try", "CUDA_COPY_ASYNC_MEM_TYPE=cuda") { - test_async_managed_mem_pool_not_registrable(); + test_async_managed_mem_pool_registrable(); } UCS_TEST_P(test_mem_alloc_device, no_current_context_vmm_mem_registrable, From a7cbb96f0bf2bedd9fb2dd4157f3b43c5e078aad Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 10 Sep 2026 13:22:59 +0000 Subject: [PATCH 17/19] UCT/CUDA_COPY: Preserve device symbol flags --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 12 ++++++++---- test/apps/test_cuda_get_symbol_address.cu | 13 +++++++++++++ 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index a8e1e23645f..f98cd3ba8bd 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -730,6 +730,7 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, * provided address and length as base address and alloc length * respectively */ mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; + cu_err = cuMemRangeGetAttribute( (void*)&pref_loc, sizeof(pref_loc), CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION, @@ -759,7 +760,7 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, * `CU_POINTER_ATTRIBUTE_IS_LEGACY_CUDA_IPC_CAPABLE` would be better * here, but due to a bug in the driver `cudaMalloc` also returns * false in that case. Therefore, checking whether the allocation - * was not allocated in a context should also allow us to + * was not allocated in a context should also allows us to * identify virtual/stream-ordered CUDA allocations. */ mem_info->type = UCS_MEMORY_TYPE_CUDA_MANAGED; *is_async_managed = 1; @@ -988,6 +989,7 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address, CUdevice avail_cuda_device = CU_DEVICE_INVALID; ucs_memory_info_t cached_mem_info; ucs_memory_info_t addr_mem_info; + ucs_memory_info_t detected_mem_info = {}; ucs_status_t cache_status; ucs_status_t status; @@ -1013,8 +1015,10 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address, return status; } - /* CUDA reports device symbols as device memory, so preserve the type - * explicitly provided by the UCM allocation event. */ + /* Preserve the driver-reported type for flags. CUDA reports device + * symbols as device memory, so use the UCM event type for callers. */ + detected_mem_info = addr_mem_info; + if ((cache_status == UCS_OK) && (cached_mem_info.type == UCS_MEMORY_TYPE_CUDA_MANAGED)) { addr_mem_info.type = cached_mem_info.type; @@ -1068,7 +1072,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, is_host_located, + md, &detected_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, diff --git a/test/apps/test_cuda_get_symbol_address.cu b/test/apps/test_cuda_get_symbol_address.cu index d9430836466..a47d60bfe44 100644 --- a/test/apps/test_cuda_get_symbol_address.cu +++ b/test/apps/test_cuda_get_symbol_address.cu @@ -10,6 +10,7 @@ #include #include +#include #include @@ -68,6 +69,7 @@ static int check_symbol_memory_type(ucp_context_h context, void *address) { ucp_mem_map_params_t map_params; ucp_mem_attr_t mem_attr; + ucs_memory_info_t cached_mem_info; ucp_mem_h memh; ucs_status_t status; int num_errors = 0; @@ -94,6 +96,17 @@ static int check_symbol_memory_type(ucp_context_h context, void *address) ++num_errors; } + status = ucs_memtype_cache_lookup(address, sizeof(device_int), + &cached_mem_info); + if (status != UCS_OK) { + printf("memtype cache lookup failed: %s\n", ucs_status_string(status)); + ++num_errors; + } else if (cached_mem_info.mem_flags & UCS_MEM_FLAG_REGISTRABLE) { + printf("unexpected cached symbol memory flags 0x%x\n", + cached_mem_info.mem_flags); + ++num_errors; + } + status = ucp_mem_unmap(context, memh); if (status != UCS_OK) { printf("ucp_mem_unmap() failed: %s\n", ucs_status_string(status)); From 437da172c8f28bca7b5607957d7430f1052a241e Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 10 Sep 2026 15:06:18 +0000 Subject: [PATCH 18/19] UCT/CUDA_COPY: Simplify managed memory checks --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 4 ++-- test/gtest/uct/cuda/test_switch_cuda_device.cc | 17 +++++++++-------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index f98cd3ba8bd..9ea5bdc1445 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -929,7 +929,7 @@ uct_cuda_copy_md_is_registrable(uct_cuda_copy_md_t *md, /* Managed memory is registered through ODP and is not dmabuf-exportable. */ if (mem_info->type == UCS_MEMORY_TYPE_CUDA_MANAGED) { - return UCS_MEM_FLAG_REGISTRABLE; + return 1; } /* Host-located CUDA VMM is registerable even if dmabuf export fails. */ @@ -987,9 +987,9 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address, int is_host_located = 0; CUdevice cur_cuda_device = CU_DEVICE_INVALID; CUdevice avail_cuda_device = CU_DEVICE_INVALID; + ucs_memory_info_t detected_mem_info = {}; ucs_memory_info_t cached_mem_info; ucs_memory_info_t addr_mem_info; - ucs_memory_info_t detected_mem_info = {}; ucs_status_t cache_status; ucs_status_t status; diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index dc4d40c84b8..4498aa25c00 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -342,16 +342,20 @@ class test_mem_alloc_device : public test_switch_cuda_device { EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); } - void query_managed_registrable(void *address, size_t size) + void check_managed_memory_context(void *address) { CUpointer_attribute attr_type = CU_POINTER_ATTRIBUTE_CONTEXT; - uct_md_mem_attr_v2_t mem_attr = {}; CUcontext cuda_mem_ctx = nullptr; ASSERT_EQ(CUDA_SUCCESS, cuPointerGetAttribute(&cuda_mem_ctx, attr_type, (CUdeviceptr)address)); EXPECT_NE(nullptr, cuda_mem_ctx); + } + + void query_managed_registrable(void *address, size_t size) + { + uct_md_mem_attr_v2_t mem_attr = {}; mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE | UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS; @@ -363,7 +367,6 @@ class test_mem_alloc_device : public test_switch_cuda_device { void test_async_managed_mem_pool_registrable() { constexpr size_t size = 192; - uct_md_mem_attr_v2_t mem_attr = {}; if (!mem_buffer::is_async_supported(UCS_MEMORY_TYPE_CUDA_MANAGED)) { UCS_TEST_SKIP_R("asynchronous CUDA managed memory is not " @@ -374,11 +377,7 @@ class test_mem_alloc_device : public test_switch_cuda_device { mem_buffer::alloc_mode::ASYNC); buffer.memset(0); - mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE | - UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS; - EXPECT_UCS_OK(uct_md_mem_query_v2(md(), buffer.ptr(), size, &mem_attr)); - EXPECT_EQ(UCS_MEMORY_TYPE_CUDA_MANAGED, mem_attr.mem_type); - EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE); + query_managed_registrable(buffer.ptr(), size); } private: @@ -486,6 +485,7 @@ UCS_TEST_P(test_mem_alloc_device, user_managed_mem_registrable, constexpr size_t size = 4 * UCS_MBYTE; mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA_MANAGED); + check_managed_memory_context(buffer.ptr()); query_managed_registrable(buffer.ptr(), size); } @@ -494,6 +494,7 @@ UCS_TEST_P(test_mem_alloc_device, uct_alloc_managed_mem_registrable, { ASSERT_UCS_OK(allocate(UCS_MEMORY_TYPE_CUDA_MANAGED)); + check_managed_memory_context(mem.address); query_managed_registrable(mem.address, mem.length); EXPECT_UCS_OK(uct_mem_free(&mem)); From fae7ff121d075e7018fae07b9212630dd1a561f8 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 10 Sep 2026 18:37:19 +0000 Subject: [PATCH 19/19] UCT/CUDA_COPY: Initialize memory query result --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 2 +- test/apps/test_cuda_get_symbol_address.cu | 13 ------------- test/gtest/uct/cuda/test_switch_cuda_device.cc | 6 +++--- 3 files changed, 4 insertions(+), 17 deletions(-) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index 9ea5bdc1445..546a8ed3a2e 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -988,8 +988,8 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address, CUdevice cur_cuda_device = CU_DEVICE_INVALID; CUdevice avail_cuda_device = CU_DEVICE_INVALID; ucs_memory_info_t detected_mem_info = {}; + ucs_memory_info_t addr_mem_info = {}; ucs_memory_info_t cached_mem_info; - ucs_memory_info_t addr_mem_info; ucs_status_t cache_status; ucs_status_t status; diff --git a/test/apps/test_cuda_get_symbol_address.cu b/test/apps/test_cuda_get_symbol_address.cu index a47d60bfe44..d9430836466 100644 --- a/test/apps/test_cuda_get_symbol_address.cu +++ b/test/apps/test_cuda_get_symbol_address.cu @@ -10,7 +10,6 @@ #include #include -#include #include @@ -69,7 +68,6 @@ static int check_symbol_memory_type(ucp_context_h context, void *address) { ucp_mem_map_params_t map_params; ucp_mem_attr_t mem_attr; - ucs_memory_info_t cached_mem_info; ucp_mem_h memh; ucs_status_t status; int num_errors = 0; @@ -96,17 +94,6 @@ static int check_symbol_memory_type(ucp_context_h context, void *address) ++num_errors; } - status = ucs_memtype_cache_lookup(address, sizeof(device_int), - &cached_mem_info); - if (status != UCS_OK) { - printf("memtype cache lookup failed: %s\n", ucs_status_string(status)); - ++num_errors; - } else if (cached_mem_info.mem_flags & UCS_MEM_FLAG_REGISTRABLE) { - printf("unexpected cached symbol memory flags 0x%x\n", - cached_mem_info.mem_flags); - ++num_errors; - } - status = ucp_mem_unmap(context, memh); if (status != UCS_OK) { printf("ucp_mem_unmap() failed: %s\n", ucs_status_string(status)); diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index 4498aa25c00..d1e4690474e 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -366,7 +366,7 @@ class test_mem_alloc_device : public test_switch_cuda_device { void test_async_managed_mem_pool_registrable() { - constexpr size_t size = 192; + constexpr size_t size = 192; if (!mem_buffer::is_async_supported(UCS_MEMORY_TYPE_CUDA_MANAGED)) { UCS_TEST_SKIP_R("asynchronous CUDA managed memory is not " @@ -480,7 +480,7 @@ UCS_TEST_P(test_mem_alloc_device, no_current_context_user_mem_registrable, } UCS_TEST_P(test_mem_alloc_device, user_managed_mem_registrable, - "CUDA_COPY_PREF_LOC=cpu") + "CUDA_COPY_PREF_LOC=gpu") { constexpr size_t size = 4 * UCS_MBYTE; mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA_MANAGED); @@ -490,7 +490,7 @@ UCS_TEST_P(test_mem_alloc_device, user_managed_mem_registrable, } UCS_TEST_P(test_mem_alloc_device, uct_alloc_managed_mem_registrable, - "CUDA_COPY_PREF_LOC=cpu") + "CUDA_COPY_PREF_LOC=gpu") { ASSERT_UCS_OK(allocate(UCS_MEMORY_TYPE_CUDA_MANAGED));