[KV Offload] Forward ownership in KV cache events (#52067)
Signed-off-by: Moein Khazraee <moein@nvidia.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@mistral.ai> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Moein Khazraee committed
Aug 31, 2026 at 00:28 UTC
f9d666f917941852b7d48ec20069f37480fafc28
4 files changed
+36
-10
tests/v1/kv_connector/unit/offloading_connector/test_events.py
+13
-2
@@ -141,12 +141,14 @@ def _stored_event(
141
keys: list[OffloadKey],
142
medium: Medium = _CPU_MEDIUM,
143
locality: Locality | None = None,
144
+ ownership: str | None = None,
145
) -> OffloadingEvent:
146
return OffloadingEvent(
147
keys=keys,
148
medium=medium,
149
removed=False,
150
locality=locality,
151
+ ownership=ownership,
152
)
153
154
@@ -154,12 +156,14 @@ def _removed_event(
156
keys: list[OffloadKey],
157
medium: Medium = _CPU_MEDIUM,
158
locality: Locality | None = None,
159
+ ownership: str | None = None,
160
) -> OffloadingEvent:
161
return OffloadingEvent(
162
keys=keys,
163
medium=medium,
164
removed=True,
165
locality=locality,
166
+ ownership=ownership,
167
)
168
169
@@ -534,8 +538,12 @@ def test_secondary_stored_event_does_not_mutate_cpu_metadata():
538
tracker, _, _, key = _lookup_chunk()
539
expected_metadata = dict(tracker._pending_event_metadata)
540
537
- stored = list(tracker.take_events([_stored_event([key], Medium.STORAGE)]))
541
+ stored = list(
542
+ tracker.take_events([_stored_event([key], Medium.STORAGE, ownership="custom")])
543
+ )
544
+
545
assert stored[0].token_ids == [1, 2, 3, 4]
546
+ assert stored[0].ownership == "custom"
547
assert tracker._pending_event_metadata == expected_metadata
548
549
@@ -548,9 +556,12 @@ def test_take_events_groups_removed_hashes_by_kv_group():
556
key0 = _record_chunks(tracker, req0, group0_config, num_chunks=1)[0]
557
key1 = _record_chunks(tracker, req1, group1_config, num_chunks=1)[0]
558
551
- removed = list(tracker.take_events([_removed_event([key0, key1])]))
559
+ removed = list(
560
+ tracker.take_events([_removed_event([key0, key1], ownership="custom")])
561
+ )
562
563
assert len(removed) == 2
564
+ assert {event.ownership for event in removed} == {"custom"}
565
by_group = {event.group_idx: event.block_hashes for event in removed}
566
assert by_group == {
567
0: [_wire_hash(_hash(0)), _wire_hash(_hash(1))],
vllm/distributed/kv_events.py
+13
-6
@@ -53,28 +53,31 @@ class BlockStored(KVCacheEvent):
53
token_ids: list[int]
54
block_size: int
55
56
- lora_id: int | None
56
"""Deprecated: use `lora_name` for KV block key hash.
57
Retained for backward compatibility.
58
"""
59
+ lora_id: int | None
60
61
medium: str | None
62
lora_name: str | None
63
64
- extra_keys: list[tuple[Any, ...] | None] | None = None
64
"""Extra keys used in block hash computation, one entry per block in
65
block_hashes. Each entry contains MM identifiers, LoRA name, cache_salt,
66
prompt embedding hashes, etc. for that specific block. Exposed for external
67
KV cache consumers to reconstruct block hashes.
68
"""
69
+ extra_keys: list[tuple[Any, ...] | None] | None = None
70
71
+ """Store events carry cache-spec metadata so consumers can classify and
72
+ filter groups as they are learned. Remove events only need group_idx+hash.
73
+ """
74
group_idx: int | None = None
72
- # Store events carry cache-spec metadata so consumers can classify and
73
- # filter groups as they are learned. Remove events only need group_idx+hash.
75
kv_cache_spec_kind: str | None = None
76
kv_cache_spec_sliding_window: int | None = None
76
- locality: str | None = None
77
"""LOCAL or REMOTE relative to the publisher; None means unspecified."""
78
+ locality: str | None = None
79
+ """Secondary offloading tier identifier, if generated by one."""
80
+ ownership: str | None = None
81
82
def __hash__(self) -> int:
83
return hash(
@@ -90,6 +93,7 @@ class BlockStored(KVCacheEvent):
93
self.kv_cache_spec_kind,
94
self.kv_cache_spec_sliding_window,
95
self.locality,
96
+ self.ownership,
97
)
98
)
99
@@ -98,8 +102,10 @@ class BlockRemoved(KVCacheEvent):
102
block_hashes: list[ExternalBlockHash]
103
medium: str | None
104
group_idx: int | None = None
101
- locality: str | None = None
105
"""LOCAL or REMOTE relative to the publisher; None means unspecified."""
106
+ locality: str | None = None
107
+ """Secondary offloading tier identifier, if generated by one."""
108
+ ownership: str | None = None
109
110
def __hash__(self) -> int:
111
return hash(
@@ -108,6 +114,7 @@ class BlockRemoved(KVCacheEvent):
114
self.medium,
115
self.group_idx,
116
self.locality,
117
+ self.ownership,
118
)
119
)
120
vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
+8
-2
@@ -79,7 +79,7 @@ def get_offloading_event_group_spec(
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
- is forwarded from the OffloadingEvent."""
82
+ 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, ...]
@@ -319,6 +319,7 @@ class OffloadingEventsTracker:
319
key: OffloadKey,
320
medium: Medium,
321
locality: str | None,
322
+ ownership: str | None,
323
) -> BlockStored:
324
return BlockStored(
325
block_hashes=[
@@ -332,6 +333,7 @@ class OffloadingEventsTracker:
333
lora_name=None,
334
group_idx=get_offload_group_idx(key),
335
locality=locality,
336
+ ownership=ownership,
337
)
338
339
def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
@@ -351,7 +353,9 @@ class OffloadingEventsTracker:
353
"groups and promotions not observed as a primary-tier "
354
"hit before translation."
355
)
354
- yield self._placeholder_stored(key, event.medium, locality)
356
+ yield self._placeholder_stored(
357
+ key, event.medium, locality, event.ownership
358
+ )
359
continue
360
361
yield BlockStored(
@@ -377,6 +381,7 @@ class OffloadingEventsTracker:
381
meta.kv_cache_spec.kv_cache_spec_sliding_window
382
),
383
locality=locality,
384
+ ownership=event.ownership,
385
)
386
387
def _take_removed_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
@@ -410,4 +415,5 @@ class OffloadingEventsTracker:
415
medium=_MEDIUM_TO_EVENT_STR[event.medium],
416
group_idx=group_idx,
417
locality=locality,
418
+ ownership=event.ownership,
419
)
vllm/v1/kv_offload/base.py
+2
@@ -157,6 +157,8 @@ class OffloadingEvent:
157
# True if blocks are removed, False if stored
158
removed: bool
159
locality: Locality | None = None
160
+ # Secondary tier identifier that generated the event, or None for primary.
161
+ ownership: str | None = None
162
163
164
"""