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 c3058eaddb..11538f0219 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -314,7 +314,10 @@ def test_apply_prefix_caching_mamba_hybrid( worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True) aligned_local, aligned_remote = worker._apply_prefix_caching( - local_block_ids, remote_block_ids, remote_physical_per_logical + local_block_ids, + remote_block_ids, + local_physical_per_logical, + remote_physical_per_logical, ) assert aligned_local == expected_local, ( @@ -413,7 +416,10 @@ def test_apply_prefix_caching_ssm_prefix_cache_hit( worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True) aligned_local, aligned_remote = worker._apply_prefix_caching( - local_block_ids, remote_block_ids, remote_physical_per_logical + local_block_ids, + remote_block_ids, + local_physical_per_logical, + remote_physical_per_logical, ) assert aligned_local == expected_local, ( @@ -442,7 +448,7 @@ def test_apply_prefix_caching_ssm_unpairable_slots_rejected(): with pytest.raises(AssertionError, match="unpairable SSM state slots"): worker._apply_prefix_caching( - [list(range(10)), [4, 5, 6, 7]], [list(range(10)), [8, 9]], 10 + [list(range(10)), [4, 5, 6, 7]], [list(range(10)), [8, 9]], 10, 10 ) @@ -527,6 +533,7 @@ def test_mismatched_physical_per_logical_fails_with_prefix_caching( aligned_local, aligned_remote = worker._apply_prefix_caching( local_block_ids, remote_block_ids, + local_physical_per_logical, remote_physical_per_logical, ) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index a5e1eabad8..f69bd544f5 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -341,6 +341,7 @@ class _StubWriterWorker(NixlPushConnectorWorker): # Single non-hybrid attention group, matching the stub block id lists. w._has_mamba = False w._group_spec_types = (FullAttentionSpec,) + w._engine_ttl = 0.0 w._engine_last_active = {} # Track _do_start_push_kv invocations. @@ -569,7 +570,7 @@ def test_do_start_push_kv_drops_request_on_handshake_failure(): the failure is logged. Blocks are reclaimed by the lease/watchdog, matching the old blocking behaviour.""" w = _StubWriterWorker.fresh() - w._logical_to_kernel_block_ids = lambda x: x + w._logical_to_kernel_block_ids = lambda x, ratio: x xfer_calls: list[dict[str, Any]] = [] w._xfer_blocks_for_req = lambda **kw: xfer_calls.append(kw) failures: list[dict[str, Any]] = [] @@ -1157,3 +1158,93 @@ class TestPushWriterMlaReplication: # All of the request's WRITE handles must be tracked together, so the # engine thread never sees a partial set and double-frees the request. assert sorted(w._sending_transfers["p-req"]) == [1000, 1001] + + +class TestPushPrefixCaching: + """Partial prefix-cache hit on D: D preallocates only its *uncomputed* + blocks, so on a partial hit it registers fewer blocks than P's full + sequence. P must WRITE only the *tail* of its sequence into D's slots -- + mirroring pull-mode ``_apply_prefix_caching`` (end-trim), never a front-trim + which would write P's cached prefix into D's uncomputed suffix slots. + + The trim runs inside ``_xfer_blocks``, so these tests drive the real + ``_xfer_blocks_for_req`` path and stub only the NIXL WRITE. With one region + and 1:1 ratios ``_compute_desc_ids`` is the identity map, so the descs + captured from ``make_prepped_xfer`` are exactly the (trimmed) block IDs. + """ + + @staticmethod + def _worker_driving_xfer(engine_id: str = "decode-engine"): + from types import SimpleNamespace + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( + TPMapping, + ) + + w = _StubWriterWorker.fresh() + w._has_mamba = False + w.use_mla = False + w.block_size = 16 + w.num_regions = 1 + w._physical_blocks_per_logical_kv_block = 1 + + w.transfer_topo = MagicMock() + w.transfer_topo.get_engine_info.return_value = SimpleNamespace( + remote_physical_blocks_per_logical=1, + remote_block_size=16, + remote_tp_size=1, + ) + w.transfer_topo.tp_ratio.return_value = 1 + w.transfer_topo.block_size_ratio.return_value = 1 + w.tp_mappings = { + engine_id: TPMapping( + source_ranks_per_group=((0,),), + all_source_ranks=(0,), + rank_to_attention_slot={0: 0}, + rank_offset_factor=0, + ) + } + w.dst_num_blocks = {engine_id: 10_000, w.engine_id: 10_000} + w.dst_xfer_side_handles = {engine_id: {0: 5000}} + w.src_xfer_handles_by_block_size = {16: 2000} + + # Stub only the NIXL WRITE; kernel expansion, prefix-cache trim, desc + # computation and count assertions all run for real. + w.nixl_wrapper = MagicMock() + w.nixl_wrapper.make_prepped_xfer.return_value = 7 + w._ensure_handshake = lambda *a, **k: None + w._logical_to_kernel_block_ids = lambda x, ratio: x + return w, engine_id + + @staticmethod + def _written_block_ids(w) -> tuple[list[int], list[int]]: + """Return the (local, remote) block IDs handed to the NIXL WRITE.""" + args, _ = w.nixl_wrapper.make_prepped_xfer.call_args + # ("WRITE", local_handle, local_descs, remote_handle, remote_descs) + return list(args[2]), list(args[4]) + + def test_partial_prefix_hit_end_trims_producer_blocks(self): + """D registered only its 2 uncomputed suffix blocks; P finished the + full 5-block sequence. P must WRITE its LAST 2 blocks into D's slots.""" + w, _ = self._worker_driving_xfer() + reg = _registration_data("req-pc", local_block_ids=([500, 501],)) + + NixlPushConnectorWorker._do_start_push_kv( + w, "req-pc", ([10, 11, 12, 13, 14],), reg + ) + + local, remote = self._written_block_ids(w) + # End-trim (suffix), NOT front-trim: [13, 14], not [10, 11]. + assert local == [13, 14] + assert remote == [500, 501] + + def test_no_prefix_hit_leaves_blocks_untrimmed(self): + """Equal counts (no prefix cache hit on D): nothing is trimmed.""" + w, _ = self._worker_driving_xfer() + reg = _registration_data("req-full", local_block_ids=([500, 501, 502],)) + + NixlPushConnectorWorker._do_start_push_kv(w, "req-full", ([10, 11, 12],), reg) + + local, remote = self._written_block_ids(w) + assert local == [10, 11, 12] + assert remote == [500, 501, 502] 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 e61b2a3921..815eaba7de 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 @@ -2397,93 +2397,101 @@ class NixlBaseConnectorWorker: def _apply_prefix_caching( self, - local_block_ids: BlockIds, - remote_block_ids: BlockIds, - remote_physical_per_logical: int, - ) -> tuple[BlockIds, list]: - """Apply prefix caching by trimming local/remote block ID lists. + decode_block_ids: BlockIds, + prefill_block_ids: BlockIds, + decode_physical_per_logical: int, + prefill_physical_per_logical: int, + ) -> tuple[BlockIds, BlockIds]: + """Trim block ID lists so only the uncomputed suffix is transferred. - For non-Mamba models: end-trim remote to match local count, so that + Inputs are *kernel* (physical) block IDs, already expanded from logical IDs + with each side's physical-per-logical ratio. Both pull and push call this after + that expansion so the trim happens at kernel granularity. + + The prefix-cache hit is always on the decode (D) side, so ``decode`` + holds only its uncomputed blocks while prefill (P) holds the full + sequence. This is mode-independent: pull passes its own (D) blocks as + ``decode``, push passes the remote D registration as ``decode``. + + For non-Mamba models: end-trim ``prefill`` to match ``decode`` count, so already-cached prefix blocks are skipped in the transfer. - For Mamba hybrid (prefix caching not yet supported): front-trim both - to the minimum count to handle kernel block count discrepancies from - logical block rounding in heterogeneous TP. + For Mamba hybrid: SSM groups pair state slots by position and FA + groups end-trim to the uncomputed suffix when physical-per-logical + matches. """ - # Partial prefix cache hit: just read uncomputed blocks. + # Partial prefix cache hit: just transfer uncomputed blocks. # Skip mamba groups — their blocks represent full state (conv+ssm), # not per-token data, so trimming would corrupt the transfer. - remote_block_ids = list(remote_block_ids) + prefill_block_ids = list(prefill_block_ids) if not self._has_mamba: - for i, remote_group in enumerate(remote_block_ids): - num_local_blocks = len(local_block_ids[i]) - assert num_local_blocks <= len(remote_group) - if num_local_blocks < len(remote_group): - remote_block_ids[i] = remote_group[-num_local_blocks:] + for i, prefill_group in enumerate(prefill_block_ids): + num_decode_blocks = len(decode_block_ids[i]) + assert num_decode_blocks <= len(prefill_group) + if num_decode_blocks < len(prefill_group): + prefill_block_ids[i] = prefill_group[-num_decode_blocks:] else: - # (NOTE: ZhanqiuHu) Mamba hybrid: no prefix caching support so far.HeteroTP - # can cause different kernel block counts due to logical block rounding. + # (NOTE: ZhanqiuHu) HeteroTP can cause different kernel block counts + # due to logical block rounding. # Example: 640 prompt tokens, kernel_block_size=64 - # remote physical_per_logical=10, local physical_per_logical=6 - # remote logical ids from kv_transfer_params = [0] - # local logical ids allocated = [0, 1] - # remote kernel blocks: [0..9] (1*10=10) - # local kernel blocks: [0..11] (2*6=12) + # prefill physical_per_logical=10, decode physical_per_logical=6 + # prefill logical ids from kv_transfer_params = [0] + # decode logical ids allocated = [0, 1] + # prefill kernel blocks: [0..9] (1*10=10) + # decode kernel blocks: [0..11] (2*6=12) # actual data blocks = ceil(640/64) = 10, trim both to 10 - # Vice versa (remote physical_per_logical=6, local=10): - # remote logical ids = [0, 1], local logical ids = [0] - # remote kernel blocks: [0..11] (2*6=12) - # local kernel blocks: [0..9] (1*10=10) + # Vice versa (prefill physical_per_logical=6, decode=10): + # prefill logical ids = [0, 1], decode logical ids = [0] + # prefill kernel blocks: [0..11] (2*6=12) + # decode kernel blocks: [0..9] (1*10=10) # actual data blocks = ceil(640/64) = 10, trim both to 10 - local_block_ids = list(local_block_ids) - for i, remote_group in enumerate(remote_block_ids): - num_local_blocks = len(local_block_ids[i]) - num_remote_blocks = len(remote_group) + decode_block_ids = list(decode_block_ids) + for i, prefill_group in enumerate(prefill_block_ids): + num_decode_blocks = len(decode_block_ids[i]) + num_prefill_blocks = len(prefill_group) if _is_ssm_spec(self._group_spec_types[i]): - if num_local_blocks == num_remote_blocks: + if num_decode_blocks == num_prefill_blocks: continue # Only state-bearing slots reach here, single-state modes # just one (see get_exchange_clipped_blocks), so differing # counts mean position-indexed "all"-mode lists. A longer - # remote list carries earlier positions the local side - # already has (prefix hit) -> read its tail; a longer local + # prefill list carries earlier positions the decode side + # already has (prefix hit) -> take its tail; a longer decode # list holds the position D recomputes itself, which gets - # no remote state. - assert num_local_blocks - num_remote_blocks <= 1, ( + # no prefill state. + assert num_decode_blocks - num_prefill_blocks <= 1, ( f"Group {i}: unpairable SSM state slots, " - f"local={num_local_blocks} remote={num_remote_blocks}" + f"decode={num_decode_blocks} prefill={num_prefill_blocks}" ) - num_blocks = min(num_local_blocks, num_remote_blocks) - if num_local_blocks < num_remote_blocks: - remote_block_ids[i] = remote_group[-num_blocks:] + num_blocks = min(num_decode_blocks, num_prefill_blocks) + if num_decode_blocks < num_prefill_blocks: + prefill_block_ids[i] = prefill_group[-num_blocks:] else: - local_block_ids[i] = local_block_ids[i][:num_blocks] + decode_block_ids[i] = decode_block_ids[i][:num_blocks] elif ( - self._physical_blocks_per_logical_kv_block - == remote_physical_per_logical - and num_local_blocks < num_remote_blocks + decode_physical_per_logical == prefill_physical_per_logical + and num_decode_blocks < num_prefill_blocks ): # Partial prefix cache hit for FA group. - remote_block_ids[i] = remote_group[-num_local_blocks:] + prefill_block_ids[i] = prefill_group[-num_decode_blocks:] else: # TODO Handle prefix caching with different block_sizes # Allocation rounding legitimately leaves up to # ppl - 1 trailing dead kernel blocks per side (plus one - # extra local block for the recomputed final token), so + # extra decode block for the recomputed final token), so # the counts may differ by up to the sum of the two # ratios; anything larger indicates mismatched lists. max_padding = ( - self._physical_blocks_per_logical_kv_block - + remote_physical_per_logical + decode_physical_per_logical + prefill_physical_per_logical ) - assert abs(num_local_blocks - num_remote_blocks) <= max_padding, ( - f"Group {i}: |{num_local_blocks} - " - f"{num_remote_blocks}| > {max_padding}" + assert abs(num_decode_blocks - num_prefill_blocks) <= max_padding, ( + f"Group {i}: |{num_decode_blocks} - " + f"{num_prefill_blocks}| > {max_padding}" ) - num_blocks = min(num_local_blocks, num_remote_blocks) - local_block_ids[i] = local_block_ids[i][:num_blocks] - remote_block_ids[i] = remote_group[:num_blocks] - return local_block_ids, remote_block_ids + num_blocks = min(num_decode_blocks, num_prefill_blocks) + decode_block_ids[i] = decode_block_ids[i][:num_blocks] + prefill_block_ids[i] = prefill_group[:num_blocks] + return decode_block_ids, prefill_block_ids def get_kv_connector_stats(self) -> KVConnectorStats | None: """ 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 cc32845a4a..ecc11c480b 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 @@ -289,9 +289,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): == len(local_block_ids) == len(self.kv_cache_config.kv_cache_groups) ) - remote_physical_per_logical = remote_info.remote_physical_blocks_per_logical local_block_ids, remote_block_ids = self._apply_prefix_caching( - local_block_ids, remote_block_ids, remote_physical_per_logical + decode_block_ids=local_block_ids, + prefill_block_ids=remote_block_ids, + decode_physical_per_logical=self._physical_blocks_per_logical_kv_block, + prefill_physical_per_logical=remote_info.remote_physical_blocks_per_logical, ) # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 1050e35b7a..7fdde27dcc 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -454,10 +454,6 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): # how pull-mode transfers touch _engine_last_active in start_load_kv. self._engine_last_active[decode_engine_id] = time.perf_counter() - # Both sides stay logical here; ``_xfer_blocks_for_req`` converts each - # to physical with its own physical-blocks-per-logical ratio -- P uses - # ``self._physical_blocks_per_logical_kv_block``, D's is learned during - # the NIXL handshake. logical_local = self._as_grouped_block_ids(local_block_ids) logical_remote = self._as_grouped_block_ids(remote_block_ids) physical_local = self._logical_to_kernel_block_ids( @@ -629,16 +625,28 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): logger.warning("No blocks to push for request %s", request_id) return None - # Align per-group block counts for push. + # Prefix caching: D allocated only uncached blocks, so on a partial hit it + # sends fewer than P's. End-trim P's blocks to that same suffix so we WRITE only + # the uncomputed tail into D's slots. Runs on kernel ids, post-expansion. + remote_block_ids, local_block_ids = self._apply_prefix_caching( + decode_block_ids=remote_block_ids, + prefill_block_ids=local_block_ids, + decode_physical_per_logical=remote_info.remote_physical_blocks_per_logical, + prefill_physical_per_logical=self._physical_blocks_per_logical_kv_block, + ) + local_block_ids = list(local_block_ids) remote_block_ids = list(remote_block_ids) - for i in range(min(len(local_block_ids), len(remote_block_ids))): - num_local = len(local_block_ids[i]) - num_remote = len(remote_block_ids[i]) - if num_local > num_remote: - local_block_ids[i] = local_block_ids[i][:num_remote] - elif num_local < num_remote: - remote_block_ids[i] = remote_block_ids[i][:num_local] + assert len(local_block_ids) == len(remote_block_ids), ( + f"push group-count mismatch for {request_id}: {len(local_block_ids)} " + f"local vs {len(remote_block_ids)} remote groups" + ) + for i in range(len(local_block_ids)): + assert len(local_block_ids[i]) == len(remote_block_ids[i]), ( + f"push block-count mismatch for {request_id} group {i}: " + f"{len(local_block_ids[i])} local vs " + f"{len(remote_block_ids[i])} remote blocks" + ) # Get descs ids. remote_block_descs_ids = self._compute_desc_ids(