[Perf][KV Offload] Avoid quadratic ARC batch eviction (#50992)
Signed-off-by: MINJUN GIL <alswnsrlf12@naver.com>
MINJUN GIL committed
Aug 5, 2026 at 16:58 UTC
b92352ca2c5625a8ca61da597046d50a121c54ad
2 files changed
+132
-21
tests/v1/kv_offload/cpu/test_manager.py
+112
@@ -1,5 +1,6 @@
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
5
from dataclasses import dataclass
6
@@ -59,6 +60,17 @@ class ExpectedPrepareStoreOutput:
60
evicted_keys: list[int]
61
62
63
+class _CountingOrderedDict(OrderedDict):
64
+ def __init__(self, *args, **kwargs):
65
+ super().__init__(*args, **kwargs)
66
+ self.items_yielded = 0
67
+
68
+ def items(self):
69
+ for item in super().items():
70
+ self.items_yielded += 1
71
+ yield item
72
+
73
+
74
def to_key(int_hash: int) -> OffloadKey:
75
return make_offload_key(str(int_hash).encode(), 0)
76
@@ -686,6 +698,106 @@ class TestARCPolicy:
698
# block 5 should be in T1
699
assert to_keys([5])[0] in arc_policy.t1
700
701
+ def test_batch_eviction_scans_t1_and_t2_once(self):
702
+ """ARC batch eviction must preserve order without restarting scans."""
703
+ cpu_manager, arc_policy = self._make_manager(
704
+ num_blocks=256, enable_events=False
705
+ )
706
+ keys = to_keys(list(range(256)))
707
+ cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX)
708
+ cpu_manager.complete_store(keys, _EMPTY_REQ_CTX)
709
+
710
+ cpu_manager.touch(keys[128:], _EMPTY_REQ_CTX)
711
+ arc_policy.target_t1_size = 64
712
+
713
+ protected = {keys[0], keys[2], keys[255]}
714
+ arc_policy.t1[keys[1]].ref_cnt = 1
715
+ arc_policy.t2[keys[254]].ref_cnt = 1
716
+
717
+ num_evictions = 124
718
+ num_t1_evictions = len(arc_policy.t1) - int(arc_policy.target_t1_size) + 1
719
+ num_t2_evictions = num_evictions - num_t1_evictions
720
+ t1_order = list(arc_policy.t1)
721
+ t2_order = list(arc_policy.t2)
722
+ expected_t1 = [
723
+ key
724
+ for key, block in arc_policy.t1.items()
725
+ if block.ref_cnt == 0 and key not in protected
726
+ ][:num_t1_evictions]
727
+ expected_t2 = [
728
+ key
729
+ for key, block in arc_policy.t2.items()
730
+ if block.ref_cnt == 0 and key not in protected
731
+ ][:num_t2_evictions]
732
+ expected_t1_scans = t1_order.index(expected_t1[-1]) + 1
733
+ expected_t2_scans = t2_order.index(expected_t2[-1]) + 1
734
+
735
+ counting_t1 = _CountingOrderedDict(arc_policy.t1)
736
+ counting_t2 = _CountingOrderedDict(arc_policy.t2)
737
+ arc_policy.t1 = counting_t1
738
+ arc_policy.t2 = counting_t2
739
+
740
+ evicted = arc_policy.evict(num_evictions, protected)
741
+
742
+ assert evicted is not None
743
+ assert [key for key, _ in evicted] == expected_t1 + expected_t2
744
+ assert counting_t1.items_yielded == expected_t1_scans
745
+ assert counting_t2.items_yielded == expected_t2_scans
746
+
747
+ def test_batch_eviction_falls_back_after_t1_iterator_exhausted(self):
748
+ """An exhausted T1 scan must keep falling back to T2."""
749
+ cpu_manager, arc_policy = self._make_manager(num_blocks=8, enable_events=False)
750
+ keys = to_keys(list(range(8)))
751
+ cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX)
752
+ cpu_manager.complete_store(keys, _EMPTY_REQ_CTX)
753
+
754
+ cpu_manager.touch(keys[6:], _EMPTY_REQ_CTX)
755
+ arc_policy.target_t1_size = 4
756
+
757
+ t1_order = list(arc_policy.t1)
758
+ t2_order = list(arc_policy.t2)
759
+ protected = set(t1_order[1:3])
760
+ for key in t1_order[3:]:
761
+ arc_policy.t1[key].ref_cnt = 1
762
+
763
+ # Selecting the sole eligible T1 entry leaves virtual_t1_size above
764
+ # the target, so each remaining selection must retry T1 then use T2.
765
+ eligible_t1 = [
766
+ key
767
+ for key, block in arc_policy.t1.items()
768
+ if block.ref_cnt == 0 and key not in protected
769
+ ]
770
+ assert eligible_t1 == t1_order[:1]
771
+ assert len(t1_order) - 1 >= int(arc_policy.target_t1_size)
772
+
773
+ counting_t1 = _CountingOrderedDict(arc_policy.t1)
774
+ counting_t2 = _CountingOrderedDict(arc_policy.t2)
775
+ arc_policy.t1 = counting_t1
776
+ arc_policy.t2 = counting_t2
777
+
778
+ evicted = arc_policy.evict(3, protected)
779
+
780
+ assert evicted is not None
781
+ assert [key for key, _ in evicted] == [t1_order[0], *t2_order]
782
+ assert counting_t1.items_yielded == len(t1_order)
783
+ assert counting_t2.items_yielded == len(t2_order)
784
+
785
+ def test_batch_eviction_failure_is_atomic(self):
786
+ """Finding only some candidates must not partially evict the cache."""
787
+ cpu_manager, arc_policy = self._make_manager(num_blocks=4, enable_events=False)
788
+ keys = to_keys(list(range(4)))
789
+ cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX)
790
+ cpu_manager.complete_store(keys, _EMPTY_REQ_CTX)
791
+
792
+ before_t1 = list(arc_policy.t1.items())
793
+ protected = set(keys[1:])
794
+
795
+ assert arc_policy.evict(2, protected) is None
796
+ assert list(arc_policy.t1.items()) == before_t1
797
+ assert not arc_policy.t2
798
+ assert not arc_policy.b1
799
+ assert not arc_policy.b2
800
+
801
def test_ghost_list_bounds(self):
802
"""
803
Tests that ghost lists (B1, B2) don't grow unbounded.
vllm/v1/kv_offload/cpu/policies/arc.py
+20
-21
@@ -1,7 +1,7 @@
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
4
+from collections.abc import Iterable, Iterator
5
6
from typing_extensions import override
7
@@ -120,37 +120,36 @@ class ARCCachePolicy(CachePolicy):
120
candidates: list[
121
tuple[OffloadKey, BlockStatus, bool]
122
] = [] # (key, block, from_t1)
123
- already_selected: set[OffloadKey] = set()
123
virtual_t1_size = len(self.t1)
124
+ # Keep the scans monotonic: restarting from the LRU end after every
125
+ # selection makes a batch eviction quadratic in the number of blocks.
126
+ t1_iter = iter(self.t1.items())
127
+ t2_iter = iter(self.t2.items())
128
+
129
+ def next_candidate(
130
+ entries: Iterator[tuple[OffloadKey, BlockStatus]],
131
+ ) -> tuple[OffloadKey, BlockStatus] | None:
132
+ for key, block in entries:
133
+ if block.ref_cnt == 0 and key not in protected:
134
+ return key, block
135
+ return None
136
137
for _ in range(n):
138
candidate: tuple[OffloadKey, BlockStatus, bool] | None = None
139
140
if virtual_t1_size >= int(self.target_t1_size):
130
- for key, block in self.t1.items():
131
- if (
132
- block.ref_cnt == 0
133
- and key not in protected
134
- and key not in already_selected
135
- ):
136
- candidate = (key, block, True)
137
- virtual_t1_size -= 1
138
- break
141
+ entry = next_candidate(t1_iter)
142
+ if entry is not None:
143
+ candidate = (*entry, True)
144
+ virtual_t1_size -= 1
145
146
if candidate is None:
141
- for key, block in self.t2.items():
142
- if (
143
- block.ref_cnt == 0
144
- and key not in protected
145
- and key not in already_selected
146
- ):
147
- candidate = (key, block, False)
148
- break
149
- if candidate is None:
147
+ entry = next_candidate(t2_iter)
148
+ if entry is None:
149
return None
150
+ candidate = (*entry, False)
151
152
candidates.append(candidate)
153
- already_selected.add(candidate[0])
153
154
# Apply all evictions now that we know n candidates exist.
155
result: list[tuple[OffloadKey, BlockStatus]] = []