[KV Offload] Support out-of-tree secondary tier managers via `module_path` (#51007)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Ronen Schaffer committed
Aug 5, 2026 at 20:30 UTC
8543522ca792de824c026e1b9e3eb51ca809550d
6 files changed
+153
-15
docs/features/kv_offloading_usage.md
+21
-1
@@ -38,7 +38,7 @@ vllm serve <model> \
38
39
## Multi-Tier Setup
40
41
-Set `spec_name` to `"TieringOffloadingSpec"` and supply a `secondary_tiers` list. Each entry is a dict with a required `type` key plus tier-specific fields. The list is ordered: tier 0 is consulted before tier 1, and so on. See [Secondary Tiers](#secondary-tiers) for tier-specific keys.
41
+Set `spec_name` to `"TieringOffloadingSpec"` and supply a `secondary_tiers` list. Each entry is a dict with a required `type` key plus tier-specific fields (and an optional `module_path` for out-of-tree tiers). The list is ordered: tier 0 is consulted before tier 1, and so on. See [Secondary Tiers](#secondary-tiers) for tier-specific keys.
42
43
```bash
44
vllm serve <model> \
@@ -272,6 +272,26 @@ Runtime handshake for a P2P (or P/D) pull, once the orchestrator has set the key
272
273
In classic **P/D mode** (`remote_prefiller` set, no `remote_kv_source`), the lookup phase (steps 2–4) is skipped: the decode consumer assumes the prefiller holds all of the request's blocks, so every block `lookup()` returns an immediate hit and the consumer jumps straight to the **`FetchMsg`** in step 5. The `LookupMsg`/`LookupRespMsg` round-trip only happens in P2P mode, where the consumer does not know in advance which blocks the peer has cached.
274
275
+### Out-of-Tree Secondary Tiers
276
+
277
+Implement `SecondaryTierManager` (`vllm/v1/kv_offload/tiering/base.py`) in your own package — no vLLM fork or patch required — and point the tier config at it directly:
278
+
279
+```json
280
+{
281
+ "spec_name": "TieringOffloadingSpec",
282
+ "cpu_bytes_to_use": 10737418240,
283
+ "secondary_tiers": [
284
+ {
285
+ "type": "MyCustomTier",
286
+ "module_path": "my_package.my_module",
287
+ "custom_param": "value"
288
+ }
289
+ ]
290
+}
291
+```
292
+
293
+`type` is checked against the built-in registry first; if it isn't a registered name, vLLM imports `module_path` and looks up `type` as a class name in that module.
294
+
295
## Tuning Tips
296
297
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
tests/v1/kv_offload/tiering/test_factory.py
+51
-2
@@ -133,13 +133,13 @@ def test_missing_tier_type_raises():
133
134
135
def test_unknown_tier_type_raises():
136
- """Unrecognized tier_type → ValueError with supported types list."""
136
+ """Unrecognized tier_type without module_path → ValueError."""
137
primary_kv_view, offloading_spec = _make_mock_args()
138
tier_config = {"type": "nonexistent_tier"}
139
140
with pytest.raises(
141
ValueError,
142
- match=r"Unknown secondary tier type.*Supported types:",
142
+ match=r"Unknown secondary tier type.*also set 'module_path'",
143
):
144
SecondaryTierFactory.create_secondary_tier(
145
tier_config, primary_kv_view, offloading_spec
@@ -150,3 +150,52 @@ def test_duplicate_registration_raises():
150
"""register_tier with existing type → ValueError."""
151
with pytest.raises(ValueError, match="is already registered"):
152
SecondaryTierFactory.register_tier("example", "some.module", "SomeClass")
153
+
154
+
155
+# ---------------------------------------------------------------------------
156
+# Out-of-tree dynamic loading via module_path
157
+# ---------------------------------------------------------------------------
158
+
159
+
160
+def test_create_tier_from_module_path():
161
+ """Full end-to-end: create_secondary_tier with module_path."""
162
+ assert "ExampleSecondaryTierManager" not in SecondaryTierFactory._registry
163
+
164
+ primary_kv_view, offloading_spec = _make_mock_args()
165
+ tier_config = {
166
+ "type": "ExampleSecondaryTierManager",
167
+ "module_path": "vllm.v1.kv_offload.tiering.example.manager",
168
+ "custom_param": 42,
169
+ }
170
+
171
+ tier = SecondaryTierFactory.create_secondary_tier(
172
+ tier_config, primary_kv_view, offloading_spec
173
+ )
174
+
175
+ assert isinstance(tier, ExampleSecondaryTierManager)
176
+ assert tier.tier_type == "ExampleSecondaryTierManager"
177
+
178
+
179
+def test_module_path_invalid_module_raises():
180
+ """Non-existent module_path → ModuleNotFoundError."""
181
+ primary_kv_view, offloading_spec = _make_mock_args()
182
+ tier_config = {
183
+ "type": "SomeTier",
184
+ "module_path": "nonexistent.module.path",
185
+ }
186
+
187
+ with pytest.raises(ModuleNotFoundError):
188
+ SecondaryTierFactory.create_secondary_tier(
189
+ tier_config, primary_kv_view, offloading_spec
190
+ )
191
+
192
+
193
+def test_module_path_not_subclass_raises():
194
+ """module_path resolving to a non-SecondaryTierManager → AssertionError."""
195
+ tier_config = {
196
+ "type": "MagicMock",
197
+ "module_path": "unittest.mock",
198
+ }
199
+
200
+ with pytest.raises(AssertionError):
201
+ SecondaryTierFactory.get_tier_class(tier_config)
vllm/v1/kv_offload/cpu/policies/factory.py
+7
-4
@@ -65,10 +65,13 @@ class CachePolicyFactory:
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.",
68
+ logger.warning_once(
69
+ "Loading out-of-tree cache policy. This API is "
70
+ "experimental and subject to change in the future "
71
+ "as we iterate the design."
72
+ )
73
+ logger.info(
74
+ "Loading out-of-tree cache policy '%s' from '%s'.",
75
name,
76
module_path,
77
)
vllm/v1/kv_offload/factory.py
+7
-4
@@ -35,10 +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.",
38
+ logger.warning_once(
39
+ "Loading out-of-tree offloading spec. This API is "
40
+ "experimental and subject to change in the future "
41
+ "as we iterate the design."
42
+ )
43
+ logger.info(
44
+ "Loading out-of-tree offloading spec '%s' from '%s'.",
45
spec_name,
46
spec_module_path,
47
)
vllm/v1/kv_offload/tiering/factory.py
+49
-3
@@ -4,13 +4,25 @@ import importlib
4
from collections.abc import Callable
5
from typing import TYPE_CHECKING
6
7
+from vllm.logger import init_logger
8
from vllm.v1.kv_offload.tiering.base import SecondaryTierManager
9
10
if TYPE_CHECKING:
11
from vllm.v1.kv_offload.base import OffloadingSpec
12
13
+logger = init_logger(__name__)
14
+
15
16
class SecondaryTierFactory:
17
+ """Registry for SecondaryTierManager implementations, resolved by type.
18
+
19
+ Mirrors OffloadingSpecFactory (vllm/v1/kv_offload/factory.py): built-in
20
+ tiers are pre-registered below. External tiers can either
21
+ register_tier() a friendly short name up front, or skip registration
22
+ entirely and pass a module_path at lookup time (out-of-tree, no vLLM
23
+ fork/patch required) -- see get_tier_class.
24
+ """
25
+
26
_registry: dict[str, Callable[[], type[SecondaryTierManager]]] = {}
27
28
@classmethod
@@ -34,6 +46,7 @@ class SecondaryTierFactory:
46
tier_cls = cls.get_tier_class(tier_config)
47
config = tier_config.copy()
48
tier_type = config.pop("type")
49
+ config.pop("module_path", None)
50
return tier_cls(
51
offloading_spec=offloading_spec,
52
primary_kv_view=primary_kv_view,
@@ -43,15 +56,48 @@ class SecondaryTierFactory:
56
57
@classmethod
58
def get_tier_class(cls, tier_config: dict) -> type[SecondaryTierManager]:
59
+ """Get a secondary tier class by type.
60
+
61
+ Args:
62
+ tier_config: Tier configuration dict. Must contain 'type' (the
63
+ tier name or out-of-tree class name). If 'type' is not a
64
+ registered tier and 'module_path' is given, the class is
65
+ imported from there -- an out-of-tree tier needs no
66
+ register_tier() call, just this module path in the config.
67
+
68
+ Returns:
69
+ The secondary tier manager class.
70
+
71
+ Raises:
72
+ ValueError: If 'type' is missing, or is neither registered nor
73
+ resolvable via module_path.
74
+ """
75
tier_type = tier_config.get("type")
76
if not tier_type:
77
raise ValueError("Secondary tier configuration must include 'type'")
49
- if tier_type not in cls._registry:
78
+ if tier_type in cls._registry:
79
+ return cls._registry[tier_type]()
80
+ module_path = tier_config.get("module_path")
81
+ if module_path is None:
82
raise ValueError(
83
f"Unknown secondary tier type: {tier_type!r}. "
52
- f"Supported types: {list(cls._registry)}"
84
+ f"Supported types: {list(cls._registry)}. "
85
+ "For an out-of-tree tier, also set 'module_path'."
86
)
54
- return cls._registry[tier_type]()
87
+ logger.warning_once(
88
+ "Loading out-of-tree secondary tier. This API is "
89
+ "experimental and subject to change in the future "
90
+ "as we iterate the design."
91
+ )
92
+ logger.info(
93
+ "Loading out-of-tree secondary tier '%s' from '%s'.",
94
+ tier_type,
95
+ module_path,
96
+ )
97
+ module = importlib.import_module(module_path)
98
+ tier_cls = getattr(module, tier_type)
99
+ assert issubclass(tier_cls, SecondaryTierManager)
100
+ return tier_cls
101
102
103
SecondaryTierFactory.register_tier(
vllm/v1/kv_offload/tiering/spec.py
+18
-1
@@ -18,7 +18,12 @@ Configuration via kv_connector_extra_config:
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")
21
+ - type: (required) Type of secondary tier (e.g., "example", "fs",
22
+ "p2p", "obj"), or the class name of an out-of-tree
23
+ SecondaryTierManager paired with module_path.
24
+ - module_path: (optional) Python import path to load 'type' from
25
+ when it names an out-of-tree SecondaryTierManager not registered
26
+ via SecondaryTierFactory.register_tier()
27
- Additional tier-specific parameters are passed directly to the tier
28
constructor. See each tier's documentation for supported parameters.
29
@@ -34,6 +39,18 @@ Example configuration:
39
}
40
]
41
}
42
+
43
+Example out-of-tree tier configuration:
44
+{
45
+ "cpu_bytes_to_use": 10737418240,
46
+ "secondary_tiers": [
47
+ {
48
+ "type": "MyCustomTier",
49
+ "module_path": "my_package.my_module",
50
+ "custom_param": "value"
51
+ }
52
+ ]
53
+}
54
"""
55
56
from typing import Any