From 8543522ca792de824c026e1b9e3eb51ca809550d Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 5 Aug 2026 20:30:53 +0300 Subject: [PATCH] [KV Offload] Support out-of-tree secondary tier managers via `module_path` (#51007) Signed-off-by: Ronen Schaffer Co-authored-by: Or Ozeri --- docs/features/kv_offloading_usage.md | 22 ++++++++- tests/v1/kv_offload/tiering/test_factory.py | 53 ++++++++++++++++++++- vllm/v1/kv_offload/cpu/policies/factory.py | 11 +++-- vllm/v1/kv_offload/factory.py | 11 +++-- vllm/v1/kv_offload/tiering/factory.py | 52 ++++++++++++++++++-- vllm/v1/kv_offload/tiering/spec.py | 19 +++++++- 6 files changed, 153 insertions(+), 15 deletions(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 72f838c7d2..0167b45930 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -38,7 +38,7 @@ vllm serve \ ## 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 \ @@ -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. diff --git a/tests/v1/kv_offload/tiering/test_factory.py b/tests/v1/kv_offload/tiering/test_factory.py index 2b65baa020..f5a576d254 100644 --- a/tests/v1/kv_offload/tiering/test_factory.py +++ b/tests/v1/kv_offload/tiering/test_factory.py @@ -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) diff --git a/vllm/v1/kv_offload/cpu/policies/factory.py b/vllm/v1/kv_offload/cpu/policies/factory.py index 88d11b270c..eedcc68076 100644 --- a/vllm/v1/kv_offload/cpu/policies/factory.py +++ b/vllm/v1/kv_offload/cpu/policies/factory.py @@ -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, ) diff --git a/vllm/v1/kv_offload/factory.py b/vllm/v1/kv_offload/factory.py index 19bc401277..33584e5bba 100644 --- a/vllm/v1/kv_offload/factory.py +++ b/vllm/v1/kv_offload/factory.py @@ -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, ) diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index 180c87d094..17f8430204 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -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( diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 7d3d84a622..fb2cdacb99 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -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