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

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

tests/v1/kv_connector/unit/test_multi_connector.py

@@ -283,12 +283,11 @@ def test_multi_example_connector_consistency():
     events = get_connector_events()
     storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])
     storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])
-    # First event is bind_gpu_block_pool from initialization, then
-    # set_xfer_handshake_metadata_pp_aware, then on_new_request when the request is
-    # enqueued, then get_num_new_matched_tokens and update_state_after_alloc from
-    # generate().
-    assert storage1_scheduler_events[:6] == [
+    # Initial events bind the block pool, query completion counts, and exchange
+    # handshake metadata before the request is enqueued.
+    assert storage1_scheduler_events[:7] == [
         "bind_gpu_block_pool",
+        "get_finished_count",
         "set_xfer_handshake_metadata_pp_aware",
         "on_new_request",
         "get_num_new_matched_tokens 0",
@@ -307,8 +306,9 @@ def test_multi_example_connector_consistency():
         "wait_for_layer_load",
         "save_kv_layer",
     ]
-    assert storage2_scheduler_events[:6] == [
+    assert storage2_scheduler_events[:7] == [
         "bind_gpu_block_pool",
+        "get_finished_count",
         "set_xfer_handshake_metadata_pp_aware",
         "on_new_request",
         "get_num_new_matched_tokens 0",

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -570,6 +570,58 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
 
 
 class TestNixlHandshake:
+    @pytest.mark.parametrize("pcp_rank", [0, 1])
+    @patch(
+        "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
+        FakeNixlWrapper,
+    )
+    def test_pcp_producer_uses_canonical_replica(
+        self, default_vllm_config, dist_init, pcp_rank
+    ):
+        """Only PCP rank zero publishes and reports sending completion."""
+        vllm_config = create_vllm_config(kv_role="kv_producer")
+        vllm_config.parallel_config.prefill_context_parallel_size = 2
+        with (
+            patch(
+                "vllm.distributed.kv_transfer.kv_connector.v1.nixl."
+                "base_worker.get_current_attn_backends",
+                return_value=[FlashAttentionBackend],
+            ),
+            patch(
+                "vllm.distributed.kv_transfer.kv_connector.v1.nixl."
+                "base_worker.get_pcp_group"
+            ) as mock_get_pcp_group,
+        ):
+            mock_get_pcp_group.return_value.rank_in_group = pcp_rank
+            connector = NixlConnector(
+                vllm_config,
+                KVConnectorRole.WORKER,
+                make_kv_cache_config(block_size=16),
+            )
+
+        worker = connector.connector_worker
+        assert worker is not None
+        assert worker.pcp_rank == pcp_rank
+
+        req_id = "req"
+        metadata = NixlConnectorMetadata()
+        metadata.reqs_in_batch.add(req_id)
+        metadata.reqs_to_send[req_id] = time.perf_counter() + 10
+        worker.start_load_kv(metadata)
+        expected_tracked = pcp_rank == 0
+        assert (req_id in worker._reqs_to_process) == expected_tracked
+        assert (req_id in worker._reqs_to_send) == expected_tracked
+
+        payload = MagicMock(spec=NixlHandshakePayload)
+        worker.xfer_handshake_metadata = payload
+        worker.get_finished = MagicMock(return_value=({"sent"}, set()))
+
+        expected_payload = payload if pcp_rank == 0 else None
+        assert connector.get_handshake_metadata() is expected_payload
+        done_sending, done_recving = connector.get_finished(set())
+        assert done_sending == ({"sent"} if pcp_rank == 0 else set())
+        assert done_recving == set()
+
     @patch(
         "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
         FakeNixlWrapper,

tests/v1/kv_connector/unit/test_nixl_push_connector.py

@@ -334,6 +334,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):
         w._reqs_to_send = {}
         w.consumer_notification_counts_by_req = defaultdict(int)
         w.tp_rank = 0
+        w.pcp_rank = 0
         w.world_size = 1
         w.engine_id = "test-decode-engine"
         w._remote_agents = {}
@@ -514,6 +515,24 @@ class TestPushWriterStartLoadKv:
         assert w._push_writer_wake.is_set()
         assert w.start_push_calls == []
 
+    def test_noncanonical_pcp_rank_skips_producer_work(self):
+        w = _StubWriterWorker.fresh()
+        w.pcp_rank = 1
+        w._send_heartbeats = MagicMock()
+
+        meta = NixlConnectorMetadata()
+        meta.push_finished_blocks["req"] = ([1, 2],)
+        meta.reqs_in_batch.add("req")
+        meta.reqs_to_send["req"] = time.perf_counter() + 10
+
+        w.start_load_kv(meta)
+
+        assert w._finished_blocks_inbox.empty()
+        assert "req" not in w._reqs_to_process
+        assert "req" not in w._reqs_to_send
+        assert not w._push_writer_wake.is_set()
+        w._send_heartbeats.assert_not_called()
+
 
 # The P→D handshake must run on the base worker's background executor, never
 # blocking the writer thread: ``_do_start_push_kv`` defers the WRITE until the

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

@@ -363,9 +363,18 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
             c.handle_preemptions(cm)
 
     def get_finished_count(self) -> int | None:
-        # TODO(https://github.com/vllm-project/vllm/issues/33400)
-        # Currently no connectors return non-None
-        return None
+        child_counts = [
+            connector.get_finished_count() for connector in self._connectors
+        ]
+        if any(count is None for count in child_counts):
+            return None
+        counts = {count for count in child_counts if count is not None}
+        if len(counts) > 1:
+            raise ValueError(
+                "MultiConnector children returned incompatible finished counts: "
+                f"{sorted(counts)}"
+            )
+        return next(iter(counts), None)
 
     def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:
         metadata_list: list[KVConnectorWorkerMetadata | None] | None = None

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

@@ -62,6 +62,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import
 )
 from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
 from vllm.distributed.parallel_state import (
+    get_pcp_group,
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
 )
@@ -369,6 +370,9 @@ class NixlBaseConnectorWorker:
         self.engine_id: EngineId = engine_id
         self.tp_rank = get_tensor_model_parallel_rank()
         self.world_size = get_tensor_model_parallel_world_size()
+        self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
+        self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size
+        self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0
 
         self.num_blocks = kv_cache_config.num_blocks
         self.enable_permute_local_kv = False
@@ -561,6 +565,23 @@ class NixlBaseConnectorWorker:
             "enforce_handshake_compat", True
         )
 
+    def _validate_remote_parallel_config(
+        self, agent_metadata: NixlAgentMetadata
+    ) -> None:
+        local_pcp_size = self.pcp_size
+        local_dcp_size = self.dcp_size
+        remote_pcp_size = agent_metadata.pcp_size
+        remote_dcp_size = agent_metadata.dcp_size
+        if (local_pcp_size > 1 and remote_dcp_size > 1) or (
+            remote_pcp_size > 1 and local_dcp_size > 1
+        ):
+            raise NotImplementedError(
+                "NixlConnector PCP requires decode_context_parallel_size=1 "
+                "on both instances. "
+                f"Local PCP/DCP={local_pcp_size}/{local_dcp_size}; "
+                f"remote PCP/DCP={remote_pcp_size}/{remote_dcp_size}."
+            )
+
     def _sync_block_size_with_kernel(self) -> None:
         backends = get_current_attn_backends(self.vllm_config)
         kernel_block_size = select_common_block_size(self.block_size, backends)
@@ -700,6 +721,8 @@ class NixlBaseConnectorWorker:
                         f"Failed to decode NixlAgentMetadata. Error: {e}"
                     ) from e
 
+                self._validate_remote_parallel_config(metadata)
+
                 # Ensure engine id matches.
                 if metadata.engine_id != expected_engine_id:
                     raise RuntimeError(
@@ -1033,6 +1056,8 @@ class NixlBaseConnectorWorker:
             physical_blocks_per_logical_kv_block=(
                 self._physical_blocks_per_logical_kv_block
             ),
+            dcp_size=self.dcp_size,
+            pcp_size=self.pcp_size,
         )
         assert self.compat_hash is not None
         encoder = msgspec.msgpack.Encoder()
@@ -1277,6 +1302,8 @@ class NixlBaseConnectorWorker:
             physical_blocks_per_logical_kv_block=(
                 self._physical_blocks_per_logical_kv_block
             ),
+            dcp_size=self.dcp_size,
+            pcp_size=self.pcp_size,
         )
         # Wrap metadata in payload with hash for defensive decoding
         assert self.compat_hash is not None

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

@@ -131,6 +131,29 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
                 "kv_role='kv_consumer' for decode instances. "
             )
 
+        parallel_config = vllm_config.parallel_config
+        pcp_size = parallel_config.prefill_context_parallel_size
+        kv_role = vllm_config.kv_transfer_config.kv_role
+        if pcp_size > 1 and kv_role in ("kv_consumer", "kv_both"):
+            raise NotImplementedError(
+                "NixlConnector PCP currently supports kv_producer only. "
+                "Consumers and kv_both require "
+                "prefill_context_parallel_size=1."
+            )
+        if pcp_size > 1 and parallel_config.decode_context_parallel_size > 1:
+            raise NotImplementedError(
+                "NixlConnector PCP producers currently require "
+                "decode_context_parallel_size=1."
+            )
+        # TODO: Support PCP with bidirectional KV transfer by tracking separate
+        # send and receive completion counts.
+        if pcp_size > 1 and vllm_config.kv_transfer_config.get_from_extra_config(
+            "bidirectional_kv_xfer", False
+        ):
+            raise NotImplementedError(
+                "NixlConnector PCP producers do not support bidirectional KV transfer."
+            )
+
         self.kv_cache_config = kv_cache_config
         self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
         self.kv_transfer_config = vllm_config.kv_transfer_config
@@ -164,6 +187,18 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
     # Scheduler Side Methods
     ############################################################
 
+    def get_finished_count(self) -> int | None:
+        parallel_config = self._vllm_config.parallel_config
+        if (
+            self.kv_transfer_config.kv_role == "kv_producer"
+            and parallel_config.prefill_context_parallel_size > 1
+        ):
+            return (
+                parallel_config.tensor_parallel_size
+                * parallel_config.pipeline_parallel_size
+            )
+        return None
+
     def get_num_new_matched_tokens(
         self, request: "Request", num_computed_tokens: int
     ) -> tuple[int | None, bool]:
@@ -244,7 +279,13 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
     def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
         """Get the finished recving and sending requests."""
         assert self.connector_worker is not None
-        return self.connector_worker.get_finished()
+        done_sending, done_recving = self.connector_worker.get_finished()
+        if (
+            self.kv_transfer_config.kv_role == "kv_producer"
+            and self.connector_worker.pcp_rank > 0
+        ):
+            done_sending.clear()
+        return done_sending, done_recving
 
     def get_block_ids_with_load_errors(self) -> set[int]:
         """Get block IDs that failed to load via NIXL."""
@@ -320,6 +361,11 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
             None if no handshake metadata is available.
         """
         assert self.connector_worker is not None
+        if (
+            self.kv_transfer_config.kv_role == "kv_producer"
+            and self.connector_worker.pcp_rank > 0
+        ):
+            return None
         return self.connector_worker.xfer_handshake_metadata
 
 

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

@@ -43,8 +43,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"
 #      clock-sync timestamp
 #   6: Validate EAGLE/MTP speculative configuration compatibility
 #   7: Include NIXL transfer mode (push vs pull) in the compatibility hash
+#   8: Add dcp_size and pcp_size to NixlAgentMetadata
 #
-NIXL_CONNECTOR_VERSION: int = 7
+NIXL_CONNECTOR_VERSION: int = 8
 
 
 @dataclass
@@ -60,6 +61,8 @@ class NixlAgentMetadata:
     ssm_sizes: tuple[int, int]
     attn_backend_name: str
     physical_blocks_per_logical_kv_block: int
+    dcp_size: int = 1
+    pcp_size: int = 1
 
 
 @dataclass

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

@@ -76,6 +76,9 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
         while not self._ready_requests.empty():
             self._read_blocks_for_req(*self._ready_requests.get_nowait())
 
+        if self.pcp_rank > 0:
+            return
+
         # Keep around the requests that have been part of a batch. This is
         # needed because async scheduling pushes the misalignment between the
         # 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):
 
     def start_load_kv(self, metadata: NixlConnectorMetadata):
         """Pre-process metadata; defer NIXL ops to the writer thread."""
+        if self.pcp_rank > 0:
+            return
+
         # D-side: track reqs waiting for P to push.
         for req_id, meta in metadata.reqs_to_recv.items():
             meta.local_physical_block_ids = self._logical_to_kernel_block_ids(