Add CachePolicyFactory for pluggable/external eviction policies (#49114)
Signed-off-by: Philip Pesic <philippesic06@gmail.com>
Philip Pesic committed
Jul 29, 2026 at 04:34 UTC
dc1be79031d948d7a18c37600881e45ca708d913
15 files changed
+269
-33
docs/features/kv_offloading_usage.md
+32
-1
@@ -70,7 +70,8 @@ vllm serve <model> \
70
| `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
71
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. |
72
| `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. |
73
-| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
73
+| `eviction_policy` | no | `lru` | both | Primary tier policy: built-in `lru`/`arc`, or a custom `CachePolicy` name (see [Custom Eviction Policies](#custom-eviction-policies)). |
74
+| `cache_policy_module_path` | no | — | both | Python import path for a custom `CachePolicy` not in the built-in registry. Required only when `eviction_policy` is not built-in and wasn't pre-registered via `CachePolicyFactory` (advanced). |
75
| `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
76
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
77
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
@@ -78,6 +79,36 @@ vllm serve <model> \
79
| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
80
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
81
82
+## Custom Eviction Policies
83
+
84
+`eviction_policy` resolves through `CachePolicyFactory` (`vllm/v1/kv_offload/cpu/policies/factory.py`), which pre-registers the built-in `lru` and `arc` policies.
85
+
86
+### Out-of-tree (recommended)
87
+
88
+Implement `CachePolicy` (`vllm/v1/kv_offload/cpu/policies/base.py`) in your own package — no vLLM fork or patch required — and point `kv_connector_extra_config` at it directly:
89
+
90
+```json
91
+{
92
+ "cpu_bytes_to_use": 10737418240,
93
+ "eviction_policy": "MyCachePolicy",
94
+ "cache_policy_module_path": "my_package.my_module"
95
+}
96
+```
97
+
98
+`eviction_policy` is checked against the built-in registry first; if it isn't a registered name, vLLM imports `cache_policy_module_path` and looks up `eviction_policy` as a class name in that module — the same fallback `spec_module_path` provides for a custom `OffloadingSpec`. No import or registration call needs to run before the server starts.
99
+
100
+### Registering a friendly short name (in-process only)
101
+
102
+If you control the process that constructs the vLLM engine (e.g. an embedding application), you can register a short name once at startup instead of repeating the module path in every config:
103
+
104
+```python
105
+from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory
106
+
107
+CachePolicyFactory.register_cache_policy("my_policy", "my_package.my_module", "MyCachePolicy")
108
+```
109
+
110
+Then set `"eviction_policy": "my_policy"` in `kv_connector_extra_config`, the same as `"lru"`/`"arc"`. This only takes effect within the process that ran the `register_cache_policy` call — it does not help when the server is launched as a separate process (e.g. via the `vllm serve` CLI), where the out-of-tree `cache_policy_module_path` config above is the only option.
111
+
112
## Secondary Tiers
113
114
Each entry in `secondary_tiers` is a dict with a required `type` field plus tier-specific fields.
tests/v1/kv_offload/cpu/__init__.py
tests/v1/kv_offload/cpu/policies/__init__.py
tests/v1/kv_offload/cpu/policies/test_factory.py
new
+111
@@ -0,0 +1,111 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Iterable
4
+
5
+import pytest
6
+
7
+from vllm.v1.kv_offload.base import OffloadKey, ReqContext
8
+from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
9
+from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
10
+from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
11
+from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory
12
+from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
13
+
14
+
15
+class _DummyCachePolicy(CachePolicy):
16
+ """Minimal CachePolicy for CachePolicyFactory registration tests. Loaded
17
+ by module path, so it must be importable at module scope (mirrors
18
+ tests/v1/kv_offload/test_factory.py's SingleArgExternalOffloadingSpec)."""
19
+
20
+ def __init__(self, cache_capacity: int) -> None:
21
+ self.cache_capacity = cache_capacity
22
+
23
+ def get(self, key: OffloadKey) -> BlockStatus | None:
24
+ return None
25
+
26
+ def insert(self, key: OffloadKey, block: BlockStatus) -> None:
27
+ pass
28
+
29
+ def remove(self, key: OffloadKey) -> None:
30
+ pass
31
+
32
+ def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
33
+ pass
34
+
35
+ def evict(
36
+ self, n: int, protected: set[OffloadKey]
37
+ ) -> list[tuple[OffloadKey, BlockStatus]] | None:
38
+ return None
39
+
40
+ def clear(self) -> None:
41
+ pass
42
+
43
+
44
+@pytest.fixture(autouse=True)
45
+def restore_cache_policy_registry():
46
+ """Save and restore CachePolicyFactory._registry between tests."""
47
+ original = dict(CachePolicyFactory._registry)
48
+ yield
49
+ CachePolicyFactory._registry = original
50
+
51
+
52
+class TestCachePolicyFactory:
53
+ """Unit tests for CachePolicyFactory (registration/resolution by name)."""
54
+
55
+ def test_pre_registered_policies_can_be_imported(self):
56
+ """If someone moves a policy module but forgets to update
57
+ factory.py, CI fails."""
58
+ for name in CachePolicyFactory._registry:
59
+ cls = CachePolicyFactory._registry[name]()
60
+ assert issubclass(cls, CachePolicy)
61
+
62
+ def test_lru_and_arc_registered(self):
63
+ assert CachePolicyFactory.get_cache_policy_cls("lru") is LRUCachePolicy
64
+ assert CachePolicyFactory.get_cache_policy_cls("arc") is ARCCachePolicy
65
+
66
+ def test_register_and_resolve_custom_policy(self):
67
+ CachePolicyFactory.register_cache_policy(
68
+ "dummy",
69
+ "tests.v1.kv_offload.cpu.policies.test_factory",
70
+ "_DummyCachePolicy",
71
+ )
72
+ policy_cls = CachePolicyFactory.get_cache_policy_cls("dummy")
73
+ assert policy_cls is _DummyCachePolicy
74
+
75
+ manager = CPUOffloadingManager(num_blocks=4, cache_policy="dummy")
76
+ assert isinstance(manager._policy, _DummyCachePolicy)
77
+
78
+ def test_unregistered_policy_raises(self):
79
+ with pytest.raises(ValueError, match="Unknown cache policy"):
80
+ CachePolicyFactory.get_cache_policy_cls("nonexistent")
81
+
82
+ def test_duplicate_registration_raises(self):
83
+ with pytest.raises(ValueError, match="is already registered"):
84
+ CachePolicyFactory.register_cache_policy("lru", "some.module", "SomeClass")
85
+
86
+ def test_dynamic_load_via_cache_policy_module_path(self):
87
+ """Out-of-tree policy loaded via cache_policy_module_path, no
88
+ register_cache_policy() call -- this is how external projects
89
+ integrate a custom CachePolicy without forking/patching vLLM.
90
+ Mirrors tests/v1/kv_offload/test_factory.py's
91
+ test_dynamic_load_via_spec_module_path."""
92
+ policy_cls = CachePolicyFactory.get_cache_policy_cls(
93
+ "_DummyCachePolicy", "tests.v1.kv_offload.cpu.policies.test_factory"
94
+ )
95
+ assert policy_cls is _DummyCachePolicy
96
+
97
+ def test_manager_resolves_policy_via_module_path(self):
98
+ """End-to-end: CPUOffloadingManager resolves an unregistered policy
99
+ purely from cache_policy_module_path."""
100
+ manager = CPUOffloadingManager(
101
+ num_blocks=4,
102
+ cache_policy="_DummyCachePolicy",
103
+ cache_policy_module_path="tests.v1.kv_offload.cpu.policies.test_factory",
104
+ )
105
+ assert isinstance(manager._policy, _DummyCachePolicy)
106
+
107
+ def test_unregistered_policy_without_module_path_raises(self):
108
+ """eviction_policy not in registry + no cache_policy_module_path ->
109
+ ValueError, same shape as the OffloadingSpecFactory error path."""
110
+ with pytest.raises(ValueError, match="Unknown cache policy"):
111
+ CachePolicyFactory.get_cache_policy_cls("nonexistent", None)
tests/v1/kv_offload/cpu/test_manager.py
+2
@@ -37,6 +37,7 @@ _EMPTY_REQ_CTX = make_req_context()
37
def make_cpu_manager(
38
num_blocks: int = 4,
39
cache_policy: str = "lru",
40
+ cache_policy_module_path: str | None = None,
41
enable_events: bool = False,
42
store_threshold: int = 0,
43
max_tracker_size: int = 64_000,
@@ -44,6 +45,7 @@ def make_cpu_manager(
45
return CPUOffloadingManager(
46
num_blocks=num_blocks,
47
cache_policy=cache_policy,
48
+ cache_policy_module_path=cache_policy_module_path,
49
enable_events=enable_events,
50
store_threshold=store_threshold,
51
max_tracker_size=max_tracker_size,
vllm/v1/kv_offload/base.py
-8
@@ -13,8 +13,6 @@ from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, NewType, TypeVar
13
import numpy as np
14
import torch
15
16
-from vllm.logger import init_logger
17
-
16
if TYPE_CHECKING:
17
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
18
OffloadingConnectorStats,
@@ -27,8 +25,6 @@ from vllm.v1.kv_offload.config import OffloadingConfig
25
# Use the helper functions below to construct / decompose keys.
26
OffloadKey = NewType("OffloadKey", bytes)
27
30
-logger = init_logger(__name__)
31
-
28
29
def make_offload_key(block_hash: bytes, group_idx: int) -> OffloadKey:
30
"""Pack a block hash and group index into an `OffloadKey`."""
@@ -540,10 +536,6 @@ class OffloadingSpec(ABC):
536
return {}
537
538
def __init__(self, config: OffloadingConfig):
543
- logger.warning(
544
- "Initializing OffloadingSpec. This API is experimental and "
545
- "subject to change in the future as we iterate the design."
546
- )
539
self.config = config
540
self.extra_config = config.extra_config
541
self.replicated_layout: bool = False
vllm/v1/kv_offload/cpu/manager.py
+9
-16
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections import OrderedDict
4
from collections.abc import Collection, Iterable
5
-from typing import Literal
5
6
from typing_extensions import override
7
@@ -24,19 +23,15 @@ from vllm.v1.kv_offload.cpu.common import (
23
CPULoadStoreSpec,
24
CPUOffloadingMetrics,
25
)
27
-from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
26
from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
29
-from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
30
-
31
-_CACHE_POLICIES: dict[str, type[CachePolicy]] = {
32
- "lru": LRUCachePolicy,
33
- "arc": ARCCachePolicy,
34
-}
27
+from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory
28
29
30
class CPUOffloadingManager(OffloadingManager):
31
"""
39
- An OffloadingManager with a pluggable CachePolicy (LRU or ARC).
32
+ An OffloadingManager with a pluggable CachePolicy, resolved by name via
33
+ CachePolicyFactory (built in: "lru", "arc"; external policies can either
34
+ register their own or be loaded out-of-tree via cache_policy_module_path).
35
36
The manager owns all shared logic: ref-counting, event emission,
37
block pool management, and the prepare_store/complete_store skeletons.
@@ -47,7 +42,8 @@ class CPUOffloadingManager(OffloadingManager):
42
def __init__(
43
self,
44
num_blocks: int,
50
- cache_policy: Literal["lru", "arc"] = "lru",
45
+ cache_policy: str = "lru",
46
+ cache_policy_module_path: str | None = None,
47
enable_events: bool = False,
48
store_threshold: int = 1,
49
max_tracker_size: int = 64_000,
@@ -57,12 +53,9 @@ class CPUOffloadingManager(OffloadingManager):
53
self._num_allocated_blocks: int = 0
54
self._free_list: list[int] = []
55
self.events: list[OffloadingEvent] | None = [] if enable_events else None
60
- policy_cls = _CACHE_POLICIES.get(cache_policy)
61
- if policy_cls is None:
62
- raise ValueError(
63
- f"Unknown cache policy: {cache_policy!r}. "
64
- f"Supported: {list(_CACHE_POLICIES)}"
65
- )
56
+ policy_cls = CachePolicyFactory.get_cache_policy_cls(
57
+ cache_policy, cache_policy_module_path
58
+ )
59
self._policy: CachePolicy = policy_cls(cache_capacity=num_blocks)
60
# Track the number of blocks in the cache that are evictable. i.e. ref_cnt 0.
61
self._num_evictable_cache_blocks: int = 0
vllm/v1/kv_offload/cpu/policies/arc.py
+1
-1
@@ -48,7 +48,7 @@ class ARCCachePolicy(CachePolicy):
48
"""
49
50
def __init__(self, cache_capacity: int):
51
- self.cache_capacity: int = cache_capacity
51
+ super().__init__(cache_capacity)
52
self.target_t1_size: float = 0.0
53
self.t1: OrderedDict[OffloadKey, BlockStatus] = OrderedDict()
54
self.t2: OrderedDict[OffloadKey, BlockStatus] = OrderedDict()
vllm/v1/kv_offload/cpu/policies/base.py
+2
-2
@@ -41,8 +41,8 @@ class CachePolicy(ABC):
41
and eviction, so they cannot be separated cleanly.
42
"""
43
44
- @abstractmethod
45
- def __init__(self, cache_capacity: int) -> None: ...
44
+ def __init__(self, cache_capacity: int) -> None:
45
+ self.cache_capacity = cache_capacity
46
47
@abstractmethod
48
def get(self, key: OffloadKey) -> BlockStatus | None:
vllm/v1/kv_offload/cpu/policies/factory.py
new
+87
@@ -0,0 +1,87 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import importlib
4
+from collections.abc import Callable
5
+
6
+from vllm.logger import init_logger
7
+from vllm.v1.kv_offload.cpu.policies.base import CachePolicy
8
+
9
+logger = init_logger(__name__)
10
+
11
+
12
+class CachePolicyFactory:
13
+ """Registry for CachePolicy implementations, resolved by name.
14
+
15
+ Mirrors OffloadingSpecFactory (vllm/v1/kv_offload/factory.py): built-in
16
+ policies are pre-registered below. External policies can either
17
+ register_cache_policy() a friendly short name up front, or skip
18
+ registration entirely and pass a module path at lookup time (out-of-tree,
19
+ no vLLM fork/patch required) -- see get_cache_policy_cls.
20
+ """
21
+
22
+ _registry: dict[str, Callable[[], type[CachePolicy]]] = {}
23
+
24
+ @classmethod
25
+ def register_cache_policy(
26
+ cls, name: str, module_path: str, class_name: str
27
+ ) -> None:
28
+ """Register a cache policy with a lazy-loading module and class name."""
29
+ if name in cls._registry:
30
+ raise ValueError(f"Cache policy '{name}' is already registered.")
31
+
32
+ def loader() -> type[CachePolicy]:
33
+ module = importlib.import_module(module_path)
34
+ return getattr(module, class_name)
35
+
36
+ cls._registry[name] = loader
37
+
38
+ @classmethod
39
+ def get_cache_policy_cls(
40
+ cls, name: str, module_path: str | None = None
41
+ ) -> type[CachePolicy]:
42
+ """Get a cache policy class by name.
43
+
44
+ Args:
45
+ name: Name of the cache policy. Checked against the registry
46
+ first; if it's not registered and `module_path` is given,
47
+ `name` is imported from there instead -- an out-of-tree
48
+ policy needs no register_cache_policy() call at all, just
49
+ this module path passed through config (mirrors
50
+ OffloadingSpecFactory.get_spec_cls's spec_module_path
51
+ fallback).
52
+ module_path: Python import path to load `name` from when it is
53
+ not a registered policy.
54
+
55
+ Returns:
56
+ The cache policy class.
57
+
58
+ Raises ValueError if the cache policy is neither registered nor
59
+ resolvable via `module_path`.
60
+ """
61
+ if name in cls._registry:
62
+ return cls._registry[name]()
63
+ if module_path is None:
64
+ raise ValueError(
65
+ f"Unknown cache policy: {name!r}. Supported: {list(cls._registry)}. "
66
+ "For an out-of-tree policy, also set cache_policy_module_path."
67
+ )
68
+ logger.warning(
69
+ "Loading out-of-tree cache policy '%s' from '%s'. This API is "
70
+ "experimental and subject to change in the future as we "
71
+ "iterate the design.",
72
+ name,
73
+ module_path,
74
+ )
75
+ module = importlib.import_module(module_path)
76
+ policy_cls = getattr(module, name)
77
+ assert issubclass(policy_cls, CachePolicy)
78
+ return policy_cls
79
+
80
+
81
+# Register built-in policies here.
82
+CachePolicyFactory.register_cache_policy(
83
+ "lru", "vllm.v1.kv_offload.cpu.policies.lru", "LRUCachePolicy"
84
+)
85
+CachePolicyFactory.register_cache_policy(
86
+ "arc", "vllm.v1.kv_offload.cpu.policies.arc", "ARCCachePolicy"
87
+)
vllm/v1/kv_offload/cpu/policies/lru.py
+1
@@ -19,6 +19,7 @@ class LRUCachePolicy(CachePolicy):
19
"""
20
21
def __init__(self, cache_capacity: int):
22
+ super().__init__(cache_capacity)
23
# Blocks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU
24
self.evictable_blocks: OrderedDict[OffloadKey, None] = OrderedDict()
25
self.blocks: dict[OffloadKey, BlockStatus] = {}
vllm/v1/kv_offload/cpu/spec.py
+5
-1
@@ -117,6 +117,9 @@ class CPUOffloadingSpec(OffloadingSpec):
117
self._worker: CPUOffloadingWorker | None = None
118
119
self.eviction_policy: str = self.extra_config.get("eviction_policy", "lru")
120
+ self.cache_policy_module_path: str | None = self.extra_config.get(
121
+ "cache_policy_module_path"
122
+ )
123
124
@override
125
def get_manager(self) -> OffloadingManager:
@@ -131,7 +134,8 @@ class CPUOffloadingSpec(OffloadingSpec):
134
135
self._manager = CPUOffloadingManager(
136
num_blocks=self.num_blocks,
134
- cache_policy=self.eviction_policy, # type: ignore[arg-type]
137
+ cache_policy=self.eviction_policy,
138
+ cache_policy_module_path=self.cache_policy_module_path,
139
enable_events=self.kv_events_config.enable_kv_cache_events,
140
store_threshold=store_threshold,
141
max_tracker_size=max_tracker_size,
vllm/v1/kv_offload/factory.py
+7
@@ -35,6 +35,13 @@ class OffloadingSpecFactory:
35
spec_module_path = extra_config.get("spec_module_path")
36
if spec_module_path is None:
37
raise ValueError(f"Unsupported spec type: {spec_name}")
38
+ logger.warning(
39
+ "Loading out-of-tree offloading spec '%s' from '%s'. This "
40
+ "API is experimental and subject to change in the future "
41
+ "as we iterate the design.",
42
+ spec_name,
43
+ spec_module_path,
44
+ )
45
spec_module = importlib.import_module(spec_module_path)
46
spec_cls = getattr(spec_module, spec_name)
47
assert issubclass(spec_cls, OffloadingSpec)
vllm/v1/kv_offload/tiering/manager.py
+3
-1
@@ -93,11 +93,13 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager):
93
num_blocks: int,
94
mmap_region: SharedOffloadRegion,
95
cache_policy: str = "lru",
96
+ cache_policy_module_path: str | None = None,
97
enable_events: bool = False,
98
):
99
super().__init__(
100
num_blocks=num_blocks,
100
- cache_policy=cache_policy, # type: ignore[arg-type]
101
+ cache_policy=cache_policy,
102
+ cache_policy_module_path=cache_policy_module_path,
103
enable_events=enable_events,
104
)
105
self._mmap_region = mmap_region
vllm/v1/kv_offload/tiering/spec.py
+9
-3
@@ -9,8 +9,13 @@ and configurable secondary tiers (e.g., Storage, Network).
9
Configuration via kv_connector_extra_config:
10
- cpu_bytes_to_use: (required) Bytes to allocate for CPU primary tier
11
- block_size: (optional) Block size for offloaded blocks (default: GPU block size)
12
- - eviction_policy: (optional) Primary tier eviction policy: "lru" or
13
- "arc" (default: "lru")
12
+ - eviction_policy: (optional) Primary tier eviction policy: built-in "lru"/
13
+ "arc", or the name of a policy registered via CachePolicyFactory, or an
14
+ out-of-tree CachePolicy class name paired with cache_policy_module_path
15
+ (default: "lru")
16
+ - cache_policy_module_path: (optional) Python import path to load
17
+ eviction_policy from when it names an out-of-tree CachePolicy not
18
+ registered via CachePolicyFactory
19
- secondary_tiers: (optional) List of secondary tier configurations
20
Each secondary tier config is a dict with:
21
- type: (required) Type of secondary tier (e.g., "example", "storage", "network")
@@ -178,7 +183,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
183
# Create primary tier (CPU-based)
184
primary_tier = CPUPrimaryTierOffloadingManager(
185
num_blocks=self.num_blocks,
181
- cache_policy=self.eviction_policy, # type: ignore[arg-type]
186
+ cache_policy=self.eviction_policy,
187
+ cache_policy_module_path=self.cache_policy_module_path,
188
enable_events=self.kv_events_config.enable_kv_cache_events,
189
mmap_region=scheduler_mmap,
190
)