@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7fe8fc803d

[NIXL][PCP][DCP] Expose PCP producer KV shards as transfer ranks (#56645)

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

Lucas Wilkinson committed Sep 13, 2026 at 16:23 UTC 7fe8fc803d904795fd10b785776e20181aee64c5

13 files changed +124 -44

tests/v1/kv_connector/unit/test_nixl_connector.py

+57 -16

@@ -580,19 +580,38 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

580

581

582 class TestNixlHandshake:

583 - @pytest.mark.parametrize("pcp_rank", [0, 1])

583 + @pytest.mark.parametrize(

584 + ("pcp_rank", "pcp_size", "dcp_size", "expected_tracked"),

585 + [

586 + (0, 2, 1, True),

587 + (1, 2, 1, False),

588 + (0, 2, 2, True),

589 + (1, 2, 2, True),

590 + (0, 4, 4, True),

591 + (1, 4, 4, True),

592 + (2, 4, 4, True),

593 + (3, 4, 4, True),

594 + ],

595 + )

596 @patch(

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

598 FakeNixlWrapper,

599 )

588 - def test_pcp_producer_uses_canonical_replica(

589 - self, default_vllm_config, dist_init, pcp_rank

600 + def test_pcp_producer_exposes_dcp_shards_or_canonical_replica(

601 + self,

602 + default_vllm_config,

603 + dist_init,

604 + pcp_rank,

605 + pcp_size,

606 + dcp_size,

607 + expected_tracked,

608 ):

591 - """Only PCP rank zero publishes, but every rank reports completion."""

609 + """Replicated PCP is canonicalized; PCP-DCP publishes every shard."""

610 from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend

611

612 vllm_config = create_vllm_config(kv_role="kv_producer")

595 - vllm_config.parallel_config.prefill_context_parallel_size = 2

613 + vllm_config.parallel_config.prefill_context_parallel_size = pcp_size

614 + vllm_config.parallel_config.decode_context_parallel_size = dcp_size

615 with (

616 patch(

617 "vllm.distributed.kv_transfer.kv_connector.v1.nixl."

@@ -614,13 +633,15 @@ class TestNixlHandshake:

633 worker = connector.connector_worker

634 assert worker is not None

635 assert worker.pcp_rank == pcp_rank

636 + assert worker.pcp_dcp_sharded is (dcp_size > 1)

637 + assert worker.transfer_tp_rank == (pcp_rank if dcp_size > 1 else 0)

638 + assert worker.transfer_tp_size == (pcp_size if dcp_size > 1 else 1)

639

640 req_id = "req"

641 metadata = NixlConnectorMetadata()

642 metadata.reqs_in_batch.add(req_id)

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

644 worker.start_load_kv(metadata)

623 - expected_tracked = pcp_rank == 0

645 assert (req_id in worker._reqs_to_process) == expected_tracked

646 assert (req_id in worker._reqs_to_send) == expected_tracked

647

@@ -628,22 +649,22 @@ class TestNixlHandshake:

649 worker.xfer_handshake_metadata = payload

650 worker.transfer_topo = MagicMock()

651 worker._get_new_notifs = MagicMock(

631 - side_effect=lambda: {"sent"} if pcp_rank == 0 else set()

652 + side_effect=lambda: {"sent"} if expected_tracked else set()

653 )

654

634 - expected_payload = payload if pcp_rank == 0 else None

655 + expected_payload = payload if expected_tracked else None

656 assert connector.get_handshake_metadata() is expected_payload

657 done_sending, done_recving = connector.get_finished(set())

637 - assert done_sending == ({"sent"} if pcp_rank == 0 else {req_id})

658 + assert done_sending == ({"sent"} if expected_tracked else {req_id})

659 assert done_recving == set()

639 - if pcp_rank > 0:

660 + if not expected_tracked:

661 assert connector.get_finished(set()) == (set(), set())

662

663 worker.get_transfer_results = MagicMock(

664 return_value=KVConnectorTransferResults(finished_sending={"sent"})

665 )

666 results = connector.get_transfer_results(set())

646 - assert results.finished_sending == ({"sent"} if pcp_rank == 0 else set())

667 + assert results.finished_sending == ({"sent"} if expected_tracked else set())

668

669 @patch(

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

@@ -3354,9 +3375,15 @@ def test_transfer_mode_changes_compatibility_hash():

3375

3376

3377 @pytest.mark.skip_global_cleanup

3357 -def test_scheduler_advertises_transfer_mode():

3358 - # Each scheduler advertises its transfer mode in kv_transfer_params so an

3359 - # external router can route pull (READ) vs push (WRITE) producers.

3378 +@pytest.mark.parametrize("mode", ["pull", "push"])

3379 +@pytest.mark.parametrize(

3380 + "tp_size,pcp_size,dcp_size,transfer_tp_size",

3381 + [(1, 1, 1, 1), (4, 1, 1, 4), (4, 1, 4, 4), (1, 4, 1, 1), (1, 4, 4, 4)],

3382 +)

3383 +def test_scheduler_advertises_transfer_topology(

3384 + mode, tp_size, pcp_size, dcp_size, transfer_tp_size

3385 +):

3386 + """The consumer must address every distinct producer KV shard."""

3387 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (

3388 NixlPullConnectorScheduler,

3389 )

@@ -3364,8 +3391,22 @@ def test_scheduler_advertises_transfer_mode():

3391 NixlPushConnectorScheduler,

3392 )

3393

3367 - assert NixlPullConnectorScheduler._TRANSFER_MODE == "pull"

3368 - assert NixlPushConnectorScheduler._TRANSFER_MODE == "push"

3394 + config = create_vllm_config(kv_role="kv_producer")

3395 + config.parallel_config.tensor_parallel_size = tp_size

3396 + config.parallel_config.prefill_context_parallel_size = pcp_size

3397 + config.parallel_config.decode_context_parallel_size = dcp_size

3398 + cls = NixlPullConnectorScheduler if mode == "pull" else NixlPushConnectorScheduler

3399 + scheduler = cls(config, "prefiller", make_kv_cache_config(block_size=16))

3400 + request = create_request(request_id=1, num_tokens=32, do_remote_decode=True)

3401 + request.status = RequestStatus.FINISHED_LENGTH_CAPPED

3402 + try:

3403 + delay, params = scheduler.request_finished(request, ([0, 1],))

3404 + assert delay

3405 + assert params["transfer_mode"] == mode

3406 + assert params["tp_size"] == transfer_tp_size

3407 + assert params["dcp_size"] == dcp_size

3408 + finally:

3409 + scheduler.shutdown()

3410

3411

3412 @pytest.mark.parametrize(

tests/v1/kv_connector/unit/test_nixl_push_connector.py

+9 -6

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

358 w.consumer_notification_counts_by_req = defaultdict(int)

359 w.tp_rank = 0

360 w.pcp_rank = 0

361 + w.pcp_dcp_sharded = False

362 w.world_size = 1

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

364 w._remote_agents = {}

@@ -545,9 +546,11 @@ class TestPushWriterStartLoadKv:

546 assert w._push_writer_wake.is_set()

547 assert w.start_push_calls == []

548

548 - def test_noncanonical_pcp_rank_skips_producer_work(self):

549 + @pytest.mark.parametrize("sharded", [False, True])

550 + def test_noncanonical_pcp_rank_only_pushes_distinct_shards(self, sharded):

551 w = _StubWriterWorker.fresh()

552 w.pcp_rank = 1

553 + w.pcp_dcp_sharded = sharded

554 w._send_heartbeats = MagicMock()

555

556 meta = NixlConnectorMetadata()

@@ -557,11 +560,11 @@ class TestPushWriterStartLoadKv:

560

561 w.start_load_kv(meta)

562

560 - assert w._finished_blocks_inbox.empty()

561 - assert "req" not in w._reqs_to_process

562 - assert "req" not in w._reqs_to_send

563 - assert not w._push_writer_wake.is_set()

564 - w._send_heartbeats.assert_not_called()

563 + assert w._finished_blocks_inbox.empty() is (not sharded)

564 + assert ("req" in w._reqs_to_process) is sharded

565 + assert ("req" in w._reqs_to_send) is sharded

566 + assert w._push_writer_wake.is_set() is sharded

567 + assert w._send_heartbeats.called is sharded

568

569

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

tests/v1/kv_connector/unit/test_tp_mapping.py

@@ -75,6 +75,7 @@ class TestTPMappingStructure:

75 @pytest.mark.parametrize(

76 "tp_rank,tp_size,remote_tp_size,dcp_size,remote_dcp_size,expected_ranks",

77 [

78 + (0, 1, 8, 1, 8, tuple(range(8))),

79 (0, 4, 4, 1, 4, (0, 1, 2, 3)),

80 (2, 4, 4, 4, 4, (2,)),

81 (0, 2, 4, 2, 4, (0, 2)),

tests/v1/kv_connector/unit/utils.py

@@ -548,6 +548,7 @@ def make_nixl_scheduler(

548 sched._reqs_need_save = {}

549 sched.use_host_buffer = False

550 sched.engine_id = "test-engine"

551 + sched.transfer_tp_size = 1

552 sched.side_channel_host = "localhost"

553 sched.side_channel_port = 5555

554 sched.blocks_per_sw = []

@@ -587,6 +588,7 @@ def make_nixl_push_scheduler(

588 sched.decoder_kv_blocks_ttl = decoder_kv_blocks_ttl

589 sched.use_host_buffer = False

590 sched.engine_id = "decode-engine"

591 + sched.transfer_tp_size = 1

592 sched.side_channel_host = "127.0.0.1"

593 sched.side_channel_port = 5600

594 sched.is_bidirectional_kv_xfer_enabled = is_bidirectional_kv_xfer_enabled

vllm/config/vllm.py

+6 -6

@@ -1329,14 +1329,14 @@ class VllmConfig:

1329 self.kv_transfer_config is not None

1330 and self.kv_transfer_config.has_connector("NixlConnector")

1331 ):

1332 - assert self.parallel_config.prefill_context_parallel_size == 1, (

1333 - "NIXL does not support prefill context parallelism."

1334 - )

1332 dcp_size = self.parallel_config.decode_context_parallel_size

1336 - tp_size = self.parallel_config.tensor_parallel_size

1337 - assert dcp_size in (1, tp_size), (

1333 + transfer_tp_size = max(

1334 + self.parallel_config.tensor_parallel_size,

1335 + self.parallel_config.prefill_context_parallel_size,

1336 + )

1337 + assert dcp_size in (1, transfer_tp_size), (

1338 f"decode_context_parallel_size={dcp_size} must be 1 or equal "

1339 - f"to tensor_parallel_size={tp_size} when using NixlConnector."

1339 + f"to the NIXL transfer parallel size={transfer_tp_size}."

1340 )

1341 if self.model_config is not None:

1342 assert self.model_config.use_mla or dcp_size == 1, (

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

@@ -63,6 +63,12 @@ class NixlBaseConnectorScheduler:

63 kv_cache_config: "KVCacheConfig",

64 ):

65 self.vllm_config = vllm_config

66 + parallel_config = vllm_config.parallel_config

67 + # TP1 PCP+DCP exposes its DCP shards as transfer ranks.

68 + self.transfer_tp_size = max(

69 + parallel_config.tensor_parallel_size,

70 + parallel_config.decode_context_parallel_size,

71 + )

72 self.block_size = vllm_config.cache_config.block_size

73 self.engine_id: EngineId = engine_id

74 self.kv_cache_config = kv_cache_config

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

+19 -9

@@ -548,9 +548,15 @@ class NixlBaseConnectorWorker:

548 self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size

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

550

551 - # DCP support is scoped to MLA, with dcp_size in (1, tp_size): either fully

552 - # replicated or fully sharded. A DCP rank is always derivable this way.

553 - self.dcp_rank = self.tp_rank % self.dcp_size

551 + # TP1 PCP+DCP owns distinct KV shards; replicated PCP uses rank zero.

552 + self.pcp_dcp_sharded = self.pcp_size > 1 and self.dcp_size > 1

553 + self.transfer_tp_size = (

554 + self.pcp_size if self.pcp_dcp_sharded else self.world_size

555 + )

556 + self.transfer_tp_rank = self.pcp_rank if self.pcp_dcp_sharded else self.tp_rank

557 +

558 + # MLA is either fully replicated or fully sharded across transfer ranks.

559 + self.dcp_rank = self.transfer_tp_rank % self.dcp_size

560

561 self.num_blocks = kv_cache_config.num_blocks

562 self.enable_permute_local_kv = False

@@ -781,12 +787,16 @@ class NixlBaseConnectorWorker:

787 local_dcp_size = self.dcp_size

788 remote_pcp_size = agent_metadata.pcp_size

789 remote_dcp_size = agent_metadata.dcp_size

784 - if (local_pcp_size > 1 and remote_dcp_size > 1) or (

785 - remote_pcp_size > 1 and local_dcp_size > 1

790 + if remote_pcp_size > 1 and remote_dcp_size not in (1, remote_pcp_size):

791 + raise NotImplementedError(

792 + "Remote NixlConnector PCP+DCP does not span the full PCP group. "

793 + f"Remote PCP/DCP={remote_pcp_size}/{remote_dcp_size}."

794 + )

795 + if (local_pcp_size > 1 and local_dcp_size == 1 and remote_dcp_size > 1) or (

796 + remote_pcp_size > 1 and remote_dcp_size == 1 and local_dcp_size > 1

797 ):

798 raise NotImplementedError(

788 - "NixlConnector PCP requires decode_context_parallel_size=1 "

789 - "on both instances. "

799 + "Replicated PCP cannot be paired with a DCP-sharded NIXL peer. "

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

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

802 )

@@ -1210,8 +1220,8 @@ class NixlBaseConnectorWorker:

1220 """Register the KV Cache data in nixl."""

1221

1222 self.transfer_topo = TransferTopology(

1213 - tp_rank=self.tp_rank,

1214 - tp_size=self.world_size,

1223 + tp_rank=self.transfer_tp_rank,

1224 + tp_size=self.transfer_tp_size,

1225 block_size=self.block_size,

1226 engine_id=self.engine_id,

1227 is_mla=self.use_mla,

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

+10 -3

@@ -108,10 +108,15 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):

108 "Consumers and kv_both require "

109 "prefill_context_parallel_size=1."

110 )

111 - if pcp_size > 1 and parallel_config.decode_context_parallel_size > 1:

111 + dcp_size = parallel_config.decode_context_parallel_size

112 + if (

113 + pcp_size > 1

114 + and dcp_size > 1

115 + and (parallel_config.tensor_parallel_size != 1 or dcp_size != pcp_size)

116 + ):

117 raise NotImplementedError(

113 - "NixlConnector PCP producers currently require "

114 - "decode_context_parallel_size=1."

118 + "NixlConnector PCP+DCP requires TP1 with DCP spanning "

119 + "the full PCP group."

120 )

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

122 # send and receive completion counts.

@@ -243,6 +248,7 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):

248 if (

249 self.kv_transfer_config.kv_role == "kv_producer"

250 and self.connector_worker.pcp_rank > 0

251 + and not self.connector_worker.pcp_dcp_sharded

252 ):

253 results.finished_sending.clear()

254 return results

@@ -324,6 +330,7 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):

330 if (

331 self.kv_transfer_config.kv_role == "kv_producer"

332 and self.connector_worker.pcp_rank > 0

333 + and not self.connector_worker.pcp_dcp_sharded

334 ):

335 return None

336 return self.connector_worker.xfer_handshake_metadata

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

+1 -1

@@ -282,7 +282,7 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):

282 remote_request_id=request.request_id,

283 remote_host=self.side_channel_host,

284 remote_port=self.side_channel_port,

285 - tp_size=self.vllm_config.parallel_config.tensor_parallel_size,

285 + tp_size=self.transfer_tp_size,

286 dcp_size=self.vllm_config.parallel_config.decode_context_parallel_size,

287 pp_size=self.vllm_config.parallel_config.pipeline_parallel_size,

288 remote_num_tokens=remote_num_tokens,

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

+1 -1

@@ -80,7 +80,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

80 while not self._ready_requests.empty():

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

83 - if self.pcp_rank > 0:

83 + if self.pcp_rank > 0 and not self.pcp_dcp_sharded:

84 # Replicated-KV PCP: only PCP rank 0 serves the KV, so this rank

85 # has nothing to send. Report the requests as sent right away so

86 # the scheduler-side aggregation (world_size workers, and any

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

+2 -1

@@ -286,7 +286,8 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):

286 remote_request_id=request.request_id,

287 remote_host=self.side_channel_host,

288 remote_port=self.side_channel_port,

289 - tp_size=self.vllm_config.parallel_config.tensor_parallel_size,

289 + tp_size=self.transfer_tp_size,

290 + dcp_size=self.vllm_config.parallel_config.decode_context_parallel_size,

291 pp_size=self.vllm_config.parallel_config.pipeline_parallel_size,

292 remote_num_tokens=remote_num_tokens,

293 transfer_mode=self._TRANSFER_MODE,

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

+1 -1

@@ -166,7 +166,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

166

167 def start_load_kv(self, metadata: NixlConnectorMetadata):

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

169 - if self.pcp_rank > 0:

169 + if self.pcp_rank > 0 and not self.pcp_dcp_sharded:

170 return

171

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

vllm/v1/worker/gpu_worker.py

@@ -45,6 +45,7 @@ from vllm.distributed.parallel_state import (

45 Handle,

46 checkpoint_prepare_distributed_state,

47 checkpoint_restore_distributed_state,

48 + get_pcp_group,

49 get_pp_group,

50 get_tp_group,

51 resume_device_comms,

@@ -703,6 +704,14 @@ class Worker(WorkerBase):

704

705 pp_rank = get_pp_group().rank_in_group

706 tp_rank = get_tp_group().rank_in_group

707 + parallel_config = self.vllm_config.parallel_config

708 + if (

709 + parallel_config.prefill_context_parallel_size > 1

710 + and parallel_config.decode_context_parallel_size > 1

711 + ):

712 + tp_rank += (

713 + get_pcp_group().rank_in_group * parallel_config.tensor_parallel_size

714 + )

715 return {(pp_rank, tp_rank): metadata}

716

717 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: