UCP/DEVICE: Add multi-lane support by Artemy-Mellanox · Pull Request #11268 · openucx/ucx

Artemy-Mellanox · 2026-03-16T16:09:06Z

No description provided.

rakhmets · 2026-04-20T08:21:32Z


+static ucp_lane_index_t ucp_device_ep_find_lane(const ucp_ep_h ep, ucp_rsc_index_t tl_id)
+{
+    ucp_ep_config_t *ep_config = ucp_ep_config(ep);


Wouldn't it better to use ucp_ep_num_lanes to get num lanes here, as you don't need ep_config var in this function?

rakhmets · 2026-04-20T08:27:48Z

+    iface_attr->ctl_device = uct_cuda_get_cuda_device(iface->cuda_dev);
+    iface_attr->dev_num_paths = 1;


Suggested change

iface_attr->ctl_device = uct_cuda_get_cuda_device(iface->cuda_dev);

iface_attr->dev_num_paths = 1;

iface_attr->ctl_device = uct_cuda_get_cuda_device(iface->cuda_dev);

iface_attr->dev_num_paths = 1;

Sign in to view

-            (remote_sys_dev != lane_key->dst_sys_dev)) {
-            ucs_trace("lane[%u] wrong destination sys_dev: dst_sys_dev=%u",
-                      lane, lane_key->dst_sys_dev);
+        if (wiface->attr.ctl_device != local_sys_dev) {


Artemy-Mellanox · 2026-04-28T07:50:32Z

Sign in to view

+    void *local_addr;
+
+    ucp_device_get_tl_bitmap(worker, tl_bitmap, local_sys_dev);
+    num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[tl_type]);


we don't need to pack anything lane specific for no-lkey lanes, so we don't allocate data for it

roiedanino · 2026-04-27T13:09:57Z

Sign in to view

                   unsigned mem_list_index, void **addr_p)
 {
-    const size_t elem_size = sizeof(uct_device_remote_mem_list_elem_t);
+    const size_t elem_size = sizeof(uct_device_remote_mem_elem_t);


Copilot

Pull request overview

Adds multi-lane support for UCP device memory lists and updates UCT device-element packing/types so device-side operations can select among multiple device-EP lanes.

Changes:

Introduce new device memory element/layout types (uct_device_mem_elem_t, per-lane TL arrays) and propagate signature updates across UCT/UCP and CUDA test kernels.
Extend UCT iface attributes/capabilities with ctl_device and UCT_IFACE_FLAG_DEVICE_LKEY to support correct device endpoint selection and packing behavior.
Rework UCP device mem-list creation/packing to build multi-lane descriptors and adjust wireup path selection behavior for device lanes.

Reviewed changes

Copilot reviewed 27 out of 27 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
test/gtest/uct/test_device.cc	Updates CUDA context handling and switches tests to the new packed device mem element type.
test/gtest/uct/cuda/test_kernels_uct.h	Updates kernel launcher prototypes to use `uct_device_mem_elem_t`.
test/gtest/uct/cuda/test_kernels_uct.cu	Updates CUDA kernels/launchers to accept `uct_device_mem_elem_t`.
test/gtest/uct/cuda/test_kernels.h	Updates test kernel APIs to use `uct_device_mem_elem_t`.
test/gtest/uct/cuda/test_kernels.cu	Updates CUDA kernels to use `uct_device_mem_elem_t`.
test/gtest/uct/cuda/test_cuda_ipc_device.cc	Updates CUDA IPC device tests to the new packed mem element type.
test/gtest/ucp/test_ucp_device.cc	Adjusts kernel test expectations/usage to match updated device structures.
test/gtest/ucp/cuda/test_kernels.h	Removes per-kernel state fields from the device-kernel result struct.
test/gtest/ucp/cuda/test_kernels.cu	Removes device-side state parsing logic tied to the old lane/layout assumptions.
src/uct/ib/mlx5/gdaki/gdaki.cuh	Updates GDA device EP APIs to accept `uct_device_mem_elem_t` and new source element semantics.
src/uct/ib/mlx5/gdaki/gdaki.c	Exposes DEVICE_LKEY capability and sets `ctl_device` for GDA iface.
src/uct/ib/mlx5/dv/ib_mlx5dv_md.c	Updates MD mem element pack signature to the new type.
src/uct/cuda/cuda_ipc/cuda_ipc_md.c	Updates CUDA IPC MD mem element pack signature to the new type.
src/uct/cuda/cuda_ipc/cuda_ipc_iface.c	Initializes `ctl_device` for CUDA IPC iface attributes.
src/uct/cuda/cuda_ipc/cuda_ipc.cuh	Updates CUDA IPC device EP functions to take `uct_device_mem_elem_t`.
src/uct/base/uct_md.h	Updates MD ops callback typedef to use `uct_device_mem_elem_t`.
src/uct/base/uct_md.c	Updates `uct_md_mem_elem_pack()` signature to use `uct_device_mem_elem_t`.
src/uct/api/v2/uct_v2.h	Updates the public v2 API signature for `uct_md_mem_elem_pack()`.
src/uct/api/uct_def.h	Renames/introduces device element typedefs for multi-lane descriptors.
src/uct/api/uct.h	Adds `UCT_IFACE_FLAG_DEVICE_LKEY` and `uct_iface_attr::ctl_device`.
src/uct/api/device/uct_device_types.h	Redefines device memory list element layouts for multi-lane support.
src/uct/api/device/uct_device_impl.h	Updates device EP put/atomic/get_ptr APIs to use `uct_device_mem_elem_t`.
src/ucp/wireup/select.c	Adjusts path-count logic for DEVICE lane selection.
src/ucp/wireup/address.c	Forces `num_paths=1` for device endpoints when gathering devices.
src/ucp/core/ucp_device.c	Reworks host-side packing/allocation of local/remote device mem list handles for multi-lane descriptors.
src/ucp/api/device/ucp_device_types.h	Adds `num_lanes` to device mem list handle headers and updates element types.
src/ucp/api/device/ucp_device_impl.h	Adds lane selection macros and updates device-side send/put/counter paths for multi-lane descriptors.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-27T12:50:23Z

+    for (i = 0; i < num_lanes;) {
+        UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, tl_bitmap) {
+            status = ucp_device_remote_mem_list_element_pack(ucp_element, tl_id,
+                                                             tl_element);
+            if (status != UCS_OK) {
+                return status;
+            }
+
+            tl_element = UCS_PTR_TYPE_OFFSET(tl_element, *tl_element);
+            i++;
+        }
+    }
+


ucp_device_remote_mem_list_fill() repeats packing the TL bitmap until i == num_lanes. If UCS_STATIC_BITMAP_POPCOUNT(*tl_bitmap) differs from num_lanes, this will either duplicate TL entries (popcount < num_lanes) or overflow the allocated tl[] array (popcount > num_lanes). It would be safer to iterate the bitmap once and assert/verify that the number of packed entries equals num_lanes.

Suggested change

for (i = 0; i < num_lanes;) {

UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, tl_bitmap) {

status = ucp_device_remote_mem_list_element_pack(ucp_element, tl_id,

tl_element);

if (status != UCS_OK) {

return status;

}

tl_element = UCS_PTR_TYPE_OFFSET(tl_element, *tl_element);

i++;

}

}

i = 0;

UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, tl_bitmap) {

if (i >= num_lanes) {

return UCS_ERR_INVALID_PARAM;

}

status = ucp_device_remote_mem_list_element_pack(ucp_element, tl_id,

tl_element);

if (status != UCS_OK) {

return status;

}

tl_element = UCS_PTR_TYPE_OFFSET(tl_element, *tl_element);

i++;

}

if (i != num_lanes) {

return UCS_ERR_INVALID_PARAM;

}

Copilot · 2026-04-27T12:50:24Z

+                }
+            }
+
+            ucs_assert(tl_type < UCP_DEVICE_TL_TYPE_LAST);


If no tl_type bitmap matches an element EP, this ucs_assert() will abort the process. This can be triggered by runtime configuration mismatches and should return an error status (propagated to the caller) rather than asserting, so applications can handle it gracefully.

Suggested change

ucs_assert(tl_type < UCP_DEVICE_TL_TYPE_LAST);

if (tl_type == UCP_DEVICE_TL_TYPE_LAST) {

ucs_error("failed to find matching tl_type for remote mem "

"list element");

status = UCS_ERR_INVALID_PARAM;

goto out;

}

Copilot · 2026-04-27T12:50:24Z

        /* Interface capability */
 #define UCT_IFACE_FLAG_INTER_NODE      UCS_BIT(54) /**< Interface is inter-node capable */
 #define UCT_IFACE_FLAG_DEVICE_EP       UCS_BIT(55) /**< Interface supports device endpoint */
+#define UCT_IFACE_FLAG_DEVICE_LKEY     UCS_BIT(56) /**< Interface require lkey for device operations */


Grammar in the flag description should be singular: "Interface requires lkey for device operations".

Suggested change

#define UCT_IFACE_FLAG_DEVICE_LKEY UCS_BIT(56) /**< Interface require lkey for device operations */

#define UCT_IFACE_FLAG_DEVICE_LKEY UCS_BIT(56) /**< Interface requires lkey for device operations */

Copilot · 2026-04-27T12:50:25Z

+    const auto mem_element = static_cast<uct_device_remote_mem_elem_t*>(
            UCS_PTR_BYTE_OFFSET(mem_list_h->mem_elements,
                                mem_list_index * elem_size));

-    return uct_device_ep_get_ptr(mem_element->device_ep,
-                                 &mem_element->uct_mem_element,
+    return uct_device_ep_get_ptr(mem_element->tl[0].ep, &mem_element->tl[0].uct,


ucp_device_get_ptr() indexes mem_elements with a fixed sizeof(uct_device_remote_mem_elem_t) stride, but remote mem list elements are now variable-sized (depend on mem_list_h->num_lanes). This will compute the wrong element address and can read out-of-bounds when num_lanes > 1. Use the same stride calculation as UCP_DEVICE_GET_ELEM(...) (or reuse the macro) to locate the mem_list_index element.

Copilot · 2026-04-27T12:50:25Z

+    int tl_type = UCP_DEVICE_TL_TYPE_LKEY;
+    ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_LAST] = {};


ucp_device_local_mem_list_create_handle() initializes tl_type to UCP_DEVICE_TL_TYPE_LKEY and never switches to UCP_DEVICE_TL_TYPE_NOLKEY if there are no DEVICE_LKEY transports. This can lead to exporting a handle with num_lanes == 0, which later causes division/modulo by zero in UCP_DEVICE_GET_LANE() and leaves no packed TL elements. Add a fallback/selection based on which bitmap is non-empty (or return an error) to guarantee num_lanes >= 1.

Copilot · 2026-04-27T12:50:25Z

+    ucp_device_get_tl_bitmap(ep->worker, tl_bitmap, local_sys_dev);
+    num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[UCP_DEVICE_TL_TYPE_LKEY]);
+    if (!num_lanes) {


num_lanes is derived only from the LKEY bitmap here, but later tl_type may be selected per-element (LKEY vs NOLKEY). If an element ends up using the NOLKEY bitmap (different popcount), the code will pack a different number of TL entries than num_lanes implies, risking duplicate entries or out-of-bounds writes. Consider selecting a single tl_type/num_lanes for the whole handle or validating that all elements match the chosen bitmap size.

roiedanino · 2026-04-27T12:59:34Z

-        uct_element = UCS_PTR_BYTE_OFFSET(uct_element, uct_elem_size);
    }

    handle->version = UCP_DEVICE_MEM_LIST_VERSION_V1;


Should we consider UCP_DEVICE_MEM_LIST_MULTI_LANE as the handle structure changed?

roiedanino · 2026-04-27T13:06:16Z

+#define UCP_DEVICE_GET_LANE(_handle, _channel_id, _lane, _uct_channel_id) \
+    _lane           = _channel_id % _handle->num_lanes; \
+    _uct_channel_id = _channel_id / _handle->num_lanes;


ucs_assertv(_handle != NULL, "dest memlist can't be null")
before deref, clearer failure for users

in other APIs we not validate objects for NULL

roiedanino · 2026-04-27T13:06:41Z

+#define UCP_DEVICE_GET_ELEM(_handle, _index, _lane) \
+    static_cast<ucs_typeof(_handle->mem_elements[0])*>(UCS_PTR_BYTE_OFFSET( \
+            _handle->mem_elements, \
+            (sizeof(_handle->mem_elements[0]) + \
+             (sizeof(_handle->mem_elements[0].tl[0]) * _handle->num_lanes)) * \
+                    _index))


same idea here

roiedanino · 2026-04-27T13:13:10Z

+    _uct_channel_id = _channel_id / _handle->num_lanes;
+
+
+#define UCP_DEVICE_GET_ELEM(_handle, _index, _lane) \


Suggested change

#define UCP_DEVICE_GET_ELEM(_handle, _index, _lane) \

#define UCP_DEVICE_GET_ELEM(_handle, _index) \

roiedanino · 2026-04-27T13:14:45Z

    return status;
 }

+static ucp_lane_index_t ucp_device_ep_find_lane(const ucp_ep_h ep, ucp_rsc_index_t tl_id)


Is it safe to assume tl_id is a device lane?

device API working with device lanes, but this function is not specific

yosefe · 2026-05-03T13:25:51Z

+#define UCP_DEVICE_GET_ELEM(_handle, _index) \
+    static_cast<ucs_typeof(_handle->mem_elements[0])*>(UCS_PTR_BYTE_OFFSET( \


can it be inline function?

yosefe · 2026-05-03T13:27:07Z

+    static_cast<ucs_typeof(_handle->mem_elements[0])*>(UCS_PTR_BYTE_OFFSET( \
+            _handle->mem_elements, \
+            (sizeof(_handle->mem_elements[0]) + \
+             (sizeof(_handle->mem_elements[0].tl[0]) * _handle->num_lanes)) * \


how much it affects performance?

yosefe · 2026-05-03T13:52:14Z

+    /**
+     * Number of lanes in each memory descriptor.
+     */
+    uint16_t                     num_lanes;


maybe uint8_t?

yosefe · 2026-05-03T13:53:33Z

+    /**
+     * Number of lanes in each memory descriptor.
+     */
+    uint16_t                    num_lanes;


consider optimizing this to be something like MAX_LANES or MAX_MDS

yosefe · 2026-05-03T13:55:18Z

+    int tl_type = UCP_DEVICE_TL_TYPE_LKEY;
+    ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_LAST] = {};


initialized vars should be first

yosefe · 2026-05-03T14:13:17Z

+    uct_element->addr = ucp_element->remote_addr;
+    tl_element        = uct_element->tl;
+    for (i = 0; i < num_lanes;) {
+        UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, tl_bitmap) {


does it mean we could pack tl which does not have a lane on the current endpoint? because tl_bitmap was not created according to ep lanes

yosefe · 2026-05-03T14:13:46Z

+        if (!(iface_attr->cap.flags & UCT_IFACE_FLAG_DEVICE_EP)) {
+            dev->num_paths = ucs_min(max_num_paths, iface_attr->dev_num_paths);


yosefe · 2026-05-03T14:15:07Z

        local_num_paths  = iface_attr->dev_num_paths;
-        remote_num_paths = ae->dev_num_paths;
+
+        if (bw_info->criteria.lane_type == UCP_LANE_TYPE_DEVICE) {


looks weird..
what is the value of iface_attr->dev_num_paths ?

yosefe · 2026-05-03T14:16:15Z

        /* Interface capability */
 #define UCT_IFACE_FLAG_INTER_NODE      UCS_BIT(54) /**< Interface is inter-node capable */
 #define UCT_IFACE_FLAG_DEVICE_EP       UCS_BIT(55) /**< Interface supports device endpoint */
+#define UCT_IFACE_FLAG_DEVICE_LKEY     UCS_BIT(56) /**< Interface requires lkey for device operations */


why needed?
and can we add it to query v2?

yosefe · 2026-05-03T14:16:19Z

                                                compared to using only a single
                                                endpoint. */
+
+    ucs_sys_device_t ctl_device;           /**< System device controlling this iface,


can we add it to query v2?

ofirfarjun7

Fix comments in next PR

Artemy-Mellanox force-pushed the multilane branch 2 times, most recently from ca9ba6b to af04b9e Compare March 16, 2026 21:42

UCP/DEVICE: Add multi-lane support

fa7e9fb

Artemy-Mellanox force-pushed the multilane branch from af04b9e to fa7e9fb Compare March 17, 2026 01:25

UCP/DEVICE: Add multi-lane support - 2

60e5be9

Artemy-Mellanox force-pushed the multilane branch from 7714a0b to 60e5be9 Compare March 18, 2026 12:37

ofirfarjun7 reviewed Mar 21, 2026

View reviewed changes

Comment thread src/ucp/api/device/ucp_device_impl.h Outdated

Comment thread src/ucp/core/ucp_device.c Outdated

Comment thread src/ucp/core/ucp_device.c Outdated

Comment thread src/ucp/core/ucp_device.c Outdated

Artemy-Mellanox added 2 commits March 23, 2026 15:10

UCP/DEVICE: Add multi-lane support - 3

82faeb9

UCP/DEVICE: Add multi-lane support - 4

1f3445e

Artemy-Mellanox force-pushed the multilane branch from 8127934 to b59dcb1 Compare March 29, 2026 04:23

UCP/DEVICE: Add multi-lane support - 5

71d7b9c

Artemy-Mellanox force-pushed the multilane branch from b59dcb1 to 71d7b9c Compare March 29, 2026 09:50

rakhmets reviewed Apr 14, 2026

View reviewed changes

UCP/DEVICE: Add multi-lane support - 6

c863ad2

rakhmets reviewed Apr 20, 2026

View reviewed changes

jeynmann reviewed Apr 23, 2026

View reviewed changes

Comment thread src/ucp/core/ucp_device.c Outdated

(remote_sys_dev != lane_key->dst_sys_dev)) {

ucs_trace("lane[%u] wrong destination sys_dev: dst_sys_dev=%u",

lane, lane_key->dst_sys_dev);

if (wiface->attr.ctl_device != local_sys_dev) {

This comment was marked as resolved.

Sign in to view

ofirfarjun7 reviewed Apr 23, 2026

View reviewed changes

Comment thread src/ucp/core/ucp_device.c Outdated

Comment thread src/ucp/core/ucp_device.c Outdated

Artemy-Mellanox added 5 commits April 24, 2026 15:36

UCP/DEVICE: Add multi-lane support - 7

e2b4ec0

UCP/DEVICE: Add multi-lane support - 8

ed5762b

Merge remote-tracking branch 'origin/master' into multilane

3e4eeea

UCP/DEVICE: Add multi-lane support - 9

7986ea3

UCP/DEVICE: Add multi-lane support - 10

72d8beb

jeynmann reviewed Apr 27, 2026

View reviewed changes

roiedanino requested a review from Copilot April 27, 2026 12:43

Copilot started reviewing on behalf of roiedanino April 27, 2026 12:44 View session

Copilot AI reviewed Apr 27, 2026

View reviewed changes

roiedanino reviewed Apr 27, 2026

View reviewed changes

ofirfarjun7 reviewed Apr 27, 2026

View reviewed changes

Comment thread src/ucp/api/device/ucp_device_impl.h Outdated

Comment thread src/ucp/api/device/ucp_device_impl.h Outdated

roiedanino reviewed Apr 27, 2026

View reviewed changes

ofirfarjun7 reviewed Apr 27, 2026

View reviewed changes

Comment thread src/ucp/core/ucp_device.c

Comment thread src/uct/cuda/cuda_ipc/cuda_ipc_iface.c

Artemy-Mellanox added 2 commits April 27, 2026 19:05

UCP/DEVICE: Add multi-lane support - 11

f0228c0

UCP/DEVICE: Add multi-lane support - 12

508fa22

yosefe requested changes May 3, 2026

View reviewed changes

openucx deleted a comment from svc-nixl May 4, 2026

Merge remote-tracking branch 'origin/master' into multilane

eb078cc

ofirfarjun7 approved these changes May 5, 2026

View reviewed changes

ofirfarjun7 merged commit b788849 into openucx:master May 6, 2026
153 checks passed

		iface_attr->ctl_device = uct_cuda_get_cuda_device(iface->cuda_dev);
		iface_attr->dev_num_paths = 1;

-            ucs_assert(tl_type < UCP_DEVICE_TL_TYPE_LAST);
+            if (tl_type == UCP_DEVICE_TL_TYPE_LAST) {
+                ucs_error("failed to find matching tl_type for remote mem "
+                          "list element");
+                status = UCS_ERR_INVALID_PARAM;
+                goto out;
+            }

	#define UCT_IFACE_FLAG_DEVICE_LKEY UCS_BIT(56) /*< Interface require lkey for device operations /
	#define UCT_IFACE_FLAG_DEVICE_LKEY UCS_BIT(56) /*< Interface requires lkey for device operations /

		int tl_type = UCP_DEVICE_TL_TYPE_LKEY;
		ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_LAST] = {};

		_uct_channel_id = _channel_id / _handle->num_lanes;


		#define UCP_DEVICE_GET_ELEM(_handle, _index, _lane) \

		#define UCP_DEVICE_GET_ELEM(_handle, _index) \
		static_cast<ucs_typeof(_handle->mem_elements[0])*>(UCS_PTR_BYTE_OFFSET( \

		if (!(iface_attr->cap.flags & UCT_IFACE_FLAG_DEVICE_EP)) {
		dev->num_paths = ucs_min(max_num_paths, iface_attr->dev_num_paths);

Uh oh!

Conversation

Artemy-Mellanox commented Mar 16, 2026 • edited by yosefe Loading Uh oh! There was an error while loading. Please reload this page.

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!

This comment was marked as resolved.

Uh oh!

Uh oh!

Uh oh!

This comment was marked as resolved.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This comment was marked as resolved.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

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!

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!

Artemy-Mellanox commented Mar 16, 2026 •

edited by yosefe

Loading