[Bugfix][Mooncake] Fix heterogeneous PP transfer completion (#56033)

Signed-off-by: wangyicong <wangyicong@bytedance.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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",

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()