| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import ctypes |
| 4 | from abc import ABC, abstractmethod |
| 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 | """Offloading status for a single chunk of KV data. |
| 27 | Holds the following information: |
| 28 | |
| 29 | ref_cnt - the current number of transfers using this chunk as a source. |
| 30 | A value of -1 indicates the chunk is not yet ready to be read. |
| 31 | chunk_id - index of the physical CPU buffer slot. |
| 32 | """ |
| 33 | |
| 34 | _fields_ = [("ref_cnt", ctypes.c_int32), ("chunk_id", ctypes.c_int64)] |
| 35 | |
| 36 | def __init__(self, chunk_id: int): |
| 37 | super().__init__() |
| 38 | # initialize chunk as "not ready" (ref_cnt = -1) |
| 39 | self.ref_cnt = -1 |
| 40 | self.chunk_id = chunk_id |
| 41 | |
| 42 | @property |
| 43 | def is_ready(self) -> bool: |
| 44 | """Returns whether the chunk is ready to be read.""" |
| 45 | return self.ref_cnt >= 0 |
| 46 | |
| 47 | |
| 48 | class CachePolicy(ABC): |
| 49 | """Encapsulates both chunk organization (data structures) and replacement |
| 50 | decisions (which chunk to evict). LRU and ARC differ in both dimensions — |
| 51 | ARC's ghost lists and target_t1_size live at the intersection of storage |
| 52 | and eviction, so they cannot be separated cleanly. |
| 53 | """ |
| 54 | |
| 55 | def __init__(self, cache_capacity: int) -> None: |
| 56 | self.cache_capacity = cache_capacity |
| 57 | |
| 58 | @abstractmethod |
| 59 | def get(self, key: OffloadKey) -> ChunkStatus | None: |
| 60 | """Find chunk in data structures. Returns None if not present.""" |
| 61 | |
| 62 | @abstractmethod |
| 63 | def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None: |
| 64 | """Add a newly allocated chunk. For ARC: also removes from ghost lists.""" |
| 65 | |
| 66 | @abstractmethod |
| 67 | def remove(self, key: OffloadKey) -> None: |
| 68 | """Remove a chunk (used to clean up after a failed store).""" |
| 69 | |
| 70 | @abstractmethod |
| 71 | def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: |
| 72 | """Mark chunks as recently used. |
| 73 | |
| 74 | Args: |
| 75 | keys: Chunks to mark as recently used. |
| 76 | req_context: Per-request context for the request touching these chunks. |
| 77 | |
| 78 | """ |
| 79 | |
| 80 | def on_store_miss( |
| 81 | self, keys: Iterable[OffloadKey], req_context: ReqContext |
| 82 | ) -> None: |
| 83 | """Observe store misses before their cache entries are inserted. |
| 84 | |
| 85 | The default delegates to ``touch`` for compatibility with external |
| 86 | policies. Policies may override this hook when a miss has distinct |
| 87 | semantics, such as ARC adapting to a ghost-list hit. |
| 88 | """ |
| 89 | self.touch(keys, req_context) |
| 90 | |
| 91 | def on_request_finished( |
| 92 | self, |
| 93 | key_groups: Sequence[Sequence[OffloadKey]], |
| 94 | insertion_only_keys: set[OffloadKey], |
| 95 | reused_keys: set[OffloadKey], |
| 96 | req_context: ReqContext, |
| 97 | ) -> None: |
| 98 | """Apply one request-scoped cache access in prefix order. |
| 99 | |
| 100 | ``key_groups`` contains keys observed for each KV cache group in |
| 101 | head-to-tail order. ``insertion_only_keys`` and ``reused_keys`` |
| 102 | distinguish chunks only created by this request from ready chunks it |
| 103 | actually reused, which matters for policies such as ARC where reuse |
| 104 | changes frequency but insertion and pending observations do not. |
| 105 | |
| 106 | The default forwards one touch per group, preserving compatibility |
| 107 | for experimental out-of-tree policies while moving those touches to |
| 108 | request finalization. Policies that distinguish insertion from reuse |
| 109 | can override this hook and inspect the access classifications. |
| 110 | """ |
| 111 | del insertion_only_keys, reused_keys |
| 112 | for keys in key_groups: |
| 113 | self.touch(keys, req_context) |
| 114 | |
| 115 | @abstractmethod |
| 116 | def evict( |
| 117 | self, n: int, protected: set[OffloadKey] |
| 118 | ) -> list[tuple[OffloadKey, ChunkStatus]] | None: |
| 119 | """Evict exactly n chunks, skipping any in protected. |
| 120 | |
| 121 | Returns a list of (key, chunk) for the evicted chunks, |
| 122 | or None if n evictions cannot be satisfied. The operation is atomic: |
| 123 | if None is returned, no state changes are made. |
| 124 | |
| 125 | For ARC: ghost list cleanup (trimming to cache_capacity) is performed |
| 126 | at the end of a successful eviction. |
| 127 | """ |
| 128 | |
| 129 | @abstractmethod |
| 130 | def clear(self) -> None: |
| 131 | """Remove ALL chunks regardless of ref_cnt. |
| 132 | |
| 133 | Ghost lists and adaptive state are also reset. |
| 134 | """ |
| 135 | |
| 136 | def mark_evictable(self, key: OffloadKey) -> None: |
| 137 | """Called when a chunk's ref_cnt transitions to 0.""" |
| 138 | return |
| 139 | |
| 140 | def mark_non_evictable(self, key: OffloadKey) -> None: |
| 141 | """Called when a chunk's ref_cnt transitions from 0.""" |
| 142 | return |