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

[KV Offload] Support self-describing KV events with TieringOffloadingSpec (#48679)

Signed-off-by: Change72 <changg@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Chang Guo committed Jul 20, 2026 at 06:41 UTC f007cceb422599330495b202632d78c7b3b89e61
6 files changed +452 -65
docs/features/kv_offloading_usage.md
+1 -1
@@ -75,7 +75,7 @@ vllm serve <model> \
75 | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
76 | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
77 | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
78 -| `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, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
78 +| `self_describing_kv_events` | no | `false` | both | 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. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
79 | `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). |
80
81 ## Secondary Tiers
tests/v1/kv_connector/unit/offloading_connector/test_events.py
+143 -14
@@ -7,7 +7,13 @@ 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 MEDIUM_CPU, MEDIUM_FS, BlockRemoved, BlockStored
10 +from vllm.distributed.kv_events import (
11 + MEDIUM_CPU,
12 + MEDIUM_FS,
13 + MEDIUM_OBJ,
14 + BlockRemoved,
15 + BlockStored,
16 +)
17 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
18 build_offloading_config,
19 )
@@ -62,8 +68,9 @@ def _wire_hash(block_hash: BlockHash):
68 return maybe_convert_block_hash(block_hash)
69
70
65 -def _request(*, block_hashes: list[BlockHash], token_count: int):
71 +def _request(*, block_hashes: list[BlockHash], token_count: int, req_id: str = "req"):
72 req = MagicMock()
73 + req.request_id = req_id
74 req.block_hashes = block_hashes
75 req.all_token_ids = list(range(1, token_count + 1))
76 req.lora_request = None
@@ -104,10 +111,32 @@ def _record_chunks(
111 return keys
112
113
114 +def _record_lookup_chunks(
115 + tracker: OffloadingEventsTracker,
116 + req,
117 + group_config: GroupOffloadConfig,
118 + num_chunks: int,
119 +) -> list[OffloadKey]:
120 + keys: list[OffloadKey] = []
121 + hbf = group_config.hashes_per_chunk
122 + for chunk_idx in range(num_chunks):
123 + tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1]
124 + assert tail_hash is not None
125 + key = make_offload_key(tail_hash, group_config.group_idx)
126 + tracker.record_lookup(
127 + req,
128 + group_config,
129 + chunk_idx,
130 + key,
131 + )
132 + keys.append(key)
133 + return keys
134 +
135 +
136 def _stored_event(
137 keys: list[OffloadKey],
109 - locality: Locality | None = None,
138 medium: str = _CPU_MEDIUM,
139 + locality: Locality | None = None,
140 ) -> OffloadingEvent:
141 return OffloadingEvent(
142 keys=keys,
@@ -119,8 +148,8 @@ def _stored_event(
148
149 def _removed_event(
150 keys: list[OffloadKey],
122 - locality: Locality | None = None,
151 medium: str = _CPU_MEDIUM,
152 + locality: Locality | None = None,
153 ) -> OffloadingEvent:
154 return OffloadingEvent(
155 keys=keys,
@@ -130,6 +159,21 @@ def _removed_event(
159 )
160
161
162 +def _lookup_chunk() -> tuple[
163 + OffloadingEventsTracker, MagicMock, GroupOffloadConfig, OffloadKey
164 +]:
165 + tracker = _tracker()
166 + req = _request(block_hashes=[_hash(0)], token_count=4)
167 + group_config = _group_config()
168 + key = _record_lookup_chunks(
169 + tracker,
170 + req,
171 + group_config,
172 + num_chunks=1,
173 + )[0]
174 + return tracker, req, group_config, key
175 +
176 +
177 def test_take_events_forwards_locality_to_rich_store():
178 tracker = _tracker()
179 req = _request(block_hashes=[_hash(0)], token_count=4)
@@ -220,18 +264,37 @@ def test_take_events_publishes_routable_block_stored():
264 assert len(tracker._pending_event_metadata) == 6
265
266
223 -def test_take_events_factor_gt_1_chunk_store_and_remove():
267 +def test_promotion_emits_full_cpu_stored_event():
268 + tracker, _, _, key = _lookup_chunk()
269 +
270 + [event] = tracker.take_events([_stored_event([key])])
271 +
272 + assert isinstance(event, BlockStored)
273 + assert event.medium == MEDIUM_CPU
274 + assert event.block_hashes == [_wire_hash(_hash(0))]
275 + assert event.parent_block_hash is None
276 + assert event.token_ids == [1, 2, 3, 4]
277 + assert event.block_size == 4
278 + assert event.lora_id is None
279 + assert event.lora_name is None
280 + assert event.extra_keys is None
281 + assert event.group_idx == 0
282 + assert event.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
283 + assert event.kv_cache_spec_sliding_window is None
284 +
285 +
286 +def test_lookup_promotion_factor_gt_1_store_and_remove():
287 block_size = 4
225 - blocks_per_chunk = 3
288 + blocks_per_chunk = 2
289 tracker = _tracker()
290 group_config = _group_config(
291 block_size=block_size, blocks_per_chunk=blocks_per_chunk
292 )
293 req = _request(
231 - block_hashes=[_hash(i) for i in range(6)],
294 + block_hashes=[_hash(i) for i in range(4)],
295 token_count=block_size * blocks_per_chunk * 2,
296 )
234 - keys = _record_chunks(tracker, req, group_config, num_chunks=2)
297 + keys = _record_lookup_chunks(tracker, req, group_config, num_chunks=2)
298
299 stored = list(tracker.take_events([_stored_event(keys)]))
300 assert len(stored) == 2
@@ -293,6 +356,7 @@ def test_take_events_opt_out_keeps_placeholders():
356 group_config = _group_config()
357 req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
358 keys = _record_chunks(tracker, req, group_config, num_chunks=3)
359 + _record_lookup_chunks(tracker, req, group_config, num_chunks=3)
360
361 assert not tracker.self_describing_enabled
362 assert not tracker._pending_event_metadata
@@ -315,11 +379,21 @@ def test_take_events_opt_out_keeps_placeholders():
379 assert len(events[3].block_hashes) == 3
380
381
318 -def test_record_store_skips_sliding_window_group():
382 +@pytest.mark.parametrize(
383 + "sliding_window_size_in_chunks",
384 + [1, 2],
385 + ids=["ssm", "sliding-window"],
386 +)
387 +def test_event_metadata_skips_non_full_attention_group(
388 + sliding_window_size_in_chunks: int,
389 +):
390 tracker = _tracker()
320 - group_config = _group_config(sliding_window_size_in_chunks=2)
391 + group_config = _group_config(
392 + sliding_window_size_in_chunks=sliding_window_size_in_chunks
393 + )
394 req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
395 keys = _record_chunks(tracker, req, group_config, num_chunks=3)
396 + _record_lookup_chunks(tracker, req, group_config, num_chunks=3)
397
398 assert not tracker._pending_event_metadata
399
@@ -329,6 +403,57 @@ def test_record_store_skips_sliding_window_group():
403 assert events[0].block_size == 0
404
405
406 +def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit():
407 + tracker = _tracker()
408 + block_hashes = [_hash(0), _hash(1)]
409 + req = _request(block_hashes=block_hashes, token_count=8)
410 + group_config = _group_config(blocks_per_chunk=2)
411 + key = _record_chunks(tracker, req, group_config, num_chunks=1)[0]
412 + confirmed_meta = tracker._pending_event_metadata[key]
413 + lookup_req = _request(
414 + block_hashes=block_hashes,
415 + token_count=8,
416 + req_id="new-request",
417 + )
418 +
419 + tracker.record_lookup(
420 + lookup_req,
421 + group_config,
422 + 0,
423 + key,
424 + )
425 + assert tracker._pending_event_metadata[key] is confirmed_meta
426 +
427 + removed = list(tracker.take_events([_removed_event([key])]))
428 + assert len(removed) == 1
429 + assert removed[0].block_hashes == [
430 + _wire_hash(_hash(0)),
431 + _wire_hash(_hash(1)),
432 + ]
433 +
434 + stored = list(tracker.take_events([_stored_event([key])]))
435 + assert len(stored) == 1
436 + assert stored[0].block_size == 0
437 + assert stored[0].token_ids == []
438 +
439 + tracker.record_lookup(lookup_req, group_config, 0, key)
440 + removed = list(tracker.take_events([_removed_event([key])]))
441 + assert removed[0].block_hashes == [
442 + _wire_hash(_hash(0)),
443 + _wire_hash(_hash(1)),
444 + ]
445 +
446 +
447 +@pytest.mark.parametrize("medium", [MEDIUM_FS, MEDIUM_OBJ])
448 +def test_secondary_stored_event_does_not_mutate_cpu_metadata(medium: str):
449 + tracker, _, _, key = _lookup_chunk()
450 + expected_metadata = dict(tracker._pending_event_metadata)
451 +
452 + stored = list(tracker.take_events([_stored_event([key], medium)]))
453 + assert stored[0].token_ids == [1, 2, 3, 4]
454 + assert tracker._pending_event_metadata == expected_metadata
455 +
456 +
457 def test_take_events_groups_removed_hashes_by_kv_group():
458 tracker = _tracker()
459 group0_config = _group_config(group_idx=0, blocks_per_chunk=2)
@@ -378,7 +503,7 @@ def test_reset_cache_clears_side_table():
503 tracker = _tracker()
504 group_config = _group_config()
505 req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
381 - _record_chunks(tracker, req, group_config, num_chunks=3)
506 + _record_lookup_chunks(tracker, req, group_config, num_chunks=3)
507
508 assert tracker._pending_event_metadata
509
@@ -387,7 +512,7 @@ def test_reset_cache_clears_side_table():
512 assert not tracker._pending_event_metadata
513
514
390 -def test_tiering_rejects_self_describing_kv_events():
515 +def test_tiering_accepts_self_describing_kv_events():
516 vllm_config = create_vllm_config(
517 block_size=4,
518 max_num_batched_tokens=16,
@@ -423,5 +548,9 @@ def test_tiering_rejects_self_describing_kv_events():
548 ],
549 )
550
426 - with pytest.raises(ValueError, match="TieringOffloadingSpec"):
427 - TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
551 + spec = TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
552 + tracker = OffloadingEventsTracker(spec.kv_events_config)
553 +
554 + assert spec.kv_events_config.enable_kv_cache_events
555 + assert spec.kv_events_config.self_describing_kv_events
556 + assert tracker.self_describing_enabled
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+259 -27
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from types import SimpleNamespace
4 -from unittest.mock import MagicMock
4 +from unittest.mock import MagicMock, call
5
6 import pytest
7 import torch
@@ -11,6 +11,7 @@ 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
14 +from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
15 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
16 OffloadingConnectorStats,
17 _ConnectorMetricName,
@@ -19,6 +20,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
20 OffloadingConnectorScheduler,
21 RequestOffloadState,
22 )
23 +from vllm.v1.core.kv_cache_utils import BlockHash
24 from vllm.v1.kv_cache_interface import (
25 FullAttentionSpec,
26 KVCacheGroupSpec,
@@ -26,6 +28,7 @@ from vllm.v1.kv_cache_interface import (
28 )
29 from vllm.v1.kv_offload.base import (
30 LookupResult,
31 + OffloadingEvent,
32 OffloadingManager,
33 OffloadPolicy,
34 ReqContext,
@@ -143,6 +146,159 @@ def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):
146 assert reduced[f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_sum"] > 0
147
148
149 +def test_max_offload_tokens_zero_does_not_record_pending_lookups(request_runner):
150 + runner = request_runner(
151 + block_size=4,
152 + num_gpu_blocks=10,
153 + async_scheduling=False,
154 + )
155 + runner.manager.lookup.return_value = LookupResult.RETRY
156 + runner.manager.take_events.return_value = []
157 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (
158 + generate_store_output(keys)
159 + )
160 +
161 + runner.new_request(
162 + token_ids=[1] * 12,
163 + kv_transfer_params={"max_offload_tokens": 0},
164 + )
165 + runner.run(decoded_tokens=[])
166 +
167 + tracker = runner.connector_scheduler._events_tracker
168 + assert runner.manager.lookup.call_count == 3
169 + assert not tracker._pending_event_metadata
170 + assert list(runner.connector_scheduler.take_events()) == []
171 +
172 + runner.manager.lookup.return_value = LookupResult.MISS
173 + runner.run(decoded_tokens=[EOS_TOKEN_ID])
174 +
175 + assert not tracker._pending_event_metadata
176 + assert list(runner.connector_scheduler.take_events()) == []
177 +
178 +
179 +def test_abort_before_hit_uses_placeholder_then_later_hit_heals_removal(
180 + request_runner,
181 +):
182 + runner = request_runner(
183 + block_size=4,
184 + num_gpu_blocks=10,
185 + async_scheduling=False,
186 + blocks_per_chunk=2,
187 + )
188 + raw_events: list[OffloadingEvent] = []
189 +
190 + def take_raw_events():
191 + yield from raw_events
192 + raw_events.clear()
193 +
194 + runner.manager.lookup.return_value = LookupResult.RETRY
195 + runner.manager.take_events.side_effect = take_raw_events
196 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (
197 + generate_store_output([])
198 + )
199 +
200 + runner.new_request(token_ids=[1] * 8)
201 + runner.run(decoded_tokens=[])
202 +
203 + tracker = runner.connector_scheduler._events_tracker
204 + assert not tracker._pending_event_metadata
205 + key = runner.manager.lookup.call_args.args[0]
206 + req_id = str(runner.req_id)
207 + req_status = runner.connector_scheduler._req_status[req_id]
208 +
209 + runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED)
210 +
211 + assert not tracker._pending_event_metadata
212 +
213 + raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=False))
214 + events = list(runner.connector_scheduler.take_events())
215 + assert len(events) == 1
216 + assert isinstance(events[0], BlockStored)
217 + assert events[0].block_size == 0
218 + assert events[0].token_ids == []
219 +
220 + runner.manager.lookup.return_value = LookupResult.HIT
221 + group_config = runner.connector_scheduler.config.kv_group_configs[0]
222 + assert (
223 + runner.connector_scheduler._maximal_prefix_lookup(
224 + [key],
225 + req_status.req_context,
226 + req_status.req,
227 + group_config,
228 + 0,
229 + )
230 + == 1
231 + )
232 + assert key in tracker._pending_event_metadata
233 +
234 + raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=True))
235 + [event] = runner.connector_scheduler.take_events()
236 + assert isinstance(event, BlockRemoved)
237 + assert event.medium == MEDIUM_CPU
238 + assert len(event.block_hashes) == 2
239 + assert key not in tracker._pending_event_metadata
240 +
241 +
242 +@pytest.mark.parametrize("blocks_per_chunk", [1, 2])
243 +def test_promotion_hit_precedes_stored_event_translation(
244 + request_runner,
245 + blocks_per_chunk: int,
246 +):
247 + runner = request_runner(
248 + block_size=4,
249 + num_gpu_blocks=10,
250 + async_scheduling=False,
251 + blocks_per_chunk=blocks_per_chunk,
252 + )
253 + token_ids = [1] * 4 * blocks_per_chunk
254 +
255 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (
256 + generate_store_output(keys)
257 + )
258 + runner.new_request(token_ids=token_ids)
259 + runner.run(
260 + decoded_tokens=[EOS_TOKEN_ID],
261 + expected_stored=tuple(range(blocks_per_chunk)),
262 + )
263 + runner.scheduler.reset_prefix_cache()
264 + runner.connector_scheduler._events_tracker.reset()
265 +
266 + raw_events: list[OffloadingEvent] = []
267 +
268 + def lookup(key, req_context):
269 + raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=False))
270 + return LookupResult.HIT
271 +
272 + def take_raw_events():
273 + yield from raw_events
274 + raw_events.clear()
275 +
276 + runner.manager.lookup.side_effect = lookup
277 + runner.manager.take_events.side_effect = take_raw_events
278 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (
279 + generate_store_output([])
280 + )
281 + publisher = MagicMock()
282 + runner.scheduler.kv_event_publisher = publisher
283 +
284 + runner.new_request(token_ids=token_ids)
285 + runner.run(
286 + decoded_tokens=[],
287 + expected_loaded=tuple(range(blocks_per_chunk)),
288 + )
289 +
290 + events = [
291 + event
292 + for publish_call in publisher.publish.call_args_list
293 + for event in publish_call.args[0].events
294 + if isinstance(event, BlockStored) and event.medium == MEDIUM_CPU
295 + ]
296 + assert len(events) == 1
297 + assert len(events[0].block_hashes) == blocks_per_chunk
298 + assert events[0].block_size == 4
299 + assert events[0].token_ids == token_ids
300 +
301 +
302 @pytest.mark.parametrize("async_scheduling", [True, False])
303 def test_offloading_connector(request_runner, async_scheduling: bool):
304 block_size = 4
@@ -241,7 +397,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
397 runner.manager.prepare_store.side_effect = lambda keys, req_context: (
398 generate_store_output([])
399 )
244 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
400 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
401 runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(0, 1, 2))
402
403 # single block lookup with a hit in a middle block
@@ -249,7 +405,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
405 runner.manager.prepare_store.side_effect = lambda keys, req_context: (
406 generate_store_output([])
407 )
252 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
408 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
409 runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(3, 4, 5))
410
411
@@ -307,7 +463,7 @@ def test_request_preemption(request_runner, async_scheduling: bool):
463
464 # request should now return from preemption
465 # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
310 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 3
466 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 3
467 runner.manager.prepare_store.side_effect = lambda keys, req_context: (
468 generate_store_output(keys)
469 )
@@ -427,7 +583,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
583 # start a request to load the first block, but don't complete
584 runner.scheduler.reset_prefix_cache()
585 runner.new_request(token_ids=[0] * tokens_per_chunk)
430 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
586 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
587 runner.run(
588 decoded_tokens=[],
589 complete_transfers=False,
@@ -439,7 +595,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
595
596 # start a new request to load the same first block
597 runner.new_request(token_ids=[0] * tokens_per_chunk)
442 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
598 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
599 runner.run(
600 decoded_tokens=[],
601 complete_transfers=False,
@@ -491,7 +647,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
647 # start a request to load the first block, but don't complete
648 runner.scheduler.reset_prefix_cache()
649 runner.new_request(token_ids=[0] * tokens_per_chunk)
494 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
650 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
651 runner.run(
652 decoded_tokens=[],
653 complete_transfers=False,
@@ -793,73 +949,144 @@ def _make_scheduler_with_lookup(
949
950 scheduler = object.__new__(OffloadingConnectorScheduler)
951 scheduler.manager = manager
952 + scheduler._events_tracker = MagicMock()
953 return scheduler
954
955
956 _EMPTY_REQ_CTX = ReqContext(req_id="")
957 +_LOOKUP_REQ = MagicMock()
958 +_LOOKUP_REQ.request_id = "req"
959 +_LOOKUP_GROUP_CONFIG = MagicMock()
960 +
961 +
962 +def _maximal_lookup(sched, keys, start_chunk_idx: int = 0):
963 + return sched._maximal_prefix_lookup(
964 + keys,
965 + _EMPTY_REQ_CTX,
966 + _LOOKUP_REQ,
967 + _LOOKUP_GROUP_CONFIG,
968 + start_chunk_idx,
969 + )
970
971
972 class TestMaximalPrefixLookup:
973 def test_all_hit(self):
974 sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
805 - assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
975 + assert _maximal_lookup(sched, to_keys([1, 2])) == 2
976 +
977 + def test_records_absolute_chunk_indices(self):
978 + keys = to_keys([1, 2])
979 + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
980 +
981 + assert _maximal_lookup(sched, keys, start_chunk_idx=3) == 2
982 + assert sched._events_tracker.record_lookup.call_args_list == [
983 + call(
984 + _LOOKUP_REQ,
985 + _LOOKUP_GROUP_CONFIG,
986 + 3,
987 + keys[0],
988 + ),
989 + call(
990 + _LOOKUP_REQ,
991 + _LOOKUP_GROUP_CONFIG,
992 + 4,
993 + keys[1],
994 + ),
995 + ]
996
997 def test_all_miss(self):
998 sched = _make_scheduler_with_lookup({})
809 - assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
999 + assert _maximal_lookup(sched, to_keys([1, 2])) == 0
1000 + sched._events_tracker.record_lookup.assert_not_called()
1001
1002 def test_partial_prefix(self):
1003 sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
813 - assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
1004 + assert _maximal_lookup(sched, to_keys([1, 2, 3])) == 2
1005
1006 def test_miss_then_hit(self):
1007 sched = _make_scheduler_with_lookup({2: LookupResult.HIT})
817 - assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
1008 + assert _maximal_lookup(sched, to_keys([1, 2])) == 0
1009
1010 def test_single_hit(self):
1011 sched = _make_scheduler_with_lookup({1: LookupResult.HIT})
821 - assert sched._maximal_prefix_lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
1012 + assert _maximal_lookup(sched, to_keys([1])) == 1
1013
1014 def test_empty(self):
1015 sched = _make_scheduler_with_lookup({})
825 - assert sched._maximal_prefix_lookup([], _EMPTY_REQ_CTX) == 0
1016 + assert _maximal_lookup(sched, []) == 0
1017 +
1018 + @pytest.mark.parametrize(
1019 + "pending_result",
1020 + [LookupResult.RETRY, LookupResult.HIT_PENDING],
1021 + )
1022 + def test_pending_result_is_not_recorded(
1023 + self,
1024 + pending_result: LookupResult,
1025 + ):
1026 + sched = _make_scheduler_with_lookup({1: pending_result})
1027 +
1028 + assert _maximal_lookup(sched, to_keys([1])) is None
1029 + sched._events_tracker.record_lookup.assert_not_called()
1030
1031 def test_retry_defers(self):
1032 + keys = to_keys([1, 2])
1033 sched = _make_scheduler_with_lookup(
1034 {1: LookupResult.RETRY, 2: LookupResult.HIT}
1035 )
831 - assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
1036 + assert _maximal_lookup(sched, keys) is None
1037 assert sched.manager.lookup.call_count == 2
1038 + sched._events_tracker.record_lookup.assert_called_once_with(
1039 + _LOOKUP_REQ,
1040 + _LOOKUP_GROUP_CONFIG,
1041 + 1,
1042 + keys[1],
1043 + )
1044
1045 def test_retry_after_hit_defers(self):
1046 + keys = to_keys([1, 2])
1047 sched = _make_scheduler_with_lookup(
1048 {1: LookupResult.HIT, 2: LookupResult.RETRY}
1049 )
838 - assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
1050 + assert _maximal_lookup(sched, keys) is None
1051 + sched._events_tracker.record_lookup.assert_called_once_with(
1052 + _LOOKUP_REQ,
1053 + _LOOKUP_GROUP_CONFIG,
1054 + 0,
1055 + keys[0],
1056 + )
1057
1058 def test_hit_pending_defers(self):
1059 + keys = to_keys([1, 2])
1060 sched = _make_scheduler_with_lookup(
1061 {1: LookupResult.HIT_PENDING, 2: LookupResult.HIT}
1062 )
844 - assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
1063 + assert _maximal_lookup(sched, keys) is None
1064 assert sched.manager.lookup.call_count == 2
1065 + sched._events_tracker.record_lookup.assert_called_once_with(
1066 + _LOOKUP_REQ,
1067 + _LOOKUP_GROUP_CONFIG,
1068 + 1,
1069 + keys[1],
1070 + )
1071
1072 def test_hit_pending_does_not_stop_scan(self):
1073 """HIT_PENDING defers but does not break — scan continues until miss."""
1074 sched = _make_scheduler_with_lookup(
1075 {1: LookupResult.HIT_PENDING, 2: LookupResult.MISS, 3: LookupResult.HIT}
1076 )
852 - assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None
1077 + assert _maximal_lookup(sched, to_keys([1, 2, 3])) is None
1078 assert sched.manager.lookup.call_count == 2
1079 + sched._events_tracker.record_lookup.assert_not_called()
1080
1081 def test_retry_stops_at_miss(self):
1082 """RETRY is treated as hit for iteration, but miss stops the scan."""
1083 sched = _make_scheduler_with_lookup(
1084 {1: LookupResult.RETRY, 2: LookupResult.MISS, 3: LookupResult.HIT}
1085 )
860 - assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None
1086 + assert _maximal_lookup(sched, to_keys([1, 2, 3])) is None
1087 # lookup should have been called for blocks 1 and 2 (stops at miss)
1088 assert sched.manager.lookup.call_count == 2
1089 + sched._events_tracker.record_lookup.assert_not_called()
1090
1091
1092 class TestSlidingWindowLookup:
@@ -1011,7 +1238,7 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling
1238
1239 # New request with 2 offloaded chunks; first matches what's in CPU.
1240 runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
1014 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
1241 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
1242 runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1243 generate_store_output(keys)
1244 )
@@ -1042,7 +1269,7 @@ def test_loads_do_not_populate_fence_index(request_runner):
1269 async_scheduling=False,
1270 )
1271 runner.new_request(token_ids=[0] * 12)
1045 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
1272 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
1273 runner.run(decoded_tokens=[], complete_transfers=False)
1274 assert runner.connector_scheduler._block_id_to_pending_jobs == {}
1275
@@ -1088,7 +1315,7 @@ def test_fence_at_update_state_after_alloc(request_runner):
1315
1316 runner.scheduler.reset_prefix_cache()
1317 runner.new_request(token_ids=[0] * 4)
1091 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
1318 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
1319 runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1320 generate_store_output([])
1321 )
@@ -1139,7 +1366,7 @@ def test_fence_at_build_store_jobs(request_runner):
1366
1367 runner.scheduler.reset_prefix_cache()
1368 runner.new_request(token_ids=[1] * 4)
1142 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 0
1369 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 0
1370 runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1371 generate_store_output([])
1372 )
@@ -1365,7 +1592,7 @@ def test_reset_cache(request_runner, async_scheduling: bool):
1592 # Leave the load in-flight so that reset_cache must flush it.
1593 runner.scheduler.reset_prefix_cache()
1594 runner.new_request(token_ids=[0] * tokens_per_chunk)
1368 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
1595 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
1596 runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1597 generate_store_output([])
1598 )
@@ -1553,9 +1780,7 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner
1780 # preemption batch's ModelRunnerOutput is consumed by update_from_output().
1781 free_block_queue.num_free_blocks = num_free_blocks_empty
1782 assert runner.scheduler.reset_prefix_cache()
1556 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: len(
1557 - key
1558 - )
1783 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: len(keys)
1784
1785 readmit_output = runner.scheduler.schedule()
1786
@@ -1669,7 +1894,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
1894 runner.scheduler.reset_prefix_cache()
1895 runner.new_request(token_ids=[0] * num_tokens + [1])
1896 runner.manager.lookup.return_value = LookupResult.HIT
1672 - runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2
1897 + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 2
1898 runner.run(
1899 decoded_tokens=[EOS_TOKEN_ID],
1900 # Group 0: full prefix lookup hits 2 offloaded chunks
@@ -1840,6 +2065,13 @@ class TestEagle:
2065 req.request_id = "test-req"
2066 req.num_tokens = num_tokens
2067 req.kv_transfer_params = None
2068 + num_hash_blocks = max(
2069 + len(hashes) * scheduler.config.kv_group_configs[idx].hashes_per_chunk
2070 + for idx, hashes in enumerate(offload_keys_per_group)
2071 + )
2072 + req.block_hashes = [BlockHash(str(i).encode()) for i in range(num_hash_blocks)]
2073 + req.all_token_ids = list(range(num_tokens))
2074 + req.lora_request = None
2075
2076 state = RequestOffloadState(
2077 config=scheduler.config,
vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
+29 -11
@@ -61,9 +61,9 @@ def get_offloading_event_group_spec(
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."""
64 + """BlockStored payload snapshot for one OffloadKey, captured while the
65 + Request is available and kept until the matching eviction event. ``medium``
66 + is forwarded from the OffloadingEvent."""
67
68 # The chunk's constituent block hashes; the last one is the OffloadKey.
69 block_hashes: tuple[BlockHash, ...]
@@ -81,10 +81,11 @@ class _OffloadEventMetadata:
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`.
84 + The scheduler calls :meth:`record_store` from ``_build_store_jobs`` and
85 + :meth:`record_lookup` for ready primary-tier hits while the ``Request`` is
86 + available. Deferred and missing lookups add no state. Under the connector's
87 + supported success-only transfer model, entries follow primary allocations
88 + until CPU removal translation or :meth:`reset`.
89 """
90
91 def __init__(self, config: OffloadingKVEventsConfig):
@@ -93,8 +94,7 @@ class OffloadingEventsTracker:
94 config.enable_kv_cache_events and config.self_describing_kv_events
95 )
96
96 - # OffloadKey -> payload snapshot, kept until the eviction event so
97 - # BlockRemoved can fan out. Bounded: one entry per offloaded chunk.
97 + # OffloadKey -> payload snapshot, kept until CPU removal or reset.
98 self._pending_event_metadata: dict[OffloadKey, _OffloadEventMetadata] = {}
99
100 def record_store(
@@ -116,6 +116,23 @@ class OffloadingEventsTracker:
116 meta = self._build_event_metadata(req, group_config, chunk_idx)
117 self._pending_event_metadata[offload_key] = meta
118
119 + def record_lookup(
120 + self,
121 + req: Request,
122 + group_config: "GroupOffloadConfig",
123 + chunk_idx: int,
124 + offload_key: OffloadKey,
125 + ) -> None:
126 + """Snapshot metadata for a ready primary-tier lookup hit."""
127 + if not self.self_describing_enabled:
128 + return
129 + if group_config.sliding_window_size_in_chunks is not None:
130 + return
131 + if offload_key not in self._pending_event_metadata:
132 + self._pending_event_metadata[offload_key] = self._build_event_metadata(
133 + req, group_config, chunk_idx
134 + )
135 +
136 def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]:
137 """Translate raw OffloadingEvents into self-describing KV events.
138
@@ -165,7 +182,7 @@ class OffloadingEventsTracker:
182 assert len(chunk_hashes) == hbf
183
184 if group_config.sliding_window_size_in_chunks is not None:
168 - # record_store filters these out before calling this helper.
185 + # The recording methods filter these out before calling this helper.
186 raise AssertionError("self-describing events only support full attention")
187
188 parent_block_hash: BlockHash | None
@@ -232,7 +249,8 @@ class OffloadingEventsTracker:
249 "OffloadingEventsTracker: no event metadata for "
250 "offload key during BlockStored emission; emitting a "
251 "placeholder payload. Expected for non-full-attention "
235 - "groups; otherwise indicates a missing populate path."
252 + "groups and promotions not observed as a primary-tier "
253 + "hit before translation."
254 )
255 yield self._placeholder_stored(key, event.medium, locality)
256 continue
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+20 -4
@@ -463,15 +463,27 @@ class OffloadingConnectorScheduler:
463 del self._req_status[req_id]
464
465 def _maximal_prefix_lookup(
466 - self, keys: Iterable[OffloadKey], req_context: ReqContext
466 + self,
467 + keys: Iterable[OffloadKey],
468 + req_context: ReqContext,
469 + req: Request,
470 + group_config: GroupOffloadConfig,
471 + start_chunk_idx: int,
472 ) -> int | None:
473 """Return the number of consecutive offloaded chunks from the start,
474 or None if the backend deferred a lookup."""
475 hit_count = 0
476 defer_lookup = False
472 - for key in keys:
473 - match self.manager.lookup(key, req_context):
477 + for local_idx, key in enumerate(keys):
478 + result = self.manager.lookup(key, req_context)
479 + match result:
480 case LookupResult.HIT:
481 + self._events_tracker.record_lookup(
482 + req,
483 + group_config,
484 + start_chunk_idx + local_idx,
485 + key,
486 + )
487 hit_count += 1
488 case LookupResult.HIT_PENDING:
489 defer_lookup = True
@@ -616,7 +628,11 @@ class OffloadingConnectorScheduler:
628 num_hit_chunks: int | None
629 if sliding_window_size_in_chunks is None:
630 num_hit_chunks = self._maximal_prefix_lookup(
619 - offload_keys, req_status.req_context
631 + offload_keys,
632 + req_status.req_context,
633 + req_status.req,
634 + group_config,
635 + start_chunk_idx,
636 )
637 else:
638 required_window = sliding_window_size_in_chunks
vllm/v1/kv_offload/tiering/spec.py
-8
@@ -136,14 +136,6 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
136 super().__init__(config)
137 # Redeclare for mypy: parent sets this but `--follow-imports skip` hides it
138 self._manager: OffloadingManager | None = None
139 - if self.kv_events_config.self_describing_kv_events:
140 - raise ValueError(
141 - "self_describing_kv_events is not supported by "
142 - "TieringOffloadingSpec. Tier promotions can emit primary-tier "
143 - "store events that do not correspond to GPU store jobs, so the "
144 - "current self-describing side table cannot describe them "
145 - "correctly."
146 - )
139
140 # Parse secondary tier configurations
141 self.secondary_tier_configs = self.extra_config.get("secondary_tiers", [])