iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[KV Offload] Support out-of-tree secondary tier managers via module_path (#51007)

Browse Source 
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
This commit is contained in:
Ronen Schaffer
2026-08-05 20:30:53 +03:00
committed by GitHub
co-authored by Or Ozeri
parent d4da0c55af
commit 8543522ca7
6 changed files with 153 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/features/kv_offloading_usage.md
+21 -1
View File
@@ -38,7 +38,7 @@ vllm serve <model> \
## Multi-Tier Setup
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.
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.
```bash
vllm serve <model> \
@@ -272,6 +272,26 @@ Runtime handshake for a P2P (or P/D) pull, once the orchestrator has set the key
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.
### Out-of-Tree Secondary Tiers
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:
```json
{
"spec_name": "TieringOffloadingSpec",
"cpu_bytes_to_use": 10737418240,
"secondary_tiers": [
{
"type": "MyCustomTier",
"module_path": "my_package.my_module",
"custom_param": "value"
}
]
}
```
`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.
## Tuning Tips
- `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
View File
@@ -133,13 +133,13 @@ def test_missing_tier_type_raises():
def test_unknown_tier_type_raises():
"""Unrecognized tier_type → ValueError with supported types list."""
"""Unrecognized tier_type without module_path → ValueError."""
primary_kv_view, offloading_spec = _make_mock_args()
tier_config = {"type": "nonexistent_tier"}
with pytest.raises(
ValueError,
match=r"Unknown secondary tier type.*Supported types:",
match=r"Unknown secondary tier type.*also set 'module_path'",
):
SecondaryTierFactory.create_secondary_tier(
tier_config, primary_kv_view, offloading_spec
@@ -150,3 +150,52 @@ def test_duplicate_registration_raises():
"""register_tier with existing type → ValueError."""
with pytest.raises(ValueError, match="is already registered"):
SecondaryTierFactory.register_tier("example", "some.module", "SomeClass")
# ---------------------------------------------------------------------------
# Out-of-tree dynamic loading via module_path
# ---------------------------------------------------------------------------
def test_create_tier_from_module_path():
"""Full end-to-end: create_secondary_tier with module_path."""
assert "ExampleSecondaryTierManager" not in SecondaryTierFactory._registry
primary_kv_view, offloading_spec = _make_mock_args()
tier_config = {
"type": "ExampleSecondaryTierManager",
"module_path": "vllm.v1.kv_offload.tiering.example.manager",
"custom_param": 42,
}
tier = SecondaryTierFactory.create_secondary_tier(
tier_config, primary_kv_view, offloading_spec
)
assert isinstance(tier, ExampleSecondaryTierManager)
assert tier.tier_type == "ExampleSecondaryTierManager"
def test_module_path_invalid_module_raises():
"""Non-existent module_path → ModuleNotFoundError."""
primary_kv_view, offloading_spec = _make_mock_args()
tier_config = {
"type": "SomeTier",
"module_path": "nonexistent.module.path",
}
with pytest.raises(ModuleNotFoundError):
SecondaryTierFactory.create_secondary_tier(
tier_config, primary_kv_view, offloading_spec
)
def test_module_path_not_subclass_raises():
"""module_path resolving to a non-SecondaryTierManager → AssertionError."""
tier_config = {
"type": "MagicMock",
"module_path": "unittest.mock",
}
with pytest.raises(AssertionError):
SecondaryTierFactory.get_tier_class(tier_config)
vllm/v1/kv_offload/cpu/policies/factory.py
+7 -4
View File
@@ -65,10 +65,13 @@ class CachePolicyFactory:
f"Unknown cache policy: {name!r}. Supported: {list(cls._registry)}. "
"For an out-of-tree policy, also set cache_policy_module_path."
)
logger.warning(
"Loading out-of-tree cache policy '%s' from '%s'. This API is "
"experimental and subject to change in the future as we "
"iterate the design.",
logger.warning_once(
"Loading out-of-tree cache policy. This API is "
"experimental and subject to change in the future "
"as we iterate the design."
)
logger.info(
"Loading out-of-tree cache policy '%s' from '%s'.",
name,
module_path,
)
vllm/v1/kv_offload/factory.py
+7 -4
View File
@@ -35,10 +35,13 @@ class OffloadingSpecFactory:
spec_module_path = extra_config.get("spec_module_path")
if spec_module_path is None:
raise ValueError(f"Unsupported spec type: {spec_name}")
logger.warning(
"Loading out-of-tree offloading spec '%s' from '%s'. This "
"API is experimental and subject to change in the future "
"as we iterate the design.",
logger.warning_once(
"Loading out-of-tree offloading spec. This API is "
"experimental and subject to change in the future "
"as we iterate the design."
)
logger.info(
"Loading out-of-tree offloading spec '%s' from '%s'.",
spec_name,
spec_module_path,
)
vllm/v1/kv_offload/tiering/factory.py
+49 -3
View File
@@ -4,13 +4,25 @@ import importlib
from collections.abc import Callable
from typing import TYPE_CHECKING
from vllm.logger import init_logger
from vllm.v1.kv_offload.tiering.base import SecondaryTierManager
if TYPE_CHECKING:
from vllm.v1.kv_offload.base import OffloadingSpec
logger = init_logger(__name__)
class SecondaryTierFactory:
"""Registry for SecondaryTierManager implementations, resolved by type.
Mirrors OffloadingSpecFactory (vllm/v1/kv_offload/factory.py): built-in
tiers are pre-registered below. External tiers can either
register_tier() a friendly short name up front, or skip registration
entirely and pass a module_path at lookup time (out-of-tree, no vLLM
fork/patch required) -- see get_tier_class.
"""
_registry: dict[str, Callable[[], type[SecondaryTierManager]]] = {}
@classmethod
@@ -34,6 +46,7 @@ class SecondaryTierFactory:
tier_cls = cls.get_tier_class(tier_config)
config = tier_config.copy()
tier_type = config.pop("type")
config.pop("module_path", None)
return tier_cls(
offloading_spec=offloading_spec,
primary_kv_view=primary_kv_view,
@@ -43,15 +56,48 @@ class SecondaryTierFactory:
@classmethod
def get_tier_class(cls, tier_config: dict) -> type[SecondaryTierManager]:
"""Get a secondary tier class by type.
Args:
tier_config: Tier configuration dict. Must contain 'type' (the
tier name or out-of-tree class name). If 'type' is not a
registered tier and 'module_path' is given, the class is
imported from there -- an out-of-tree tier needs no
register_tier() call, just this module path in the config.
Returns:
The secondary tier manager class.
Raises:
ValueError: If 'type' is missing, or is neither registered nor
resolvable via module_path.
"""
tier_type = tier_config.get("type")
if not tier_type:
raise ValueError("Secondary tier configuration must include 'type'")
if tier_type not in cls._registry:
if tier_type in cls._registry:
return cls._registry[tier_type]()
module_path = tier_config.get("module_path")
if module_path is None:
raise ValueError(
f"Unknown secondary tier type: {tier_type!r}. "
f"Supported types: {list(cls._registry)}"
f"Supported types: {list(cls._registry)}. "
"For an out-of-tree tier, also set 'module_path'."
)
return cls._registry[tier_type]()
logger.warning_once(
"Loading out-of-tree secondary tier. This API is "
"experimental and subject to change in the future "
"as we iterate the design."
)
logger.info(
"Loading out-of-tree secondary tier '%s' from '%s'.",
tier_type,
module_path,
)
module = importlib.import_module(module_path)
tier_cls = getattr(module, tier_type)
assert issubclass(tier_cls, SecondaryTierManager)
return tier_cls
SecondaryTierFactory.register_tier(
vllm/v1/kv_offload/tiering/spec.py
+18 -1
View File
@@ -18,7 +18,12 @@ Configuration via kv_connector_extra_config:
registered via CachePolicyFactory
- secondary_tiers: (optional) List of secondary tier configurations
Each secondary tier config is a dict with:
- type: (required) Type of secondary tier (e.g., "example", "storage", "network")
- type: (required) Type of secondary tier (e.g., "example", "fs",
"p2p", "obj"), or the class name of an out-of-tree
SecondaryTierManager paired with module_path.
- module_path: (optional) Python import path to load 'type' from
when it names an out-of-tree SecondaryTierManager not registered
via SecondaryTierFactory.register_tier()
- Additional tier-specific parameters are passed directly to the tier
constructor. See each tier's documentation for supported parameters.
@@ -34,6 +39,18 @@ Example configuration:
}
]
}
Example out-of-tree tier configuration:
{
"cpu_bytes_to_use": 10737418240,
"secondary_tiers": [
{
"type": "MyCustomTier",
"module_path": "my_package.my_module",
"custom_param": "value"
}
]
}
"""
from typing import Any