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

[Bugfix][KV Offload] Track cache recency once per request (#51787)

Signed-off-by: mindungil <alswnsrlf12@naver.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

MINJUN GIL committed Sep 16, 2026 at 16:47 UTC f12fe10c5c377ccd401b66c39df3ffb394fa067c
10 files changed +918 -171
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+18 -53
@@ -258,6 +258,9 @@ def test_partial_tail_store_uses_attention_and_recurrent_cow_sources():
258 jobs = scheduler._build_partial_tail_store_jobs(output)
259
260 assert len(jobs) == 1
261 + offered_keys = scheduler.manager.prepare_store.call_args.args[0]
262 + req_context = scheduler._req_status["req"].req_context
263 + assert all(req_context.get_offload_key_position(key) == 28 for key in offered_keys)
264 [job_id] = jobs
265 src_spec = jobs[job_id].src_spec
266 assert isinstance(src_spec, GPULoadStoreSpec)
@@ -380,6 +383,13 @@ def test_normal_store_excludes_align_mode_mamba_sources():
383 req_status.group_states[0].block_ids[:] = [11]
384 req_status.group_states[1].block_ids[:] = [99]
385 req_status.update_offload_keys()
386 + for group_config, group_state in zip(
387 + scheduler.config.kv_group_configs, req_status.group_states
388 + ):
389 + assert (
390 + req_status.req_context.get_offload_key_position(group_state.offload_keys[0])
391 + == group_config.tokens_per_chunk
392 + )
393 scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (
394 generate_store_output(keys)
395 )
@@ -921,7 +931,6 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
931 runner.run(decoded_tokens=[0] * (tokens_per_chunk + 1))
932
933 # 1 more block (+ token for kicking off offloading)
924 - # now check touch was called with all 6 blocks
934 runner.manager.prepare_store.side_effect = lambda keys, req_context: (
935 generate_store_output(keys)
936 )
@@ -929,9 +938,6 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
938 decoded_tokens=[0] * (tokens_per_chunk + 1),
939 expected_stored=(15, 16, 17),
940 )
932 - runner.manager.touch.assert_called()
933 - block_hashes1 = list(runner.manager.touch.call_args.args[0])
934 - assert len(block_hashes1) == 6
941
942 # terminate request
943 runner.run(decoded_tokens=[EOS_TOKEN_ID])
@@ -939,13 +945,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
945 # create a new request differing only on the last token
946 runner.new_request(token_ids=[0] * (tokens_per_chunk * 6 - 1) + [1])
947 runner.run(decoded_tokens=[0])
942 - runner.manager.touch.assert_called()
943 - block_hashes2 = list(runner.manager.touch.call_args.args[0])
944 - assert len(block_hashes2) == 6
945 -
946 - # verify hashes are the same, except for the last block
947 - assert block_hashes1[:5] == block_hashes2[:5]
948 - assert block_hashes1[5] != block_hashes2[5]
948 + runner.manager.touch.assert_not_called()
949
950 # terminate request
951 runner.run(
@@ -1353,14 +1353,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
1353 generate_store_output(keys)
1354 )
1355 runner.run(decoded_tokens=[0])
1356 - # _touch called from get_num_new_matched_tokens (2 groups) and
1357 - # _get_reqs_to_store (2 groups) → 4 touch calls total.
1358 - touch_calls = runner.manager.touch.call_args_list
1359 - assert len(touch_calls) == 4
1360 - assert len(touch_calls[0].args[0]) == 3
1361 - assert len(touch_calls[1].args[0]) == 3
1362 - assert len(touch_calls[2].args[0]) == 3
1363 - assert len(touch_calls[3].args[0]) == 3
1356 + runner.manager.touch.assert_not_called()
1357
1358 # store 3 more block
1359 runner.manager.prepare_store.side_effect = lambda keys, req_context: (
@@ -1371,9 +1364,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
1364 expected_stored=(0, 1, 2, 3, 4, 5),
1365 )
1366
1374 - # touch called from _get_reqs_to_store * 3 blocks, once for each group
1375 - touch_calls = runner.manager.touch.call_args_list
1376 - assert len(touch_calls) == 6
1367 + runner.manager.touch.assert_not_called()
1368
1369 # EOS lands in the last slot of the 7th block (offset 6). No forward pass
1370 # writes that slot, so the finishing step declines the block.
@@ -1392,13 +1383,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
1383 expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)),
1384 )
1385
1395 - # 2 touch calls from get_num_new_matched_tokens (2 groups)
1396 - touch_calls = runner.manager.touch.call_args_list
1397 - assert len(touch_calls) == 2
1398 - # full attention group touched all 3 blocks
1399 - assert len(touch_calls[0].args[0]) == 3
1400 - # sliding window group touched just the last 2 blocks
1401 - assert len(touch_calls[1].args[0]) == 2
1386 + runner.manager.touch.assert_not_called()
1387
1388 # 3 blocks are hit on GPU [0, 1, 2]
1389 # 1 block loaded [3,]
@@ -1464,15 +1449,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
1449 generate_store_output(keys)
1450 )
1451 runner.run(decoded_tokens=[0])
1467 - # _touch called from get_num_new_matched_tokens (2 groups) and
1468 - # _get_reqs_to_store (2 groups) → 4 touch calls total.
1469 - # Group 0 has 2 offload keys, group 1 has 1.
1470 - touch_calls = runner.manager.touch.call_args_list
1471 - assert len(touch_calls) == 4
1472 - assert len(touch_calls[0].args[0]) == 2
1473 - assert len(touch_calls[1].args[0]) == 1
1474 - assert len(touch_calls[2].args[0]) == 2
1475 - assert len(touch_calls[3].args[0]) == 1
1452 + runner.manager.touch.assert_not_called()
1453
1454 # Get to 31 tokens
1455 # No further blocks offloaded
@@ -1482,11 +1459,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
1459 # Group 0 blocks: [0, 1], ending_token_offset = 24
1460 # Group 1 blocks: [0, 1], ending_token_offset = 32
1461 runner.run(decoded_tokens=[0])
1485 - # _get_reqs_to_store touch: only group 1 has a new block to store
1486 - touch_calls = runner.manager.touch.call_args_list
1487 - assert len(touch_calls) == 2
1488 - assert len(touch_calls[0].args[0]) == 2
1489 - assert len(touch_calls[1].args[0]) == 2
1462 + runner.manager.touch.assert_not_called()
1463
1464 # Get to 35 tokens
1465 # No further blocks offloaded
@@ -1496,11 +1469,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
1469 # Group 0 blocks: [0, 1, 2], ending_token_offset = 36
1470 # Group 1 blocks: [0, 1], ending_token_offset = 32
1471 runner.run(decoded_tokens=[0])
1499 - # _get_reqs_to_store touch: only group 0 has a new block to store
1500 - touch_calls = runner.manager.touch.call_args_list
1501 - assert len(touch_calls) == 2
1502 - assert len(touch_calls[0].args[0]) == 3
1503 - assert len(touch_calls[1].args[0]) == 2
1472 + runner.manager.touch.assert_not_called()
1473
1474 # Get to 47 tokens
1475 # No further blocks offloaded
@@ -1510,11 +1479,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
1479 # Group 0 blocks: [0, 1, 2, 3], ending_token_offset = 4
1480 # Group 1 blocks: [0, 1, 2], ending_token_offset = 48
1481 runner.run(decoded_tokens=[0])
1513 - # _get_reqs_to_store touch: both groups have a new block, each with 1 key
1514 - touch_calls = runner.manager.touch.call_args_list
1515 - assert len(touch_calls) == 2
1516 - assert len(touch_calls[0].args[0]) == 4
1517 - assert len(touch_calls[1].args[0]) == 3
1482 + runner.manager.touch.assert_not_called()
1483
1484 runner.run(decoded_tokens=[0], expected_stored=((0, 3), (1, 2)))
1485
tests/v1/kv_offload/cpu/test_manager.py
+428 -8
@@ -23,6 +23,7 @@ from vllm.v1.kv_offload.cpu.common import (
23 )
24 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
25 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
26 +from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
27
28
29 def make_req_context(
@@ -430,31 +431,34 @@ def test_cpu_manager():
431 assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is LookupResult.MISS
432
433 # prepare load [2, 3]
433 - prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
434 + load_ctx = make_req_context("load-2-3")
435 + prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), load_ctx)
436 verify_load_output(prepare_load_output, [1, 2])
437
438 # prepare store with no space ([2, 3] is being loaded)
439 assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None
440
441 # complete load [2, 3]. Load changes the eviction list, making 2, 3 recent.
440 - cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
442 + cpu_manager.complete_load(to_keys([2, 3]), load_ctx)
443 + cpu_manager.on_request_finished(load_ctx)
444
442 - # prepare store [6, 7, 8] -> evicts [4, 5, 2] (oldest)
445 + # prepare store [6, 7, 8] -> evicts [4, 5, 3] (oldest). Within
446 + # the accessed prefix [2, 3], the tail is less valuable than the head.
447 prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
448 verify_store_output(
449 prepare_store_output,
450 ExpectedPrepareStoreOutput(
451 keys_to_store=[6, 7, 8],
448 - store_chunk_ids=[1, 0, 3],
449 - evicted_keys=[4, 5, 2],
452 + store_chunk_ids=[2, 0, 3],
453 + evicted_keys=[4, 5, 3],
454 ),
455 )
456
457 # complete store [6, 7, 8]
458 cpu_manager.complete_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
459
456 - # touch [3, 6, 7] (move to end of LRU order)
457 - cpu_manager.touch(to_keys([3, 6, 7]), _EMPTY_REQ_CTX)
460 + # touch [2, 6, 7] (move to end of LRU order)
461 + cpu_manager.touch(to_keys([2, 6, 7]), _EMPTY_REQ_CTX)
462
463 # prepare store [7, 9] -> evicts [8] (oldest following previous touch)
464 prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX)
@@ -477,7 +481,7 @@ def test_cpu_manager():
481 verify_events(
482 cpu_manager.take_events(),
483 expected_stores=({3, 4, 5}, {6, 7, 8}),
480 - expected_evictions=({4, 5, 2}, {8}),
484 + expected_evictions=({4, 5, 3}, {8}),
485 )
486
487
@@ -518,6 +522,65 @@ def test_prepare_load_preserves_key_order():
522 manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant
523
524
525 +def test_lru_batch_eviction_failure_is_atomic():
526 + manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
527 + policy = manager._policy
528 + assert isinstance(policy, LRUCachePolicy)
529 + keys = to_keys([1, 2, 3, 4])
530 + assert manager.prepare_store(keys, _EMPTY_REQ_CTX) is not None
531 + manager.complete_store(keys, _EMPTY_REQ_CTX)
532 +
533 + protected = set(keys[1:])
534 + assert policy.evict(2, protected) is None
535 + assert all(manager.lookup(key, _EMPTY_REQ_CTX) is LookupResult.HIT for key in keys)
536 +
537 + evicted = policy.evict(1, protected)
538 + assert evicted is not None
539 + assert [key for key, _ in evicted] == [keys[0]]
540 +
541 +
542 +def test_lru_repeated_pin_cycles_compact_lazy_heap_entries():
543 + manager = make_cpu_manager(num_chunks=1, cache_policy="lru")
544 + policy = manager._policy
545 + assert isinstance(policy, LRUCachePolicy)
546 + key = to_key(1)
547 + assert manager.prepare_store([key], _EMPTY_REQ_CTX) is not None
548 + manager.complete_store([key], _EMPTY_REQ_CTX)
549 +
550 + for _ in range(128):
551 + manager.prepare_load([key], _EMPTY_REQ_CTX)
552 + manager.complete_load([key], _EMPTY_REQ_CTX)
553 +
554 + assert len(policy._heap) < 64
555 + evicted = policy.evict(1, set())
556 + assert evicted is not None
557 + assert [evicted_key for evicted_key, _ in evicted] == [key]
558 +
559 +
560 +def test_reset_discards_stale_request_access_classification():
561 + manager = make_cpu_manager(num_chunks=1, cache_policy="arc")
562 + policy = manager._policy
563 + assert isinstance(policy, ARCCachePolicy)
564 + key = to_key(1)
565 + resumed_ctx = make_req_context("resumed")
566 +
567 + assert manager.prepare_store([key], resumed_ctx) is not None
568 + manager.complete_store([key], resumed_ctx)
569 + manager.reset_cache()
570 +
571 + replacement_ctx = make_req_context("replacement")
572 + assert manager.prepare_store([key], replacement_ctx) is not None
573 + manager.complete_store([key], replacement_ctx)
574 + manager.on_request_finished(replacement_ctx)
575 + assert key in policy.t1
576 +
577 + manager.prepare_load([key], resumed_ctx)
578 + manager.complete_load([key], resumed_ctx)
579 + manager.on_request_finished(resumed_ctx)
580 +
581 + assert key in policy.t2
582 +
583 +
584 class TestARCPolicy:
585 """Unit tests for CPUOffloadingManager with ARC eviction policy."""
586
@@ -1042,6 +1105,48 @@ def test_filter_reused_manager_oversized_offer_makes_progress():
1105 assert stored_keys == set(keys)
1106
1107
1108 +def test_store_threshold_does_not_hide_ready_reuse():
1109 + manager = make_cpu_manager(
1110 + num_chunks=2,
1111 + cache_policy="lru",
1112 + store_threshold=2,
1113 + max_tracker_size=2,
1114 + )
1115 + head, tail, noise_1, noise_2, replacement = to_keys(list(range(5)))
1116 +
1117 + seed_ctx = make_req_context("seed")
1118 + output = manager.prepare_store([head, tail], seed_ctx)
1119 + assert output is not None
1120 + assert not output.keys_to_store
1121 + output = manager.prepare_store([head, tail], seed_ctx)
1122 + assert output is not None
1123 + assert output.keys_to_store == [head, tail]
1124 + manager.complete_store([head, tail], seed_ctx)
1125 + manager.on_request_finished(seed_ctx)
1126 +
1127 + # Age both resident keys out of the admission tracker without storing the
1128 + # noise keys. Reusing the tail recreates its counter at one, below the
1129 + # threshold, but must still refresh its cache recency.
1130 + output = manager.prepare_store([noise_1, noise_2], make_req_context("noise"))
1131 + assert output is not None
1132 + assert not output.keys_to_store
1133 + reuse_ctx = make_req_context("reuse")
1134 + skipped_before_reuse = manager.stores_skipped_in_current_batch
1135 + output = manager.prepare_store([tail], reuse_ctx)
1136 + assert output is not None
1137 + assert not output.keys_to_store
1138 + assert manager.stores_skipped_in_current_batch == skipped_before_reuse
1139 + manager.on_request_finished(reuse_ctx)
1140 +
1141 + replacement_ctx = make_req_context("replacement")
1142 + output = manager.prepare_store([replacement], replacement_ctx)
1143 + assert output is not None
1144 + assert not output.keys_to_store
1145 + output = manager.prepare_store([replacement], replacement_ctx)
1146 + assert output is not None
1147 + assert output.evicted_keys == [head]
1148 +
1149 +
1150 def test_evictable_cache_chunk_count():
1151 """
1152 Verifies _num_evictable_cache_chunks is maintained correctly through the
@@ -1156,3 +1261,318 @@ def test_touch_forwards_req_context_to_policy(monkeypatch):
1261 assert len(received) == 1
1262 assert received[0][0] == keys
1263 assert received[0][1] is ctx
1264 +
1265 +
1266 +@pytest.mark.parametrize("finish_before_completion", [False, True])
1267 +def test_request_finish_orders_lru_prefix_independent_of_store_completion(
1268 + finish_before_completion: bool,
1269 +):
1270 + """A transfer's completion order must not make the prefix head LRU."""
1271 + manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
1272 + prefix = to_keys([1, 2, 3, 4])
1273 + ctx = make_req_context("prefix")
1274 +
1275 + output = manager.prepare_store(prefix, ctx)
1276 + assert output is not None
1277 + if finish_before_completion:
1278 + manager.on_request_finished(ctx)
1279 + manager.complete_store(set(prefix), ctx)
1280 + if not finish_before_completion:
1281 + manager.on_request_finished(ctx)
1282 +
1283 + output = manager.prepare_store(to_keys([5]), make_req_context("evict"))
1284 + assert output is not None
1285 + assert output.evicted_keys == to_keys([4])
1286 +
1287 +
1288 +@pytest.mark.parametrize("finish_before_completion", [False, True])
1289 +def test_request_finish_orders_lru_prefix_independent_of_load_completion(
1290 + finish_before_completion: bool,
1291 +):
1292 + manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
1293 + prefix = to_keys([1, 2, 3, 4])
1294 +
1295 + seed_ctx = make_req_context("seed")
1296 + output = manager.prepare_store(prefix, seed_ctx)
1297 + assert output is not None
1298 + manager.complete_store(prefix, seed_ctx)
1299 +
1300 + reuse_ctx = make_req_context("reuse")
1301 + manager.prepare_load(prefix, reuse_ctx)
1302 + if finish_before_completion:
1303 + manager.on_request_finished(reuse_ctx)
1304 + manager.complete_load(set(prefix), reuse_ctx)
1305 + if not finish_before_completion:
1306 + manager.on_request_finished(reuse_ctx)
1307 +
1308 + output = manager.prepare_store(to_keys([5]), make_req_context("evict"))
1309 + assert output is not None
1310 + assert output.evicted_keys == to_keys([4])
1311 +
1312 +
1313 +def test_late_old_completion_does_not_override_newer_request_recency():
1314 + manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
1315 + old_keys = to_keys([1, 2])
1316 + new_keys = to_keys([3, 4])
1317 +
1318 + old_ctx = make_req_context("old")
1319 + assert manager.prepare_store(old_keys, old_ctx) is not None
1320 + manager.on_request_finished(old_ctx)
1321 +
1322 + new_ctx = make_req_context("new")
1323 + assert manager.prepare_store(new_keys, new_ctx) is not None
1324 + manager.complete_store(new_keys, new_ctx)
1325 + manager.on_request_finished(new_ctx)
1326 +
1327 + # Physical completion is late, but the logical access is still older.
1328 + manager.complete_store(old_keys, old_ctx)
1329 +
1330 + output = manager.prepare_store(to_keys([5]), make_req_context("evict"))
1331 + assert output is not None
1332 + assert output.evicted_keys == [old_keys[-1]]
1333 +
1334 +
1335 +def test_arc_counts_reuse_once_and_keeps_insertions_in_t1():
1336 + manager = make_cpu_manager(num_chunks=4, cache_policy="arc")
1337 + policy = manager._policy
1338 + assert isinstance(policy, ARCCachePolicy)
1339 + prefix = to_keys([1, 2, 3, 4])
1340 +
1341 + seed_ctx = make_req_context("seed")
1342 + assert manager.prepare_store(prefix, seed_ctx) is not None
1343 + manager.complete_store(prefix, seed_ctx)
1344 + # Repeated store offers from the same request must not become ARC hits.
1345 + assert manager.prepare_store(prefix, seed_ctx) is not None
1346 + # Nor should an internal read used to cascade the new chunks to a tier.
1347 + manager.prepare_load(prefix, seed_ctx)
1348 + manager.complete_load(prefix, seed_ctx)
1349 + manager.on_request_finished(seed_ctx)
1350 +
1351 + assert list(policy.t1) == list(reversed(prefix))
1352 + assert not policy.t2
1353 +
1354 + reuse_ctx = make_req_context("reuse")
1355 + manager.prepare_load(prefix, reuse_ctx)
1356 + manager.complete_load(prefix, reuse_ctx)
1357 + # A duplicate internal pin, as used by tiering cascades, is deduplicated.
1358 + manager.prepare_load(prefix, reuse_ctx)
1359 + manager.complete_load(prefix, reuse_ctx)
1360 + manager.on_request_finished(reuse_ctx)
1361 +
1362 + assert not policy.t1
1363 + assert list(policy.t2) == list(reversed(prefix))
1364 +
1365 +
1366 +def test_arc_reuse_wins_if_same_request_later_reinserts_key():
1367 + manager = make_cpu_manager(num_chunks=1, cache_policy="arc")
1368 + policy = manager._policy
1369 + assert isinstance(policy, ARCCachePolicy)
1370 + key_1, key_2 = to_keys([1, 2])
1371 +
1372 + seed_ctx = make_req_context("seed")
1373 + assert manager.prepare_store([key_1], seed_ctx) is not None
1374 + manager.complete_store([key_1], seed_ctx)
1375 + manager.on_request_finished(seed_ctx)
1376 +
1377 + reuse_ctx = make_req_context("reuse-and-reinsert")
1378 + manager.prepare_load([key_1], reuse_ctx)
1379 + manager.complete_load([key_1], reuse_ctx)
1380 +
1381 + other_ctx = make_req_context("evict-reused-key")
1382 + assert manager.prepare_store([key_2], other_ctx) is not None
1383 + manager.complete_store([key_2], other_ctx)
1384 + manager.on_request_finished(other_ctx)
1385 +
1386 + assert manager.prepare_store([key_1], reuse_ctx) is not None
1387 + manager.complete_store([key_1], reuse_ctx)
1388 + manager.on_request_finished(reuse_ctx)
1389 +
1390 + assert key_1 not in policy.t1
1391 + assert list(policy.t2) == [key_1]
1392 +
1393 +
1394 +def test_request_finish_forwards_grouped_deduplicated_accesses(monkeypatch):
1395 + manager = make_cpu_manager(num_chunks=5)
1396 + existing = make_offload_key(b"existing", 0)
1397 + group_0_new = make_offload_key(b"group-0-new", 0)
1398 + group_1_head = make_offload_key(b"group-1-head", 1)
1399 + group_1_tail = make_offload_key(b"group-1-tail", 1)
1400 +
1401 + seed_ctx = make_req_context("seed")
1402 + assert manager.prepare_store([existing], seed_ctx) is not None
1403 + manager.complete_store([existing], seed_ctx)
1404 +
1405 + ctx = make_req_context("grouped")
1406 + # Group 1 is observed first, but policy finalization follows group index.
1407 + offered = [group_1_head, existing, group_0_new]
1408 + assert manager.prepare_store(offered, ctx) is not None
1409 + # Later offers extend a group's prefix and may repeat earlier keys.
1410 + assert (
1411 + manager.prepare_store([group_0_new, group_1_head, group_1_tail], ctx)
1412 + is not None
1413 + )
1414 +
1415 + received = []
1416 +
1417 + def on_request_finished(key_groups, insertion_only_keys, reused_keys, req_context):
1418 + received.append((key_groups, insertion_only_keys, reused_keys, req_context))
1419 +
1420 + monkeypatch.setattr(manager._policy, "on_request_finished", on_request_finished)
1421 + manager.on_request_finished(ctx)
1422 + manager.on_request_finished(ctx)
1423 +
1424 + assert received == [
1425 + (
1426 + ((existing, group_0_new), (group_1_head, group_1_tail)),
1427 + {group_0_new, group_1_head, group_1_tail},
1428 + {existing},
1429 + ctx,
1430 + )
1431 + ]
1432 +
1433 +
1434 +def test_failed_store_records_access_once_and_uses_store_miss_hook(monkeypatch):
1435 + manager = make_cpu_manager(num_chunks=1)
1436 + resident, candidate = to_keys([1, 2])
1437 +
1438 + seed_ctx = make_req_context("seed")
1439 + assert manager.prepare_store([resident], seed_ctx) is not None
1440 + manager.complete_store([resident], seed_ctx)
1441 +
1442 + pin_ctx = make_req_context("pin")
1443 + manager.prepare_load([resident], pin_ctx)
1444 +
1445 + record_calls = 0
1446 + original_record = manager._record_request_cache_access
1447 +
1448 + def record_access(*args, **kwargs):
1449 + nonlocal record_calls
1450 + record_calls += 1
1451 + return original_record(*args, **kwargs)
1452 +
1453 + store_miss_calls = []
1454 +
1455 + def on_store_miss(keys, req_context):
1456 + store_miss_calls.append((list(keys), req_context))
1457 +
1458 + finalized = []
1459 +
1460 + def on_request_finished(key_groups, insertion_only_keys, reused_keys, req_context):
1461 + finalized.append((key_groups, insertion_only_keys, reused_keys, req_context))
1462 +
1463 + monkeypatch.setattr(manager, "_record_request_cache_access", record_access)
1464 + monkeypatch.setattr(manager._policy, "on_store_miss", on_store_miss)
1465 + monkeypatch.setattr(manager._policy, "on_request_finished", on_request_finished)
1466 +
1467 + ctx = make_req_context("failed-store")
1468 + assert manager.prepare_store([candidate, resident], ctx) is None
1469 + assert record_calls == 1
1470 + assert store_miss_calls == [([candidate], ctx)]
1471 +
1472 + manager.on_request_finished(ctx)
1473 + assert finalized == [
1474 + (((candidate, resident),), set(), {resident}, ctx),
1475 + ]
1476 +
1477 + manager.complete_load([resident], pin_ctx)
1478 +
1479 +
1480 +def test_arc_ghost_hit_adapts_once_per_request_before_insertion():
1481 + manager = make_cpu_manager(num_chunks=2, cache_policy="arc")
1482 + policy = manager._policy
1483 + assert isinstance(policy, ARCCachePolicy)
1484 + keys = to_keys([1, 2, 3])
1485 +
1486 + seed_ctx = make_req_context("seed")
1487 + assert manager.prepare_store(keys[:2], seed_ctx) is not None
1488 + manager.complete_store(keys[:2], seed_ctx)
1489 + manager.on_request_finished(seed_ctx)
1490 +
1491 + evict_ctx = make_req_context("create-ghost")
1492 + output = manager.prepare_store([keys[2]], evict_ctx)
1493 + assert output is not None
1494 + [ghost_key] = output.evicted_keys
1495 + manager.complete_store([keys[2]], evict_ctx)
1496 + manager.on_request_finished(evict_ctx)
1497 + assert ghost_key in policy.b1
1498 +
1499 + resident_keys = [key for key in keys if key != ghost_key]
1500 + pin_ctx = make_req_context("pin-residents")
1501 + manager.prepare_load(resident_keys, pin_ctx)
1502 +
1503 + ghost_ctx = make_req_context("ghost-hit")
1504 + target_before = policy.target_t1_size
1505 + assert manager.prepare_store([ghost_key], ghost_ctx) is None
1506 + target_after_first_offer = policy.target_t1_size
1507 + assert target_after_first_offer > target_before
1508 + assert manager.prepare_store([ghost_key], ghost_ctx) is None
1509 + assert policy.target_t1_size == target_after_first_offer
1510 +
1511 + manager.complete_load(resident_keys, pin_ctx)
1512 + output = manager.prepare_store([ghost_key], ghost_ctx)
1513 + assert output is not None
1514 + assert policy.target_t1_size == target_after_first_offer
1515 + manager.complete_store([ghost_key], ghost_ctx)
1516 + manager.on_request_finished(ghost_ctx)
1517 + assert ghost_key in policy.t1
1518 + assert ghost_key not in policy.t2
1519 +
1520 +
1521 +def test_arc_pending_chunk_is_not_a_frequency_hit():
1522 + manager = make_cpu_manager(num_chunks=2, cache_policy="arc")
1523 + policy = manager._policy
1524 + assert isinstance(policy, ARCCachePolicy)
1525 + key = to_key(1)
1526 +
1527 + writer_ctx = make_req_context("writer")
1528 + assert manager.prepare_store([key], writer_ctx) is not None
1529 +
1530 + observer_ctx = make_req_context("pending-observer")
1531 + output = manager.prepare_store([key], observer_ctx)
1532 + assert output is not None and not output.keys_to_store
1533 + manager.on_request_finished(observer_ctx)
1534 +
1535 + chunk = policy.get(key)
1536 + assert chunk is not None and not chunk.is_ready
1537 + assert key in policy.t1
1538 + assert key not in policy.t2
1539 +
1540 +
1541 +def test_request_key_positions_override_store_observation_order():
1542 + manager = make_cpu_manager(num_chunks=2, cache_policy="lru")
1543 + head, tail = to_keys([1, 2])
1544 + ctx = make_req_context("out-of-order-store")
1545 + ctx.set_offload_key_position(head, 16)
1546 + ctx.set_offload_key_position(tail, 32)
1547 +
1548 + # A backfill or partial-tail path may offer a later key first.
1549 + assert manager.prepare_store([tail], ctx) is not None
1550 + assert manager.prepare_store([head], ctx) is not None
1551 + manager.on_request_finished(ctx)
1552 + manager.complete_store({head, tail}, ctx)
1553 +
1554 + output = manager.prepare_store([to_key(3)], make_req_context("evict-tail"))
1555 + assert output is not None
1556 + assert output.evicted_keys == [tail]
1557 +
1558 +
1559 +def test_request_key_positions_order_tails_across_kv_groups():
1560 + manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
1561 + group_0_head = make_offload_key(b"group-0-head", 0)
1562 + group_0_tail = make_offload_key(b"group-0-tail", 0)
1563 + group_1_head = make_offload_key(b"group-1-head", 1)
1564 + group_1_tail = make_offload_key(b"group-1-tail", 1)
1565 + keys = [group_0_head, group_0_tail, group_1_head, group_1_tail]
1566 + ctx = make_req_context("hybrid-groups")
1567 + for key in (group_0_head, group_1_head):
1568 + ctx.set_offload_key_position(key, 16)
1569 + for key in (group_0_tail, group_1_tail):
1570 + ctx.set_offload_key_position(key, 32)
1571 +
1572 + assert manager.prepare_store(keys, ctx) is not None
1573 + manager.complete_store(keys, ctx)
1574 + manager.on_request_finished(ctx)
1575 +
1576 + output = manager.prepare_store(to_keys([10, 11]), make_req_context("evict"))
1577 + assert output is not None
1578 + assert set(output.evicted_keys) == {group_0_tail, group_1_tail}
tests/v1/kv_offload/tiering/test_tiering_offloading.py
+56 -5
@@ -40,6 +40,7 @@ from vllm.v1.kv_offload.base import (
40 TierMatcher,
41 make_offload_key,
42 )
43 +from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
44 from vllm.v1.kv_offload.tiering.base import (
45 JobResult,
46 SecondaryTierManager,
@@ -753,14 +754,64 @@ class TestTieringOffloadingManager:
754 # Touch chunks
755 self.manager.touch(chunks, _CTX)
756
756 - # Verify touch was called on primary tier (check LRU order)
757 - primary_keys = list(self.primary_tier._policy.evictable_chunks.keys())
758 - assert primary_keys[-3:] == list(reversed(chunks))
757 + # Verify touch was called on the primary tier (head is most recent).
758 + policy = self.primary_tier._policy
759 + assert isinstance(policy, LRUCachePolicy)
760 + ranks = policy._ranks
761 + assert ranks[chunks[2]] < ranks[chunks[1]] < ranks[chunks[0]]
762
763 # Verify touch was propagated to all secondary tiers
764 self.secondary_tier1.touch.assert_called_once_with(chunks, _CTX)
765 self.secondary_tier2.touch.assert_called_once_with(chunks, _CTX)
766
767 + def test_request_order_survives_late_cascade_completion(self, manager_setup):
768 + """Cascade completion must not replace request order with I/O order."""
769 + chunks = to_keys(range(5))
770 + self._start_request()
771 + output = self.manager.prepare_store(chunks, _CTX)
772 + assert output is not None
773 + self.manager.complete_store(chunks, _CTX, success=True)
774 +
775 + # Finalize while both secondary-tier cascades still pin primary chunks.
776 + self.manager.on_request_finished(_CTX)
777 + self._simulate_on_schedule_end()
778 + self._simulate_on_schedule_end()
779 +
780 + evict_ctx = ReqContext(req_id="evict")
781 + self._start_request(evict_ctx)
782 + output = self.manager.prepare_store(to_keys([5]), evict_ctx)
783 + assert output is not None
784 + assert output.evicted_keys == [chunks[-1]]
785 +
786 + def test_late_old_store_submission_does_not_override_newer_request_order(
787 + self, manager_setup
788 + ):
789 + """Cascade submission after finish is not a new primary access."""
790 + old_blocks = to_keys(range(2))
791 + old_ctx = ReqContext(req_id="old")
792 + self._start_request(old_ctx)
793 + assert self.manager.prepare_store(old_blocks, old_ctx) is not None
794 + self.manager.on_request_finished(old_ctx)
795 +
796 + new_blocks = to_keys(range(2, 5))
797 + new_ctx = ReqContext(req_id="new")
798 + self._start_request(new_ctx)
799 + assert self.manager.prepare_store(new_blocks, new_ctx) is not None
800 + self.manager.complete_store(new_blocks, new_ctx, success=True)
801 + self.manager.on_request_finished(new_ctx)
802 +
803 + # The old GPU->primary write lands after the newer request finished.
804 + # Its cascade pins must not turn it into the more recent request.
805 + self.manager.complete_store(old_blocks, old_ctx, success=True)
806 + self._simulate_on_schedule_end()
807 + self._simulate_on_schedule_end()
808 +
809 + evict_ctx = ReqContext(req_id="evict")
810 + self._start_request(evict_ctx)
811 + output = self.manager.prepare_store(to_keys([5]), evict_ctx)
812 + assert output is not None
813 + assert output.evicted_keys == [old_blocks[-1]]
814 +
815 def test_failed_store_no_cascade(self, manager_setup):
816 """Test that failed GPU→primary store doesn't cascade."""
817 chunks = to_keys(range(3))
@@ -873,10 +924,10 @@ class TestTieringOffloadingManager:
924 job_metadata = self.secondary_tier1.submit_store.call_args.args[0]
925 assert job_metadata.req_context is ctx
926
876 - def test_on_request_finished_delays_secondary_until_store_submitted(
927 + def test_on_request_finished_delays_only_secondary_until_store_submitted(
928 self, manager_setup
929 ):
879 - """Manager hook is eager; secondary hooks wait for cascade submission."""
930 + """Primary order commits immediately; secondary cleanup waits."""
931 chunks = to_keys(range(2))
932 ctx = ReqContext(req_id="req_delayed_secondary")
933 calls: list[tuple[str, str]] = []
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+27 -43
@@ -427,8 +427,11 @@ class RequestOffloadState:
427 None,
428 group_config.hashes_per_chunk,
429 ):
430 - group_state.offload_keys.append(
431 - make_offload_key(req_block_hash, group_config.group_idx)
430 + key = make_offload_key(req_block_hash, group_config.group_idx)
431 + group_state.offload_keys.append(key)
432 + self.req_context.set_offload_key_position(
433 + key,
434 + len(group_state.offload_keys) * group_config.tokens_per_chunk,
435 )
436
437 def update_block_id_groups(
@@ -742,37 +745,6 @@ class OffloadingConnectorScheduler:
745 )
746 return None if defer_lookup or pending_in_window else consecutive_hits
747
745 - def _touch(self, req_status: RequestOffloadState):
746 - for group_config, group_state in zip(
747 - self.config.kv_group_configs, req_status.group_states
748 - ):
749 - if group_config.sliding_window_size_in_chunks is None:
750 - self.manager.touch(group_state.offload_keys, req_status.req_context)
751 - else:
752 - # Keep only chunks needed to hit the original request, plus
753 - # decoded chunks.
754 - chunks_to_skip = max(
755 - 0,
756 - group_state.num_hit_chunks
757 - - group_config.sliding_window_size_in_chunks,
758 - )
759 - self.manager.touch(
760 - group_state.offload_keys[chunks_to_skip:],
761 - req_status.req_context,
762 - )
763 - if req_status.partial_tail_boundary is not None:
764 - self.manager.touch(
765 - tuple(
766 - self._make_boundary_key(
767 - req_status.req,
768 - group.group_idx,
769 - req_status.partial_tail_boundary,
770 - )
771 - for group in self.config.kv_group_configs
772 - ),
773 - req_status.req_context,
774 - )
775 -
748 def _lookup_complete_chunks(
749 self,
750 req_status: RequestOffloadState,
@@ -977,10 +949,16 @@ class OffloadingConnectorScheduler:
949 return num_hit_tokens
950
951 def _make_boundary_key(
980 - self, request: Request, group_idx: int, boundary_tokens: int
952 + self,
953 + request: Request,
954 + group_idx: int,
955 + boundary_tokens: int,
956 + req_context: ReqContext,
957 ) -> OffloadKey:
958 hash_idx = boundary_tokens // self.config.tokens_per_hash - 1
983 - return make_offload_key(request.block_hashes[hash_idx], group_idx)
959 + key = make_offload_key(request.block_hashes[hash_idx], group_idx)
960 + req_context.set_offload_key_position(key, boundary_tokens)
961 + return key
962
963 def _lookup(
964 self,
@@ -1015,7 +993,10 @@ class OffloadingConnectorScheduler:
993 boundary_keys = []
994 for group_config in self.config.kv_group_configs:
995 key = self._make_boundary_key(
1018 - req_status.req, group_config.group_idx, boundary
996 + req_status.req,
997 + group_config.group_idx,
998 + boundary,
999 + req_status.req_context,
1000 )
1001 boundary_keys.append(key)
1002 result = self.manager.lookup(key, req_status.req_context)
@@ -1108,8 +1089,6 @@ class OffloadingConnectorScheduler:
1089 self._maybe_observe_lookup_async_delay(req_status)
1090 req_status.update_num_hit_chunks(num_computed_tokens + (num_hit_tokens or 0))
1091
1111 - self._touch(req_status)
1112 -
1092 return num_hit_tokens, bool(num_hit_tokens)
1093
1094 def update_state_after_alloc(
@@ -1183,7 +1162,10 @@ class OffloadingConnectorScheduler:
1162 if partial_tail_boundary is not None:
1163 keys_to_load.append(
1164 self._make_boundary_key(
1186 - request, group_config.group_idx, partial_tail_boundary
1165 + request,
1166 + group_config.group_idx,
1167 + partial_tail_boundary,
1168 + req_status.req_context,
1169 )
1170 )
1171
@@ -1307,7 +1289,9 @@ class OffloadingConnectorScheduler:
1289 ):
1290 continue
1291
1310 - key = self._make_boundary_key(req, group_idx, boundary)
1292 + key = self._make_boundary_key(
1293 + req, group_idx, boundary, req_status.req_context
1294 + )
1295 store_output = self.manager.prepare_store([key], req_status.req_context)
1296 if store_output is None:
1297 self._connector_stats.increase_counter(
@@ -1402,7 +1386,9 @@ class OffloadingConnectorScheduler:
1386 ):
1387 continue
1388 keys = [
1405 - self._make_boundary_key(req, group.group_idx, boundary)
1389 + self._make_boundary_key(
1390 + req, group.group_idx, boundary, req_status.req_context
1391 + )
1392 for group in self.config.kv_group_configs
1393 ]
1394 block_ids = [
@@ -1689,8 +1675,6 @@ class OffloadingConnectorScheduler:
1675 req_status.advance_stored_idx(num_offloadable_tokens)
1676 continue
1677
1692 - self._touch(req_status)
1693 -
1678 keys_to_store = set(store_output.keys_to_store)
1679
1680 group_sizes: list[int] = []
vllm/v1/kv_offload/base.py
+12
@@ -96,6 +96,12 @@ class ReqContext:
96 # kv_transfer_params once (in on_new_request) and read the result back
97 # on later calls for the same request.
98 _state: dict[type, Any] = field(default_factory=dict, repr=False, init=False)
99 + # End-token position for each key in this request. The scheduler records
100 + # these positions so managers can recover prefix order even when store
101 + # calls arrive out of order (for example, SWA backfills).
102 + _offload_key_positions: dict[OffloadKey, int] = field(
103 + default_factory=dict, repr=False, init=False
104 + )
105
106 def set_state(self, val: Any) -> None:
107 self._state[type(val)] = val
@@ -103,6 +109,12 @@ class ReqContext:
109 def get_state(self, cls: type[_T]) -> _T | None:
110 return self._state.get(cls)
111
112 + def set_offload_key_position(self, key: OffloadKey, end_token: int) -> None:
113 + self._offload_key_positions[key] = end_token
114 +
115 + def get_offload_key_position(self, key: OffloadKey) -> int | None:
116 + return self._offload_key_positions.get(key)
117 +
118
119 class LookupResult(Enum):
120 """Result of OffloadingManager.lookup()."""
vllm/v1/kv_offload/cpu/manager.py
+141 -5
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections import OrderedDict
4 from collections.abc import Collection, Iterable
5 +from dataclasses import dataclass, field
6
7 from typing_extensions import override
8
@@ -18,6 +19,7 @@ from vllm.v1.kv_offload.base import (
19 PrepareStoreOutput,
20 ReqContext,
21 RequestOffloadingContext,
22 + get_offload_group_idx,
23 )
24 from vllm.v1.kv_offload.cpu.common import (
25 CPULoadStoreSpec,
@@ -27,6 +29,20 @@ from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus
29 from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory
30
31
32 +@dataclass(slots=True)
33 +class _RequestCacheAccess:
34 + """Cache keys observed by one request, grouped in prefix order."""
35 +
36 + owner: object
37 + cache_generation: int
38 + key_groups: dict[int, list[OffloadKey]] = field(default_factory=dict)
39 + seen_keys: set[OffloadKey] = field(default_factory=set)
40 + inserted_keys: set[OffloadKey] = field(default_factory=set)
41 + reused_keys: set[OffloadKey] = field(default_factory=set)
42 + store_miss_keys: set[OffloadKey] = field(default_factory=set)
43 + finished: bool = False
44 +
45 +
46 class CPUOffloadingManager(OffloadingManager):
47 """
48 An OffloadingManager with a pluggable CachePolicy, resolved by name via
@@ -66,6 +82,7 @@ class CPUOffloadingManager(OffloadingManager):
82 self.max_tracker_size: int = max_tracker_size
83 self.stores_skipped_in_current_batch: int = 0
84 self.allocation_sizes_in_current_batch: list[int] = []
85 + self._cache_generation = 0
86
87 # Number of chunk references. It is ordered so can evict the LRU entry in O(1).
88 self.counts: OrderedDict[OffloadKey, int] | None = (
@@ -124,10 +141,48 @@ class CPUOffloadingManager(OffloadingManager):
141 num_unprotected -= 1
142 self.counts[key] = 1
143
144 + def _get_request_cache_access(self, req_context: ReqContext) -> _RequestCacheAccess:
145 + state = req_context.get_state(_RequestCacheAccess)
146 + if (
147 + state is None
148 + or state.owner is not self
149 + or state.cache_generation != self._cache_generation
150 + ):
151 + state = _RequestCacheAccess(
152 + owner=self, cache_generation=self._cache_generation
153 + )
154 + req_context.set_state(state)
155 + return state
156 +
157 + def _record_request_cache_access(
158 + self,
159 + keys: Iterable[OffloadKey],
160 + req_context: ReqContext,
161 + inserted_keys: Iterable[OffloadKey] = (),
162 + reused_keys: Iterable[OffloadKey] = (),
163 + ) -> None:
164 + state = self._get_request_cache_access(req_context)
165 + for key in keys:
166 + if key in state.seen_keys:
167 + continue
168 + assert not state.finished, (
169 + "New cache keys observed after request finalization"
170 + )
171 + group_idx = get_offload_group_idx(key)
172 + state.key_groups.setdefault(group_idx, []).append(key)
173 + state.seen_keys.add(key)
174 + state.inserted_keys.update(inserted_keys)
175 + # Re-reading a chunk inserted by this request is an internal transfer
176 + # (for example, a tiering cascade), not a second cache access.
177 + state.reused_keys.update(
178 + key for key in reused_keys if key not in state.inserted_keys
179 + )
180 +
181 # --- OffloadingManager interface ---
182
183 @override
184 def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
185 + self._get_request_cache_access(req_context)
186 return RequestOffloadingContext()
187
188 @override
@@ -144,6 +199,15 @@ class CPUOffloadingManager(OffloadingManager):
199 self,
200 keys: Collection[OffloadKey],
201 req_context: ReqContext,
202 + ) -> LoadStoreSpec:
203 + return self._prepare_load(keys, req_context, record_access=True)
204 +
205 + def _prepare_load(
206 + self,
207 + keys: Collection[OffloadKey],
208 + req_context: ReqContext,
209 + *,
210 + record_access: bool,
211 ) -> LoadStoreSpec:
212 chunks = []
213 for key in keys:
@@ -156,6 +220,8 @@ class CPUOffloadingManager(OffloadingManager):
220 assert self._num_evictable_cache_chunks >= 0
221 chunk.ref_cnt += 1
222 chunks.append(chunk)
223 + if record_access:
224 + self._record_request_cache_access(keys, req_context, reused_keys=keys)
225 return self._get_load_store_spec(keys, chunks)
226
227 @override
@@ -181,13 +247,52 @@ class CPUOffloadingManager(OffloadingManager):
247 keys: Collection[OffloadKey],
248 req_context: ReqContext,
249 ) -> PrepareStoreOutput | None:
250 + keys = list(keys)
251 if self.counts is not None:
185 - num_keys = len(keys)
252 self._record_accesses(keys)
187 - keys = [k for k in keys if self.counts.get(k, 0) >= self.store_threshold]
188 - self.stores_skipped_in_current_batch += num_keys - len(keys)
189 - # filter out chunks that are already stored
190 - keys_to_store = [k for k in keys if self._policy.get(k) is None]
253 + # Partition the original offer before admission filtering. Ready
254 + # resident chunks are request reuses even if their tracker entry was
255 + # aged out; the threshold applies only to new store candidates.
256 + keys_to_store: list[OffloadKey] = []
257 + ready_existing_keys: list[OffloadKey] = []
258 + for key in keys:
259 + chunk = self._policy.get(key)
260 + if chunk is None:
261 + keys_to_store.append(key)
262 + else:
263 + if chunk.is_ready:
264 + ready_existing_keys.append(key)
265 +
266 + if self.counts is not None:
267 + num_store_candidates = len(keys_to_store)
268 + keys_to_store = [
269 + key
270 + for key in keys_to_store
271 + if self.counts.get(key, 0) >= self.store_threshold
272 + ]
273 + self.stores_skipped_in_current_batch += num_store_candidates - len(
274 + keys_to_store
275 + )
276 +
277 + state = self._get_request_cache_access(req_context)
278 + new_store_misses = [
279 + key for key in keys_to_store if key not in state.store_miss_keys
280 + ]
281 + if new_store_misses:
282 + # ARC learns from B1/B2 before insert() removes the ghost entry.
283 + # Deduplication makes this one policy observation per request.
284 + self._policy.on_store_miss(new_store_misses, req_context)
285 + state.store_miss_keys.update(new_store_misses)
286 +
287 + recorded_keys = set(keys_to_store)
288 + recorded_keys.update(ready_existing_keys)
289 + # Record once before any allocation-related early return. Store
290 + # candidates are classified as insertions only after allocation succeeds.
291 + self._record_request_cache_access(
292 + (key for key in keys if key in recorded_keys),
293 + req_context,
294 + reused_keys=ready_existing_keys,
295 + )
296
297 if not keys_to_store:
298 return PrepareStoreOutput(
@@ -239,6 +344,7 @@ class CPUOffloadingManager(OffloadingManager):
344 for key, chunk in zip(keys_to_store, chunks):
345 self._policy.insert(key, chunk)
346 self._num_write_pending_chunks += len(keys_to_store)
347 + state.inserted_keys.update(keys_to_store)
348
349 # build store specs for allocated chunks
350 store_spec = self._get_load_store_spec(keys_to_store, chunks)
@@ -284,6 +390,35 @@ class CPUOffloadingManager(OffloadingManager):
390 )
391 )
392
393 + @override
394 + def on_request_finished(self, req_context: ReqContext) -> None:
395 + state = req_context.get_state(_RequestCacheAccess)
396 + if (
397 + state is None
398 + or state.owner is not self
399 + or state.cache_generation != self._cache_generation
400 + or state.finished
401 + ):
402 + return
403 + state.finished = True
404 + key_groups = []
405 + for group_idx in sorted(state.key_groups):
406 + keys = state.key_groups[group_idx]
407 + positions = {
408 + key: position
409 + for key in keys
410 + if (position := req_context.get_offload_key_position(key)) is not None
411 + }
412 + if len(positions) == len(keys):
413 + keys = sorted(keys, key=positions.__getitem__)
414 + key_groups.append(tuple(keys))
415 + self._policy.on_request_finished(
416 + tuple(key_groups),
417 + state.inserted_keys - state.reused_keys,
418 + state.reused_keys,
419 + req_context,
420 + )
421 +
422 @override
423 def reset_cache(self) -> None:
424 # Clear ALL chunks unconditionally. The scheduler's _stale_job_threshold
@@ -292,6 +427,7 @@ class CPUOffloadingManager(OffloadingManager):
427 # flushes in-flight load job IDs to the workers before any new stores
428 # can begin, preventing a cross-direction data race on reused offload chunk IDs.
429 self._policy.clear()
430 + self._cache_generation += 1
431 self._num_evictable_cache_chunks = 0
432 self._num_write_pending_chunks = 0
433
vllm/v1/kv_offload/cpu/policies/arc.py
+59 -22
@@ -1,12 +1,16 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections import OrderedDict
4 -from collections.abc import Iterable, Iterator
4 +from collections.abc import Iterable, Iterator, Sequence
5
6 from typing_extensions import override
7
8 from vllm.v1.kv_offload.base import OffloadKey, ReqContext
9 -from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus
9 +from vllm.v1.kv_offload.cpu.policies.base import (
10 + CachePolicy,
11 + ChunkStatus,
12 + order_request_keys,
13 +)
14
15
16 class ARCCachePolicy(CachePolicy):
@@ -24,12 +28,11 @@ class ARCCachePolicy(CachePolicy):
28 Searches T1 and T2 for chunk hashes and counts consecutive hits
29 until a miss or non-ready chunk is encountered.
30
27 - 2. Cache touch (touch) - Adaptive Learning:
28 - For each key (in reverse order):
29 - - If in T1: Move to T2 (promotion from recent to frequent).
30 - - If in T2: Move to MRU position (end of queue).
31 - - If in B1 ghost list: Increase target_t1_size.
32 - - If in B2 ghost list: Decrease target_t1_size.
31 + 2. Request access - Adaptive Learning:
32 + - Ready chunks reused by a request move from T1 to T2 once, or
33 + move to the MRU end of T2.
34 + - B1/B2 misses adjust target_t1_size once per request before
35 + insertion removes their ghost entries.
36
37 3. Chunk eviction (evict) - Adaptive Replacement:
38 Determines eviction source based on adaptive target:
@@ -39,7 +42,7 @@ class ARCCachePolicy(CachePolicy):
42
43 4. Chunk insertion (insert):
44 New chunks are always inserted into T1 and removed from B1/B2 if
42 - present. Chunks may later be promoted to T2 during touch operations.
45 + present. A later request reuse may promote them to T2.
46
47 Adaptive Behavior:
48 The algorithm self-tunes the recency vs. frequency trade-off:
@@ -71,6 +74,26 @@ class ARCCachePolicy(CachePolicy):
74 if self.t1.pop(key, None) is None:
75 self.t2.pop(key, None)
76
77 + def _adapt_to_ghost_hit(self, key: OffloadKey) -> bool:
78 + if key in self.b1:
79 + delta = max(1, len(self.b2) / len(self.b1))
80 + self.target_t1_size = min(self.target_t1_size + delta, self.cache_capacity)
81 + self.b1.move_to_end(key)
82 + return True
83 + if key in self.b2:
84 + delta = max(1, len(self.b1) / len(self.b2))
85 + self.target_t1_size = max(self.target_t1_size - delta, 0)
86 + self.b2.move_to_end(key)
87 + return True
88 + return False
89 +
90 + @override
91 + def on_store_miss(
92 + self, keys: Iterable[OffloadKey], req_context: ReqContext
93 + ) -> None:
94 + for key in reversed(list(keys)):
95 + self._adapt_to_ghost_hit(key)
96 +
97 @override
98 def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
99 for key in reversed(list(keys)):
@@ -86,19 +109,33 @@ class ARCCachePolicy(CachePolicy):
109 elif key in self.t2:
110 self.t2.move_to_end(key)
111
89 - elif key in self.b1:
90 - delta = max(1, len(self.b2) / len(self.b1))
91 - self.target_t1_size = min(
92 - self.target_t1_size + delta, self.cache_capacity
93 - )
94 - # move to MRU position (end) to keep it fresh in the ghost list
95 - self.b1.move_to_end(key)
96 -
97 - elif key in self.b2:
98 - delta = max(1, len(self.b1) / len(self.b2))
99 - self.target_t1_size = max(self.target_t1_size - delta, 0)
100 - # move to MRU position (end) to keep it fresh in the ghost list
101 - self.b2.move_to_end(key)
112 + else:
113 + self._adapt_to_ghost_hit(key)
114 +
115 + @override
116 + def on_request_finished(
117 + self,
118 + key_groups: Sequence[Sequence[OffloadKey]],
119 + insertion_only_keys: set[OffloadKey],
120 + reused_keys: set[OffloadKey],
121 + req_context: ReqContext,
122 + ) -> None:
123 + for key in reversed(order_request_keys(key_groups, req_context)):
124 + if key in insertion_only_keys:
125 + # A store is not a frequency hit. Preserve T1 membership
126 + # while still restoring tail-to-head recency.
127 + if key in self.t1:
128 + self.t1.move_to_end(key)
129 + elif key in self.t2:
130 + self.t2.move_to_end(key)
131 + continue
132 +
133 + # Ready chunks reused by this request count as one access.
134 + if key in reused_keys and key in self.t1:
135 + chunk = self.t1.pop(key)
136 + self.t2[key] = chunk
137 + elif key in reused_keys and key in self.t2:
138 + self.t2.move_to_end(key)
139
140 @override
141 def clear(self) -> None:
vllm/v1/kv_offload/cpu/policies/base.py
+51 -1
@@ -2,11 +2,26 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import ctypes
4 from abc import ABC, abstractmethod
5 -from collections.abc import Iterable
5 +from collections.abc import Iterable, Sequence
6
7 from vllm.v1.kv_offload.base import OffloadKey, ReqContext
8
9
10 +def order_request_keys(
11 + key_groups: Sequence[Sequence[OffloadKey]], req_context: ReqContext
12 +) -> list[OffloadKey]:
13 + """Return one head-to-tail order across all KV cache groups."""
14 + keys = [key for group in key_groups for key in group]
15 + positions = {
16 + key: position
17 + for key in keys
18 + if (position := req_context.get_offload_key_position(key)) is not None
19 + }
20 + if len(positions) == len(keys):
21 + keys.sort(key=positions.__getitem__)
22 + return keys
23 +
24 +
25 class ChunkStatus(ctypes.Structure):
26 """
27 Offloading status for a single chunk of KV data.
@@ -66,6 +81,41 @@ class CachePolicy(ABC):
81 req_context: Per-request context for the request touching these chunks.
82 """
83
84 + def on_store_miss(
85 + self, keys: Iterable[OffloadKey], req_context: ReqContext
86 + ) -> None:
87 + """Observe store misses before their cache entries are inserted.
88 +
89 + The default delegates to ``touch`` for compatibility with external
90 + policies. Policies may override this hook when a miss has distinct
91 + semantics, such as ARC adapting to a ghost-list hit.
92 + """
93 + self.touch(keys, req_context)
94 +
95 + def on_request_finished(
96 + self,
97 + key_groups: Sequence[Sequence[OffloadKey]],
98 + insertion_only_keys: set[OffloadKey],
99 + reused_keys: set[OffloadKey],
100 + req_context: ReqContext,
101 + ) -> None:
102 + """Apply one request-scoped cache access in prefix order.
103 +
104 + ``key_groups`` contains keys observed for each KV cache group in
105 + head-to-tail order. ``insertion_only_keys`` and ``reused_keys``
106 + distinguish chunks only created by this request from ready chunks it
107 + actually reused, which matters for policies such as ARC where reuse
108 + changes frequency but insertion and pending observations do not.
109 +
110 + The default forwards one touch per group, preserving compatibility
111 + for experimental out-of-tree policies while moving those touches to
112 + request finalization. Policies that distinguish insertion from reuse
113 + can override this hook and inspect the access classifications.
114 + """
115 + del insertion_only_keys, reused_keys
116 + for keys in key_groups:
117 + self.touch(keys, req_context)
118 +
119 @abstractmethod
120 def evict(
121 self, n: int, protected: set[OffloadKey]
vllm/v1/kv_offload/cpu/policies/lru.py
+113 -29
@@ -1,28 +1,74 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from collections import OrderedDict
4 -from collections.abc import Iterable
3 +import heapq
4 +from collections.abc import Iterable, Sequence
5
6 from typing_extensions import override
7
8 from vllm.v1.kv_offload.base import OffloadKey, ReqContext
9 -from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus
9 +from vllm.v1.kv_offload.cpu.policies.base import (
10 + CachePolicy,
11 + ChunkStatus,
12 + order_request_keys,
13 +)
14
15
16 class LRUCachePolicy(CachePolicy):
17 """
14 - LRU Caching policy that keeps a dedicated evictable list for fast eviction.
18 + LRU caching policy with logical recency independent of transfer pinning.
19 +
20 + Evictable chunks live in a lazy-invalidating min-heap. A chunk's recency
21 + can therefore be updated while it is pinned; when it later becomes
22 + evictable, it enters the heap with the order assigned by the request
23 + rather than its transfer-completion time.
24 +
25 A use is indicated by,
26 - First time the key is added (store).
17 - - Load job completion
18 - - touch
27 + - A request-scoped access.
28 """
29
30 def __init__(self, cache_capacity: int):
31 super().__init__(cache_capacity)
23 - # Chunks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU
24 - self.evictable_chunks: OrderedDict[OffloadKey, None] = OrderedDict()
32 self.chunks: dict[OffloadKey, ChunkStatus] = {}
33 + self._ranks: dict[OffloadKey, int] = {}
34 + self._evictable: set[OffloadKey] = set()
35 + self._heap: list[tuple[int, OffloadKey]] = []
36 + self._next_rank = 0
37 +
38 + def _assign_new_rank(self, key: OffloadKey) -> bool:
39 + if key not in self.chunks:
40 + return False
41 + self._next_rank += 1
42 + self._ranks[key] = self._next_rank
43 + return key in self._evictable
44 +
45 + def _push_evictable(self, key: OffloadKey) -> None:
46 + heapq.heappush(
47 + self._heap,
48 + (self._ranks[key], key),
49 + )
50 +
51 + def _is_current(self, entry: tuple[int, OffloadKey]) -> bool:
52 + rank, key = entry
53 + return key in self._evictable and self._ranks.get(key) == rank
54 +
55 + def _maybe_compact_heap(self) -> None:
56 + if len(self._heap) < max(64, 2 * len(self._evictable)):
57 + return
58 + self._heap = [(self._ranks[key], key) for key in self._evictable]
59 + heapq.heapify(self._heap)
60 +
61 + def _update_recency(self, keys: Iterable[OffloadKey]) -> None:
62 + updated_evictable = [key for key in keys if self._assign_new_rank(key)]
63 + # Rebuilding is linear and substantially cheaper than k heap pushes
64 + # for a long prefix. Small updates retain the incremental path.
65 + if len(updated_evictable) >= max(64, len(self._evictable) // 4):
66 + self._heap = [(self._ranks[key], key) for key in self._evictable]
67 + heapq.heapify(self._heap)
68 + else:
69 + for key in updated_evictable:
70 + self._push_evictable(key)
71 + self._maybe_compact_heap()
72
73 @override
74 def get(self, key: OffloadKey) -> ChunkStatus | None:
@@ -31,26 +77,40 @@ class LRUCachePolicy(CachePolicy):
77 @override
78 def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None:
79 self.chunks[key] = chunk
80 + self._next_rank += 1
81 + self._ranks[key] = self._next_rank
82 if chunk.ref_cnt == 0:
35 - self.evictable_chunks[key] = None
83 + self._evictable.add(key)
84 + self._push_evictable(key)
85
86 @override
87 def remove(self, key: OffloadKey) -> None:
88 del self.chunks[key]
40 - self.evictable_chunks.pop(key, None)
89 + self._ranks.pop(key, None)
90 + self._evictable.discard(key)
91
92 @override
93 def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
44 - for key in reversed(list(keys)):
45 - if key in self.evictable_chunks:
46 - self.evictable_chunks.move_to_end(key)
47 - # active chunks are untouched as they are non-evictable now. They
48 - # will eventually reach the end of evictable_chunks when they finish.
94 + self._update_recency(reversed(list(keys)))
95 +
96 + @override
97 + def on_request_finished(
98 + self,
99 + key_groups: Sequence[Sequence[OffloadKey]],
100 + insertion_only_keys: set[OffloadKey],
101 + reused_keys: set[OffloadKey],
102 + req_context: ReqContext,
103 + ) -> None:
104 + del insertion_only_keys, reused_keys
105 + self._update_recency(reversed(order_request_keys(key_groups, req_context)))
106
107 @override
108 def clear(self) -> None:
52 - self.evictable_chunks.clear()
109 self.chunks.clear()
110 + self._ranks.clear()
111 + self._evictable.clear()
112 + self._heap.clear()
113 + self._next_rank = 0
114
115 @override
116 def evict(
@@ -59,22 +119,44 @@ class LRUCachePolicy(CachePolicy):
119 if n == 0:
120 return []
121
62 - candidates: list[tuple[OffloadKey, ChunkStatus]] = []
63 - for key, _ in self.evictable_chunks.items():
122 + selected: list[tuple[tuple[int, OffloadKey], ChunkStatus]] = []
123 + selected_keys: set[OffloadKey] = set()
124 + deferred: list[tuple[int, OffloadKey]] = []
125 + while self._heap and len(selected) < n:
126 + entry = heapq.heappop(self._heap)
127 + if not self._is_current(entry):
128 + continue
129 + key = entry[1]
130 + # Re-pinning without a recency change can leave an equivalent
131 + # lazy entry behind. Select each cache key at most once.
132 + if key in selected_keys:
133 + continue
134 if key in protected:
135 + deferred.append(entry)
136 continue
66 -
137 chunk = self.chunks[key]
138 assert chunk.ref_cnt == 0
69 - candidates.append((key, chunk))
70 - if len(candidates) == n:
71 - break
72 -
73 - if len(candidates) < n:
139 + selected.append((entry, chunk))
140 + selected_keys.add(key)
141 +
142 + if len(selected) < n:
143 + for entry, _ in selected:
144 + heapq.heappush(self._heap, entry)
145 + for entry in deferred:
146 + heapq.heappush(self._heap, entry)
147 return None
75 - for key, _ in candidates:
76 - del self.evictable_chunks[key]
148 +
149 + for entry in deferred:
150 + heapq.heappush(self._heap, entry)
151 +
152 + candidates: list[tuple[OffloadKey, ChunkStatus]] = []
153 + for entry, chunk in selected:
154 + key = entry[1]
155 + self._evictable.remove(key)
156 del self.chunks[key]
157 + del self._ranks[key]
158 + candidates.append((key, chunk))
159 + self._maybe_compact_heap()
160 return candidates
161
162 @override
@@ -82,9 +164,11 @@ class LRUCachePolicy(CachePolicy):
164 # chunks can become evictable when,
165 # store completes - i.e. ref_cnt -1 -> 0 # not in evictable list
166 # all loads complete - i.e ref_cnt 1 -> 0 # not in evictable list
85 - self.evictable_chunks[key] = None
167 + self._evictable.add(key)
168 + self._push_evictable(key)
169 + self._maybe_compact_heap()
170
171 @override
172 def mark_non_evictable(self, key: OffloadKey) -> None:
89 - # key must have been in the evictable list.
90 - del self.evictable_chunks[key]
173 + # key must have been evictable before it was pinned.
174 + self._evictable.remove(key)
vllm/v1/kv_offload/tiering/manager.py
+13 -5
@@ -110,13 +110,22 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager):
110 # read/write is for CPU<->secondary transfers,
111 # load/store is for CPU<->GPU transfers.
112 # These aliases avoid calling prepare_load inside a store path.
113 - self.prepare_read = self.prepare_load
113 self.complete_read = self.complete_load
114 self.prepare_write = self.prepare_store
115 self.complete_write = self.complete_store
116
117 self._kv_memoryview = mmap_region.create_kv_memoryview()
118
119 + def prepare_read(
120 + self, keys: Collection[OffloadKey], req_context: ReqContext
121 + ) -> LoadStoreSpec:
122 + """Pin chunks for a CPU-to-secondary transfer.
123 +
124 + Cascade reads are implementation details of tiering, not additional
125 + request accesses, so they must not alter request-scoped recency.
126 + """
127 + return self._prepare_load(keys, req_context, record_access=False)
128 +
129 def get_kv_memoryview(self) -> memoryview:
130 """Return the memoryview over the primary tier's KV cache buffer.
131
@@ -781,11 +790,10 @@ class TieringOffloadingManager(OffloadingManager):
790 req_id: str,
791 exclude_tier_idx: int | None = None,
792 ) -> None:
784 - """Finalize secondary tiers once no more store cascades can be submitted.
793 + """Finalize secondary tiers once no more cascades can be submitted.
794
786 - Finalization means forwarding on_request_finished() to secondary tiers.
787 - It is delayed until pending GPU->primary stores finish, since their
788 - complete_store() callbacks may still submit primary->secondary stores.
795 + Their finalization is delayed until pending GPU->primary stores
796 + finish, since those callbacks may still submit secondary stores.
797 """
798 state = self._req_state[req_id]
799 if not state.is_finished: