updated

Signed-off-by: [email protected] <[email protected]>

Author: robertgshaw2-redhat

vllm/distributed/kv_transfer_state.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import TYPE_CHECKING, Union
+
+from vllm import envs
+from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
+from vllm.distributed.kv_transfer.kv_connector.factory import (
+    KVConnectorFactory)
+from vllm.distributed.kv_transfer.kv_connector.v1 import (KVConnectorBase_V1,
+                                                          KVConnectorRole)
+from vllm.distributed.parallel_state import get_world_group
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+_KV_CONNECTOR_AGENT: Union[KVConnectorBase, KVConnectorBase_V1, None] = None
+
+
+def get_kv_transfer_group() -> Union[KVConnectorBase, KVConnectorBase_V1]:
+    assert _KV_CONNECTOR_AGENT is not None, (
+        "disaggregated KV cache transfer parallel group is not initialized")
+    return _KV_CONNECTOR_AGENT
+
+
+def has_kv_transfer_group() -> bool:
+    return _KV_CONNECTOR_AGENT is not None
+
+
+def is_v1_kv_transfer_group(
+    connector: Union[KVConnectorBase_V1, KVConnectorBase,
+                     None] = None) -> bool:
+    """Check if the KV connector is the v1 connector.
+    If the argument is None, it will check the global KV connector
+
+    Args:
+        connector: The KV connector to check. If None, it will check the
+            global KV connector.
+
+    Note:
+        This function will no-longer be needed after the v1 KV connector
+        becomes the default.
+    """
+    if connector is None:
+        connector = _KV_CONNECTOR_AGENT
+
+    if connector is None:
+        return False
+
+    return isinstance(connector, KVConnectorBase_V1)
+
+
+def ensure_kv_transfer_initialized(vllm_config: "VllmConfig") -> None:
+    """
+    Initialize KV cache transfer parallel group.
+    """
+
+    global _KV_CONNECTOR_AGENT
+
+    if vllm_config.kv_transfer_config is None:
+        return
+
+    if all([
+            vllm_config.kv_transfer_config.is_kv_transfer_instance,
+            _KV_CONNECTOR_AGENT is None
+    ]):
+
+        if envs.VLLM_USE_V1:
+            _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector_v1(
+                config=vllm_config, role=KVConnectorRole.WORKER)
+        else:
+            _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector_v0(
+                rank=get_world_group().rank,
+                local_rank=get_world_group().local_rank,
+                config=vllm_config,
+            )

vllm/distributed/parallel_state.py

@@ -40,15 +40,13 @@
 import vllm.envs as envs
 from vllm.distributed.device_communicators.base_device_communicator import (
     DeviceCommunicatorBase)
-from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
-from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
 from vllm.distributed.utils import StatelessProcessGroup
 from vllm.logger import init_logger
 from vllm.utils import (direct_register_custom_op, resolve_obj_by_qualname,
                         supports_custom_op)
 
 if TYPE_CHECKING:
-    from vllm.config import VllmConfig
+    pass
 
 
 @dataclass
@@ -771,44 +769,6 @@ def get_pp_group() -> GroupCoordinator:
 # kept for backward compatibility
 get_pipeline_model_parallel_group = get_pp_group
 
-# TODO: once we deprecate V0 KV transfer, we can move this to
-# be a non-global object.
-_KV_CONNECTOR_AGENT: Union[KVConnectorBase, KVConnectorBase_V1, None] = None
-
-
-def get_kv_transfer_group() -> Union[KVConnectorBase, KVConnectorBase_V1]:
-    assert _KV_CONNECTOR_AGENT is not None, (
-        "disaggregated KV cache transfer parallel group is not initialized")
-    return _KV_CONNECTOR_AGENT
-
-
-def has_kv_transfer_group() -> bool:
-    return _KV_CONNECTOR_AGENT is not None
-
-
-def is_v1_kv_transfer_group(
-    connector: Union[KVConnectorBase_V1, KVConnectorBase,
-                     None] = None) -> bool:
-    """Check if the KV connector is the v1 connector.
-    If the argument is None, it will check the global KV connector
-
-    Args:
-        connector: The KV connector to check. If None, it will check the
-            global KV connector.
-
-    Note:
-        This function will no-longer be needed after the v1 KV connector
-        becomes the default.
-    """
-    if connector is None:
-        connector = _KV_CONNECTOR_AGENT
-
-    if connector is None:
-        # Global KV connector is not set
-        return False
-
-    return isinstance(connector, KVConnectorBase_V1)
-
 
 @contextmanager
 def graph_capture(device: torch.device):
@@ -991,37 +951,6 @@ def initialize_model_parallel(
         _DP.rank_in_group, _PP.rank_in_group, _TP.rank_in_group)
 
 
-def ensure_kv_transfer_initialized(vllm_config: "VllmConfig") -> None:
-    """
-    Initialize KV cache transfer parallel group.
-    """
-
-    global _KV_CONNECTOR_AGENT
-
-    if vllm_config.kv_transfer_config is None:
-        return
-
-    if all([
-            vllm_config.kv_transfer_config.is_kv_transfer_instance,
-            _KV_CONNECTOR_AGENT is None
-    ]):
-        from vllm.distributed.kv_transfer.kv_connector.factory import (
-            KVConnectorFactory)
-        from vllm.distributed.kv_transfer.kv_connector.v1 import (
-            KVConnectorRole as KVConnectorRole_V1)
-
-        kwargs = {
-            "rank": get_world_group().rank,
-            "local_rank": get_world_group().local_rank,
-            "config": vllm_config,
-            # NOTE(Kuntai):
-            # Parallel state is initialized in v1 worker,
-            # so this connector is for sure worker connector.
-            "role": KVConnectorRole_V1.WORKER,
-        }
-        _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector(**kwargs)
-
-
 def ensure_model_parallel_initialized(
     tensor_model_parallel_size: int,
     pipeline_model_parallel_size: int,