diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_events.py b/tests/v1/kv_connector/unit/offloading_connector/test_events.py index 29fc29b299..9fdb1a9fed 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_events.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py @@ -142,6 +142,7 @@ def _stored_event( medium: Medium = _CPU_MEDIUM, locality: Locality | None = None, ownership: str | None = None, + removal_expected: bool = False, ) -> OffloadingEvent: return OffloadingEvent( keys=keys, @@ -149,6 +150,7 @@ def _stored_event( removed=False, locality=locality, ownership=ownership, + removal_expected=removal_expected, ) @@ -514,6 +516,9 @@ def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit(): ) assert tracker._pending_event_metadata[key] is confirmed_meta + list( + tracker.take_events([_stored_event([key], Medium.STORAGE, ownership="custom")]) + ) removed = list(tracker.take_events([_removed_event([key])])) assert len(removed) == 1 assert removed[0].block_hashes == [ @@ -534,17 +539,36 @@ def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit(): ] -def test_secondary_stored_event_does_not_mutate_cpu_metadata(): - tracker, _, _, key = _lookup_chunk() - expected_metadata = dict(tracker._pending_event_metadata) - - stored = list( - tracker.take_events([_stored_event([key], Medium.STORAGE, ownership="custom")]) +@pytest.mark.parametrize( + ("record_method", "position"), + [("record_store", 0), ("record_partial_store", 4)], +) +def test_reoffload_preserves_secondary_residency(record_method, position): + tracker, req, group_config, key = _lookup_chunk() + [stored] = tracker.take_events( + [ + _stored_event( + [key], + Medium.STORAGE, + ownership="custom", + removal_expected=True, + ) + ] ) - assert stored[0].token_ids == [1, 2, 3, 4] - assert stored[0].ownership == "custom" - assert tracker._pending_event_metadata == expected_metadata + getattr(tracker, record_method)(req, group_config, position, key) + + assert tracker._pending_event_metadata[key].active_residencies == { + (Medium.CPU, None), + (Medium.STORAGE, "custom"), + } + [removed] = tracker.take_events( + [_removed_event([key], Medium.STORAGE, ownership="custom")] + ) + + assert stored.token_ids == [1, 2, 3, 4] + assert stored.ownership == removed.ownership == "custom" + assert key in tracker._pending_event_metadata def test_take_events_groups_removed_hashes_by_kv_group(): @@ -556,12 +580,9 @@ def test_take_events_groups_removed_hashes_by_kv_group(): key0 = _record_chunks(tracker, req0, group0_config, num_chunks=1)[0] key1 = _record_chunks(tracker, req1, group1_config, num_chunks=1)[0] - removed = list( - tracker.take_events([_removed_event([key0, key1], ownership="custom")]) - ) + removed = list(tracker.take_events([_removed_event([key0, key1])])) assert len(removed) == 2 - assert {event.ownership for event in removed} == {"custom"} by_group = {event.group_idx: event.block_hashes for event in removed} assert by_group == { 0: [_wire_hash(_hash(0)), _wire_hash(_hash(1))], diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py index 9c1b6542e6..98720b5b70 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py @@ -78,8 +78,8 @@ def get_offloading_event_group_spec( @dataclass(slots=True) class _OffloadEventMetadata: """BlockStored payload snapshot for one OffloadKey, captured while the - Request is available and kept until the matching eviction event. ``medium`` - and ``ownership`` are forwarded from the OffloadingEvent.""" + Request is available and kept until the final matching removal event. + ``medium`` and ``ownership`` are forwarded from the OffloadingEvent.""" # The chunk's constituent block hashes; the last one is the OffloadKey. block_hashes: tuple[BlockHash, ...] @@ -92,6 +92,7 @@ class _OffloadEventMetadata: extra_keys: tuple[tuple[Any, ...] | None, ...] | None group_idx: int kv_cache_spec: OffloadingEventGroupSpec + active_residencies: set[tuple[Medium, str | None]] class OffloadingEventsTracker: @@ -100,8 +101,8 @@ class OffloadingEventsTracker: The scheduler calls :meth:`record_store` from ``_build_store_jobs`` and :meth:`record_lookup` for ready primary-tier hits while the ``Request`` is available. Deferred and missing lookups add no state. Under the connector's - supported success-only transfer model, entries follow primary allocations - until CPU removal translation or :meth:`reset`. + supported success-only transfer model, entries remain until the final + observed residency removal or :meth:`reset`. """ def __init__(self, config: OffloadingKVEventsConfig): @@ -110,7 +111,7 @@ class OffloadingEventsTracker: config.enable_kv_cache_events and config.self_describing_kv_events ) - # OffloadKey -> payload snapshot, kept until CPU removal or reset. + # OffloadKey -> payload snapshot, kept until final removal or reset. self._pending_event_metadata: dict[OffloadKey, _OffloadEventMetadata] = {} def record_store( @@ -130,6 +131,8 @@ class OffloadingEventsTracker: if group_config.sliding_window_size_in_chunks is not None: return meta = self._build_event_metadata(req, group_config, chunk_idx) + if existing := self._pending_event_metadata.get(offload_key): + meta.active_residencies.update(existing.active_residencies) self._pending_event_metadata[offload_key] = meta def record_lookup( @@ -216,7 +219,7 @@ class OffloadingEventsTracker: lora_id = req.lora_request.adapter_id if req.lora_request is not None else None lora_name = req.lora_request.name if req.lora_request is not None else None - self._pending_event_metadata[offload_key] = _OffloadEventMetadata( + meta = _OffloadEventMetadata( block_hashes=block_hashes, parent_block_hash=parent_block_hash, token_ids=tuple(req.all_token_ids[chunk_start:boundary_tokens]), @@ -226,7 +229,11 @@ class OffloadingEventsTracker: extra_keys=None, group_idx=group_config.group_idx, kv_cache_spec=group_config.kv_event_group_spec, + active_residencies={(Medium.CPU, None)}, ) + if existing := self._pending_event_metadata.get(offload_key): + meta.active_residencies.update(existing.active_residencies) + self._pending_event_metadata[offload_key] = meta def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]: """Translate raw OffloadingEvents into self-describing KV events. @@ -312,6 +319,7 @@ class OffloadingEventsTracker: extra_keys=None, group_idx=group_config.group_idx, kv_cache_spec=group_config.kv_event_group_spec, + active_residencies={(Medium.CPU, None)}, ) def _placeholder_stored( @@ -358,6 +366,8 @@ class OffloadingEventsTracker: ) continue + if event.removal_expected: + meta.active_residencies.add((event.medium, event.ownership)) yield BlockStored( block_hashes=list( maybe_convert_block_hash(h) for h in meta.block_hashes @@ -389,12 +399,15 @@ class OffloadingEventsTracker: locality = event.locality.value if event.locality is not None else None by_group: dict[int, list] = {} for key in event.keys: - meta = self._pending_event_metadata.pop(key, None) + meta = self._pending_event_metadata.get(key) if meta is not None: group_idx = meta.group_idx by_group.setdefault(group_idx, []).extend( maybe_convert_block_hash(h) for h in meta.block_hashes ) + meta.active_residencies.discard((event.medium, event.ownership)) + if not meta.active_residencies: + self._pending_event_metadata.pop(key) else: if self.self_describing_enabled: logger.warning_once( diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index b3b4df8fd7..dede311446 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -159,6 +159,7 @@ class OffloadingEvent: locality: Locality | None = None # Secondary tier identifier that generated the event, or None for primary. ownership: str | None = None + removal_expected: bool = False """