diff --git a/docs/source/user_guide/feature_guide/context_parallel.md b/docs/source/user_guide/feature_guide/context_parallel.md index 1f936d18a317..d5e33ee52fc4 100644 --- a/docs/source/user_guide/feature_guide/context_parallel.md +++ b/docs/source/user_guide/feature_guide/context_parallel.md @@ -19,10 +19,10 @@ PCP support is experimental and available only with ModelRunner V2. The followin | Attention Backend | Basic PCP | Prefix Caching + PCP | Chunked Prefill + PCP | MLAPO + PCP | Speculative Decoding + PCP | P/D Disaggregation + PCP | Sequence Parallelism (SP) + PCP | | --- | --- | --- | --- | --- | --- | --- | --- | -| MLA | ✅ Full compatibility | ✅ Full compatibility | ✅ Full compatibility | ✅ Full compatibility | 🟠 Partial compatibility (MTP, eager and `FULL_DECODE_ONLY`) | ❌ No compatibility | ❌ No compatibility | -| GQA | ✅ Full compatibility | ✅ Full compatibility | ✅ Full compatibility | — Not applicable | 🟠 Partial compatibility (Eagle3, eager and `FULL_DECODE_ONLY`) | ❌ No compatibility | ❌ No compatibility | -| SFA | ✅ Full compatibility | ✅ Full compatibility | ✅ Full compatibility | ❌ No compatibility | ❌ No compatibility | ❌ No compatibility | ❌ No compatibility | -| DSA | ✅ Full compatibility | ✅ Full compatibility | ✅ Full compatibility | — Not applicable | 🟠 Partial compatibility (MTP and DSpark, eager and `FULL_DECODE_ONLY`) | ❌ No compatibility | ❌ No compatibility | +| MLA | ✅ Full compatibility | ✅ Full compatibility | ✅ Full compatibility | ✅ Full compatibility | 🟠 Partial compatibility (MTP, eager and `FULL_DECODE_ONLY`) | ✅ Full compatibility (`MooncakeConnectorV1`) | ❌ No compatibility | +| GQA | ✅ Full compatibility | ✅ Full compatibility | ✅ Full compatibility | — Not applicable | 🟠 Partial compatibility (Eagle3, eager and `FULL_DECODE_ONLY`) | ✅ Full compatibility (`MooncakeConnectorV1`) | ❌ No compatibility | +| SFA | ✅ Full compatibility | ✅ Full compatibility | ✅ Full compatibility | ❌ No compatibility | ❌ No compatibility | ✅ Full compatibility (`MooncakeConnectorV1`) | ❌ No compatibility | +| DSA | ✅ Full compatibility | ✅ Full compatibility | ✅ Full compatibility | — Not applicable | 🟠 Partial compatibility (MTP and DSpark, eager and `FULL_DECODE_ONLY`) | ✅ Full compatibility (`MooncakeHybridConnector`) | ❌ No compatibility | - ✅ **Full compatibility**: The basic path or feature combination is supported. - 🟠 **Partial compatibility**: The basic path or feature combination is supported with the stated limitations. @@ -109,6 +109,7 @@ For either method, remove `--enforce-eager` and add the following option to use #### Constraints - PCP is supported only with ModelRunner V2. +- In P/D disaggregation, enable PCP only on the prefill (`kv_producer`) engine; the decode (`kv_consumer`) engine must use `prefill_context_parallel_size=1`. - PCP speculative decoding supports MTP with MLA and DSA models, Eagle3 with GQA models, and DSpark with DeepSeek-V4 DSA models. - Draft sampling must use the greedy method. diff --git a/tests/ut/kv_offload/test_mooncake_connector.py b/tests/ut/kv_offload/test_mooncake_connector.py index c6a857a6baaa..6ab99952b84b 100644 --- a/tests/ut/kv_offload/test_mooncake_connector.py +++ b/tests/ut/kv_offload/test_mooncake_connector.py @@ -796,6 +796,9 @@ def test_hybrid_group_pulls_metadata_filters_groups_per_remote_card(self): class TestKVCacheRecvingThreadBasic(unittest.TestCase): def setUp(self): + set_device_patch = patch("torch.npu.set_device") + set_device_patch.start() + self.addCleanup(set_device_patch.stop) self.engine = MagicMock() self.ready_event = threading.Event() self.vllm_config = MockVllmConfig() @@ -1141,15 +1144,22 @@ def _configure_mock_mamba_transfer(self): @patch.object(KVCacheRecvingThread, "_transfer_kv_cache_all_groups") @patch.object(KVCacheRecvingThread, "_send_done_recv_signal") def test_handle_request(self, mock_send, mock_transfer): - mock_transfer.return_value = None - mock_send.return_value = None - - self.thread._handle_request(self.test_req) - - mock_transfer.assert_called_once_with(self.test_req) - mock_send.assert_called_once_with("req1", "localhost", 6666, {6666: 1}) - cast(Any, self.thread.task_tracker).update_done_task_count.assert_called_once_with("req1") - self.mock_queue.task_done.assert_called_once() + for transfer_error in (None, RuntimeError("transfer failed")): + with self.subTest(transfer_error=transfer_error): + mock_send.reset_mock() + mock_transfer.reset_mock() + mock_transfer.side_effect = transfer_error + self.thread.task_tracker.reset_mock() + self.mock_queue.reset_mock() + + self.thread._handle_request(self.test_req) + + mock_transfer.assert_called_once_with(self.test_req) + mock_send.assert_called_once_with("req1", "localhost", 6666, {6666: 1}) + self.thread.task_tracker.update_done_task_count.assert_called_once_with("req1") + self.mock_queue.task_done.assert_called_once() + expected_errors = {1, 2} if transfer_error else set() + self.assertEqual(self.thread.get_and_clear_invalid_block_ids(), expected_errors) @patch.object(KVCacheRecvingThread, "_send_done_signal_to_free_remote_port") @patch.object(KVCacheRecvingThread, "_send_done_recv_signal") @@ -1722,31 +1732,22 @@ def setUp(self): ) self.thread.request_queue = queue.Queue() - @patch.object(KVCacheRecvingThread, "_handle_request") - def test_run_loop_normal(self, mock_handle): + def test_run_loop_normal(self): test_request = { "request_id": "req1", - "local_block_ids": [1, 2], - "remote_block_ids": [3, 4], - "remote_engine_id": "remote_engine", "remote_host": "localhost", "remote_handshake_port": 6666, - "remote_transfer_port": 7777, - "offset": 0, - "tp_num_need_pulls": 2, "all_task_done": False, } - - self.thread.request_queue.put(test_request) - self.thread.request_queue.put(None) - - self.thread.start() - time.sleep(0.1) - self.thread.join(timeout=1.0) - + # Test queue dispatch here; executor ordering is covered by the peer tests. + with ( + patch.object(self.thread.request_queue, "get", side_effect=[test_request, None, KeyboardInterrupt]), + patch.object(self.thread, "_submit_request") as mock_submit, + self.assertRaises(KeyboardInterrupt), + ): + self.thread.run() self.assertTrue(self.thread.ready_event.is_set()) - mock_handle.assert_called_once_with(test_request) - self.assertTrue(self.thread.request_queue.empty()) + mock_submit.assert_called_once_with(test_request) class MockVllmConfig: @@ -2300,21 +2301,6 @@ def test_get_transfer_block_ids_uses_compressed_prompt_len(self): self.assertEqual(block_ids, ([30, 31],)) - def test_get_transfer_block_ids_uses_cp_grouped_block_len(self): - self.scheduler.pcp_size = 1 - self.scheduler.dcp_size = 4 - self.scheduler.group_transfer_info = [ - types.SimpleNamespace( # type: ignore[list-item] - tokens_per_block=16, - blocks_per_window=0, - is_state_group=False, - ) - ] - - block_ids = self.scheduler._get_transfer_block_ids(([10, 11, 12, 13, 14],), prompt_len=65) - - self.assertEqual(block_ids, ([10, 11],)) - def test_get_transfer_block_ids_trims_sliding_window_mtp_blocks(self): self.scheduler.group_transfer_info = [ types.SimpleNamespace( # type: ignore[list-item] @@ -2387,28 +2373,6 @@ def test_request_finished_trims_mtp_blocks_in_params(self): self.assertEqual(params["num_prompt_blocks"], 3) self.assertIn("req_mtp", self.scheduler._reqs_need_send) - def test_request_finished_trims_cp_grouped_mtp_blocks_in_params(self): - self.scheduler.pcp_size = 1 - self.scheduler.dcp_size = 4 - self.scheduler.group_transfer_info = [ - types.SimpleNamespace( - tokens_per_block=16, - blocks_per_window=0, - is_state_group=False, - ) - ] - request = self._make_remote_decode_request(prompt_len=65, request_id="req_cp_mtp") - - delay_free, params = self.scheduler.request_finished(request, ([10, 11, 12, 13, 14],)) - - self.assertTrue(delay_free) - self.assertIsNotNone(params) - assert params is not None - self.assertEqual(params["remote_block_ids"], ([10, 11],)) - # num_prompt_blocks stays in no-CP units for worker-side CP distribution. - self.assertEqual(params["num_prompt_blocks"], 5) - self.assertIn("req_cp_mtp", self.scheduler._reqs_need_send) - def test_request_finished_clips_sliding_window_blocks_in_params(self): self.scheduler.group_transfer_info = [ types.SimpleNamespace( @@ -2447,49 +2411,43 @@ def test_request_finished_trims_mtp_before_swa_tail_clip(self): self.assertEqual(params["num_prompt_blocks"], 4) self.assertIn("req_mtp_swa", self.scheduler._reqs_need_send) - def test_request_finished_handles_mtp_swa_and_state_groups_together(self): + def test_request_finished_respects_kv_group_layout(self): self.scheduler.vllm_config.cache_config.mamba_cache_mode = "align" self.scheduler.group_transfer_info = [ - types.SimpleNamespace( - tokens_per_block=16, - blocks_per_window=0, - is_state_group=False, - ), - types.SimpleNamespace( - tokens_per_block=16, - blocks_per_window=3, - is_state_group=False, - ), - types.SimpleNamespace( - tokens_per_block=16, - blocks_per_window=0, - is_state_group=True, - ), + types.SimpleNamespace(tokens_per_block=16, blocks_per_window=0, is_state_group=False), + types.SimpleNamespace(tokens_per_block=32, blocks_per_window=0, is_state_group=False), + types.SimpleNamespace(tokens_per_block=16, blocks_per_window=3, is_state_group=False), + types.SimpleNamespace(tokens_per_block=16, blocks_per_window=0, is_state_group=True), + ] + blocks = ( + [100, 101, 102, 103, 104, 105], + [200, 201, 202, 203], + [0, 300, 301, 302, 303, 304], + [400, 401, 402, 403, 404, 405], + ) + # PCP keeps complete replicas. Only DCP changes attention block counts; + # compressed attention, sliding windows and aligned state retain their layouts. + cases = [ + (1, 1, ([100, 101, 102, 103, 104], [200, 201, 202], [301, 302, 303], [404])), + (4, 1, ([100, 101, 102, 103, 104], [200, 201, 202], [301, 302, 303], [404])), + (1, 2, ([100, 101, 102], [200, 201], [300, 301], [404])), + (1, 4, ([100, 101], [200], [300], [404])), ] - request = self._make_remote_decode_request(prompt_len=64, request_id="req_mixed_groups") + for pcp_size, dcp_size, expected in cases: + with self.subTest(pcp_size=pcp_size, dcp_size=dcp_size): + self.scheduler.pcp_size = pcp_size + self.scheduler.dcp_size = dcp_size + request = self._make_remote_decode_request(prompt_len=65) - delay_free, params = self.scheduler.request_finished( - request, - ( - [100, 101, 102, 103, 104], - [0, 200, 201, 202, 203, 204], - [300, 301, 302, 303, 304], - ), - ) + delay_free, params = self.scheduler.request_finished(request, blocks) - self.assertTrue(delay_free) - self.assertIsNotNone(params) - assert params is not None - self.assertEqual( - params["remote_block_ids"], - ( - [100, 101, 102, 103], - [200, 201, 202], - [303], - ), - ) - self.assertEqual(params["num_prompt_blocks"], 4) - self.assertIn("req_mixed_groups", self.scheduler._reqs_need_send) + self.assertTrue(delay_free) + assert params is not None + self.assertEqual(params["remote_block_ids"], expected) + self.assertEqual(params["remote_pcp_size"], pcp_size) + self.assertEqual(params["remote_dcp_size"], dcp_size) + self.assertEqual(params["num_prompt_blocks"], 5) + self.assertIn(request.request_id, self.scheduler._reqs_need_send) class TestUtils(unittest.TestCase): @@ -2627,10 +2585,6 @@ def mock_get_ip(): return "127.0.0.1" -def mock_string_to_int64_hash(s): - return hash(s) - - def make_cpu_kv_cache(kv_heads: int = 8, head_dim: int = 16): return ( torch.empty((10, 16, kv_heads, head_dim), device="cpu"), @@ -2668,10 +2622,6 @@ def setUp(self): return_value=0, ), patch("vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_connector.get_ip", mock_get_ip), - patch( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_connector.string_to_int64_hash", - mock_string_to_int64_hash, - ), patch( "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_connector.global_te.get_transfer_engine", return_value=self.mock_transfer_engine, @@ -2986,8 +2936,8 @@ def get_kv_split_metadata( meta.remote_port = remote_port meta.remote_block_ids = (remote_block_ids,) meta.local_block_ids = (local_block_ids,) - meta.num_external_tokens = pcp_size * dcp_size * len(local_block_ids) * worker.block_size - meta.num_prompt_blocks = pcp_size * dcp_size * len(local_block_ids) + meta.num_external_tokens = dcp_size * len(local_block_ids) * worker.block_size + meta.num_prompt_blocks = dcp_size * len(local_block_ids) meta.num_computed_tokens = 0 meta.remote_engine_id = remote_engine_id meta.remote_host = "localhost" @@ -3014,62 +2964,6 @@ def get_kv_split_metadata( ), ) - self.assertEqual( - get_kv_split_metadata(False, 1, 1, 8, 1, 0, 8, 2, 8, 30000, [1], [1], 0), - ( - [ - [30001], - [30002], - [30003], - [30004], - [30005], - [30006], - [30007], - [30008], - [30009], - [30010], - [30011], - [30012], - [30013], - [30014], - [30015], - [30000], - ], - [[], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [1]], - [[], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [1]], - ), - ) - - self.assertEqual( - get_kv_split_metadata(True, 1, 1, 8, 1, 0, 8, 2, 2, 30000, [1], [1], 0), - ([[30001], [30008], [30009], [30000]], [[], [], [], [1]], [[], [], [], [1]]), - ) - - self.assertEqual( - get_kv_split_metadata(False, 1, 1, 8, 1, 0, 8, 2, 2, 30000, [1], [1], 0), - ([[30001], [30008], [30009], [30000]], [[], [], [], [1]], [[], [], [], [1]]), - ) - - self.assertEqual( - get_kv_split_metadata(True, 1, 2, 8, 1, 0, 8, 2, 2, 30000, [1], [1], 0), - ([[30000], [30008]], [[1], []], [[1], []]), - ) - - self.assertEqual( - get_kv_split_metadata(False, 1, 2, 8, 1, 0, 8, 2, 2, 30000, [1], [1], 0), - ([[30000], [30008]], [[1], []], [[1], []]), - ) - - # D rank0 holds 5 external blocks [1,2,3,4,5]; P stores blocks interleaved - # across 4 cp ranks (cp0: global 0,4,8 -> D local idx 0,2,4 = blocks 1,3,5; - # cp2: global 2,6 -> D local idx 1,3 = blocks 2,4). Expansion now happens in - # _get_kv_split_metadata (scale 1 => kernel == block), so each shard's local - # list is the chunk-selected kernels: shard0 -> [1,3,5], shard1 -> [2,4]. - self.assertEqual( - get_kv_split_metadata(True, 1, 2, 8, 0, 0, 8, 2, 2, 30000, [1, 2, 3], [1, 2, 3, 4, 5], 0)[:3], - ([[30000], [30008]], [[1, 3, 5], [2, 4]], [[1, 2, 3], [1, 2]]), - ) - # check remote ptp size self.assertEqual( get_kv_split_metadata(True, 1, 1, 8, 1, 0, 8, 1, 8, 30000, [1], [1], 0, 16), @@ -3079,13 +2973,9 @@ def get_kv_split_metadata( get_kv_split_metadata(False, 1, 1, 8, 1, 0, 8, 1, 8, 30000, [1], [1], 0, 16), get_kv_split_metadata(False, 1, 1, 8, 1, 0, 16, 1, 8, 30000, [1], [1], 0), ) - self.assertEqual( - get_kv_split_metadata(False, 1, 1, 8, 1, 0, 8, 2, 8, 30000, [1], [1], 0, 16), - get_kv_split_metadata(False, 1, 1, 8, 1, 0, 16, 2, 8, 30000, [1], [1], 0), - ) - def test_get_kv_split_metadata_unequal_block_size_with_decode_cp(self): - """Bd=2*Bp with D-side CP: P cp ranks 0,1 -> D rank0; cp ranks 2,3 -> D rank1.""" + def test_get_kv_split_metadata_unequal_block_size_with_decode_dcp(self): + """Bd=2*Bp: P DCP ranks 0,1 map to D rank 0; ranks 2,3 map to D rank 1.""" for dcp_rank in (0, 1): with self.subTest(dcp_rank=dcp_rank): with patch( @@ -3116,8 +3006,8 @@ def test_get_kv_split_metadata_unequal_block_size_with_decode_cp(self): } meta = types.SimpleNamespace( - remote_pcp_size=2, - remote_dcp_size=2, + remote_pcp_size=1, + remote_dcp_size=4, remote_ptp_size=4, remote_port=30000, remote_block_ids=([10, 11, 12, 13],), @@ -3141,10 +3031,10 @@ def test_get_kv_split_metadata_unequal_block_size_with_decode_cp(self): if dcp_rank == 0: self.assertEqual(ports, [[30000], [30001]]) else: - self.assertEqual(ports, [[30004], [30005]]) + self.assertEqual(ports, [[30002], [30003]]) - def test_get_kv_split_metadata_cp_with_prefix_cache_skips_prefix(self): - """CP + prefix cache hit (P0>0): remote ids must start past the prefix + def test_get_kv_split_metadata_dcp_with_prefix_cache_skips_prefix(self): + """DCP + prefix cache hit (P0>0): remote ids must start past the prefix blocks (remote_first), aligned with local_chunk_token_starts.""" worker = MooncakeConnectorWorker(self.vllm_config, self.engine_id, MockKVCacheConfig()) @@ -3168,10 +3058,10 @@ def test_get_kv_split_metadata_cp_with_prefix_cache_skips_prefix(self): 0: ({"kv_cache_spec_type": "FullAttentionSpec"}, [0]), } - # 6 prompt blocks, 4 external (P0 = 2 prefix-cached blocks), remote PCP=2. + # 6 prompt blocks, 4 external (P0 = 2 prefix-cached blocks), remote DCP=2. meta = types.SimpleNamespace( - remote_pcp_size=2, - remote_dcp_size=1, + remote_pcp_size=1, + remote_dcp_size=2, remote_ptp_size=8, remote_port=30000, remote_block_ids=([50, 51, 52],), @@ -3283,17 +3173,20 @@ def test_get_kv_split_metadata_non_cp_uses_compress_ratio(self): self.assertEqual(remote_ids, [([7, 8, 9],)]) def _build_worker_for_pd_case(self, case, tp_rank, pcp_rank=0, dcp_rank=0): - with patch.object( - self.vllm_config.kv_transfer_config, - "get_from_extra_config", - side_effect=lambda k, d=None, case=case: { - "prefill": { - "tp_size": case["prefill_tp_size"], - "dp_size": 1, - "pp_size": case["prefill_pp_size"], - }, - "decode": {"tp_size": case["decode_tp_size"], "dp_size": 1, "pp_size": 1}, - }.get(k, d), + with ( + patch.object(self.vllm_config.parallel_config, "tensor_parallel_size", case["decode_tp_size"]), + patch.object( + self.vllm_config.kv_transfer_config, + "get_from_extra_config", + side_effect=lambda k, d=None, case=case: { + "prefill": { + "tp_size": case["prefill_tp_size"], + "dp_size": 1, + "pp_size": case["prefill_pp_size"], + }, + "decode": {"tp_size": case["decode_tp_size"], "dp_size": 1, "pp_size": 1}, + }.get(k, d), + ), ): self.vllm_config.model_config.is_deepseek_mla = case["use_mla"] self.vllm_config.model_config.hf_text_config.num_key_value_heads = case["num_key_value_heads"] @@ -3327,7 +3220,7 @@ def _assert_group_pull_finish_flags(self, ports, group_pulls, expected_group_ids self.assertEqual(len(group_pulls), len(ports)) finish_count_by_group = {group_id: 0 for group_id in expected_group_ids} - for pcp_dcp_rank, (remote_ports, port_group_pulls) in enumerate(zip(ports, group_pulls)): + for shard_idx, (remote_ports, port_group_pulls) in enumerate(zip(ports, group_pulls)): self.assertEqual(len(port_group_pulls), len(remote_ports)) for remote_port_idx, pulls in enumerate(port_group_pulls): self.assertEqual({pull.group_id for pull in pulls}, expected_group_ids) @@ -3342,7 +3235,7 @@ def _assert_group_pull_finish_flags(self, ports, group_pulls, expected_group_ids finish_count_by_group[pull.group_id] += 1 if len(remote_ports) == 1: - expected_offset = pcp_dcp_rank % pulls[0].num_group_pulls + expected_offset = shard_idx % pulls[0].num_group_pulls else: expected_offset = remote_port_idx % pulls[0].num_group_pulls self.assertTrue(all(pull.remote_tp_offset == expected_offset for pull in pulls)) @@ -3385,8 +3278,8 @@ def test_pd_disaggregated_split_cross_covers_prefix_tp_cp_pp(self): "prefill_tp_size": 8, "decode_tp_size": 4, "prefill_pp_size": 2, - "remote_pcp_size": 2, - "remote_dcp_size": 2, + "remote_pcp_size": 1, + "remote_dcp_size": 4, "pcp_size": 1, "dcp_size": 2, "remote_block_ids": ([10, 11], [10, 11]), @@ -3461,6 +3354,49 @@ def test_pd_disaggregated_split_cross_covers_prefix_tp_cp_pp(self): self.assertEqual(sum(len(ids[0]) for ids in local_ids), per_rank_external_blocks) self._assert_group_pull_finish_flags(ports, group_pulls, {0, 1}) + def test_pd_disaggregated_hybrid_dcp_shards_attention_and_final_state(self): + case = dict( + use_mla=False, + num_key_value_heads=1, + prefill_tp_size=4, + decode_tp_size=2, + prefill_pp_size=1, + pcp_size=1, + dcp_size=1, + ) + for tp_rank in range(2): + with self.subTest(tp_rank=tp_rank): + worker = self._build_worker_for_pd_case(case, tp_rank) + worker._is_hma_required = True + worker.kv_group2layeridx[1] = ({"kv_cache_spec_type": "MambaSpec"}, [1]) + meta = types.SimpleNamespace( + remote_pcp_size=1, + remote_dcp_size=2, + remote_ptp_size=4, + remote_port=31000, + remote_block_ids=([10, 11, 12], [30]), + local_block_ids=([20, 21, 22, 23, 24], [40]), + num_external_tokens=80, + num_prompt_blocks=5, + num_computed_tokens=0, + remote_block_size=16, + remote_engine_id="prefill", + remote_host="localhost", + remote_multi_nodes_meta_mapping={}, + ) + + ports, local_ids, remote_ids = worker._get_kv_split_metadata("req-dcp", cast(ReqMeta, meta)) + group_pulls = worker._get_group_pulls_metadata("req-dcp", ports, 4, 31000, 1, 2) + + self.assertEqual(sorted(block for ids in local_ids for block in ids[0]), [20, 21, 22, 23, 24]) + self.assertEqual([ids[1] for ids in local_ids], [[], [40]]) + self.assertEqual([ids[1] for ids in remote_ids], [[], [30]]) + self._assert_hybrid_group_pull_finish_flags(ports, group_pulls, {0, 1}, {0: 2, 1: 1}) + state_ports = { + port for port, pulls in zip(ports[-1], group_pulls[-1]) if any(pull.group_id == 1 for pull in pulls) + } + self.assertEqual(state_ports, {31000 + tp_rank * 2, 31001 + tp_rank * 2}) + def test_pd_disaggregated_hybrid_prefix_tp_and_pp_unequal(self): for tp_rank in range(2): with self.subTest(tp_rank=tp_rank): @@ -3527,6 +3463,8 @@ def test_pd_disaggregated_hybrid_prefix_tp_and_pp_unequal(self): self.assertEqual(local_ids, [([70, 71], [80, 81, 82])]) self.assertEqual(remote_ids, [([50, 51], [60, 61, 62])]) self.assertGreater(len(ports[0]), 1) + for port, pulls in zip(ports[0], group_pulls[0]): + self.assertTrue(all(pull.prefill_pp_rank == (port - 31000) // 4 for pull in pulls)) self._assert_hybrid_group_pull_finish_flags( ports, group_pulls, @@ -3534,83 +3472,125 @@ def test_pd_disaggregated_hybrid_prefix_tp_and_pp_unequal(self): expected_finishes={0: worker._prefill_pp_size, 1: worker._prefill_pp_size}, ) - def test_pd_disaggregated_hybrid_remote_pcp_splits_attention_and_final_mamba_state(self): - for tp_rank in range(2): - with self.subTest(tp_rank=tp_rank): - with patch.object( - self.vllm_config.kv_transfer_config, - "get_from_extra_config", - side_effect=lambda k, d=None: { - "prefill": {"tp_size": 4, "dp_size": 1, "pp_size": 1}, - "decode": {"tp_size": 2, "dp_size": 1, "pp_size": 1}, - }.get(k, d), - ): - self.vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False - self.vllm_config.model_config.is_deepseek_mla = False - self.vllm_config.model_config.hf_text_config.num_key_value_heads = 8 - worker = MooncakeConnectorWorker(self.vllm_config, self.engine_id, MockKVCacheConfig()) - - worker._is_hma_required = True - worker.use_mla = False - worker.use_sparse = False - worker.num_key_value_heads = 8 - worker.tp_size = 2 - worker.tp_rank = tp_rank - worker.pcp_size = 1 - worker.dcp_size = 1 - worker.pcp_rank = 0 - worker.dcp_rank = 0 - worker._decode_tp_size = 2 - worker._prefill_tp_size = 4 - worker._prefill_pp_size = 1 - worker.side_channel_port = 5000 - worker.handshake_port = worker.side_channel_port + tp_rank - worker.local_remote_block_port_mapping = {} - worker.remote_port_send_num = {} - worker.block_size_scale = [[1], [1], [1]] - worker.kv_group2layeridx = { - 0: ( - { - "kv_cache_spec_type": "FullAttentionSpec", - "kv_cache_spec": {"num_kv_heads": 8}, - }, - [0, 1], - ), - 1: ({"kv_cache_spec_type": "MambaSpec"}, [2]), - } - - meta = types.SimpleNamespace( - remote_pcp_size=2, - remote_dcp_size=1, - remote_ptp_size=4, - remote_port=31000, - remote_block_ids=([50, 51, 52, 53], [60, 61, 62, 63]), - local_block_ids=([70, 71, 72, 73], [80, 81, 82, 83]), - num_external_tokens=4 * worker.block_size, - num_prompt_blocks=4, - remote_engine_id=f"remote_hybrid_pcp_{tp_rank}", - remote_host="localhost", - remote_multi_nodes_meta_mapping={}, - remote_block_size=16, - ) - ports, local_ids, remote_ids = worker._get_kv_split_metadata("req_hybrid_pcp", cast(ReqMeta, meta)) - group_pulls = worker._get_group_pulls_metadata( - "req_hybrid_pcp", ports, 4, 31000, meta.remote_pcp_size, meta.remote_dcp_size - ) - - self.assertEqual(len(ports), 2) - # Attention (group 0) is expanded in metadata (scale 1): the 4 external - # blocks are interleaved across the 2 PCP shards, 2 kernels each. - self.assertEqual([len(ids[0]) for ids in local_ids], [2, 2]) - self.assertEqual([ids[1] for ids in local_ids], [[], [80, 81, 82, 83]]) - self.assertEqual([ids[1] for ids in remote_ids], [[], [60, 61, 62, 63]]) - self.assertTrue(worker.remote_port_send_num[meta.remote_engine_id]) - self._assert_hybrid_group_pull_finish_flags( - ports, - group_pulls, - expected_group_ids={0, 1}, - expected_finishes={0: 2, 1: 1}, + def test_start_load_kv_replica_routing_and_completion(self): + """D pulls and P completion tracking must agree across layouts and batches.""" + cases = [ + (1, 2, 2, "gqa"), + (2, 2, 2, "gqa"), + (2, 2, 1, "gqa"), + (1, 4, 2, "mla"), + (2, 4, 2, "mla"), + (1, 4, 2, "hybrid"), + (2, 4, 2, "hybrid"), + ] + for pcp_size, prefill_tp_size, decode_tp_size, layout in cases: + with self.subTest( + pcp_size=pcp_size, prefill_tp_size=prefill_tp_size, decode_tp_size=decode_tp_size, layout=layout + ): + case = dict( + use_mla=layout != "gqa", + num_key_value_heads=8 if layout == "gqa" else 1, + prefill_tp_size=prefill_tp_size, + decode_tp_size=decode_tp_size, + prefill_pp_size=1, + pcp_size=1, + dcp_size=1, ) + decoders = [self._build_worker_for_pd_case(case, rank) for rank in range(decode_tp_size)] + senders = {} + for pcp_rank in range(pcp_size): + for tp_rank in range(prefill_tp_size): + sender = self._build_worker_for_pd_case( + {**case, "decode_tp_size": prefill_tp_size, "pcp_size": pcp_size}, tp_rank, pcp_rank + ) + sender._decode_tp_size = decode_tp_size + sender.kv_send_thread = KVCacheSendingThread.__new__(KVCacheSendingThread) + sender.kv_send_thread.task_tracker = KVCacheTaskTracker() + senders[31000 + pcp_rank * prefill_tp_size + tp_rank] = sender + + for worker in [*decoders, *senders.values()]: + worker._is_hma_required = layout == "hybrid" + worker.enable_sfa_dcp_replicated_indexer = False + worker.kv_group2layeridx[0][0]["kv_cache_spec"] = {"num_kv_heads": case["num_key_value_heads"]} + if layout == "hybrid": + worker.kv_group2layeridx[1] = ({"kv_cache_spec_type": "MambaSpec"}, [1]) + worker.block_size_scale = [[2], [1]] + for worker in decoders: + worker.kv_recv_thread = MagicMock() + + # Reuse workers across requests; D and P have different request IDs. + for request_id in ("req-0", "req-1", "req-2", "req-3"): + with self.subTest(request_id=request_id): + metadata = MooncakeConnectorMetadata() + metadata.add_new_req( + request_id=f"decode-{request_id}", + local_block_ids=([20, 21], [40, 41, 42]), + num_external_tokens=32, + kv_transfer_params=dict( + remote_request_id=request_id, + remote_engine_id="prefill", + remote_host="prefill-host", + remote_port=31000, + remote_pcp_size=pcp_size, + remote_dcp_size=1, + remote_ptp_size=prefill_tp_size, + remote_block_ids=([10, 11, 12], [30, 31, 32]), + remote_block_size=16, + num_prompt_blocks=3, + num_computed_tokens=16, + ), + ) + metadata.reqs_in_batch = {f"decode-{request_id}"} + source_ports = set() + for decoder in decoders: + decoder.kv_recv_thread.reset_mock() + decoder.start_load_kv(metadata) + pulls = [call.kwargs for call in decoder.kv_recv_thread.add_request.call_args_list] + self.assertTrue(pulls) + self.assertEqual(sum(pull["all_task_done"] for pull in pulls), 1) + group_pulls = [group for pull in pulls for group in pull["group_pulls"]] + for group_id in (0, 1): + group = [pull for pull in group_pulls if pull.group_id == group_id] + num_pulls = ( + prefill_tp_size // decode_tp_size + if layout == "gqa" or (layout == "hybrid" and group_id == 1) + else 1 + ) + self.assertEqual( + sorted((pull.remote_tp_offset, pull.is_group_transfer_end) for pull in group), + [(offset, offset == num_pulls - 1) for offset in range(num_pulls)], + ) + for pull in pulls: + source_ports.add(pull["remote_handshake_port"]) + self.assertEqual(pull["remote_request_id"], request_id) + self.assertEqual( + pull["local_block_ids"], + ([40, 41, 42, 43], [40, 41, 42] if layout == "hybrid" else [40, 41]), + ) + self.assertEqual( + pull["remote_block_ids"], + ([22, 23, 24, 25], [30, 31, 32] if layout == "hybrid" else [31, 32]), + ) + self.assertIsNone(pull["remote_port_send_num"]) + self.assertEqual(pull["shard_idx"], 0) + + self.assertEqual(len({(port - 31000) // prefill_tp_size for port in source_ports}), 1) + self.assertTrue(source_ports.issubset(senders)) + send_metadata = MooncakeConnectorMetadata() + send_metadata.reqs_in_batch = {request_id} + send_metadata.requests_to_send = {request_id: time.time()} + for port, sender in senders.items(): + sender.start_load_kv(send_metadata) + tracker = sender.kv_send_thread.task_tracker + if port in source_ports: + self.assertEqual(tracker.get_and_clear_finished_requests(), set()) + self.assertEqual(set(tracker.delayed_free_requests), {request_id}) + # The selected source completes only after D's DONE message. + tracker.update_done_task_count(request_id) + self.assertEqual(tracker.get_and_clear_finished_requests(), {request_id}) + self.assertEqual(tracker.get_and_clear_finished_requests(), set()) + self.assertFalse(tracker.delayed_free_requests) + self.assertFalse(tracker.reqs_to_process) def test_hybrid_no_cp_uses_kv_cache_group_ids_for_split_transfer_groups(self): with patch.object( @@ -3787,6 +3767,8 @@ def test_start_load_kv_puts_replicated_indexer_on_existing_transfer_port(self): worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) worker.kv_send_thread = None worker.kv_recv_thread = MagicMock() + worker.pcp_size = 1 + worker.dcp_size = 1 worker._prefill_tp_size = 4 worker.remote_port_send_num = {"remote_engine": {31001: {"num": 1, "host": "localhost"}}} worker._get_sfa_replicate_k_block_ids = MagicMock(return_value=(([40],), ([20],))) @@ -3809,7 +3791,7 @@ def test_start_load_kv_puts_replicated_indexer_on_existing_transfer_port(self): remote_engine_id="remote_engine", remote_host="localhost", remote_port=31000, - remote_pcp_size=2, + remote_pcp_size=1, remote_dcp_size=2, remote_ptp_size=4, remote_multi_nodes_meta_mapping={}, diff --git a/tests/ut/kv_offload/test_mooncake_hybrid_connector.py b/tests/ut/kv_offload/test_mooncake_hybrid_connector.py index 6623465b387b..1b12e6db787a 100644 --- a/tests/ut/kv_offload/test_mooncake_hybrid_connector.py +++ b/tests/ut/kv_offload/test_mooncake_hybrid_connector.py @@ -1,3 +1,4 @@ +import queue import sys import threading import time @@ -14,11 +15,16 @@ fake_engine.TransferEngine = MagicMock() # type: ignore[attr-defined] sys.modules["mooncake.engine"] = fake_engine +from vllm.v1.kv_cache_interface import MambaSpec # noqa: E402 from vllm.v1.request import RequestStatus # noqa: E402 from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_hybrid_connector import ( # noqa: E402 MAX_REQUESTS_PER_PEER_HANDLER, KVCacheRecvingThread, + KVCacheSendingThread, + KVCacheTaskTracker, + MooncakeAgentMetadata, + MooncakeConnectorMetadata, MooncakeConnectorScheduler, MooncakeConnectorWorker, ) @@ -55,6 +61,78 @@ def _make_thread(self): thread.request_task_counts_lock = threading.Lock() return thread + def test_group_transfer_and_completion(self): + thread = self._make_thread() + self.addCleanup(thread.executor.shutdown, wait=True) + thread.use_hybrid = True + thread.tp_rank = 0 + thread._prefill_pp_size = 1 + thread.hma_group_size = 2 + thread.kv_cache_specs = [MagicMock(), MagicMock(spec=MambaSpec)] + thread.task_tracker = KVCacheTaskTracker() + thread.request_queue = queue.Queue() + thread.proc_not_transfer_request = {} + thread.proc_not_transfer_request_lock = threading.Lock() + thread.side_channel_port = thread.local_handshake_port = 32000 + thread.local_engine_id = "decode" + thread.remote_metadata_lock = threading.Lock() + thread.kv_caches_base_addr = { + "decode": {32000: [0x1000, 0x2000, 0x3000]}, + "prefill": {31002: [0x4000, 0x5000, 0x6000]}, + } + thread.remote_te_port = {"prefill": {31002: 7777}} + # Two buffers share the first group's blocks; the state group has its own stride. + thread.addr_group_idx = [[0], [0], [1]] + thread.block_len_per_addr = [16, 8, 32] + thread.block_stride_per_addr = [32, 16, 64] + thread.engine = MagicMock() + thread._send_done_recv_signal = MagicMock() + + for outcome in ("success", "empty", "failure"): + with self.subTest(outcome=outcome): + request_id = f"decode-{outcome}" + remote_request_id = f"prefill-{outcome}" + thread.engine.reset_mock() + thread.engine.batch_transfer_sync_read.return_value = -1 if outcome == "failure" else 0 + thread._send_done_recv_signal.reset_mock() + thread.task_tracker.add_req_to_process(request_id) + thread.add_request( + request_id=request_id, + remote_request_id=remote_request_id, + local_block_ids=([], []) if outcome == "empty" else ([2, 3], [4]), + remote_block_ids=([1, 2], [3]), + remote_engine_id="prefill", + remote_host="192.0.2.1", + remote_handshake_port=31002, + offset=0, + tp_num_need_pulls=1, + all_task_done=True, + ) + req_meta = thread.request_queue.get_nowait() + thread._mark_request_task_submitted(req_meta) + with patch( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_hybrid_connector.logger.exception" + ) as log_exception: + thread._handle_request(req_meta) + self.assertEqual(log_exception.call_count, int(outcome == "failure")) + if outcome == "empty": + thread.engine.batch_transfer_sync_read.assert_not_called() + else: + thread.engine.batch_transfer_sync_read.assert_called_once_with( + "192.0.2.1:7777", + [0x1040, 0x2020, 0x3100], + [0x4020, 0x5010, 0x60C0], + [32, 16, 32], + ) + thread._send_done_recv_signal.assert_called_once_with(remote_request_id, "192.0.2.1", 31002, {}) + self.assertEqual(thread.get_and_clear_finished_requests(), {request_id}) + self.assertEqual(thread.get_and_clear_finished_requests(), set()) + self.assertFalse(thread.task_tracker.reqs_to_process) + self.assertFalse(thread.request_task_counts) + self.assertFalse(thread.finished_request_markers) + self.assertFalse(thread.proc_not_transfer_request) + self.assertEqual(thread.request_queue.unfinished_tasks, 0) + def test_executor_workers_bind_kv_cache_device_before_handling_requests(self): expected_device_index = 5 kv_cache = MagicMock(device=expected_device_index) @@ -242,6 +320,205 @@ def handle_request(req_meta: dict[str, Any]): thread.executor.submit.assert_called_once_with(thread._handle_peer_requests, peer_key) +class TestMooncakeHybridConnectorWorker(unittest.TestCase): + def setUp(self): + for patcher in ( + patch.dict("os.environ"), + patch.multiple( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_hybrid_connector", + init_ascend_config=MagicMock(), + get_ascend_config=MagicMock(), + get_transfer_timeout_value=MagicMock(return_value=30), + get_ip=MagicMock(return_value="127.0.0.1"), + get_tp_group=MagicMock(), + get_pp_group=MagicMock(return_value=types.SimpleNamespace(rank_in_group=0)), + global_te=MagicMock(), + ), + ): + patcher.start() + self.addCleanup(patcher.stop) + + def _make_worker(self, prefill_tp_size, decode_tp_size, tp_rank, pcp_size, pcp_rank, dp_rank, use_mamba, role): + tp_size = prefill_tp_size if role == "kv_producer" else decode_tp_size + dp_size = 2 if role == "kv_producer" else 1 + extra_config = { + "prefill": {"tp_size": prefill_tp_size, "dp_size": 2}, + "decode": {"tp_size": decode_tp_size, "dp_size": 1}, + } + config = types.SimpleNamespace( + parallel_config=types.SimpleNamespace( + tensor_parallel_size=tp_size, + prefill_context_parallel_size=pcp_size, + decode_context_parallel_size=1, + pipeline_parallel_size=1, + data_parallel_rank=dp_rank, + data_parallel_size=dp_size, + data_parallel_rank_local=dp_rank, + data_parallel_size_local=dp_size, + ), + kv_transfer_config=types.SimpleNamespace( + kv_role=role, + kv_port=31000 if role == "kv_producer" else 32000, + get_from_extra_config=extra_config.get, + ), + model_config=types.SimpleNamespace( + is_deepseek_mla=not use_mamba, + hf_config=types.SimpleNamespace(**({} if use_mamba else {"compress_ratios": [1, 4]})), + hf_text_config=types.SimpleNamespace(num_key_value_heads=8), + ), + cache_config=types.SimpleNamespace(block_size=128), + scheduler_config=types.SimpleNamespace(disable_hybrid_kv_cache_manager=False), + ) + state_spec = ( + MagicMock(spec=MambaSpec, block_size=128, shapes=((2, 2), (2, 2)), dtypes=(torch.float32, torch.float32)) + if use_mamba + else types.SimpleNamespace(block_size=128) + ) + cache_config = types.SimpleNamespace( + kv_cache_groups=[ + types.SimpleNamespace(kv_cache_spec=types.SimpleNamespace(block_size=512), layer_names=["layer.0"]), + types.SimpleNamespace(kv_cache_spec=state_spec, layer_names=["layer.1"]), + ] + ) + host = f"192.0.2.{pcp_rank * tp_size + tp_rank + 1}" + engine_id = f"{role}-{dp_rank}-{pcp_rank}-{tp_rank}" + with patch.multiple( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_hybrid_connector", + get_tensor_model_parallel_rank=MagicMock(return_value=tp_rank), + get_tensor_model_parallel_world_size=MagicMock(return_value=tp_size), + get_pcp_group=MagicMock(return_value=types.SimpleNamespace(rank_in_group=pcp_rank, world_size=pcp_size)), + get_ip=MagicMock(return_value=host), + ): + worker = MooncakeConnectorWorker(config, engine_id, cache_config) + worker.use_sparse = False + if role == "kv_producer": + metadata = MooncakeAgentMetadata(engine_id, 7777, 128, [], 0, [], (0, 0), local_ip=host) + worker.kv_send_thread = KVCacheSendingThread( + config, + tp_rank, + prefill_tp_size, + engine_id, + host, + worker.side_channel_port, + metadata, + threading.Event(), + {}, + pcp_rank, + ) + else: + worker.kv_recv_thread = MagicMock() + return worker + + def test_start_load_kv_replica_routing_and_completion(self): + cases = [ + # PCP, P-TP, D-TP, P-DP rank, Mamba receive branch. + (1, 2, 2, 0, False), + (2, 2, 2, 0, False), + (1, 4, 2, 1, False), + (4, 4, 2, 1, False), + (1, 2, 2, 1, True), + (2, 2, 2, 1, True), + ] + for pcp_size, prefill_tp_size, decode_tp_size, dp_rank, use_mamba in cases: + with self.subTest(pcp_size=pcp_size, prefill_tp_size=prefill_tp_size, mamba=use_mamba): + senders: dict[int, MooncakeConnectorWorker] = {} + handshake_metadata = {} + for pcp_rank in range(pcp_size): + for tp_rank in range(prefill_tp_size): + worker = self._make_worker( + prefill_tp_size, + decode_tp_size, + tp_rank, + pcp_size, + pcp_rank, + dp_rank, + use_mamba, + "kv_producer", + ) + self.assertNotIn(worker.handshake_port, senders) + senders[worker.handshake_port] = worker + key = (0, tp_rank) if pcp_size == 1 else (0, pcp_rank, tp_rank) + handshake_metadata[key] = worker.kv_send_thread.metadata + with ( + patch( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_hybrid_connector.zmq_ctx" + ) as ctx, + patch.object(worker.kv_send_thread, "run_busy_loop") as run_busy_loop, + ): + worker.kv_send_thread.run() + self.assertEqual( + ctx.call_args.args[1], f"tcp://{worker.side_channel_host}:{worker.handshake_port}" + ) + run_busy_loop.assert_called_once() + + scheduler = MooncakeConnectorScheduler(worker.vllm_config, "prefill", worker.kv_cache_config) + scheduler.set_xfer_handshake_metadata_from_workers(handshake_metadata) + base_port = 31000 + dp_rank * prefill_tp_size * pcp_size + self.assertEqual(scheduler.side_channel_port, base_port) + self.assertEqual(set(senders), set(range(base_port, base_port + prefill_tp_size * pcp_size))) + decoders = [ + self._make_worker(prefill_tp_size, decode_tp_size, rank, 1, 0, 0, use_mamba, "kv_consumer") + for rank in range(decode_tp_size) + ] + + # Reuse the workers across requests, with different P and D request IDs. + for request_id in ("req-0", "req-1", "req-2", "req-3"): + with self.subTest(request_id=request_id): + request = MockRequest( + request_id, + list(range(513)), + {"do_remote_decode": True}, + RequestStatus.FINISHED_LENGTH_CAPPED, + ) + delay_free, params = scheduler.request_finished_all_groups(request, ([10, 11, 12], [30, 31])) + self.assertTrue(delay_free) + self.assertEqual(params["remote_block_ids"], ([10, 11], [30, 31])) + if pcp_size == 1: + # Old request metadata defaults to one replica. + params.pop("remote_pcp_size") + metadata = MooncakeConnectorMetadata() + metadata.add_new_req(f"decode-{request_id}", ([20, 21], [40, 41]), 513, params) + metadata.reqs_in_batch = {f"decode-{request_id}"} + source_ports = set() + for decoder in decoders: + decoder.kv_recv_thread.reset_mock() + decoder.start_load_kv(metadata) + decoder.kv_recv_thread.add_request.assert_called_once() + pull = decoder.kv_recv_thread.add_request.call_args.kwargs + port = pull["remote_handshake_port"] + self.assertIn(port, senders) + source_ports.add(port) + self.assertEqual(pull["remote_host"], senders[port].side_channel_host) + self.assertEqual(pull["remote_engine_id"], senders[port].engine_id) + self.assertEqual(pull["request_id"], f"decode-{request_id}") + self.assertEqual(pull["remote_request_id"], request_id) + self.assertEqual(pull["local_block_ids"], ([20, 21], [40, 41])) + self.assertEqual(pull["remote_block_ids"], ([10, 11], [30, 31])) + self.assertEqual( + (pull["offset"], pull["tp_num_need_pulls"], pull["all_task_done"]), (0, 1, True) + ) + self.assertIsNone(pull.get("remote_port_send_num")) + if prefill_tp_size == decode_tp_size: + self.assertEqual((port - base_port) % prefill_tp_size, decoder.tp_rank) + self.assertEqual(len(source_ports), decode_tp_size) + self.assertEqual(len({(port - base_port) // prefill_tp_size for port in source_ports}), 1) + + scheduler._reqs_in_batch.add(request_id) + send_metadata = scheduler.build_connector_meta(MagicMock()) + for port, sender in senders.items(): + sender.start_load_kv(send_metadata) + tracker = sender.kv_send_thread.task_tracker + if port in source_ports: + self.assertEqual(tracker.get_and_clear_finished_requests(), set()) + self.assertEqual(set(tracker.delayed_free_requests), {request_id}) + # Only actual sources wait for D's DONE message. + tracker.update_done_task_count(request_id) + self.assertEqual(tracker.get_and_clear_finished_requests(), {request_id}) + self.assertEqual(tracker.get_and_clear_finished_requests(), set()) + self.assertFalse(tracker.delayed_free_requests) + self.assertFalse(tracker.reqs_to_process) + + class TestMooncakeHybridConnectorRegistration(unittest.TestCase): def test_hybrid_registration_uses_actual_merged_tensor_ranges(self): alignment = 2 * 1024 * 1024 @@ -313,6 +590,7 @@ def _make_scheduler(self): scheduler.side_channel_host = "127.0.0.1" scheduler.side_channel_port = 12345 scheduler.tp_size = 1 + scheduler.pcp_size = 1 scheduler.multi_nodes_meta_mapping = {} return scheduler @@ -324,10 +602,10 @@ def test_compute_transfer_block_ids_trims_swa_groups(self): self.assertEqual(transfer_block_ids, ([0], [100, 101])) - def test_request_finished_trims_logical_compressed_group_spans(self): + def test_request_finished_preserves_group_layout_with_pcp(self): scheduler = self._make_scheduler() - scheduler.group_block_size = [512, 16384] - scheduler.num_swa_blocks = [0, 0] + scheduler.group_block_size = [512, 16384, 128] + scheduler.num_swa_blocks = [0, 0, 2] request = MockRequest( "req-compressed", prompt_token_ids=list(range(513)), @@ -335,17 +613,22 @@ def test_request_finished_trims_logical_compressed_group_spans(self): status=RequestStatus.FINISHED_LENGTH_CAPPED, ) - delay_free, params = scheduler.request_finished_all_groups( - request, - ([10, 11, 12], [20, 21]), - ) - - self.assertTrue(delay_free) - self.assertIsNotNone(params) - assert params is not None - self.assertEqual(params["remote_block_ids"], ([10, 11], [20])) - # This unused compatibility field remains in the legacy physical-block unit. - self.assertEqual(params["num_prompt_blocks"], 5) + for pcp_size in (1, 2, 4): + with self.subTest(pcp_size=pcp_size): + scheduler.pcp_size = pcp_size + delay_free, params = scheduler.request_finished_all_groups( + request, + ([10, 11, 12], [20, 21], [30, 31, 32, 33, 34, 35]), + ) + self.assertTrue(delay_free) + self.assertIsNotNone(params) + assert params is not None + self.assertEqual(params["remote_block_ids"], ([10, 11], [20], [33, 34])) + self.assertEqual(params["remote_pcp_size"], pcp_size) + self.assertEqual(params["remote_ptp_size"], scheduler.tp_size) + self.assertIn(request.request_id, scheduler._reqs_need_send) + # This unused compatibility field remains in the legacy physical-block unit. + self.assertEqual(params["num_prompt_blocks"], 5) def test_request_finished_trims_before_swa_clip(self): scheduler = self._make_scheduler() diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py index 0405804420fe..3b1cb2ff9722 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py @@ -1822,7 +1822,6 @@ def _get_transfer_block_ids(self, block_ids: BlockIds, prompt_len: int) -> Block assert len(block_ids) == len(self.group_transfer_info), "Number of KV cache groups must match" transfer_block_ids = [] - cp_size = max(1, self.pcp_size * self.dcp_size) for blocks, group_info in zip(block_ids, self.group_transfer_info): is_aligned_state_group = group_info.is_state_group and ( getattr(self.vllm_config.cache_config, "mamba_cache_mode", None) == "align" @@ -1830,9 +1829,9 @@ def _get_transfer_block_ids(self, block_ids: BlockIds, prompt_len: int) -> Block if group_info.is_state_group and not is_aligned_state_group: transfer_block_ids.append(blocks) elif is_aligned_state_group: - # Mamba state is not CP-sharded like attention KV. Its aligned + # Mamba state is not DCP-sharded like attention KV. Its aligned # block index is derived from the actual (already truncated) - # prompt length, without multiplying by the CP size. + # prompt length, without multiplying by the DCP size. num_prompt_state_blocks = cdiv(prompt_len, group_info.tokens_per_block) if num_prompt_state_blocks <= 0 or num_prompt_state_blocks > len(blocks): raise RuntimeError( @@ -1842,10 +1841,9 @@ def _get_transfer_block_ids(self, block_ids: BlockIds, prompt_len: int) -> Block ) transfer_block_ids.append(blocks[num_prompt_state_blocks - 1 : num_prompt_state_blocks]) else: - # In context parallelism, each scheduler-visible block id is a - # CP-grouped/virtual block shared by all CP ranks. It therefore - # covers cp_size times the token span of one no-CP block. - num_prompt_blocks = cdiv(prompt_len, group_info.tokens_per_block * cp_size) + # Each scheduler-visible block id is a DCP-grouped virtual + # block shared by all DCP ranks. + num_prompt_blocks = cdiv(prompt_len, group_info.tokens_per_block * self.dcp_size) transfer_block_ids.append(blocks[:num_prompt_blocks]) return tuple(transfer_block_ids) @@ -2140,6 +2138,11 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str, kv_cache_config: KVC self.max_device_id = self.tp_size * self.dp_size * self.pcp_size * self.pp_size self.kv_role = vllm_config.kv_transfer_config.kv_role + if self.kv_role == "kv_consumer" and self.pcp_size > 1: + raise ValueError( + "In P/D disaggregation, Mooncake supports PCP only on the prefill (kv_producer) engine. " + "Set prefill_context_parallel_size=1 on the decode (kv_consumer) engine." + ) self.num_key_value_heads = self.vllm_config.model_config.hf_text_config.num_key_value_heads # kv cache config @@ -2936,21 +2939,24 @@ def _get_kv_split_metadata( * remote_block_ids_list[i]: remote kernel block ids, grouped by KV cache group, where blocks are read from. - In PCP/DCP scenarios, prompt blocks can be split across multiple remote - P workers. This method also accounts for unequal P/D prefix-cache hits - by reducing the number of remote blocks that still need to be pulled. + PCP selects one complete P-side KV replica; DCP splits prompt blocks + across P workers. Unequal P/D prefix-cache hits reduce the number of + remote blocks that still need to be pulled. """ prefill_tp_size: int = meta.remote_ptp_size if meta.remote_ptp_size is not None else self._prefill_tp_size - if meta.remote_pcp_size * meta.remote_dcp_size * self.pcp_size * self.dcp_size == 1: + if self.dcp_size == meta.remote_dcp_size == 1: if self._is_hma_required: chosen_rank_list, _ = self._get_hybrid_remote_rank_group_pulls(req_id, prefill_tp_size) else: chosen_rank_list = self._get_remote_rank(req_id, prefill_tp_size) - remote_handshake_port_list = [[x + meta.remote_port for x in chosen_rank_list]] - # No CP: expand logical blocks into kernel blocks here so the transfer - # stage consumes kernel-level ids directly (chunk_starts no longer needed). + # Select the same TP rank in the chosen PCP replica. + # E.g. TP2/PP1, PCP rank 1, TP rank 1: offset = 2, port = base + 3. + pcp_offset = self._get_selected_pcp_rank(req_id, meta.remote_pcp_size) * prefill_tp_size + remote_handshake_port_list = [[x + meta.remote_port + pcp_offset for x in chosen_rank_list]] + # Complete KV replicas use the same logical-to-kernel block mapping + # as the non-CP path. use_transfer_group_block_ids = transfer_groups_need_independent_block_ids( self.kv_group2layeridx, self.block_size_scale, @@ -3336,9 +3342,8 @@ def _set_hma_shared_port(prefill_tp_size, meta, remote_handshake_port_list, req_ return remote_handshake_port_list, local_block_ids_list, remote_block_ids_list - def _get_cp_shard_pulls(self, remote_handshake_port_list, prefill_tp_size, remote_base_port, remote_pcp_size): - # CP case: `group_pulls` is derived from `port` (which already includes the random selection result), - # eliminating the need for a table lookup. + def _get_dcp_shard_pulls(self, remote_handshake_port_list, prefill_tp_size, remote_base_port): + """Build group pulls from the selected ports of each DCP shard.""" mamba_num = prefill_tp_size // self.tp_size attn_num = self._get_tp_num_need_pulls(prefill_tp_size) attn_gids = [ @@ -3355,9 +3360,8 @@ def _get_cp_shard_pulls(self, remote_handshake_port_list, prefill_tp_size, remot for port_idx, port in enumerate(ports): pulls = [] port_tp = (port - remote_base_port) % prefill_tp_size - # PCP and PP are mutually exclusive; when PCP > 1, pp_rank is always 0. - pp_rank = 0 if remote_pcp_size > 1 else (port - remote_base_port) // prefill_tp_size - # The first attn_num ports of each shard (i.e., the original ports with randomly substituted TPs). + pp_rank = (port - remote_base_port) // prefill_tp_size + # Attention uses the leading ports selected for each DCP shard. if port_idx < attn_num: pulls += [ GroupPull( @@ -3369,7 +3373,7 @@ def _get_cp_shard_pulls(self, remote_handshake_port_list, prefill_tp_size, remot ) for g in attn_gids ] - # Mamba: Only applicable to the final shard; the offset is back-calculated from the port's TP ID. + # Transfer Mamba state only on the final DCP shard, using the port's TP rank to derive its offset. if is_final: m_off = port_tp - self.tp_rank * mamba_num if 0 <= m_off < mamba_num: @@ -3419,19 +3423,21 @@ def _get_group_pulls_metadata( this pull is the final pull for the group. The final-pull flag is used by the receiver to decide when group reformatting can run. """ - cp_transfer = remote_pcp_size * remote_dcp_size * self.pcp_size * self.dcp_size > 1 + dcp_transfer = remote_dcp_size * self.dcp_size > 1 if self._is_hma_required: - if not cp_transfer: - # Non-CP case: port = base + chosen_rank, which has a one-to-one correspondence - # with the table keys, maintaining the original logic. + if not dcp_transfer: + # The table uses TP/PP ranks without PCP replica offsets. + # Undo the offset added by _get_kv_split_metadata, keeping the PP stage. + # E.g. TP2/PP1, PCP rank 1: port base + 3 maps back to rank 3 - 2 = 1. _, rank_group_pulls = self._get_hybrid_remote_rank_group_pulls(req_id, prefill_tp_size) - return [[rank_group_pulls[p - remote_base_port] for p in ports] for ports in remote_handshake_port_list] + pcp_offset = self._get_selected_pcp_rank(req_id, remote_pcp_size) * prefill_tp_size + return [ + [rank_group_pulls[p - remote_base_port - pcp_offset] for p in ports] + for ports in remote_handshake_port_list + ] - # CP case: `group_pulls` is derived from `port` (which already includes the random selection result), - # eliminating the need for a table lookup. - return self._get_cp_shard_pulls( - remote_handshake_port_list, prefill_tp_size, remote_base_port, remote_pcp_size - ) + # The DCP path has already selected the source ports for each shard. + return self._get_dcp_shard_pulls(remote_handshake_port_list, prefill_tp_size, remote_base_port) tp_num_need_pulls = self._get_tp_num_need_pulls(prefill_tp_size) group_ids = [group_id for group_id, (_, layer_indices) in self.kv_group2layeridx.items() if layer_indices] @@ -3449,9 +3455,9 @@ def make_group_pulls(remote_tp_offset: int, prefill_pp_rank: int) -> list[GroupP ] group_pulls_list = [] - for pcp_dcp_rank, remote_ports in enumerate(remote_handshake_port_list): + for shard_idx, remote_ports in enumerate(remote_handshake_port_list): if len(remote_ports) == 1: - remote_tp_offsets = [pcp_dcp_rank % tp_num_need_pulls] + remote_tp_offsets = [shard_idx % tp_num_need_pulls] prefill_pp_ranks = [ ((remote_ports[0] - remote_base_port) % (prefill_tp_size * self._prefill_pp_size)) // prefill_tp_size @@ -3743,7 +3749,7 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata): meta.remote_dcp_size, ) - for pcp_dcp_rank, remote_ports in enumerate(remote_handshake_port_list): + for shard_idx, remote_ports in enumerate(remote_handshake_port_list): for remote_tp_offset, remote_handshake_port in enumerate(remote_ports): assert self.kv_recv_thread is not None remote_host, remote_engine_id = self._get_remote_host_info_by_port( @@ -3754,9 +3760,7 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata): meta.remote_multi_nodes_meta_mapping, ) remote_port_send_num = ( - self.remote_port_send_num[meta.remote_engine_id] - if meta.remote_pcp_size * meta.remote_dcp_size > 1 - else None + self.remote_port_send_num[meta.remote_engine_id] if meta.remote_dcp_size > 1 else None ) local_block_ids_replicate_k_for_port = ( local_block_ids_replicate_k @@ -3771,32 +3775,35 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata): self.kv_recv_thread.add_request( request_id=req_id, remote_request_id=remote_req_id, - local_block_ids=local_block_ids_list[pcp_dcp_rank], - remote_block_ids=remote_block_ids_list[pcp_dcp_rank], - group_pulls=group_pulls_list[pcp_dcp_rank][remote_tp_offset], + local_block_ids=local_block_ids_list[shard_idx], + remote_block_ids=remote_block_ids_list[shard_idx], + group_pulls=group_pulls_list[shard_idx][remote_tp_offset], remote_engine_id=remote_engine_id, remote_host=remote_host, remote_handshake_port=remote_handshake_port, remote_port_send_num=remote_port_send_num, num_computed_tokens=meta.num_computed_tokens, all_task_done=( - pcp_dcp_rank == len(remote_handshake_port_list) - 1 + shard_idx == len(remote_handshake_port_list) - 1 and remote_tp_offset == len(remote_ports) - 1 ), - shard_idx=pcp_dcp_rank, + shard_idx=shard_idx, remote_block_size=meta.remote_block_size, local_block_ids_replicate_k=local_block_ids_replicate_k_for_port, remote_block_ids_replicate_k=remote_block_ids_replicate_k_for_port, ) - if self.kv_send_thread is not None and self.pcp_size * self.dcp_size == 1: + if self.kv_send_thread is not None and self.dcp_size == 1: for req_id, delay_start_time in metadata.requests_to_send.items(): - if self.tp_rank in self._prefill_get_remote_rank(req_id): + # Unused PCP replicas report completion locally; only transfer + # sources wait for the D-side completion signal. + selected_pcp_rank = self._get_selected_pcp_rank(req_id, self.pcp_size) + if self.pcp_rank == selected_pcp_rank and self.tp_rank in self._prefill_get_remote_rank(req_id): self.kv_send_thread.add_delayed_request(req_id, delay_start_time) else: self.kv_send_thread.add_not_transfer_request(req_id) - if self.kv_send_thread is not None and self.pcp_size * self.dcp_size > 1: + if self.kv_send_thread is not None and self.dcp_size > 1: for req_id, delay_start_time in metadata.requests_to_send.items(): self.kv_send_thread.add_delayed_request(req_id, delay_start_time) @@ -3815,6 +3822,14 @@ def _get_tp_num_need_pulls(self, prefill_tp_size: int | None) -> int: tp_num_need_pulls = num_d_block_heads // num_p_block_heads return tp_num_need_pulls + @staticmethod + def _get_selected_pcp_rank(req_id: str, pcp_size: int) -> int: + if pcp_size == 1: + return 0 + # P and D use the P request ID to select the same replica, independently + # of TP routing. + return random.Random(string_to_int64_hash(f"pcp:{req_id}")).randrange(pcp_size) + def _get_remote_host_info_by_port( self, base_port: int, diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py index 7bff7a86ce4b..00d13923d917 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py @@ -32,6 +32,7 @@ SupportsHMA, ) from vllm.distributed.parallel_state import ( + get_pcp_group, get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -102,6 +103,7 @@ class ReqMeta: remote_ptp_size: int | None remote_multi_nodes_meta_mapping: dict[str, dict[str, Any]] num_prompt_blocks: int + remote_pcp_size: int = 1 @dataclass @@ -215,11 +217,14 @@ def __init__( metadata: MooncakeAgentMetadata, ready_event: threading.Event, kv_caches: dict[str, Any], + pcp_rank: int, ): super().__init__(daemon=True, name="KVCacheSendingThread") self.tp_rank = tp_rank self.prefill_tp_size = prefill_tp_size self.pp_rank = get_pp_group().rank_in_group + self.pcp_size = get_pcp_group().world_size + self.pcp_rank = pcp_rank self.pp_size = vllm_config.parallel_config.pipeline_parallel_size self.tp_size = get_tensor_model_parallel_world_size() self.local_engine_id = local_engine_id @@ -253,7 +258,7 @@ def run(self): # to have a unique port. This hack to keeps us moving. We will # switch when moving to etcd or where we have a single ZMQ socket in # the scheduler. - device_index = self.pp_rank * self.tp_size + self.tp_rank + device_index = (self.pp_rank * self.pcp_size + self.pcp_rank) * self.tp_size + self.tp_rank handshake_port = self.side_channel_port + device_index path = make_zmq_path("tcp", self.side_channel_host, handshake_port) logger.info( @@ -1110,6 +1115,7 @@ def add_new_req( remote_ptp_size=kv_transfer_params.get("remote_ptp_size"), remote_multi_nodes_meta_mapping=kv_transfer_params.get("remote_multi_nodes_meta_mapping", {}), num_prompt_blocks=kv_transfer_params.get("num_prompt_blocks", 0), + remote_pcp_size=kv_transfer_params.get("remote_pcp_size", 1), ) @@ -1234,7 +1240,7 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str, kv_cache_config: KVC self.tp_size = vllm_config.parallel_config.tensor_parallel_size self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size - assert self.pcp_size * self.dcp_size == 1, "Mooncake Hybrid Connector only support cp_world_size == 1. " + assert self.dcp_size == 1, "Mooncake Hybrid Connector requires decode_context_parallel_size=1." self.max_device_id = ( vllm_config.parallel_config.tensor_parallel_size * vllm_config.parallel_config.data_parallel_size @@ -1247,6 +1253,7 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str, kv_cache_config: KVC + vllm_config.parallel_config.data_parallel_rank * vllm_config.parallel_config.tensor_parallel_size * vllm_config.parallel_config.pipeline_parallel_size + * vllm_config.parallel_config.prefill_context_parallel_size ) # Requests that need to start recv. # New requests are added by update_state_after_alloc in @@ -1494,6 +1501,7 @@ def request_finished_all_groups( remote_host=self.side_channel_host, remote_port=self.side_channel_port, remote_ptp_size=self.tp_size, + remote_pcp_size=self.pcp_size, last_token_id=request.output_token_ids[-1], remote_multi_nodes_meta_mapping=self.multi_nodes_meta_mapping, num_prompt_blocks=num_prompt_blocks, @@ -1516,8 +1524,8 @@ def set_xfer_handshake_metadata_from_workers( metadata, ) -> None: """Store worker metadata keyed by handshake port offset - (pp_rank * tp_size + tp_rank), matching how the recv thread resolves - peers in _get_remote_host_info_by_port.""" + ((pp_rank * pcp_size + pcp_rank) * tp_size + tp_rank), matching how + the recv thread resolves peers in _get_remote_host_info_by_port.""" for metadata_key, rank_metadata in metadata.items(): offset = self._port_offset_from_handshake_metadata(rank_metadata, metadata_key) self.multi_nodes_meta_mapping[str(offset)] = { @@ -1563,13 +1571,20 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str, kv_cache_config: KVC self.dp_size = vllm_config.parallel_config.data_parallel_size_local self.pp_size = vllm_config.parallel_config.pipeline_parallel_size self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size + self.pcp_rank = get_pcp_group().rank_in_group self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size - assert self.pcp_size * self.dcp_size == 1, "Mooncake Hybrid Connector only support cp_world_size == 1. " + assert self.dcp_size == 1, "Mooncake Hybrid Connector requires decode_context_parallel_size=1." + assert not (self.pp_size > 1 and self.pcp_size > 1), "pp and pcp cannot open in same time" self.kv_caches: dict[str, torch.Tensor] = {} self.side_channel_host = get_ip() self.max_device_id = self.tp_size * self.dp_size * self.pp_size self.kv_role = vllm_config.kv_transfer_config.kv_role + if self.kv_role == "kv_consumer" and self.pcp_size > 1: + raise ValueError( + "In P/D disaggregation, Mooncake supports PCP only on the prefill (kv_producer) engine. " + "Set prefill_context_parallel_size=1 on the decode (kv_consumer) engine." + ) self.num_key_value_heads = self.vllm_config.model_config.hf_text_config.num_key_value_heads # kv cache config @@ -1615,8 +1630,9 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str, kv_cache_config: KVC + vllm_config.parallel_config.data_parallel_rank * vllm_config.parallel_config.tensor_parallel_size * vllm_config.parallel_config.pipeline_parallel_size + * vllm_config.parallel_config.prefill_context_parallel_size ) - device_index = self.pp_rank * self.tp_size + self.tp_rank + device_index = (self.pp_rank * self.pcp_size + self.pcp_rank) * self.tp_size + self.tp_rank self.handshake_port = self.side_channel_port + device_index self.sockets: dict = {} inject_qos(vllm_config.kv_transfer_config.get_from_extra_config("qos_priority", PD_QOS_DEFAULT)) @@ -1835,6 +1851,7 @@ def _align_up(addr: int) -> int: metadata, ready_event, self.kv_caches, + self.pcp_rank, ) self.kv_send_thread.start() else: @@ -1908,11 +1925,13 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata): prefill_tp_size = meta.remote_ptp_size if getattr(meta, "remote_ptp_size", None) else self._prefill_tp_size tp_num_need_pulls = self._get_tp_num_need_pulls(prefill_tp_size) remote_req_id = meta.remote_request_id + # PCP selects a complete replica; TP offsets and group block IDs stay unchanged. + pcp_offset = self._get_selected_pcp_rank(remote_req_id, meta.remote_pcp_size) * prefill_tp_size if self.use_mamba: assert self.kv_recv_thread is not None chosen_rank_list = self._get_remote_rank(remote_req_id, prefill_tp_size) - remote_handshake_port_list = [[x + meta.remote_port] for x in chosen_rank_list] + remote_handshake_port_list = [[x + meta.remote_port + pcp_offset] for x in chosen_rank_list] # Iterate all remote peers like the non-mamba branch; the old # code only pulled from the first peer, so with P-side PP>1 # the later stages never transferred their KV. @@ -1947,7 +1966,7 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata): ) else: # TODO: support prefill context parallel and pipeline parallel open at the same time chosen_rank_list = self._get_remote_rank(remote_req_id, prefill_tp_size) - remote_handshake_port_list = [[x + meta.remote_port] for x in chosen_rank_list] + remote_handshake_port_list = [[x + meta.remote_port + pcp_offset] for x in chosen_rank_list] for i in range(tp_num_need_pulls * self._prefill_pp_size): assert self.kv_recv_thread is not None remote_host, remote_engine_id = self._get_remote_host_info_by_port( @@ -1978,11 +1997,20 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata): if self.kv_send_thread is not None: for req_id, delay_start_time in metadata.requests_to_send.items(): - if self.tp_rank in self._prefill_get_remote_rank(req_id): + # Only selected sources wait for DONE; unused PCP replicas finish locally. + selected_pcp_rank = self._get_selected_pcp_rank(req_id, self.pcp_size) + if self.pcp_rank == selected_pcp_rank and self.tp_rank in self._prefill_get_remote_rank(req_id): self.kv_send_thread.add_delayed_request(req_id, delay_start_time) else: self.kv_send_thread.add_not_transfer_request(req_id) + @staticmethod + def _get_selected_pcp_rank(req_id: str, pcp_size: int) -> int: + if pcp_size == 1: + return 0 + # Use the P request ID, independently of the TP routing seed. + return random.Random(string_to_int64_hash(f"pcp:{req_id}")).randrange(pcp_size) + def _get_tp_num_need_pulls(self, prefill_tp_size: int) -> int: if self.use_mamba: assert prefill_tp_size == self.tp_size, "Mooncake connector does not support different TP size with Mamba." diff --git a/vllm_ascend/platform.py b/vllm_ascend/platform.py index d889127b078b..85f2efe57790 100644 --- a/vllm_ascend/platform.py +++ b/vllm_ascend/platform.py @@ -1301,22 +1301,22 @@ def _validate_sfa_dcp_kv_sp(vllm_config: VllmConfig) -> None: cache_config = vllm_config.cache_config model_config = vllm_config.model_config - cp_size = parallel_config.prefill_context_parallel_size * parallel_config.decode_context_parallel_size + dcp_enabled = parallel_config.decode_context_parallel_size > 1 use_sparse = model_uses_sfa_sparse(model_config) if ( vllm_config.kv_transfer_config is not None and cache_config.block_size != parallel_config.cp_kv_cache_interleave_size - and cp_size > 1 + and dcp_enabled ): raise AssertionError( f"cp_kv_cache_interleave_size({parallel_config.cp_kv_cache_interleave_size}) " f"and block_size({cache_config.block_size}) " - "needs to be equal if PCP or DCP is enabled in P/D disaggregate and kv pool scenario." + "needs to be equal if DCP is enabled in P/D disaggregate and kv pool scenario." ) - if use_sparse and cp_size > 1 and parallel_config.cp_kv_cache_interleave_size != cache_config.block_size: + if use_sparse and dcp_enabled and parallel_config.cp_kv_cache_interleave_size != cache_config.block_size: logger.warning_once( - "The current SFA context-parallel implementation requires " + "The current SFA decode-context-parallel implementation requires " f"cp_kv_cache_interleave_size({parallel_config.cp_kv_cache_interleave_size})" f" == block_size({cache_config.block_size}). " f"Override cp_kv_cache_interleave_size to {cache_config.block_size}."