| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """NixlTransport: Data-plane transport for RDMA-based KV block transfers via NIXL.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import itertools |
| 8 | from collections.abc import Iterable |
| 9 | from typing import Any, NamedTuple |
| 10 | |
| 11 | import numpy as np |
| 12 | |
| 13 | from vllm.distributed.nixl_utils import NixlWrapper as _NixlAgent |
| 14 | from vllm.distributed.nixl_utils import nixl_agent_config as _NixlAgentConfig |
| 15 | from vllm.logger import init_logger |
| 16 | from vllm.v1.kv_offload.tiering.p2p.data.base import ( |
| 17 | CancelMode, |
| 18 | DataTransport, |
| 19 | PollResult, |
| 20 | ) |
| 21 | |
| 22 | logger = init_logger(__name__) |
| 23 | |
| 24 | # Shared sentinel returned by poll() in the steady state (no inflight, or |
| 25 | # no transfer changed state since the last poll). Tuples make it immutable; |
| 26 | # callers only iterate / membership-test / equality-check. |
| 27 | _EMPTY_POLL_RESULT: PollResult = PollResult(done=(), failed=()) |
| 28 | |
| 29 | |
| 30 | class _Inflight(NamedTuple): |
| 31 | """A submitted-but-not-yet-drained transfer. |
| 32 | |
| 33 | ``peer_id`` lets poll() scope to a single owning session, since the |
| 34 | transport is shared across all peer sessions of the engine. |
| 35 | """ |
| 36 | |
| 37 | peer_id: str |
| 38 | handle: object |
| 39 | |
| 40 | |
| 41 | class NixlTransport(DataTransport): |
| 42 | """Manages a NIXL agent, memory registration, and block transfers. |
| 43 | |
| 44 | Wraps the NIXL C library behind a Python interface so the rest of the |
| 45 | P2P tier code never touches NIXL types directly. Tracks inflight |
| 46 | handles internally and returns completed/failed tags on poll. |
| 47 | """ |
| 48 | |
| 49 | def __init__( |
| 50 | self, |
| 51 | agent_name: str, |
| 52 | view: memoryview, |
| 53 | config_fields: dict | None = None, |
| 54 | backends: list[str] | None = None, |
| 55 | num_threads: int = 4, |
| 56 | ) -> None: |
| 57 | super().__init__(view, config_fields=config_fields) |
| 58 | self._agent_name = agent_name |
| 59 | self._backends = list(backends) if backends else ["UCX"] |
| 60 | self._num_threads = num_threads |
| 61 | self._agent: Any = None |
| 62 | self._reg: Any = None |
| 63 | self._local_dlist: Any = None |
| 64 | self._remote_dlists: dict[str, object] = {} |
| 65 | self._peer_nixl_names: dict[str, str] = {} |
| 66 | # transfer_id → _Inflight(peer_id, handle). |
| 67 | self._inflight: dict[int, _Inflight] = {} |
| 68 | self._next_id = itertools.count() |
| 69 | |
| 70 | self._init(view) |
| 71 | |
| 72 | @property |
| 73 | def available(self) -> bool: |
| 74 | return self._agent is not None |
| 75 | |
| 76 | def _init(self, view: memoryview) -> None: |
| 77 | if _NixlAgent is None: |
| 78 | return |
| 79 | |
| 80 | non_ucx_backends = [b for b in self._backends if b != "UCX"] |
| 81 | if non_ucx_backends: |
| 82 | cfg = _NixlAgentConfig(backends=self._backends, capture_telemetry=True) |
| 83 | logger.info( |
| 84 | "NixlTransport %s: NIXL backends=%s", |
| 85 | self._agent_name, |
| 86 | self._backends, |
| 87 | ) |
| 88 | else: |
| 89 | cfg = _NixlAgentConfig( |
| 90 | num_threads=self._num_threads, capture_telemetry=True |
| 91 | ) |
| 92 | logger.info( |
| 93 | "NixlTransport %s: NIXL backends=[UCX] num_threads=%d", |
| 94 | self._agent_name, |
| 95 | self._num_threads, |
| 96 | ) |
| 97 | self._agent = _NixlAgent(self._agent_name, cfg) |
| 98 | |
| 99 | total_size = self._num_blocks * self._block_len |
| 100 | reg_descs = [(self._base_addr, total_size, 0, "")] |
| 101 | self._reg = self._agent.register_memory(reg_descs, mem_type="DRAM") |
| 102 | |
| 103 | block_tuples = [ |
| 104 | (self._base_addr + i * self._block_len, self._block_len, 0) |
| 105 | for i in range(self._num_blocks) |
| 106 | ] |
| 107 | xfer_dlist = self._agent.get_xfer_descs(block_tuples, mem_type="DRAM") |
| 108 | self._local_dlist = self._agent.prep_xfer_dlist("NIXL_INIT_AGENT", xfer_dlist) |
| 109 | logger.info( |
| 110 | "NixlTransport %s: registered %d blocks", self._agent_name, self._num_blocks |
| 111 | ) |
| 112 | |
| 113 | def get_agent_metadata(self) -> bytes: |
| 114 | assert self._agent is not None |
| 115 | return self._agent.get_agent_metadata() |
| 116 | |
| 117 | # ------------------------------------------------------------------ |
| 118 | # Peer management |
| 119 | # ------------------------------------------------------------------ |
| 120 | |
| 121 | def add_remote_peer( |
| 122 | self, |
| 123 | peer_id: str, |
| 124 | agent_metadata: bytes, |
| 125 | base_addr: int, |
| 126 | num_blocks: int, |
| 127 | block_len: int, |
| 128 | ) -> None: |
| 129 | nixl_name = self._agent.add_remote_agent(agent_metadata) |
| 130 | block_descs = [ |
| 131 | (base_addr + i * block_len, block_len, 0) for i in range(num_blocks) |
| 132 | ] |
| 133 | xfer_dlist = self._agent.get_xfer_descs(block_descs, mem_type="DRAM") |
| 134 | remote_dlist = self._agent.prep_xfer_dlist(nixl_name, xfer_dlist) |
| 135 | self._peer_nixl_names[peer_id] = nixl_name |
| 136 | self._remote_dlists[peer_id] = remote_dlist |
| 137 | |
| 138 | def remove_remote_peer(self, peer_id: str) -> None: |
| 139 | nixl_name = self._peer_nixl_names.pop(peer_id, None) |
| 140 | dlist = self._remote_dlists.pop(peer_id, None) |
| 141 | if self._agent is not None: |
| 142 | if dlist is not None: |
| 143 | self._agent.release_dlist_handle(dlist) |
| 144 | if nixl_name: |
| 145 | self._agent.remove_remote_agent(nixl_name) |
| 146 | |
| 147 | # ------------------------------------------------------------------ |
| 148 | # Transfer submission and polling |
| 149 | # ------------------------------------------------------------------ |
| 150 | |
| 151 | def write_blocks( |
| 152 | self, |
| 153 | peer_id: str, |
| 154 | local_idxs: list[int], |
| 155 | remote_idxs: list[int], |
| 156 | ) -> int | None: |
| 157 | """Submit a WRITE transfer to *peer_id*. |
| 158 | |
| 159 | Returns a transfer ID, or None if the peer is not registered. |
| 160 | The ID is returned via poll() when the transfer completes or fails. |
| 161 | """ |
| 162 | remote_dlist = self._remote_dlists.get(peer_id) |
| 163 | if remote_dlist is None: |
| 164 | logger.warning( |
| 165 | "NixlTransport %s: write_blocks NO REMOTE DLIST for peer=%s " |
| 166 | "(known peers=%s)", |
| 167 | self._agent_name, |
| 168 | peer_id, |
| 169 | list(self._remote_dlists.keys()), |
| 170 | ) |
| 171 | return None |
| 172 | logger.debug( |
| 173 | "NixlTransport %s: write_blocks NIXL.transfer peer=%s blocks=%d", |
| 174 | self._agent_name, |
| 175 | peer_id, |
| 176 | len(local_idxs), |
| 177 | ) |
| 178 | handle = self._agent.make_prepped_xfer( |
| 179 | "WRITE", |
| 180 | self._local_dlist, |
| 181 | np.asarray(local_idxs, dtype=np.int32), |
| 182 | remote_dlist, |
| 183 | np.asarray(remote_idxs, dtype=np.int32), |
| 184 | ) |
| 185 | self._agent.transfer(handle) |
| 186 | transfer_id = next(self._next_id) |
| 187 | self._inflight[transfer_id] = _Inflight(peer_id, handle) |
| 188 | return transfer_id |
| 189 | |
| 190 | def poll(self, peer_id: str | None = None) -> PollResult: |
| 191 | """Poll inflight transfers. |
| 192 | |
| 193 | When *peer_id* is given, only transfers submitted for that peer_id are |
| 194 | checked and drained — the transport is shared across peer sessions, so |
| 195 | an unscoped poll by one session would consume and discard siblings' |
| 196 | completions. *peer_id* None polls every peer (shutdown drain only). |
| 197 | |
| 198 | Returns PollResult(done=..., failed=...) with transfer IDs. |
| 199 | Completed handles are released automatically. |
| 200 | """ |
| 201 | if not self._inflight: |
| 202 | return _EMPTY_POLL_RESULT |
| 203 | |
| 204 | done_ids: list[int] | None = None |
| 205 | failed_ids: list[int] | None = None |
| 206 | |
| 207 | for transfer_id, entry in self._inflight.items(): |
| 208 | if peer_id is not None and entry.peer_id != peer_id: |
| 209 | continue |
| 210 | try: |
| 211 | state = self._agent.check_xfer_state(entry.handle) |
| 212 | except Exception as exc: |
| 213 | logger.warning( |
| 214 | "NixlTransport %s: check_xfer_state failed for transfer_id=%d: %s", |
| 215 | self._agent_name, |
| 216 | transfer_id, |
| 217 | exc, |
| 218 | ) |
| 219 | continue |
| 220 | if state == "DONE": |
| 221 | if done_ids is None: |
| 222 | done_ids = [] |
| 223 | done_ids.append(transfer_id) |
| 224 | elif state not in ("PROC", "PEND"): |
| 225 | if failed_ids is None: |
| 226 | failed_ids = [] |
| 227 | failed_ids.append(transfer_id) |
| 228 | |
| 229 | if done_ids is None and failed_ids is None: |
| 230 | return _EMPTY_POLL_RESULT |
| 231 | |
| 232 | handles_to_release = [] |
| 233 | for tid in done_ids or (): |
| 234 | handles_to_release.append(self._inflight.pop(tid).handle) |
| 235 | for tid in failed_ids or (): |
| 236 | handles_to_release.append(self._inflight.pop(tid).handle) |
| 237 | self._release_handles(handles_to_release) |
| 238 | |
| 239 | return PollResult( |
| 240 | done=done_ids if done_ids is not None else _EMPTY_POLL_RESULT.done, |
| 241 | failed=failed_ids if failed_ids is not None else _EMPTY_POLL_RESULT.failed, |
| 242 | ) |
| 243 | |
| 244 | def cancel( |
| 245 | self, |
| 246 | transfer_ids: Iterable[int], |
| 247 | mode: CancelMode = "immediate", |
| 248 | ) -> list[int]: |
| 249 | """Cancel inflight transfers by their IDs. |
| 250 | |
| 251 | See ``DataTransport.cancel`` for the contract. In "wait" mode, |
| 252 | transfers whose ``release_xfer_handle`` raises (NIXL could not |
| 253 | complete the abort because the backend is still draining) stay |
| 254 | in ``self._inflight`` so a later ``poll()`` will observe them. |
| 255 | """ |
| 256 | if mode == "immediate": |
| 257 | handles = [ |
| 258 | self._inflight.pop(tid).handle |
| 259 | for tid in transfer_ids |
| 260 | if tid in self._inflight |
| 261 | ] |
| 262 | self._release_handles(handles) |
| 263 | return [] |
| 264 | |
| 265 | still_inflight: list[int] = [] |
| 266 | for tid in transfer_ids: |
| 267 | entry = self._inflight.get(tid) |
| 268 | if entry is None: |
| 269 | continue |
| 270 | try: |
| 271 | self._agent.release_xfer_handle(entry.handle) |
| 272 | except Exception as exc: |
| 273 | logger.debug( |
| 274 | "NixlTransport %s: cancel pending for transfer_id=%d: %s", |
| 275 | self._agent_name, |
| 276 | tid, |
| 277 | exc, |
| 278 | ) |
| 279 | still_inflight.append(tid) |
| 280 | continue |
| 281 | del self._inflight[tid] |
| 282 | return still_inflight |
| 283 | |
| 284 | # ------------------------------------------------------------------ |
| 285 | # Lifecycle |
| 286 | # ------------------------------------------------------------------ |
| 287 | |
| 288 | def close(self) -> None: |
| 289 | if self._agent is None: |
| 290 | return |
| 291 | self._release_handles([entry.handle for entry in self._inflight.values()]) |
| 292 | self._inflight.clear() |
| 293 | for peer_id in list(self._remote_dlists): |
| 294 | self.remove_remote_peer(peer_id) |
| 295 | if self._local_dlist is not None: |
| 296 | self._agent.release_dlist_handle(self._local_dlist) |
| 297 | self._local_dlist = None |
| 298 | if self._reg is not None: |
| 299 | self._agent.deregister_memory(self._reg) |
| 300 | self._reg = None |
| 301 | self._agent = None |
| 302 | |
| 303 | # ------------------------------------------------------------------ |
| 304 | # Internal |
| 305 | # ------------------------------------------------------------------ |
| 306 | |
| 307 | def _release_handles(self, handles: list[object]) -> None: |
| 308 | if self._agent is None: |
| 309 | return |
| 310 | for handle in handles: |
| 311 | try: |
| 312 | self._agent.release_xfer_handle(handle) |
| 313 | except Exception as exc: |
| 314 | logger.warning( |
| 315 | "NixlTransport %s: release_xfer_handle failed: %s", |
| 316 | self._agent_name, |
| 317 | exc, |
| 318 | ) |