DEVICE/API: Change addresses to offsets and add channel_id

src/api/gpu/ucx/nixl_device.cuh

@@ -74,29 +74,29 @@ nixlGpuConvertUcsStatus(ucs_status_t status) {
 /**
  * @brief Post a memory transfer request to the GPU.
  *
- * @param req_hndl    [in]  Request handle.
- * @param index       [in]  Index of the memory descriptor in the transfer request.
- * @param addr        [in]  Local address of the memory to be transferred.
- * @param remote_addr [in]  Remote address of the memory to be transferred to.
- * @param size        [in]  Size of the memory to be transferred.
- * @param is_no_delay [in]  Whether to use no-delay mode.
- * @param xfer_status [out] Status of the transfer. If null, the status is not reported.
+ * @param req_hndl      [in]  Request handle.
+ * @param index         [in]  Index of the memory descriptor in the transfer request.
+ * @param local_offset  [in]  Local offset of the memory to be transferred.
+ * @param remote_offset [in]  Remote offset of the memory to be transferred to.
+ * @param size          [in]  Size of the memory to be transferred.
+ * @param is_no_delay   [in]  Whether to use no-delay mode.
+ * @param xfer_status   [out] Status of the transfer. If null, the status is not reported.
  *
  * @return nixl_status_t    Error code if call was not successful
  */
 template<nixl_gpu_level_t level = nixl_gpu_level_t::THREAD>
 __device__ nixl_status_t
 nixlGpuPostSingleWriteXferReq(nixlGpuXferReqH req_hndl,
                               unsigned index,
-                              const void *addr,
-                              uint64_t remote_addr,
+                              const void *local_offset,
+                              uint64_t remote_offset,
                               size_t size,
                               bool is_no_delay = true,
                               nixlGpuXferStatusH *xfer_status = nullptr) {
     const nixlGpuXferReqParams params{req_hndl, is_no_delay, xfer_status};
 
     ucs_status_t status = ucp_device_put_single<static_cast<ucs_device_level_t>(level)>(
-        params.mem_list, index, addr, remote_addr, size, params.flags, params.ucp_request);
+        params.mem_list, index, local_offset, remote_offset, size, params.flags, params.ucp_request);
 
     return nixlGpuConvertUcsStatus(status);
 }
@@ -105,8 +105,8 @@ nixlGpuPostSingleWriteXferReq(nixlGpuXferReqH req_hndl,
  * @brief Post a signal transfer request to the GPU.
  *
  * @param req_hndl           [in]  Request handle.
- * @param index              [in]  Index of the signal to be transferred.
  * @param signal             [in]  Signal to be sent.
+ * @param signal_index       [in]  Index of the signal to be transferred.
  * @param is_no_delay        [in]  Whether to use no-delay mode.
  * @param xfer_status        [out] Status of the transfer. If null, the status is not reported.
  *
@@ -115,14 +115,14 @@ nixlGpuPostSingleWriteXferReq(nixlGpuXferReqH req_hndl,
 template<nixl_gpu_level_t level = nixl_gpu_level_t::THREAD>
 __device__ nixl_status_t
 nixlGpuPostSignalXferReq(nixlGpuXferReqH req_hndl,
-                         unsigned index,
                          const nixlGpuSignal &signal,
+                         unsigned signal_index,
                          bool is_no_delay = true,
                          nixlGpuXferStatusH *xfer_status = nullptr) {
     const nixlGpuXferReqParams params{req_hndl, is_no_delay, xfer_status};
 
     ucs_status_t status = ucp_device_counter_inc<static_cast<ucs_device_level_t>(level)>(
-        params.mem_list, index, signal.inc, signal.remote_addr, params.flags, params.ucp_request);
+        params.mem_list, signal_index, signal.inc, signal.remote_addr, params.flags, params.ucp_request);
 
     return nixlGpuConvertUcsStatus(status);
 }
@@ -132,12 +132,13 @@ nixlGpuPostSignalXferReq(nixlGpuXferReqH req_hndl,
  *
  * @param req_hndl           [in]  Request handle.
  * @param count              [in]  Number of blocks to send. This is also the length of the arrays
- *                                 @a indices, @a sizes, @a addrs, and @a remote_addrs.
+ *                                 @a indices, @a sizes, @a local_offsets, and @a remote_offsets.
  * @param indices            [in]  Indices of the blocks to send.
  * @param sizes              [in]  Sizes of the blocks to send.
- * @param addrs              [in]  Addresses of the blocks to send.
- * @param remote_addrs       [in]  Remote addresses of the blocks to send to.
+ * @param local_offsets      [in]  Local offsets of the blocks to send.
+ * @param remote_offsets     [in]  Remote offsets of the blocks to send to.
  * @param signal             [in]  Signal to be sent.
+ * @param signal_index       [in]  Index of the signal to be sent.
  * @param is_no_delay        [in]  Whether to use no-delay mode.
  * @param xfer_status        [out] Status of the transfer. If null, the status is not reported.
  *
@@ -149,8 +150,8 @@ nixlGpuPostPartialWriteXferReq(nixlGpuXferReqH req_hndl,
                                size_t count,
                                const unsigned *indices,
                                const size_t *sizes,
-                               void *const *addrs,
-                               const uint64_t *remote_addrs,
+                               void *const *local_offsets,
+                               const uint64_t *remote_offsets,
                                const nixlGpuSignal &signal,
                                unsigned signal_index,
                                bool is_no_delay = true,
@@ -161,8 +162,8 @@ nixlGpuPostPartialWriteXferReq(nixlGpuXferReqH req_hndl,
         ucp_device_put_multi_partial<static_cast<ucs_device_level_t>(level)>(params.mem_list,
                                                                              indices,
                                                                              count,
-                                                                             addrs,
-                                                                             remote_addrs,
+                                                                             local_offsets,
+                                                                             remote_offsets,
                                                                              sizes,
                                                                              signal_index,
                                                                              signal.inc,
@@ -178,13 +179,13 @@ nixlGpuPostPartialWriteXferReq(nixlGpuXferReqH req_hndl,
  *
  * @param req_hndl           [in]  Request handle.
  * @param sizes              [in]  Sizes of the blocks to send.
- * @param addrs              [in]  Addresses of the blocks to send.
- * @param remote_addrs       [in]  Remote addresses of the blocks to send to.
+ * @param offsets            [in]  Offsets of the blocks to send.
+ * @param remote_offsets     [in]  Remote offsets of the blocks to send to.
  * @param signal             [in]  Signal to be sent.
  * @param is_no_delay        [in]  Whether to use no-delay mode.
  * @param xfer_status        [out] Status of the transfer. If null, the status is not reported.
  *
- * @note The arrays @a sizes, @a addrs, and @a remote_addrs must have the same length, which
+ * @note The arrays @a sizes, @a offsets, and @a remote_offsets must have the same length, which
  *       corresponds to the number of blocks to transfer as specified in @a req_hndl.
  *
  * @return nixl_status_t    Error code if call was not successful
@@ -193,17 +194,17 @@ template<nixl_gpu_level_t level = nixl_gpu_level_t::THREAD>
 __device__ nixl_status_t
 nixlGpuPostWriteXferReq(nixlGpuXferReqH req_hndl,
                         const size_t *sizes,
-                        void *const *addrs,
-                        const uint64_t *remote_addrs,
+                        void *const *offsets,
+                        const uint64_t *remote_offsets,
                         const nixlGpuSignal &signal,
                         bool is_no_delay = true,
                         nixlGpuXferStatusH *xfer_status = nullptr) {
     const nixlGpuXferReqParams params{req_hndl, is_no_delay, xfer_status};
 
     ucs_status_t status =
         ucp_device_put_multi<static_cast<ucs_device_level_t>(level)>(params.mem_list,
-                                                                     addrs,
-                                                                     remote_addrs,
+                                                                     offsets,
+                                                                     remote_offsets,
                                                                      sizes,
                                                                      signal.inc,
                                                                      signal.remote_addr,

src/plugins/ucx/ucx_backend.cpp

@@ -1654,19 +1654,22 @@ nixlUcxEngine::createGpuXferReq(const nixlBackendReqH &req_hndl,
 
     std::vector<nixlUcxMem> local_mems;
     std::vector<const nixl::ucx::rkey *> remote_rkeys;
+    std::vector<uint64_t> remote_addrs;
     local_mems.reserve(local_descs.descCount());
     remote_rkeys.reserve(remote_descs.descCount());
+    remote_addrs.reserve(remote_descs.descCount());
 
     for (size_t i = 0; i < static_cast<size_t>(local_descs.descCount()); i++) {
         auto localMd = static_cast<nixlUcxPrivateMetadata *>(local_descs[i].metadataP);
         auto remoteMdDesc = static_cast<nixlUcxPublicMetadata *>(remote_descs[i].metadataP);
 
         local_mems.push_back(localMd->mem);
         remote_rkeys.push_back(&remoteMdDesc->getRkey(workerId));
+        remote_addrs.push_back(static_cast<uint64_t>(remote_descs[i].addr));
     }
 
     try {
-        gpu_req_hndl = nixl::ucx::createGpuXferReq(*ep, local_mems, remote_rkeys);
+        gpu_req_hndl = nixl::ucx::createGpuXferReq(*ep, local_mems, remote_rkeys, remote_addrs);
         return NIXL_SUCCESS;
     }
     catch (const std::exception &e) {

src/utils/ucx/gpu_xfer_req_h.cpp

@@ -34,29 +34,34 @@ namespace nixl::ucx {
 nixlGpuXferReqH
 createGpuXferReq(const nixlUcxEp &ep,
                  const std::vector<nixlUcxMem> &local_mems,
-                 const std::vector<const nixl::ucx::rkey *> &remote_rkeys) {
+                 const std::vector<const nixl::ucx::rkey *> &remote_rkeys,
+                 const std::vector<uint64_t> &remote_addrs) {
     nixl_status_t status = ep.checkTxState();
     if (status != NIXL_SUCCESS) {
         throw std::runtime_error("Endpoint not in valid state for creating memory list");
     }
 
-    if (local_mems.empty() || remote_rkeys.empty()) {
-        throw std::invalid_argument("Empty memh or rkey lists provided");
+    if (local_mems.empty() || remote_rkeys.empty() || remote_addrs.empty()) {
+        throw std::invalid_argument("Empty memory, rkey, or address lists provided");
     }
 
-    if (local_mems.size() != remote_rkeys.size()) {
-        throw std::invalid_argument("Local memh and remote rkey lists must have same size");
+    if (local_mems.size() != remote_rkeys.size() || local_mems.size() != remote_addrs.size()) {
+        throw std::invalid_argument(
+            "Local memory, remote rkey, and remote address lists must have same size");
     }
 
     std::vector<ucp_device_mem_list_elem_t> ucp_elements;
     ucp_elements.reserve(local_mems.size());
 
     for (size_t i = 0; i < local_mems.size(); i++) {
         ucp_device_mem_list_elem_t ucp_elem;
-        ucp_elem.field_mask =
-            UCP_DEVICE_MEM_LIST_ELEM_FIELD_MEMH | UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY;
+        ucp_elem.field_mask = UCP_DEVICE_MEM_LIST_ELEM_FIELD_MEMH |
+            UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY | UCP_DEVICE_MEM_LIST_ELEM_FIELD_L_ADDR |
+            UCP_DEVICE_MEM_LIST_ELEM_FIELD_R_ADDR;
         ucp_elem.memh = local_mems[i].getMemh();
         ucp_elem.rkey = remote_rkeys[i]->get();
+        ucp_elem.l_addr = local_mems[i].getBase();
+        ucp_elem.r_addr = remote_addrs[i];
         ucp_elements.push_back(ucp_elem);
     }
 

src/utils/ucx/gpu_xfer_req_h.h

@@ -29,9 +29,10 @@ namespace nixl::ucx {
 class rkey;
 
 nixlGpuXferReqH
-createGpuXferReq(const nixlUcxEp &,
-                 const std::vector<nixlUcxMem> &,
-                 const std::vector<const nixl::ucx::rkey *> &);
+createGpuXferReq(const nixlUcxEp &ep,
+                 const std::vector<nixlUcxMem> &local_mems,
+                 const std::vector<const nixl::ucx::rkey *> &remote_rkeys,
+                 const std::vector<uint64_t> &remote_addrs);
 
 void
 releaseGpuXferReq(nixlGpuXferReqH gpu_req) noexcept;

src/utils/ucx/ucx_utils.h

@@ -165,6 +165,11 @@ class nixlUcxMem {
         return memh;
     }
 
+    [[nodiscard]] void *
+    getBase() const noexcept {
+        return base;
+    }
+
     friend class nixlUcxWorker;
     friend class nixlUcxContext;
     friend class nixlUcxEp;