iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[KV Offload] Preserve KV event metadata until final residency removal (#52068)

Browse Source 
Signed-off-by: Moein Khazraee <moein@nvidia.com>
This commit is contained in:
Moein Khazraee
2026-08-31 13:28:30 +03:00
committed by GitHub
parent eeb549a74d
commit 9acbc5360a
3 changed files with 55 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/offloading_connector/test_events.py
+34 -13
View File
@@ -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))],
vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
+20 -7
View File
@@ -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(
vllm/v1/kv_offload/base.py
+1
View File
@@ -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
"""