Cpu proxy ucx backend#1931
Draft
tomerdav wants to merge 18 commits into
Draft
Conversation
Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
|
👋 Hi tomerdav! Thank you for contributing to ai-dynamo/nixl. Your PR reviewers will review your contribution then trigger the CI to test your changes. 🚀 |
Factor the existing UCX-only GPU device headers into a backend-agnostic common layer (nixl_device_types, nixl_device_api, nixl_device_wrappers) and move the UCX implementation detail into ucx/nixl_device_impl.cuh. Add a proxy backend variant under src/api/gpu/proxy/ that compiles into the nixl_gpu_proxy_lib static library. The proxy headers depend on proxy_protocol.h (the shared GPU/CPU ring-buffer protocol), which is introduced here alongside the GPU layer it serves. The top-level src/api/gpu/nixl_device.cuh dispatches to the selected backend implementation via the NIXL_GPU_DEVICE_BACKEND_PROXY macro. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Introduce the CPU-side proxy machinery under src/core/device_proxy/: - proxy_runtime.h/cpp: ProxyRuntime orchestrates channel state and the ProxyMemViewRegistry that maps GPU-visible proxy handles to backend memory-view descriptors. - proxy_worker.h/cpp: ProxyWorker is a CPU thread that drains GPU submission rings and dispatches transfers through a backend adapter. - backend_adapter.h: pure-virtual interface for proxy transport backends. - backend_provider.h: factory interface for backend engines that can supply a proxy adapter. Both new .cpp files are compiled into nixl_lib (src/core/meson.build). Unit tests for ProxyMemViewRegistry and ProxyRuntime are added under test/gtest/unit/. They depend only on internal core headers (no public nixl.h API) and run against a stub backend. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Make nixlUcxEngine inherit from DeviceProxyBackendProvider and implement createDeviceProxyBackendAdapter(), which returns a nixlUcxProxyBackend. nixlUcxProxyBackend (src/plugins/ucx/device_proxy/) implements the DeviceProxyBackendAdapter interface: it translates PreparedProxyTransferDesc objects into UCX postXfer calls, using the engine's existing UCX endpoint and memory-registration infrastructure. The UCX plugin meson.build is updated to compile ucx_proxy_backend.cpp into the plugin DSO. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Add three fields to nixlAgentConfig: enableDeviceProxy, proxyWorkerCount, and proxyChannelCount, all defaulting to disabled/1. Add nixlAgent::getProxyDeviceContext() to the public API, returning the ProxyDeviceContextData pointer that GPU kernels need to enqueue transfers. Wire proxy lifecycle into nixlAgentData: - createProxyRuntime(): called when a backend that implements DeviceProxyBackendProvider is registered and enableDeviceProxy is set. Initialises ProxyRuntime and starts worker threads. - shutdownProxyRuntime(): called from the agent destructor. nixl_listener.cpp: propagate remote connection info into the proxy runtime (proxyRuntime->loadRemoteConnInfo) so proxy workers can post to remote peers after metadata exchange. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Add proxy_write_test.cu: an end-to-end device-API test that stands up a ProxyRuntime with a stub backend, publishes the device context to the GPU via nixlProxyPublishContext(), enqueues PUT operations from a CUDA kernel, and polls for completion. Extend single_write_test.cu with conditional proxy-backend support: when compiled with -DNIXL_GPU_DEVICE_BACKEND_PROXY the test uses the proxy code path (nixlProxyPublishContext / nixlProxyClearContext) instead of the UCX device API path. Introduce the gtest_proxy executable in test/gtest/meson.build. It links nixl_gpu_proxy_lib directly (so Meson includes nixl_device_proxy.cu.o in the RDC device-link step) and is compiled with -DNIXL_GPU_DEVICE_BACKEND_PROXY. Gate it on cuda_dep.found() rather than the UCX GPU flag used by gtest. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Change implicit behavior of UCX device api selection to explicit by adding a new option `gpu_device_api_backend`: auto -> choose default backend if available ucx -> require UCX GPU Device API proxy -> require CUDA/proxy support none -> disable default GPU Device API backend Also centralize backend-specific compile args, deps, flags and link inputs so tests and examples will use the same resolved GPU Device API selected backend. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Route proxy PUT through a lower-level UCX RMA write helper while preserving the existing request handle and flush-backed completion semantics. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
And make sure it is reflected in proxy memview registry as it only requires remote memview Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Publish proxy submissions by writing the payload first and committing op_idx last, then have the CPU worker acquire-poll op_idx instead of ready_flag. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Narrow protocol-only submission fields and assert the 64-byte record layout while keeping op_idx at the first word for readiness publication. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
…emory Publish the proxy context through constant memory and force-inline device helper paths used by proxy API calls. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Move the GPU-owned producer index to device memory, derive operation ids from it, and keep mapped host aliases local to allocation. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Cache the host-pinned consumer index in device memory and refresh from host only when the ring appears full. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Keep the ready-word publish at GPU scope for host-pinned proxy records. The enqueue path writes the payload first and commits op_idx last; the CPU worker acquire-polls op_idx before copying the record. This avoids the enqueue overhead of a system-scope release while documenting the D2H FIFO contract explicitly. Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
tomerdav
force-pushed
the
cpu-proxy-ucx-backend
branch
from
July 13, 2026 15:08
dcb1e30 to
95a5503
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What?
Add a CPU-proxy backend for the NIXL GPU Device API, allowing GPU kernels to initiate transfers (PUT, atomic add) on transports without native device-initiated support by delegating submission to CPU worker threads.
src/api/gpu/into a backend-agnostic common layer withucx/andproxy/implementations, selected via theNIXL_GPU_DEVICE_BACKEND_PROXYmacro and a newgpu_device_api_backendmeson option (auto/ucx/proxy/none).src/core/device_proxy/):ProxyRuntimemanages per-channel GPU-to-CPU submission rings and a memview registry;ProxyWorkerthreads drain the rings and dispatch through aDeviceProxyBackendAdapterinterface.nixlUcxEngineimplements the backend-provider interface;nixlUcxProxyBackendtranslates proxy submissions into direct UCX RMA puts and atomic adds.nixlAgentConfigfields (enableDeviceProxy,proxyWorkerCount,proxyChannelCount) andnixlAgent::getProxyDeviceContext()for publishing the device context to kernels; remote connection info is propagated to the proxy runtime after metadata exchange.proxy_write_test.cuand proxy-mode coverage insingle_write_test.cu.Why?
The GPU Device API currently requires UCX device-side support (GDAKI-capable transports). A CPU proxy makes device-initiated transfers available on any backend that can serve as a proxy adapter, with UCX as the first provider.
How?
op_idxlast with a GPU-scope release; the CPU worker acquire-pollsop_idxas the readiness word, dispatches through the backend adapter, and publishes completions back through a pinned completion slot that the GPU polls.