diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index a4dfc168f3..3be1730ec3 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -58,6 +58,7 @@ def region_pull_worker(): worker._recving_transfers = {} worker.use_mla, worker._has_mamba = True, False worker.dcp_size = 1 + worker.dcp_rank = 0 spec = MLAAttentionSpec( block_size=64, num_kv_heads=1, head_size=8, dtype=torch.bfloat16 ) @@ -165,6 +166,82 @@ def test_region_pull_ignores_allocation_padding( ] +@pytest.mark.cpu_test +@pytest.mark.parametrize("num_pages", [0, 3, 19]) +@pytest.mark.parametrize("region_groups", [(0, 1), (1, 0, 1)]) +def test_dcp_region_pull(region_pull_worker, num_pages, region_groups): + """Read each uncached page from its DCP rank; notify ranks with no pages.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + RemoteMeta, + ReqMeta, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping + + worker = region_pull_worker + worker.region_group_ids = list(region_groups) + worker.num_regions = len(region_groups) + worker.dst_region_group_ids = {"P": [0] * worker.num_regions} + worker.dst_region_num_blocks = { + engine: [100] * worker.num_regions for engine in ("P", "D") + } + worker.block_len_per_layer = [1024] * worker.num_regions + worker.dcp_rank = 0 + remote = worker.transfer_topo.get_engine_info.return_value + remote.remote_tp_size = remote.remote_dcp_size = 8 + remote.remote_physical_blocks_per_logical = 1 + ranks = tuple(range(8)) + worker.tp_mappings["P"] = TPMapping((ranks, ranks), ranks, {r: r for r in ranks}, 8) + worker.dst_xfer_side_handles = {"P": {r: 1000 + r for r in ranks}} + worker._remote_agents = {"P": {(0, r): f"P-rank{r}" for r in ranks}} + # Two regions with different prefix hits, plus one padding page each. + local = ( + [list(range(31, 31 + num_pages)), list(range(42, 41 + num_pages))] + if num_pages + else [] + ) + meta = ReqMeta( + local, + local, + tp_size=8, + local_num_computed_blocks=(0, 1, 2), + remote=RemoteMeta( + [[10, 11, 12]], + "localhost", + 1, + "P", + "request-P", + num_tokens=max(0, num_pages * 64 - 7), + ), + ) + worker._read_blocks_for_req("request", meta) + reads = { + call.kwargs["remote_xfer_side_handle"] - 1000: call.kwargs + for call in worker._read_blocks_mixed.call_args_list + } + notified = {call.args[0] for call in worker.nixl_wrapper.send_notif.call_args_list} + for rank in ranks: + expected = [ + (region * 100 + 30 + group * 10 + page, region * 100 + 10 + page // 8) + for region, group in enumerate(region_groups) + for page in range(group + 1, num_pages) + if page % 8 == rank + ] + if expected: + read = reads[rank] + assert ( + list(zip(read["local_block_descs_ids"], read["remote_block_descs_ids"])) + == expected + ) + assert f"P-rank{rank}" not in notified + else: + assert rank not in reads and f"P-rank{rank}" in notified + assert meta.region_blocks_to_zero == ( + [[30 + group * 10 + num_pages] for group in region_groups] + if num_pages + else None + ) + + @pytest.mark.cpu_test def test_region_pull_completion_zeros_only_own_padding(region_pull_worker): """Completion must preserve neighboring layers and already imported pages.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index c2b20ffc59..0f13a265df 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -75,6 +75,7 @@ from vllm.distributed.parallel_state import ( from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import make_zmq_path from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.kv_cache_interface import ( @@ -2430,6 +2431,18 @@ class NixlBaseConnectorWorker: remote_physical_per_logical = ( nixl_agent_meta.physical_blocks_per_logical_kv_block ) + if ( + (self.dcp_size > 1 or remote_dcp_size > 1) + and self.region_group_ids + and self.region_group_ids != self.dst_region_group_ids[remote_engine_id] + and ( + self._physical_blocks_per_logical_kv_block != 1 + or remote_physical_per_logical != 1 + ) + ): + raise NotImplementedError( + "DCP region pulls require matching logical and physical block sizes" + ) if ( self._has_mamba and remote_physical_per_logical @@ -3234,11 +3247,14 @@ class NixlBaseConnectorWorker: local_dcp_rank: int, remote_dcp_size: int, local_num_computed_blocks: int, + num_remote_blocks: int | None = None, ) -> tuple[list[int], list[int]]: """Match local and remote block IDs by global DCP position. ``local_ids`` excludes blocks already satisfied by the local prefix cache, while ``remote_ids`` contains the full transferable remote list. + ``num_remote_blocks``, when provided, is the global block count before + DCP partitioning and excludes allocation padding. Example: local DCP size 2, rank 0 owns *global positions*=[0, 2, 4, 6] @@ -3254,6 +3270,9 @@ class NixlBaseConnectorWorker: """ local_size, remote_size = local_dcp_size, remote_dcp_size + if num_remote_blocks is not None: + num_local = cdiv(num_remote_blocks - local_dcp_rank, local_size) + local_ids = local_ids[: max(0, num_local - local_num_computed_blocks)] if local_size == remote_size: local_slice = local_ids @@ -3274,6 +3293,8 @@ class NixlBaseConnectorWorker: remote_slice = remote_ids[start_remote::k] matched_blocks = min(len(local_slice), len(remote_slice)) + if num_remote_blocks is not None and matched_blocks < len(local_slice): + raise ValueError("Remote KV pages do not cover the requested range") return local_slice[:matched_blocks], remote_slice[:matched_blocks] @staticmethod @@ -3290,13 +3311,15 @@ class NixlBaseConnectorWorker: block_ids_by_region.append(list(block_ids[group_id])) return block_ids_by_region - @staticmethod def _apply_prefix_caching_by_region( + self, decode_block_ids: BlockIds, prefill_block_ids: BlockIds, *, num_computed_blocks: list[int] | None = None, num_remote_blocks: int | None = None, + remote_rank: int = 0, + remote_dcp_size: int = 1, ) -> tuple[BlockIds, BlockIds]: """Pair an uncached decode suffix with the same prefill regions.""" assert len(decode_block_ids) == len(prefill_block_ids) @@ -3310,13 +3333,18 @@ class NixlBaseConnectorWorker: for decode_region, prefill_region, start in zip( decode_block_ids, prefill_block_ids, num_computed_blocks, strict=True ): - count = min(len(decode_region), max(num_remote_blocks - start, 0)) - if start + count > len(prefill_region): - raise ValueError("Remote KV pages do not cover the requested range") - # An extra producer page can be allocation padding, not a - # cached prefix. Select by token position instead of list length. - matched_decode.append(list(decode_region[:count])) - matched_prefill.append(list(prefill_region[start : start + count])) + local, remote = self._apply_dcp_prefix_caching( + decode_region, + prefill_region, + remote_rank, + self.dcp_size, + self.dcp_rank, + remote_dcp_size, + start, + num_remote_blocks=num_remote_blocks, + ) + matched_decode.append(local) + matched_prefill.append(remote) return matched_decode, matched_prefill trimmed_prefill: list[list[int]] = [] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 43fbbe0c3c..6539b7e57e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -191,7 +191,6 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): "Different NIXL cache-group layouts are only supported for " "pure MLA models" ) - assert len(plan.all_source_ranks) == 1 if self.block_size != remote_info.remote_block_size: raise NotImplementedError( "Region-mapped NIXL transfers require matching physical block sizes" @@ -213,7 +212,6 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): and meta.local_num_computed_blocks and all(group >= 0 for group in local_region_groups) and all(group >= 0 for group in remote_region_groups) - and not dcp_active ): transfer_groups = self.kv_cache_config.transfer_group_ids num_computed_blocks = [ @@ -224,32 +222,34 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): num_remote_blocks = cdiv( meta.remote.num_tokens, remote_info.remote_block_size ) - elif ( - remote_info.remote_physical_blocks_per_logical + elif any(local_by_region) and ( + dcp_active + or remote_info.remote_physical_blocks_per_logical != self._physical_blocks_per_logical_kv_block ): raise NotImplementedError( - "Region-mapped pulls with different logical block sizes require " - "remote_num_tokens, per-group prefix counts, unshared regions " - "and DCP=1" + "Region-mapped pulls with DCP or different logical block sizes " + "require remote_num_tokens, per-group prefix counts " + "and unshared regions" ) - matched_local, matched_remote = self._apply_prefix_caching_by_region( - local_by_region, - remote_by_region, - num_computed_blocks=num_computed_blocks, - num_remote_blocks=num_remote_blocks, - ) - meta.region_blocks_to_zero = [ - list(blocks[len(matched) :]) - for blocks, matched in zip(local_by_region, matched_local, strict=True) - ] read_specs = [ ReadSpec( - remote_rank=plan.all_source_ranks[0], - local_block_ids=matched_local, - remote_block_ids=matched_remote, + rank, + *self._apply_prefix_caching_by_region( + local_by_region, + remote_by_region, + num_computed_blocks=num_computed_blocks, + num_remote_blocks=num_remote_blocks, + remote_rank=rank, + remote_dcp_size=remote_info.remote_dcp_size, + ), block_ids_by_region=True, ) + for rank in plan.all_source_ranks + ] + meta.region_blocks_to_zero = [ + list(blocks[sum(len(spec.local_block_ids[r]) for spec in read_specs) :]) + for r, blocks in enumerate(local_by_region) ] else: remote_logical_block_ids = meta.remote.block_ids