[jambow0320]

Motivation

Step 2 of #34510 for NIXL: move the NIXL engine calls behind a NixlTransport on the interface proposed in #46455.

Stacked on #46460 (MooncakeTransport);

Merge order: #46460 should be merged first. This PR is then rebased onto main.

Modifications

• nixl/transport.py: NixlTransport owns the NIXL agent and makes every engine call: memory registration, peer import and reload, writes, prepared descriptor lists and transfer status. It also keeps every prepared descriptor list; nixl/conn.py keeps only the geometry it computes indices from, and no longer calls the agent directly or holds a list.
• Refactor only: the wire format and transfer behavior are unchanged. A rejected post still raises _NixlPostError, and the failure messages are kept. Two failure paths of the peer reload differ; see Notes.
• Aligned with the Mooncake side so that the protocol code can be shared later: nixl/buffers.py matches mooncake/buffers.py apart from the key class name, and nixl/conn.py follows MooncakeKVManager (init_engine(), BufferKeys built once per side, PeerBuffers, staging keys by base address, _transfer_copies(), src_keys/dst_keys on the send paths).
• transport/base.py and the Mooncake files are untouched. NIXL-only types stay under nixl/; shared types are extracted once every backend has a Transport.

What NixlTransport adds over MooncakeTransport

Registration, the per-region registration bookkeeping, write(), write_all(), status() and close() have the same shape as in MooncakeTransport. The rest follows from how the NIXL engine differs from Mooncake's. Apart from the interface methods, each item is engine handling that nixl/conn.py already had before this PR and that moved into the Transport.

Added or different in NixlTransport	Why NIXL needs it
NixlTransport.create()	Mooncake borrows the process-wide engine. NIXL creates one agent per KVManager, with a thread-safe sync mode, backend thread parameters, and the GPU selected in the thread that creates the backend. create() makes the agent and returns the Transport that holds it; the KVManager keeps no reference. Moved from NixlKVManager.__init__; init_engine() calls it.
agent_name, agent_metadata	Carried by the registration message, as session_id is for Mooncake.
import_registration(agent_name, metadata, bases, devices)	NIXL must load the peer's metadata (add_remote_agent) before any transfer, and peer descriptors need each buffer's memory kind and device id. Mooncake only needs the session id and base addresses.
is_loaded(), reload(), unload()	NIXL drops a peer's metadata after a remote disconnect; the peer reload from #41691 goes through these. A list prepared over the peer is invalid once the peer is loaded again, so reload() rebuilds those lists from their kept runs and keeps them out of reach until rebuilt. Mooncake uses a session blacklist and send_probe() instead.
Per-memory-kind register_all(), deregister() and write_all()	A NIXL registration or descriptor list holds a single memory kind (VRAM or DRAM), and each descriptor carries a device id. Registrations are therefore grouped by kind, and one write_all() becomes one NIXL transfer per (source kind, destination kind). Mooncake registers and transfers raw addresses in a single batch.
prepare(), write_prepared(), StridedRun, supports_strided_runs	NIXL prepared descriptor lists: a list is prepared once, and later transfers select its blocks by index (#37827 strided runs, #44740 hybrid-model state). The Transport keeps each list with the runs it was built from, by side (this side or a peer) and layout name, and write_prepared() selects both lists by name; nixl/conn.py still builds the runs and computes the indices. NIXL without strided-run support needs one descriptor per block, so prepare() keeps that fallback; state lists, which can hold millions of slots, are never expanded (allow_expansion=False).
status() polls and caches the final result	NIXL posts asynchronously and the protocol polls check_xfer_state. Mooncake's write_all() blocks, so its handle already holds the final result.
NixlResult.quiescent, NixlHandle.error	nixl/conn.py handles a transfer that NIXL never created or reported as ERR differently from a rejected post or an unreadable state, and builds its failure messages from the NIXL exception. These two fields carry that information across the Transport boundary. Mooncake cannot make this distinction: a timed-out batch_transfer_sync may still be accessing memory.
close() also drops the prepared lists and unloads peers	The Transport owns the agent instead of borrowing it.

Notes

• Registration skips zero-length regions of every kind and registers a region shared by several keys once, as MooncakeTransport does and deregister() in transport/base.py requires. nixl/conn.py still leaves zero-pointer state tensors unregistered, as before.
• On NIXL 1.3.0 the bindings raise for error statuses, so check_xfer_state() does not return "ERR", and a failed transfer takes the unreadable-state path, as it does on main.
• Peer reload: nixl/conn.py used to drop a reloaded peer's lists and rerun the registration-time preparation; reload() now rebuilds the lists prepared over the peer from their kept runs, with the same descriptors. This side's lists stay valid across a peer reload and are kept. Two failure paths differ from main:
  • A state list that could not be prepared at registration is not retried on a reload; that peer's state stays on the generic path.
  • A state list that fails to rebuild leaves the whole peer dropped until a later reload succeeds; main kept the peer and moved its state to the generic path.
• deregister() and close() implement the abstract interface; nixl/conn.py has no teardown and does not call them yet.

Tests

test/registered/unit/disaggregation/ passes on CPU (407 passed, 17 skipped). Not yet run with a real NIXL backend.

---

CI States

Latest PR Test (Base): ❌ Run #41037302856
Latest PR Test (Extra): ❌ Run #41037302734
Latest PR Test (AMD ROCm 10): ❌ Run #41037302824

[livingshade]

#46747

CC

[jambow0320]

@livingshade Thanks for the heads-up. #46747 calls the NIXL agent directly in _add_remote_peer, and this PR moves those calls behind NixlTransport. On top of this PR, the warm-up becomes transport.write_prepared() followed by polling transport.status(), with no interface change. NIXL 1.3 releases a transfer handle when the handle object is collected, so no separate release call is needed. If #46747 lands first, I'll rebase this PR and route its warm-up through the Transport.