[WIP][feat] optimize NIXL-UCX payload transfer for SimpleStorage

[Ginray]

Motivation

PR #163 introduced the PayloadTransfer abstraction and the optional NIXL-UCX Host-to-Host payload backend for SimpleStorage.

This PR further optimizes and hardens the NIXL-UCX path for RL workloads. The main goals are to reduce avoidable Host-side copies, support large frame-native payloads, avoid blocking SimpleStorage workers during GET transfers, and make buffer / transfer lifecycle handling safer under concurrency and failure.

This work follows the optimization plan outlined in the Payload Transfer roadmap.

Summary

• Replace the packed NIXL payload format with frame-native scatter WRITE, so encoded frames can be transferred directly without rebuilding a payload-sized packed buffer.
• Reuse registered receive buffers and keep them alive with a lease while decoded tensors or views still reference the underlying memory.
• Reuse existing receive MRs as send sources when possible; writable contiguous Host frames can otherwise be registered directly without an extra staging copy.
• Add deferred GET responses, allowing the SimpleStorage worker to continue processing requests while the NIXL transfer completes.
• Use a single-worker executor per remote peer: transfers to the same peer remain serialized, while different peers can progress concurrently.
• Harden failure and shutdown handling with receive-buffer quarantine, retained send resources, failed-peer isolation, and ordered runtime teardown.

Design

ZMQ remains the control plane, while non-empty Host payloads use NIXL-UCX.

SimpleStorage
     |
     | prepare / commit / response
     v
    ZMQ
     |
     +------------------------------+
                                    |
encoded frames                      |
     |                              |
     v                              v
NIXL-UCX scatter WRITE      registered receive buffer
                                    |
                                    v
                           zero-copy frame views
                                    |
                                    v
                                  decode

The implementation intentionally keeps the scope limited to the SimpleStorage Host payload fast path rather than introducing a general-purpose NIXL transport or MR management subsystem.

Compatibility

• Existing SimpleStorage public APIs remain unchanged.
• ZMQ remains the default payload backend.
• NIXL-UCX remains opt-in.
• Existing PUT / GET / CLEAR semantics are unchanged.
• NIXL transfer failures are reported explicitly and do not silently fall back to ZMQ.

Validation

Validation was performed on the current PR head with both two-node NIXL-UCX tests and local RL integration tests.

• Two-node tests verified the payload path over rma(rc_verbs/hns_0:1), including PUT / GET / CLEAR, datatype and nested-object handling, receive-buffer lease reuse, same-peer / different-peer concurrency, and controlled peer failure.
• PUT / GET correctness passed for payloads from 1 MiB through single-frame 5 GiB and multi-frame 8 GiB (3+3+2 GiB). No 32-bit packed-length truncation or StorageUnit splitting was observed.
• TQ-side source / staging instrumentation did not observe payload-sized copies; only small metadata-frame staging was recorded.
• Local VIME integration verified the TransferQueue rollout / fetch / clear / training data path.
• Focused regression tests passed: 32 passed.
• Full regression result: 624 passed, 10 skipped.

I will supplement more experimental results and data later.

Follow-up

1. I still need to further evaluate and validate whether the current design is appropriate and whether any parts are over-engineered.

2. Validate the optimized NIXL-UCX path in real RL framework workloads to confirm that the expected performance improvements are realized and that the current design remains appropriate for practical RL usage.

3. Consider adding a resource limit to the idle registered receive pool. This may depend on the results of integration with RL frameworks, and I have not yet decided on the right approach.

4. Connect the runtime diagnostics to the existing TransferQueue metrics / Prometheus path.

[ascend-robot]

CLA Signature Pass

Ginray, thanks for your pull request. All authors of the commits have signed the CLA. 👍

[ascend-robot]

CLA Signature Pass

Ginray, thanks for your pull request. All authors of the commits have signed the CLA. 👍