[KV Offload] Preserve KV event metadata until final residency removal (#52068)
Signed-off-by: Moein Khazraee <moein@nvidia.com>
Moein Khazraee committed
Aug 31, 2026 at 03:28 UTC
9acbc5360aa50d148fcfcbd597703514aa9a3e22
3 files changed
+54
-19
tests/v1/kv_connector/unit/offloading_connector/test_events.py
+33
-12
@@ -142,6 +142,7 @@ def _stored_event(
142
medium: Medium = _CPU_MEDIUM,
143
locality: Locality | None = None,
144
ownership: str | None = None,
145
+ removal_expected: bool = False,
146
) -> OffloadingEvent:
147
return OffloadingEvent(
148
keys=keys,
@@ -149,6 +150,7 @@ def _stored_event(
150
removed=False,
151
locality=locality,
152
ownership=ownership,
153
+ removal_expected=removal_expected,
154
)
155
156
@@ -514,6 +516,9 @@ def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit():
516
)
517
assert tracker._pending_event_metadata[key] is confirmed_meta
518
519
+ list(
520
+ tracker.take_events([_stored_event([key], Medium.STORAGE, ownership="custom")])
521
+ )
522
removed = list(tracker.take_events([_removed_event([key])]))
523
assert len(removed) == 1
524
assert removed[0].block_hashes == [
@@ -534,17 +539,36 @@ def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit():
539
]
540
541
537
-def test_secondary_stored_event_does_not_mutate_cpu_metadata():
538
- tracker, _, _, key = _lookup_chunk()
539
- expected_metadata = dict(tracker._pending_event_metadata)
542
+@pytest.mark.parametrize(
543
+ ("record_method", "position"),
544
+ [("record_store", 0), ("record_partial_store", 4)],
545
+)
546
+def test_reoffload_preserves_secondary_residency(record_method, position):
547
+ tracker, req, group_config, key = _lookup_chunk()
548
+ [stored] = tracker.take_events(
549
+ [
550
+ _stored_event(
551
+ [key],
552
+ Medium.STORAGE,
553
+ ownership="custom",
554
+ removal_expected=True,
555
+ )
556
+ ]
557
+ )
558
541
- stored = list(
542
- tracker.take_events([_stored_event([key], Medium.STORAGE, ownership="custom")])
559
+ getattr(tracker, record_method)(req, group_config, position, key)
560
+
561
+ assert tracker._pending_event_metadata[key].active_residencies == {
562
+ (Medium.CPU, None),
563
+ (Medium.STORAGE, "custom"),
564
+ }
565
+ [removed] = tracker.take_events(
566
+ [_removed_event([key], Medium.STORAGE, ownership="custom")]
567
)
568
545
- assert stored[0].token_ids == [1, 2, 3, 4]
546
- assert stored[0].ownership == "custom"
547
- assert tracker._pending_event_metadata == expected_metadata
569
+ assert stored.token_ids == [1, 2, 3, 4]
570
+ assert stored.ownership == removed.ownership == "custom"
571
+ assert key in tracker._pending_event_metadata
572
573
574
def test_take_events_groups_removed_hashes_by_kv_group():
@@ -556,12 +580,9 @@ def test_take_events_groups_removed_hashes_by_kv_group():
580
key0 = _record_chunks(tracker, req0, group0_config, num_chunks=1)[0]
581
key1 = _record_chunks(tracker, req1, group1_config, num_chunks=1)[0]
582
559
- removed = list(
560
- tracker.take_events([_removed_event([key0, key1], ownership="custom")])
561
- )
583
+ removed = list(tracker.take_events([_removed_event([key0, key1])]))
584
585
assert len(removed) == 2
564
- assert {event.ownership for event in removed} == {"custom"}
586
by_group = {event.group_idx: event.block_hashes for event in removed}
587
assert by_group == {
588
0: [_wire_hash(_hash(0)), _wire_hash(_hash(1))],
vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
+20
-7
@@ -78,8 +78,8 @@ def get_offloading_event_group_spec(
78
@dataclass(slots=True)
79
class _OffloadEventMetadata:
80
"""BlockStored payload snapshot for one OffloadKey, captured while the
81
- Request is available and kept until the matching eviction event. ``medium``
82
- and ``ownership`` are forwarded from the OffloadingEvent."""
81
+ Request is available and kept until the final matching removal event.
82
+ ``medium`` and ``ownership`` are forwarded from the OffloadingEvent."""
83
84
# The chunk's constituent block hashes; the last one is the OffloadKey.
85
block_hashes: tuple[BlockHash, ...]
@@ -92,6 +92,7 @@ class _OffloadEventMetadata:
92
extra_keys: tuple[tuple[Any, ...] | None, ...] | None
93
group_idx: int
94
kv_cache_spec: OffloadingEventGroupSpec
95
+ active_residencies: set[tuple[Medium, str | None]]
96
97
98
class OffloadingEventsTracker:
@@ -100,8 +101,8 @@ class OffloadingEventsTracker:
101
The scheduler calls :meth:`record_store` from ``_build_store_jobs`` and
102
:meth:`record_lookup` for ready primary-tier hits while the ``Request`` is
103
available. Deferred and missing lookups add no state. Under the connector's
103
- supported success-only transfer model, entries follow primary allocations
104
- until CPU removal translation or :meth:`reset`.
104
+ supported success-only transfer model, entries remain until the final
105
+ observed residency removal or :meth:`reset`.
106
"""
107
108
def __init__(self, config: OffloadingKVEventsConfig):
@@ -110,7 +111,7 @@ class OffloadingEventsTracker:
111
config.enable_kv_cache_events and config.self_describing_kv_events
112
)
113
113
- # OffloadKey -> payload snapshot, kept until CPU removal or reset.
114
+ # OffloadKey -> payload snapshot, kept until final removal or reset.
115
self._pending_event_metadata: dict[OffloadKey, _OffloadEventMetadata] = {}
116
117
def record_store(
@@ -130,6 +131,8 @@ class OffloadingEventsTracker:
131
if group_config.sliding_window_size_in_chunks is not None:
132
return
133
meta = self._build_event_metadata(req, group_config, chunk_idx)
134
+ if existing := self._pending_event_metadata.get(offload_key):
135
+ meta.active_residencies.update(existing.active_residencies)
136
self._pending_event_metadata[offload_key] = meta
137
138
def record_lookup(
@@ -216,7 +219,7 @@ class OffloadingEventsTracker:
219
220
lora_id = req.lora_request.adapter_id if req.lora_request is not None else None
221
lora_name = req.lora_request.name if req.lora_request is not None else None
219
- self._pending_event_metadata[offload_key] = _OffloadEventMetadata(
222
+ meta = _OffloadEventMetadata(
223
block_hashes=block_hashes,
224
parent_block_hash=parent_block_hash,
225
token_ids=tuple(req.all_token_ids[chunk_start:boundary_tokens]),
@@ -226,7 +229,11 @@ class OffloadingEventsTracker:
229
extra_keys=None,
230
group_idx=group_config.group_idx,
231
kv_cache_spec=group_config.kv_event_group_spec,
232
+ active_residencies={(Medium.CPU, None)},
233
)
234
+ if existing := self._pending_event_metadata.get(offload_key):
235
+ meta.active_residencies.update(existing.active_residencies)
236
+ self._pending_event_metadata[offload_key] = meta
237
238
def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]:
239
"""Translate raw OffloadingEvents into self-describing KV events.
@@ -312,6 +319,7 @@ class OffloadingEventsTracker:
319
extra_keys=None,
320
group_idx=group_config.group_idx,
321
kv_cache_spec=group_config.kv_event_group_spec,
322
+ active_residencies={(Medium.CPU, None)},
323
)
324
325
def _placeholder_stored(
@@ -358,6 +366,8 @@ class OffloadingEventsTracker:
366
)
367
continue
368
369
+ if event.removal_expected:
370
+ meta.active_residencies.add((event.medium, event.ownership))
371
yield BlockStored(
372
block_hashes=list(
373
maybe_convert_block_hash(h) for h in meta.block_hashes
@@ -389,12 +399,15 @@ class OffloadingEventsTracker:
399
locality = event.locality.value if event.locality is not None else None
400
by_group: dict[int, list] = {}
401
for key in event.keys:
392
- meta = self._pending_event_metadata.pop(key, None)
402
+ meta = self._pending_event_metadata.get(key)
403
if meta is not None:
404
group_idx = meta.group_idx
405
by_group.setdefault(group_idx, []).extend(
406
maybe_convert_block_hash(h) for h in meta.block_hashes
407
)
408
+ meta.active_residencies.discard((event.medium, event.ownership))
409
+ if not meta.active_residencies:
410
+ self._pending_event_metadata.pop(key)
411
else:
412
if self.self_describing_enabled:
413
logger.warning_once(
vllm/v1/kv_offload/base.py
+1
@@ -159,6 +159,7 @@ class OffloadingEvent:
159
locality: Locality | None = None
160
# Secondary tier identifier that generated the event, or None for primary.
161
ownership: str | None = None
162
+ removal_expected: bool = False
163
164
165
"""