From 3e0d435027f7585afbbc8b44ca1e386c0a88ae47 Mon Sep 17 00:00:00 2001 From: David Ben-David Date: Sat, 17 May 2025 09:40:39 +0300 Subject: [PATCH] [P/D][V1] Support dynamic loading of external KV connector implementations (#18142) Signed-off-by: David Ben-David Co-authored-by: David Ben-David --- vllm/config.py | 4 ++++ .../distributed/kv_transfer/kv_connector/factory.py | 13 +++++++++++-- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/vllm/config.py b/vllm/config.py index d07a1ff052348..3a1598f574262 100644 --- a/vllm/config.py +++ b/vllm/config.py @@ -3524,6 +3524,10 @@ class KVTransferConfig: kv_connector_extra_config: dict[str, Any] = field(default_factory=dict) """any extra config that the connector may need.""" + kv_connector_module_path: Optional[str] = None + """The Python module path to dynamically load the KV connector from. + Only supported in V1.""" + def compute_hash(self) -> str: """ WARNING: Whenever a new field is added to this config, diff --git a/vllm/distributed/kv_transfer/kv_connector/factory.py b/vllm/distributed/kv_transfer/kv_connector/factory.py index f998f5dd7b158..06b3983ed68bd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/factory.py +++ b/vllm/distributed/kv_transfer/kv_connector/factory.py @@ -58,8 +58,17 @@ class KVConnectorFactory: raise ValueError("Attempting to initialize a V1 Connector, " f"but found {envs.VLLM_USE_V1=}") - connector_name = config.kv_transfer_config.kv_connector - connector_cls = cls._registry[connector_name]() + kv_transfer_config = config.kv_transfer_config + connector_name = kv_transfer_config.kv_connector + if connector_name in cls._registry: + connector_cls = cls._registry[connector_name]() + else: + connector_module_path = kv_transfer_config.kv_connector_module_path + if connector_module_path is None: + raise ValueError( + f"Unsupported connector type: {connector_name}") + connector_module = importlib.import_module(connector_module_path) + connector_cls = getattr(connector_module, connector_name) assert issubclass(connector_cls, KVConnectorBase_V1) logger.info("Creating v1 connector with name: %s", connector_name) # NOTE(Kuntai): v1 connector is explicitly separated into two roles.