UCP/UCT/UCS: Add memory registration flags by tomerg-nvidia · Pull Request #11553 · openucx/ucx

tomerg-nvidia · 2026-06-11T15:50:11Z

What?

Add memory registration flags to UCS memory information and propagate them through UCT mem_query_v2 and UCP memory/protocol selection paths.

This adds a REGISTERABLE memory flag, lets CUDA async allocations classified as CUDA managed report that they are not registerable, and makes UCP avoid registration-dependent MDs/protocols when that flag is missing.

Why?

Some CUDA memory can be detected as a memory type that normally has registration-capable transports, while the specific allocation cannot actually be registered on IB MDs. Without carrying this per-allocation information into UCP, protocol selection can choose zcopy paths that require registration and later fail.

How?

Extend UCS memory info and the memtype cache to carry memory flags.
Extend UCT mem_query_v2 helpers/providers to return memory flags, defaulting to REGISTERABLE.
Mark CUDA async memory classified as CUDA_MANAGED as not registerable.
Add required_mem_flags to UCT MD attributes and require REGISTERABLE for IB MD registration paths.
Propagate memory flags through UCP memory detection, rcache/memh compatibility, datatype iterators, RNDV, tag offload, and proto selection.
Add focused gtests for memtype-cache flags, IB registration failure for CUDA async memory, and UCP protocol filtering.

Add memory flags to UCS memory information and propagate them through UCT mem_query_v2 so providers can report whether a memory region is safe to register. Use the CAN_REGISTER flag in UCP memory registration and protocol selection paths to avoid selecting registration-dependent protocols for non-registerable memory. Mark CUDA async allocations classified as CUDA managed as non-registerable, and require CAN_REGISTER for IB memory domains that need registration. Add focused gtests for memtype cache flag handling, IB registration failure for CUDA async memory, and UCP protocol filtering.

guy-ealey-morag · 2026-06-15T13:46:34Z

Also update NVIDIA copyright years on changes files

tvegas1 · 2026-06-15T18:20:49Z

+    ucp_md_index_t i, md_index;
+    uint8_t sg_count;
+
+    if (!mem_buffer::is_mem_type_supported(UCS_MEMORY_TYPE_CUDA)) {


remove if we check for async after?

is_async_supported doesn't check we have a gpu, just the compile-time cuda library version

i see, can we merge the two checks in is_async_supported?

since all uses of is_async_supported (even in master) first check for cuda mem type supported, we can call is_mem_type_supported(UCS_MEMORY_TYPE_CUDA) inside is_async_supported(), and only keep this one everywhere.

Add the mem_flags to the rcache key. Add more tests. Apply CR fixes.

tvegas1 · 2026-06-16T14:50:43Z

-    ucp_memory_detect(context, address, length, &info);
-    ucp_memh_init(memh, context, memh_flags, uct_flags, method, mem_type,
-                  info.sys_dev);
+    ucp_memh_init(memh, context, memh_flags, uct_flags, method, mem_info->type,


can we simply pass mem_info here, instead of type, sys_dev and mem_flags?

brminich

maybe we do not keep non-registrable memory in rcache?

brminich · 2026-06-16T16:31:34Z

+        .parent          = &ucp_mem_dummy_handle.memh,
+        .mem_type        = UCS_MEMORY_TYPE_HOST,
+        .sys_dev         = UCS_SYS_DEVICE_ID_UNKNOWN,
+        .super.mem_flags = UCS_MEM_FLAG_REGISTRABLE,


is there a reason to keep flags separately in rcache?

brminich · 2026-06-17T06:40:00Z

+        reg_mem_info.type      = internal_mem_info.type;
+        reg_mem_info.sys_dev   = internal_mem_info.sys_dev;
+        reg_mem_info.mem_flags = internal_mem_info.mem_flags;
+        ucp_mem_info_update_for_type(&reg_mem_info, mem_type);


why do we need it here? Is there a use case that mem_type can be different from internal_mem_info.type?

brminich · 2026-06-17T06:45:12Z

+    }
+
+    ucp_memory_detect(context, address, length, mem_info);
+    ucp_mem_info_update_for_type(mem_info, mem_type);


why do we need it ? what is the use case that we havee to change the info filled by detect function?

brminich · 2026-06-17T06:46:21Z

+        ucp_mem_info_detect_for_type(context, buffer, length, mem_type,
+                                     &mem_info);


why do we need to redetect here?

brminich · 2026-06-17T07:13:35Z


    status = ucs_rcache_get(rcache, address, length, alignment,
-                            PROT_READ | PROT_WRITE, &reg_ctx, &rregion);
+                            PROT_READ | PROT_WRITE, mem_info->mem_flags,


do we really need to pass flags? Can we assume that we always need registrable?

actually that's interesting, so we would never cache memory without REGISTRABLE flag?

brminich · 2026-06-17T07:14:20Z

+        mem_info = ucp_memory_info_from_memh(memh);
+        status = ucp_memh_get(context, address, length, &mem_info, cache_md_map,


minor: align

tvegas1 · 2026-06-17T07:35:17Z


    status = ucs_rcache_get(rcache, address, length, alignment,
-                            PROT_READ | PROT_WRITE, &reg_ctx, &rregion);
+                            PROT_READ | PROT_WRITE, mem_info->mem_flags,


actually that's interesting, so we would never cache memory without REGISTRABLE flag?

Also fix CR comments

…m v1 Also assert on different mem_type from detected mem_type.

brminich · 2026-06-22T12:09:40Z

+    *md_idx_p             = context->alloc_md[alloc_mem_type].md_index;
+    mem_info_p->type      = alloc_mem_type;
+    mem_info_p->flags     = context->alloc_md[alloc_mem_type].mem_flags;


brminich · 2026-06-22T14:42:57Z

+    CUdevice cuda_device            = CU_DEVICE_INVALID;
+    CUpointer_attribute attr_type[] = {CU_POINTER_ATTRIBUTE_MEMORY_TYPE,
+                                       CU_POINTER_ATTRIBUTE_IS_MANAGED,
+                                       CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL,


why do we need cuda_device?

also can you avoid doing this query and cache the result from the previous query in the mem_query function?

you're right, removed this function and reused the result from the earlier query

brminich · 2026-06-22T14:50:35Z

+        return UCS_MEM_FLAG_REGISTRABLE;
+    }
+
+    if (!md->config.dmabuf_supported) {


why? what if it is cuda_managed memory?

cuda_managed is handled above

Don't query twice the pointer for asyncness Refactoring tests

* UCP/UCT/UCS: Add memory registration flags Add memory flags to UCS memory information and propagate them through UCT mem_query_v2 so providers can report whether a memory region is safe to register. Use the CAN_REGISTER flag in UCP memory registration and protocol selection paths to avoid selecting registration-dependent protocols for non-registerable memory. Mark CUDA async allocations classified as CUDA managed as non-registerable, and require CAN_REGISTER for IB memory domains that need registration. Add focused gtests for memtype cache flag handling, IB registration failure for CUDA async memory, and UCP protocol filtering.

tomerg-nvidia added 3 commits June 11, 2026 17:45

UCP/CORE: Fix release build issues

340465c

UCP/UCT/UCS: Fix format issues

425cab7

tomerg-nvidia marked this pull request as ready for review June 12, 2026 08:02

tomerg-nvidia requested review from brminich, tvegas1 and yosefe June 12, 2026 08:02

UCT/CUDA/CUDA_COPY: add dmabuf export based flag detection

dbb7009

guy-ealey-morag reviewed Jun 15, 2026

View reviewed changes

UCP/UCS/UCT: CR fixes

33a94f0

tvegas1 reviewed Jun 15, 2026

View reviewed changes

tomerg-nvidia added 2 commits June 16, 2026 09:32

Merge branch 'master' into ucs_mem_flags_pr

fe4680c

UCP/UCT/UCS: CR fixes

7553779

Add the mem_flags to the rcache key. Add more tests. Apply CR fixes.

tvegas1 reviewed Jun 16, 2026

View reviewed changes

brminich reviewed Jun 17, 2026

View reviewed changes

tvegas1 reviewed Jun 17, 2026

View reviewed changes

UCP/UCT/UCS: Remove the flags from the rcache key

f2f99ab

Also fix CR comments

brminich reviewed Jun 17, 2026

View reviewed changes

Comment thread src/ucp/core/ucp_mm.c Resolved

Comment thread src/ucp/core/ucp_mm.c Resolved

Comment thread src/ucp/dt/dt.h Resolved

Comment thread src/ucp/rma/rma_rndv.c Resolved

Comment thread src/uct/api/v2/uct_v2.h Resolved

Comment thread src/uct/cuda/cuda_copy/cuda_copy_md.c Resolved

UCP/UCT: Rename mem_info.mem_flags to flags, separate v2 api enum fro…

64954ec

…m v1 Also assert on different mem_type from detected mem_type.

tomerg-nvidia requested review from brminich and tvegas1 June 22, 2026 10:31

brminich reviewed Jun 22, 2026

View reviewed changes

tvegas1 reviewed Jun 22, 2026

View reviewed changes

Comment thread src/ucs/memory/memtype_cache.h Resolved

Comment thread test/gtest/ucp/test_ucp_proto.cc Resolved

Comment thread test/gtest/ucp/test_ucp_rma.cc Resolved

UCP/UCT/UCS: CR fixes

d07bfdc

Don't query twice the pointer for asyncness Refactoring tests

tomerg-nvidia requested review from brminich and tvegas1 June 23, 2026 06:59

brminich approved these changes Jun 23, 2026

View reviewed changes

tvegas1 approved these changes Jun 23, 2026

View reviewed changes

brminich merged commit b6b8531 into openucx:master Jun 23, 2026
160 checks passed

svc-ucx mentioned this pull request Aug 24, 2026

GTEST/UCS: Increase rwlock test sleep to avoid CI races #11809

Merged

guy-ealey-morag mentioned this pull request Oct 7, 2026

UCP/RNDV: Use local memory detection for RTR_REQ responder #12041

Merged

TomAugspurger mentioned this pull request Oct 20, 2026

Flaky segmentation fault in cudf libcudf_streaming_ucxx_tests tests rapidsai/ucxx#783

Closed

		ucp_mem_info_detect_for_type(context, buffer, length, mem_type,
		&mem_info);

		mem_info = ucp_memory_info_from_memh(memh);
		status = ucp_memh_get(context, address, length, &mem_info, cache_md_map,

Conversation

tomerg-nvidia commented Jun 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What?

Why?

How?

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

guy-ealey-morag commented Jun 15, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

brminich left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tomerg-nvidia commented Jun 11, 2026 •

edited

Loading