@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a9f7b2d41c

[feature][kv_offload] Self-describing KV events for OffloadingConnector (#43468)

Signed-off-by: Change72 <changg@nvidia.com> Co-authored-by: Claude <noreply@anthropic.com>

Change72 committed Jun 22, 2026 at 00:27 UTC a9f7b2d41c5e92f0d60ebdd3dbe04e627019c179
9 files changed +720 -62
docs/features/kv_offloading_usage.md
+1
@@ -74,6 +74,7 @@ vllm serve <model> \
74 | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
75 | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
76 | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
77 +| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
78 | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
79
80 ## Secondary Tiers
tests/v1/kv_connector/unit/offloading_connector/test_events.py new
+355
@@ -0,0 +1,355 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from unittest.mock import MagicMock
4 +
5 +import pytest
6 +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
11 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import (
12 + OffloadingEventGroupSpec,
13 + OffloadingEventsTracker,
14 +)
15 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
16 + GroupOffloadConfig,
17 +)
18 +from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
19 +from vllm.v1.kv_cache_interface import (
20 + FullAttentionSpec,
21 + KVCacheConfig,
22 + KVCacheGroupSpec,
23 + KVCacheSpecKind,
24 +)
25 +from vllm.v1.kv_offload.base import (
26 + OffloadingEvent,
27 + OffloadingKVEventsConfig,
28 + OffloadKey,
29 + make_offload_key,
30 +)
31 +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
32 +from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
33 +
34 +_CPU_MEDIUM = CPULoadStoreSpec.medium()
35 +_FULL_ATTENTION_EVENT_SPEC = OffloadingEventGroupSpec(
36 + kv_cache_spec_kind=KVCacheSpecKind.FULL_ATTENTION.value,
37 + kv_cache_spec_sliding_window=None,
38 +)
39 +
40 +
41 +def _tracker(
42 + *,
43 + enable_kv_cache_events: bool = True,
44 + self_describing_kv_events: bool = True,
45 +) -> OffloadingEventsTracker:
46 + return OffloadingEventsTracker(
47 + OffloadingKVEventsConfig(
48 + enable_kv_cache_events=enable_kv_cache_events,
49 + self_describing_kv_events=self_describing_kv_events,
50 + )
51 + )
52 +
53 +
54 +def _hash(i: int) -> BlockHash:
55 + return BlockHash(str(i).encode())
56 +
57 +
58 +def _wire_hash(block_hash: BlockHash):
59 + return maybe_convert_block_hash(block_hash)
60 +
61 +
62 +def _request(*, block_hashes: list[BlockHash], token_count: int):
63 + req = MagicMock()
64 + req.block_hashes = block_hashes
65 + req.all_token_ids = list(range(1, token_count + 1))
66 + req.lora_request = None
67 + return req
68 +
69 +
70 +def _group_config(
71 + *,
72 + group_idx: int = 0,
73 + block_size: int = 4,
74 + block_size_factor: int = 1,
75 + sliding_window_size_in_blocks: int | None = None,
76 +) -> GroupOffloadConfig:
77 + return GroupOffloadConfig(
78 + group_idx=group_idx,
79 + gpu_block_size=block_size,
80 + offloaded_block_size=block_size * block_size_factor,
81 + hash_block_size_factor=block_size_factor,
82 + sliding_window_size_in_blocks=sliding_window_size_in_blocks,
83 + kv_event_group_spec=_FULL_ATTENTION_EVENT_SPEC,
84 + )
85 +
86 +
87 +def _record_chunks(
88 + tracker: OffloadingEventsTracker,
89 + req,
90 + group_config: GroupOffloadConfig,
91 + num_chunks: int,
92 +) -> list[OffloadKey]:
93 + keys: list[OffloadKey] = []
94 + hbf = group_config.hash_block_size_factor
95 + for chunk_idx in range(num_chunks):
96 + tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1]
97 + assert tail_hash is not None
98 + key = make_offload_key(tail_hash, group_config.group_idx)
99 + tracker.record_store(req, group_config, chunk_idx, key)
100 + keys.append(key)
101 + return keys
102 +
103 +
104 +def _stored_event(keys: list[OffloadKey]) -> OffloadingEvent:
105 + return OffloadingEvent(keys=keys, medium=_CPU_MEDIUM, removed=False)
106 +
107 +
108 +def _removed_event(keys: list[OffloadKey]) -> OffloadingEvent:
109 + return OffloadingEvent(keys=keys, medium=_CPU_MEDIUM, removed=True)
110 +
111 +
112 +def test_take_events_publishes_routable_block_stored():
113 + block_size = 4
114 + tracker = _tracker()
115 + group_config = _group_config(block_size=block_size)
116 + req = _request(
117 + block_hashes=[_hash(i) for i in range(6)],
118 + token_count=block_size * 6,
119 + )
120 + keys = _record_chunks(tracker, req, group_config, num_chunks=6)
121 +
122 + batch1 = list(tracker.take_events([_stored_event(keys[:3])]))
123 + assert len(batch1) == 3
124 +
125 + for i, event in enumerate(batch1):
126 + assert isinstance(event, BlockStored)
127 + assert event.medium == _CPU_MEDIUM
128 + assert event.block_hashes == [_wire_hash(_hash(i))]
129 + assert event.block_size == block_size
130 + assert event.token_ids == list(
131 + range(i * block_size + 1, (i + 1) * block_size + 1)
132 + )
133 + if i == 0:
134 + assert event.parent_block_hash is None
135 + else:
136 + assert event.parent_block_hash == _wire_hash(_hash(i - 1))
137 + assert event.lora_id is None
138 + assert event.lora_name is None
139 + assert event.extra_keys is None
140 + assert event.group_idx == 0
141 + assert event.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
142 + assert event.kv_cache_spec_sliding_window is None
143 +
144 + batch2 = list(tracker.take_events([_stored_event(keys[3:])]))
145 + assert len(batch2) == 3
146 + assert batch2[0].parent_block_hash == batch1[-1].block_hashes[-1]
147 +
148 + assert len(tracker._pending_event_metadata) == 6
149 +
150 +
151 +def test_take_events_factor_gt_1_chunk_store_and_remove():
152 + block_size = 4
153 + block_size_factor = 3
154 + tracker = _tracker()
155 + group_config = _group_config(
156 + block_size=block_size, block_size_factor=block_size_factor
157 + )
158 + req = _request(
159 + block_hashes=[_hash(i) for i in range(6)],
160 + token_count=block_size * block_size_factor * 2,
161 + )
162 + keys = _record_chunks(tracker, req, group_config, num_chunks=2)
163 +
164 + stored = list(tracker.take_events([_stored_event(keys)]))
165 + assert len(stored) == 2
166 +
167 + expected_hashes = []
168 + for chunk_idx, event in enumerate(stored):
169 + assert isinstance(event, BlockStored)
170 + expected_chunk_hashes = [
171 + _wire_hash(_hash(i))
172 + for i in range(
173 + chunk_idx * block_size_factor,
174 + (chunk_idx + 1) * block_size_factor,
175 + )
176 + ]
177 + assert event.block_hashes == expected_chunk_hashes
178 + assert event.block_size == block_size
179 + assert len(event.token_ids) == block_size * block_size_factor
180 + if chunk_idx == 0:
181 + assert event.parent_block_hash is None
182 + else:
183 + assert event.parent_block_hash == _wire_hash(_hash(block_size_factor - 1))
184 + expected_hashes.extend(expected_chunk_hashes)
185 +
186 + assert len(tracker._pending_event_metadata) == 2
187 +
188 + removed = list(tracker.take_events([_removed_event(keys)]))
189 + assert len(removed) == 1
190 + assert isinstance(removed[0], BlockRemoved)
191 + assert removed[0].block_hashes == expected_hashes
192 + assert removed[0].medium == _CPU_MEDIUM
193 + assert removed[0].group_idx == 0
194 + assert not tracker._pending_event_metadata
195 +
196 +
197 +def test_take_events_factor_gt_1_store_is_order_independent():
198 + block_size_factor = 3
199 + tracker = _tracker()
200 + group_config = _group_config(block_size_factor=block_size_factor)
201 + req = _request(
202 + block_hashes=[_hash(i) for i in range(6)],
203 + token_count=4 * block_size_factor * 2,
204 + )
205 + keys = _record_chunks(tracker, req, group_config, num_chunks=2)
206 + unknown_key = make_offload_key(_hash(12345), 0)
207 +
208 + events = list(tracker.take_events([_stored_event([keys[1], unknown_key, keys[0]])]))
209 +
210 + assert len(events) == 3
211 + chunk1, placeholder, chunk0 = events
212 + assert [len(event.block_hashes) for event in events] == [3, 1, 3]
213 + assert placeholder.block_size == 0
214 + assert placeholder.token_ids == []
215 + assert chunk0.parent_block_hash is None
216 + assert chunk1.parent_block_hash == chunk0.block_hashes[-1]
217 +
218 +
219 +def test_take_events_opt_out_keeps_placeholders():
220 + tracker = _tracker(self_describing_kv_events=False)
221 + group_config = _group_config()
222 + req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
223 + keys = _record_chunks(tracker, req, group_config, num_chunks=3)
224 +
225 + assert not tracker.self_describing_enabled
226 + assert not tracker._pending_event_metadata
227 +
228 + events = list(
229 + tracker.take_events(
230 + [
231 + _stored_event(keys),
232 + _removed_event(keys),
233 + ]
234 + )
235 + )
236 + assert len(events) == 4
237 + for event in events[:3]:
238 + assert isinstance(event, BlockStored)
239 + assert event.block_size == 0
240 + assert event.token_ids == []
241 + assert event.parent_block_hash is None
242 + assert isinstance(events[3], BlockRemoved)
243 + assert len(events[3].block_hashes) == 3
244 +
245 +
246 +def test_record_store_skips_sliding_window_group():
247 + tracker = _tracker()
248 + group_config = _group_config(sliding_window_size_in_blocks=2)
249 + req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
250 + keys = _record_chunks(tracker, req, group_config, num_chunks=3)
251 +
252 + assert not tracker._pending_event_metadata
253 +
254 + events = list(tracker.take_events([_stored_event(keys[:1])]))
255 + assert len(events) == 1
256 + assert isinstance(events[0], BlockStored)
257 + assert events[0].block_size == 0
258 +
259 +
260 +def test_take_events_groups_removed_hashes_by_kv_group():
261 + tracker = _tracker()
262 + group0_config = _group_config(group_idx=0, block_size_factor=2)
263 + group1_config = _group_config(group_idx=1, block_size_factor=2)
264 + req0 = _request(block_hashes=[_hash(0), _hash(1)], token_count=8)
265 + req1 = _request(block_hashes=[_hash(10), _hash(11)], token_count=8)
266 + key0 = _record_chunks(tracker, req0, group0_config, num_chunks=1)[0]
267 + key1 = _record_chunks(tracker, req1, group1_config, num_chunks=1)[0]
268 +
269 + removed = list(tracker.take_events([_removed_event([key0, key1])]))
270 +
271 + assert len(removed) == 2
272 + by_group = {event.group_idx: event.block_hashes for event in removed}
273 + assert by_group == {
274 + 0: [_wire_hash(_hash(0)), _wire_hash(_hash(1))],
275 + 1: [_wire_hash(_hash(10)), _wire_hash(_hash(11))],
276 + }
277 +
278 +
279 +def test_take_events_supports_restore_after_eviction():
280 + block_size = 4
281 + tracker = _tracker()
282 + group_config = _group_config(block_size=block_size)
283 + req = _request(block_hashes=[_hash(0)], token_count=block_size)
284 + key = _record_chunks(tracker, req, group_config, num_chunks=1)[0]
285 +
286 + first_store = list(tracker.take_events([_stored_event([key])]))
287 + assert len(first_store) == 1
288 + assert isinstance(first_store[0], BlockStored)
289 + assert first_store[0].token_ids == [1, 2, 3, 4]
290 +
291 + removed = list(tracker.take_events([_removed_event([key])]))
292 + assert len(removed) == 1
293 + assert isinstance(removed[0], BlockRemoved)
294 + assert not tracker._pending_event_metadata
295 +
296 + req.all_token_ids = [5, 6, 7, 8]
297 + tracker.record_store(req, group_config, offload_block_idx=0, offload_key=key)
298 +
299 + second_store = list(tracker.take_events([_stored_event([key])]))
300 + assert len(second_store) == 1
301 + assert isinstance(second_store[0], BlockStored)
302 + assert second_store[0].token_ids == [5, 6, 7, 8]
303 +
304 +
305 +def test_reset_cache_clears_side_table():
306 + tracker = _tracker()
307 + group_config = _group_config()
308 + req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
309 + _record_chunks(tracker, req, group_config, num_chunks=3)
310 +
311 + assert tracker._pending_event_metadata
312 +
313 + tracker.reset()
314 +
315 + assert not tracker._pending_event_metadata
316 +
317 +
318 +def test_tiering_rejects_self_describing_kv_events():
319 + vllm_config = create_vllm_config(
320 + block_size=4,
321 + max_num_batched_tokens=16,
322 + disable_hybrid_kv_cache_manager=False,
323 + )
324 + vllm_config.kv_transfer_config = KVTransferConfig(
325 + kv_connector="OffloadingConnector",
326 + kv_role="kv_both",
327 + kv_connector_extra_config={
328 + "spec_name": "TieringOffloadingSpec",
329 + "cpu_bytes_to_use": 1 << 20,
330 + "self_describing_kv_events": True,
331 + "secondary_tiers": [{"type": "example"}],
332 + },
333 + )
334 + vllm_config.kv_events_config = KVEventsConfig(
335 + enable_kv_cache_events=True,
336 + publisher="null",
337 + )
338 + kv_cache_config = KVCacheConfig(
339 + num_blocks=0,
340 + kv_cache_tensors=[],
341 + kv_cache_groups=[
342 + KVCacheGroupSpec(
343 + ["layer"],
344 + FullAttentionSpec(
345 + block_size=4,
346 + num_kv_heads=1,
347 + head_size=1,
348 + dtype=torch.float32,
349 + ),
350 + )
351 + ],
352 + )
353 +
354 + with pytest.raises(ValueError, match="TieringOffloadingSpec"):
355 + TieringOffloadingSpec(vllm_config, kv_cache_config)
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
-29
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from collections.abc import Iterable
3 from types import SimpleNamespace
4 from unittest.mock import MagicMock
5
@@ -12,19 +11,16 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (
11 to_keys,
12 )
13 from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
15 -from vllm.distributed.kv_events import BlockRemoved, BlockStored
14 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
15 OffloadingConnectorScheduler,
16 RequestOffloadState,
17 )
20 -from vllm.v1.core.kv_cache_utils import BlockHash
18 from vllm.v1.kv_cache_interface import (
19 FullAttentionSpec,
20 KVCacheGroupSpec,
21 SlidingWindowSpec,
22 )
23 from vllm.v1.kv_offload.base import (
27 - OffloadingEvent,
24 OffloadingManager,
25 OffloadPolicy,
26 ReqContext,
@@ -146,31 +142,6 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
142 runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
143 runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(3, 4, 5))
144
149 - # test take_events
150 - def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
151 - return [BlockHash(str(i).encode()) for i in int_hashes]
152 -
153 - def take_events() -> Iterable[OffloadingEvent]:
154 - yield OffloadingEvent(keys=to_keys([1, 2, 3]), medium="A", removed=False)
155 - yield OffloadingEvent(keys=to_keys([4, 5, 6]), medium="B", removed=True)
156 -
157 - runner.manager.take_events.side_effect = take_events
158 - events = list(runner.scheduler_connector.take_events())
159 - assert len(events) == 2
160 - event = events[0]
161 - assert isinstance(event, BlockStored)
162 - assert event.block_hashes == to_hashes([1, 2, 3])
163 - assert event.block_size == 0
164 - assert event.medium == "A"
165 - assert event.token_ids == []
166 - assert event.parent_block_hash is None
167 - assert event.lora_id is None
168 - assert event.lora_name is None
169 - event = events[1]
170 - assert isinstance(event, BlockRemoved)
171 - assert event.block_hashes == to_hashes([4, 5, 6])
172 - assert event.medium == "B"
173 -
145
146 @pytest.mark.parametrize("async_scheduling", [True, False])
147 def test_request_preemption(request_runner, async_scheduling: bool):
tests/v1/kv_connector/unit/offloading_connector/utils.py
+16 -1
@@ -14,7 +14,12 @@ from tests.v1.kv_connector.unit.utils import (
14 create_vllm_config,
15 )
16 from vllm import SamplingParams
17 -from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config
17 +from vllm.config import (
18 + KVEventsConfig,
19 + KVTransferConfig,
20 + VllmConfig,
21 + set_current_vllm_config,
22 +)
23 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
24 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
25 OffloadingConnectorMetadata,
@@ -198,6 +203,9 @@ class RequestRunner:
203 "spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils", # noqa: E501
204 # Preserve legacy behavior for tests; new opt-in tests override.
205 "offload_prompt_only": False,
206 + # Exercise the self-describing KV events path by default;
207 + # opt-out tests override this to cover the legacy placeholders.
208 + "self_describing_kv_events": True,
209 }
210 if block_size_factor > 1:
211 extra_config["block_size"] = block_size * block_size_factor
@@ -209,6 +217,13 @@ class RequestRunner:
217 kv_role="kv_both",
218 kv_connector_extra_config=extra_config,
219 )
220 + vllm_config.kv_events_config = KVEventsConfig(
221 + # Enable so the offloading events tracker is active, but use the
222 + # null publisher: these tests drain take_events directly and a
223 + # real ZMQ publisher would bind a port per test.
224 + enable_kv_cache_events=True,
225 + publisher="null",
226 + )
227
228 if kv_cache_groups is None:
229 kv_cache_groups = [
vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py new
+286
@@ -0,0 +1,286 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Self-describing KV cache events for the offloading connector.
4 +
5 +The OffloadingManager identifies an offloaded chunk only by its OffloadKey,
6 +so its raw events carry no token ids, parent hash, or block size.
7 +:class:`OffloadingEventsTracker` snapshots each chunk's full ``BlockStored``
8 +payload while the ``Request`` is alive and publishes stores as block-granular
9 +payloads: a chunk event may carry multiple constituent per-block hashes, and
10 +evictions fan out to the same hashes. Chunks overlapping a non-chunk-aligned
11 +shared prefix re-announce the shared hashes once per chunk; consumers are
12 +expected to deduplicate (reference-count) repeated store/remove announcements
13 +of the same hash. Opt-in via
14 +``kv_connector_extra_config["self_describing_kv_events"]``; inert unless
15 +KV cache events are enabled. See the PR description for the full design.
16 +"""
17 +
18 +from collections.abc import Iterable
19 +from dataclasses import dataclass
20 +from typing import TYPE_CHECKING, Any, NamedTuple
21 +
22 +from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
23 +from vllm.logger import init_logger
24 +from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
25 +from vllm.v1.kv_cache_interface import (
26 + KVCacheGroupSpec,
27 + get_kv_cache_spec_kind,
28 + get_kv_cache_spec_sliding_window,
29 +)
30 +from vllm.v1.kv_offload.base import (
31 + OffloadingEvent,
32 + OffloadingKVEventsConfig,
33 + OffloadKey,
34 + get_offload_block_hash,
35 + get_offload_group_idx,
36 +)
37 +from vllm.v1.request import Request
38 +
39 +if TYPE_CHECKING:
40 + from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
41 + GroupOffloadConfig,
42 + )
43 +
44 +logger = init_logger(__name__)
45 +
46 +
47 +class OffloadingEventGroupSpec(NamedTuple):
48 + kv_cache_spec_kind: str | None
49 + kv_cache_spec_sliding_window: int | None
50 +
51 +
52 +def get_offloading_event_group_spec(
53 + kv_cache_group: KVCacheGroupSpec,
54 +) -> OffloadingEventGroupSpec:
55 + kv_cache_spec = kv_cache_group.kv_cache_spec
56 + return OffloadingEventGroupSpec(
57 + kv_cache_spec_kind=get_kv_cache_spec_kind(kv_cache_spec).value,
58 + kv_cache_spec_sliding_window=get_kv_cache_spec_sliding_window(kv_cache_spec),
59 + )
60 +
61 +
62 +@dataclass(slots=True)
63 +class _OffloadEventMetadata:
64 + """BlockStored payload snapshot for one OffloadKey, captured at store
65 + time and kept until the matching eviction event. ``medium`` is forwarded
66 + from the OffloadingEvent."""
67 +
68 + # The chunk's constituent block hashes; the last one is the OffloadKey.
69 + block_hashes: tuple[BlockHash, ...]
70 + parent_block_hash: BlockHash | None
71 + token_ids: tuple[int, ...]
72 + block_size: int
73 + lora_id: int | None
74 + lora_name: str | None
75 + # Deferred: needs the same incremental curr_mm_idx handling as GPU events.
76 + extra_keys: tuple[tuple[Any, ...] | None, ...] | None
77 + group_idx: int
78 + kv_cache_spec: OffloadingEventGroupSpec
79 +
80 +
81 +class OffloadingEventsTracker:
82 + """Tracks offloaded chunks' KV event payloads from store to eviction.
83 +
84 + The scheduler calls :meth:`record_store` from ``_build_store_jobs``
85 + while the ``Request`` is available, and routes the manager's raw
86 + :class:`OffloadingEvent` stream through :meth:`take_events`. All state
87 + is bounded by the CPU pool capacity and cleared by :meth:`reset`.
88 + """
89 +
90 + def __init__(self, config: OffloadingKVEventsConfig):
91 + self.config = config
92 + self.self_describing_enabled = (
93 + config.enable_kv_cache_events and config.self_describing_kv_events
94 + )
95 +
96 + # OffloadKey -> payload snapshot, kept until the eviction event so
97 + # BlockRemoved can fan out. Bounded: one entry per offloaded chunk.
98 + self._pending_event_metadata: dict[OffloadKey, _OffloadEventMetadata] = {}
99 +
100 + def record_store(
101 + self,
102 + req: Request,
103 + group_config: "GroupOffloadConfig",
104 + offload_block_idx: int,
105 + offload_key: OffloadKey,
106 + ) -> None:
107 + """Snapshot the KV cache event payload for one offloaded chunk.
108 +
109 + No-op when self-describing event capture is disabled or for
110 + sliding-window / SSM groups, which keep the legacy placeholder payload.
111 + """
112 + if not self.self_describing_enabled:
113 + return
114 + if group_config.sliding_window_size_in_blocks is not None:
115 + return
116 + meta = self._build_event_metadata(req, group_config, offload_block_idx)
117 + self._pending_event_metadata[offload_key] = meta
118 +
119 + def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]:
120 + """Translate raw OffloadingEvents into self-describing KV events.
121 +
122 + Complete metadata is available only for full-attention groups when
123 + the tracker is enabled. Other shapes retain the legacy placeholder
124 + payload so consumers can ignore them.
125 +
126 + Yields:
127 + ``BlockStored`` or ``BlockRemoved`` events corresponding to
128 + the underlying :class:`OffloadingEvent` stream.
129 + """
130 + for event in events:
131 + if event.removed:
132 + yield from self._take_removed_event(event)
133 + else:
134 + yield from self._take_stored_event(event)
135 +
136 + def reset(self) -> None:
137 + """Drop all tracked state; pending payloads are stale after a
138 + manager cache reset."""
139 + self._pending_event_metadata.clear()
140 +
141 + def _build_event_metadata(
142 + self,
143 + req: Request,
144 + group_config: "GroupOffloadConfig",
145 + offload_block_idx: int,
146 + ) -> _OffloadEventMetadata:
147 + """Build the payload snapshot for one offloaded chunk: its
148 + constituent per-block hashes, the whole chunk's tokens, and the
149 + per-block ``block_size``."""
150 + hbf = group_config.hash_block_size_factor
151 + assert hbf > 0
152 + assert offload_block_idx >= 0
153 + # per-block token count (= the GPU/hash block size)
154 + sub_block_size = group_config.offloaded_block_size // hbf
155 + # chunk c covers hash-blocks [c*hbf, (c+1)*hbf); its tail block's hash
156 + # is the chunk's OffloadKey.
157 + first_hash_idx = offload_block_idx * hbf
158 + last_hash_idx = first_hash_idx + hbf
159 + assert first_hash_idx >= 0
160 + assert last_hash_idx <= len(req.block_hashes)
161 + chunk_hashes: list[BlockHash] = []
162 + for block_hash in req.block_hashes[first_hash_idx:last_hash_idx]:
163 + assert block_hash is not None
164 + chunk_hashes.append(block_hash)
165 + assert len(chunk_hashes) == hbf
166 +
167 + if group_config.sliding_window_size_in_blocks is not None:
168 + # record_store filters these out before calling this helper.
169 + raise AssertionError("self-describing events only support full attention")
170 +
171 + parent_block_hash: BlockHash | None
172 + if first_hash_idx == 0:
173 + parent_block_hash = None
174 + else:
175 + parent_block_hash = req.block_hashes[first_hash_idx - 1]
176 + assert parent_block_hash is not None
177 +
178 + tok_start = offload_block_idx * group_config.offloaded_block_size
179 + tok_end = tok_start + group_config.offloaded_block_size
180 + assert tok_end <= len(req.all_token_ids)
181 + token_ids = tuple(req.all_token_ids[tok_start:tok_end])
182 +
183 + lora_id: int | None = None
184 + lora_name: str | None = None
185 + if req.lora_request is not None:
186 + lora_id = req.lora_request.adapter_id
187 + lora_name = req.lora_request.name
188 +
189 + return _OffloadEventMetadata(
190 + block_hashes=tuple(chunk_hashes),
191 + parent_block_hash=parent_block_hash,
192 + token_ids=token_ids,
193 + block_size=sub_block_size,
194 + lora_id=lora_id,
195 + lora_name=lora_name,
196 + extra_keys=None,
197 + group_idx=group_config.group_idx,
198 + kv_cache_spec=group_config.kv_event_group_spec,
199 + )
200 +
201 + def _placeholder_stored(self, key: OffloadKey, medium: str) -> BlockStored:
202 + return BlockStored(
203 + block_hashes=[
204 + maybe_convert_block_hash(BlockHash(get_offload_block_hash(key)))
205 + ],
206 + parent_block_hash=None,
207 + token_ids=[],
208 + lora_id=None,
209 + block_size=0,
210 + medium=medium,
211 + lora_name=None,
212 + group_idx=get_offload_group_idx(key),
213 + )
214 +
215 + def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
216 + # Metadata is read, NOT popped: the entry must survive until the
217 + # eviction event so BlockRemoved can fan out to the same hashes.
218 + # Events are self-contained (own parent), so key order is free.
219 + for key in event.keys:
220 + meta = self._pending_event_metadata.get(key)
221 + if meta is None:
222 + if self.self_describing_enabled:
223 + # Expected for unsupported shapes; warn once only.
224 + logger.warning_once(
225 + "OffloadingEventsTracker: no event metadata for "
226 + "offload key during BlockStored emission; emitting a "
227 + "placeholder payload. Expected for non-full-attention "
228 + "groups; otherwise indicates a missing populate path."
229 + )
230 + yield self._placeholder_stored(key, event.medium)
231 + continue
232 +
233 + yield BlockStored(
234 + block_hashes=list(
235 + maybe_convert_block_hash(h) for h in meta.block_hashes
236 + ),
237 + parent_block_hash=(
238 + maybe_convert_block_hash(meta.parent_block_hash)
239 + if meta.parent_block_hash is not None
240 + else None
241 + ),
242 + token_ids=list(meta.token_ids),
243 + block_size=meta.block_size,
244 + lora_id=meta.lora_id,
245 + medium=event.medium,
246 + lora_name=meta.lora_name,
247 + extra_keys=(
248 + list(meta.extra_keys) if meta.extra_keys is not None else None
249 + ),
250 + group_idx=meta.group_idx,
251 + kv_cache_spec_kind=meta.kv_cache_spec.kv_cache_spec_kind,
252 + kv_cache_spec_sliding_window=(
253 + meta.kv_cache_spec.kv_cache_spec_sliding_window
254 + ),
255 + )
256 +
257 + def _take_removed_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
258 + # Keep group_idx unambiguous if a manager batch spans groups.
259 + by_group: dict[int, list] = {}
260 + for key in event.keys:
261 + meta = self._pending_event_metadata.pop(key, None)
262 + if meta is not None:
263 + group_idx = meta.group_idx
264 + by_group.setdefault(group_idx, []).extend(
265 + maybe_convert_block_hash(h) for h in meta.block_hashes
266 + )
267 + else:
268 + if self.self_describing_enabled:
269 + logger.warning_once(
270 + "OffloadingEventsTracker: no event metadata for "
271 + "offload key during BlockRemoved emission; emitting a "
272 + "placeholder removal. Expected if the matching store "
273 + "used the legacy placeholder payload; otherwise "
274 + "indicates missing store metadata."
275 + )
276 + group_idx = get_offload_group_idx(key)
277 + by_group.setdefault(group_idx, []).append(
278 + maybe_convert_block_hash(BlockHash(get_offload_block_hash(key)))
279 + )
280 +
281 + for group_idx, hashes in by_group.items():
282 + yield BlockRemoved(
283 + block_hashes=hashes,
284 + medium=event.medium,
285 + group_idx=group_idx,
286 + )
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+32 -19
@@ -5,7 +5,7 @@ from dataclasses import dataclass, field
5 from itertools import islice
6 from typing import Any, NamedTuple
7
8 -from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
8 +from vllm.distributed.kv_events import KVCacheEvent
9 from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
10 from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
11 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
@@ -14,6 +14,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
14 ReqId,
15 TransferJob,
16 )
17 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import (
18 + OffloadingEventGroupSpec,
19 + OffloadingEventsTracker,
20 + get_offloading_event_group_spec,
21 +)
22 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
23 OffloadingConnectorStats,
24 _TransferMetricName,
@@ -36,7 +41,6 @@ from vllm.v1.kv_offload.base import (
41 OffloadPolicy,
42 ReqContext,
43 RequestOffloadingContext,
39 - get_offload_block_hash,
44 make_offload_key,
45 )
46 from vllm.v1.outputs import KVConnectorOutput
@@ -69,6 +73,9 @@ class GroupOffloadConfig(NamedTuple):
73 gpu_block_size: int
74 offloaded_block_size: int
75 hash_block_size_factor: int
76 + # KV cache spec metadata propagated onto emitted BlockStored events so
77 + # KV-aware consumers can classify and filter the group.
78 + kv_event_group_spec: OffloadingEventGroupSpec
79 # None below means full attention
80 sliding_window_size_in_blocks: int | None
81 # Number of this group's offloaded blocks per full-attention alignment
@@ -200,6 +207,9 @@ class SchedulerOffloadConfig(NamedTuple):
207 alignment_block_count=_alignment_block_count(
208 gpu_block_size * spec.block_size_factor, sw
209 ),
210 + kv_event_group_spec=get_offloading_event_group_spec(
211 + spec.kv_cache_config.kv_cache_groups[idx]
212 + ),
213 is_eagle_group=idx in eagle_groups,
214 )
215 for idx, gpu_block_size in enumerate(spec.gpu_block_size)
@@ -361,6 +371,8 @@ class OffloadingConnectorScheduler:
371 # be freed before a request finishes).
372 self._block_id_to_pending_jobs: dict[int, set[int]] = {}
373
374 + self._events_tracker = OffloadingEventsTracker(spec.kv_events_config)
375 +
376 def _generate_job_id(self) -> int:
377 job_id = self._job_counter
378 self._job_counter += 1
@@ -934,6 +946,11 @@ class OffloadingConnectorScheduler:
946 continue
947
948 offloaded_block_idx = start_block_idx + idx
949 +
950 + self._events_tracker.record_store(
951 + req, group_config, offloaded_block_idx, offload_key
952 + )
953 +
954 gpu_block_idx = offloaded_block_idx * block_size_factor
955 for i in range(block_size_factor):
956 block_id = block_ids[gpu_block_idx + i]
@@ -1184,25 +1201,17 @@ class OffloadingConnectorScheduler:
1201 return False, None
1202
1203 def take_events(self) -> Iterable[KVCacheEvent]:
1187 - """Take the KV cache events from the connector.
1204 + """Drain pending KV cache events.
1205
1189 - Returns:
1190 - A list of KV cache events.
1206 + Complete metadata is available only when self-describing KV events
1207 + are enabled, and only for full-attention groups. Other shapes retain
1208 + the previous placeholder payload so consumers can ignore them.
1209 +
1210 + Yields:
1211 + ``BlockStored`` or ``BlockRemoved`` events corresponding to
1212 + the underlying :class:`OffloadingEvent` stream.
1213 """
1192 - for event in self.manager.take_events():
1193 - block_hashes = [get_offload_block_hash(key) for key in event.keys]
1194 - if event.removed:
1195 - yield BlockRemoved(block_hashes=block_hashes, medium=event.medium)
1196 - else:
1197 - yield BlockStored(
1198 - block_hashes=block_hashes,
1199 - parent_block_hash=None,
1200 - token_ids=[],
1201 - lora_id=None,
1202 - block_size=0,
1203 - medium=event.medium,
1204 - lora_name=None,
1205 - )
1214 + yield from self._events_tracker.take_events(self.manager.take_events())
1215
1216 def reset_cache(self) -> None:
1217 """Reset the offloading manager cache, evicting all stored blocks."""
@@ -1238,6 +1247,10 @@ class OffloadingConnectorScheduler:
1247 self._jobs.clear()
1248 self._block_id_to_pending_jobs.clear()
1249
1250 + # The manager pool is empty; pending event payloads and announced
1251 + # reference counts are stale.
1252 + self._events_tracker.reset()
1253 +
1254 # Note: _current_batch_jobs_to_flush is intentionally NOT cleared.
1255 # The load flush IDs collected above must be delivered to workers.
1256 if self._blocks_being_loaded is not None:
vllm/v1/kv_offload/base.py
+19
@@ -147,6 +147,16 @@ class OffloadingHistogramMetadata(OffloadingMetricMetadata):
147 buckets: tuple[float, ...] | None = None
148
149
150 +@dataclass(frozen=True)
151 +class OffloadingKVEventsConfig:
152 + # Global vLLM KV event publishing flag. When false, connector-specific
153 + # event capture must stay inert because take_events() is not drained.
154 + enable_kv_cache_events: bool
155 + # OffloadingConnector opt-in for self-describing BlockStored payloads.
156 + # Effective only when enable_kv_cache_events is true.
157 + self_describing_kv_events: bool
158 +
159 +
160 class OffloadingManager(ABC):
161 @abstractmethod
162 def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
@@ -443,6 +453,15 @@ class OffloadingSpec(ABC):
453 kv_transfer_config = vllm_config.kv_transfer_config
454 assert kv_transfer_config is not None
455 self.extra_config = kv_transfer_config.kv_connector_extra_config
456 + kv_events_config = vllm_config.kv_events_config
457 + self.kv_events_config = OffloadingKVEventsConfig(
458 + enable_kv_cache_events=(
459 + kv_events_config is not None and kv_events_config.enable_kv_cache_events
460 + ),
461 + self_describing_kv_events=bool(
462 + self.extra_config.get("self_describing_kv_events", False)
463 + ),
464 + )
465
466 # When True, only prompt (prefill) blocks are offloaded; decode-phase
467 # blocks (KV generated after the prompt) are skipped. Useful when prior
vllm/v1/kv_offload/cpu/spec.py
+1 -6
@@ -114,11 +114,6 @@ class CPUOffloadingSpec(OffloadingSpec):
114 @override
115 def get_manager(self) -> OffloadingManager:
116 if not self._manager:
117 - kv_events_config = self.vllm_config.kv_events_config
118 - enable_events = (
119 - kv_events_config is not None and kv_events_config.enable_kv_cache_events
120 - )
121 -
117 # store_threshold: how many times a block must appear in lookup()
118 # before it is eligible for CPU offloading. Values < 2 disable
119 # filtering (a threshold of 1 equals no filter; 0 is the default).
@@ -130,7 +125,7 @@ class CPUOffloadingSpec(OffloadingSpec):
125 self._manager = CPUOffloadingManager(
126 num_blocks=self.num_blocks,
127 cache_policy=self.eviction_policy, # type: ignore[arg-type]
133 - enable_events=enable_events,
128 + enable_events=self.kv_events_config.enable_kv_cache_events,
129 store_threshold=store_threshold,
130 max_tracker_size=max_tracker_size,
131 )
vllm/v1/kv_offload/tiering/spec.py
+10 -7
@@ -69,6 +69,14 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
69 super().__init__(vllm_config, kv_cache_config)
70 # Redeclare for mypy: parent sets this but `--follow-imports skip` hides it
71 self._manager: OffloadingManager | None = None
72 + if self.kv_events_config.self_describing_kv_events:
73 + raise ValueError(
74 + "self_describing_kv_events is not supported by "
75 + "TieringOffloadingSpec. Tier promotions can emit primary-tier "
76 + "store events that do not correspond to GPU store jobs, so the "
77 + "current self-describing side table cannot describe them "
78 + "correctly."
79 + )
80
81 # Parse secondary tier configurations
82 self.secondary_tier_configs = self.extra_config.get("secondary_tiers", [])
@@ -91,11 +99,6 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
99 TieringOffloadingManager instance
100 """
101 if not self._manager:
94 - kv_events_config = self.vllm_config.kv_events_config
95 - enable_events = (
96 - kv_events_config is not None and kv_events_config.enable_kv_cache_events
97 - )
98 -
102 # Create scheduler-side SharedOffloadRegion (rank=None) so the
103 # primary tier can eagerly create a memoryview over _base.
104 scheduler_mmap = SharedOffloadRegion(
@@ -111,7 +114,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
114 primary_tier = CPUPrimaryTierOffloadingManager(
115 num_blocks=self.num_blocks,
116 cache_policy=self.eviction_policy, # type: ignore[arg-type]
114 - enable_events=enable_events,
117 + enable_events=self.kv_events_config.enable_kv_cache_events,
118 mmap_region=scheduler_mmap,
119 )
120
@@ -143,7 +146,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
146 tiering_manager = TieringOffloadingManager(
147 primary_tier=primary_tier,
148 secondary_tiers=secondary_tiers,
146 - enable_events=enable_events,
149 + enable_events=self.kv_events_config.enable_kv_cache_events,
150 )
151 if int(self.extra_config.get("store_threshold", 0)) >= 2:
152 raise ValueError(