From 87e86bf897c3a1aabc96205f8a4ae433f70751f1 Mon Sep 17 00:00:00 2001
From: Sohom Chakraborty <16609933+sohom-cs@users.noreply.github.com>
Date: Tue, 29 Sep 2026 20:32:23 +0000
Subject: [PATCH] [Bugfix][NIXL] Stop clearing finished_sending for
 replicated-PCP producer ranks

#53903 made replicated-PCP producer ranks > 0 report every request they
were asked to send as done, so that the world_size aggregation in
KVOutputAggregator completes, and removed the code that hid those
completions from get_finished(). #53781 then added
NixlBaseConnector.get_transfer_results(), which the model runner now
calls instead of get_finished(), with the old clear() for those ranks
carried over. The synthetic completions were dropped again, the
aggregator waited for pcp_size reports and received one, and the
prefill side never freed a finished request's blocks.

Return the worker's results unchanged, as get_finished() does. Update
the handshake test that asserted the empty set, and add an aggregator
test for a PCP=2 replicated producer.

Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: Sohom Chakraborty <16609933+sohom-cs@users.noreply.github.com>
---
 .../kv_connector/unit/test_nixl_connector.py  | 70 ++++++++++++++++++-
 .../kv_connector/v1/nixl/connector.py         |  9 +--
 2 files changed, 70 insertions(+), 9 deletions(-)

diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py
index 017dee10d8b1..13ccfa808289 100644
--- a/tests/v1/kv_connector/unit/test_nixl_connector.py
+++ b/tests/v1/kv_connector/unit/test_nixl_connector.py
@@ -722,8 +722,76 @@ def test_pcp_producer_exposes_dcp_shards_or_canonical_replica(
         worker.get_transfer_results = MagicMock(
             return_value=KVConnectorTransferResults(finished_sending={"sent"})
         )
+        # The runner reads completions through get_transfer_results, so it must
+        # agree with get_finished: replicated ranks report their synthetic
+        # completions too, or the world_size aggregation never finishes.
         results = connector.get_transfer_results(set())
-        assert results.finished_sending == ({"sent"} if expected_tracked else set())
+        assert results.finished_sending == {"sent"}
+
+    @patch(
+        "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
+        FakeNixlWrapper,
+    )
+    def test_replicated_pcp_producer_send_aggregation_completes(
+        self, default_vllm_config, dist_init
+    ):
+        """PCP=2 replicated producer: D reads from the canonical rank 0 and
+        notifies only it; rank 1 reports a synthetic completion. The
+        world_size=2 aggregation must report the request as sent so the
+        scheduler frees its blocks on P."""
+        from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+
+        vllm_config = create_vllm_config(kv_role="kv_producer")
+        vllm_config.parallel_config.prefill_context_parallel_size = 2
+        vllm_config.parallel_config.decode_context_parallel_size = 1
+        req_id = "req"
+        connectors = []
+        for pcp_rank in (0, 1):
+            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
+            worker.transfer_topo = MagicMock()
+            metadata = NixlConnectorMetadata()
+            metadata.reqs_in_batch.add(req_id)
+            metadata.reqs_to_send[req_id] = time.perf_counter() + 30
+            worker.start_load_kv(metadata)
+            connectors.append(connector)
+        connectors[0].connector_worker._get_new_notifs = MagicMock(
+            return_value={req_id}
+        )
+
+        aggregator = KVOutputAggregator.from_connector(connectors[0], world_size=2)
+        outputs = [
+            ModelRunnerOutput(
+                req_ids=[],
+                req_id_to_index={},
+                sampled_token_ids=[],
+                logprobs=None,
+                prompt_logprobs_dict={},
+                pooler_output=[],
+                kv_connector_output=KVConnectorOutput(
+                    finished_sending=c.get_transfer_results(set()).finished_sending
+                ),
+            )
+            for c in connectors
+        ]
+        aggregated = aggregator.aggregate(outputs)
+        assert aggregated.kv_connector_output.finished_sending == {req_id}
 
     @patch(
         "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
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 ac929f7fcd94..1696ab77e172 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
@@ -244,14 +244,7 @@ def get_transfer_results(
         self, finished_req_ids: set[str]
     ) -> KVConnectorTransferResults:
         assert self.connector_worker is not None
-        results = self.connector_worker.get_transfer_results()
-        if (
-            self.kv_transfer_config.kv_role == "kv_producer"
-            and self.connector_worker.pcp_rank > 0
-            and not self.connector_worker.pcp_dcp_sharded
-        ):
-            results.finished_sending.clear()
-        return results
+        return self.connector_worker.get_transfer_results()
 
     def get_block_ids_with_load_errors(self) -> set[int]:
         """Get block IDs that failed to load via NIXL."""