Skip to content

UCP/CORE: Don't force blocking registration of managed RNDV frag pool. #12263

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rakhmets wants to merge 2 commits into
base: master
Choose a base branch
from
+11 −2
Open

UCP/CORE: Don't force blocking registration of managed RNDV frag pool. #12263

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7571a2e
UCP/CORE: Don't force blocking registration of managed RNDV frag pool.
rakhmets Nov 5, 2026
7aadb96
UCP/CORE: Addressed review comments.
rakhmets Nov 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions src/ucp/core/ucp_mm.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1530,10 +1530,20 @@ ucp_rndv_frag_malloc_mpools(ucs_mpool_t *mp, size_t *size_p, void **chunk_p)
ucs_memory_type_t mem_type = mpriv->mem_type;
ucs_sys_device_t sys_dev = mpriv->sys_dev;
size_t frag_size = context->config.ext.rndv_frag_size[mem_type];
unsigned uct_flags = UCT_MD_MEM_ACCESS_RMA;
ucp_rndv_frag_mp_chunk_hdr_t *chunk_hdr;
ucs_status_t status;
unsigned num_elems;

/* CUDA managed memory cannot be pinned, it can only be registered on demand

@svc-nvidia-pr-review svc-nvidia-pr-review Nov 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we cover the other managed types here? if IB_ODP_MEM_TYPES is set to rocm-managed or ze-managed, those frag pools fail the same way, since reg_mem_types gets the ODP types in uct_ib_md_check_odp().

@tvegas1 tvegas1 Nov 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The fix seems odd: the frag mpool hardcodes a memory type to bypass LOCK, so other callers (e.g. ucp_mem_map() with UCP_MEM_MAP_LOCK on managed memory) still fail. Note that LOCK was never implemented by IB MD; it only uses NONBLOCK to select ODP vs pinned registration.

Suggestions WDYT:

  1. to instead have the memory detector report it? cuda_copy already identifies managed memory by CU_POINTER_ATTRIBUTE_IS_MANAGED and assumes it is registered via ODP (cuda_copy_md.c:893), so it could set a new UCS_MEM_FLAG_NOT_PINNABLE in uct_cuda_copy_md_detect_mem_flags(), and ucp_memh_register_internal() would replace LOCK with NONBLOCK for such memh. This covers all callers and can be extended to rocm/ze-managed. It needs OR semantics for this flag on rcache merge (ucp_mm.c:1765), since a negative flag is lost with AND.

  2. Otherwise, maybe fix it in IB MD: UCP already passes mem_type for cuda-managed, so uct_ib_verbs_mem_reg() and uct_ib_mlx5_devx_mem_reg() can do:

/* CUDA managed memory cannot be pinned, it can be registered only by ODP */
if ((mem_type == UCS_MEMORY_TYPE_CUDA_MANAGED) &&
    (md->reg_nonblock_mem_types & UCS_BIT(mem_type))) {
    flags |= UCT_MD_MEM_FLAG_NONBLOCK;
}

* (e.g. by ODP), regardless of UCX_REG_NONBLOCK_MEM_TYPES. Request
* blocking registration for any other memory type. */
if (UCP_MEM_IS_CUDA_MANAGED(mem_type)) {

@svc-nvidia-pr-review svc-nvidia-pr-review Nov 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why not handle this in ucp_memh_register_internal()? ucp_mem_map() with UCP_MEM_MAP_LOCK on cuda-managed memory hits the same failure, and a central fix would cover all callers.

uct_flags |= UCT_MD_MEM_FLAG_NONBLOCK;
} else {
uct_flags |= UCT_MD_MEM_FLAG_LOCK;
}

/* metadata */
chunk_hdr = ucs_malloc(sizeof(*chunk_hdr) + *size_p, "chunk_hdr");
if (chunk_hdr == NULL) {
Expand All @@ -1545,8 +1555,7 @@ ucp_rndv_frag_malloc_mpools(ucs_mpool_t *mp, size_t *size_p, void **chunk_p)

/* payload; need to get default flags from ucp_mem_map_params2uct_flags() */
status = ucp_memh_alloc(context, NULL, frag_size * num_elems, mem_type,
sys_dev, UCP_MEMH_FLAG_NO_RCACHE,
UCT_MD_MEM_ACCESS_RMA | UCT_MD_MEM_FLAG_LOCK, 0,
sys_dev, UCP_MEMH_FLAG_NO_RCACHE, uct_flags, 0,
ucs_mpool_name(mp), &chunk_hdr->memh);
if (status != UCS_OK) {
return status;

@svc-nvidia-pr-review svc-nvidia-pr-review Nov 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: chunk_hdr is leaked here; pls free it before returning. Pre-existing, but this is exactly the failure path this PR fixes.

Expand Down
Loading