From ce08bb5b3463fd423be90d5f22c2b41142834bf6 Mon Sep 17 00:00:00 2001 From: Yicong Wang Date: Fri, 11 Sep 2026 10:57:19 +0800 Subject: [PATCH] [Bugfix][Mooncake] Fix heterogeneous PP transfer completion (#56033) Signed-off-by: wangyicong Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../unit/test_mooncake_connector.py | 62 +++++++++++++++++++ .../v1/mooncake/mooncake_connector.py | 26 +++++++- 2 files changed, 85 insertions(+), 3 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index f989fb4e32..43de313821 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -892,6 +892,68 @@ def test_resolve_need_send_accounts_for_remote_tp_fanout(): assert send_meta.need_send == 2 +@pytest.mark.asyncio +async def test_heterogeneous_pp_waits_for_consumer_without_shared_layers(): + """P4/D2 must retain KV until both D stages finish, including a no-op pull.""" + config = create_vllm_config(kv_connector="MooncakeConnector", kv_role="kv_producer") + with ( + set_current_vllm_config(config), + patch_worker_dependencies(), + patch.object(MooncakeConnectorWorker, "_sync_block_size_with_kernel"), + ): + worker = MooncakeConnector( + config, KVConnectorRole.WORKER, _make_test_kv_cache_config() + ).connector_worker + try: + worker.pp_size = 4 + worker.kv_caches_base_addr = [0x1000] + worker.block_len_per_layer = [256] + worker.kv_block_len_per_layer = [256] + worker.registered_layer_names = ["model.layers.0.self_attn"] + worker.registered_layer_indices = [0] + send_meta = SendBlockMeta( + p_req_id="p-req", + transfer_id="transfer", + local_block_ids=[[1]], + ready=asyncio.Event(), + ) + send_meta.ready.set() + worker.reqs_need_send["transfer"] = send_meta + sock = AsyncMock(spec=zmq.asyncio.Socket, send_multipart=AsyncMock()) + with ( + patch.object(worker, "sender_loop", asyncio.get_running_loop()), + patch.object(worker, "_send_blocks", return_value=0) as send, + ): + for pp_rank in range(2): + metadata = MooncakeXferMetadata( + remote_hostname="consumer", + remote_port=1234 + pp_rank, + remote_tp_size=1, + remote_tp_rank=0, + remote_pp_size=2, + req_blocks={"d-req": ("transfer", [[2]])}, + kv_caches_base_addr=[0x2000], + block_lens=[256], + kv_block_lens=[256], + registered_layer_names=[f"model.layers.{pp_rank}.self_attn"], + registered_layer_indices=[pp_rank], + ) + await worker.send_kv_to_decode(b"consumer", sock, metadata) + response = worker._xfer_resp_decoder.decode( + sock.send_multipart.call_args.args[0][1] + ) + assert response.status == MooncakeXferResponseStatus.FINISH + assert response.ok_reqs == ["d-req"] + assert ("transfer" in worker.reqs_need_send) == (pp_rank == 0) + assert worker.finished_sending_reqs == ( + set() if pp_rank == 0 else {"p-req"} + ) + send.assert_called_once_with("consumer:1234", [0x1100], [0x2200], [256]) + finally: + worker.shutdown() + worker.is_kv_consumer = True + + @pytest.mark.asyncio @patch( "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index dc574a6219..0a899b7b62 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -307,6 +307,8 @@ def _validate_asymmetric_region_lengths( def _align_transfer_regions( local_regions: list[TransferRegion], remote_regions: list[TransferRegion], + *, + allow_partial_layers: bool = False, ) -> tuple[list[TransferRegion], list[TransferRegion], str | None]: """Align KV transfer regions by registered layer-name occurrence. @@ -335,6 +337,11 @@ def _align_transfer_regions( for key, local_region in local_keyed: remote_region = remote_by_key.get(key) if remote_region is None: + if ( + allow_partial_layers + and (local_region.layer_name, 0) not in remote_by_key + ): + continue return ( [], [], @@ -393,6 +400,7 @@ class MooncakeXferMetadata( registered_layer_names: list[str] = msgspec.field(default_factory=list) registered_layer_indices: list[int] = msgspec.field(default_factory=list) registered_group_indices: list[int] = msgspec.field(default_factory=list) + remote_pp_size: int = 1 class MooncakeXferResponseStatus(IntEnum): @@ -1232,7 +1240,11 @@ class MooncakeConnectorWorker: meta.registered_group_indices, ) local_regions, remote_regions, align_err = _align_transfer_regions( - local_regions, remote_regions + local_regions, + remote_regions, + allow_partial_layers=( + meta.remote_pp_size > 1 and meta.remote_pp_size != self.pp_size + ), ) if align_err is not None: response = MooncakeXferResponse( @@ -1322,7 +1334,9 @@ class MooncakeConnectorWorker: # Mark it sending to avoid expiration. send_meta.sending += 1 if not send_meta.need_send: - self.resolve_need_send(send_meta, remote_tp_ranks) + self.resolve_need_send( + send_meta, remote_tp_ranks, meta.remote_pp_size + ) ready_reqs.append((d_req_id, send_meta)) else: # Otherwise (expired, very unlikely), just forget it. @@ -1398,11 +1412,16 @@ class MooncakeConnectorWorker: self, send_meta: SendBlockMeta, remote_tp_ranks: list[int], + remote_pp_size: int = 1, ): # Prepare for heterogeneous TP (one P pairs to multiple D) send_meta.need_send = len(remote_tp_ranks) + if remote_pp_size > 1 and remote_pp_size != self.pp_size: + # Each consumer PP stage pulls every producer stage, including + # peers with no shared layers. + send_meta.need_send *= remote_pp_size logger.debug( - "Mooncake request %s will be served by %d consumer TP workers: TP ranks=%s", + "Mooncake request %s will be served by %d consumer workers: TP ranks=%s", send_meta.transfer_id, send_meta.need_send, remote_tp_ranks, @@ -1851,6 +1870,7 @@ class MooncakeConnectorWorker: remote_port=self.rpc_port, remote_tp_size=self.tp_size, remote_tp_rank=self.tp_rank, + remote_pp_size=self.pp_size, req_blocks={ req_id: (pull_meta.transfer_id, pull_meta.local_block_ids) for req_id, pull_meta in pull_metas.items()