From a7cb466f36e85e688005b8de0679f5b638002772 Mon Sep 17 00:00:00 2001 From: nwpu-zxr <978074293@qq.com> Date: Thu, 20 Aug 2026 00:19:28 +0800 Subject: [PATCH 1/3] add basic DCP support for mooncake_connector Signed-off-by: nwpu-zxr <978074293@qq.com> --- .../unit/test_mooncake_connector.py | 628 ++++++++++++++++++ .../v1/mooncake/mooncake_connector.py | 387 +++++++++-- .../v1/mooncake/mooncake_utils.py | 15 +- 3 files changed, 982 insertions(+), 48 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index afe626cba82b..7dab8dbb1434 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -42,6 +42,11 @@ _has_opaque_packed_row, _pp_mismatch_hides_packed_layers, _validate_asymmetric_region_lengths, + + + _filter_dcp_block_ids, + _get_decode_ranks_for_prefill, + _get_prefill_ranks_for_decode, get_mooncake_bootstrap_addr, should_launch_bootstrap_server, ) @@ -319,6 +324,193 @@ def test_align_transfer_regions_uses_layer_name_occurrences(): assert mismatch_err is not None +@pytest.mark.parametrize( + ( + "p_block_ids", + "d_block_ids", + "p_dcp_rank", + "p_dcp_size", + "d_dcp_rank", + "d_dcp_size", + "expected_p_blocks", + "expected_d_blocks", + ), + [ + # P splits a full D block range into finer DCP slices. + ([10], list(range(20, 28)), 0, 8, 0, 1, [10], [20]), + ([10], list(range(20, 28)), 7, 8, 0, 1, [10], [27]), + # P keeps full blocks while D asks for one finer DCP slice. + (list(range(10, 18)), [20], 0, 1, 4, 8, [14], [20]), + # Different DCP sizes: match by sequence-block position. + ([10, 11, 12, 13], [20, 21], 1, 4, 5, 8, [11, 13], [20, 21]), + ([10, 11], [20, 21, 22, 23], 5, 8, 1, 4, [10, 11], [21, 23]), + # Partial prefix hit: D only requests the tail sequence span. + ([10], [20, 21, 22, 23], 4, 8, 0, 1, [10], [20]), + ([10], [20, 21, 22, 23], 0, 8, 0, 1, [], []), + # Short P span is clipped to the overlapping sequence-block range. + ([10], list(range(20, 30)), 0, 8, 0, 1, [10], [22]), + ([10], [20], 1, 4, 1, 8, [], []), + ([10], [20], 1, 4, 5, 8, [10], [20]), + ], +) +def test_filter_dcp_block_ids( + p_block_ids: list[int], + d_block_ids: list[int], + p_dcp_rank: int, + p_dcp_size: int, + d_dcp_rank: int, + d_dcp_size: int, + expected_p_blocks: list[int], + expected_d_blocks: list[int], +): + actual_p_blocks, actual_d_blocks = _filter_dcp_block_ids( + p_block_ids=p_block_ids, + d_block_ids=d_block_ids, + p_dcp_rank=p_dcp_rank, + p_dcp_size=p_dcp_size, + d_dcp_rank=d_dcp_rank, + d_dcp_size=d_dcp_size, + ) + assert actual_p_blocks == expected_p_blocks + assert actual_d_blocks == expected_d_blocks + + +@pytest.mark.asyncio +async def test_build_transfer_params_filters_dcp_blocks(): + worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) + worker.tp_rank = 4 + worker.tp_size = 8 + worker.dcp_size = 8 + worker.use_mla = True + worker._total_num_kv_heads = 1 + worker.kv_cache_config = _make_test_kv_cache_config() + worker._physical_blocks_per_logical_kv_block = 1 + worker.transfer_topo = SimpleNamespace(local_replicates_kv_cache=False) + + block_len = 256 + local_regions = [ + TransferRegion( + layer_name="model.layers.0.self_attn", + layer_index=0, + base_addr=0x1000, + block_len=block_len, + kv_block_len=block_len, + ) + ] + remote_regions = [ + TransferRegion( + layer_name="model.layers.0.self_attn", + layer_index=0, + base_addr=0x2000, + block_len=block_len, + kv_block_len=block_len, + ) + ] + send_meta = SendBlockMeta( + p_req_id="p-req-dcp-filter", + transfer_id="xfer-dcp-filter", + local_block_ids=[[10]], + ready=asyncio.Event(), + ) + xfer_meta = MooncakeXferMetadata( + remote_hostname="consumer-host", + remote_port=54321, + remote_tp_size=1, + remote_tp_rank=0, + req_blocks={"d-req-dcp-filter": ("xfer-dcp-filter", [[20, 21, 22, 23]])}, + kv_caches_base_addr=[0x2000], + block_lens=[block_len], + kv_block_lens=[block_len], + registered_layer_names=["model.layers.0.self_attn"], + registered_layer_indices=[0], + remote_dcp_size=1, + remote_block_size=16, + ) + + src_ptrs, dst_ptrs, lengths, err_reqs, err_msg = ( + await worker._build_transfer_params( + ready_reqs=[("d-req-dcp-filter", send_meta)], + agent_meta=xfer_meta, + local_regions=local_regions, + remote_regions=remote_regions, + ) + ) + + assert err_reqs == [] + assert err_msg is None + assert src_ptrs == [0x1000 + 10 * block_len] + assert dst_ptrs == [0x2000 + 20 * block_len] + assert lengths == [block_len] + + +@pytest.mark.asyncio +async def test_build_transfer_params_returns_empty_for_unowned_dcp_blocks(): + worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) + worker.tp_rank = 0 + worker.tp_size = 8 + worker.dcp_size = 8 + worker.use_mla = True + worker._total_num_kv_heads = 1 + worker.kv_cache_config = _make_test_kv_cache_config() + worker._physical_blocks_per_logical_kv_block = 1 + worker.transfer_topo = SimpleNamespace(local_replicates_kv_cache=False) + + block_len = 256 + local_regions = [ + TransferRegion( + layer_name="model.layers.0.self_attn", + layer_index=0, + base_addr=0x1000, + block_len=block_len, + kv_block_len=block_len, + ) + ] + remote_regions = [ + TransferRegion( + layer_name="model.layers.0.self_attn", + layer_index=0, + base_addr=0x2000, + block_len=block_len, + kv_block_len=block_len, + ) + ] + send_meta = SendBlockMeta( + p_req_id="p-req-dcp-empty", + transfer_id="xfer-dcp-empty", + local_block_ids=[[10]], + ready=asyncio.Event(), + ) + xfer_meta = MooncakeXferMetadata( + remote_hostname="consumer-host", + remote_port=54321, + remote_tp_size=1, + remote_tp_rank=0, + req_blocks={"d-req-dcp-empty": ("xfer-dcp-empty", [[20, 21, 22, 23]])}, + kv_caches_base_addr=[0x2000], + block_lens=[block_len], + kv_block_lens=[block_len], + registered_layer_names=["model.layers.0.self_attn"], + registered_layer_indices=[0], + remote_dcp_size=1, + remote_block_size=16, + ) + + src_ptrs, dst_ptrs, lengths, err_reqs, err_msg = ( + await worker._build_transfer_params( + ready_reqs=[("d-req-dcp-empty", send_meta)], + agent_meta=xfer_meta, + local_regions=local_regions, + remote_regions=remote_regions, + ) + ) + + assert src_ptrs == [] + assert dst_ptrs == [] + assert lengths == [] + assert err_reqs == [] + assert err_msg is None + + @pytest.mark.asyncio async def test_build_transfer_params_separates_prefill_pp_layers(): """Each producer PP stage should send only its registered layer shard.""" @@ -771,6 +963,7 @@ async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer): "tp_rank": 0, "pp_rank": 0, "addr": "tcp://1.1.1.1:1111", + "dcp_size": 2, } async with httpx.AsyncClient() as client: response = await client.post(f"{base_url}/register", json=payload1) @@ -783,6 +976,7 @@ async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer): "tp_rank": 0, "pp_rank": 1, "addr": "tcp://2.2.2.2:2222", + "dcp_size": 2, } async with httpx.AsyncClient() as client: response = await client.post(f"{base_url}/register", json=payload2) @@ -796,6 +990,7 @@ async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer): data = response.json() assert "0" in data assert data["0"]["engine_id"] == "eng-1" + assert data["0"]["dcp_size"] == 2 assert data["0"]["worker_addr"]["0"]["0"] == "tcp://1.1.1.1:1111" assert data["0"]["worker_addr"]["0"]["1"] == "tcp://2.2.2.2:2222" @@ -2813,3 +3008,436 @@ def test_sleep_unregisters_and_restores_kv_caches(): worker.restore_kv_caches() worker.restore_kv_caches() reg.assert_called_once_with([raw.data_ptr()], [raw.nbytes]) + + +@pytest.mark.asyncio +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake." + "mooncake_connector.TransferEngine", + FakeMooncakeWrapper, +) +async def test_kv_producer_dcp_uses_effective_tp_for_offsets(monkeypatch): + """DCP transfer planning should slice by TP//DCP, not physical TP.""" + + P_TP_SIZE = 4 + P_TP_RANK = 2 + P_DCP_SIZE = 2 + D_TP_SIZE = 1 + D_TP_RANK = 0 + D_DCP_SIZE = 1 + LOCAL_BLOCK_LEN = 4096 + REMOTE_BLOCK_LEN = LOCAL_BLOCK_LEN * 2 + + monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") + vllm_config = create_vllm_config( + kv_connector="MooncakeConnector", kv_role="kv_producer" + ) + + with set_current_vllm_config(vllm_config), patch_worker_dependencies(): + prefill_connector = MooncakeConnector( + vllm_config, + KVConnectorRole.WORKER, + _make_test_kv_cache_config(), + ) + prefill_worker = prefill_connector.connector_worker + + prefill_worker.tp_rank = P_TP_RANK + prefill_worker.tp_size = P_TP_SIZE + prefill_worker.dcp_size = P_DCP_SIZE + prefill_worker.use_mla = False + prefill_worker._total_num_kv_heads = 2 + prefill_worker._tp_size[prefill_worker.engine_id] = P_TP_SIZE + prefill_worker._dcp_size[prefill_worker.engine_id] = P_DCP_SIZE + prefill_worker.transfer_topo.tp_rank = P_TP_RANK + prefill_worker.transfer_topo.tp_size = P_TP_SIZE + + prefill_worker.block_size = 16 + prefill_worker.kv_caches_base_addr = [0x1000] + prefill_worker.block_len_per_layer = [LOCAL_BLOCK_LEN] + prefill_worker.kv_block_len_per_layer = [LOCAL_BLOCK_LEN] + prefill_worker.registered_layer_names = ["model.layers.0.self_attn"] + prefill_worker.registered_layer_indices = [0] + + origin_sender_loop = prefill_worker.sender_loop + prefill_worker.sender_loop = asyncio.get_event_loop() + + transfer_id = "xfer-dcp-1" + local_block_ids = [[10]] + send_meta = SendBlockMeta( + p_req_id="p-req-dcp-1", + transfer_id=transfer_id, + local_block_ids=local_block_ids, + ready=asyncio.Event(), + ) + prefill_worker.reqs_need_send[transfer_id] = send_meta + send_meta.ready.set() + + xfer_meta = MooncakeXferMetadata( + remote_hostname="consumer-host", + remote_port=54321, + remote_tp_size=D_TP_SIZE, + remote_tp_rank=D_TP_RANK, + req_blocks={"d-req-dcp-1": (transfer_id, [[20, 21]])}, + kv_caches_base_addr=[0x2000], + block_lens=[REMOTE_BLOCK_LEN], + kv_block_lens=[REMOTE_BLOCK_LEN], + registered_layer_names=["model.layers.0.self_attn"], + registered_layer_indices=[0], + remote_dcp_size=D_DCP_SIZE, + remote_block_size=prefill_worker.block_size, + ) + + mock_socket = AsyncMock(spec=zmq.asyncio.Socket) + mock_socket.send_multipart = AsyncMock() + identity = b"consumer-dcp" + + with patch.object(prefill_worker, "_send_blocks", return_value=0) as mock_send: + await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta) + + mock_send.assert_called_once() + _, src_ptrs, dst_ptrs, lengths = mock_send.call_args[0] + assert src_ptrs == [0x1000 + 10 * LOCAL_BLOCK_LEN] + assert dst_ptrs == [0x2000 + 20 * REMOTE_BLOCK_LEN + LOCAL_BLOCK_LEN] + assert lengths == [LOCAL_BLOCK_LEN] + + mock_socket.send_multipart.assert_called_once() + _, sent_payload = mock_socket.send_multipart.call_args[0][0] + response = prefill_worker._xfer_resp_decoder.decode(sent_payload) + assert response.status == MooncakeXferResponseStatus.FINISH + assert response.ok_reqs == ["d-req-dcp-1"] + + prefill_worker.sender_loop = origin_sender_loop + prefill_worker.shutdown() + + +def test_get_prefill_ranks_for_decode_dcp_ratios(): + """DCP rank routing should follow modulo mapping up to dcp16.""" + + def route(d_dcp_size: int, p_dcp_size: int, rank: int = 0): + return _get_prefill_ranks_for_decode( + d_tp_rank=rank, + d_tp_size=d_dcp_size, + d_dcp_size=d_dcp_size, + p_tp_size=p_dcp_size, + p_dcp_size=p_dcp_size, + total_num_kv_heads=1, + is_mla=True, + ) + + assert route(1, 8) == list(range(8)) + assert route(8, 1, rank=0) == [0] + assert route(8, 1, rank=7) == [0] + assert route(4, 8, rank=0) == [0, 4] + assert route(4, 8, rank=3) == [3, 7] + assert route(8, 4, rank=0) == [0] + assert route(8, 4, rank=7) == [3] + assert route(1, 16) == list(range(16)) + assert route(16, 1, rank=15) == [0] + + +@pytest.mark.parametrize( + ("d_tp_rank", "d_tp_size", "d_dcp_size", "p_tp_size", "p_dcp_size", "expected"), + [ + (0, 1, 1, 8, 8, list(range(8))), + (7, 8, 8, 1, 1, [0]), + (0, 4, 4, 8, 8, [0, 4]), + (3, 4, 4, 8, 8, [3, 7]), + (0, 8, 8, 4, 4, [0]), + (7, 8, 8, 4, 4, [3]), + ], +) +def test_get_prefill_ranks_for_decode_mla_cases( + d_tp_rank: int, + d_tp_size: int, + d_dcp_size: int, + p_tp_size: int, + p_dcp_size: int, + expected: list[int], +): + assert _get_prefill_ranks_for_decode( + d_tp_rank=d_tp_rank, + d_tp_size=d_tp_size, + d_dcp_size=d_dcp_size, + p_tp_size=p_tp_size, + p_dcp_size=p_dcp_size, + total_num_kv_heads=1, + is_mla=True, + ) == expected + + +@pytest.mark.parametrize( + ( + "d_tp_rank", + "d_tp_size", + "d_dcp_size", + "p_tp_size", + "p_dcp_size", + "num_kv_heads", + "expected", + ), + [ + # Same TP/head groups: D keeps full blocks, P splits each head group by DCP. + (0, 16, 1, 16, 4, 4, [0, 1, 2, 3]), + (3, 16, 1, 16, 4, 4, [0, 1, 2, 3]), + (4, 16, 1, 16, 4, 4, [4, 5, 6, 7]), + # Same TP/head groups: D splits by DCP, P keeps one representative copy. + (5, 16, 4, 16, 1, 4, [4]), + # Same TP/head groups and same DCP: point-to-point by dcp rank. + (6, 16, 4, 16, 4, 4, [6]), + # P has a larger DCP size: one D slice maps to multiple finer P slices. + (5, 16, 4, 32, 8, 4, [9, 13]), + # D has a larger DCP size: multiple D slices map to the same coarser P slice. + (13, 32, 8, 16, 4, 4, [5]), + # TP does not fully split KV heads on D; one D rank maps to + # multiple P head groups. + (0, 4, 1, 16, 2, 8, [0, 1, 2, 3]), + (3, 4, 1, 16, 2, 8, [12, 13, 14, 15]), + # D has more KV-head groups than P; adjacent D groups share one P group. + (0, 8, 2, 2, 1, 4, [0]), + (2, 8, 2, 2, 1, 4, [0]), + (4, 8, 2, 2, 1, 4, [1]), + ], +) +def test_get_prefill_ranks_for_decode_gqa_and_heterogeneous_tp( + d_tp_rank: int, + d_tp_size: int, + d_dcp_size: int, + p_tp_size: int, + p_dcp_size: int, + num_kv_heads: int, + expected: list[int], +): + assert _get_prefill_ranks_for_decode( + d_tp_rank=d_tp_rank, + d_tp_size=d_tp_size, + d_dcp_size=d_dcp_size, + p_tp_size=p_tp_size, + p_dcp_size=p_dcp_size, + total_num_kv_heads=num_kv_heads, + is_mla=False, + ) == expected + + +@pytest.mark.parametrize( + ( + "p_tp_rank", + "p_tp_size", + "p_dcp_size", + "d_tp_size", + "d_dcp_size", + "num_kv_heads", + "is_mla", + "expected", + ), + [ + # D keeps full MLA blocks, P splits them across all DCP ranks. + (0, 8, 8, 1, 1, 1, True, [0]), + (7, 8, 8, 1, 1, 1, True, [0]), + # P keeps a full MLA block, all D DCP ranks query the same P rank. + (0, 1, 1, 8, 8, 1, True, list(range(8))), + # GQA same TP/head groups: each P DCP rank serves all D replicas + # of that KV-head group when D does not use DCP. + (0, 16, 4, 16, 1, 4, False, [0, 1, 2, 3]), + (4, 16, 4, 16, 1, 4, False, [4, 5, 6, 7]), + # GQA DCP size differs: finer P slices can serve the same D rank. + (9, 32, 8, 16, 4, 4, False, [5]), + (13, 32, 8, 16, 4, 4, False, [5]), + # Coarser P slices can be queried by multiple finer D ranks. + (5, 16, 4, 32, 8, 4, False, [9, 13]), + # D TP does not fully split KV heads; one D rank maps to multiple + # P KV-head groups and each matching P rank expects that D rank. + (0, 16, 2, 4, 1, 8, False, [0]), + (3, 16, 2, 4, 1, 8, False, [0]), + (12, 16, 2, 4, 1, 8, False, [3]), + ], +) +def test_get_decode_ranks_for_prefill_inverse_mapping( + p_tp_rank: int, + p_tp_size: int, + p_dcp_size: int, + d_tp_size: int, + d_dcp_size: int, + num_kv_heads: int, + is_mla: bool, + expected: list[int], +): + assert _get_decode_ranks_for_prefill( + p_tp_rank=p_tp_rank, + p_tp_size=p_tp_size, + p_dcp_size=p_dcp_size, + d_tp_size=d_tp_size, + d_dcp_size=d_dcp_size, + total_num_kv_heads=num_kv_heads, + is_mla=is_mla, + ) == expected + + +def test_process_pulling_result_waits_for_all_dcp_acks(): + worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) + worker.finished_recving_reqs = set() + pull_meta = PullReqMeta( + d_req_id="d-req-dcp-acks", + transfer_id="xfer-dcp-acks", + local_block_ids=[[0]], + remote_engine_id="p-engine", + remote_bootstrap_addr="http://bootstrap:33333", + pull_tasks_count=2, + ) + pull_metas = {pull_meta.d_req_id: pull_meta} + + response = MooncakeXferResponse( + status=MooncakeXferResponseStatus.CONTINUE, + ok_reqs=[pull_meta.d_req_id], + ) + worker.process_pulling_result(response, pull_metas) + assert pull_meta.pull_tasks_count == 1 + assert pull_meta.d_req_id not in worker.finished_recving_reqs + + response = MooncakeXferResponse( + status=MooncakeXferResponseStatus.FINISH, + ok_reqs=[pull_meta.d_req_id], + ) + worker.process_pulling_result(response, pull_metas) + assert pull_meta.pull_tasks_count == 0 + assert pull_meta.d_req_id in worker.finished_recving_reqs + + +@pytest.mark.asyncio +async def test_receive_kv_selects_remote_dcp_peer_ranks(): + """Decode should query the producer DCP representative rank.""" + + vllm_config = create_vllm_config( + kv_connector="MooncakeConnector", kv_role="kv_consumer" + ) + + with set_current_vllm_config(vllm_config), patch_worker_dependencies(): + decode_connector = MooncakeConnector( + vllm_config, + KVConnectorRole.WORKER, + _make_test_kv_cache_config(), + ) + decode_worker = decode_connector.connector_worker + decode_worker.use_mla = True + decode_worker.tp_rank = 1 + decode_worker.tp_size = 2 + decode_worker.dcp_size = 2 + decode_worker.pp_size = 1 + decode_worker.pp_rank = 0 + decode_worker._remote_agents = { + "p-engine": { + rank: {0: f"tcp://producer-tp{rank}:1234"} + for rank in range(4) + } + } + decode_worker._tp_size["p-engine"] = 4 + decode_worker._dcp_size["p-engine"] = 2 + + pull_metas = { + "d-req-dcp-rank": PullReqMeta( + d_req_id="d-req-dcp-rank", + transfer_id="xfer-dcp-rank", + local_block_ids=[[100, 101]], + remote_engine_id="p-engine", + remote_bootstrap_addr="http://bootstrap:33333", + ) + } + seen_addrs: list[str] = [] + + async def fake_receive(worker_addr: str, metas: dict[str, PullReqMeta]): + seen_addrs.append(worker_addr) + for meta in metas.values(): + meta.pull_tasks_count -= 1 + + with patch.object( + decode_worker, + "receive_kv_from_single_worker", + side_effect=fake_receive, + ): + decode_worker.receive_kv("p-engine", pull_metas) + await asyncio.sleep(0) + + assert seen_addrs == ["tcp://producer-tp1:1234"] + assert pull_metas["d-req-dcp-rank"].pull_tasks_count == 0 + + +@pytest.mark.asyncio +async def test_receive_kv_routes_metadata_by_dcp_rank_modulo(): + """DCP metadata routing follows modulo dcp-rank mapping.""" + + vllm_config = create_vllm_config( + kv_connector="MooncakeConnector", kv_role="kv_consumer" + ) + + with set_current_vllm_config(vllm_config), patch_worker_dependencies(): + decode_connector = MooncakeConnector( + vllm_config, + KVConnectorRole.WORKER, + _make_test_kv_cache_config(), + ) + decode_worker = decode_connector.connector_worker + decode_worker.use_mla = True + decode_worker.tp_rank = 0 + decode_worker.tp_size = 2 + decode_worker.dcp_size = 2 + decode_worker.pp_size = 1 + decode_worker.pp_rank = 0 + decode_worker._remote_agents = { + "p-engine": { + rank: {0: f"tcp://producer-tp{rank}:1234"} + for rank in range(4) + } + } + decode_worker._tp_size["p-engine"] = 4 + decode_worker._dcp_size["p-engine"] = 4 + + pull_metas = { + "d-req-block0": PullReqMeta( + d_req_id="d-req-block0", + transfer_id="xfer-block0", + local_block_ids=[[0]], + remote_engine_id="p-engine", + remote_bootstrap_addr="http://bootstrap:33333", + ), + "d-req-block1": PullReqMeta( + d_req_id="d-req-block1", + transfer_id="xfer-block1", + local_block_ids=[[1]], + remote_engine_id="p-engine", + remote_bootstrap_addr="http://bootstrap:33333", + ), + "d-req-both": PullReqMeta( + d_req_id="d-req-both", + transfer_id="xfer-both", + local_block_ids=[[0, 1]], + remote_engine_id="p-engine", + remote_bootstrap_addr="http://bootstrap:33333", + ), + } + seen: dict[str, set[str]] = {} + + async def fake_receive(worker_addr: str, metas: dict[str, PullReqMeta]): + seen[worker_addr] = set(metas) + for meta in metas.values(): + meta.pull_tasks_count -= 1 + + with patch.object( + decode_worker, + "receive_kv_from_single_worker", + side_effect=fake_receive, + ): + decode_worker.receive_kv("p-engine", pull_metas) + await asyncio.sleep(0) + + assert seen == { + "tcp://producer-tp0:1234": { + "d-req-block0", + "d-req-block1", + "d-req-both", + }, + "tcp://producer-tp2:1234": { + "d-req-block0", + "d-req-block1", + "d-req-both", + }, + } + assert all(meta.pull_tasks_count == 0 for meta in pull_metas.values()) 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 1d12bd1ab19a..43eba5b7c9bb 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 @@ -144,6 +144,137 @@ def _block_ids_for_region( return list(block_ids_by_group[group_index]) +def _validate_dcp_size_for_kv_head_groups( + tp_size: int, + dcp_size: int, + total_num_kv_heads: int, + is_mla: bool, + role: str, +) -> None: + if dcp_size == 1: + return + kv_head_group_count = 1 if is_mla else max(1, min(tp_size, total_num_kv_heads)) + expected_dcp_size = tp_size // kv_head_group_count + if dcp_size != expected_dcp_size: + raise ValueError( + f"{role} DCP size must consume all redundant KV-cache replicas: " + f"dcp_size={dcp_size}, expected={expected_dcp_size}, " + f"tp_size={tp_size}, kv_head_group_count={kv_head_group_count}." + ) + + + + +def _get_kv_head_group_rank_and_count( + tp_rank: int, + tp_size: int, + total_num_kv_heads: int, + is_mla: bool, +) -> tuple[int, int]: + kv_head_count = 1 if is_mla else total_num_kv_heads + kv_head_group_count = max(1, min(tp_size, kv_head_count)) + ranks_per_kv_head_group = tp_size // kv_head_group_count + return tp_rank // ranks_per_kv_head_group, kv_head_group_count + + +def _get_prefill_ranks_for_decode( + d_tp_rank: int, + d_tp_size: int, + d_dcp_size: int, + p_tp_size: int, + p_dcp_size: int, + total_num_kv_heads: int, + is_mla: bool, +) -> list[int]: + d_kv_head_group_rank, d_kv_head_group_count = ( + _get_kv_head_group_rank_and_count( + d_tp_rank, d_tp_size, total_num_kv_heads, is_mla + ) + ) + _, p_kv_head_group_count = _get_kv_head_group_rank_and_count( + 0, p_tp_size, total_num_kv_heads, is_mla + ) + p_ranks_per_kv_head_group = p_tp_size // p_kv_head_group_count + + if d_kv_head_group_count >= p_kv_head_group_count: + head_group_ratio = d_kv_head_group_count // p_kv_head_group_count + p_kv_head_group_ranks = [d_kv_head_group_rank // head_group_ratio] + else: + head_group_ratio = p_kv_head_group_count // d_kv_head_group_count + p_start_head_group_rank = d_kv_head_group_rank * head_group_ratio + p_kv_head_group_ranks = list( + range( + p_start_head_group_rank, + p_start_head_group_rank + head_group_ratio, + ) + ) + + d_dcp_rank = d_tp_rank % d_dcp_size + if p_dcp_size >= d_dcp_size: + p_dcp_ranks = list(range(d_dcp_rank, p_dcp_size, d_dcp_size)) + else: + p_dcp_ranks = [d_dcp_rank % p_dcp_size] + + return sorted( + p_kv_head_group_rank * p_ranks_per_kv_head_group + p_dcp_rank + for p_kv_head_group_rank in p_kv_head_group_ranks + for p_dcp_rank in p_dcp_ranks + ) + + +def _get_decode_ranks_for_prefill( + p_tp_rank: int, + p_tp_size: int, + p_dcp_size: int, + d_tp_size: int, + d_dcp_size: int, + total_num_kv_heads: int, + is_mla: bool, +) -> list[int]: + d_tp_ranks: list[int] = [] + for d_tp_rank in range(d_tp_size): + p_tp_ranks = _get_prefill_ranks_for_decode( + d_tp_rank=d_tp_rank, + d_tp_size=d_tp_size, + d_dcp_size=d_dcp_size, + p_tp_size=p_tp_size, + p_dcp_size=p_dcp_size, + total_num_kv_heads=total_num_kv_heads, + is_mla=is_mla, + ) + if p_tp_rank in p_tp_ranks: + d_tp_ranks.append(d_tp_rank) + return d_tp_ranks + + +def _filter_dcp_block_ids( + p_block_ids: list[int], + d_block_ids: list[int], + p_dcp_rank: int, + p_dcp_size: int, + d_dcp_rank: int, + d_dcp_size: int, +) -> tuple[list[int], list[int]]: + p_sequence_span = len(p_block_ids) * p_dcp_size + d_sequence_span = len(d_block_ids) * d_dcp_size + d_start_sequence_block = p_sequence_span - d_sequence_span + d_sequence_block_to_cache_block = { + d_start_sequence_block + i * d_dcp_size + d_dcp_rank: block_id + for i, block_id in enumerate(d_block_ids) + } + + filtered_p_block_ids: list[int] = [] + filtered_d_block_ids: list[int] = [] + for i, p_block_id in enumerate(p_block_ids): + p_sequence_block = i * p_dcp_size + p_dcp_rank + d_block_id = d_sequence_block_to_cache_block.get(p_sequence_block) + if d_block_id is not None: + filtered_p_block_ids.append(p_block_id) + filtered_d_block_ids.append(d_block_id) + + return filtered_p_block_ids, filtered_d_block_ids + + def _get_tp_ratio(local_tp_size: int, remote_tp_size: int) -> int: """Return the TP ratio used by heterogeneous TP transfer planning. @@ -630,6 +761,8 @@ class MooncakeXferMetadata( # Random per register_kv_caches, so ids never repeat across restarts. # 0 means the peer sends the layout on every pull. layout_id: int = 0 + remote_dcp_size: int = 1 + remote_block_size: int = 0 class MooncakeXferResponseStatus(IntEnum): @@ -1297,6 +1430,7 @@ def __init__( assert (parallel_config := vllm_config.parallel_config) dp_rank = parallel_config.data_parallel_index dp_local_rank = parallel_config.data_parallel_rank_local + self.dcp_size = parallel_config.decode_context_parallel_size or 1 self.dp_rank = dp_local_rank if parallel_config.local_engines_only else dp_rank self.pp_size = vllm_config.parallel_config.pipeline_parallel_size self.pp_rank = get_pp_group().rank_in_group @@ -1358,6 +1492,24 @@ def __init__( self.cache_config = vllm_config.cache_config self.kv_cache_config = kv_cache_config self.use_mla = self.model_config.use_mla + self._total_num_kv_heads = 1 if self.use_mla else ( + self.model_config.get_total_num_kv_heads() + ) + if self.dcp_size > 1 and len(kv_cache_config.kv_cache_groups) != 1: + msg = ( + "Mooncake DCP currently supports exactly one KV cache group; " + f"got {len(kv_cache_config.kv_cache_groups)} groups. " + "Hybrid/group-aware DCP transfer is not supported yet." + ) + logger.error(msg) + raise ValueError(msg) + _validate_dcp_size_for_kv_head_groups( + tp_size=self.tp_size, + dcp_size=self.dcp_size, + total_num_kv_heads=self._total_num_kv_heads, + is_mla=self.use_mla, + role="Local", + ) self._physical_blocks_per_logical_kv_block = 1 self._sync_block_size_with_kernel() @@ -1368,14 +1520,29 @@ def __init__( ) self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size} + self._dcp_size: dict[EngineId, int] = {self.engine_id: self.dcp_size} self._layer_specs: dict[str, KVCacheSpec] = {} for group in kv_cache_config.transfer_groups: group_spec = group.kv_cache_spec specs_by_layer = getattr(group_spec, "kv_cache_specs", {}) for layer_name in group.layer_names: - self._layer_specs[layer_name] = specs_by_layer.get( - layer_name, group_spec - ) + layer_spec = specs_by_layer.get(layer_name, group_spec) + if self.dcp_size > 1: + is_swa = isinstance(layer_spec, SlidingWindowSpec) or ( + isinstance(layer_spec, FullAttentionSpec) + and getattr(layer_spec, "sliding_window", None) is not None + ) + if isinstance(layer_spec, MambaSpec) or is_swa: + msg = ( + "Mooncake DCP currently supports only single-group " + "GQA/full-attention or MLA KV cache. Mamba/GDN and " + "sliding-window attention are not supported by this " + f"DCP path; found {type(layer_spec).__name__} " + f"at layer {layer_name}." + ) + logger.error(msg) + raise ValueError(msg) + self._layer_specs[layer_name] = layer_spec self.transfer_topo = TransferTopology( tp_rank=self.tp_rank, tp_size=self.tp_size, @@ -1444,6 +1611,7 @@ async def register_worker_with_bootstrap(self): tp_rank=self.tp_rank, pp_rank=self.pp_rank, addr=worker_addr, + dcp_size=self.dcp_size, ) timeout = envs.VLLM_MOONCAKE_CONNECTOR_TIMEOUT max_attempts = _BOOTSTRAP_MAX_ATTEMPTS @@ -1547,13 +1715,43 @@ async def send_kv_to_decode( self, identity: bytes, sock: zmq.asyncio.Socket, meta: MooncakeXferMetadata ): pending_reqs: dict[ReqId, SendBlockMeta] = {} - remote_tp_ranks = self.transfer_topo.handshake_target_ranks(meta.remote_tp_size) - if meta.remote_tp_rank not in remote_tp_ranks: + if meta.remote_block_size and meta.remote_block_size != self.block_size: + msg = ( + "Mooncake DCP transfer currently requires matching block sizes: " + f"local={self.block_size}, remote={meta.remote_block_size}." + ) + logger.error(msg) + response = MooncakeXferResponse( + status=MooncakeXferResponseStatus.ERROR, + err_msg=msg, + ) + await sock.send_multipart((identity, self._encoder.encode(response))) + return + no_dcp = self.dcp_size == 1 and meta.remote_dcp_size == 1 + if no_dcp: + remote_tp_ranks = self.transfer_topo.handshake_target_ranks( + meta.remote_tp_size + ) + should_handle_request = meta.remote_tp_rank in remote_tp_ranks + expected_peer_ranks = remote_tp_ranks + else: + remote_tp_ranks = _get_decode_ranks_for_prefill( + p_tp_rank=self.tp_rank, + p_tp_size=self.tp_size, + p_dcp_size=self.dcp_size, + d_tp_size=meta.remote_tp_size, + d_dcp_size=meta.remote_dcp_size, + total_num_kv_heads=self._total_num_kv_heads, + is_mla=self.use_mla, + ) + should_handle_request = meta.remote_tp_rank in remote_tp_ranks + expected_peer_ranks = remote_tp_ranks + if not should_handle_request: # This D worker does not pair with the P worker. msg = ( "This D tp_rank " f"{meta.remote_tp_rank} is not paired with P tp_rank " - f"{self.tp_rank}; expected one of {remote_tp_ranks}." + f"{self.tp_rank}; expected peer ranks {expected_peer_ranks}." ) logger.error(msg) response = MooncakeXferResponse( @@ -1606,12 +1804,30 @@ async def send_kv_to_decode( ) await sock.send_multipart((identity, self._encoder.encode(response))) return + if no_dcp: + local_transfer_tp_size = self.tp_size + remote_transfer_tp_size = meta.remote_tp_size + producer_cache_replicated = self._producer_cache_is_replicated() + else: + _, local_transfer_tp_size = _get_kv_head_group_rank_and_count( + self.tp_rank, + self.tp_size, + self._total_num_kv_heads, + self.use_mla, + ) + _, remote_transfer_tp_size = _get_kv_head_group_rank_and_count( + meta.remote_tp_rank, + meta.remote_tp_size, + self._total_num_kv_heads, + self.use_mla, + ) + producer_cache_replicated = False validation_err = _validate_asymmetric_region_lengths( local_regions=local_regions, remote_regions=remote_regions, - local_tp_size=self.tp_size, - remote_tp_size=meta.remote_tp_size, - producer_cache_replicated=self._producer_cache_is_replicated(), + local_tp_size=local_transfer_tp_size, + remote_tp_size=remote_transfer_tp_size, + producer_cache_replicated=producer_cache_replicated, total_num_kv_heads=( None if self.use_mla or self.kv_cache_config.has_mamba_layers @@ -1848,6 +2064,9 @@ async def _build_transfer_params( local_block_ids_by_group: list[list[int]] = [] remote_block_ids_by_group: list[list[int]] = [] has_block_error = False + uses_dcp = self.dcp_size > 1 or agent_meta.remote_dcp_size > 1 + p_dcp_rank = self.tp_rank % self.dcp_size + d_dcp_rank = agent_meta.remote_tp_rank % agent_meta.remote_dcp_size group_specs = self.kv_cache_config.transfer_groups for group_index, (local_group, remote_group) in enumerate( zip(send_meta.local_block_ids, remote_block_ids_per_group) @@ -1871,22 +2090,40 @@ async def _build_transfer_params( if block_id != NULL_BLOCK_ID ] - n_local = len(local_group) - n_remote = len(remote_group) - if n_local < n_remote: - logger.error( - "req %s: local blocks(%d) < remote blocks(%d) " - "in a KV cache group (is_mamba_group=%s)", - d_req_id, - n_local, - n_remote, - is_mamba_group, + if uses_dcp: + if is_mamba_group: + logger.error( + "req %s: Mooncake DCP transfer does not support " + "Mamba/GDN KV cache groups.", + d_req_id, + ) + has_block_error = True + break + local_group, remote_group = _filter_dcp_block_ids( + p_block_ids=local_group, + d_block_ids=remote_group, + p_dcp_rank=p_dcp_rank, + p_dcp_size=self.dcp_size, + d_dcp_rank=d_dcp_rank, + d_dcp_size=agent_meta.remote_dcp_size, ) - has_block_error = True - break - elif n_local > n_remote: - # Partial prefix cache hit: just read uncomputed blocks. - local_group = local_group[-n_remote:] if n_remote > 0 else [] + else: + n_local = len(local_group) + n_remote = len(remote_group) + if n_local < n_remote: + logger.error( + "req %s: local blocks(%d) < remote blocks(%d) " + "in a KV cache group (is_mamba_group=%s)", + d_req_id, + n_local, + n_remote, + is_mamba_group, + ) + has_block_error = True + break + elif n_local > n_remote: + # Partial prefix cache hit: just read uncomputed blocks. + local_group = local_group[-n_remote:] if n_remote > 0 else [] local_block_ids_by_group.append(local_group) remote_block_ids_by_group.append(remote_group) @@ -1938,6 +2175,7 @@ async def _build_transfer_params( remote_kv_block_len=remote_region.kv_block_len, remote_tp_rank=agent_meta.remote_tp_rank, remote_tp_size=agent_meta.remote_tp_size, + remote_dcp_size=agent_meta.remote_dcp_size, ) if not should_transfer: # Replicated KV cache: only one producer rank in the TP group @@ -2577,6 +2815,7 @@ async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str): for tp_rank, tp_entry in dp_entry["worker_addr"].items() } self._tp_size[remote_engine_id] = len(dp_entry["worker_addr"]) + self._dcp_size[remote_engine_id] = dp_entry.get("dcp_size", 1) except Exception as e: logger.error( "Failed to connect to bootstrap server %s (%s): %s", @@ -2594,33 +2833,58 @@ def receive_kv( remote_engine_id: EngineId, pull_metas: dict[ReqId, PullReqMeta], ): - remote_tp_ranks = self.transfer_topo.handshake_target_ranks( - self._tp_size[remote_engine_id] + remote_tp_size = self._tp_size[remote_engine_id] + remote_dcp_size = self._dcp_size.get(remote_engine_id, 1) + no_dcp = self.dcp_size == 1 and remote_dcp_size == 1 + total_num_kv_heads = ( + 1 if self.use_mla else self.model_config.get_total_num_kv_heads() ) - worker_addrs: list[str] = [] + worker_addr_to_metas: dict[str, dict[ReqId, PullReqMeta]] = defaultdict(dict) + req_remote_tp_ranks: dict[ReqId, list[int]] = {} selected_remote_pp: dict[int, list[int]] = {} - for remote_tp_rank in remote_tp_ranks: - pp_to_addr = self._remote_agents[remote_engine_id][remote_tp_rank] - if self.pp_size == len(pp_to_addr) and self.pp_rank in pp_to_addr: - pp_ranks = [self.pp_rank] + + for req_id, pull_meta in pull_metas.items(): + if no_dcp: + remote_tp_ranks = self.transfer_topo.handshake_target_ranks( + remote_tp_size + ) else: - pp_ranks = sorted(pp_to_addr) - selected_remote_pp[remote_tp_rank] = pp_ranks - worker_addrs.extend(pp_to_addr[pp_rank] for pp_rank in pp_ranks) + remote_tp_ranks = _get_prefill_ranks_for_decode( + d_tp_rank=self.tp_rank, + d_tp_size=self.tp_size, + d_dcp_size=self.dcp_size, + p_tp_size=remote_tp_size, + p_dcp_size=remote_dcp_size, + total_num_kv_heads=total_num_kv_heads, + is_mla=self.use_mla, + ) + req_remote_tp_ranks[req_id] = remote_tp_ranks + + count = 0 + for remote_tp_rank in remote_tp_ranks: + pp_to_addr = self._remote_agents[remote_engine_id][remote_tp_rank] + if self.pp_size == len(pp_to_addr) and self.pp_rank in pp_to_addr: + pp_ranks = [self.pp_rank] + else: + pp_ranks = sorted(pp_to_addr) + selected_remote_pp[remote_tp_rank] = pp_ranks + for pp_rank in pp_ranks: + worker_addr = pp_to_addr[pp_rank] + worker_addr_to_metas[worker_addr][req_id] = pull_meta + count += 1 + pull_meta.pull_tasks_count = count - count = len(worker_addrs) logger.debug( - "Receiving Mooncake KV for engine %s from producer TP ranks %s " - "and PP ranks %s", + "Receiving Mooncake KV for engine %s from producer TP ranks by req %s " + "and PP ranks %s, remote_dcp_size=%d", remote_engine_id, - remote_tp_ranks, + req_remote_tp_ranks, selected_remote_pp, + remote_dcp_size, ) - for pull_meta in pull_metas.values(): - pull_meta.pull_tasks_count = count - for worker_addr in worker_addrs: + for worker_addr, worker_pull_metas in worker_addr_to_metas.items(): asyncio.create_task( - self.receive_kv_from_single_worker(worker_addr, pull_metas) + self.receive_kv_from_single_worker(worker_addr, worker_pull_metas) ) async def handle_new_engine_id( @@ -2699,6 +2963,8 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata): ) def _producer_cache_is_replicated(self) -> bool: + if self.dcp_size > 1: + return False return self.transfer_topo.local_replicates_kv_cache def _prepare_transfer_regions( @@ -2960,15 +3226,42 @@ def _get_sender_transfer_plan( remote_kv_block_len: int, remote_tp_rank: int, remote_tp_size: int, + remote_dcp_size: int, ) -> tuple[bool, int, int, int]: + no_dcp = self.dcp_size == 1 and remote_dcp_size == 1 + if no_dcp: + local_transfer_tp_rank = self.tp_rank + local_transfer_tp_size = self.tp_size + remote_transfer_tp_rank = remote_tp_rank + remote_transfer_tp_size = remote_tp_size + producer_cache_replicated = self._producer_cache_is_replicated() + else: + local_transfer_tp_rank, local_transfer_tp_size = ( + _get_kv_head_group_rank_and_count( + self.tp_rank, + self.tp_size, + self._total_num_kv_heads, + self.use_mla, + ) + ) + remote_transfer_tp_rank, remote_transfer_tp_size = ( + _get_kv_head_group_rank_and_count( + remote_tp_rank, + remote_tp_size, + self._total_num_kv_heads, + self.use_mla, + ) + ) + producer_cache_replicated = False + return _compute_sender_transfer_plan( - local_tp_rank=self.tp_rank, - local_tp_size=self.tp_size, - remote_tp_rank=remote_tp_rank, - remote_tp_size=remote_tp_size, + local_tp_rank=local_transfer_tp_rank, + local_tp_size=local_transfer_tp_size, + remote_tp_rank=remote_transfer_tp_rank, + remote_tp_size=remote_transfer_tp_size, local_kv_block_len=local_kv_block_len, remote_kv_block_len=remote_kv_block_len, - producer_cache_replicated=self._producer_cache_is_replicated(), + producer_cache_replicated=producer_cache_replicated, total_num_kv_heads=( None if self.use_mla or self.kv_cache_config.has_mamba_layers diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py index 5f5fba157f1c..df99eb249a3c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py @@ -23,6 +23,7 @@ class RegisterWorkerPayload(BaseModel): tp_rank: int pp_rank: int addr: WorkerAddr + dcp_size: int = 1 @dataclass @@ -30,6 +31,7 @@ class EngineEntry: engine_id: EngineId # {tp_rank: {pp_rank: worker_addr}} worker_addr: dict[int, dict[int, WorkerAddr]] + dcp_size: int = 1 class MooncakeBootstrapServer: @@ -130,6 +132,7 @@ async def register_worker(self, payload: RegisterWorkerPayload): self.workers[payload.dp_rank] = EngineEntry( engine_id=payload.engine_id, worker_addr={}, + dcp_size=payload.dcp_size, ) dp_entry = self.workers[payload.dp_rank] @@ -141,6 +144,14 @@ async def register_worker(self, payload: RegisterWorkerPayload): f"expected {dp_entry.engine_id}, got {payload.engine_id}" ), ) + if dp_entry.dcp_size != payload.dcp_size: + raise HTTPException( + status_code=400, + detail=( + f"DCP size mismatch for dp_rank={payload.dp_rank}: " + f"expected {dp_entry.dcp_size}, got {payload.dcp_size}" + ), + ) if payload.tp_rank not in dp_entry.worker_addr: dp_entry.worker_addr[payload.tp_rank] = {} @@ -164,11 +175,13 @@ async def register_worker(self, payload: RegisterWorkerPayload): tp_entry[payload.pp_rank] = payload.addr logger.debug( - "Registered worker: engine_id=%s, dp_rank=%d, tp_rank=%d, pp_rank=%d at %s", + "Registered worker: engine_id=%s, dp_rank=%d, tp_rank=%d, " + "pp_rank=%d, dcp_size=%d at %s", payload.engine_id, payload.dp_rank, payload.tp_rank, payload.pp_rank, + payload.dcp_size, payload.addr, ) From b2f90022fa5a2de193057b988b7c5c2959ff49f1 Mon Sep 17 00:00:00 2001 From: nwpu-zxr <978074293@qq.com> Date: Fri, 21 Aug 2026 10:41:38 +0800 Subject: [PATCH 2/3] bugfix: fix decode block start loc Signed-off-by: nwpu-zxr <978074293@qq.com> --- .../unit/test_mooncake_connector.py | 152 ++++++++++-------- .../v1/mooncake/mooncake_connector.py | 97 ++++++----- 2 files changed, 139 insertions(+), 110 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index 7dab8dbb1434..d8ab8c390bf5 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -39,14 +39,11 @@ _block_ids_for_region, _coalesce_contiguous_transfer_regions, _compute_sender_transfer_plan, + _get_decode_ranks_for_prefill, + _get_prefill_ranks_for_decode, _has_opaque_packed_row, _pp_mismatch_hides_packed_layers, _validate_asymmetric_region_lengths, - - - _filter_dcp_block_ids, - _get_decode_ranks_for_prefill, - _get_prefill_ranks_for_decode, get_mooncake_bootstrap_addr, should_launch_bootstrap_server, ) @@ -332,25 +329,26 @@ def test_align_transfer_regions_uses_layer_name_occurrences(): "p_dcp_size", "d_dcp_rank", "d_dcp_size", + "num_computed_tokens", "expected_p_blocks", "expected_d_blocks", ), [ # P splits a full D block range into finer DCP slices. - ([10], list(range(20, 28)), 0, 8, 0, 1, [10], [20]), - ([10], list(range(20, 28)), 7, 8, 0, 1, [10], [27]), + ([10], list(range(20, 28)), 0, 8, 0, 1, 0, [10], [20]), + ([10], list(range(20, 28)), 7, 8, 0, 1, 0, [10], [27]), # P keeps full blocks while D asks for one finer DCP slice. - (list(range(10, 18)), [20], 0, 1, 4, 8, [14], [20]), + (list(range(10, 18)), [20], 0, 1, 4, 8, 0, [14], [20]), # Different DCP sizes: match by sequence-block position. - ([10, 11, 12, 13], [20, 21], 1, 4, 5, 8, [11, 13], [20, 21]), - ([10, 11], [20, 21, 22, 23], 5, 8, 1, 4, [10, 11], [21, 23]), - # Partial prefix hit: D only requests the tail sequence span. - ([10], [20, 21, 22, 23], 4, 8, 0, 1, [10], [20]), - ([10], [20, 21, 22, 23], 0, 8, 0, 1, [], []), - # Short P span is clipped to the overlapping sequence-block range. - ([10], list(range(20, 30)), 0, 8, 0, 1, [10], [22]), - ([10], [20], 1, 4, 1, 8, [], []), - ([10], [20], 1, 4, 5, 8, [10], [20]), + ([10, 11, 12, 13], [20, 21], 1, 4, 5, 8, 0, [11, 13], [20, 21]), + ([10, 11], [20, 21, 22, 23], 5, 8, 1, 4, 0, [10, 11], [21, 23]), + # Partial prefix hit: D only requests a non-zero sequence span. + ([10], [20, 21, 22, 23], 4, 8, 0, 1, 64, [10], [20]), + ([10], [20, 21, 22, 23], 0, 8, 0, 1, 64, [], []), + # Full-prefill short spans align from sequence block zero. + ([10], list(range(20, 30)), 0, 8, 0, 1, 0, [10], [20]), + ([10], [20], 1, 4, 1, 8, 0, [10], [20]), + ([10], [20], 1, 4, 5, 8, 0, [], []), ], ) def test_filter_dcp_block_ids( @@ -360,16 +358,20 @@ def test_filter_dcp_block_ids( p_dcp_size: int, d_dcp_rank: int, d_dcp_size: int, + num_computed_tokens: int, expected_p_blocks: list[int], expected_d_blocks: list[int], ): - actual_p_blocks, actual_d_blocks = _filter_dcp_block_ids( + worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) + worker.block_size = 16 + actual_p_blocks, actual_d_blocks = worker._filter_dcp_block_ids( p_block_ids=p_block_ids, d_block_ids=d_block_ids, p_dcp_rank=p_dcp_rank, p_dcp_size=p_dcp_size, d_dcp_rank=d_dcp_rank, d_dcp_size=d_dcp_size, + num_computed_tokens=num_computed_tokens, ) assert actual_p_blocks == expected_p_blocks assert actual_d_blocks == expected_d_blocks @@ -417,7 +419,7 @@ async def test_build_transfer_params_filters_dcp_blocks(): remote_port=54321, remote_tp_size=1, remote_tp_rank=0, - req_blocks={"d-req-dcp-filter": ("xfer-dcp-filter", [[20, 21, 22, 23]])}, + req_blocks={"d-req-dcp-filter": ("xfer-dcp-filter", [[20, 21, 22, 23]], 64)}, kv_caches_base_addr=[0x2000], block_lens=[block_len], kv_block_lens=[block_len], @@ -427,13 +429,17 @@ async def test_build_transfer_params_filters_dcp_blocks(): remote_block_size=16, ) - src_ptrs, dst_ptrs, lengths, err_reqs, err_msg = ( - await worker._build_transfer_params( - ready_reqs=[("d-req-dcp-filter", send_meta)], - agent_meta=xfer_meta, - local_regions=local_regions, - remote_regions=remote_regions, - ) + ( + src_ptrs, + dst_ptrs, + lengths, + err_reqs, + err_msg, + ) = await worker._build_transfer_params( + ready_reqs=[("d-req-dcp-filter", send_meta)], + agent_meta=xfer_meta, + local_regions=local_regions, + remote_regions=remote_regions, ) assert err_reqs == [] @@ -485,7 +491,7 @@ async def test_build_transfer_params_returns_empty_for_unowned_dcp_blocks(): remote_port=54321, remote_tp_size=1, remote_tp_rank=0, - req_blocks={"d-req-dcp-empty": ("xfer-dcp-empty", [[20, 21, 22, 23]])}, + req_blocks={"d-req-dcp-empty": ("xfer-dcp-empty", [[20, 21, 22, 23]], 64)}, kv_caches_base_addr=[0x2000], block_lens=[block_len], kv_block_lens=[block_len], @@ -495,13 +501,17 @@ async def test_build_transfer_params_returns_empty_for_unowned_dcp_blocks(): remote_block_size=16, ) - src_ptrs, dst_ptrs, lengths, err_reqs, err_msg = ( - await worker._build_transfer_params( - ready_reqs=[("d-req-dcp-empty", send_meta)], - agent_meta=xfer_meta, - local_regions=local_regions, - remote_regions=remote_regions, - ) + ( + src_ptrs, + dst_ptrs, + lengths, + err_reqs, + err_msg, + ) = await worker._build_transfer_params( + ready_reqs=[("d-req-dcp-empty", send_meta)], + agent_meta=xfer_meta, + local_regions=local_regions, + remote_regions=remote_regions, ) assert src_ptrs == [] @@ -1850,7 +1860,7 @@ async def test_kv_consumuer(monkeypatch): assert sent_meta.remote_hostname == "127.0.0.1" assert sent_meta.remote_port == 54321 - assert sent_meta.req_blocks["d-req-1"] == ("xfer-req-1", [[100, 101]]) + assert sent_meta.req_blocks["d-req-1"] == ("xfer-req-1", [[100, 101]], 0) assert sent_meta.kv_caches_base_addr == [0x1000] assert sent_meta.block_lens == [4096] assert sent_meta.registered_layer_names == ["model.layers.0.self_attn"] @@ -3018,7 +3028,6 @@ def test_sleep_unregisters_and_restores_kv_caches(): ) async def test_kv_producer_dcp_uses_effective_tp_for_offsets(monkeypatch): """DCP transfer planning should slice by TP//DCP, not physical TP.""" - P_TP_SIZE = 4 P_TP_RANK = 2 P_DCP_SIZE = 2 @@ -3154,15 +3163,18 @@ def test_get_prefill_ranks_for_decode_mla_cases( p_dcp_size: int, expected: list[int], ): - assert _get_prefill_ranks_for_decode( - d_tp_rank=d_tp_rank, - d_tp_size=d_tp_size, - d_dcp_size=d_dcp_size, - p_tp_size=p_tp_size, - p_dcp_size=p_dcp_size, - total_num_kv_heads=1, - is_mla=True, - ) == expected + assert ( + _get_prefill_ranks_for_decode( + d_tp_rank=d_tp_rank, + d_tp_size=d_tp_size, + d_dcp_size=d_dcp_size, + p_tp_size=p_tp_size, + p_dcp_size=p_dcp_size, + total_num_kv_heads=1, + is_mla=True, + ) + == expected + ) @pytest.mark.parametrize( @@ -3207,15 +3219,18 @@ def test_get_prefill_ranks_for_decode_gqa_and_heterogeneous_tp( num_kv_heads: int, expected: list[int], ): - assert _get_prefill_ranks_for_decode( - d_tp_rank=d_tp_rank, - d_tp_size=d_tp_size, - d_dcp_size=d_dcp_size, - p_tp_size=p_tp_size, - p_dcp_size=p_dcp_size, - total_num_kv_heads=num_kv_heads, - is_mla=False, - ) == expected + assert ( + _get_prefill_ranks_for_decode( + d_tp_rank=d_tp_rank, + d_tp_size=d_tp_size, + d_dcp_size=d_dcp_size, + p_tp_size=p_tp_size, + p_dcp_size=p_dcp_size, + total_num_kv_heads=num_kv_heads, + is_mla=False, + ) + == expected + ) @pytest.mark.parametrize( @@ -3261,15 +3276,18 @@ def test_get_decode_ranks_for_prefill_inverse_mapping( is_mla: bool, expected: list[int], ): - assert _get_decode_ranks_for_prefill( - p_tp_rank=p_tp_rank, - p_tp_size=p_tp_size, - p_dcp_size=p_dcp_size, - d_tp_size=d_tp_size, - d_dcp_size=d_dcp_size, - total_num_kv_heads=num_kv_heads, - is_mla=is_mla, - ) == expected + assert ( + _get_decode_ranks_for_prefill( + p_tp_rank=p_tp_rank, + p_tp_size=p_tp_size, + p_dcp_size=p_dcp_size, + d_tp_size=d_tp_size, + d_dcp_size=d_dcp_size, + total_num_kv_heads=num_kv_heads, + is_mla=is_mla, + ) + == expected + ) def test_process_pulling_result_waits_for_all_dcp_acks(): @@ -3305,7 +3323,6 @@ def test_process_pulling_result_waits_for_all_dcp_acks(): @pytest.mark.asyncio async def test_receive_kv_selects_remote_dcp_peer_ranks(): """Decode should query the producer DCP representative rank.""" - vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer" ) @@ -3325,8 +3342,7 @@ async def test_receive_kv_selects_remote_dcp_peer_ranks(): decode_worker.pp_rank = 0 decode_worker._remote_agents = { "p-engine": { - rank: {0: f"tcp://producer-tp{rank}:1234"} - for rank in range(4) + rank: {0: f"tcp://producer-tp{rank}:1234"} for rank in range(4) } } decode_worker._tp_size["p-engine"] = 4 @@ -3363,7 +3379,6 @@ async def fake_receive(worker_addr: str, metas: dict[str, PullReqMeta]): @pytest.mark.asyncio async def test_receive_kv_routes_metadata_by_dcp_rank_modulo(): """DCP metadata routing follows modulo dcp-rank mapping.""" - vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer" ) @@ -3383,8 +3398,7 @@ async def test_receive_kv_routes_metadata_by_dcp_rank_modulo(): decode_worker.pp_rank = 0 decode_worker._remote_agents = { "p-engine": { - rank: {0: f"tcp://producer-tp{rank}:1234"} - for rank in range(4) + rank: {0: f"tcp://producer-tp{rank}:1234"} for rank in range(4) } } decode_worker._tp_size["p-engine"] = 4 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 43eba5b7c9bb..07d2f8d26817 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 @@ -163,8 +163,6 @@ def _validate_dcp_size_for_kv_head_groups( ) - - def _get_kv_head_group_rank_and_count( tp_rank: int, tp_size: int, @@ -186,10 +184,8 @@ def _get_prefill_ranks_for_decode( total_num_kv_heads: int, is_mla: bool, ) -> list[int]: - d_kv_head_group_rank, d_kv_head_group_count = ( - _get_kv_head_group_rank_and_count( - d_tp_rank, d_tp_size, total_num_kv_heads, is_mla - ) + d_kv_head_group_rank, d_kv_head_group_count = _get_kv_head_group_rank_and_count( + d_tp_rank, d_tp_size, total_num_kv_heads, is_mla ) _, p_kv_head_group_count = _get_kv_head_group_rank_and_count( 0, p_tp_size, total_num_kv_heads, is_mla @@ -247,34 +243,6 @@ def _get_decode_ranks_for_prefill( return d_tp_ranks -def _filter_dcp_block_ids( - p_block_ids: list[int], - d_block_ids: list[int], - p_dcp_rank: int, - p_dcp_size: int, - d_dcp_rank: int, - d_dcp_size: int, -) -> tuple[list[int], list[int]]: - p_sequence_span = len(p_block_ids) * p_dcp_size - d_sequence_span = len(d_block_ids) * d_dcp_size - d_start_sequence_block = p_sequence_span - d_sequence_span - d_sequence_block_to_cache_block = { - d_start_sequence_block + i * d_dcp_size + d_dcp_rank: block_id - for i, block_id in enumerate(d_block_ids) - } - - filtered_p_block_ids: list[int] = [] - filtered_d_block_ids: list[int] = [] - for i, p_block_id in enumerate(p_block_ids): - p_sequence_block = i * p_dcp_size + p_dcp_rank - d_block_id = d_sequence_block_to_cache_block.get(p_sequence_block) - if d_block_id is not None: - filtered_p_block_ids.append(p_block_id) - filtered_d_block_ids.append(d_block_id) - - return filtered_p_block_ids, filtered_d_block_ids - - def _get_tp_ratio(local_tp_size: int, remote_tp_size: int) -> int: """Return the TP ratio used by heterogeneous TP transfer planning. @@ -744,7 +712,7 @@ class MooncakeXferMetadata( remote_port: int remote_tp_size: int remote_tp_rank: int - req_blocks: dict[ReqId, tuple[TransferId, list[list[int]]]] + req_blocks: dict[ReqId, tuple[TransferId, list[list[int]], int]] kv_caches_base_addr: list[int] block_lens: list[int] kv_block_lens: list[int] @@ -797,6 +765,7 @@ class PullReqMeta: local_block_ids: list[list[int]] remote_engine_id: EngineId remote_bootstrap_addr: str + num_computed_tokens: int = 0 # Set expire time to avoid infinitely sending requests. expire_time: float = float("inf") # Designed for one D pairing to multiple P @@ -842,6 +811,7 @@ def add_new_req( remote_engine_id=remote_engine_id, remote_bootstrap_addr=kv_transfer_params["remote_bootstrap_addr"], transfer_id=transfer_id, + num_computed_tokens=kv_transfer_params.get("num_computed_tokens", 0), ) else: self.reqs_to_send[request_id] = (transfer_id, local_block_ids) @@ -1192,6 +1162,7 @@ def get_num_new_matched_tokens( # Remote prefill: get all prompt blocks from remote. assert not self.is_kv_producer token_ids = request.prompt_token_ids or [] + params["num_computed_tokens"] = num_computed_tokens count = self._get_remote_prefill_token_count(len(token_ids)) - ( num_computed_tokens ) @@ -1492,8 +1463,8 @@ def __init__( self.cache_config = vllm_config.cache_config self.kv_cache_config = kv_cache_config self.use_mla = self.model_config.use_mla - self._total_num_kv_heads = 1 if self.use_mla else ( - self.model_config.get_total_num_kv_heads() + self._total_num_kv_heads = ( + 1 if self.use_mla else (self.model_config.get_total_num_kv_heads()) ) if self.dcp_size > 1 and len(kv_cache_config.kv_cache_groups) != 1: msg = ( @@ -1841,7 +1812,8 @@ async def send_kv_to_decode( ) await sock.send_multipart((identity, self._encoder.encode(response))) return - for d_req_id, (transfer_id, _) in meta.req_blocks.items(): + for d_req_id, req_block_entry in meta.req_blocks.items(): + transfer_id = req_block_entry[0] if transfer_id not in self.reqs_need_send: # This req is not enqueued in P side yet, create it here. self.reqs_need_send[transfer_id] = SendBlockMeta( @@ -2023,6 +1995,36 @@ def _logical_to_kernel_block_ids( for i, group in enumerate(block_ids) ] + def _filter_dcp_block_ids( + self, + p_block_ids: list[int], + d_block_ids: list[int], + p_dcp_rank: int, + p_dcp_size: int, + d_dcp_rank: int, + d_dcp_size: int, + num_computed_tokens: int, + ) -> tuple[list[int], list[int]]: + d_start_sequence_block = num_computed_tokens // self.block_size + d_first_sequence_block = d_start_sequence_block + ( + (d_dcp_rank - d_start_sequence_block) % d_dcp_size + ) + d_sequence_block_to_cache_block = { + d_first_sequence_block + i * d_dcp_size: block_id + for i, block_id in enumerate(d_block_ids) + } + + filtered_p_block_ids: list[int] = [] + filtered_d_block_ids: list[int] = [] + for i, p_block_id in enumerate(p_block_ids): + p_sequence_block = i * p_dcp_size + p_dcp_rank + d_block_id = d_sequence_block_to_cache_block.get(p_sequence_block) + if d_block_id is not None: + filtered_p_block_ids.append(p_block_id) + filtered_d_block_ids.append(d_block_id) + + return filtered_p_block_ids, filtered_d_block_ids + async def _build_transfer_params( self, ready_reqs: list[tuple[ReqId, SendBlockMeta]], @@ -2038,7 +2040,11 @@ async def _build_transfer_params( remote_session = f"{agent_meta.remote_hostname}:{agent_meta.remote_port}" for d_req_id, send_meta in ready_reqs: - _, remote_block_ids_per_group = agent_meta.req_blocks[d_req_id] + req_block_entry = agent_meta.req_blocks[d_req_id] + remote_block_ids_per_group = req_block_entry[1] + remote_num_computed_tokens = ( + req_block_entry[2] if len(req_block_entry) > 2 else 0 + ) if not remote_block_ids_per_group or all( len(g) == 0 for g in remote_block_ids_per_group @@ -2099,13 +2105,14 @@ async def _build_transfer_params( ) has_block_error = True break - local_group, remote_group = _filter_dcp_block_ids( + local_group, remote_group = self._filter_dcp_block_ids( p_block_ids=local_group, d_block_ids=remote_group, p_dcp_rank=p_dcp_rank, p_dcp_size=self.dcp_size, d_dcp_rank=d_dcp_rank, d_dcp_size=agent_meta.remote_dcp_size, + num_computed_tokens=remote_num_computed_tokens, ) else: n_local = len(local_group) @@ -3132,7 +3139,11 @@ def _build_pull_metadata( include_layout: bool, ) -> MooncakeXferMetadata: req_blocks = { - req_id: (pull_meta.transfer_id, pull_meta.local_block_ids) + req_id: ( + pull_meta.transfer_id, + pull_meta.local_block_ids, + pull_meta.num_computed_tokens, + ) for req_id, pull_meta in pull_metas.items() } if not include_layout: @@ -3141,6 +3152,8 @@ def _build_pull_metadata( remote_port=self.rpc_port, remote_tp_size=self.tp_size, remote_tp_rank=self.tp_rank, + remote_dcp_size=self.dcp_size, + remote_block_size=self.block_size, remote_pp_size=self.pp_size, layout_id=self._kv_layout_id, req_blocks=req_blocks, @@ -3153,6 +3166,8 @@ def _build_pull_metadata( remote_port=self.rpc_port, remote_tp_size=self.tp_size, remote_tp_rank=self.tp_rank, + remote_dcp_size=self.dcp_size, + remote_block_size=self.block_size, remote_pp_size=self.pp_size, layout_id=self._kv_layout_id, req_blocks=req_blocks, From 8622ec1f439163dddfbdad1764f2336e4316f219 Mon Sep 17 00:00:00 2001 From: nwpu-zxr <978074293@qq.com> Date: Mon, 19 Oct 2026 19:10:31 +0800 Subject: [PATCH 3/3] [Mooncake] Align DCP transfer constraints and replica completion Require whole physical blocks and equal P/D block sizes for DCP transfers. Resolve Mooncake PD interleave consistently with NIXL and reject unsupported DCP ratios. Include all producer replicas in completion accounting while selecting one replica for payload transfer when producer DCP is disabled. Co-authored-by: Codex Signed-off-by: nwpu-zxr <978074293@qq.com> --- .../unit/test_mooncake_connector.py | 6 +- vllm/config/vllm.py | 15 +++- .../v1/mooncake/mooncake_connector.py | 75 ++++++++++++++++--- 3 files changed, 81 insertions(+), 15 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index d8ab8c390bf5..f5ff5649479b 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -530,6 +530,7 @@ async def test_build_transfer_params_separates_prefill_pp_layers(): worker.is_kv_producer = True worker.tp_rank = 0 worker.tp_size = 1 + worker.dcp_size = 1 worker.use_mla = False worker.kv_cache_config = _make_test_kv_cache_config() worker._physical_blocks_per_logical_kv_block = 1 @@ -2160,6 +2161,7 @@ async def test_build_transfer_params_sends_packed_region_once( worker.is_kv_producer = True worker.tp_rank = 0 worker.tp_size = 1 + worker.dcp_size = 1 worker.use_mla = True worker.kv_cache_config = _make_packed_mla_kv_cache_config( num_blocks=4, num_groups=len(local_ids) @@ -3192,8 +3194,8 @@ def test_get_prefill_ranks_for_decode_mla_cases( (0, 16, 1, 16, 4, 4, [0, 1, 2, 3]), (3, 16, 1, 16, 4, 4, [0, 1, 2, 3]), (4, 16, 1, 16, 4, 4, [4, 5, 6, 7]), - # Same TP/head groups: D splits by DCP, P keeps one representative copy. - (5, 16, 4, 16, 1, 4, [4]), + # D queries every P replica for completion; only the first sends data. + (5, 16, 4, 16, 1, 4, [4, 5, 6, 7]), # Same TP/head groups and same DCP: point-to-point by dcp rank. (6, 16, 4, 16, 4, 4, [6]), # P has a larger DCP size: one D slice maps to multiple finer P slices. diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 7312ef9c4e48..25fc9db937e2 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3331,8 +3331,9 @@ def adjust_dcp_kv_cache_interleave_size( "deprecated when PCP is fully supported." ) - if self.kv_transfer_config is None or not self.kv_transfer_config.has_connector( - "NixlConnector" + if self.kv_transfer_config is None or ( + not self.kv_transfer_config.has_connector("NixlConnector") + and not self.kv_transfer_config.has_connector("MooncakeConnector") ): return if not self.parallel_config._allow_auto_resolve_cp_interleave_size: @@ -3370,7 +3371,15 @@ def validate_block_size(self) -> None: self.kv_transfer_config is not None and self.kv_transfer_config.has_connector("NixlConnector") ) - if self.parallel_config.decode_context_parallel_size > 1 and not nixl_pd_active: + mooncake_pd_active = ( + self.kv_transfer_config is not None + and self.kv_transfer_config.has_connector("MooncakeConnector") + ) + if ( + self.parallel_config.decode_context_parallel_size > 1 + and not nixl_pd_active + and not mooncake_pd_active + ): assert ( self.parallel_config.cp_kv_cache_interleave_size <= block_size and block_size % self.parallel_config.cp_kv_cache_interleave_size == 0 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 07d2f8d26817..71f26d1d01dc 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 @@ -184,6 +184,11 @@ def _get_prefill_ranks_for_decode( total_num_kv_heads: int, is_mla: bool, ) -> list[int]: + if max(p_dcp_size, d_dcp_size) % min(p_dcp_size, d_dcp_size): + raise ValueError( + "Mooncake DCP transfer requires mutually divisible P/D DCP sizes: " + f"P={p_dcp_size}, D={d_dcp_size}." + ) d_kv_head_group_rank, d_kv_head_group_count = _get_kv_head_group_rank_and_count( d_tp_rank, d_tp_size, total_num_kv_heads, is_mla ) @@ -206,7 +211,10 @@ def _get_prefill_ranks_for_decode( ) d_dcp_rank = d_tp_rank % d_dcp_size - if p_dcp_size >= d_dcp_size: + if p_dcp_size == 1: + # All replicas must acknowledge completion; only one sends the payload. + p_dcp_ranks = list(range(p_ranks_per_kv_head_group)) + elif p_dcp_size >= d_dcp_size: p_dcp_ranks = list(range(d_dcp_rank, p_dcp_size, d_dcp_size)) else: p_dcp_ranks = [d_dcp_rank % p_dcp_size] @@ -731,6 +739,7 @@ class MooncakeXferMetadata( layout_id: int = 0 remote_dcp_size: int = 1 remote_block_size: int = 0 + remote_physical_blocks_per_logical: int = 1 class MooncakeXferResponseStatus(IntEnum): @@ -1483,6 +1492,7 @@ def __init__( ) self._physical_blocks_per_logical_kv_block = 1 self._sync_block_size_with_kernel() + self._validate_dcp_interleave() self.attn_backends = get_current_attn_backends(vllm_config) logger.debug( @@ -1530,6 +1540,24 @@ def __init__( self._xfer_meta_decoder = msgspec.msgpack.Decoder(MooncakeXferMetadata) self._xfer_resp_decoder = msgspec.msgpack.Decoder(MooncakeXferResponse) + def _validate_dcp_interleave(self) -> None: + if self.dcp_size == 1: + return + if self._physical_blocks_per_logical_kv_block != 1: + raise ValueError( + "Mooncake DCP currently requires logical block size to equal " + "physical kernel block size; split logical blocks are unsupported." + ) + interleave = self.vllm_config.parallel_config.cp_kv_cache_interleave_size + if interleave != self.block_size: + raise ValueError( + "Mooncake DCP requires cp_kv_cache_interleave_size to equal " + f"the physical kernel block size ({self.block_size}); " + f"got {interleave}. Set --cp-kv-cache-interleave-size " + f"{self.block_size} and use an attention backend supporting " + "block-aligned DCP." + ) + def _sync_block_size_with_kernel(self) -> None: # When speculative decoding (e.g. Eagle) is enabled, the main model # and draft model may use different attention backends with different @@ -1686,6 +1714,18 @@ async def send_kv_to_decode( self, identity: bytes, sock: zmq.asyncio.Socket, meta: MooncakeXferMetadata ): pending_reqs: dict[ReqId, SendBlockMeta] = {} + if (self.dcp_size > 1 or meta.remote_dcp_size > 1) and ( + self._physical_blocks_per_logical_kv_block != 1 + or meta.remote_physical_blocks_per_logical != 1 + ): + response = MooncakeXferResponse( + status=MooncakeXferResponseStatus.ERROR, + err_msg=( + "Mooncake DCP requires unsplit logical blocks on both P and D." + ), + ) + await sock.send_multipart((identity, self._encoder.encode(response))) + return if meta.remote_block_size and meta.remote_block_size != self.block_size: msg = ( "Mooncake DCP transfer currently requires matching block sizes: " @@ -2856,15 +2896,19 @@ def receive_kv( remote_tp_size ) else: - remote_tp_ranks = _get_prefill_ranks_for_decode( - d_tp_rank=self.tp_rank, - d_tp_size=self.tp_size, - d_dcp_size=self.dcp_size, - p_tp_size=remote_tp_size, - p_dcp_size=remote_dcp_size, - total_num_kv_heads=total_num_kv_heads, - is_mla=self.use_mla, - ) + try: + remote_tp_ranks = _get_prefill_ranks_for_decode( + d_tp_rank=self.tp_rank, + d_tp_size=self.tp_size, + d_dcp_size=self.dcp_size, + p_tp_size=remote_tp_size, + p_dcp_size=remote_dcp_size, + total_num_kv_heads=total_num_kv_heads, + is_mla=self.use_mla, + ) + except ValueError as exc: + self._handle_failed_recv(pull_metas, list(pull_metas), str(exc)) + return req_remote_tp_ranks[req_id] = remote_tp_ranks count = 0 @@ -3154,6 +3198,9 @@ def _build_pull_metadata( remote_tp_rank=self.tp_rank, remote_dcp_size=self.dcp_size, remote_block_size=self.block_size, + remote_physical_blocks_per_logical=( + self._physical_blocks_per_logical_kv_block + ), remote_pp_size=self.pp_size, layout_id=self._kv_layout_id, req_blocks=req_blocks, @@ -3168,6 +3215,9 @@ def _build_pull_metadata( remote_tp_rank=self.tp_rank, remote_dcp_size=self.dcp_size, remote_block_size=self.block_size, + remote_physical_blocks_per_logical=( + self._physical_blocks_per_logical_kv_block + ), remote_pp_size=self.pp_size, layout_id=self._kv_layout_id, req_blocks=req_blocks, @@ -3259,6 +3309,11 @@ def _get_sender_transfer_plan( self.use_mla, ) ) + if ( + self.dcp_size == 1 + and self.tp_rank % (self.tp_size // local_transfer_tp_size) != 0 + ): + return False, 0, 0, 0 remote_transfer_tp_rank, remote_transfer_tp_size = ( _get_kv_head_group_rank_and_count( remote_tp_rank,