[kv_offload] Establish tier-owned KV event handling (#46544)
Signed-off-by: Change72 <changg@nvidia.com> Signed-off-by: Chang Guo <changg@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Chang Guo committed
Jul 7, 2026 at 06:55 UTC
c46ced1ee3fa0516ec5a8c4c7c328373eaa43a03
11 files changed
+61
-45
tests/v1/kv_connector/unit/offloading_connector/test_events.py
+2
-3
@@ -7,7 +7,7 @@ import torch
7
8
from tests.v1.kv_connector.unit.utils import create_vllm_config
9
from vllm.config import KVEventsConfig, KVTransferConfig
10
-from vllm.distributed.kv_events import BlockRemoved, BlockStored
10
+from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
11
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import (
12
OffloadingEventGroupSpec,
13
OffloadingEventsTracker,
@@ -28,10 +28,9 @@ from vllm.v1.kv_offload.base import (
28
OffloadKey,
29
make_offload_key,
30
)
31
-from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
31
from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
32
34
-_CPU_MEDIUM = CPULoadStoreSpec.medium()
33
+_CPU_MEDIUM = MEDIUM_CPU
34
_FULL_ATTENTION_EVENT_SPEC = OffloadingEventGroupSpec(
35
kv_cache_spec_kind=KVCacheSpecKind.FULL_ATTENTION.value,
36
kv_cache_spec_sliding_window=None,
tests/v1/kv_connector/unit/offloading_connector/utils.py
-4
@@ -73,10 +73,6 @@ class MockLoadStoreSpec(LoadStoreSpec):
73
def __init__(self, offload_keys: Iterable[OffloadKey]):
74
self.offload_keys: list[OffloadKey] = list(offload_keys)
75
76
- @staticmethod
77
- def medium() -> str:
78
- return "Mock"
79
-
76
def __repr__(self) -> str:
77
return repr(self.offload_keys)
78
tests/v1/kv_offload/cpu/test_manager.py
+21
-1
@@ -6,6 +6,7 @@ from dataclasses import dataclass
6
import numpy as np
7
import pytest
8
9
+from vllm.distributed.kv_events import MEDIUM_CPU
10
from vllm.v1.kv_offload.base import (
11
LoadStoreSpec,
12
LookupResult,
@@ -99,7 +100,7 @@ def verify_events(
100
stores: list[set[OffloadKey]] = []
101
evictions: list[set[OffloadKey]] = []
102
for event in events:
102
- assert event.medium == CPULoadStoreSpec.medium()
103
+ assert event.medium == MEDIUM_CPU
104
if event.removed:
105
evictions.append(set(event.keys))
106
else:
@@ -114,6 +115,25 @@ def verify_events(
115
assert tuple(stores) == to_key_sets(expected_stores)
116
117
118
+def test_cpu_eviction_removed_precedes_stored():
119
+ """An eviction is announced before the store that reuses its capacity."""
120
+ manager = make_cpu_manager(num_blocks=2, enable_events=True)
121
+
122
+ manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
123
+ manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
124
+ list(manager.take_events())
125
+
126
+ manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX)
127
+ manager.complete_store(to_keys([3]), _EMPTY_REQ_CTX)
128
+
129
+ events = list(manager.take_events())
130
+ removed_idx = [i for i, event in enumerate(events) if event.removed]
131
+ stored_idx = [i for i, event in enumerate(events) if not event.removed]
132
+ assert removed_idx and stored_idx, events
133
+ assert max(removed_idx) < min(stored_idx)
134
+ assert all(event.medium == manager.medium for event in events)
135
+
136
+
137
@pytest.mark.parametrize("eviction_policy", ["lru", "arc"])
138
def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
139
"""
tests/v1/kv_offload/tiering/test_tiering_offloading.py
+19
@@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
23
from vllm.v1.kv_offload.base import (
24
LookupResult,
25
OffloadingCounterMetadata,
26
+ OffloadingEvent,
27
OffloadKey,
28
OffloadPolicy,
29
ReqContext,
@@ -242,6 +243,24 @@ class TestTieringOffloadingManager:
243
if req_context.req_id not in self.manager._req_state:
244
self.manager.on_new_request(req_context)
245
246
+ def test_take_events_aggregates_tier_owned_events(self, manager_setup):
247
+ primary_event = OffloadingEvent(to_keys([1]), "CPU", removed=False)
248
+ secondary_event1 = OffloadingEvent(to_keys([2]), "tier-1", removed=False)
249
+ secondary_event2 = OffloadingEvent(to_keys([3]), "tier-2", removed=True)
250
+
251
+ self.primary_tier.take_events = MagicMock(return_value=[primary_event])
252
+ self.secondary_tier1.take_events = MagicMock(return_value=[secondary_event1])
253
+ self.secondary_tier2.take_events = MagicMock(return_value=[secondary_event2])
254
+
255
+ assert list(self.manager.take_events()) == [
256
+ primary_event,
257
+ secondary_event1,
258
+ secondary_event2,
259
+ ]
260
+ self.primary_tier.take_events.assert_called_once_with()
261
+ self.secondary_tier1.take_events.assert_called_once_with()
262
+ self.secondary_tier2.take_events.assert_called_once_with()
263
+
264
def test_basic_store_to_primary(self, manager_setup):
265
"""Test basic store operation to primary tier."""
266
blocks = to_keys(range(3))
vllm/distributed/kv_events.py
+1
@@ -43,6 +43,7 @@ class KVCacheEvent(
43
44
45
MEDIUM_GPU = "GPU"
46
+MEDIUM_CPU = "CPU"
47
48
49
class BlockStored(KVCacheEvent):
vllm/v1/kv_offload/base.py
+6
-17
@@ -12,7 +12,6 @@ from typing import TYPE_CHECKING, Any, NamedTuple, NewType
12
13
import numpy as np
14
import torch
15
-from typing_extensions import override
15
16
from vllm.logger import init_logger
17
from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
@@ -85,21 +84,12 @@ class ScheduleEndContext(NamedTuple):
84
preempted_req_ids: Collection[str]
85
86
88
-class LoadStoreSpec(ABC):
87
+class LoadStoreSpec:
88
"""
90
- Abstract metadata that encapsulates information allowing a worker
89
+ Metadata that encapsulates information allowing a worker
90
to load, and optionally also to store, blocks of KV data.
91
"""
92
94
- @staticmethod
95
- @abstractmethod
96
- def medium() -> str:
97
- """
98
- Returns a string representation of the medium type
99
- this store/load targets.
100
- """
101
- pass
102
-
93
94
@dataclass
95
class PrepareStoreOutput:
@@ -313,6 +303,10 @@ class OffloadingManager(ABC):
303
"""
304
Take the offloading events from the manager.
305
306
+ A tier manager emits only events for storage state it owns. A
307
+ composing manager may aggregate child event streams, but should not
308
+ synthesize events on behalf of a child tier.
309
+
310
Yields:
311
New OffloadingEvents collected since the last call.
312
"""
@@ -392,11 +386,6 @@ class GPULoadStoreSpec(BlockIDsLoadStoreSpec):
386
self.group_sizes: Sequence[int] = group_sizes
387
self.block_indices: Sequence[int] = block_indices
388
395
- @staticmethod
396
- @override
397
- def medium() -> str:
398
- return "GPU"
399
-
389
390
@dataclass
391
class CanonicalKVCacheTensor:
vllm/v1/kv_offload/cpu/common.py
-7
@@ -1,7 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing_extensions import override
4
-
3
from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec
4
5
@@ -15,8 +13,3 @@ class CPULoadStoreSpec(BlockIDsLoadStoreSpec):
13
"""
14
Spec for loading/storing a KV block to CPU memory.
15
"""
18
-
19
- @staticmethod
20
- @override
21
- def medium() -> str:
22
- return "CPU"
vllm/v1/kv_offload/cpu/manager.py
+2
-1
@@ -6,6 +6,7 @@ from typing import Literal
6
7
from typing_extensions import override
8
9
+from vllm.distributed.kv_events import MEDIUM_CPU
10
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
11
OffloadingConnectorStats,
12
)
@@ -51,7 +52,7 @@ class CPUOffloadingManager(OffloadingManager):
52
store_threshold: int = 1,
53
max_tracker_size: int = 64_000,
54
):
54
- self.medium: str = CPULoadStoreSpec.medium()
55
+ self.medium: str = MEDIUM_CPU
56
self._num_blocks: int = num_blocks
57
self._num_allocated_blocks: int = 0
58
self._free_list: list[int] = []
vllm/v1/kv_offload/tiering/base.py
+5
@@ -13,6 +13,7 @@ import numpy as np
13
14
from vllm.v1.kv_offload.base import (
15
LookupResult,
16
+ OffloadingEvent,
17
OffloadingMetricMetadata,
18
OffloadKey,
19
ReqContext,
@@ -208,6 +209,10 @@ class SecondaryTierManager(ABC):
209
"""
210
return False
211
212
+ def take_events(self) -> Iterable[OffloadingEvent]:
213
+ """Take KV events for storage state owned by this tier."""
214
+ return ()
215
+
216
def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
217
"""
218
Mark blocks as recently used for eviction policy.
vllm/v1/kv_offload/tiering/manager.py
+5
-11
@@ -170,7 +170,6 @@ class TieringOffloadingManager(OffloadingManager):
170
self,
171
primary_tier: CPUPrimaryTierOffloadingManager,
172
secondary_tiers: list[SecondaryTierManager] | None = None,
173
- enable_events: bool = False,
173
):
174
"""
175
Initialize the TieringOffloadingManager.
@@ -179,14 +178,11 @@ class TieringOffloadingManager(OffloadingManager):
178
primary_tier: The primary tier manager (CPU-based).
179
secondary_tiers: List of secondary tier managers (e.g., Storage,
180
Network). Can be None or empty list.
182
- enable_events: Whether to track offloading events
181
"""
182
self.primary_tier: CPUPrimaryTierOffloadingManager = primary_tier
183
self.secondary_tiers = secondary_tiers or []
184
185
self._job_id_counter: int = 0
188
- self.events: list[OffloadingEvent] | None = [] if enable_events else None
189
-
186
# Job tracking: maps job_id to metadata for all in-flight transfers.
187
# JobMetadata.is_promotion distinguishes direction:
188
# True: secondary → primary (promotion)
@@ -539,7 +535,7 @@ class TieringOffloadingManager(OffloadingManager):
535
keys: Collection[OffloadKey],
536
req_context: ReqContext,
537
success: bool = True,
542
- ):
538
+ ) -> None:
539
"""
540
Mark blocks as done storing from GPU to primary tier.
541
@@ -707,16 +703,14 @@ class TieringOffloadingManager(OffloadingManager):
703
704
@override
705
def take_events(self) -> Iterable[OffloadingEvent]:
710
- """Yield offloading events collected since the last call.
706
+ """Yield events owned by the primary and secondary tiers.
707
708
Yields:
713
- New OffloadingEvents collected since the last call.
709
+ New OffloadingEvents collected by each tier since the last call.
710
"""
715
- if self.events is not None:
716
- yield from self.events
717
- self.events.clear()
718
-
711
yield from self.primary_tier.take_events()
712
+ for tier in self.secondary_tiers:
713
+ yield from tier.take_events()
714
715
@override
716
def reset_cache(self) -> None:
vllm/v1/kv_offload/tiering/spec.py
-1
@@ -168,7 +168,6 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
168
tiering_manager = TieringOffloadingManager(
169
primary_tier=primary_tier,
170
secondary_tiers=secondary_tiers,
171
- enable_events=self.kv_events_config.enable_kv_cache_events,
171
)
172
if int(self.extra_config.get("store_threshold", 0)) >= 2:
173
raise ValueError(