PLUGINS/UCX: Add scatter-gather list (SGL) put path

meson.build

@@ -422,6 +422,16 @@ if ucx_dep.found() and cuda_dep.found() and nvcc_prog.found()
   }, section: 'UCX GPU Device API', bool_yn: true)
 endif
 
+# UCX SGL (scatter-gather list) API detection
+if ucx_dep.found()
+    have_ucx_sgl = cpp.has_type('ucp_dt_local_sgl_t',
+        prefix: '#include <ucp/api/ucp.h>',
+        dependencies: ucx_dep)
+    if have_ucx_sgl
+        add_project_arguments('-DHAVE_UCX_SGL_API', language: 'cpp')
+    endif
+endif
+
 if get_option('disable_gds_backend')
     add_project_arguments('-DDISABLE_GDS_BACKEND', language: 'cpp')
 endif

src/plugins/ucx/meson.build

@@ -22,6 +22,7 @@ ucx_backend_sources = ['config.cpp',
                        'mem_list.cpp',
                        'rkey.cpp',
                        'ucx_backend.cpp',
+                       'ucx_sgl.cpp',
                        'ucx_enums.cpp',
                        'ucx_plugin.cpp',
                        'ucx_utils.cpp']

src/plugins/ucx/ucx_backend.cpp

@@ -16,9 +16,11 @@
  */
 
 #include "ucx_backend.h"
+#include "ucx_sgl.h"
 #include "common/nixl_log.h"
 #include "serdes/serdes.h"
 #include "common/backend.h"
+#include "common/configuration.h"
 #include "common/nixl_log.h"
 
 #include <optional>
@@ -39,8 +41,26 @@ epCloseFlags(const nixl_b_params_t *custom_params) {
         UCP_EP_CLOSE_FLAG_FORCE :
         0;
 }
+
+[[nodiscard]] bool
+sglEnabledFromConfig() {
+    const bool enabled = nixl::config::getValueDefaulted("NIXL_UCX_SGL_ENABLE", false);
+#ifdef HAVE_UCX_SGL_API
+    NIXL_DEBUG << "UCX SGL offload " << (enabled ? "enabled" : "disabled");
+    return enabled;
+#else
+    if (enabled) {
+        NIXL_WARN << "NIXL_UCX_SGL_ENABLE is set but NIXL was built without UCX SGL support";
+    }
+    return false;
+#endif
+}
 } // namespace
 
+// A transfer to a single endpoint posts at most three requests:
+// one data request, one flush request, and one notification request.
+constexpr size_t single_ep_request_count = 3;
+
 /****************************************
  * Backend request management
 *****************************************/
@@ -83,6 +103,10 @@ class nixlUcxBackendReqH : public nixlBackendReqH {
 
     std::optional<Notif> notif;
 
+#ifdef HAVE_UCX_SGL_API
+    std::optional<nixl::ucx::sglXfer> sgl;
+#endif
+
     explicit nixlUcxBackendReqH(nixlUcxWorker *worker) {
         setWorker(worker);
     }
@@ -764,7 +788,8 @@ nixlUcxEngine::create(const nixlBackendInitParams &init_params) {
 
 nixlUcxEngine::nixlUcxEngine(const nixlBackendInitParams &init_params, size_t num_dedicated_workers)
     : nixlBackendEngine(&init_params),
-      sharedWorkerIndex_(1) {
+      sharedWorkerIndex_(1),
+      sglEnabled_(sglEnabledFromConfig()) {
     std::vector<std::string> devs; /* Empty vector */
     nixl_b_params_t *custom_params = init_params.customParams;
 
@@ -1063,6 +1088,12 @@ nixl_status_t nixlUcxEngine::prepXfer (const nixl_xfer_op_t &operation,
     /* TODO: try to get from a pool first */
     handle = new nixlUcxBackendReqH(getSharedWorker(worker_id).get());
 
+#ifdef HAVE_UCX_SGL_API
+    if (sglEnabled_ && operation == NIXL_WRITE) {
+        return prepXferSgl(local, remote, handle);
+    }
+#endif
+
     return NIXL_SUCCESS;
 }
 
@@ -1173,6 +1204,46 @@ nixlUcxEngine::sendXferRangeBatch(nixlUcxEp &ep,
     return result;
 }
 
+#ifdef HAVE_UCX_SGL_API
+nixl_status_t
+nixlUcxEngine::prepXferSgl(const nixl_meta_dlist_t &local,
+                           const nixl_meta_dlist_t &remote,
+                           nixlBackendReqH *handle) const {
+    NIXL_ASSERT(local.descCount() == remote.descCount());
+
+    const auto int_handle = static_cast<nixlUcxBackendReqH *>(handle);
+    int_handle->sgl.emplace(local, remote, int_handle->getWorkerId(), 0, local.descCount());
+    return NIXL_SUCCESS;
+}
+
+nixl_status_t
+nixlUcxEngine::sendXferSgl(nixlBackendReqH *handle) const {
+    const auto int_handle = static_cast<nixlUcxBackendReqH *>(handle);
+    NIXL_ASSERT(int_handle->sgl);
+    auto &sgl = *int_handle->sgl;
+
+    const ucx_connection_ptr_t &conn = sgl.conn();
+
+    auto &ep = conn->getEp(int_handle->getWorkerId());
+
+    int_handle->reserve(single_ep_request_count);
+
+    nixlUcxReq req;
+    const nixl_status_t post_ret = sgl.post(*ep, req);
+    if (int_handle->append(post_ret, req, conn) != NIXL_SUCCESS) {
+        return post_ret;
+    }
+
+    nixlUcxReq flush_req;
+    const nixl_status_t flush_ret = ep->flushEp(flush_req);
+    if (int_handle->append(flush_ret, flush_req, conn) != NIXL_SUCCESS) {
+        return flush_ret;
+    }
+
+    return NIXL_SUCCESS;
+}
+#endif
+
 nixl_status_t
 nixlUcxEngine::sendXferRange(const nixl_xfer_op_t &operation,
                              const nixl_meta_dlist_t &local,
@@ -1188,9 +1259,16 @@ nixlUcxEngine::sendXferRange(const nixl_xfer_op_t &operation,
         return NIXL_ERR_INVALID_PARAM;
     }
 
-    /* Assuming we have a single EP, we need 3 requests: one pending request,
-     * one flush request, and one notification request */
-    int_handle->reserve(3);
+#ifdef HAVE_UCX_SGL_API
+    if (sglEnabled_ && operation == NIXL_WRITE) {
+        if (!int_handle->sgl) {
+            int_handle->sgl.emplace(local, remote, worker_id, start_idx, end_idx);
+        }
+        return sendXferSgl(handle);
+    }
+#endif
+
+    int_handle->reserve(single_ep_request_count);
 
     for (size_t i = start_idx; i < end_idx;) {
         /* Send requests to a single EP */

src/plugins/ucx/ucx_backend.h

@@ -287,6 +287,16 @@ class nixlUcxEngine : public nixlBackendEngine {
                        size_t start_idx,
                        size_t end_idx);
 
+#ifdef HAVE_UCX_SGL_API
+    nixl_status_t
+    prepXferSgl(const nixl_meta_dlist_t &local,
+                const nixl_meta_dlist_t &remote,
+                nixlBackendReqH *handle) const;
+
+    nixl_status_t
+    sendXferSgl(nixlBackendReqH *handle) const;
+#endif
+
     /**
      * Get the worker ID from the optional arguments.
      * Returns std::nullopt if the 'worker_id' option extraction fails.
@@ -300,6 +310,7 @@ class nixlUcxEngine : public nixlBackendEngine {
     size_t numSharedWorkers_;
     std::string workerAddr;
     mutable std::atomic<size_t> sharedWorkerIndex_;
+    const bool sglEnabled_;
 
     // Map of agent name to saved nixlUcxConnection info
     std::unordered_map<std::string, ucx_connection_ptr_t> remoteConnMap;

src/plugins/ucx/ucx_sgl.cpp

@@ -0,0 +1,72 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "ucx_sgl.h"
+
+#ifdef HAVE_UCX_SGL_API
+
+#include "common/nixl_log.h"
+
+namespace nixl::ucx {
+
+sglXfer::sglXfer(const nixl_meta_dlist_t &local,
+                 const nixl_meta_dlist_t &remote,
+                 size_t worker_id,
+                 size_t start_idx,
+                 size_t end_idx) {
+    NIXL_ASSERT(end_idx > start_idx);
+    size_ = end_idx - start_idx;
+    conn_ = static_cast<nixlUcxPublicMetadata *>(remote[start_idx].metadataP)->conn;
+
+    static_assert(sizeof(void *) == sizeof(uint64_t) && sizeof(size_t) == sizeof(uint64_t) &&
+                      sizeof(ucp_mem_h) == sizeof(uint64_t) &&
+                      sizeof(ucp_rkey_h) == sizeof(uint64_t),
+                  "sglXfer assumes all field types are 8 bytes wide");
+    constexpr size_t field_bytes = sizeof(uint64_t);
+    const size_t array_bytes = size_ * field_bytes;
+
+    storage_ = std::make_unique_for_overwrite<std::byte[]>(num_fields * array_bytes);
+    std::byte *const base = storage_.get();
+    localAddrs_ = reinterpret_cast<void **>(base + 0 * array_bytes);
+    remoteAddrs_ = reinterpret_cast<uint64_t *>(base + 1 * array_bytes);
+    lengths_ = reinterpret_cast<size_t *>(base + 2 * array_bytes);
+    memhs_ = reinterpret_cast<ucp_mem_h *>(base + 3 * array_bytes);
+    rkeys_ = reinterpret_cast<ucp_rkey_h *>(base + 4 * array_bytes);
+
+    for (size_t i = start_idx; i < end_idx; ++i) {
+        const size_t out = i - start_idx;
+        const auto lmd = static_cast<nixlUcxPrivateMetadata *>(local[i].metadataP);
+        const auto rmd = static_cast<nixlUcxPublicMetadata *>(remote[i].metadataP);
+        NIXL_ASSERT(local[i].len == remote[i].len);
+        NIXL_ASSERT(rmd->conn == conn_);
+
+        localAddrs_[out] = reinterpret_cast<void *>(local[i].addr);
+        remoteAddrs_[out] = static_cast<uint64_t>(remote[i].addr);
+        lengths_[out] = local[i].len;
+        memhs_[out] = lmd->getMem().getMemh();
+        rkeys_[out] = rmd->getRkey(worker_id).get();
+    }
+}
+
+nixl_status_t
+sglXfer::post(nixlUcxEp &ep, nixlUcxReq &req) const {
+    return ep.postSgl(localView(), remoteView(), size_, req);
+}
+
+} // namespace nixl::ucx
+
+#endif // HAVE_UCX_SGL_API

src/plugins/ucx/ucx_sgl.h

@@ -0,0 +1,89 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef NIXL_SRC_PLUGINS_UCX_UCX_SGL_H
+#define NIXL_SRC_PLUGINS_UCX_UCX_SGL_H
+
+#ifdef HAVE_UCX_SGL_API
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+extern "C" {
+#include <ucp/api/ucp.h>
+}
+
+#include "ucx_backend.h"
+
+namespace nixl::ucx {
+
+class sglXfer {
+public:
+    sglXfer(const nixl_meta_dlist_t &local,
+            const nixl_meta_dlist_t &remote,
+            size_t worker_id,
+            size_t start_idx,
+            size_t end_idx);
+
+    [[nodiscard]] const ucx_connection_ptr_t &
+    conn() const noexcept {
+        return conn_;
+    }
+
+    [[nodiscard]] nixl_status_t
+    post(nixlUcxEp &ep, nixlUcxReq &req) const;
+
+private:
+    [[nodiscard]] ucp_dt_local_sgl_t
+    localView() const noexcept {
+        return {
+            .field_mask = UCP_DT_LOCAL_SGL_FIELD_BUFFERS | UCP_DT_LOCAL_SGL_FIELD_LENGTHS |
+                UCP_DT_LOCAL_SGL_FIELD_MEMHS,
+            .buffers = localAddrs_,
+            .lengths = lengths_,
+            .memhs = memhs_,
+        };
+    }
+
+    [[nodiscard]] ucp_dt_remote_sgl_t
+    remoteView() const noexcept {
+        return {
+            .field_mask = UCP_DT_REMOTE_SGL_FIELD_REMOTE_ADDRS | UCP_DT_REMOTE_SGL_FIELD_LENGTHS |
+                UCP_DT_REMOTE_SGL_FIELD_RKEYS,
+            .remote_addrs = remoteAddrs_,
+            .lengths = lengths_,
+            .rkeys = rkeys_,
+        };
+    }
+
+    static constexpr size_t num_fields = 5;
+
+    std::unique_ptr<std::byte[]> storage_;
+    void **localAddrs_;
+    uint64_t *remoteAddrs_;
+    size_t *lengths_;
+    ucp_mem_h *memhs_;
+    ucp_rkey_h *rkeys_;
+    size_t size_ = 0;
+    ucx_connection_ptr_t conn_;
+};
+
+} // namespace nixl::ucx
+
+#endif // HAVE_UCX_SGL_API
+
+#endif // NIXL_SRC_PLUGINS_UCX_UCX_SGL_H

src/plugins/ucx/ucx_utils.cpp

@@ -362,6 +362,38 @@ nixlUcxEp::flushEp(nixlUcxReq &req) {
     return nixl::ucx::ucsToNixlStatus(UCS_PTR_STATUS(request));
 }
 
+#ifdef HAVE_UCX_SGL_API
+nixl_status_t
+nixlUcxEp::postSgl(const ucp_dt_local_sgl_t &local,
+                   const ucp_dt_remote_sgl_t &remote,
+                   size_t count,
+                   nixlUcxReq &req) {
+    const nixl_status_t status = checkTxState();
+    if (status != NIXL_SUCCESS) {
+        return status;
+    }
+
+    const ucp_request_param_t param = {
+        .op_attr_mask = UCP_OP_ATTR_FIELD_DATATYPE | UCP_OP_ATTR_FIELD_REMOTE_DATATYPE |
+            UCP_OP_ATTR_FIELD_REMOTE | UCP_OP_ATTR_FIELD_REMOTE_COUNT,
+        .datatype = ucp_dt_make_sgl(),
+        .remote_datatype = ucp_dt_make_sgl(),
+        .remote = &remote,
+        .remote_count = count,
+    };
+
+    const ucs_status_ptr_t request =
+        ucp_put_nbx(eph, &local, count, UCP_REMOTE_ADDR_INVALID, UCP_RKEY_INVALID, &param);
+    if (UCS_PTR_IS_PTR(request)) {
+        req = static_cast<nixlUcxReq>(request);
+        return NIXL_IN_PROG;
+    }
+
+    req = nullptr;
+    return nixl::ucx::ucsToNixlStatus(UCS_PTR_STATUS(request));
+}
+#endif
+
 bool
 nixlUcxMtLevelIsSupported(const nixl::ucx::mt_mode_t mt_type) noexcept {
     ucp_lib_attr_t attr;