@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0a3a4f7f35

[Bugfix][KV Connector][NIXL] Support PCP producers (#52779)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Lucas Wilkinson committed Aug 21, 2026 at 16:22 UTC 0a3a4f7f35b7e6e70a1804742dec65b6bec1d2fa

9 files changed +173 -11

tests/v1/kv_connector/unit/test_multi_connector.py

+6 -6

@@ -283,12 +283,11 @@ def test_multi_example_connector_consistency():

283 events = get_connector_events()

284 storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])

285 storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])

286 - # First event is bind_gpu_block_pool from initialization, then

287 - # set_xfer_handshake_metadata_pp_aware, then on_new_request when the request is

288 - # enqueued, then get_num_new_matched_tokens and update_state_after_alloc from

289 - # generate().

290 - assert storage1_scheduler_events[:6] == [

286 + # Initial events bind the block pool, query completion counts, and exchange

287 + # handshake metadata before the request is enqueued.

288 + assert storage1_scheduler_events[:7] == [

289 "bind_gpu_block_pool",

290 + "get_finished_count",

291 "set_xfer_handshake_metadata_pp_aware",

292 "on_new_request",

293 "get_num_new_matched_tokens 0",

@@ -307,8 +306,9 @@ def test_multi_example_connector_consistency():

306 "wait_for_layer_load",

307 "save_kv_layer",

308 ]

310 - assert storage2_scheduler_events[:6] == [

309 + assert storage2_scheduler_events[:7] == [

310 "bind_gpu_block_pool",

311 + "get_finished_count",

312 "set_xfer_handshake_metadata_pp_aware",

313 "on_new_request",

314 "get_num_new_matched_tokens 0",

tests/v1/kv_connector/unit/test_nixl_connector.py

+52

@@ -570,6 +570,58 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

570

571

572 class TestNixlHandshake:

573 + @pytest.mark.parametrize("pcp_rank", [0, 1])

574 + @patch(

575 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",

576 + FakeNixlWrapper,

577 + )

578 + def test_pcp_producer_uses_canonical_replica(

579 + self, default_vllm_config, dist_init, pcp_rank

580 + ):

581 + """Only PCP rank zero publishes and reports sending completion."""

582 + vllm_config = create_vllm_config(kv_role="kv_producer")

583 + vllm_config.parallel_config.prefill_context_parallel_size = 2

584 + with (

585 + patch(

586 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl."

587 + "base_worker.get_current_attn_backends",

588 + return_value=[FlashAttentionBackend],

589 + ),

590 + patch(

591 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl."

592 + "base_worker.get_pcp_group"

593 + ) as mock_get_pcp_group,

594 + ):

595 + mock_get_pcp_group.return_value.rank_in_group = pcp_rank

596 + connector = NixlConnector(

597 + vllm_config,

598 + KVConnectorRole.WORKER,

599 + make_kv_cache_config(block_size=16),

600 + )

601 +

602 + worker = connector.connector_worker

603 + assert worker is not None

604 + assert worker.pcp_rank == pcp_rank

605 +

606 + req_id = "req"

607 + metadata = NixlConnectorMetadata()

608 + metadata.reqs_in_batch.add(req_id)

609 + metadata.reqs_to_send[req_id] = time.perf_counter() + 10

610 + worker.start_load_kv(metadata)

611 + expected_tracked = pcp_rank == 0

612 + assert (req_id in worker._reqs_to_process) == expected_tracked

613 + assert (req_id in worker._reqs_to_send) == expected_tracked

614 +

615 + payload = MagicMock(spec=NixlHandshakePayload)

616 + worker.xfer_handshake_metadata = payload

617 + worker.get_finished = MagicMock(return_value=({"sent"}, set()))

618 +

619 + expected_payload = payload if pcp_rank == 0 else None

620 + assert connector.get_handshake_metadata() is expected_payload

621 + done_sending, done_recving = connector.get_finished(set())

622 + assert done_sending == ({"sent"} if pcp_rank == 0 else set())

623 + assert done_recving == set()

624 +

625 @patch(

626 "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",

627 FakeNixlWrapper,

tests/v1/kv_connector/unit/test_nixl_push_connector.py

+19

@@ -334,6 +334,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):

334 w._reqs_to_send = {}

335 w.consumer_notification_counts_by_req = defaultdict(int)

336 w.tp_rank = 0

337 + w.pcp_rank = 0

338 w.world_size = 1

339 w.engine_id = "test-decode-engine"

340 w._remote_agents = {}

@@ -514,6 +515,24 @@ class TestPushWriterStartLoadKv:

515 assert w._push_writer_wake.is_set()

516 assert w.start_push_calls == []

517

518 + def test_noncanonical_pcp_rank_skips_producer_work(self):

519 + w = _StubWriterWorker.fresh()

520 + w.pcp_rank = 1

521 + w._send_heartbeats = MagicMock()

522 +

523 + meta = NixlConnectorMetadata()

524 + meta.push_finished_blocks["req"] = ([1, 2],)

525 + meta.reqs_in_batch.add("req")

526 + meta.reqs_to_send["req"] = time.perf_counter() + 10

527 +

528 + w.start_load_kv(meta)

529 +

530 + assert w._finished_blocks_inbox.empty()

531 + assert "req" not in w._reqs_to_process

532 + assert "req" not in w._reqs_to_send

533 + assert not w._push_writer_wake.is_set()

534 + w._send_heartbeats.assert_not_called()

535 +

536

537 # The P→D handshake must run on the base worker's background executor, never

538 # blocking the writer thread: ``_do_start_push_kv`` defers the WRITE until the

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

+12 -3

@@ -363,9 +363,18 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):

363 c.handle_preemptions(cm)

364

365 def get_finished_count(self) -> int | None:

366 - # TODO(https://github.com/vllm-project/vllm/issues/33400)

367 - # Currently no connectors return non-None

368 - return None

366 + child_counts = [

367 + connector.get_finished_count() for connector in self._connectors

368 + ]

369 + if any(count is None for count in child_counts):

370 + return None

371 + counts = {count for count in child_counts if count is not None}

372 + if len(counts) > 1:

373 + raise ValueError(

374 + "MultiConnector children returned incompatible finished counts: "

375 + f"{sorted(counts)}"

376 + )

377 + return next(iter(counts), None)

378

379 def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:

380 metadata_list: list[KVConnectorWorkerMetadata | None] | None = None

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

+27

@@ -62,6 +62,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import

62 )

63 from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config

64 from vllm.distributed.parallel_state import (

65 + get_pcp_group,

66 get_tensor_model_parallel_rank,

67 get_tensor_model_parallel_world_size,

68 )

@@ -369,6 +370,9 @@ class NixlBaseConnectorWorker:

370 self.engine_id: EngineId = engine_id

371 self.tp_rank = get_tensor_model_parallel_rank()

372 self.world_size = get_tensor_model_parallel_world_size()

373 + self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size

374 + self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size

375 + self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0

376

377 self.num_blocks = kv_cache_config.num_blocks

378 self.enable_permute_local_kv = False

@@ -561,6 +565,23 @@ class NixlBaseConnectorWorker:

565 "enforce_handshake_compat", True

566 )

567

568 + def _validate_remote_parallel_config(

569 + self, agent_metadata: NixlAgentMetadata

570 + ) -> None:

571 + local_pcp_size = self.pcp_size

572 + local_dcp_size = self.dcp_size

573 + remote_pcp_size = agent_metadata.pcp_size

574 + remote_dcp_size = agent_metadata.dcp_size

575 + if (local_pcp_size > 1 and remote_dcp_size > 1) or (

576 + remote_pcp_size > 1 and local_dcp_size > 1

577 + ):

578 + raise NotImplementedError(

579 + "NixlConnector PCP requires decode_context_parallel_size=1 "

580 + "on both instances. "

581 + f"Local PCP/DCP={local_pcp_size}/{local_dcp_size}; "

582 + f"remote PCP/DCP={remote_pcp_size}/{remote_dcp_size}."

583 + )

584 +

585 def _sync_block_size_with_kernel(self) -> None:

586 backends = get_current_attn_backends(self.vllm_config)

587 kernel_block_size = select_common_block_size(self.block_size, backends)

@@ -700,6 +721,8 @@ class NixlBaseConnectorWorker:

721 f"Failed to decode NixlAgentMetadata. Error: {e}"

722 ) from e

723

724 + self._validate_remote_parallel_config(metadata)

725 +

726 # Ensure engine id matches.

727 if metadata.engine_id != expected_engine_id:

728 raise RuntimeError(

@@ -1033,6 +1056,8 @@ class NixlBaseConnectorWorker:

1056 physical_blocks_per_logical_kv_block=(

1057 self._physical_blocks_per_logical_kv_block

1058 ),

1059 + dcp_size=self.dcp_size,

1060 + pcp_size=self.pcp_size,

1061 )

1062 assert self.compat_hash is not None

1063 encoder = msgspec.msgpack.Encoder()

@@ -1277,6 +1302,8 @@ class NixlBaseConnectorWorker:

1302 physical_blocks_per_logical_kv_block=(

1303 self._physical_blocks_per_logical_kv_block

1304 ),

1305 + dcp_size=self.dcp_size,

1306 + pcp_size=self.pcp_size,

1307 )

1308 # Wrap metadata in payload with hash for defensive decoding

1309 assert self.compat_hash is not None

vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py

+47 -1

@@ -131,6 +131,29 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):

131 "kv_role='kv_consumer' for decode instances. "

132 )

133

134 + parallel_config = vllm_config.parallel_config

135 + pcp_size = parallel_config.prefill_context_parallel_size

136 + kv_role = vllm_config.kv_transfer_config.kv_role

137 + if pcp_size > 1 and kv_role in ("kv_consumer", "kv_both"):

138 + raise NotImplementedError(

139 + "NixlConnector PCP currently supports kv_producer only. "

140 + "Consumers and kv_both require "

141 + "prefill_context_parallel_size=1."

142 + )

143 + if pcp_size > 1 and parallel_config.decode_context_parallel_size > 1:

144 + raise NotImplementedError(

145 + "NixlConnector PCP producers currently require "

146 + "decode_context_parallel_size=1."

147 + )

148 + # TODO: Support PCP with bidirectional KV transfer by tracking separate

149 + # send and receive completion counts.

150 + if pcp_size > 1 and vllm_config.kv_transfer_config.get_from_extra_config(

151 + "bidirectional_kv_xfer", False

152 + ):

153 + raise NotImplementedError(

154 + "NixlConnector PCP producers do not support bidirectional KV transfer."

155 + )

156 +

157 self.kv_cache_config = kv_cache_config

158 self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id

159 self.kv_transfer_config = vllm_config.kv_transfer_config

@@ -164,6 +187,18 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):

187 # Scheduler Side Methods

188 ############################################################

189

190 + def get_finished_count(self) -> int | None:

191 + parallel_config = self._vllm_config.parallel_config

192 + if (

193 + self.kv_transfer_config.kv_role == "kv_producer"

194 + and parallel_config.prefill_context_parallel_size > 1

195 + ):

196 + return (

197 + parallel_config.tensor_parallel_size

198 + * parallel_config.pipeline_parallel_size

199 + )

200 + return None

201 +

202 def get_num_new_matched_tokens(

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

204 ) -> tuple[int | None, bool]:

@@ -244,7 +279,13 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):

279 def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:

280 """Get the finished recving and sending requests."""

281 assert self.connector_worker is not None

247 - return self.connector_worker.get_finished()

282 + done_sending, done_recving = self.connector_worker.get_finished()

283 + if (

284 + self.kv_transfer_config.kv_role == "kv_producer"

285 + and self.connector_worker.pcp_rank > 0

286 + ):

287 + done_sending.clear()

288 + return done_sending, done_recving

289

290 def get_block_ids_with_load_errors(self) -> set[int]:

291 """Get block IDs that failed to load via NIXL."""

@@ -320,6 +361,11 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):

361 None if no handshake metadata is available.

362 """

363 assert self.connector_worker is not None

364 + if (

365 + self.kv_transfer_config.kv_role == "kv_producer"

366 + and self.connector_worker.pcp_rank > 0

367 + ):

368 + return None

369 return self.connector_worker.xfer_handshake_metadata

370

371

vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py

+4 -1

@@ -43,8 +43,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"

43 # clock-sync timestamp

44 # 6: Validate EAGLE/MTP speculative configuration compatibility

45 # 7: Include NIXL transfer mode (push vs pull) in the compatibility hash

46 +# 8: Add dcp_size and pcp_size to NixlAgentMetadata

47 #

47 -NIXL_CONNECTOR_VERSION: int = 7

48 +NIXL_CONNECTOR_VERSION: int = 8

51 @dataclass

@@ -60,6 +61,8 @@ class NixlAgentMetadata:

61 ssm_sizes: tuple[int, int]

62 attn_backend_name: str

63 physical_blocks_per_logical_kv_block: int

64 + dcp_size: int = 1

65 + pcp_size: int = 1

68 @dataclass

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py

@@ -76,6 +76,9 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

76 while not self._ready_requests.empty():

77 self._read_blocks_for_req(*self._ready_requests.get_nowait())

79 + if self.pcp_rank > 0:

80 + return

81 +

82 # Keep around the requests that have been part of a batch. This is

83 # needed because async scheduling pushes the misalignment between the

84 # moment in which requests expiration is set (P side) and the moment in

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py

@@ -159,6 +159,9 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

159

160 def start_load_kv(self, metadata: NixlConnectorMetadata):

161 """Pre-process metadata; defer NIXL ops to the writer thread."""

162 + if self.pcp_rank > 0:

163 + return

164 +

165 # D-side: track reqs waiting for P to push.

166 for req_id, meta in metadata.reqs_to_recv.items():

167 meta.local_physical_block_ids = self._logical_to_kernel_block_ids(