Cpu proxy ucx backend

examples/device/ep/meson.build

@@ -18,6 +18,10 @@ if not cuda_dep.found()
     subdir_done()
 endif
 
+if gpu_device_api_backend_resolved != 'ucx'
+    error('nixl_ep requires gpu_device_api_backend=ucx because it currently uses the UCX GPU Device API')
+endif
+
 py = import('python').find_installation('python3', pure: false)
 
 # Check if PyTorch is available
@@ -72,7 +76,7 @@ nixl_ep_sources = [
 
 nixl_ep_inc_dirs = [
     nixl_inc_dirs,
-    nixl_gpu_inc_dirs,
+    nixl_gpu_device_api_inc_dirs,
     include_directories('csrc'),
     include_directories('csrc/kernels'),
     torch_inc_dirs,
@@ -96,7 +100,7 @@ nixl_ep_cpp_args = [
     '-Wno-sign-compare',
     '-Wno-reorder',
     '-Wno-attributes',
-]
+] + gpu_device_api_cpp_args
 
 nixl_ep_cuda_args = [
     '-DHAVE_CUDA',
@@ -109,7 +113,7 @@ nixl_ep_cuda_args = [
     '-Xcompiler', '-Wno-sign-compare',
     '-Xcompiler', '-Wno-reorder',
     '-Xcompiler', '-Wno-attributes',
-]
+] + gpu_device_api_cuda_args
 
 topk_idx_bits = meson.get_external_property('topk_idx_bits', '32')
 nixl_ep_cpp_args += ['-DTOPK_IDX_BITS=' + topk_idx_bits]

meson.build

@@ -396,6 +396,8 @@ endif
 # UCX GPU device API detection
 nvcc_prog = find_program('nvcc', required: false)
 ucx_gpu_device_api_available = false
+have_gpu_side = false
+have_host_side = false
 if ucx_dep.found() and cuda_dep.found() and nvcc_prog.found()
     cuda = meson.get_compiler('cuda')
     # TODO: Expose doca_gpunetio_dep through UCX
@@ -412,16 +414,67 @@ if ucx_dep.found() and cuda_dep.found() and nvcc_prog.found()
         ucx_gpu_device_api_available = true
         add_project_arguments('-DHAVE_UCX_GPU_DEVICE_API', language: ['cpp', 'cuda'])
     endif
+endif
 
-    summary({
-    'UCX GPU Device API'    : ucx_gpu_device_api_available,
-    'GPU-side compile'      : have_gpu_side,
-    'Host-side compile'     : have_host_side,
-    'nvcc available'        : nvcc_prog.found(),
-    'DOCA GPUNETIO found'   : doca_gpunetio_dep.found(),
-  }, section: 'UCX GPU Device API', bool_yn: true)
+gpu_device_api_backend_requested = get_option('gpu_device_api_backend')
+gpu_device_api_backend_resolved = gpu_device_api_backend_requested
+
+if gpu_device_api_backend_requested == 'auto'
+    if ucx_gpu_device_api_available
+        gpu_device_api_backend_resolved = 'ucx'
+    else
+        gpu_device_api_backend_resolved = 'none'
+    endif
+elif gpu_device_api_backend_requested == 'ucx'
+    if not ucx_gpu_device_api_available
+        error('gpu_device_api_backend=ucx requires UCX GPU Device API support')
+    endif
+elif gpu_device_api_backend_requested == 'proxy'
+    if not cuda_dep.found()
+        error('gpu_device_api_backend=proxy requires CUDA')
+    endif
 endif
 
+gpu_device_api_define = ''
+gpu_device_api_supported = gpu_device_api_backend_resolved != 'none'
+gpu_device_api_is_ucx = gpu_device_api_backend_resolved == 'ucx'
+gpu_device_api_is_proxy = gpu_device_api_backend_resolved == 'proxy'
+gpu_device_api_cpp_args = []
+gpu_device_api_cuda_args = []
+gpu_device_api_deps = []
+gpu_device_api_link_with = []
+
+if gpu_device_api_is_ucx
+    gpu_device_api_define = '-DNIXL_GPU_DEVICE_BACKEND_UCX'
+    gpu_device_api_deps = [ucx_dep, doca_gpunetio_dep]
+elif gpu_device_api_is_proxy
+    gpu_device_api_define = '-DNIXL_GPU_DEVICE_BACKEND_PROXY'
+    gpu_device_api_deps = [cuda_dep]
+endif
+
+if gpu_device_api_define != ''
+    gpu_device_api_cpp_args += [gpu_device_api_define]
+    gpu_device_api_cuda_args += [gpu_device_api_define]
+endif
+
+summary({
+    'Requested backend' : gpu_device_api_backend_requested,
+    'Resolved backend'  : gpu_device_api_backend_resolved,
+}, section: 'GPU Device API')
+
+summary({
+    'UCX backend available'   : ucx_gpu_device_api_available,
+    'Proxy backend available' : cuda_dep.found(),
+}, section: 'GPU Device API', bool_yn: true)
+
+summary({
+    'UCX GPU Device API'  : ucx_gpu_device_api_available,
+    'GPU-side compile'    : have_gpu_side,
+    'Host-side compile'   : have_host_side,
+    'nvcc available'      : nvcc_prog.found(),
+    'DOCA GPUNETIO found' : doca_gpunetio_dep.found(),
+}, section: 'UCX GPU Device API', bool_yn: true)
+
 if get_option('disable_gds_backend')
     add_project_arguments('-DDISABLE_GDS_BACKEND', language: 'cpp')
 endif
@@ -509,7 +562,9 @@ if get_option('buildtype') == 'debug'
 endif
 
 nixl_inc_dirs = include_directories('src/api/cpp', 'src/api/cpp/backend', 'src/infra', 'src/core', 'src/core/telemetry')
+nixl_gpu_device_api_inc_dirs = include_directories('src/api/gpu')
 nixl_gpu_inc_dirs = include_directories('src/api/gpu/ucx')
+nixl_gpu_proxy_inc_dirs = include_directories('src/api/gpu/proxy')
 plugins_inc_dirs = include_directories('src/plugins')
 utils_inc_dirs = include_directories('src/utils')
 

meson_options.txt

@@ -28,6 +28,7 @@ option('cudapath_lib', type: 'string', value: '', description: 'Library path for
 option('cudapath_stub', type: 'string', value: '', description: 'Extra Stub path for CUDA')
 option('nixl_cuda_arch_list', type: 'string', value: 'auto', description: 'Comma-separated CUDA SM targets (e.g. 90,100), or auto to select defaults (sm_80+ normally, sm_90+ when build_nixl_ep is enabled). Use a single target like 100 for faster single-GPU builds.')
 option('wheel_variant', type: 'string', value: '', description: 'Override wheel variant suffix (e.g. -Dwheel_variant=rocm yields nixl_rocm). Empty = autodetect from CUDA major version.')
+option('gpu_device_api_backend', type: 'combo', choices: ['auto', 'ucx', 'proxy', 'none'], value: 'auto', description: 'GPU Device API backend to select for default device-side builds')
 option('static_plugins', type: 'string', value: '', description: 'Plugins to be built-in, comma-separated')
 option('enable_plugins', type: 'string', value: '', description: 'Comma-separated list of plugins to build. Default (empty) builds all available plugins.')
 option('disable_plugins', type: 'string', value: '', description: 'Comma-separated list of plugins to exclude from build. Cannot be used with enable_plugins.')

src/api/cpp/backend/backend_engine.h

@@ -27,6 +27,8 @@
 #include "backend_aux.h"
 #include "telemetry_event.h"
 
+class nixlDeviceProxyBackendAdapter;
+
 constexpr size_t MAX_TELEMETRY_QUEUE_SIZE = 1000;
 
 // Base backend engine class for different backend implementations
@@ -111,6 +113,12 @@ class nixlBackendEngine {
         // pure virtual, and return errors, as parent shouldn't call if supportsNotif is false.
         virtual bool supportsNotif() const = 0;
 
+        // Determines if a backend supports device proxy runtime creation.
+        virtual bool
+        supportsProxy() const {
+            return false;
+        }
+
         virtual nixl_mem_list_t getSupportedMems() const = 0;  // TODO: Return by const-reference and mark noexcept?
 
 
@@ -175,6 +183,12 @@ class nixlBackendEngine {
         virtual void
         releaseMemView(nixlMemViewH) const {}
 
+        virtual nixl_status_t
+        createDeviceProxyBackendAdapter(const nixlBackendInitParams &,
+                                        std::unique_ptr<nixlDeviceProxyBackendAdapter> &) {
+            return NIXL_ERR_NOT_SUPPORTED;
+        }
+
         // *** Needs to be implemented if supportsRemote() is true *** //
 
         // Gets serialized form of public metadata

src/api/cpp/nixl.h

@@ -561,6 +561,15 @@ class nixlAgent {
         checkRemoteMD (const std::string remote_name,
                        const nixl_xfer_dlist_t &descs) const;
 
+        /**
+         * @brief  Return the proxy runtime's device context data pointer,
+         *         or nullptr when proxy mode is not active.
+         *
+         * The caller can pass this to nixlProxyPublishContext() from a CUDA TU
+         * to make the context visible to GPU kernels.
+         */
+        void *
+        getProxyDeviceContext() const;
 };
 
 #endif

src/api/cpp/nixl_params.h

@@ -33,6 +33,9 @@ struct nixlAgentConfig {
     static constexpr nixl_thread_sync_t kDefaultSyncMode =
         nixl_thread_sync_t::NIXL_THREAD_SYNC_DEFAULT;
     static constexpr bool kDefaultCaptureTelemetry = false;
+    static constexpr bool kDefaultEnableDeviceProxy = false;
+    static constexpr uint32_t kDefaultProxyWorkerCount = 1;
+    static constexpr uint32_t kDefaultProxyChannelCount = 1;
     static constexpr uint64_t kDefaultPthrDelayUs = 0;
     static constexpr uint64_t kDefaultLthrDelayUs = 100000;
     static constexpr std::chrono::microseconds kDefaultEtcdWatchTimeout =
@@ -48,6 +51,21 @@ struct nixlAgentConfig {
     nixl_thread_sync_t syncMode = kDefaultSyncMode;
     /** @var Capture telemetry info regardless of environment variables*/
     bool captureTelemetry = kDefaultCaptureTelemetry;
+    /** @var Enable the device proxy orchestration skeleton. */
+    bool enableDeviceProxy = kDefaultEnableDeviceProxy;
+    /**
+     * @var Desired number of proxy workers per proxy runtime.
+     *      The UCX proxy backend currently shares a nixlUcxEngine across proxy
+     *      workers; keep this at 1 unless the backend has been validated for
+     *      concurrent postXfer calls.
+     */
+    uint32_t proxyWorkerCount = kDefaultProxyWorkerCount;
+    /**
+     * @var Desired number of proxy channels per proxy runtime.
+     *      Channels may exceed workers; this is the supported way to expose
+     *      multiple GPU submission queues while using one proxy worker.
+     */
+    uint32_t proxyChannelCount = kDefaultProxyChannelCount;
 
     /**
      * @var Progress thread event waiting timeout.
@@ -94,12 +112,18 @@ struct nixlAgentConfig {
         uint64_t pthr_delay_us = kDefaultPthrDelayUs,
         uint64_t lthr_delay_us = kDefaultLthrDelayUs,
         bool capture_telemetry = kDefaultCaptureTelemetry,
+        bool enable_device_proxy = kDefaultEnableDeviceProxy,
+        uint32_t proxy_worker_count = kDefaultProxyWorkerCount,
+        uint32_t proxy_channel_count = kDefaultProxyChannelCount,
         std::chrono::microseconds etcd_watch_timeout = kDefaultEtcdWatchTimeout) noexcept
         : useProgThread(use_prog_thread),
           useListenThread(use_listen_thread),
           listenPort(port),
           syncMode(sync_mode),
           captureTelemetry(capture_telemetry),
+          enableDeviceProxy(enable_device_proxy),
+          proxyWorkerCount(proxy_worker_count),
+          proxyChannelCount(proxy_channel_count),
           pthrDelay(pthr_delay_us),
           lthrDelay(lthr_delay_us),
           etcdWatchTimeout(etcd_watch_timeout) {}

src/api/gpu/common/nixl_device_api.cuh

@@ -0,0 +1,74 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 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_API_GPU_COMMON_NIXL_DEVICE_API_CUH
+#define NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_API_CUH
+
+#include "nixl_device_types.cuh"
+
+#if defined(NIXL_GPU_DEVICE_BACKEND_PROXY)
+#include "../proxy/nixl_device_impl.cuh"
+
+namespace nixl::gpu {
+namespace selected_impl = proxy_impl;
+}
+#elif defined(NIXL_GPU_DEVICE_BACKEND_UCX)
+#include "../ucx/nixl_device_impl.cuh"
+
+namespace nixl::gpu {
+namespace selected_impl = ucx_impl;
+}
+#else
+#error "No GPU device backend implementation selected"
+#endif
+
+namespace nixl::gpu::api {
+
+template<nixl_gpu_level_t level = nixl_gpu_level_t::THREAD>
+__device__ inline nixl_status_t
+get_xfer_status(nixlGpuXferStatusH &xfer_status) {
+    return selected_impl::get_xfer_status<level>(xfer_status);
+}
+
+template<nixl_gpu_level_t level = nixl_gpu_level_t::THREAD>
+__device__ inline nixl_status_t
+put(const nixlMemViewElem &src,
+    const nixlMemViewElem &dst,
+    size_t size,
+    unsigned channel_id = 0,
+    uint64_t flags = 0,
+    nixlGpuXferStatusH *xfer_status = nullptr) {
+    return selected_impl::put<level>(src, dst, size, channel_id, flags, xfer_status);
+}
+
+template<nixl_gpu_level_t level = nixl_gpu_level_t::THREAD>
+__device__ inline nixl_status_t
+atomic_add(uint64_t value,
+           const nixlMemViewElem &counter,
+           unsigned channel_id = 0,
+           uint64_t flags = 0,
+           nixlGpuXferStatusH *xfer_status = nullptr) {
+    return selected_impl::atomic_add<level>(value, counter, channel_id, flags, xfer_status);
+}
+
+__device__ inline void *
+get_ptr(nixlMemViewH mvh, size_t index) {
+    return selected_impl::get_ptr(mvh, index);
+}
+
+} // namespace nixl::gpu::api
+
+#endif // NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_API_CUH

src/api/gpu/common/nixl_device_types.cuh

@@ -0,0 +1,41 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 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_API_GPU_COMMON_NIXL_DEVICE_TYPES_CUH
+#define NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_TYPES_CUH
+
+#include <cstddef>
+#include <cstdint>
+
+#include <nixl_types.h>
+
+struct nixlGpuXferStatusH {
+    alignas(16) unsigned char storage[64] = {};
+};
+
+enum class nixl_gpu_level_t : uint64_t { THREAD = 0, WARP = 1, BLOCK = 2, GRID = 3 };
+
+namespace nixl_gpu_flags {
+constexpr uint64_t defer = 1;
+} // namespace nixl_gpu_flags
+
+struct nixlMemViewElem {
+    nixlMemViewH mvh;
+    size_t index; /**< Index in the memory view */
+    size_t offset; /**< Offset within the buffer */
+};
+
+#endif // NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_TYPES_CUH