[NIXL][PCP] Report replicated-PCP ranks > 0 as done sending instead of hiding them (#53903)
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@mistral.ai>
Lucas Wilkinson committed
Sep 10, 2026 at 06:58 UTC
93911fcf68117f25a6a13942804f8c51dc31d465
4 files changed
+20
-19
tests/v1/kv_connector/unit/test_nixl_connector.py
+8
-3
@@ -584,7 +584,7 @@ class TestNixlHandshake:
584
def test_pcp_producer_uses_canonical_replica(
585
self, default_vllm_config, dist_init, pcp_rank
586
):
587
- """Only PCP rank zero publishes and reports sending completion."""
587
+ """Only PCP rank zero publishes, but every rank reports completion."""
588
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
589
590
vllm_config = create_vllm_config(kv_role="kv_producer")
@@ -622,13 +622,18 @@ class TestNixlHandshake:
622
623
payload = MagicMock(spec=NixlHandshakePayload)
624
worker.xfer_handshake_metadata = payload
625
- worker.get_finished = MagicMock(return_value=({"sent"}, set()))
625
+ worker.transfer_topo = MagicMock()
626
+ worker._get_new_notifs = MagicMock(
627
+ side_effect=lambda: {"sent"} if pcp_rank == 0 else set()
628
+ )
629
630
expected_payload = payload if pcp_rank == 0 else None
631
assert connector.get_handshake_metadata() is expected_payload
632
done_sending, done_recving = connector.get_finished(set())
630
- assert done_sending == ({"sent"} if pcp_rank == 0 else set())
633
+ assert done_sending == ({"sent"} if pcp_rank == 0 else {req_id})
634
assert done_recving == set()
635
+ if pcp_rank > 0:
636
+ assert connector.get_finished(set()) == (set(), set())
637
638
@patch(
639
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+6
@@ -656,6 +656,9 @@ class NixlBaseConnectorWorker:
656
self._recving_transfers = defaultdict[ReqId, list[TransferHandle]](list)
657
# Track the expiration time of requests that are waiting to be sent.
658
self._reqs_to_send: dict[ReqId, float] = {}
659
+ # Replicated-KV PCP ranks > 0 never transfer; requests they are asked
660
+ # to send are reported finished immediately (see get_finished).
661
+ self._replicated_pcp_done_sending: set[ReqId] = set()
662
# Set of requests that have been part of a batch, regardless of status.
663
self._reqs_to_process: set[ReqId] = set()
664
@@ -2543,6 +2546,9 @@ class NixlBaseConnectorWorker:
2546
done_sending = self._get_new_notifs()
2547
done_recving = self._pop_done_transfers(self._recving_transfers)
2548
2549
+ done_sending.update(self._replicated_pcp_done_sending)
2550
+ self._replicated_pcp_done_sending.clear()
2551
+
2552
# Drain queue of requests where handshake or transfer setup failed.
2553
failed_recv_reqs = set[ReqId]()
2554
while not self._failed_recv_reqs.empty():
vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
+1
-16
@@ -156,15 +156,6 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
156
############################################################
157
158
def get_finished_count(self) -> int | None:
159
- parallel_config = self._vllm_config.parallel_config
160
- if (
161
- self.kv_transfer_config.kv_role == "kv_producer"
162
- and parallel_config.prefill_context_parallel_size > 1
163
- ):
164
- return (
165
- parallel_config.tensor_parallel_size
166
- * parallel_config.pipeline_parallel_size
167
- )
159
return None
160
161
def get_num_new_matched_tokens(
@@ -241,13 +232,7 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
232
def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
233
"""Get the finished recving and sending requests."""
234
assert self.connector_worker is not None
244
- done_sending, done_recving = self.connector_worker.get_finished()
245
- if (
246
- self.kv_transfer_config.kv_role == "kv_producer"
247
- and self.connector_worker.pcp_rank > 0
248
- ):
249
- done_sending.clear()
250
- return done_sending, done_recving
235
+ return self.connector_worker.get_finished()
236
237
def get_block_ids_with_load_errors(self) -> set[int]:
238
"""Get block IDs that failed to load via NIXL."""
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+5
@@ -80,6 +80,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
80
self._read_blocks_for_req(*self._ready_requests.get_nowait())
81
82
if self.pcp_rank > 0:
83
+ # Replicated-KV PCP: only PCP rank 0 serves the KV, so this rank
84
+ # has nothing to send. Report the requests as sent right away so
85
+ # the scheduler-side aggregation (world_size workers, and any
86
+ # sibling connector inside a MultiConnector) still completes.
87
+ self._replicated_pcp_done_sending.update(metadata.reqs_to_send)
88
return
89
90
# Keep around the requests that have been part of a batch. This is