diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index afe626cba82b..f5ff5649479b 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -39,6 +39,8 @@ _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, @@ -319,6 +321,206 @@ 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", + "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, 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, 0, [14], [20]), + # Different DCP sizes: match by sequence-block position. + ([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( + 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, + expected_p_blocks: list[int], + expected_d_blocks: list[int], +): + 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 + + +@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]], 64)}, + 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]], 64)}, + 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.""" @@ -328,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 @@ -771,6 +974,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 +987,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 +1001,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" @@ -1655,7 +1861,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"] @@ -1955,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) @@ -2813,3 +3020,440 @@ 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]), + # 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. + (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/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 1d12bd1ab19a..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 @@ -144,6 +144,113 @@ 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]: + 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 + ) + _, 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 == 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] + + 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 _get_tp_ratio(local_tp_size: int, remote_tp_size: int) -> int: """Return the TP ratio used by heterogeneous TP transfer planning. @@ -613,7 +720,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] @@ -630,6 +737,9 @@ 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 + remote_physical_blocks_per_logical: int = 1 class MooncakeXferResponseStatus(IntEnum): @@ -664,6 +774,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 @@ -709,6 +820,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) @@ -1059,6 +1171,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 ) @@ -1297,6 +1410,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,8 +1472,27 @@ 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() + self._validate_dcp_interleave() self.attn_backends = get_current_attn_backends(vllm_config) logger.debug( @@ -1368,14 +1501,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, @@ -1392,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 @@ -1444,6 +1610,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 +1714,55 @@ 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 (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: " + 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 +1815,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 @@ -1625,7 +1852,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( @@ -1807,6 +2035,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]], @@ -1822,7 +2080,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 @@ -1848,6 +2110,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 +2136,41 @@ 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 = 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, ) - 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 +2222,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 +2862,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 +2880,62 @@ 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) + 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 + 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 +3014,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( @@ -2866,7 +3183,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: @@ -2875,6 +3196,11 @@ 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_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, @@ -2887,6 +3213,11 @@ 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_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, @@ -2960,15 +3291,47 @@ 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, + ) + ) + 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, + 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, )