[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>
Yicong Wang committed
Sep 11, 2026 at 10:57 UTC
ce08bb5b3463fd423be90d5f22c2b41142834bf6
2 files changed
+85
-3
tests/v1/kv_connector/unit/test_mooncake_connector.py
+62
@@ -892,6 +892,68 @@ def test_resolve_need_send_accounts_for_remote_tp_fanout():
892
assert send_meta.need_send == 2
893
894
895
+@pytest.mark.asyncio
896
+async def test_heterogeneous_pp_waits_for_consumer_without_shared_layers():
897
+ """P4/D2 must retain KV until both D stages finish, including a no-op pull."""
898
+ config = create_vllm_config(kv_connector="MooncakeConnector", kv_role="kv_producer")
899
+ with (
900
+ set_current_vllm_config(config),
901
+ patch_worker_dependencies(),
902
+ patch.object(MooncakeConnectorWorker, "_sync_block_size_with_kernel"),
903
+ ):
904
+ worker = MooncakeConnector(
905
+ config, KVConnectorRole.WORKER, _make_test_kv_cache_config()
906
+ ).connector_worker
907
+ try:
908
+ worker.pp_size = 4
909
+ worker.kv_caches_base_addr = [0x1000]
910
+ worker.block_len_per_layer = [256]
911
+ worker.kv_block_len_per_layer = [256]
912
+ worker.registered_layer_names = ["model.layers.0.self_attn"]
913
+ worker.registered_layer_indices = [0]
914
+ send_meta = SendBlockMeta(
915
+ p_req_id="p-req",
916
+ transfer_id="transfer",
917
+ local_block_ids=[[1]],
918
+ ready=asyncio.Event(),
919
+ )
920
+ send_meta.ready.set()
921
+ worker.reqs_need_send["transfer"] = send_meta
922
+ sock = AsyncMock(spec=zmq.asyncio.Socket, send_multipart=AsyncMock())
923
+ with (
924
+ patch.object(worker, "sender_loop", asyncio.get_running_loop()),
925
+ patch.object(worker, "_send_blocks", return_value=0) as send,
926
+ ):
927
+ for pp_rank in range(2):
928
+ metadata = MooncakeXferMetadata(
929
+ remote_hostname="consumer",
930
+ remote_port=1234 + pp_rank,
931
+ remote_tp_size=1,
932
+ remote_tp_rank=0,
933
+ remote_pp_size=2,
934
+ req_blocks={"d-req": ("transfer", [[2]])},
935
+ kv_caches_base_addr=[0x2000],
936
+ block_lens=[256],
937
+ kv_block_lens=[256],
938
+ registered_layer_names=[f"model.layers.{pp_rank}.self_attn"],
939
+ registered_layer_indices=[pp_rank],
940
+ )
941
+ await worker.send_kv_to_decode(b"consumer", sock, metadata)
942
+ response = worker._xfer_resp_decoder.decode(
943
+ sock.send_multipart.call_args.args[0][1]
944
+ )
945
+ assert response.status == MooncakeXferResponseStatus.FINISH
946
+ assert response.ok_reqs == ["d-req"]
947
+ assert ("transfer" in worker.reqs_need_send) == (pp_rank == 0)
948
+ assert worker.finished_sending_reqs == (
949
+ set() if pp_rank == 0 else {"p-req"}
950
+ )
951
+ send.assert_called_once_with("consumer:1234", [0x1100], [0x2200], [256])
952
+ finally:
953
+ worker.shutdown()
954
+ worker.is_kv_consumer = True
955
+
956
+
957
@pytest.mark.asyncio
958
@patch(
959
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine",
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+23
-3
@@ -307,6 +307,8 @@ def _validate_asymmetric_region_lengths(
307
def _align_transfer_regions(
308
local_regions: list[TransferRegion],
309
remote_regions: list[TransferRegion],
310
+ *,
311
+ allow_partial_layers: bool = False,
312
) -> tuple[list[TransferRegion], list[TransferRegion], str | None]:
313
"""Align KV transfer regions by registered layer-name occurrence.
314
@@ -335,6 +337,11 @@ def _align_transfer_regions(
337
for key, local_region in local_keyed:
338
remote_region = remote_by_key.get(key)
339
if remote_region is None:
340
+ if (
341
+ allow_partial_layers
342
+ and (local_region.layer_name, 0) not in remote_by_key
343
+ ):
344
+ continue
345
return (
346
[],
347
[],
@@ -393,6 +400,7 @@ class MooncakeXferMetadata(
400
registered_layer_names: list[str] = msgspec.field(default_factory=list)
401
registered_layer_indices: list[int] = msgspec.field(default_factory=list)
402
registered_group_indices: list[int] = msgspec.field(default_factory=list)
403
+ remote_pp_size: int = 1
404
405
406
class MooncakeXferResponseStatus(IntEnum):
@@ -1232,7 +1240,11 @@ class MooncakeConnectorWorker:
1240
meta.registered_group_indices,
1241
)
1242
local_regions, remote_regions, align_err = _align_transfer_regions(
1235
- local_regions, remote_regions
1243
+ local_regions,
1244
+ remote_regions,
1245
+ allow_partial_layers=(
1246
+ meta.remote_pp_size > 1 and meta.remote_pp_size != self.pp_size
1247
+ ),
1248
)
1249
if align_err is not None:
1250
response = MooncakeXferResponse(
@@ -1322,7 +1334,9 @@ class MooncakeConnectorWorker:
1334
# Mark it sending to avoid expiration.
1335
send_meta.sending += 1
1336
if not send_meta.need_send:
1325
- self.resolve_need_send(send_meta, remote_tp_ranks)
1337
+ self.resolve_need_send(
1338
+ send_meta, remote_tp_ranks, meta.remote_pp_size
1339
+ )
1340
ready_reqs.append((d_req_id, send_meta))
1341
else:
1342
# Otherwise (expired, very unlikely), just forget it.
@@ -1398,11 +1412,16 @@ class MooncakeConnectorWorker:
1412
self,
1413
send_meta: SendBlockMeta,
1414
remote_tp_ranks: list[int],
1415
+ remote_pp_size: int = 1,
1416
):
1417
# Prepare for heterogeneous TP (one P pairs to multiple D)
1418
send_meta.need_send = len(remote_tp_ranks)
1419
+ if remote_pp_size > 1 and remote_pp_size != self.pp_size:
1420
+ # Each consumer PP stage pulls every producer stage, including
1421
+ # peers with no shared layers.
1422
+ send_meta.need_send *= remote_pp_size
1423
logger.debug(
1405
- "Mooncake request %s will be served by %d consumer TP workers: TP ranks=%s",
1424
+ "Mooncake request %s will be served by %d consumer workers: TP ranks=%s",
1425
send_meta.transfer_id,
1426
send_meta.need_send,
1427
remote_tp_ranks,
@@ -1851,6 +1870,7 @@ class MooncakeConnectorWorker:
1870
remote_port=self.rpc_port,
1871
remote_tp_size=self.tp_size,
1872
remote_tp_rank=self.tp_rank,
1873
+ remote_pp_size=self.pp_size,
1874
req_blocks={
1875
req_id: (pull_meta.transfer_id, pull_meta.local_block_ids)
1876
for req_id, pull_meta in pull_metas.items()