UCP/DEVICE: Make memh and local_addr optional for counter elements #10945

michal-shalev · 2025-10-11T22:12:27Z

What?

Make memh and local_addr optional for counter elements in ucp_device_mem_list_create.

Why?

Counter elements only require remote addressing for atomic operations (ucp_device_counter_inc, or as part of ucp_device_put_multi/put_multi_partial). Requiring local memory registration (memh/local_addr) for these elements is unnecessary overhead.

How?

When memh is not provided, detect local_sys_dev by allocating a temporary buffer on the current CUDA context device (similar to ucp_ep_rma_batch_export)
Updated tests and documentation

ofirfarjun7 · 2025-10-12T12:23:58Z

test/gtest/ucp/test_ucp_device.cc

+                               UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR |
+                               UCP_DEVICE_MEM_LIST_ELEM_FIELD_LENGTH;
+            elem.memh        = NULL;
+            elem.local_addr  = NULL;


can be default to simplify the if

initialized elem to 0

ofirfarjun7 · 2025-10-12T12:24:18Z

src/ucp/core/ucp_device.c

        return status;
    }

+    if (local_sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) {


Maybe move to ucp_device_mem_list_params_check

ofirfarjun7 · 2025-10-12T12:24:35Z

src/ucp/core/ucp_device.c

+                *local_md_map  = memh->md_map;
+                *mem_type      = memh->mem_type;
+            } else {
+                *mem_type      = rkey->mem_type;


Should be cuda for now

ofirfarjun7 · 2025-10-12T12:25:00Z

src/ucp/api/device/ucp_host.h

-    UCP_DEVICE_MEM_LIST_ELEM_FIELD_LENGTH      = UCS_BIT(4)  /**< Length of the local buffer in bytes */
+    UCP_DEVICE_MEM_LIST_ELEM_FIELD_LOCAL_ADDR  = UCS_BIT(2), /**< Local address (optional for counter elements) */
+    UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR = UCS_BIT(3), /**< Remote address */
+    UCP_DEVICE_MEM_LIST_ELEM_FIELD_LENGTH      = UCS_BIT(4)  /**< Length of the local buffer in bytes (optional for counter elements) */


memh, laddr and length should be optional.

length is only optional for partial

ofirfarjun7 · 2025-10-13T09:38:13Z

src/ucp/api/device/ucp_host.h

    UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY        = UCS_BIT(1), /**< Unpacked remote memory key */
-    UCP_DEVICE_MEM_LIST_ELEM_FIELD_LOCAL_ADDR  = UCS_BIT(2), /**< Local address */
-    UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR = UCS_BIT(3),  /**< Remote address */
+    UCP_DEVICE_MEM_LIST_ELEM_FIELD_LOCAL_ADDR  = UCS_BIT(2), /**< Local address (optional for counter elements) */


Actually, it can be optional also for data elements, only rkey is required (we only check for rkey in ucp_device_mem_list_params_check), the rest are optional.
We do need local address for ucp_device_put_multi, but mem list is not bound to a specific API func.

src/ucp/core/ucp_device.c

src/tools/perf/cuda/ucp_cuda_kernel.cu

rakhmets · 2025-10-13T11:59:17Z

src/tools/perf/cuda/ucp_cuda_kernel.cu

-            elems[i].memh       = perf.ucp.send_memh;
-            elems[i].rkey       = perf.ucp.rkey;
-            elems[i].local_addr = UCS_PTR_BYTE_OFFSET(perf.send_buffer, offset);
+            bool is_counter = (i == count - 1);


for (size_t i = 0; i < count - 1; ++i) { elems[i].field_mask = UCP_DEVICE_MEM_LIST_ELEM_FIELD_MEMH | UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY | UCP_DEVICE_MEM_LIST_ELEM_FIELD_LOCAL_ADDR | UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR | UCP_DEVICE_MEM_LIST_ELEM_FIELD_LENGTH; elems[i].memh = perf.ucp.send_memh; elems[i].rkey = perf.ucp.rkey; elems[i].local_addr = UCS_PTR_BYTE_OFFSET(perf.send_buffer, offset); elems[i].remote_addr = perf.ucp.remote_addr + offset; elems[i].length = perf.params.msg_size_list[i]; offset += elems[i].length; } elems[count - 1].field_mask = UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY | UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR | UCP_DEVICE_MEM_LIST_ELEM_FIELD_LENGTH; elems[count - 1].rkey = perf.ucp.rkey; elems[count - 1].remote_addr = perf.ucp.remote_addr + offset; elems[count - 1].length = ONESIDED_SIGNAL_SIZE;

We will not always have a counter, check UCX_PERF_CMD_PUT_SINGLE

size_t data_count = perf.params.msg_size_cnt; for (size_t i = 0; i < data_count; ++i) { elems[i].field_mask = UCP_DEVICE_MEM_LIST_ELEM_FIELD_MEMH | UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY | UCP_DEVICE_MEM_LIST_ELEM_FIELD_LOCAL_ADDR | UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR | UCP_DEVICE_MEM_LIST_ELEM_FIELD_LENGTH; elems[i].memh = perf.ucp.send_memh; elems[i].rkey = perf.ucp.rkey; elems[i].local_addr = UCS_PTR_BYTE_OFFSET(perf.send_buffer, offset); elems[i].remote_addr = perf.ucp.remote_addr + offset; elems[i].length = perf.params.msg_size_list[i]; offset += elems[i].length; } if (m_has_counter) { elems[data_count].field_mask = UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY | UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR | UCP_DEVICE_MEM_LIST_ELEM_FIELD_LENGTH; elems[data_count].rkey = perf.ucp.rkey; elems[data_count].remote_addr = perf.ucp.remote_addr + offset; elems[data_count].length = ONESIDED_SIGNAL_SIZE; }

So you're suggesting to set rkey and remote_addr twice to the same value?

Otherwise, there will be a condition that is checked at each iteration.

src/tools/perf/cuda/ucp_cuda_kernel.cu

rakhmets · 2025-10-13T12:47:42Z

src/tools/perf/cuda/ucp_cuda_kernel.cu

-        /* +1 for the counter */
-        size_t count  = perf.params.msg_size_cnt + 1;
+        size_t count  = perf.params.msg_size_cnt + (m_has_counter ? 1 : 0);
        size_t offset = 0;


Not an issue of this PR.
The variable (offset) can be deleted.

UCP/DEVICE: Make memh and local_addr optional for counter operations

06c380a

michal-shalev self-assigned this Oct 11, 2025

michal-shalev added the Ready for Review label Oct 11, 2025

michal-shalev requested review from Artemy-Mellanox, brminich, iyastreb, ofirfarjun7 and tvegas1 October 11, 2025 22:12

This was referenced Oct 12, 2025

DEVICE/API: Simplify createGpuXferReq to single-step API - WIP ai-dynamo/nixl#816

Closed

DEVICE/API: Simplify createGpuXferReq to single-step API - WIP ai-dynamo/nixl#894

Open

ofirfarjun7 reviewed Oct 12, 2025

View reviewed changes

michal-shalev added 2 commits October 13, 2025 10:21

Merge branch 'master' into optional-memh-localaddr-counter-ops

a494c52

UCP/DEVICE: PR fixes

60a1131

michal-shalev requested a review from ofirfarjun7 October 13, 2025 09:16

michal-shalev added 2 commits October 13, 2025 12:30

UCP/DEVICE: PR fixes 2.0

6a3eb78

UCP/DEVICE: PR fixes 3.0

205b01e

ofirfarjun7 reviewed Oct 13, 2025

View reviewed changes

UCP/DEVICE: PR fixes 4.0

589b0b4

michal-shalev requested a review from ofirfarjun7 October 13, 2025 10:03

michal-shalev added 2 commits October 13, 2025 13:34

UCP/DEVICE: PR fixes 5.0

4573621

UCP/DEVICE: Update perftest

b0eff73

rakhmets reviewed Oct 13, 2025

View reviewed changes

UCP/DEVICE: Update perftest 2.0

10646cd

rakhmets reviewed Oct 13, 2025

View reviewed changes

src/tools/perf/cuda/ucp_cuda_kernel.cu Outdated Show resolved Hide resolved

rakhmets reviewed Oct 13, 2025

View reviewed changes

michal-shalev added 3 commits October 13, 2025 16:39

UCP/DEVICE: PR fixes 6.0

534057b

UCP/DEVICE: CI fix

dc11fd6

UCP/DEVICE: PR fixes 7.0

c718cec

UCP/DEVICE: Make memh and local_addr optional for counter elements #10945

Are you sure you want to change the base?

UCP/DEVICE: Make memh and local_addr optional for counter elements #10945

Conversation

michal-shalev commented Oct 11, 2025

What?

Why?

How?

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!

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!

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!

rakhmets Oct 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

rakhmets Oct 13, 2025 •

edited

Loading