@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d6517be3cd

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 13, 23:15 UTC

[Bugfix] Missing NIXL metadata for handshake initialization if instance spans multi-node (#26338)

Signed-off-by: Guan Luo <gluo@nvidia.com> Signed-off-by: GuanLuo <41310872+GuanLuo@users.noreply.github.com> Signed-off-by: Guan Luo <41310872+GuanLuo@users.noreply.github.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>

GuanLuo committed Nov 1, 2025 at 01:16 UTC d6517be3cd06111ada0a603acaeab28dd4580641

7 files changed +321 -95

docs/features/nixl_connector_usage.md

+1 -1

@@ -81,7 +81,7 @@ python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \

81 - Default: 5600

82 - **Required for both prefiller and decoder instances**

83 - Each vLLM worker needs a unique port on its host; using the same port number across different hosts is fine

84 - - For TP/DP deployments, each worker's port on a node is computed as: base_port + dp_rank * tp_size + tp_rank (e.g., with `--tensor-parallel-size=4` and base_port=5600, tp_rank 0..3 use ports 5600, 5601, 5602, 5603 on that node).

84 + - For TP/DP deployments, each worker's port on a node is computed as: base_port + dp_rank (e.g., with `--data-parallel-size=2` and base_port=5600, dp_rank 0..1 use port 5600, 5601 on that node).

85 - Used for the initial NIXL handshake between the prefiller and the decoder

87 - `VLLM_NIXL_SIDE_CHANNEL_HOST`: Host for side channel communication

tests/v1/kv_connector/unit/test_nixl_connector.py

+102 -4

@@ -27,6 +27,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (

27 NixlAgentMetadata,

28 NixlConnector,

29 NixlConnectorMetadata,

30 + NixlConnectorScheduler,

31 NixlConnectorWorker,

32 NixlKVConnectorStats,

33 )

@@ -283,6 +284,92 @@ def test_prompt_less_than_block_size():

284 assert len(scheduler_output.scheduled_new_reqs) == 0

285

286

287 +@patch(

288 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

289 + FakeNixlWrapper,

290 +)

291 +def test_kv_transfer_handshake(dist_init):

292 + """Unit test for basic NixlConnector interface functionality."""

293 +

294 + # Test setup, we creates a scheduler that contains a NixlConnector

295 + # of role SCHEDULER, and expect it to be serving NixlAgentMetadata from

296 + # all workers of the instance.

297 + vllm_config = create_vllm_config()

298 + # in case the test runs on non-GPU machine

299 + vllm_config.kv_transfer_config.kv_buffer_device = "cpu"

300 + scheduler = create_scheduler(vllm_config)

301 +

302 + # Create two NixlConnector of role WORKER, one is the worker of

303 + # the scheduler (prefill), the other is a worker of decode instance.

304 +

305 + # Prefill connector will register KV cache to populate proper handshake

306 + # metadata.

307 + prefill_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

308 + kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(

309 + num_blocks=2, block_size=16, num_kv_heads=4, head_size=64

310 + )

311 + shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)

312 + unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)

313 + kv_caches = {

314 + "layer0": shared_tensor,

315 + "layer1": unique_tensor,

316 + "layer2": shared_tensor,

317 + }

318 + prefill_connector.register_kv_caches(kv_caches)

319 +

320 + # Simulate EngineCore initialization that would

321 + # gather connector metadata from all workers, the scheduler connector

322 + # expects metadata to be in dict[int, KVConnectorHandshakeMetadata],

323 + # where the first key is the dp_rank, the second key is the tp_rank.

324 + metadata = {0: prefill_connector.get_handshake_metadata()}

325 + scheduler_connector = scheduler.get_kv_connector()

326 + scheduler_connector.set_xfer_handshake_metadata(metadata)

327 +

328 + # Simulate a request that finishes prefill, which returns

329 + # corresponding NixlConnectorMetadata for decode instance.

330 + BLOCK_SIZE = vllm_config.cache_config.block_size

331 + NUM_EXTERNAL_FULL_BLOCKS = 2

332 + NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))

333 +

334 + request = create_request(

335 + request_id=1,

336 + block_size=BLOCK_SIZE,

337 + num_tokens=NUM_TOKENS,

338 + do_remote_decode=True,

339 + )

340 + request.status = RequestStatus.FINISHED_LENGTH_CAPPED

341 + delay, kv_connector_metadata = scheduler.get_kv_connector().request_finished(

342 + request, [0, 1, 2]

343 + )

344 + assert delay

345 +

346 + # Decode connector will be able to create handshake with the prefill connector.

347 + decode_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

348 +

349 + # Here we are testing the retrieval of NIXLAgentMetadata.

350 + # Knowing the implementation detail, we override the add_remote_agent

351 + # to validate the metadata received is the same as the one in prefill_connector.

352 + with patch.object(

353 + decode_connector.connector_worker, "add_remote_agent"

354 + ) as mock_add_remote_agent:

355 + mock_add_remote_agent.return_type = "remote_agent"

356 +

357 + decode_connector.connector_worker._nixl_handshake(

358 + kv_connector_metadata["remote_host"],

359 + kv_connector_metadata["remote_port"],

360 + kv_connector_metadata["tp_size"],

361 + kv_connector_metadata["remote_engine_id"],

362 + )

363 +

364 + received_metadata = mock_add_remote_agent.call_args.args

365 + assert received_metadata[1] == 0 # remote_tp_rank

366 + assert received_metadata[2] == 1 # remote_tp_size

367 + assert metadata[0] == received_metadata[0]

368 +

369 + # Need to shutdown the background thread to release NIXL side channel port

370 + scheduler_connector.shutdown()

371 +

372 +

373 class FakeNixlConnectorWorker(NixlConnectorWorker):

374 REMOTE_ENGINE_ID = "remote_engine"

375

@@ -313,6 +400,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

400 engine_id=self.REMOTE_ENGINE_ID,

401 agent_metadata=FakeNixlWrapper.AGENT_METADATA,

402 kv_caches_base_addr=[0],

403 + device_id=0,

404 num_blocks=1,

405 block_lens=self.block_len_per_layer,

406 attn_backend_name=self.backend_name,

@@ -559,6 +647,7 @@ class TestNixlHandshake:

647 engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

648 agent_metadata=FakeNixlWrapper.AGENT_METADATA,

649 kv_caches_base_addr=[0],

650 + device_id=0,

651 num_blocks=1,

652 block_lens=worker.block_len_per_layer,

653 attn_backend_name=worker.backend_name,

@@ -611,6 +700,7 @@ class TestNixlHandshake:

700 engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

701 agent_metadata=FakeNixlWrapper.AGENT_METADATA,

702 kv_caches_base_addr=[0],

703 + device_id=0,

704 num_blocks=1,

705 # prefill TP=1, decode TP=2, remote block_lens is double to local

706 block_lens=[i * 2 for i in worker.block_len_per_layer],

@@ -1005,6 +1095,8 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):

1095 _ = llm.generate([f"What is the capital of France? {padding}"], sampling_params)

1096 # Request-0 times out and is cleared!

1097 assert "0" not in req_to_blocks

1098 + # Need to shutdown the background thread to release NIXL side channel port

1099 + llm.llm_engine.engine_core.shutdown()

1100

1101

1102 def test_register_kv_caches(dist_init):

@@ -1177,13 +1269,15 @@ def test_shutdown_cleans_up_resources(dist_init):

1269 """Test that shutdown() properly cleans up all resources."""

1270 vllm_config = create_vllm_config()

1271

1272 + scheduler = NixlConnectorScheduler(

1273 + vllm_config, vllm_config.kv_transfer_config.engine_id

1274 + )

1275 worker = NixlConnectorWorker(vllm_config, vllm_config.kv_transfer_config.engine_id)

1276 nixl_wrapper = worker.nixl_wrapper

1277

1278 with (

1279 patch.object(worker, "_handshake_initiation_executor") as mock_exec,

1185 - patch.object(worker, "_nixl_handshake_listener_t") as mock_listener,

1186 - patch.object(worker, "_nixl_handshake_listener_stop_event") as mock_event,

1280 + patch.object(scheduler, "_nixl_handshake_listener_t") as mock_listener,

1281 patch.object(nixl_wrapper, "release_xfer_handle") as mock_rel_xfer,

1282 patch.object(nixl_wrapper, "release_dlist_handle") as mock_rel_dlist,

1283 patch.object(nixl_wrapper, "remove_remote_agent") as mock_rem_agent,

@@ -1204,8 +1298,12 @@ def test_shutdown_cleans_up_resources(dist_init):

1298 worker.shutdown()

1299

1300 mock_exec.shutdown.assert_called_with(wait=False)

1207 - mock_event.set.assert_called_once()

1208 - mock_listener.join.assert_called_once_with(timeout=1.0)

1301 +

1302 + # Same sequence on scheduler.shutdown()

1303 + scheduler.shutdown()

1304 + scheduler.shutdown()

1305 + scheduler.shutdown()

1306 + mock_listener.join.assert_called_once()

1307

1308 mock_rel_xfer.assert_called_once_with(123)

1309 assert mock_rel_dlist.call_count == 2

vllm/distributed/kv_transfer/kv_connector/v1/base.py

+32

@@ -122,6 +122,15 @@ class KVConnectorRole(enum.Enum):

122 WORKER = 1

123

124

125 +class KVConnectorHandshakeMetadata(ABC): # noqa: B024

126 + """

127 + Metadata used for out of band connector handshake between

128 + P/D workers. This needs to serializeable.

129 + """

130 +

131 + pass

132 +

133 +

134 class KVConnectorMetadata(ABC): # noqa: B024

135 """

136 Abstract Metadata used to communicate between the

@@ -320,6 +329,18 @@ class KVConnectorBase_V1(ABC):

329 """

330 return None

331

332 + def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:

333 + """

334 + Get the KVConnector handshake metadata for this connector.

335 + This metadata is used for out-of-band connector handshake

336 + between P/D workers.

337 +

338 + Returns:

339 + KVConnectorHandshakeMetadata: the handshake metadata.

340 + None if no handshake metadata is available.

341 + """

342 + return None

343 +

344 # ==============================

345 # Scheduler-side methods

346 # ==============================

@@ -477,6 +498,17 @@ class KVConnectorBase_V1(ABC):

498 """

499 return None

500

501 + def set_xfer_handshake_metadata(

502 + self, metadata: dict[int, KVConnectorHandshakeMetadata]

503 + ) -> None:

504 + """

505 + Set the KV connector handshake metadata for this connector.

506 +

507 + Args:

508 + metadata (KVConnectorHandshakeMetadata): the handshake metadata to set.

509 + """

510 + return None

511 +

512 @classmethod

513 def build_prom_metrics(

514 cls,

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+136 -88

@@ -27,6 +27,7 @@ from vllm.config import VllmConfig

27 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

28 CopyBlocksOp,

29 KVConnectorBase_V1,

30 + KVConnectorHandshakeMetadata,

31 KVConnectorMetadata,

32 KVConnectorRole,

33 )

@@ -93,15 +94,12 @@ _NIXL_SUPPORTED_DEVICE = {

94 _NIXL_SUPPORTED_DEVICE.update(current_platform.get_nixl_supported_devices())

96 -class NixlAgentMetadata(

97 - msgspec.Struct,

98 - omit_defaults=True, # type: ignore[call-arg]

99 - # required for @cached_property.

100 - dict=True,

101 -):

97 +@dataclass

98 +class NixlAgentMetadata(KVConnectorHandshakeMetadata):

99 engine_id: str

100 agent_metadata: bytes

101 kv_caches_base_addr: list[int]

102 + device_id: int

103 num_blocks: int

104 block_lens: list[int]

105 attn_backend_name: str

@@ -223,6 +221,18 @@ class NixlConnector(KVConnectorBase_V1):

221 assert self.connector_scheduler is not None

222 return self.connector_scheduler.request_finished(request, block_ids)

223

224 + def set_xfer_handshake_metadata(

225 + self, metadata: dict[int, KVConnectorHandshakeMetadata]

226 + ) -> None:

227 + """

228 + Set the KV connector handshake metadata for this connector.

229 +

230 + Args:

231 + metadata (dict): the handshake metadata to set.

232 + """

233 + assert self.connector_scheduler is not None

234 + self.connector_scheduler.set_xfer_handshake_metadata(metadata)

235 +

236 ############################################################

237 # Worker Side Methods

238 ############################################################

@@ -299,6 +309,21 @@ class NixlConnector(KVConnectorBase_V1):

309 def shutdown(self):

310 if self.connector_worker is not None:

311 self.connector_worker.shutdown()

312 + if self.connector_scheduler is not None:

313 + self.connector_scheduler.shutdown()

314 +

315 + def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:

316 + """

317 + Get the KVConnector handshake metadata for this connector.

318 + This metadata is used for out-of-band connector handshake

319 + between P/D workers.

320 +

321 + Returns:

322 + KVConnectorHandshakeMetadata: the handshake metadata.

323 + None if no handshake metadata is available.

324 + """

325 + assert self.connector_worker is not None

326 + return self.connector_worker.xfer_handshake_metadata

327

328

329 class NixlConnectorScheduler:

@@ -312,12 +337,16 @@ class NixlConnectorScheduler:

337 self.side_channel_port = (

338 envs.VLLM_NIXL_SIDE_CHANNEL_PORT

339 + vllm_config.parallel_config.data_parallel_rank

315 - * vllm_config.parallel_config.tensor_parallel_size

340 )

341 assert vllm_config.kv_transfer_config is not None

342 self.use_host_buffer = vllm_config.kv_transfer_config.kv_buffer_device == "cpu"

343 logger.info("Initializing NIXL Scheduler %s", engine_id)

344

345 + # Background thread for handling new handshake requests.

346 + self._nixl_handshake_listener_t: threading.Thread | None = None

347 + self._encoded_xfer_handshake_metadata: dict[int, Any] = {}

348 + self._stop_event = threading.Event()

349 +

350 # Requests that need to start recv/send.

351 # New requests are added by update_state_after_alloc in

352 # the scheduler. Used to make metadata passed to Worker.

@@ -330,6 +359,89 @@ class NixlConnectorScheduler:

359 # remote prefill or aborted.

360 self._reqs_not_processed: set[ReqId] = set()

361

362 + def shutdown(self):

363 + self._stop_event.set()

364 + if self._nixl_handshake_listener_t is not None:

365 + self._nixl_handshake_listener_t.join()

366 + self._nixl_handshake_listener_t = None

367 +

368 + def set_xfer_handshake_metadata(

369 + self, metadata: dict[int, KVConnectorHandshakeMetadata]

370 + ) -> None:

371 + """

372 + Set the KV connector handshake metadata for this connector.

373 +

374 + Args:

375 + metadata (dict): the handshake metadata to set.

376 + """

377 + encoded_data: dict[int, bytes] = {}

378 + encoder = msgspec.msgpack.Encoder()

379 + for tp_rank, rank_metadata in metadata.items():

380 + if not isinstance(rank_metadata, NixlAgentMetadata):

381 + raise ValueError(

382 + "NixlConnectorScheduler expects NixlAgentMetadata for "

383 + "handshake metadata."

384 + )

385 + encoded_data[tp_rank] = encoder.encode(rank_metadata)

386 + logger.debug(

387 + "Tp rank %d: encoded NixlAgentMetadata size: %s bytes",

388 + tp_rank,

389 + str(len(encoded_data[tp_rank])),

390 + )

391 + self._encoded_xfer_handshake_metadata = encoded_data

392 +

393 + # Only start the listener when we have metadata to serve.

394 + if self._nixl_handshake_listener_t is None:

395 + ready_event = threading.Event()

396 + self._nixl_handshake_listener_t = threading.Thread(

397 + target=self._nixl_handshake_listener,

398 + args=(

399 + encoded_data,

400 + ready_event,

401 + self._stop_event,

402 + self.side_channel_port,

403 + ),

404 + daemon=True,

405 + name="nixl_handshake_listener",

406 + )

407 + self._nixl_handshake_listener_t.start()

408 + ready_event.wait() # Wait for listener ZMQ socket to be ready.

409 +

410 + @staticmethod

411 + def _nixl_handshake_listener(

412 + encoded_data: dict[int, Any],

413 + ready_event: threading.Event,

414 + stop_event: threading.Event,

415 + port: int,

416 + ):

417 + """Background thread for getting new NIXL handshakes."""

418 + # NOTE(rob): this is a simple implementation. We will move

419 + # to a better approach via HTTP endpoint soon.

420 +

421 + # Listen for new requests for metadata.

422 + host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST

423 + path = make_zmq_path("tcp", host, port)

424 + logger.debug("Starting listening on path: %s", path)

425 + with zmq_ctx(zmq.ROUTER, path) as sock:

426 + sock.setsockopt(zmq.RCVTIMEO, 1000)

427 + ready_event.set()

428 + while True:

429 + try:

430 + identity, _, msg = sock.recv_multipart()

431 + except zmq.Again:

432 + if stop_event.is_set():

433 + break

434 + continue

435 + # Decode the message which contains (GET_META_MSG, rank)

436 + msg, target_tp_rank = msgspec.msgpack.decode(msg)

437 + logger.debug(

438 + "Received message for tp rank %s",

439 + target_tp_rank,

440 + )

441 + if msg != GET_META_MSG:

442 + logger.warning("Connection listener got unexpected message %s", msg)

443 + sock.send_multipart((identity, b"", encoded_data[target_tp_rank]))

444 +

445 def get_num_new_matched_tokens(

446 self, request: "Request", num_computed_tokens: int

447 ) -> tuple[int, bool]:

@@ -537,8 +649,6 @@ class NixlConnectorScheduler:

649 class NixlConnectorWorker:

650 """Implementation of Worker side methods"""

651

540 - _POLL_TIMEOUT = 0.1 # Handshake thread polls for stop event every 100ms

541 -

652 @dataclass

653 class TpKVTopology:

654 """

@@ -651,16 +761,6 @@ class NixlConnectorWorker:

761 # Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}.

762 self._remote_agents: dict[EngineId, dict[int, str]] = defaultdict(dict)

763

654 - # NIXL handshake port.

655 - # NOTE(rob): Within a DP group, each DP rank gets its own

656 - # base port (which is sent in the KVTransferParams).

657 - # Each TP rank listens/queries on the base_port + tp_rank.

658 - self.side_channel_port: int = (

659 - envs.VLLM_NIXL_SIDE_CHANNEL_PORT

660 - + vllm_config.parallel_config.data_parallel_rank

661 - * vllm_config.parallel_config.tensor_parallel_size

662 - )

663 -

764 # Metadata.

765 self.engine_id: EngineId = engine_id

766 self.tp_rank = get_tensor_model_parallel_rank()

@@ -706,6 +806,7 @@ class NixlConnectorWorker:

806 # Map of engine_id -> kv_caches_base_addr. For TP case, each local

807 # rank will still only pull from a single remote TP worker.

808 self.kv_caches_base_addr: dict[EngineId, list[int]] = {}

809 + self.device_id: int = 0

810

811 # Number of NIXL regions. Currently one region per cache

812 # (so 1 per layer for MLA, otherwise 2 per layer)

@@ -736,9 +837,8 @@ class NixlConnectorWorker:

837 # requests that skipped transfer (handshake or transfer failures)

838 self._failed_recv_reqs: set[ReqId] = set()

839

739 - # Background thread for handling new handshake requests.

740 - self._nixl_handshake_listener_t: threading.Thread | None = None

741 - self._nixl_handshake_listener_stop_event: threading.Event | None = None

840 + # Handshake metadata of this worker for NIXL transfers.

841 + self.xfer_handshake_metadata: NixlAgentMetadata | None = None

842 # Background thread for initializing new NIXL handshakes.

843 self._handshake_initiation_executor = ThreadPoolExecutor(

844 # NIXL is not guaranteed to be thread-safe, limit 1 worker.

@@ -790,42 +890,6 @@ class NixlConnectorWorker:

890 total_num_kv_heads=self.model_config.get_total_num_kv_heads(),

891 )

892

793 - @staticmethod

794 - def _nixl_handshake_listener(

795 - metadata: NixlAgentMetadata,

796 - ready_event: threading.Event,

797 - stop_event: threading.Event,

798 - base_port: int,

799 - tp_rank: int,

800 - ):

801 - """Background thread for getting new NIXL handshakes."""

802 - # NOTE(rob): this is a simple implementation. We will move

803 - # to a better approach via HTTP endpoint soon.

804 -

805 - encoder = msgspec.msgpack.Encoder()

806 - encoded_data = encoder.encode(metadata)

807 - size_in_bytes = len(encoded_data)

808 - logger.debug("Size of encoded NixlAgentMetadata: %s bytes", str(size_in_bytes))

809 -

810 - # Listen for new requests for metadata.

811 - host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST

812 - path = make_zmq_path("tcp", host, base_port + tp_rank)

813 - logger.debug("Starting listening on path: %s", path)

814 - with zmq_ctx(zmq.ROUTER, path) as sock:

815 - ready_event.set()

816 - poller = zmq.Poller()

817 - poller.register(sock, zmq.POLLIN)

818 - while not stop_event.is_set():

819 - events = dict(

820 - poller.poll(timeout=NixlConnectorWorker._POLL_TIMEOUT * 1000)

821 - )

822 - if sock not in events:

823 - continue

824 - identity, _, msg = sock.recv_multipart()

825 - if msg != GET_META_MSG:

826 - logger.warning("Connection listener got unexpected message %s", msg)

827 - sock.send_multipart((identity, b"", encoded_data))

828 -

893 def _nixl_handshake(

894 self,

895 host: str,

@@ -844,16 +908,17 @@ class NixlConnectorWorker:

908 # Handshake only with the remote TP rank that current local rank will

909 # pull from. With homogeneous TP it happens to be the same rank_i.

910 p_remote_rank = self.kv_topo.get_target_remote_rank(remote_tp_size)

847 - path = make_zmq_path("tcp", host, port + p_remote_rank)

911 + path = make_zmq_path("tcp", host, port)

912 logger.debug(

849 - "Querying metadata on path: %s at remote rank %s", path, p_remote_rank

913 + "Querying metadata on path: %s at remote tp rank %s", path, p_remote_rank

914 )

915

916 # Send query for the request.

917 with zmq_ctx(zmq.REQ, path) as sock:

918 + msg = msgspec.msgpack.encode((GET_META_MSG, p_remote_rank))

919 # Set receive timeout to 5 seconds to avoid hanging on dead server

920 sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds

856 - sock.send(GET_META_MSG)

921 + sock.send(msg)

922 metadata_bytes = sock.recv()

923 decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)

924 metadata = decoder.decode(metadata_bytes)

@@ -1042,6 +1107,10 @@ class NixlConnectorWorker:

1107 assert tensor_size_bytes == curr_tensor_size_bytes, (

1108 "All kv cache tensors must have the same size"

1109 )

1110 + # Need to make sure the device ID is non-negative for NIXL,

1111 + # Torch uses -1 to indicate CPU tensors while NIXL uses explicit

1112 + # memory type.

1113 + self.device_id = max(cache.get_device(), 0)

1114 caches_data.append(

1115 (base_addr, curr_tensor_size_bytes, self.device_id, "")

1116 )

@@ -1139,10 +1208,11 @@ class NixlConnectorWorker:

1208 assert len(self.block_window_per_layer) == self.num_layers

1209

1210 # After KV Caches registered, listen for new connections.

1142 - metadata = NixlAgentMetadata(

1211 + self.xfer_handshake_metadata = NixlAgentMetadata(

1212 engine_id=self.engine_id,

1213 agent_metadata=self.nixl_wrapper.get_agent_metadata(),

1214 kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id],

1215 + device_id=self.device_id,

1216 num_blocks=self.num_blocks,

1217 block_lens=self.block_len_per_layer,

1218 attn_backend_name=self.backend_name,

@@ -1150,22 +1220,6 @@ class NixlConnectorWorker:

1220 if not self.use_host_buffer

1221 else self.host_buffer_kv_cache_layout,

1222 )

1153 - ready_event, stop_event = threading.Event(), threading.Event()

1154 - self._nixl_handshake_listener_t = threading.Thread(

1155 - target=self._nixl_handshake_listener,

1156 - args=(

1157 - metadata,

1158 - ready_event,

1159 - stop_event,

1160 - self.side_channel_port,

1161 - self.tp_rank,

1162 - ),

1163 - daemon=True,

1164 - name="nixl_handshake_listener",

1165 - )

1166 - self._nixl_handshake_listener_t.start()

1167 - self._nixl_handshake_listener_stop_event = stop_event

1168 - ready_event.wait() # Wait for listener ZMQ socket to be ready.

1223

1224 def add_remote_agent(

1225 self,

@@ -1267,7 +1321,7 @@ class NixlConnectorWorker:

1321 # self.block_len == remote_block_len//tp_ratio bytes.

1322 addr = base_addr + block_offset + rank_offset

1323 # (addr, len, device id)

1270 - blocks_data.append((addr, kv_block_len, remote_tp_rank))

1324 + blocks_data.append((addr, kv_block_len, nixl_agent_meta.device_id))

1325

1326 if self._use_flashinfer:

1327 # With FlashInfer index V separately to allow head splitting.

@@ -1275,7 +1329,9 @@ class NixlConnectorWorker:

1329 block_offset = block_id * nixl_agent_meta.block_lens[i]

1330 addr = base_addr + block_offset + rank_offset

1331 v_addr = addr + nixl_agent_meta.block_lens[i] // 2

1278 - blocks_data.append((v_addr, kv_block_len, remote_tp_rank))

1332 + blocks_data.append(

1333 + (v_addr, kv_block_len, nixl_agent_meta.device_id)

1334 + )

1335

1336 logger.debug(

1337 "Created %s blocks for dst engine %s with remote rank %s and local rank %s",

@@ -1843,14 +1899,6 @@ class NixlConnectorWorker:

1899 def shutdown(self):

1900 """Shutdown the connector worker."""

1901 self._handshake_initiation_executor.shutdown(wait=False)

1846 - if self._nixl_handshake_listener_stop_event is not None:

1847 - self._nixl_handshake_listener_stop_event.set()

1848 - self._nixl_handshake_listener_stop_event = None

1849 - if self._nixl_handshake_listener_t is not None:

1850 - # Generous timeout to allow the thread to exit

1851 - self._nixl_handshake_listener_t.join(timeout=self._POLL_TIMEOUT * 10)

1852 - assert not self._nixl_handshake_listener_t.is_alive()

1853 - self._nixl_handshake_listener_t = None

1902 for handles in self._recving_transfers.values():

1903 for handle, _ in handles:

1904 self.nixl_wrapper.release_xfer_handle(handle)

vllm/v1/engine/core.py

+22 -1

@@ -163,6 +163,27 @@ class EngineCore:

163 vllm_config, mm_registry

164 )

165

166 + # If a KV connector is initialized for scheduler, we want to collect

167 + # handshake metadata from all workers so the connector in the scheduler

168 + # will have the full context

169 + kv_connector = self.scheduler.get_kv_connector()

170 + if kv_connector is not None:

171 + # Collect and store KV connector xfer metadata from workers

172 + # (after KV cache registration)

173 + xfer_handshake_metadata = (

174 + self.model_executor.get_kv_connector_handshake_metadata()

175 + )

176 +

177 + if xfer_handshake_metadata:

178 + # xfer_handshake_metadata is list of dicts from workers

179 + # Each dict already has structure {tp_rank: metadata}

180 + # Merge all worker dicts into a single dict

181 + content: dict[int, Any] = {}

182 + for worker_dict in xfer_handshake_metadata:

183 + if worker_dict is not None:

184 + content.update(worker_dict)

185 + kv_connector.set_xfer_handshake_metadata(content)

186 +

187 # Setup batch queue for pipeline parallelism.

188 # Batch queue for scheduled batches. This enables us to asynchronously

189 # schedule and execute batches, and is required by pipeline parallelism

@@ -178,7 +199,7 @@ class EngineCore:

199 self.request_block_hasher: Callable[[Request], list[BlockHash]] | None = None

200 if (

201 self.vllm_config.cache_config.enable_prefix_caching

181 - or self.scheduler.get_kv_connector() is not None

202 + or kv_connector is not None

203 ):

204 caching_hash_fn = get_hash_fn_by_name(

205 vllm_config.cache_config.prefix_caching_hash_algo

vllm/v1/executor/abstract.py

@@ -9,6 +9,9 @@ from typing import TYPE_CHECKING, Literal, TypeVar, overload

10 from vllm.config import VllmConfig

11 from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator

12 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (

13 + KVConnectorHandshakeMetadata,

14 +)

15 from vllm.logger import init_logger

16 from vllm.lora.request import LoRARequest

17 from vllm.tasks import SupportedTask

@@ -177,6 +180,11 @@ class Executor(ABC):

180 ):

181 raise NotImplementedError

182

183 + def get_kv_connector_handshake_metadata(

184 + self,

185 + ) -> list[dict[int, KVConnectorHandshakeMetadata]]:

186 + return self.collective_rpc("get_kv_connector_handshake_metadata")

187 +

188 @overload

189 def execute_model(

190 self,

vllm/v1/worker/gpu_worker.py

+20 -1

@@ -19,7 +19,11 @@ from vllm.distributed import (

19 init_distributed_environment,

20 set_custom_all_reduce,

21 )

22 -from vllm.distributed.kv_transfer import ensure_kv_transfer_initialized

22 +from vllm.distributed.kv_transfer import (

23 + ensure_kv_transfer_initialized,

24 + get_kv_transfer_group,

25 + has_kv_transfer_group,

26 +)

27 from vllm.distributed.parallel_state import (

28 get_pp_group,

29 get_tp_group,

@@ -348,6 +352,21 @@ class Worker(WorkerBase):

352

353 return int(self.available_kv_cache_memory_bytes)

354

355 + def get_kv_connector_handshake_metadata(self) -> dict | None:

356 + """Get KV connector metadata from this worker if available."""

357 +

358 + if not has_kv_transfer_group():

359 + return None

360 +

361 + connector = get_kv_transfer_group()

362 + # Return None for connectors that don't need to exchange handshake

363 + # metadata across workers.

364 + if (metadata := connector.get_handshake_metadata()) is None:

365 + return None

366 +

367 + tp_rank = get_tp_group().rank_in_group

368 + return {tp_rank: metadata}

369 +

370 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:

371 return self.model_runner.get_kv_cache_spec()

372