From 804e5377aa0b0cea9544061beb2ebff162543c72 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Wed, 16 Sep 2026 18:39:07 +0200 Subject: [PATCH] [Bugfix][HiSparse][PD] Align region-mapped pulls across logical block sizes (#57077) Signed-off-by: NickLucche --- .../unit/test_nixl_connector_hma.py | 169 ++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 46 ++++- .../kv_connector/v1/nixl/metadata.py | 5 + .../kv_connector/v1/nixl/pull_worker.py | 58 ++++-- 4 files changed, 267 insertions(+), 11 deletions(-) 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 ae2ddece7f..1e34340d0f 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -3,6 +3,8 @@ """Unit tests for NixlConnectorScheduler with HMA and Mamba N-1 prefill.""" import gc +import queue +from types import SimpleNamespace from unittest.mock import MagicMock, patch import msgspec @@ -39,6 +41,173 @@ from .utils import ( ) +@pytest.fixture +def region_pull_worker(): + """Real page mapping and descriptors, with only the transport mocked.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping + + worker = object.__new__(NixlConnectorWorker) + worker.engine_id = "D" + worker.tp_rank = 0 + worker.block_size = 64 + worker._physical_blocks_per_logical_kv_block = 1 + worker.enable_permute_local_kv = False + worker.enable_heterogeneous_attn_post_process = False + worker._engine_last_active = {} + worker._bidirectional_kv_xfer_enabled = False + worker._recving_transfers = {} + worker.use_mla, worker._has_mamba = True, False + worker.dcp_size = 1 + spec = MLAAttentionSpec( + block_size=64, num_kv_heads=1, head_size=8, dtype=torch.bfloat16 + ) + worker.kv_cache_config = KVCacheConfig( + num_blocks=100, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["resident"], spec, enable_kv_transfer=False), + KVCacheGroupSpec(["host"], spec), + KVCacheGroupSpec(["indexer"], spec), + ], + ) + worker._group_spec_types = (MLAAttentionSpec, MLAAttentionSpec) + worker.region_group_ids = [0, 1] + worker.dst_region_group_ids = {"P": [0, 0]} + worker.num_regions = 2 + worker.block_len_per_layer = [1024, 1024] + worker.transfer_topo = MagicMock() + worker.transfer_topo.tp_ratio.return_value = 1 + worker.transfer_topo.block_size_ratio.return_value = 1 + worker.transfer_topo.get_engine_info.return_value = SimpleNamespace( + remote_tp_size=4, + remote_dcp_size=1, + remote_block_size=64, + remote_physical_blocks_per_logical=2, + ) + worker.tp_mappings = {"P": TPMapping(((0,), (0,)), (0,), {0: 0}, 0)} + worker._mixed_mem_types = True + worker.src_xfer_handles_by_block_size = {64: 1} + worker._dram_src_handles_by_block_size = {64: 2} + worker.dst_xfer_side_handles = {"P": {0: 3}} + worker.dst_num_blocks = {"P": 100, "D": 100} + worker.dst_region_num_blocks = {"P": [100, 100], "D": [100, 100]} + worker._remote_agents = {"P": {(0, 0): "P-rank0"}} + worker._read_blocks_mixed = MagicMock() + worker.nixl_wrapper = MagicMock() + return worker + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("local_ratio,remote_ratio", [(1, 2), (2, 1), (2, 2)]) +@pytest.mark.parametrize("num_tokens", [160, 224]) +@pytest.mark.parametrize("cached", [0, 1]) +def test_region_pull_ignores_allocation_padding( + region_pull_worker, local_ratio, remote_ratio, num_tokens, cached +): + """Padded capacity must not shift either the host or indexer token range.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ) + from vllm.utils.math_utils import cdiv + + worker = region_pull_worker + worker._physical_blocks_per_logical_kv_block = local_ratio + worker.transfer_topo.get_engine_info.return_value = SimpleNamespace( + remote_tp_size=4, + remote_dcp_size=1, + remote_block_size=64, + remote_physical_blocks_per_logical=remote_ratio, + ) + local_count = cdiv(num_tokens, 64 * local_ratio) + remote_count = cdiv(num_tokens, 64 * remote_ratio) + # The host has a cached prefix, while the indexer still needs the full range. + local = [ + list(range(30 + cached, 30 + local_count)), + list(range(40, 40 + local_count)), + ] + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="request", + local_block_ids=local, + local_num_computed_blocks=(0, cached, 0), + awaiting_kvs=True, + kv_transfer_params={ + "remote_engine_id": "P", + "remote_request_id": "request-P", + "remote_host": "localhost", + "remote_port": 1, + "remote_num_tokens": num_tokens, + "remote_block_ids": [list(range(10, 10 + remote_count))], + }, + ) + meta = metadata.reqs_to_recv["request"] + meta.local_physical_block_ids = worker._logical_to_kernel_block_ids( + local, local_ratio + ) + worker._read_blocks_for_req("request", meta) + read = worker._read_blocks_mixed.call_args.kwargs + valid_pages = cdiv(num_tokens, 64) + source_start = 10 * remote_ratio + assert read["remote_block_descs_ids"].tolist() == ( + list(range(source_start + cached * local_ratio, source_start + valid_pages)) + + list(range(100 + source_start, 100 + source_start + valid_pages)) + ) + assert read["local_block_descs_ids"].tolist() == ( + list(range((30 + cached) * local_ratio, 30 * local_ratio + valid_pages)) + + list(range(100 + 40 * local_ratio, 100 + 40 * local_ratio + valid_pages)) + ) + assert meta.region_blocks_to_zero == [ + list( + range(base * local_ratio + valid_pages, (base + local_count) * local_ratio) + ) + for base in (30, 40) + ] + + +@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.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + RemoteMeta, + ReqMeta, + ) + + worker = region_pull_worker + # Two strided layer regions in one backing, with nonzero layer offsets. + backing = torch.full((4, 3, 8), 7, dtype=torch.uint8) + caches = {"host": backing[:, 1], "indexer": backing[:, 2]} + worker.device_kv_caches = caches + worker.region_names = list(caches) + worker.kv_caches_base_addr = {"D": {0: [c.data_ptr() for c in caches.values()]}} + worker.region_num_blocks = [4, 4] + worker.block_len_per_layer = [8, 8] + worker.block_stride_per_layer = [24, 24] + worker._recving_metadata = { + "request": ReqMeta( + local_block_ids=[[1, 2], [0, 1]], + local_physical_block_ids=[[1, 2], [0, 1]], + tp_size=1, + remote=RemoteMeta([[3]], "localhost", 1, "P", "request-P"), + region_blocks_to_zero=[[2], [1]], + ) + } + worker._get_new_notifs = MagicMock(return_value=set()) + worker._pop_done_transfers = MagicMock(return_value=({"request"}, set())) + worker._replicated_pcp_done_sending = set() + worker._failed_recv_reqs = queue.Queue() + worker._recv_failures = set() + worker._send_pending_recv_notifs = MagicMock() + worker._sync_device_after_direct_recv = MagicMock() + worker.use_host_buffer = False + worker._reqs_to_send = {} + result = worker.get_transfer_results() + expected = torch.full_like(backing, 7) + expected[2, 1] = 0 + expected[1, 2] = 0 + assert torch.equal(backing, expected) + assert result.finished_recving == {"request"} + + @pytest.mark.cpu_test @pytest.mark.parametrize( "swa_enabled,expected_sw_sizes", 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 2704de589f..db392078ca 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 @@ -2785,6 +2785,25 @@ class NixlBaseConnectorWorker: indices=indices, ) + def _zero_region_blocks(self, block_ids: BlockIds) -> None: + """Clear clipped physical pages in their owning region only.""" + if not any(block_ids): + return + bases = self.kv_caches_base_addr[self.engine_id][self.tp_rank] + for region, blocks in enumerate(block_ids): + if not blocks: + continue + cache = self.device_kv_caches[self.region_names[region]] + storage = cache.untyped_storage() + pages = torch.empty(0, dtype=torch.uint8, device=cache.device).set_( + storage, + bases[region] - storage.data_ptr(), + (self.region_num_blocks[region], self.block_len_per_layer[region]), + (self.block_stride_per_layer[region], 1), + ) + for block in blocks: + pages[block].zero_() + def get_transfer_results(self) -> KVConnectorTransferResults: """ Get transfers that completed on this specific worker. @@ -2857,6 +2876,12 @@ class NixlBaseConnectorWorker: direct_device_recving.add(req_id) + if meta.region_blocks_to_zero is not None: + # P/D group positions differ. Use the actual region read plan, + # including any local allocation padding the read did not cover. + self._zero_region_blocks(meta.region_blocks_to_zero) + continue + # Post processing for heteroblocksize/layout, and for blocks the # transfer clipped. The latter happens either at remote-block # granularity (block_size_ratio > 1) or at kernel-block @@ -3254,13 +3279,32 @@ class NixlBaseConnectorWorker: @staticmethod def _apply_prefix_caching_by_region( - decode_block_ids: BlockIds, prefill_block_ids: BlockIds + decode_block_ids: BlockIds, + prefill_block_ids: BlockIds, + *, + num_computed_blocks: list[int] | None = None, + num_remote_blocks: int | None = None, ) -> tuple[BlockIds, BlockIds]: """Pair an uncached decode suffix with the same prefill regions.""" assert len(decode_block_ids) == len(prefill_block_ids) if not any(decode_block_ids): return [], prefill_block_ids + if num_computed_blocks is not None: + assert num_remote_blocks is not None + matched_decode, matched_prefill = [], [] + 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])) + return matched_decode, matched_prefill + trimmed_prefill: list[list[int]] = [] for decode_region, prefill_region in zip( decode_block_ids, prefill_block_ids, strict=True diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index c2b4d9dce7..1ae0ed72ce 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -230,6 +230,7 @@ class RemoteMeta: engine_id: str request_id: str blocks_expiry_time: float | None = None + num_tokens: int | None = None @dataclass @@ -252,6 +253,9 @@ class ReqMeta: # True only when the scheduler parked the request in WAITING_FOR_REMOTE_KVS # and expects it in finished_recving; notify-only recvs must not be reported. awaiting_kvs: bool = False + # Worker-only, per-region physical pages to zero after a successful pull. + # None selects group-based completion; empty lists mean no zeroing. + region_blocks_to_zero: BlockIds | None = None class NixlConnectorMetadata(KVConnectorMetadata): @@ -326,5 +330,6 @@ class NixlConnectorMetadata(KVConnectorMetadata): host=kv_transfer_params["remote_host"], port=kv_transfer_params["remote_port"], blocks_expiry_time=kv_transfer_params.get("remote_blocks_expiry_time"), + num_tokens=kv_transfer_params.get("remote_num_tokens"), ) self.reqs_to_recv[request_id] = req 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 04caebbd1f..1530b9fb6a 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 @@ -20,6 +20,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( _is_attention_spec, ) from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv if TYPE_CHECKING: from vllm.config import VllmConfig @@ -176,6 +177,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): "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" + ) remote_physical_block_ids = self._logical_to_kernel_block_ids( meta.remote.block_ids, remote_info.remote_physical_blocks_per_logical, @@ -183,13 +188,51 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): remote_by_region = self._block_ids_by_region( remote_physical_block_ids, remote_region_groups ) + local_by_region = self._block_ids_by_region( + local_block_ids, local_region_groups + ) + num_computed_blocks = None + num_remote_blocks = None + if ( + meta.remote.num_tokens is not None + 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 = [ + meta.local_num_computed_blocks[transfer_groups[group]] + * self._physical_blocks_per_logical_kv_block + for group in local_region_groups + ] + num_remote_blocks = cdiv( + meta.remote.num_tokens, remote_info.remote_block_size + ) + elif ( + 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" + ) + 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=self._block_ids_by_region( - local_block_ids, local_region_groups - ), - remote_block_ids=remote_by_region, + local_block_ids=matched_local, + remote_block_ids=matched_remote, block_ids_by_region=True, ) ] @@ -396,12 +439,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): self._recving_transfers.setdefault(request_id, []) return True - if read_spec.block_ids_by_region: - local_block_ids, remote_block_ids = self._apply_prefix_caching_by_region( - decode_block_ids=local_block_ids, - prefill_block_ids=remote_block_ids, - ) - else: + if not read_spec.block_ids_by_region: assert ( len(remote_block_ids) == len(local_block_ids)