| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import importlib |
| 5 | from collections.abc import Callable |
| 6 | from typing import TYPE_CHECKING, cast |
| 7 | |
| 8 | from vllm.config.kv_transfer import KVTransferConfig |
| 9 | from vllm.distributed.kv_transfer.kv_connector.base import ( |
| 10 | KVConnectorBase, |
| 11 | KVConnectorBaseType, |
| 12 | ) |
| 13 | from vllm.distributed.kv_transfer.kv_connector.v1 import ( |
| 14 | KVConnectorRole, |
| 15 | supports_hma, |
| 16 | ) |
| 17 | from vllm.logger import init_logger |
| 18 | from vllm.utils.func_utils import supports_kw |
| 19 | |
| 20 | if TYPE_CHECKING: |
| 21 | from vllm.config import VllmConfig |
| 22 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 23 | |
| 24 | logger = init_logger(__name__) |
| 25 | |
| 26 | |
| 27 | class KVConnectorFactory: |
| 28 | _registry: dict[str, Callable[[], type[KVConnectorBase]]] = {} |
| 29 | |
| 30 | @classmethod |
| 31 | def register_connector(cls, name: str, module_path: str, class_name: str) -> None: |
| 32 | """Register a connector with a lazy-loading module and class name.""" |
| 33 | if name in cls._registry: |
| 34 | raise ValueError(f"Connector '{name}' is already registered.") |
| 35 | |
| 36 | def loader() -> type[KVConnectorBase]: |
| 37 | module = importlib.import_module(module_path) |
| 38 | return getattr(module, class_name) |
| 39 | |
| 40 | cls._registry[name] = loader |
| 41 | |
| 42 | @classmethod |
| 43 | def create_connector( |
| 44 | cls, |
| 45 | config: "VllmConfig", |
| 46 | role: KVConnectorRole, |
| 47 | kv_cache_config: "KVCacheConfig", |
| 48 | ) -> KVConnectorBase: |
| 49 | kv_transfer_config = config.kv_transfer_config |
| 50 | if kv_transfer_config is None: |
| 51 | raise ValueError("kv_transfer_config must be set to create a connector") |
| 52 | connector_cls = cls.get_connector_class(kv_transfer_config) |
| 53 | |
| 54 | # check if the connector supports HMA |
| 55 | hma_enabled = not config.scheduler_config.disable_hybrid_kv_cache_manager |
| 56 | if hma_enabled and not cls.supports_hma_config(kv_transfer_config): |
| 57 | raise ValueError( |
| 58 | f"Connector {connector_cls.__name__} does not support HMA but " |
| 59 | f"HMA is enabled. Please set `--disable-hybrid-kv-cache-manager`." |
| 60 | ) |
| 61 | |
| 62 | logger.info( |
| 63 | "Creating v1 connector with name: %s and engine_id: %s", |
| 64 | connector_cls.__name__, |
| 65 | kv_transfer_config.engine_id, |
| 66 | ) |
| 67 | # NOTE(Kuntai): v1 connector is explicitly separated into two roles. |
| 68 | # Scheduler connector: |
| 69 | # - Co-locate with scheduler process |
| 70 | # - Should only be used inside the Scheduler class |
| 71 | # Worker connector: |
| 72 | # - Co-locate with worker process |
| 73 | # - Should only be used inside the forward context & attention layer |
| 74 | # We build separately to enforce strict separation |
| 75 | return connector_cls(config, role, kv_cache_config) |
| 76 | |
| 77 | @classmethod |
| 78 | def get_connector_class_by_name( |
| 79 | cls, connector_name: str |
| 80 | ) -> type[KVConnectorBaseType]: |
| 81 | """Get a registered connector class by name. |
| 82 | |
| 83 | Raises ValueError if the connector is not registered. |
| 84 | |
| 85 | Args: |
| 86 | connector_name: Name of the registered connector. |
| 87 | |
| 88 | Returns: |
| 89 | The connector class. |
| 90 | |
| 91 | """ |
| 92 | if connector_name not in cls._registry: |
| 93 | raise ValueError(f"Connector '{connector_name}' is not registered.") |
| 94 | return cls._registry[connector_name]() |
| 95 | |
| 96 | @classmethod |
| 97 | def get_connector_class( |
| 98 | cls, kv_transfer_config: "KVTransferConfig" |
| 99 | ) -> type[KVConnectorBaseType]: |
| 100 | connector_name = kv_transfer_config.kv_connector |
| 101 | if connector_name is None: |
| 102 | raise ValueError("Connector name is not set in KVTransferConfig") |
| 103 | connector_module_path = kv_transfer_config.kv_connector_module_path |
| 104 | if connector_module_path is not None and not connector_module_path: |
| 105 | raise ValueError("kv_connector_module_path cannot be an empty string.") |
| 106 | if connector_module_path: |
| 107 | # External module path takes priority over internal registry. |
| 108 | connector_module = importlib.import_module(connector_module_path) |
| 109 | try: |
| 110 | connector_cls = getattr(connector_module, connector_name) |
| 111 | except AttributeError as e: |
| 112 | raise AttributeError( |
| 113 | f"Class {connector_name} not found in {connector_module_path}" |
| 114 | ) from e |
| 115 | connector_cls = cast(type[KVConnectorBaseType], connector_cls) |
| 116 | if not supports_kw(connector_cls, "kv_cache_config"): |
| 117 | msg = ( |
| 118 | f"Connector {connector_cls.__name__} uses deprecated " |
| 119 | "2-argument constructor signature. External v1 KV " |
| 120 | "connectors must accept kv_cache_config as the third " |
| 121 | "constructor argument and pass it to super().__init__()." |
| 122 | ) |
| 123 | logger.error(msg) |
| 124 | raise ValueError(msg) |
| 125 | elif connector_name in cls._registry: |
| 126 | connector_cls = cls._registry[connector_name]() |
| 127 | else: |
| 128 | raise ValueError(f"Unsupported connector type: {connector_name}") |
| 129 | return connector_cls |
| 130 | |
| 131 | @classmethod |
| 132 | def supports_hma_config(cls, kv_transfer_config: "KVTransferConfig") -> bool: |
| 133 | """Return whether this KV transfer config supports HMA. |
| 134 | |
| 135 | MultiConnector is a special case: the wrapper class implements |
| 136 | SupportsHMA, but effective support depends on every configured child. |
| 137 | """ |
| 138 | connector_cls = cls.get_connector_class(kv_transfer_config) |
| 139 | if kv_transfer_config.kv_connector != "MultiConnector": |
| 140 | return supports_hma(connector_cls) |
| 141 | |
| 142 | from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( |
| 143 | MultiConnector, |
| 144 | ) |
| 145 | |
| 146 | return MultiConnector.all_children_support_hma(kv_transfer_config) |
| 147 | |
| 148 | |
| 149 | # Register various connectors here. |
| 150 | # The registration should not be done in each individual file, as we want to |
| 151 | # only load the files corresponding to the current connector. |
| 152 | |
| 153 | KVConnectorFactory.register_connector( |
| 154 | "ExampleConnector", |
| 155 | "vllm.distributed.kv_transfer.kv_connector.v1.example_connector", |
| 156 | "ExampleConnector", |
| 157 | ) |
| 158 | |
| 159 | KVConnectorFactory.register_connector( |
| 160 | "ExampleHiddenStatesConnector", |
| 161 | "vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector", |
| 162 | "ExampleHiddenStatesConnector", |
| 163 | ) |
| 164 | |
| 165 | KVConnectorFactory.register_connector( |
| 166 | "LMCacheConnectorV1", |
| 167 | "vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector", |
| 168 | "LMCacheConnectorV1", |
| 169 | ) |
| 170 | |
| 171 | KVConnectorFactory.register_connector( |
| 172 | "LMCacheMPConnector", |
| 173 | "vllm.distributed.kv_transfer.kv_connector.v1.lmcache_mp_connector", |
| 174 | "LMCacheMPConnector", |
| 175 | ) |
| 176 | |
| 177 | KVConnectorFactory.register_connector( |
| 178 | "NixlConnector", |
| 179 | "vllm.distributed.kv_transfer.kv_connector.v1.nixl", |
| 180 | "NixlConnector", |
| 181 | ) |
| 182 | |
| 183 | KVConnectorFactory.register_connector( |
| 184 | "NixlPullConnector", |
| 185 | "vllm.distributed.kv_transfer.kv_connector.v1.nixl", |
| 186 | "NixlPullConnector", |
| 187 | ) |
| 188 | |
| 189 | KVConnectorFactory.register_connector( |
| 190 | "NixlPushConnector", |
| 191 | "vllm.distributed.kv_transfer.kv_connector.v1.nixl", |
| 192 | "NixlPushConnector", |
| 193 | ) |
| 194 | |
| 195 | KVConnectorFactory.register_connector( |
| 196 | "MultiConnector", |
| 197 | "vllm.distributed.kv_transfer.kv_connector.v1.multi_connector", |
| 198 | "MultiConnector", |
| 199 | ) |
| 200 | |
| 201 | KVConnectorFactory.register_connector( |
| 202 | "HiSparseConnector", |
| 203 | "vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector", |
| 204 | "HiSparseConnector", |
| 205 | ) |
| 206 | |
| 207 | KVConnectorFactory.register_connector( |
| 208 | "MoRIIOConnector", |
| 209 | "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector", |
| 210 | "MoRIIOConnector", |
| 211 | ) |
| 212 | |
| 213 | KVConnectorFactory.register_connector( |
| 214 | "OffloadingConnector", |
| 215 | "vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector", |
| 216 | "OffloadingConnector", |
| 217 | ) |
| 218 | |
| 219 | KVConnectorFactory.register_connector( |
| 220 | "DecodeBenchConnector", |
| 221 | "vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector", |
| 222 | "DecodeBenchConnector", |
| 223 | ) |
| 224 | |
| 225 | KVConnectorFactory.register_connector( |
| 226 | "MooncakeConnector", |
| 227 | "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector", |
| 228 | "MooncakeConnector", |
| 229 | ) |
| 230 | KVConnectorFactory.register_connector( |
| 231 | "MooncakeStoreConnector", |
| 232 | "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.connector", |
| 233 | "MooncakeStoreConnector", |
| 234 | ) |
| 235 | KVConnectorFactory.register_connector( |
| 236 | "FlexKVConnectorV1", |
| 237 | "vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector", |
| 238 | "FlexKVConnectorV1", |
| 239 | ) |
| 240 | KVConnectorFactory.register_connector( |
| 241 | "SimpleCPUOffloadConnector", |
| 242 | "vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector", |
| 243 | "SimpleCPUOffloadConnector", |
| 244 | ) |
| 245 | KVConnectorFactory.register_connector( |
| 246 | "HF3FSKVConnector", |
| 247 | "vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_connector", |
| 248 | "HF3FSKVConnector", |
| 249 | ) |