From 0a3a4f7f35b7e6e70a1804742dec65b6bec1d2fa Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Fri, 21 Aug 2026 16:22:01 -0400 Subject: [PATCH] [Bugfix][KV Connector][NIXL] Support PCP producers (#52779) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- .../kv_connector/unit/test_multi_connector.py | 12 ++--- .../kv_connector/unit/test_nixl_connector.py | 52 +++++++++++++++++++ .../unit/test_nixl_push_connector.py | 19 +++++++ .../kv_connector/v1/multi_connector.py | 15 ++++-- .../kv_connector/v1/nixl/base_worker.py | 27 ++++++++++ .../kv_connector/v1/nixl/connector.py | 48 ++++++++++++++++- .../kv_connector/v1/nixl/metadata.py | 5 +- .../kv_connector/v1/nixl/pull_worker.py | 3 ++ .../kv_connector/v1/nixl/push_worker.py | 3 ++ 9 files changed, 173 insertions(+), 11 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index ceeaa2cf49..28d2d0dc13 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/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", diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index a37b12f62a..4b2ae31ba5 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/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, diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index f69bd544f5..e7326eddcb 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 957b15ea40..743a23e791 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index d2c073676d..7334a87466 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index 0b7a5d5728..77caec602c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index ed2e030c73..12cfa51d21 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index ecc11c480b..ea04702e1b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 0721bf9d1d..78fe49bb57 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/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(