vllm.v1.kv_offload.tiering.p2p.data

Modules:

• base –

  Abstract base class for the P2P data-plane transport.

• nixl –

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

Classes:

• DataTransport –

  Abstract data-plane transport for RDMA-style block transfers.

• NixlTransport –

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

• PollResult –

  Result of polling inflight transfers.

DataTransport

Bases: ABC

Abstract data-plane transport for RDMA-style block transfers.

Owns the local KV block memory region and manages transfers to/from registered remote peers.

Construction

view: A 2D memoryview (num_blocks × block_len bytes) over the local KV cache block storage. config_fields: Dict of model config values used to compute the compatibility fingerprint. None → empty fingerprint (compatible with any peer).

Methods:

• add_remote_peer –

  Register a remote peer for block transfers.

• add_remote_peer_async –

  Register a remote peer without blocking the caller.

• cancel –

  Cancel inflight transfers by their IDs.

• close –

  Release all resources (registrations, handles, memory).

• get_agent_metadata –

  Return opaque metadata needed by remote peers to connect.

• poll –

  Poll inflight transfers for completion.

• remove_remote_peer –

  Unregister a remote peer and release associated resources.

• write_blocks –

  Submit a WRITE transfer: local blocks → remote peer's blocks.

Attributes:

• base_addr (int) –

  Base address of the local block memory region.

• block_len (int) –

  Size of each block in bytes.

• config_fingerprint (str) –

  Content-hash of the model configuration (hex string).

• num_blocks (int) –

  Number of blocks in the local region.

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

class DataTransport(ABC):
    """Abstract data-plane transport for RDMA-style block transfers.

    Owns the local KV block memory region and manages transfers to/from
    registered remote peers.

    Construction:
        view: A 2D memoryview (num_blocks × block_len bytes) over the
              local KV cache block storage.
        config_fields: Dict of model config values used to compute the
                       compatibility fingerprint. None → empty fingerprint
                       (compatible with any peer).
    """

    def __init__(self, view: memoryview, config_fields: dict | None = None) -> None:
        assert view.shape is not None
        self._view = view
        self._base_addr = ctypes.addressof(ctypes.c_char.from_buffer(view))
        self._num_blocks = view.shape[0]
        self._block_len = view.shape[1]
        self._config_fingerprint = self._compute_fingerprint(config_fields)

    @property
    def base_addr(self) -> int:
        """Base address of the local block memory region."""
        return self._base_addr

    @property
    def num_blocks(self) -> int:
        """Number of blocks in the local region."""
        return self._num_blocks

    @property
    def block_len(self) -> int:
        """Size of each block in bytes."""
        return self._block_len

    @property
    def config_fingerprint(self) -> str:
        """Content-hash of the model configuration (hex string).

        Peers must have matching fingerprints to exchange blocks.
        Empty string means no fingerprint (always compatible).
        """
        return self._config_fingerprint

    @staticmethod
    def _compute_fingerprint(config_fields: dict | None) -> str:
        if not config_fields:
            return ""
        canonical = json.dumps(config_fields, sort_keys=True, separators=(",", ":"))
        return hashlib.sha256(canonical.encode()).hexdigest()[:16]

    @abstractmethod
    def get_agent_metadata(self) -> bytes:
        """Return opaque metadata needed by remote peers to connect.

        The returned bytes are sent during the control-plane handshake
        and passed to the remote peer's add_remote_peer().
        """
        ...

    @abstractmethod
    def add_remote_peer(
        self,
        peer_id: str,
        agent_metadata: bytes,
        base_addr: int,
        num_blocks: int,
        block_len: int,
    ) -> None:
        """Register a remote peer for block transfers.

        Must be called before write_blocks() to this peer.

        Args:
            peer_id: Unique identifier for the remote peer.
            agent_metadata: Opaque bytes from the peer's get_agent_metadata().
            base_addr: Base address of the peer's block memory region.
            num_blocks: Number of blocks in the peer's region.
            block_len: Size of each block (must match local block_len).

        """
        ...

    def add_remote_peer_async(
        self,
        peer_id: str,
        agent_metadata: bytes,
        base_addr: int,
        num_blocks: int,
        block_len: int,
    ) -> Future[None]:
        """Register a remote peer without blocking the caller.

        Registration is O(num_blocks) and can take seconds for a large
        secondary tier, so the session starts it here and sends
        ``ConnectAckMsg`` only once the returned future resolves. The
        acknowledgement is therefore the peer's guarantee that this side
        is ready to serve, and no ``write_blocks`` can arrive before it.

        The default implementation runs ``add_remote_peer`` inline and
        returns an already-resolved future; transports whose registration
        is cheap need not override it.

        Args:
            peer_id: Unique identifier for the remote peer.
            agent_metadata: Opaque bytes from the peer's get_agent_metadata().
            base_addr: Base address of the peer's block memory region.
            num_blocks: Number of blocks in the peer's region.
            block_len: Size of each block (must match local block_len).

        Returns:
            A future resolving to None on success. A failed future means
            the peer must be rejected rather than acknowledged.

        """
        future: Future[None] = Future()
        try:
            self.add_remote_peer(
                peer_id,
                agent_metadata=agent_metadata,
                base_addr=base_addr,
                num_blocks=num_blocks,
                block_len=block_len,
            )
        except Exception as exc:
            future.set_exception(exc)
        else:
            future.set_result(None)
        return future

    @abstractmethod
    def remove_remote_peer(self, peer_id: str) -> None:
        """Unregister a remote peer and release associated resources.

        Inflight transfers to this peer should be cancelled first.
        """
        ...

    @abstractmethod
    def write_blocks(
        self,
        peer_id: str,
        local_idxs: list[int],
        remote_idxs: list[int],
    ) -> int | None:
        """Submit a WRITE transfer: local blocks → remote peer's blocks.

        Args:
            peer_id: Target peer (must be registered via add_remote_peer).
            local_idxs: Indexes of local blocks to read from.
            remote_idxs: Indexes of remote blocks to write to.
                         Must be same length as local_idxs.

        Returns:
            A unique transfer_id (int) to track this transfer, or
            None if the peer is not registered or submission failed.

        """
        ...

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

        Args:
            peer_id: If given, only poll (and drain) transfers submitted for
                this peer_id — the value passed to ``write_blocks``. This is
                required when a single transport is shared across multiple
                peer sessions: ``poll()`` pops completed handles, so an
                unscoped poll by one session would consume and discard the
                completions of its siblings, starving them. ``None`` polls
                every peer's transfers (used only for the shutdown drain).

        Returns:
            PollResult with lists of completed and failed transfer_ids.
            Completed/failed transfers are removed from the inflight set.

        Must be called periodically to drive progress checking.

        """
        ...

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

        Best-effort: transfers that already completed are ignored.

        Args:
            transfer_ids: IDs to cancel. Unknown IDs are ignored.
            mode:
                "immediate" (default): pop and release each handle and
                    return []. Matches the legacy fire-and-forget
                    behavior — the caller does not wait for the
                    underlying transfer to drain.
                "wait": release only transfers that have stopped. A
                    transfer still PROC/PEND (or whose release fails)
                    stays in the inflight set and its id is included in
                    the returned list. The caller is expected to keep
                    calling poll() until every returned id surfaces in
                    done/failed.

        Returns:
            For mode="wait", the subset of *transfer_ids* still
            tracked as inflight after the cancel attempt. For
            mode="immediate", always [].

        """
        ...

    @abstractmethod
    def close(self) -> None:
        """Release all resources (registrations, handles, memory).

        Cancels any remaining inflight transfers. Idempotent.
        After close(), no other methods may be called.
        """
        ...

base_addr property

Base address of the local block memory region.

block_len property

Size of each block in bytes.

config_fingerprint property

Content-hash of the model configuration (hex string).

Peers must have matching fingerprints to exchange blocks. Empty string means no fingerprint (always compatible).

num_blocks property

Number of blocks in the local region.

add_remote_peer(peer_id, agent_metadata, base_addr, num_blocks, block_len) abstractmethod

Register a remote peer for block transfers.

Must be called before write_blocks() to this peer.

Parameters:

• peer_id

  (str) –

  Unique identifier for the remote peer.

• agent_metadata

  (bytes) –

  Opaque bytes from the peer's get_agent_metadata().

• base_addr

  (int) –

  Base address of the peer's block memory region.

• num_blocks

  (int) –

  Number of blocks in the peer's region.

• block_len

  (int) –

  Size of each block (must match local block_len).

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

@abstractmethod
def add_remote_peer(
    self,
    peer_id: str,
    agent_metadata: bytes,
    base_addr: int,
    num_blocks: int,
    block_len: int,
) -> None:
    """Register a remote peer for block transfers.

    Must be called before write_blocks() to this peer.

    Args:
        peer_id: Unique identifier for the remote peer.
        agent_metadata: Opaque bytes from the peer's get_agent_metadata().
        base_addr: Base address of the peer's block memory region.
        num_blocks: Number of blocks in the peer's region.
        block_len: Size of each block (must match local block_len).

    """
    ...

add_remote_peer_async(peer_id, agent_metadata, base_addr, num_blocks, block_len)

Register a remote peer without blocking the caller.

Registration is O(num_blocks) and can take seconds for a large secondary tier, so the session starts it here and sends ConnectAckMsg only once the returned future resolves. The acknowledgement is therefore the peer's guarantee that this side is ready to serve, and no write_blocks can arrive before it.

The default implementation runs add_remote_peer inline and returns an already-resolved future; transports whose registration is cheap need not override it.

Parameters:

• peer_id

  (str) –

  Unique identifier for the remote peer.

• agent_metadata

  (bytes) –

  Opaque bytes from the peer's get_agent_metadata().

• base_addr

  (int) –

  Base address of the peer's block memory region.

• num_blocks

  (int) –

  Number of blocks in the peer's region.

• block_len

  (int) –

  Size of each block (must match local block_len).

Returns:

• Future[None] –

  A future resolving to None on success. A failed future means

• Future[None] –

  the peer must be rejected rather than acknowledged.

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

def add_remote_peer_async(
    self,
    peer_id: str,
    agent_metadata: bytes,
    base_addr: int,
    num_blocks: int,
    block_len: int,
) -> Future[None]:
    """Register a remote peer without blocking the caller.

    Registration is O(num_blocks) and can take seconds for a large
    secondary tier, so the session starts it here and sends
    ``ConnectAckMsg`` only once the returned future resolves. The
    acknowledgement is therefore the peer's guarantee that this side
    is ready to serve, and no ``write_blocks`` can arrive before it.

    The default implementation runs ``add_remote_peer`` inline and
    returns an already-resolved future; transports whose registration
    is cheap need not override it.

    Args:
        peer_id: Unique identifier for the remote peer.
        agent_metadata: Opaque bytes from the peer's get_agent_metadata().
        base_addr: Base address of the peer's block memory region.
        num_blocks: Number of blocks in the peer's region.
        block_len: Size of each block (must match local block_len).

    Returns:
        A future resolving to None on success. A failed future means
        the peer must be rejected rather than acknowledged.

    """
    future: Future[None] = Future()
    try:
        self.add_remote_peer(
            peer_id,
            agent_metadata=agent_metadata,
            base_addr=base_addr,
            num_blocks=num_blocks,
            block_len=block_len,
        )
    except Exception as exc:
        future.set_exception(exc)
    else:
        future.set_result(None)
    return future

cancel(transfer_ids, mode='immediate') abstractmethod

Cancel inflight transfers by their IDs.

Best-effort: transfers that already completed are ignored.

Parameters:

• transfer_ids

  (Iterable[int]) –

  IDs to cancel. Unknown IDs are ignored.

• mode

  (CancelMode, default: 'immediate' ) –

  "immediate" (default): pop and release each handle and return []. Matches the legacy fire-and-forget behavior — the caller does not wait for the underlying transfer to drain. "wait": release only transfers that have stopped. A transfer still PROC/PEND (or whose release fails) stays in the inflight set and its id is included in the returned list. The caller is expected to keep calling poll() until every returned id surfaces in done/failed.

Returns:

• list[int] –

  For mode="wait", the subset of transfer_ids still

• list[int] –

  tracked as inflight after the cancel attempt. For

• list[int] –

  mode="immediate", always [].

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

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

    Best-effort: transfers that already completed are ignored.

    Args:
        transfer_ids: IDs to cancel. Unknown IDs are ignored.
        mode:
            "immediate" (default): pop and release each handle and
                return []. Matches the legacy fire-and-forget
                behavior — the caller does not wait for the
                underlying transfer to drain.
            "wait": release only transfers that have stopped. A
                transfer still PROC/PEND (or whose release fails)
                stays in the inflight set and its id is included in
                the returned list. The caller is expected to keep
                calling poll() until every returned id surfaces in
                done/failed.

    Returns:
        For mode="wait", the subset of *transfer_ids* still
        tracked as inflight after the cancel attempt. For
        mode="immediate", always [].

    """
    ...

close() abstractmethod

Release all resources (registrations, handles, memory).

Cancels any remaining inflight transfers. Idempotent. After close(), no other methods may be called.

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

@abstractmethod
def close(self) -> None:
    """Release all resources (registrations, handles, memory).

    Cancels any remaining inflight transfers. Idempotent.
    After close(), no other methods may be called.
    """
    ...

get_agent_metadata() abstractmethod

Return opaque metadata needed by remote peers to connect.

The returned bytes are sent during the control-plane handshake and passed to the remote peer's add_remote_peer().

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

@abstractmethod
def get_agent_metadata(self) -> bytes:
    """Return opaque metadata needed by remote peers to connect.

    The returned bytes are sent during the control-plane handshake
    and passed to the remote peer's add_remote_peer().
    """
    ...

poll(peer_id=None) abstractmethod

Poll inflight transfers for completion.

Parameters:

• peer_id

  (str | None, default: None ) –

  If given, only poll (and drain) transfers submitted for this peer_id — the value passed to write_blocks. This is required when a single transport is shared across multiple peer sessions: poll() pops completed handles, so an unscoped poll by one session would consume and discard the completions of its siblings, starving them. None polls every peer's transfers (used only for the shutdown drain).

Returns:

• PollResult –

  PollResult with lists of completed and failed transfer_ids.

• PollResult –

  Completed/failed transfers are removed from the inflight set.

Must be called periodically to drive progress checking.

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

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

    Args:
        peer_id: If given, only poll (and drain) transfers submitted for
            this peer_id — the value passed to ``write_blocks``. This is
            required when a single transport is shared across multiple
            peer sessions: ``poll()`` pops completed handles, so an
            unscoped poll by one session would consume and discard the
            completions of its siblings, starving them. ``None`` polls
            every peer's transfers (used only for the shutdown drain).

    Returns:
        PollResult with lists of completed and failed transfer_ids.
        Completed/failed transfers are removed from the inflight set.

    Must be called periodically to drive progress checking.

    """
    ...

remove_remote_peer(peer_id) abstractmethod

Unregister a remote peer and release associated resources.

Inflight transfers to this peer should be cancelled first.

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

@abstractmethod
def remove_remote_peer(self, peer_id: str) -> None:
    """Unregister a remote peer and release associated resources.

    Inflight transfers to this peer should be cancelled first.
    """
    ...

write_blocks(peer_id, local_idxs, remote_idxs) abstractmethod

Submit a WRITE transfer: local blocks → remote peer's blocks.

Parameters:

• peer_id

  (str) –

  Target peer (must be registered via add_remote_peer).

• local_idxs

  (list[int]) –

  Indexes of local blocks to read from.

• remote_idxs

  (list[int]) –

  Indexes of remote blocks to write to. Must be same length as local_idxs.

Returns:

• int | None –

  A unique transfer_id (int) to track this transfer, or

• int | None –

  None if the peer is not registered or submission failed.

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

@abstractmethod
def write_blocks(
    self,
    peer_id: str,
    local_idxs: list[int],
    remote_idxs: list[int],
) -> int | None:
    """Submit a WRITE transfer: local blocks → remote peer's blocks.

    Args:
        peer_id: Target peer (must be registered via add_remote_peer).
        local_idxs: Indexes of local blocks to read from.
        remote_idxs: Indexes of remote blocks to write to.
                     Must be same length as local_idxs.

    Returns:
        A unique transfer_id (int) to track this transfer, or
        None if the peer is not registered or submission failed.

    """
    ...

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()
        # Registration is O(num_blocks) NIXL work — add_remote_agent plus a
        # prep_xfer_dlist over every block of the peer's region — so it runs
        # off the scheduler thread, mirroring the NIXL connector's
        # _handshake_initiation_executor. One worker only: NIXL is not
        # guaranteed to be thread-safe, so registrations stay serialized
        # against each other.
        self._reg_executor = ThreadPoolExecutor(
            max_workers=1, thread_name_prefix="vllm-p2p-nixl-reg"
        )
        # peer_id → registration generation, bumped on every add and remove.
        # A worker that finishes building handles under a stale generation
        # releases them instead of publishing, so a late completion cannot
        # revive a peer that was reaped or superseded meanwhile.
        self._peer_gen: dict[str, int] = {}
        # Guards _peer_gen together with _peer_nixl_names/_remote_dlists so
        # the generation check and the publish are one atomic step.
        self._peer_lock = threading.Lock()

        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:
        gen = self._begin_registration(peer_id)
        nixl_name, remote_dlist = self._build_peer_handles(
            agent_metadata, base_addr, num_blocks, block_len
        )
        self._publish_peer(peer_id, gen, nixl_name, remote_dlist)

    def add_remote_peer_async(
        self,
        peer_id: str,
        agent_metadata: bytes,
        base_addr: int,
        num_blocks: int,
        block_len: int,
    ) -> Future[None]:
        gen = self._begin_registration(peer_id)

        def register() -> None:
            nixl_name, remote_dlist = self._build_peer_handles(
                agent_metadata, base_addr, num_blocks, block_len
            )
            if not self._publish_peer(peer_id, gen, nixl_name, remote_dlist):
                raise RuntimeError(
                    f"registration for peer {peer_id} was superseded "
                    "(peer removed or re-registered while registering)"
                )

        return self._reg_executor.submit(register)

    def remove_remote_peer(self, peer_id: str) -> None:
        with self._peer_lock:
            # Invalidate any registration still being built for this peer so
            # its worker releases the handles instead of publishing them.
            self._peer_gen[peer_id] = self._peer_gen.get(peer_id, 0) + 1
            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)

    def _begin_registration(self, peer_id: str) -> int:
        """Claim the next registration generation for *peer_id*."""
        with self._peer_lock:
            gen = self._peer_gen.get(peer_id, 0) + 1
            self._peer_gen[peer_id] = gen
            return gen

    def _build_peer_handles(
        self,
        agent_metadata: bytes,
        base_addr: int,
        num_blocks: int,
        block_len: int,
    ) -> tuple[str, Any]:
        """The expensive half: agent wireup plus one descriptor per block.

        Touches no shared state, so it runs unlocked — holding the peer
        lock across it would stall the scheduler thread for exactly as
        long as the registration it is meant to move off that thread.
        """
        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)
        return nixl_name, remote_dlist

    def _publish_peer(
        self,
        peer_id: str,
        gen: int,
        nixl_name: str,
        remote_dlist: Any,
    ) -> bool:
        """Install freshly built handles, unless *gen* is stale.

        Returns False when the peer was removed or re-registered while the
        handles were being built; the handles are released here since
        nothing else ever saw them.
        """
        with self._peer_lock:
            fresh = self._peer_gen.get(peer_id) == gen
            if fresh:
                self._peer_nixl_names[peer_id] = nixl_name
                self._remote_dlists[peer_id] = remote_dlist
        if fresh:
            return True
        logger.info(
            "NixlTransport %s: discarding superseded registration for peer=%s",
            self._agent_name,
            peer_id,
        )
        if self._agent is not None:
            try:
                self._agent.release_dlist_handle(remote_dlist)
                self._agent.remove_remote_agent(nixl_name)
            except Exception as exc:
                logger.warning(
                    "NixlTransport %s: releasing superseded registration "
                    "for peer=%s failed: %s",
                    self._agent_name,
                    peer_id,
                    exc,
                )
        return False

    # ------------------------------------------------------------------
    # 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,
                )
                if not _is_remote_disconnect(exc):
                    continue
                state = "ERR"
            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, a
        transfer is released only once NIXL reports it terminal. Releasing
        a PROC/PEND handle does not stop the WRITE: the UCX backend frees
        the request without cancelling RMA, and the bytes keep landing. So
        PROC/PEND transfers, and those whose state check or
        ``release_xfer_handle`` raises, stay in ``self._inflight`` for a
        later ``poll()`` to observe. A state check that fails because the
        peer disconnected counts as terminal.
        """
        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:
                state = self._agent.check_xfer_state(entry.handle)
            except Exception as exc:
                logger.debug(
                    "NixlTransport %s: cancel state check failed for "
                    "transfer_id=%d: %s",
                    self._agent_name,
                    tid,
                    exc,
                )
                if not _is_remote_disconnect(exc):
                    still_inflight.append(tid)
                    continue
                # The handle may no longer be releasable; it is terminal
                # either way, so drop it instead of retrying forever.
                self._release_handles([entry.handle])
                del self._inflight[tid]
                continue
            if state in ("PROC", "PEND"):
                still_inflight.append(tid)
                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:
        # Stop accepting registrations and let any in-flight one finish
        # before the agent goes away underneath it.
        self._reg_executor.shutdown(wait=True)
        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,
                )

_begin_registration(peer_id)

Claim the next registration generation for peer_id.

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

def _begin_registration(self, peer_id: str) -> int:
    """Claim the next registration generation for *peer_id*."""
    with self._peer_lock:
        gen = self._peer_gen.get(peer_id, 0) + 1
        self._peer_gen[peer_id] = gen
        return gen

_build_peer_handles(agent_metadata, base_addr, num_blocks, block_len)

The expensive half: agent wireup plus one descriptor per block.

Touches no shared state, so it runs unlocked — holding the peer lock across it would stall the scheduler thread for exactly as long as the registration it is meant to move off that thread.

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

def _build_peer_handles(
    self,
    agent_metadata: bytes,
    base_addr: int,
    num_blocks: int,
    block_len: int,
) -> tuple[str, Any]:
    """The expensive half: agent wireup plus one descriptor per block.

    Touches no shared state, so it runs unlocked — holding the peer
    lock across it would stall the scheduler thread for exactly as
    long as the registration it is meant to move off that thread.
    """
    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)
    return nixl_name, remote_dlist

_publish_peer(peer_id, gen, nixl_name, remote_dlist)

Install freshly built handles, unless gen is stale.

Returns False when the peer was removed or re-registered while the handles were being built; the handles are released here since nothing else ever saw them.

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

def _publish_peer(
    self,
    peer_id: str,
    gen: int,
    nixl_name: str,
    remote_dlist: Any,
) -> bool:
    """Install freshly built handles, unless *gen* is stale.

    Returns False when the peer was removed or re-registered while the
    handles were being built; the handles are released here since
    nothing else ever saw them.
    """
    with self._peer_lock:
        fresh = self._peer_gen.get(peer_id) == gen
        if fresh:
            self._peer_nixl_names[peer_id] = nixl_name
            self._remote_dlists[peer_id] = remote_dlist
    if fresh:
        return True
    logger.info(
        "NixlTransport %s: discarding superseded registration for peer=%s",
        self._agent_name,
        peer_id,
    )
    if self._agent is not None:
        try:
            self._agent.release_dlist_handle(remote_dlist)
            self._agent.remove_remote_agent(nixl_name)
        except Exception as exc:
            logger.warning(
                "NixlTransport %s: releasing superseded registration "
                "for peer=%s failed: %s",
                self._agent_name,
                peer_id,
                exc,
            )
    return False

cancel(transfer_ids, mode='immediate')

Cancel inflight transfers by their IDs.

See DataTransport.cancel for the contract. In "wait" mode, a transfer is released only once NIXL reports it terminal. Releasing a PROC/PEND handle does not stop the WRITE: the UCX backend frees the request without cancelling RMA, and the bytes keep landing. So PROC/PEND transfers, and those whose state check or release_xfer_handle raises, stay in self._inflight for a later poll() to observe. A state check that fails because the peer disconnected counts as terminal.

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, a
    transfer is released only once NIXL reports it terminal. Releasing
    a PROC/PEND handle does not stop the WRITE: the UCX backend frees
    the request without cancelling RMA, and the bytes keep landing. So
    PROC/PEND transfers, and those whose state check or
    ``release_xfer_handle`` raises, stay in ``self._inflight`` for a
    later ``poll()`` to observe. A state check that fails because the
    peer disconnected counts as terminal.
    """
    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:
            state = self._agent.check_xfer_state(entry.handle)
        except Exception as exc:
            logger.debug(
                "NixlTransport %s: cancel state check failed for "
                "transfer_id=%d: %s",
                self._agent_name,
                tid,
                exc,
            )
            if not _is_remote_disconnect(exc):
                still_inflight.append(tid)
                continue
            # The handle may no longer be releasable; it is terminal
            # either way, so drop it instead of retrying forever.
            self._release_handles([entry.handle])
            del self._inflight[tid]
            continue
        if state in ("PROC", "PEND"):
            still_inflight.append(tid)
            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,
            )
            if not _is_remote_disconnect(exc):
                continue
            state = "ERR"
        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

PollResult

Bases: NamedTuple

Result of polling inflight transfers.

Attributes:

• done (Sequence[int]) –

  Transfer IDs that completed successfully.

• failed (Sequence[int]) –

  Transfer IDs that failed (error, timeout, etc.).

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

class PollResult(NamedTuple):
    """Result of polling inflight transfers.

    Attributes:
        done: Transfer IDs that completed successfully.
        failed: Transfer IDs that failed (error, timeout, etc.).

    """

    done: Sequence[int]
    failed: Sequence[int]