vllm.v1.kv_offload.tiering.p2p.data.nixl

NixlTransport: Data-plane transport for RDMA-based KV block transfers via NIXL.

Threading: the agent is reached from the scheduler thread and, when the tiering manager runs one, from its control-plane thread. The manager's lock serializes them strictly -- one Python thread inside any agent call at a time, with a full memory barrier between.

That lock is the only protection, not a belt over NIXL's own. nixl_agent picks NIXL_THREAD_SYNC_STRICT only when enable_listen is set, which this transport does not set, so the agent runs with NIXL_THREAD_SYNC_NONE -- under which NIXL's internal agent lock compiles down to no-ops, leaving every agent entry point unguarded. So the manager lock must cover all of them, and the lifetime of every handle passed between calls. The exception is peer registration moved onto a worker thread, as NixlConnector's handshake executor does: it runs outside the lock and relies, like that connector, on NIXL tolerating it beside transfers.

Serialized use from a second thread is supported (checked against NIXL 1.4.1 sources): NIXL never creates its UCX worker in UCS_THREAD_MODE_SINGLE, the only creator-thread-affine mode. The weakest mode it ever requests is UCS_THREAD_MODE_SERIALIZED, documented as "multiple threads can access, but only one at a time", and because enable_prog_thread defaults on here it actually requests UCS_THREAD_MODE_MULTI and fails construction if the UCX build cannot provide it.

Two caveats. NIXL binds a shared UCX worker per calling thread (tlsSharedWorkerMap), which is harmless only while there is exactly one shared worker -- true here, but it would stop being true if a num_workers UCX backend param were ever set. And NIXL 1.4.1 does accept an explicit sync_mode (including NIXL_THREAD_SYNC_RW), which would make its own locking real; that parameter does not exist before 1.x, so passing it would need a version gate, the way capture_telemetry does.

Classes:

• NixlTransport –

  Manages a NIXL agent, memory registration, and block transfers.

NixlTransport

Bases: DataTransport

Manages a NIXL agent, memory registration, and block transfers.

Wraps the NIXL C library behind a Python interface so the rest of the P2P tier code never touches NIXL types directly. Tracks inflight handles internally and returns completed/failed tags on poll.

Methods:

• cancel –

  Cancel inflight transfers by their IDs.

• poll –

  Poll inflight transfers.

• write_blocks –

  Submit a WRITE transfer to peer_id.

Source code in vllm/v1/kv_offload/tiering/p2p/data/nixl.py

class NixlTransport(DataTransport):
    """Manages a NIXL agent, memory registration, and block transfers.

    Wraps the NIXL C library behind a Python interface so the rest of the
    P2P tier code never touches NIXL types directly. Tracks inflight
    handles internally and returns completed/failed tags on poll.
    """

    def __init__(
        self,
        agent_name: str,
        view: memoryview,
        config_fields: dict | None = None,
        backends: list[str] | None = None,
        num_threads: int = 4,
    ) -> None:
        super().__init__(view, config_fields=config_fields)
        self._agent_name = agent_name
        self._backends = list(backends) if backends else ["UCX"]
        self._num_threads = num_threads
        self._agent: Any = None
        self._reg: Any = None
        self._local_dlist: Any = None
        self._remote_dlists: dict[str, object] = {}
        self._peer_nixl_names: dict[str, str] = {}
        # transfer_id → _Inflight(peer_id, handle).
        self._inflight: dict[int, _Inflight] = {}
        self._next_id = itertools.count()

        self._init(view)

    @property
    def available(self) -> bool:
        return self._agent is not None

    def _init(self, view: memoryview) -> None:
        if _NixlAgent is None:
            return

        non_ucx_backends = [b for b in self._backends if b != "UCX"]
        if non_ucx_backends:
            cfg = _NixlAgentConfig(backends=self._backends, capture_telemetry=True)
            logger.info(
                "NixlTransport %s: NIXL backends=%s",
                self._agent_name,
                self._backends,
            )
        else:
            cfg = _NixlAgentConfig(
                num_threads=self._num_threads, capture_telemetry=True
            )
            logger.info(
                "NixlTransport %s: NIXL backends=[UCX] num_threads=%d",
                self._agent_name,
                self._num_threads,
            )
        self._agent = _NixlAgent(self._agent_name, cfg)

        total_size = self._num_blocks * self._block_len
        reg_descs = [(self._base_addr, total_size, 0, "")]
        self._reg = self._agent.register_memory(reg_descs, mem_type="DRAM")

        block_tuples = [
            (self._base_addr + i * self._block_len, self._block_len, 0)
            for i in range(self._num_blocks)
        ]
        xfer_dlist = self._agent.get_xfer_descs(block_tuples, mem_type="DRAM")
        self._local_dlist = self._agent.prep_xfer_dlist("NIXL_INIT_AGENT", xfer_dlist)
        logger.info(
            "NixlTransport %s: registered %d blocks", self._agent_name, self._num_blocks
        )

    def get_agent_metadata(self) -> bytes:
        assert self._agent is not None
        return self._agent.get_agent_metadata()

    # ------------------------------------------------------------------
    # Peer management
    # ------------------------------------------------------------------

    def add_remote_peer(
        self,
        peer_id: str,
        agent_metadata: bytes,
        base_addr: int,
        num_blocks: int,
        block_len: int,
    ) -> None:
        nixl_name = self._agent.add_remote_agent(agent_metadata)
        block_descs = [
            (base_addr + i * block_len, block_len, 0) for i in range(num_blocks)
        ]
        xfer_dlist = self._agent.get_xfer_descs(block_descs, mem_type="DRAM")
        remote_dlist = self._agent.prep_xfer_dlist(nixl_name, xfer_dlist)
        self._peer_nixl_names[peer_id] = nixl_name
        self._remote_dlists[peer_id] = remote_dlist

    def remove_remote_peer(self, peer_id: str) -> None:
        nixl_name = self._peer_nixl_names.pop(peer_id, None)
        dlist = self._remote_dlists.pop(peer_id, None)
        if self._agent is not None:
            if dlist is not None:
                self._agent.release_dlist_handle(dlist)
            if nixl_name:
                self._agent.remove_remote_agent(nixl_name)

    # ------------------------------------------------------------------
    # Transfer submission and polling
    # ------------------------------------------------------------------

    def write_blocks(
        self,
        peer_id: str,
        local_idxs: list[int],
        remote_idxs: list[int],
    ) -> int | None:
        """Submit a WRITE transfer to *peer_id*.

        Returns a transfer ID, or None if the peer is not registered.
        The ID is returned via poll() when the transfer completes or fails.
        """
        remote_dlist = self._remote_dlists.get(peer_id)
        if remote_dlist is None:
            logger.warning(
                "NixlTransport %s: write_blocks NO REMOTE DLIST for peer=%s "
                "(known peers=%s)",
                self._agent_name,
                peer_id,
                list(self._remote_dlists.keys()),
            )
            return None
        logger.debug(
            "NixlTransport %s: write_blocks NIXL.transfer peer=%s blocks=%d",
            self._agent_name,
            peer_id,
            len(local_idxs),
        )
        handle = self._agent.make_prepped_xfer(
            "WRITE",
            self._local_dlist,
            np.asarray(local_idxs, dtype=np.int32),
            remote_dlist,
            np.asarray(remote_idxs, dtype=np.int32),
        )
        self._agent.transfer(handle)
        transfer_id = next(self._next_id)
        self._inflight[transfer_id] = _Inflight(peer_id, handle)
        return transfer_id

    def poll(self, peer_id: str | None = None) -> PollResult:
        """Poll inflight transfers.

        When *peer_id* is given, only transfers submitted for that peer_id are
        checked and drained — the transport is shared across peer sessions, so
        an unscoped poll by one session would consume and discard siblings'
        completions. *peer_id* None polls every peer (shutdown drain only).

        Returns PollResult(done=..., failed=...) with transfer IDs.
        Completed handles are released automatically.
        """
        if not self._inflight:
            return _EMPTY_POLL_RESULT

        done_ids: list[int] | None = None
        failed_ids: list[int] | None = None

        for transfer_id, entry in self._inflight.items():
            if peer_id is not None and entry.peer_id != peer_id:
                continue
            try:
                state = self._agent.check_xfer_state(entry.handle)
            except Exception as exc:
                logger.warning(
                    "NixlTransport %s: check_xfer_state failed for transfer_id=%d: %s",
                    self._agent_name,
                    transfer_id,
                    exc,
                )
                continue
            if state == "DONE":
                if done_ids is None:
                    done_ids = []
                done_ids.append(transfer_id)
            elif state not in ("PROC", "PEND"):
                if failed_ids is None:
                    failed_ids = []
                failed_ids.append(transfer_id)

        if done_ids is None and failed_ids is None:
            return _EMPTY_POLL_RESULT

        handles_to_release = []
        for tid in done_ids or ():
            handles_to_release.append(self._inflight.pop(tid).handle)
        for tid in failed_ids or ():
            handles_to_release.append(self._inflight.pop(tid).handle)
        self._release_handles(handles_to_release)

        return PollResult(
            done=done_ids if done_ids is not None else _EMPTY_POLL_RESULT.done,
            failed=failed_ids if failed_ids is not None else _EMPTY_POLL_RESULT.failed,
        )

    def cancel(
        self,
        transfer_ids: Iterable[int],
        mode: CancelMode = "immediate",
    ) -> list[int]:
        """Cancel inflight transfers by their IDs.

        See ``DataTransport.cancel`` for the contract. In "wait" mode,
        transfers whose ``release_xfer_handle`` raises (NIXL could not
        complete the abort because the backend is still draining) stay
        in ``self._inflight`` so a later ``poll()`` will observe them.
        """
        if mode == "immediate":
            handles = [
                self._inflight.pop(tid).handle
                for tid in transfer_ids
                if tid in self._inflight
            ]
            self._release_handles(handles)
            return []

        still_inflight: list[int] = []
        for tid in transfer_ids:
            entry = self._inflight.get(tid)
            if entry is None:
                continue
            try:
                self._agent.release_xfer_handle(entry.handle)
            except Exception as exc:
                logger.debug(
                    "NixlTransport %s: cancel pending for transfer_id=%d: %s",
                    self._agent_name,
                    tid,
                    exc,
                )
                still_inflight.append(tid)
                continue
            del self._inflight[tid]
        return still_inflight

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def close(self) -> None:
        if self._agent is None:
            return
        self._release_handles([entry.handle for entry in self._inflight.values()])
        self._inflight.clear()
        for peer_id in list(self._remote_dlists):
            self.remove_remote_peer(peer_id)
        if self._local_dlist is not None:
            self._agent.release_dlist_handle(self._local_dlist)
            self._local_dlist = None
        if self._reg is not None:
            self._agent.deregister_memory(self._reg)
            self._reg = None
        self._agent = None

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _release_handles(self, handles: list[object]) -> None:
        if self._agent is None:
            return
        for handle in handles:
            try:
                self._agent.release_xfer_handle(handle)
            except Exception as exc:
                logger.warning(
                    "NixlTransport %s: release_xfer_handle failed: %s",
                    self._agent_name,
                    exc,
                )

cancel(transfer_ids, mode='immediate')

Cancel inflight transfers by their IDs.

See DataTransport.cancel for the contract. In "wait" mode, transfers whose release_xfer_handle raises (NIXL could not complete the abort because the backend is still draining) stay in self._inflight so a later poll() will observe them.

Source code in vllm/v1/kv_offload/tiering/p2p/data/nixl.py

def cancel(
    self,
    transfer_ids: Iterable[int],
    mode: CancelMode = "immediate",
) -> list[int]:
    """Cancel inflight transfers by their IDs.

    See ``DataTransport.cancel`` for the contract. In "wait" mode,
    transfers whose ``release_xfer_handle`` raises (NIXL could not
    complete the abort because the backend is still draining) stay
    in ``self._inflight`` so a later ``poll()`` will observe them.
    """
    if mode == "immediate":
        handles = [
            self._inflight.pop(tid).handle
            for tid in transfer_ids
            if tid in self._inflight
        ]
        self._release_handles(handles)
        return []

    still_inflight: list[int] = []
    for tid in transfer_ids:
        entry = self._inflight.get(tid)
        if entry is None:
            continue
        try:
            self._agent.release_xfer_handle(entry.handle)
        except Exception as exc:
            logger.debug(
                "NixlTransport %s: cancel pending for transfer_id=%d: %s",
                self._agent_name,
                tid,
                exc,
            )
            still_inflight.append(tid)
            continue
        del self._inflight[tid]
    return still_inflight

poll(peer_id=None)

Poll inflight transfers.

When peer_id is given, only transfers submitted for that peer_id are checked and drained — the transport is shared across peer sessions, so an unscoped poll by one session would consume and discard siblings' completions. peer_id None polls every peer (shutdown drain only).

Returns PollResult(done=..., failed=...) with transfer IDs. Completed handles are released automatically.

Source code in vllm/v1/kv_offload/tiering/p2p/data/nixl.py

def poll(self, peer_id: str | None = None) -> PollResult:
    """Poll inflight transfers.

    When *peer_id* is given, only transfers submitted for that peer_id are
    checked and drained — the transport is shared across peer sessions, so
    an unscoped poll by one session would consume and discard siblings'
    completions. *peer_id* None polls every peer (shutdown drain only).

    Returns PollResult(done=..., failed=...) with transfer IDs.
    Completed handles are released automatically.
    """
    if not self._inflight:
        return _EMPTY_POLL_RESULT

    done_ids: list[int] | None = None
    failed_ids: list[int] | None = None

    for transfer_id, entry in self._inflight.items():
        if peer_id is not None and entry.peer_id != peer_id:
            continue
        try:
            state = self._agent.check_xfer_state(entry.handle)
        except Exception as exc:
            logger.warning(
                "NixlTransport %s: check_xfer_state failed for transfer_id=%d: %s",
                self._agent_name,
                transfer_id,
                exc,
            )
            continue
        if state == "DONE":
            if done_ids is None:
                done_ids = []
            done_ids.append(transfer_id)
        elif state not in ("PROC", "PEND"):
            if failed_ids is None:
                failed_ids = []
            failed_ids.append(transfer_id)

    if done_ids is None and failed_ids is None:
        return _EMPTY_POLL_RESULT

    handles_to_release = []
    for tid in done_ids or ():
        handles_to_release.append(self._inflight.pop(tid).handle)
    for tid in failed_ids or ():
        handles_to_release.append(self._inflight.pop(tid).handle)
    self._release_handles(handles_to_release)

    return PollResult(
        done=done_ids if done_ids is not None else _EMPTY_POLL_RESULT.done,
        failed=failed_ids if failed_ids is not None else _EMPTY_POLL_RESULT.failed,
    )

write_blocks(peer_id, local_idxs, remote_idxs)

Submit a WRITE transfer to peer_id.

Returns a transfer ID, or None if the peer is not registered. The ID is returned via poll() when the transfer completes or fails.

Source code in vllm/v1/kv_offload/tiering/p2p/data/nixl.py

def write_blocks(
    self,
    peer_id: str,
    local_idxs: list[int],
    remote_idxs: list[int],
) -> int | None:
    """Submit a WRITE transfer to *peer_id*.

    Returns a transfer ID, or None if the peer is not registered.
    The ID is returned via poll() when the transfer completes or fails.
    """
    remote_dlist = self._remote_dlists.get(peer_id)
    if remote_dlist is None:
        logger.warning(
            "NixlTransport %s: write_blocks NO REMOTE DLIST for peer=%s "
            "(known peers=%s)",
            self._agent_name,
            peer_id,
            list(self._remote_dlists.keys()),
        )
        return None
    logger.debug(
        "NixlTransport %s: write_blocks NIXL.transfer peer=%s blocks=%d",
        self._agent_name,
        peer_id,
        len(local_idxs),
    )
    handle = self._agent.make_prepped_xfer(
        "WRITE",
        self._local_dlist,
        np.asarray(local_idxs, dtype=np.int32),
        remote_dlist,
        np.asarray(remote_idxs, dtype=np.int32),
    )
    self._agent.transfer(handle)
    transfer_id = next(self._next_id)
    self._inflight[transfer_id] = _Inflight(peer_id, handle)
    return transfer_id

_Inflight

Bases: NamedTuple

A submitted-but-not-yet-drained transfer.

peer_id lets poll() scope to a single owning session, since the transport is shared across all peer sessions of the engine.

Source code in vllm/v1/kv_offload/tiering/p2p/data/nixl.py

class _Inflight(NamedTuple):
    """A submitted-but-not-yet-drained transfer.

    ``peer_id`` lets poll() scope to a single owning session, since the
    transport is shared across all peer sessions of the engine.
    """

    peer_id: str
    handle: object