| 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, 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 ( |
| 10 | CachePolicy, |
| 11 | ChunkStatus, |
| 12 | order_request_keys, |
| 13 | ) |
| 14 | |
| 15 | |
| 16 | class ARCCachePolicy(CachePolicy): |
| 17 | """ARC (Adaptive Replacement Cache) cache policy. |
| 18 | |
| 19 | Data Structures: |
| 20 | T1: Recent cache containing chunks accessed once. |
| 21 | T2: Frequent cache containing chunks accessed multiple times. |
| 22 | B1/B2: Ghost lists tracking recently evicted chunks from T1/T2. |
| 23 | target_t1_size: Adaptive target size for the T1 partition. |
| 24 | |
| 25 | Algorithm Flow: |
| 26 | 1. Cache lookup (lookup): |
| 27 | Searches T1 and T2 for chunk hashes and counts consecutive hits |
| 28 | until a miss or non-ready chunk is encountered. |
| 29 | |
| 30 | 2. Request access - Adaptive Learning: |
| 31 | - Ready chunks reused by a request move from T1 to T2 once, or |
| 32 | move to the MRU end of T2. |
| 33 | - B1/B2 misses adjust target_t1_size once per request before |
| 34 | insertion removes their ghost entries. |
| 35 | |
| 36 | 3. Chunk eviction (evict) - Adaptive Replacement: |
| 37 | Determines eviction source based on adaptive target: |
| 38 | - If T1 size >= target_t1_size: Evict from T1, add to B1. |
| 39 | - Otherwise: Evict from T2, add to B2. |
| 40 | Finally, bound each ghost list size. |
| 41 | |
| 42 | 4. Chunk insertion (insert): |
| 43 | New chunks are always inserted into T1 and removed from B1/B2 if |
| 44 | present. A later request reuse may promote them to T2. |
| 45 | |
| 46 | Adaptive Behavior: |
| 47 | The algorithm self-tunes the recency vs. frequency trade-off: |
| 48 | - B1 hit: Recent access patterns matter more → increase T1. |
| 49 | - B2 hit: Frequent access patterns matter more → decrease T1. |
| 50 | """ |
| 51 | |
| 52 | def __init__(self, cache_capacity: int): |
| 53 | super().__init__(cache_capacity) |
| 54 | self.target_t1_size: float = 0.0 |
| 55 | self.t1: OrderedDict[OffloadKey, ChunkStatus] = OrderedDict() |
| 56 | self.t2: OrderedDict[OffloadKey, ChunkStatus] = OrderedDict() |
| 57 | # key -> None (only care about presence) |
| 58 | self.b1: OrderedDict[OffloadKey, None] = OrderedDict() |
| 59 | self.b2: OrderedDict[OffloadKey, None] = OrderedDict() |
| 60 | |
| 61 | @override |
| 62 | def get(self, key: OffloadKey) -> ChunkStatus | None: |
| 63 | return self.t1.get(key) or self.t2.get(key) |
| 64 | |
| 65 | @override |
| 66 | def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None: |
| 67 | self.t1[key] = chunk |
| 68 | self.b1.pop(key, None) |
| 69 | self.b2.pop(key, None) |
| 70 | |
| 71 | @override |
| 72 | def remove(self, key: OffloadKey) -> None: |
| 73 | if self.t1.pop(key, None) is None: |
| 74 | self.t2.pop(key, None) |
| 75 | |
| 76 | def _adapt_to_ghost_hit(self, key: OffloadKey) -> bool: |
| 77 | if key in self.b1: |
| 78 | delta = max(1, len(self.b2) / len(self.b1)) |
| 79 | self.target_t1_size = min(self.target_t1_size + delta, self.cache_capacity) |
| 80 | self.b1.move_to_end(key) |
| 81 | return True |
| 82 | if key in self.b2: |
| 83 | delta = max(1, len(self.b1) / len(self.b2)) |
| 84 | self.target_t1_size = max(self.target_t1_size - delta, 0) |
| 85 | self.b2.move_to_end(key) |
| 86 | return True |
| 87 | return False |
| 88 | |
| 89 | @override |
| 90 | def on_store_miss( |
| 91 | self, keys: Iterable[OffloadKey], req_context: ReqContext |
| 92 | ) -> None: |
| 93 | for key in reversed(list(keys)): |
| 94 | self._adapt_to_ghost_hit(key) |
| 95 | |
| 96 | @override |
| 97 | def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: |
| 98 | for key in reversed(list(keys)): |
| 99 | if key in self.t1: |
| 100 | chunk = self.t1.pop(key) |
| 101 | if not chunk.is_ready: |
| 102 | # chunk was just prepared to be stored, not really touched |
| 103 | # twice — keep it in T1 and mark as most recently used |
| 104 | self.t1[key] = chunk |
| 105 | else: |
| 106 | self.t2[key] = chunk |
| 107 | |
| 108 | elif key in self.t2: |
| 109 | self.t2.move_to_end(key) |
| 110 | |
| 111 | else: |
| 112 | self._adapt_to_ghost_hit(key) |
| 113 | |
| 114 | @override |
| 115 | def on_request_finished( |
| 116 | self, |
| 117 | key_groups: Sequence[Sequence[OffloadKey]], |
| 118 | insertion_only_keys: set[OffloadKey], |
| 119 | reused_keys: set[OffloadKey], |
| 120 | req_context: ReqContext, |
| 121 | ) -> None: |
| 122 | for key in reversed(order_request_keys(key_groups, req_context)): |
| 123 | if key in insertion_only_keys: |
| 124 | # A store is not a frequency hit. Preserve T1 membership |
| 125 | # while still restoring tail-to-head recency. |
| 126 | if key in self.t1: |
| 127 | self.t1.move_to_end(key) |
| 128 | elif key in self.t2: |
| 129 | self.t2.move_to_end(key) |
| 130 | continue |
| 131 | |
| 132 | # Ready chunks reused by this request count as one access. |
| 133 | if key in reused_keys and key in self.t1: |
| 134 | chunk = self.t1.pop(key) |
| 135 | self.t2[key] = chunk |
| 136 | elif key in reused_keys and key in self.t2: |
| 137 | self.t2.move_to_end(key) |
| 138 | |
| 139 | @override |
| 140 | def clear(self) -> None: |
| 141 | self.t1.clear() |
| 142 | self.t2.clear() |
| 143 | self.b1.clear() |
| 144 | self.b2.clear() |
| 145 | self.target_t1_size = 0.0 |
| 146 | |
| 147 | @override |
| 148 | def evict( |
| 149 | self, n: int, protected: set[OffloadKey] |
| 150 | ) -> list[tuple[OffloadKey, ChunkStatus]] | None: |
| 151 | if n == 0: |
| 152 | return [] |
| 153 | |
| 154 | # Collect candidates atomically: simulate T1 size changes as we select, |
| 155 | # but do not modify actual data structures until all n are found. |
| 156 | candidates: list[ |
| 157 | tuple[OffloadKey, ChunkStatus, bool] |
| 158 | ] = [] # (key, chunk, from_t1) |
| 159 | virtual_t1_size = len(self.t1) |
| 160 | # Keep the scans monotonic: restarting from the LRU end after every |
| 161 | # selection makes a batch eviction quadratic in the number of chunks. |
| 162 | t1_iter = iter(self.t1.items()) |
| 163 | t2_iter = iter(self.t2.items()) |
| 164 | |
| 165 | def next_candidate( |
| 166 | entries: Iterator[tuple[OffloadKey, ChunkStatus]], |
| 167 | ) -> tuple[OffloadKey, ChunkStatus] | None: |
| 168 | for key, chunk in entries: |
| 169 | if chunk.ref_cnt == 0 and key not in protected: |
| 170 | return key, chunk |
| 171 | return None |
| 172 | |
| 173 | for _ in range(n): |
| 174 | candidate: tuple[OffloadKey, ChunkStatus, bool] | None = None |
| 175 | |
| 176 | if virtual_t1_size >= int(self.target_t1_size): |
| 177 | entry = next_candidate(t1_iter) |
| 178 | if entry is not None: |
| 179 | candidate = (*entry, True) |
| 180 | virtual_t1_size -= 1 |
| 181 | |
| 182 | if candidate is None: |
| 183 | entry = next_candidate(t2_iter) |
| 184 | if entry is not None: |
| 185 | candidate = (*entry, False) |
| 186 | |
| 187 | if candidate is None: |
| 188 | entry = next_candidate(t1_iter) |
| 189 | if entry is None: |
| 190 | return None |
| 191 | candidate = (*entry, True) |
| 192 | virtual_t1_size -= 1 |
| 193 | |
| 194 | candidates.append(candidate) |
| 195 | |
| 196 | # Apply all evictions now that we know n candidates exist. |
| 197 | result: list[tuple[OffloadKey, ChunkStatus]] = [] |
| 198 | for key, chunk, from_t1 in candidates: |
| 199 | if from_t1: |
| 200 | del self.t1[key] |
| 201 | self.b1[key] = None |
| 202 | else: |
| 203 | del self.t2[key] |
| 204 | self.b2[key] = None |
| 205 | result.append((key, chunk)) |
| 206 | |
| 207 | # Trim ghost lists to cache_capacity. |
| 208 | for ghost in (self.b1, self.b2): |
| 209 | for _ in range(len(ghost) - self.cache_capacity): |
| 210 | ghost.popitem(last=False) |
| 211 | |
| 212 | return result |