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

[KV Offloading] Replace `bool|None` lookup return with LookupResult enum (#46363)

Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

Ronen Schaffer committed Jun 24, 2026 at 18:06 UTC bb61177e49dd781645b39a99bb2abdb9d37552cd

14 files changed +323 -172

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

+125 -43

@@ -21,6 +21,7 @@ from vllm.v1.kv_cache_interface import (

21 SlidingWindowSpec,

22 )

23 from vllm.v1.kv_offload.base import (

24 + LookupResult,

25 OffloadingManager,

26 OffloadPolicy,

27 ReqContext,

@@ -484,7 +485,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo

485

486 # full 3 blocks hit [0, 1, 2]

487 runner.new_request(token_ids=[0] * (block_size * 3 + 1))

487 - runner.manager.lookup.return_value = True

488 + runner.manager.lookup.return_value = LookupResult.HIT

489 runner.run(

490 decoded_tokens=[EOS_TOKEN_ID],

491 # Group 0 (full attn): prefix lookup hits 3 → loads blocks 0,1,2

@@ -504,7 +505,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo

505 # 3 blocks are hit on GPU [0, 1, 2]

506 # 1 block loaded [3,]

507 runner.new_request(token_ids=[0] * (block_size * 4 + 1))

507 - runner.manager.lookup.return_value = True

508 + runner.manager.lookup.return_value = LookupResult.HIT

509 runner.run(

510 decoded_tokens=[EOS_TOKEN_ID],

511 # Group 0 (full attn): prefix lookup hits 3 → loads blocks 0,1,2

@@ -632,7 +633,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool

633 # 48 tokens (3 block) from the second group

634 # Total 48 tokens can be loaded

635 runner.new_request(token_ids=[0] * 48)

635 - runner.manager.lookup.return_value = True

636 + runner.manager.lookup.return_value = LookupResult.HIT

637 runner.manager.prepare_store.side_effect = lambda keys, req_context: (

638 generate_store_output([])

639 )

@@ -648,7 +649,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool

649 # extra tokens [0, 36] (blocks [4, 5, 6]) from the first group

650 # extra tokens [0, 32] (block [3, 4]) from the second group

651 runner.new_request(token_ids=[0] * (48 + 37))

651 - runner.manager.lookup.return_value = True

652 + runner.manager.lookup.return_value = LookupResult.HIT

653 runner.manager.prepare_store.side_effect = lambda keys, req_context: (

654 generate_store_output([])

655 )

@@ -665,12 +666,12 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool

666

667

668 def _make_scheduler_with_lookup(

668 - lookup_results: dict[int, bool | None],

669 + lookup_results: dict[int, LookupResult],

670 ) -> OffloadingConnectorScheduler:

671 """Create an OffloadingConnectorScheduler with a mocked manager.lookup."""

672 manager = MagicMock(spec=OffloadingManager)

673 manager.lookup.side_effect = lambda key, req_context: lookup_results.get(

673 - int(get_offload_block_hash(key).decode()), False

674 + int(get_offload_block_hash(key).decode()), LookupResult.MISS

675 )

676

677 scheduler = object.__new__(OffloadingConnectorScheduler)

@@ -683,7 +684,7 @@ _EMPTY_REQ_CTX = ReqContext(req_id="")

684

685 class TestMaximalPrefixLookup:

686 def test_all_hit(self):

686 - sched = _make_scheduler_with_lookup({1: True, 2: True})

687 + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})

688 assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2

689

690 def test_all_miss(self):

@@ -691,32 +692,54 @@ class TestMaximalPrefixLookup:

692 assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0

693

694 def test_partial_prefix(self):

694 - sched = _make_scheduler_with_lookup({1: True, 2: True})

695 + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})

696 assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2

697

698 def test_miss_then_hit(self):

698 - sched = _make_scheduler_with_lookup({2: True})

699 + sched = _make_scheduler_with_lookup({2: LookupResult.HIT})

700 assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0

701

702 def test_single_hit(self):

702 - sched = _make_scheduler_with_lookup({1: True})

703 + sched = _make_scheduler_with_lookup({1: LookupResult.HIT})

704 assert sched._maximal_prefix_lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1

705

706 def test_empty(self):

707 sched = _make_scheduler_with_lookup({})

708 assert sched._maximal_prefix_lookup([], _EMPTY_REQ_CTX) == 0

709

709 - def test_none_defers(self):

710 - sched = _make_scheduler_with_lookup({1: None, 2: True})

710 + def test_retry_defers(self):

711 + sched = _make_scheduler_with_lookup(

712 + {1: LookupResult.RETRY, 2: LookupResult.HIT}

713 + )

714 + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None

715 + assert sched.manager.lookup.call_count == 2

716 +

717 + def test_retry_after_hit_defers(self):

718 + sched = _make_scheduler_with_lookup(

719 + {1: LookupResult.HIT, 2: LookupResult.RETRY}

720 + )

721 assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None

722

713 - def test_none_after_hit_defers(self):

714 - sched = _make_scheduler_with_lookup({1: True, 2: None})

723 + def test_hit_pending_defers(self):

724 + sched = _make_scheduler_with_lookup(

725 + {1: LookupResult.HIT_PENDING, 2: LookupResult.HIT}

726 + )

727 assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None

728 + assert sched.manager.lookup.call_count == 2

729 +

730 + def test_hit_pending_does_not_stop_scan(self):

731 + """HIT_PENDING defers but does not break — scan continues until miss."""

732 + sched = _make_scheduler_with_lookup(

733 + {1: LookupResult.HIT_PENDING, 2: LookupResult.MISS, 3: LookupResult.HIT}

734 + )

735 + assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None

736 + assert sched.manager.lookup.call_count == 2

737

717 - def test_none_stops_at_miss(self):

718 - """None is treated as hit for iteration, but miss stops the scan."""

719 - sched = _make_scheduler_with_lookup({1: None, 2: False, 3: True})

738 + def test_retry_stops_at_miss(self):

739 + """RETRY is treated as hit for iteration, but miss stops the scan."""

740 + sched = _make_scheduler_with_lookup(

741 + {1: LookupResult.RETRY, 2: LookupResult.MISS, 3: LookupResult.HIT}

742 + )

743 assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None

744 # lookup should have been called for blocks 1 and 2 (stops at miss)

745 assert sched.manager.lookup.call_count == 2

@@ -724,7 +747,7 @@ class TestMaximalPrefixLookup:

747

748 class TestSlidingWindowLookup:

749 def test_all_hit_exact_window(self):

727 - sched = _make_scheduler_with_lookup({1: True, 2: True})

750 + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})

751 assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2

752

753 def test_all_miss(self):

@@ -732,25 +755,27 @@ class TestSlidingWindowLookup:

755 assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 0

756

757 def test_window_at_end(self):

735 - sched = _make_scheduler_with_lookup({2: True, 3: True})

758 + sched = _make_scheduler_with_lookup({2: LookupResult.HIT, 3: LookupResult.HIT})

759 assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 2, _EMPTY_REQ_CTX) == 3

760

761 def test_window_in_middle(self):

739 - sched = _make_scheduler_with_lookup({2: True, 3: True})

762 + sched = _make_scheduler_with_lookup({2: LookupResult.HIT, 3: LookupResult.HIT})

763 assert (

764 sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 3

765 )

766

767 def test_no_full_window_falls_back_to_prefix(self):

745 - sched = _make_scheduler_with_lookup({1: True, 2: True})

768 + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})

769 assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 3, _EMPTY_REQ_CTX) == 2

770

771 def test_single_block_window(self):

749 - sched = _make_scheduler_with_lookup({2: True, 3: True})

772 + sched = _make_scheduler_with_lookup({2: LookupResult.HIT, 3: LookupResult.HIT})

773 assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 3

774

775 def test_gap_resets_consecutive(self):

753 - sched = _make_scheduler_with_lookup({2: True, 3: True, 4: True})

776 + sched = _make_scheduler_with_lookup(

777 + {2: LookupResult.HIT, 3: LookupResult.HIT, 4: LookupResult.HIT}

778 + )

779 # [1, 2, 3, 0, 4] — gap at 0 resets, window of 2 found at [2,3]

780 assert (

781 sched._sliding_window_lookup(to_keys([1, 2, 3, 0, 4]), 2, _EMPTY_REQ_CTX)

@@ -758,7 +783,14 @@ class TestSlidingWindowLookup:

783 )

784

785 def test_window_prefers_rightmost(self):

761 - sched = _make_scheduler_with_lookup({1: True, 2: True, 4: True, 5: True})

786 + sched = _make_scheduler_with_lookup(

787 + {

788 + 1: LookupResult.HIT,

789 + 2: LookupResult.HIT,

790 + 4: LookupResult.HIT,

791 + 5: LookupResult.HIT,

792 + }

793 + )

794 # two valid windows: [1,2] at positions 0-1 and [4,5] at positions 3-4

795 # scans right-to-left, finds [4,5] first

796 assert (

@@ -767,7 +799,14 @@ class TestSlidingWindowLookup:

799 )

800

801 def test_prefix_fallback_with_gap(self):

770 - sched = _make_scheduler_with_lookup({2: True, 3: True, 4: True, 5: True})

802 + sched = _make_scheduler_with_lookup(

803 + {

804 + 2: LookupResult.HIT,

805 + 3: LookupResult.HIT,

806 + 4: LookupResult.HIT,

807 + 5: LookupResult.HIT,

808 + }

809 + )

810 # window of 4 not found contiguously (gap at 1)

811 assert (

812 sched._sliding_window_lookup(to_keys([2, 1, 3, 4, 5]), 4, _EMPTY_REQ_CTX)

@@ -778,20 +817,47 @@ class TestSlidingWindowLookup:

817 sched = _make_scheduler_with_lookup({})

818 assert sched._sliding_window_lookup([], 1, _EMPTY_REQ_CTX) == 0

819

781 - def test_none_defers(self):

782 - sched = _make_scheduler_with_lookup({1: True, 2: None})

820 + def test_retry_defers(self):

821 + sched = _make_scheduler_with_lookup(

822 + {1: LookupResult.HIT, 2: LookupResult.RETRY}

823 + )

824 assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) is None

825

785 - def test_none_with_full_window_still_defers(self):

786 - """Even if a real window is found after a None, result is deferred."""

787 - # Scan right-to-left: 4(True), 3(None) resets, 2(True), 1(True) = window

788 - # but block 3 was None so defer_lookup is set

789 - sched = _make_scheduler_with_lookup({1: True, 2: True, 3: None, 4: True})

826 + def test_retry_with_full_window_still_defers(self):

827 + """Even if a real window is found after a RETRY, result is deferred."""

828 + # Scan right-to-left: 4(HIT), 3(RETRY) resets, 2(HIT), 1(HIT) = window

829 + # but block 3 was RETRY so defer_lookup is set

830 + sched = _make_scheduler_with_lookup(

831 + {

832 + 1: LookupResult.HIT,

833 + 2: LookupResult.HIT,

834 + 3: LookupResult.RETRY,

835 + 4: LookupResult.HIT,

836 + }

837 + )

838 assert (

839 sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX)

840 is None

841 )

842

843 + def test_hit_pending_counts_as_hit(self):

844 + """HIT_PENDING counts toward the consecutive-hit streak."""

845 + sched = _make_scheduler_with_lookup(

846 + {1: LookupResult.HIT, 2: LookupResult.HIT_PENDING}

847 + )

848 + # window=2: both count as hits, but defer_lookup is set

849 + assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) is None

850 +

851 + def test_hit_pending_does_not_break_streak(self):

852 + """HIT_PENDING in the middle of a window doesn't reset the streak."""

853 + sched = _make_scheduler_with_lookup(

854 + {1: LookupResult.HIT, 2: LookupResult.HIT_PENDING, 3: LookupResult.HIT}

855 + )

856 + # window=3: right-to-left finds 3(HIT),2(HIT_PENDING),1(HIT) = 3 consecutive

857 + assert (

858 + sched._sliding_window_lookup(to_keys([1, 2, 3]), 3, _EMPTY_REQ_CTX) is None

859 + )

860 +

861

862 @pytest.mark.parametrize("async_scheduling", [True, False])

863 def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling: bool):

@@ -1485,7 +1551,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):

1551 # Verify that loads still work correctly for the stored SWA blocks.

1552 runner.scheduler.reset_prefix_cache()

1553 runner.new_request(token_ids=[0] * num_tokens + [1])

1488 - runner.manager.lookup.return_value = True

1554 + runner.manager.lookup.return_value = LookupResult.HIT

1555 runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2

1556 runner.run(

1557 decoded_tokens=[EOS_TOKEN_ID],

@@ -1701,7 +1767,9 @@ class TestEagle:

1767 kv_cache_groups=groups,

1768 )

1769 runner.manager.lookup.side_effect = lambda key, req_context: (

1704 - int(get_offload_block_hash(key).decode()) in {1, 2, 3}

1770 + LookupResult.HIT

1771 + if int(get_offload_block_hash(key).decode()) in {1, 2, 3}

1772 + else LookupResult.MISS

1773 )

1774 sched = runner.connector_scheduler

1775 req_status = self._make_req_status(

@@ -1732,7 +1800,9 @@ class TestEagle:

1800 kv_cache_groups=groups,

1801 )

1802 runner.manager.lookup.side_effect = lambda key, req_context: (

1735 - int(get_offload_block_hash(key).decode()) in {1}

1803 + LookupResult.HIT

1804 + if int(get_offload_block_hash(key).decode()) in {1}

1805 + else LookupResult.MISS

1806 )

1807 sched = runner.connector_scheduler

1808 req_status = self._make_req_status(

@@ -1762,7 +1832,7 @@ class TestEagle:

1832 async_scheduling=False,

1833 kv_cache_groups=groups,

1834 )

1765 - runner.manager.lookup.return_value = False

1835 + runner.manager.lookup.return_value = LookupResult.MISS

1836 sched = runner.connector_scheduler

1837 req_status = self._make_req_status(

1838 sched, num_tokens=8, offload_keys_per_group=[[1, 2]]

@@ -1803,7 +1873,9 @@ class TestEagle:

1873 kv_cache_groups=groups,

1874 )

1875 runner.manager.lookup.side_effect = lambda key, req_context: (

1806 - int(get_offload_block_hash(key).decode()) in {1, 2, 3, 4}

1876 + LookupResult.HIT

1877 + if int(get_offload_block_hash(key).decode()) in {1, 2, 3, 4}

1878 + else LookupResult.MISS

1879 )

1880 sched = runner.connector_scheduler

1881

@@ -1858,7 +1930,9 @@ class TestEagle:

1930 kv_cache_groups=groups,

1931 )

1932 runner.manager.lookup.side_effect = lambda key, req_context: (

1861 - int(get_offload_block_hash(key).decode()) in {1, 2}

1933 + LookupResult.HIT

1934 + if int(get_offload_block_hash(key).decode()) in {1, 2}

1935 + else LookupResult.MISS

1936 )

1937 sched = runner.connector_scheduler

1938 req_status = self._make_req_status(

@@ -1891,7 +1965,9 @@ class TestEagle:

1965 kv_cache_groups=groups,

1966 )

1967 runner.manager.lookup.side_effect = lambda key, req_context: (

1894 - int(get_offload_block_hash(key).decode()) in {1, 2, 3}

1968 + LookupResult.HIT

1969 + if int(get_offload_block_hash(key).decode()) in {1, 2, 3}

1970 + else LookupResult.MISS

1971 )

1972 sched = runner.connector_scheduler

1973 # num_tokens=13 → max_hit=13-1=12, query_max=min(12+4,12)=12

@@ -1940,7 +2016,9 @@ class TestEagle:

2016 kv_cache_groups=groups,

2017 )

2018 runner.manager.lookup.side_effect = lambda key, req_context: (

1943 - int(get_offload_block_hash(key).decode()) in {1, 2, 3}

2019 + LookupResult.HIT

2020 + if int(get_offload_block_hash(key).decode()) in {1, 2, 3}

2021 + else LookupResult.MISS

2022 )

2023 sched = runner.connector_scheduler

2024 req_status = self._make_req_status(

@@ -1995,7 +2073,9 @@ class TestEagle:

2073 # Group 0 keys [10,11,12]: only 10 hits.

2074 # Group 1 keys [1,2,3]: all hit.

2075 runner.manager.lookup.side_effect = lambda key, req_context: (

1998 - int(get_offload_block_hash(key).decode()) in {10, 1, 2, 3}

2076 + LookupResult.HIT

2077 + if int(get_offload_block_hash(key).decode()) in {10, 1, 2, 3}

2078 + else LookupResult.MISS

2079 )

2080 sched = runner.connector_scheduler

2081 req_status = self._make_req_status(

@@ -2046,7 +2126,9 @@ class TestEagle:

2126 kv_cache_groups=groups,

2127 )

2128 runner.manager.lookup.side_effect = lambda key, req_context: (

2049 - int(get_offload_block_hash(key).decode()) in {1, 2, 3}

2129 + LookupResult.HIT

2130 + if int(get_offload_block_hash(key).decode()) in {1, 2, 3}

2131 + else LookupResult.MISS

2132 )

2133 sched = runner.connector_scheduler

2134 req_status = self._make_req_status(

@@ -2273,7 +2355,7 @@ class TestEagle:

2355 runner.scheduler.reset_prefix_cache()

2356

2357 runner.new_request(token_ids=[0] * offloaded_block_size * 3 + [1])

2276 - runner.manager.lookup.return_value = True

2358 + runner.manager.lookup.return_value = LookupResult.HIT

2359 runner.manager.prepare_store.side_effect = lambda keys, req_context: (

2360 generate_store_output([])

2361 )

tests/v1/kv_connector/unit/offloading_connector/utils.py

+2 -2

@@ -46,6 +46,7 @@ from vllm.v1.kv_offload.base import (

46 CanonicalKVCaches,

47 GPULoadStoreSpec,

48 LoadStoreSpec,

49 + LookupResult,

50 OffloadingManager,

51 OffloadingSpec,

52 OffloadingWorker,

@@ -129,9 +130,8 @@ class MockOffloadingSpec(OffloadingSpec):

130 super().__init__(vllm_config, kv_cache_config)

131

132 self.manager = MagicMock(spec=OffloadingManager)

132 - self.manager.lookup.return_value = 0

133 self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)

134 - self.manager.lookup.return_value = False

134 + self.manager.lookup.return_value = LookupResult.MISS

135 self.manager.on_new_request.return_value = RequestOffloadingContext()

136 self.handler = MockOffloadingWorker()

137

tests/v1/kv_offload/cpu/test_manager.py

+29 -28

@@ -8,6 +8,7 @@ import pytest

9 from vllm.v1.kv_offload.base import (

10 LoadStoreSpec,

11 + LookupResult,

12 OffloadingEvent,

13 OffloadKey,

14 PrepareStoreOutput,

@@ -160,7 +161,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):

161 manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX)

162

163 # block 2 must still be present in the cache

163 - assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True

164 + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT

165

166

167 def test_filter_reused_manager_reports_stores_skipped_counter():

@@ -242,8 +243,8 @@ def test_cpu_manager():

243 )

244

245 # lookup [1, 2] -> write in-flight, not yet ready

245 - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is None

246 - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is None

246 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING

247 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING

248

249 # no events so far

250 assert list(cpu_manager.take_events()) == []

@@ -253,9 +254,9 @@ def test_cpu_manager():

254 verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))

255

256 # lookup [1, 2]

256 - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is True

257 - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True

258 - assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False

257 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT

258 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT

259 + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS

260

261 # prepare store [2, 3, 4, 5] -> evicts [1]

262 prepare_store_output = cpu_manager.prepare_store(

@@ -280,12 +281,12 @@ def test_cpu_manager():

281 cpu_manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX)

282

283 # lookup (now that we have [2, 3, 4, 5])

283 - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False

284 - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True

285 - assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is True

286 - assert cpu_manager.lookup(to_key(4), _EMPTY_REQ_CTX) is True

287 - assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is True

288 - assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is False

284 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS

285 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT

286 + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.HIT

287 + assert cpu_manager.lookup(to_key(4), _EMPTY_REQ_CTX) is LookupResult.HIT

288 + assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is LookupResult.HIT

289 + assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is LookupResult.MISS

290

291 # prepare load [2, 3]

292 prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)

@@ -329,8 +330,8 @@ def test_cpu_manager():

330 cpu_manager.complete_store(to_keys([7, 9]), _EMPTY_REQ_CTX, success=False)

331

332 # assert [7] is still stored, but [9] is not

332 - assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is True

333 - assert cpu_manager.lookup(to_key(9), _EMPTY_REQ_CTX) is False

333 + assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is LookupResult.HIT

334 + assert cpu_manager.lookup(to_key(9), _EMPTY_REQ_CTX) is LookupResult.MISS

335

336 verify_events(

337 cpu_manager.take_events(),

@@ -412,8 +413,8 @@ class TestARCPolicy:

413 )

414

415 # lookup [1, 2] -> write in-flight, not yet ready

415 - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is None

416 - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is None

416 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING

417 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING

418

419 # no events so far

420 assert list(cpu_manager.take_events()) == []

@@ -423,9 +424,9 @@ class TestARCPolicy:

424 verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))

425

426 # lookup [1, 2]

426 - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is True

427 - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True

428 - assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False

427 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT

428 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT

429 + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS

430

431 # blocks should be in T1 (recent)

432 assert len(arc_policy.t1) == 2

@@ -629,7 +630,7 @@ class TestARCPolicy:

630 cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX, success=False)

631

632 # block 5 should not be in cache

632 - assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is False

633 + assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is LookupResult.MISS

634 # block 5 should not be in T1 or T2

635 assert to_keys([5])[0] not in arc_policy.t1

636 assert to_keys([5])[0] not in arc_policy.t2

@@ -670,8 +671,8 @@ class TestARCPolicy:

671 cpu_manager.complete_store(to_keys([6]), _EMPTY_REQ_CTX)

672

673 # verify blocks 2, 3 (in T2) are still present

673 - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True

674 - assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is True

674 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT

675 + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.HIT

676

677 # verify events

678 events = list(cpu_manager.take_events())

@@ -691,8 +692,8 @@ def test_filter_reused_manager():

692 )

693

694 # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet

694 - assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False

695 - assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False

695 + assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS

696 + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.MISS

697

698 # prepare store [1, 2] -> should be filtered

699 prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)

@@ -700,7 +701,7 @@ def test_filter_reused_manager():

701 assert prepare_store_output.keys_to_store == []

702

703 # Lookup [1] -> 2nd time, eligible now

703 - assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False

704 + assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS

705

706 # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered

707 prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)

@@ -709,13 +710,13 @@ def test_filter_reused_manager():

710

711 # Lookup [3, 4] -> 1st time

712 # (evicts [2] from tracker since max_size is 3 and tracker has [1])

712 - assert manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False

713 - assert manager.lookup(to_key(4), _EMPTY_REQ_CTX) is False

713 + assert manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS

714 + assert manager.lookup(to_key(4), _EMPTY_REQ_CTX) is LookupResult.MISS

715 # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4])

716 assert to_keys([2])[0] not in manager.counts

717

718 # Lookup [2] again -> (this adds [2] back to the tracker as 1st time)

718 - assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False

719 + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.MISS

720 # Verify [2] was re-added with count=1 (not eligible yet)

721 assert manager.counts.get(to_keys([2])[0]) == 1

722

tests/v1/kv_offload/tiering/test_fs_tier.py

+20 -6

@@ -18,7 +18,12 @@ import numpy as np

18 import pytest

19 import torch

21 -from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key

21 +from vllm.v1.kv_offload.base import (

22 + LookupResult,

23 + OffloadKey,

24 + ReqContext,

25 + make_offload_key,

26 +)

27 from vllm.v1.kv_offload.tiering.base import JobMetadata

28 from vllm.v1.kv_offload.tiering.fs.manager import (

29 FileSystemTierManager,

@@ -166,7 +171,7 @@ def fs_tier(tmp_path):

171 def test_lookup_empty_tier(fs_tier):

172 tier, _ = fs_tier

173 results = lookup_and_wait(tier, [key(1), key(2)])

169 - assert results == [False, False]

174 + assert results == [LookupResult.MISS, LookupResult.MISS]

175

176

177 def test_store_creates_file_and_lookup_succeeds(fs_tier):

@@ -176,7 +181,7 @@ def test_store_creates_file_and_lookup_succeeds(fs_tier):

181 results = drain(tier)

182 assert len(results) == 1

183 assert results[0].success

179 - assert lookup_and_wait(tier, [key(1)]) == [True]

184 + assert lookup_and_wait(tier, [key(1)]) == [LookupResult.HIT]

185 dest = tier.file_mapper.get_file_name(key(1))

186 assert os.path.exists(dest), f"Expected file at {dest}"

187

@@ -188,14 +193,20 @@ def test_store_then_load_roundtrip(fs_tier):

193 store_results = drain(tier)

194 assert all(r.success for r in store_results)

195

191 - assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True]

196 + assert lookup_and_wait(tier, [key(1), key(2)]) == [

197 + LookupResult.HIT,

198 + LookupResult.HIT,

199 + ]

200

201 job_l = make_job(2, [key(1), key(2)], [2, 3], is_promotion=True)

202 tier.submit_load(job_l)

203 load_results = drain(tier)

204 assert all(r.success for r in load_results)

205 # Blocks stay on disk after load

198 - assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True]

206 + assert lookup_and_wait(tier, [key(1), key(2)]) == [

207 + LookupResult.HIT,

208 + LookupResult.HIT,

209 + ]

210

211

212 def test_invalid_path_raises_at_construction():

@@ -231,7 +242,10 @@ def test_multiple_jobs_tracked_independently(fs_tier):

242 results = drain(tier)

243 job_ids = {r.job_id for r in results}

244 assert job_ids == {1, 2}

234 - assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True]

245 + assert lookup_and_wait(tier, [key(1), key(2)]) == [

246 + LookupResult.HIT,

247 + LookupResult.HIT,

248 + ]

249

250

251 def test_multi_block_job_partial_failure(fs_tier):

tests/v1/kv_offload/tiering/test_obj_tier.py

+16 -7

@@ -17,7 +17,12 @@ from unittest.mock import MagicMock, patch

17 import numpy as np

18 import torch

20 -from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key

20 +from vllm.v1.kv_offload.base import (

21 + LookupResult,

22 + OffloadKey,

23 + ReqContext,

24 + make_offload_key,

25 +)

26 from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult

27 from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager

@@ -236,19 +241,19 @@ class TestMockObjTierBasic:

241 self.tier, self.agent = _make_tier(num_blocks=4)

242

243 def test_lookup_empty_tier(self):

239 - assert lookup_and_wait(self.tier, [key(1)]) == [False]

244 + assert lookup_and_wait(self.tier, [key(1)]) == [LookupResult.MISS]

245

246 def test_store_and_lookup(self):

247 self.tier.submit_store(make_job(1, [key(1)], [0]))

248 results = drain(self.tier)

249 assert len(results) == 1

250 assert results[0].success

246 - assert lookup_and_wait(self.tier, [key(1)]) == [True]

251 + assert lookup_and_wait(self.tier, [key(1)]) == [LookupResult.HIT]

252

253 def test_lookup_unrelated_key_returns_false(self):

254 self.tier.submit_store(make_job(1, [key(1)], [0]))

255 drain(self.tier)

251 - assert lookup_and_wait(self.tier, [key(999)]) == [False]

256 + assert lookup_and_wait(self.tier, [key(999)]) == [LookupResult.MISS]

257

258 def test_store_then_load_roundtrip(self):

259 self.tier.submit_store(make_job(1, [key(1), key(2)], [0, 1]))

@@ -327,13 +332,17 @@ class TestMockObjTierMultiBlock:

332 results = drain(tier)

333 assert len(results) == 1

334 assert results[0].success

330 - assert lookup_and_wait(tier, keys) == [True] * 8

335 + assert lookup_and_wait(tier, keys) == [LookupResult.HIT] * 8

336

337 def test_partial_block_lookup(self):

338 tier, _ = _make_tier(num_blocks=4)

339 tier.submit_store(make_job(1, [key(0), key(1)], [0, 1]))

340 drain(tier)

336 - assert lookup_and_wait(tier, [key(0), key(1), key(2)]) == [True, True, False]

341 + assert lookup_and_wait(tier, [key(0), key(1), key(2)]) == [

342 + LookupResult.HIT,

343 + LookupResult.HIT,

344 + LookupResult.MISS,

345 + ]

346

347

348 class TestMockObjTierFailures:

@@ -342,7 +351,7 @@ class TestMockObjTierFailures:

351 agent.query_memory = lambda *a, **k: (_ for _ in ()).throw(

352 RuntimeError("backend error")

353 )

345 - assert lookup_and_wait(tier, [key(1)]) == [False]

354 + assert lookup_and_wait(tier, [key(1)]) == [LookupResult.MISS]

355

356 def test_submit_store_register_memory_failure_reported_in_get_finished(self):

357 tier, agent = _make_tier(num_blocks=4)

tests/v1/kv_offload/tiering/test_tiering_offloading.py

+30 -21

@@ -21,6 +21,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (

21 OffloadingConnectorStats,

22 )

23 from vllm.v1.kv_offload.base import (

24 + LookupResult,

25 OffloadingCounterMetadata,

26 OffloadKey,

27 OffloadPolicy,

@@ -60,15 +61,15 @@ def to_keys(int_ids: Iterable[int]) -> list[OffloadKey]:

61 def count_hits(manager, keys: list[OffloadKey]) -> int | None:

62 """Count consecutive lookup hits from the start of keys.

63 - Returns the count of leading True results, or None if any lookup

64 - returns None (retry-later signal).

64 + Returns the count of leading HIT results, or None if any lookup

65 + returns HIT_PENDING or RETRY.

66 """

67 count = 0

68 for key in keys:

69 result = manager.lookup(key, _CTX)

69 - if result is None:

70 + if result in (LookupResult.HIT_PENDING, LookupResult.RETRY):

71 return None

71 - if not result:

72 + if result is not LookupResult.HIT:

73 break

74 count += 1

75 return count

@@ -185,18 +186,18 @@ class TestExampleSecondaryTierManager:

186

187 # Initially empty

188 blocks = to_keys(range(3))

188 - assert tier.lookup(blocks[0], _CTX) is False

189 + assert tier.lookup(blocks[0], _CTX) is LookupResult.MISS

190

191 # Store blocks (simulate with direct insertion for testing)

192 tier.blocks[blocks[0]] = True

193 tier.blocks[blocks[1]] = True

194

195 # Lookup should find first two blocks

195 - assert tier.lookup(blocks[0], _CTX) is True

196 - assert tier.lookup(blocks[1], _CTX) is True

196 + assert tier.lookup(blocks[0], _CTX) is LookupResult.HIT

197 + assert tier.lookup(blocks[1], _CTX) is LookupResult.HIT

198

199 # Third block not present

199 - assert tier.lookup(blocks[2], _CTX) is False

200 + assert tier.lookup(blocks[2], _CTX) is LookupResult.MISS

201

202

203 class TestTieringOffloadingManager:

@@ -283,8 +284,12 @@ class TestTieringOffloadingManager:

284 assert self.secondary_tier2.get_num_blocks() == 3

285

286 # Verify blocks are present

286 - assert all(self.secondary_tier1.lookup(b, _CTX) for b in blocks)

287 - assert all(self.secondary_tier2.lookup(b, _CTX) for b in blocks)

287 + assert all(

288 + self.secondary_tier1.lookup(b, _CTX) is LookupResult.HIT for b in blocks

289 + )

290 + assert all(

291 + self.secondary_tier2.lookup(b, _CTX) is LookupResult.HIT for b in blocks

292 + )

293

294 def test_ref_cnt_protection_during_cascade(self, manager_setup):

295 """Test that ref_cnt protects blocks during cascade."""

@@ -355,7 +360,7 @@ class TestTieringOffloadingManager:

360 # Lookup each block to initiate promotion for all of them

361 for block in blocks:

362 result = self.manager.lookup(block, _CTX)

358 - assert result is None # Retry later (promotion initiated)

363 + assert result is LookupResult.RETRY # promotion initiated

364

365 # End of step 1: flushes deferred submit_load() calls

366 self._simulate_on_schedule_end()

@@ -470,11 +475,11 @@ class TestTieringOffloadingManager:

475 ctx_a = ReqContext(req_id="req_a")

476 ctx_b = ReqContext(req_id="req_b")

477

473 - # All lookups return None: secondary hit triggers promotion (in-flight)

474 - assert self.manager.lookup(blocks[0], ctx_a) is None

475 - assert self.manager.lookup(blocks[1], ctx_a) is None

476 - assert self.manager.lookup(blocks[2], ctx_b) is None

477 - assert self.manager.lookup(blocks[3], ctx_b) is None

478 + # All lookups return RETRY: secondary hit triggers promotion

479 + assert self.manager.lookup(blocks[0], ctx_a) is LookupResult.RETRY

480 + assert self.manager.lookup(blocks[1], ctx_a) is LookupResult.RETRY

481 + assert self.manager.lookup(blocks[2], ctx_b) is LookupResult.RETRY

482 + assert self.manager.lookup(blocks[3], ctx_b) is LookupResult.RETRY

483

484 # submit_load must not fire during lookup - only at end of step

485 self.secondary_tier1.submit_load.assert_not_called()

@@ -511,9 +516,10 @@ class TestTieringOffloadingManager:

516 result_a = self.manager.lookup(shared_block, ctx_a)

517 result_b = self.manager.lookup(shared_block, ctx_b)

518

514 - # Both see None (in-flight), but promotion is only queued once

515 - assert result_a is None

516 - assert result_b is None

519 + # First lookup triggers promotion (RETRY), second finds block

520 + # already in primary with write in-flight (HIT_PENDING).

521 + assert result_a is LookupResult.RETRY

522 + assert result_b is LookupResult.HIT_PENDING

523

524 self._simulate_on_schedule_end()

525

@@ -796,7 +802,10 @@ class TestTieringOffloadingManager:

802 # the lookup that staged it).

803 promo_block = to_keys([99])[0]

804 self.secondary_tier1.blocks[promo_block] = True

799 - assert self.manager.lookup(promo_block, ReqContext(req_id="pending")) is None

805 + assert (

806 + self.manager.lookup(promo_block, ReqContext(req_id="pending"))

807 + is LookupResult.RETRY

808 + )

809 assert self.manager._pending_load_submissions

810

811 # Request-level tier registration.

@@ -829,7 +838,7 @@ class TestTieringOffloadingManager:

838 assert self.primary_tier._num_allocated_blocks == 0

839 assert self.primary_tier._free_list == []

840 for block in blocks:

832 - assert self.primary_tier.lookup(block, _CTX) is False

841 + assert self.primary_tier.lookup(block, _CTX) is LookupResult.MISS

842

843 # Pending submission was dropped, not submitted.

844 self.secondary_tier1.submit_load.assert_not_called()

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+32 -21

@@ -35,6 +35,7 @@ from vllm.v1.kv_cache_interface import (

35 )

36 from vllm.v1.kv_offload.base import (

37 GPULoadStoreSpec,

38 + LookupResult,

39 OffloadingManager,

40 OffloadingSpec,

41 OffloadKey,

@@ -393,15 +394,18 @@ class OffloadingConnectorScheduler:

394 hit_count = 0

395 defer_lookup = False

396 for key in keys:

396 - result = self.manager.lookup(key, req_context)

397 - if result is None:

398 - defer_lookup = True

399 - # continue lookup to allow manager to kick-off async lookups

400 - # for all blocks (until a miss is detected)

401 - result = True

402 - if not result:

403 - break

404 - hit_count += 1

397 + match self.manager.lookup(key, req_context):

398 + case LookupResult.HIT:

399 + hit_count += 1

400 + case LookupResult.HIT_PENDING:

401 + defer_lookup = True

402 + hit_count += 1

403 + case LookupResult.RETRY:

404 + # Don't break: keep scanning to let manager kick off

405 + # async lookups (until a miss is detected).

406 + defer_lookup = True

407 + case LookupResult.MISS:

408 + break

409 return hit_count if not defer_lookup else None

410

411 def _sliding_window_lookup(

@@ -416,18 +420,25 @@ class OffloadingConnectorScheduler:

420 defer_lookup = False

421 consecutive_hits = 0

422 for idx in range(len(keys) - 1, -1, -1):

419 - result = self.manager.lookup(keys[idx], req_context)

420 - if result is None:

421 - defer_lookup = True

422 - # continue lookup to allow manager to kick-off async lookups

423 - # for all blocks (until a hit is detected)

424 - result = False

425 - if not result:

426 - consecutive_hits = 0

427 - else:

428 - consecutive_hits += 1

429 - if consecutive_hits == sliding_window_size:

430 - return idx + sliding_window_size if not defer_lookup else None

423 + match self.manager.lookup(keys[idx], req_context):

424 + case LookupResult.HIT:

425 + consecutive_hits += 1

426 + case LookupResult.HIT_PENDING:

427 + # Block is in cache, just not readable yet — counts

428 + # as hit for the consecutive streak. Don't break:

429 + # keep scanning to let manager kick off async lookups.

430 + defer_lookup = True

431 + consecutive_hits += 1

432 + case LookupResult.RETRY:

433 + # Block location uncertain — does not count as hit.

434 + # Don't break: keep scanning to let manager kick off

435 + # async lookups.

436 + defer_lookup = True

437 + consecutive_hits = 0

438 + case LookupResult.MISS:

439 + consecutive_hits = 0

440 + if consecutive_hits == sliding_window_size:

441 + return idx + sliding_window_size if not defer_lookup else None

442 return consecutive_hits if not defer_lookup else None

443

444 def _touch(self, req_status: RequestOffloadState):

vllm/v1/kv_offload/base.py

+14 -6

@@ -7,7 +7,7 @@ Core abstractions for KV cache offloading in vLLM v1.

7 from abc import ABC, abstractmethod

8 from collections.abc import Collection, Iterable, Sequence

9 from dataclasses import dataclass

10 -from enum import Enum

10 +from enum import Enum, auto

11 from typing import TYPE_CHECKING, Any, NewType

13 import numpy as np

@@ -53,6 +53,15 @@ class ReqContext:

53 kv_transfer_params: dict[str, Any] | None = None

56 +class LookupResult(Enum):

57 + """Result of OffloadingManager.lookup()."""

58 +

59 + MISS = auto()

60 + HIT = auto()

61 + HIT_PENDING = auto()

62 + RETRY = auto()

63 +

64 +

65 class OffloadPolicy(Enum):

66 # Offload only newly-computed blocks as they arrive; prefix-hit

67 # blocks (already offloaded by a prior request) are skipped.

@@ -158,7 +167,7 @@ class OffloadingKVEventsConfig:

167

168 class OffloadingManager(ABC):

169 @abstractmethod

161 - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

170 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:

171 """

172 Checks whether a single block is offloaded and ready to be read.

173

@@ -167,10 +176,9 @@ class OffloadingManager(ABC):

176 req_context: per-request context (e.g. kv_transfer_params).

177

178 Returns:

170 - True if the block is offloaded and ready, False if not,

171 - or None if the lookup should be retried later.

172 - Returning None will delay the request handling by the vLLM

173 - scheduler.

179 + HIT if the block is offloaded and ready, MISS if not found,

180 + HIT_PENDING if found but not yet readable, or RETRY if the

181 + lookup should be retried later.

182 """

183 pass

184

vllm/v1/kv_offload/cpu/manager.py

+5 -4

@@ -11,6 +11,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (

11 )

12 from vllm.v1.kv_offload.base import (

13 LoadStoreSpec,

14 + LookupResult,

15 OffloadingEvent,

16 OffloadingManager,

17 OffloadKey,

@@ -112,7 +113,7 @@ class CPUOffloadingManager(OffloadingManager):

113 return RequestOffloadingContext()

114

115 @override

115 - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

116 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:

117 if self.counts is not None:

118 if key in self.counts:

119 self.counts.move_to_end(key)

@@ -123,10 +124,10 @@ class CPUOffloadingManager(OffloadingManager):

124 self.counts[key] = 1

125 block = self._policy.get(key)

126 if block is None:

126 - return False

127 + return LookupResult.MISS

128 if not block.is_ready:

128 - return None # write in-flight; caller should retry

129 - return True

129 + return LookupResult.HIT_PENDING

130 + return LookupResult.HIT

131

132 @override

133 def prepare_load(

vllm/v1/kv_offload/tiering/base.py

+5 -4

@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any

12 import numpy as np

14 from vllm.v1.kv_offload.base import (

15 + LookupResult,

16 OffloadingMetricMetadata,

17 OffloadKey,

18 ReqContext,

@@ -79,7 +80,7 @@ class SecondaryTierManager(ABC):

80 self.tier_type = tier_type

82 @abstractmethod

82 - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

83 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:

84 """

85 Check whether a block exists in this secondary tier.

@@ -88,9 +89,9 @@ class SecondaryTierManager(ABC):

89 req_context: per-request context (e.g. kv_transfer_params).

91 Returns:

91 - True if the block is present and ready,

92 - False if not found,

93 - or None if the block is being transferred (retry later).

92 + HIT if the block is present and ready,

93 + MISS if not found,

94 + or RETRY if the block is being transferred (retry later).

95 """

96 pass

vllm/v1/kv_offload/tiering/example/manager.py

+9 -4

@@ -15,7 +15,12 @@ from typing import TYPE_CHECKING

16 from typing_extensions import override

18 -from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext

18 +from vllm.v1.kv_offload.base import (

19 + LookupResult,

20 + OffloadKey,

21 + ReqContext,

22 + RequestOffloadingContext,

23 +)

24 from vllm.v1.kv_offload.tiering.base import (

25 JobMetadata,

26 JobResult,

@@ -67,7 +72,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager):

72 self.completed_jobs: list[JobResult] = []

74 @override

70 - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

75 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:

76 """

77 Check whether a block exists in this secondary tier.

@@ -76,9 +81,9 @@ class ExampleSecondaryTierManager(SecondaryTierManager):

81 req_context: Per-request context.

83 Returns:

79 - True if the block is present, False if not found.

84 + HIT if the block is present, MISS if not found.

85 """

81 - return key in self.blocks

86 + return LookupResult.HIT if key in self.blocks else LookupResult.MISS

88 @override

89 def submit_store(self, job_metadata: JobMetadata) -> None:

vllm/v1/kv_offload/tiering/fs/manager.py

+6 -3

@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING

24 from typing_extensions import override

26 from vllm.logger import init_logger

27 -from vllm.v1.kv_offload.base import OffloadKey, ReqContext

27 +from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext

28 from vllm.v1.kv_offload.file_mapper import FileMapper

29 from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager

30 from vllm.v1.kv_offload.tiering.base import (

@@ -137,8 +137,11 @@ class FileSystemTierManager(SecondaryTierManager):

137 return RequestOffloadingContext()

138

139 @override

140 - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

141 - return self._lookup_manager.lookup(key, req_context)

140 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:

141 + result = self._lookup_manager.lookup(key, req_context)

142 + if result is None:

143 + return LookupResult.RETRY

144 + return LookupResult.HIT if result else LookupResult.MISS

145

146 @override

147 def submit_store(self, job_metadata: JobMetadata) -> None:

vllm/v1/kv_offload/tiering/manager.py

+24 -20

@@ -32,6 +32,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (

32 from vllm.logger import init_logger

33 from vllm.v1.kv_offload.base import (

34 LoadStoreSpec,

35 + LookupResult,

36 OffloadingEvent,

37 OffloadingManager,

38 OffloadKey,

@@ -233,7 +234,7 @@ class TieringOffloadingManager(OffloadingManager):

234 )

235

236 @override

236 - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

237 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:

238 """

239 Check whether a single block is offloaded and ready.

240

@@ -248,33 +249,34 @@ class TieringOffloadingManager(OffloadingManager):

249 req_context: Per-request context.

250

251 Returns:

251 - True — block is ready in the primary tier.

252 - None — block found but not yet ready (primary in-flight,

253 - promotion started, or a secondary tier is busy).

254 - False — block not found in any tier, or primary is full

255 - and cannot accept a promotion.

252 + HIT — block is ready in the primary tier.

253 + HIT_PENDING — block found but not yet readable (write

254 + in-flight on the primary tier).

255 + RETRY — promotion started or a secondary tier is busy.

256 + MISS — block not found in any tier, or primary is full

257 + and cannot accept a promotion.

258 """

259 self._maybe_process_finished_jobs()

260

261 primary_hit = self.primary_tier.lookup(key, req_context)

260 - if primary_hit is True:

261 - return True

262 - if primary_hit is None:

263 - return None

262 + if primary_hit is LookupResult.HIT:

263 + return LookupResult.HIT

264 + if primary_hit is LookupResult.HIT_PENDING:

265 + return LookupResult.HIT_PENDING

266

265 - any_none = False

267 + any_retry = False

268 for tier in self.secondary_tiers:

269 result = tier.lookup(key, req_context)

268 - if result is True:

270 + if result is LookupResult.HIT:

271 if not self._initiate_promotion(tier, key, req_context):

270 - return False # primary full, block unavailable

271 - return None # promotion started, retry later

272 - if result is None:

273 - any_none = True

272 + return LookupResult.MISS

273 + return LookupResult.RETRY

274 + if result is LookupResult.RETRY:

275 + any_retry = True

276

275 - if any_none:

276 - return None

277 - return False

277 + if any_retry:

278 + return LookupResult.RETRY

279 + return LookupResult.MISS

280

281 def _initiate_promotion(

282 self,

@@ -467,7 +469,9 @@ class TieringOffloadingManager(OffloadingManager):

469 """

470 # Filter out keys that are not ready in primary (e.g. in-flight)

471 ready_keys = tuple(

470 - k for k in keys if self.primary_tier.lookup(k, req_context) is True

472 + k

473 + for k in keys

474 + if self.primary_tier.lookup(k, req_context) is LookupResult.HIT

475 )

476 if not ready_keys:

477 return

vllm/v1/kv_offload/tiering/obj/manager.py

+6 -3

@@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, NamedTuple

10 from vllm.distributed.nixl_utils import NixlWrapper as nixl_agent

11 from vllm.distributed.nixl_utils import nixl_agent_config

12 from vllm.logger import init_logger

13 -from vllm.v1.kv_offload.base import OffloadKey, ReqContext

13 +from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext

14 from vllm.v1.kv_offload.file_mapper import FileMapper

15 from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager

16 from vllm.v1.kv_offload.tiering.base import (

@@ -221,8 +221,11 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

221

222 self._transfers[job_id] = TransferEntry(xfer_handle, files_desc, obj_handle)

223

224 - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

225 - return self._lookup_manager.lookup(key, req_context)

224 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:

225 + result = self._lookup_manager.lookup(key, req_context)

226 + if result is None:

227 + return LookupResult.RETRY

228 + return LookupResult.HIT if result else LookupResult.MISS

229

230 def submit_store(self, job_metadata: JobMetadata) -> None:

231 obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys)