vllm.distributed.kv_transfer.kv_connector.factory

logger module-attribute

logger = init_logger(__name__)

KVConnectorFactory

Source code in vllm/distributed/kv_transfer/kv_connector/factory.py

class KVConnectorFactory:
    _registry: dict[str, Callable[[], type[KVConnectorBaseType]]] = {}

    @classmethod
    def register_connector(cls, name: str, module_path: str,
                           class_name: str) -> None:
        """Register a connector with a lazy-loading module and class name."""
        if name in cls._registry:
            raise ValueError(f"Connector '{name}' is already registered.")

        def loader() -> type[KVConnectorBaseType]:
            module = importlib.import_module(module_path)
            return getattr(module, class_name)

        cls._registry[name] = loader

    @classmethod
    def create_connector_v0(cls, rank: int, local_rank: int,
                            config: "VllmConfig") -> KVConnectorBase:
        if envs.VLLM_USE_V1:
            raise ValueError("Attempting to initialize a V0 Connector, "
                             f"but found {envs.VLLM_USE_V1=}")

        connector_name = config.kv_transfer_config.kv_connector
        if connector_name not in cls._registry:
            raise ValueError(f"Unsupported connector type: {connector_name}")

        connector_cls = cls._registry[connector_name]()
        assert issubclass(connector_cls, KVConnectorBase)
        return connector_cls(rank, local_rank, config)

    @classmethod
    def create_connector_v1(
        cls,
        config: "VllmConfig",
        role: KVConnectorRole,
    ) -> KVConnectorBase_V1:
        if not envs.VLLM_USE_V1:
            raise ValueError("Attempting to initialize a V1 Connector, "
                             f"but found {envs.VLLM_USE_V1=}")

        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 and engine_id: %s",
                    connector_name, kv_transfer_config.engine_id)
        # NOTE(Kuntai): v1 connector is explicitly separated into two roles.
        # Scheduler connector:
        # - Co-locate with scheduler process
        # - Should only be used inside the Scheduler class
        # Worker connector:
        # - Co-locate with worker process
        # - Should only be used inside the forward context & attention layer
        # We build separately to enforce strict separation
        return connector_cls(config, role)

_registry class-attribute instance-attribute

_registry: dict[
    str, Callable[[], type[KVConnectorBaseType]]
] = {}

create_connector_v0 classmethod

create_connector_v0(
    rank: int, local_rank: int, config: VllmConfig
) -> KVConnectorBase

Source code in vllm/distributed/kv_transfer/kv_connector/factory.py

@classmethod
def create_connector_v0(cls, rank: int, local_rank: int,
                        config: "VllmConfig") -> KVConnectorBase:
    if envs.VLLM_USE_V1:
        raise ValueError("Attempting to initialize a V0 Connector, "
                         f"but found {envs.VLLM_USE_V1=}")

    connector_name = config.kv_transfer_config.kv_connector
    if connector_name not in cls._registry:
        raise ValueError(f"Unsupported connector type: {connector_name}")

    connector_cls = cls._registry[connector_name]()
    assert issubclass(connector_cls, KVConnectorBase)
    return connector_cls(rank, local_rank, config)

create_connector_v1 classmethod

create_connector_v1(
    config: VllmConfig, role: KVConnectorRole
) -> KVConnectorBase_V1

Source code in vllm/distributed/kv_transfer/kv_connector/factory.py

@classmethod
def create_connector_v1(
    cls,
    config: "VllmConfig",
    role: KVConnectorRole,
) -> KVConnectorBase_V1:
    if not envs.VLLM_USE_V1:
        raise ValueError("Attempting to initialize a V1 Connector, "
                         f"but found {envs.VLLM_USE_V1=}")

    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 and engine_id: %s",
                connector_name, kv_transfer_config.engine_id)
    # NOTE(Kuntai): v1 connector is explicitly separated into two roles.
    # Scheduler connector:
    # - Co-locate with scheduler process
    # - Should only be used inside the Scheduler class
    # Worker connector:
    # - Co-locate with worker process
    # - Should only be used inside the forward context & attention layer
    # We build separately to enforce strict separation
    return connector_cls(config, role)

register_connector classmethod

register_connector(
    name: str, module_path: str, class_name: str
) -> None

Register a connector with a lazy-loading module and class name.

Source code in vllm/distributed/kv_transfer/kv_connector/factory.py

@classmethod
def register_connector(cls, name: str, module_path: str,
                       class_name: str) -> None:
    """Register a connector with a lazy-loading module and class name."""
    if name in cls._registry:
        raise ValueError(f"Connector '{name}' is already registered.")

    def loader() -> type[KVConnectorBaseType]:
        module = importlib.import_module(module_path)
        return getattr(module, class_name)

    cls._registry[name] = loader