diff --git a/setup.py b/setup.py index 351efa1b3..0795304b5 100644 --- a/setup.py +++ b/setup.py @@ -391,7 +391,9 @@ setup( extras_require={}, entry_points={ "vllm.platform_plugins": ["ascend = vllm_ascend:register"], - "vllm.general_plugins": - ["ascend_enhanced_model = vllm_ascend:register_model"], + "vllm.general_plugins": [ + "ascend_enhanced_model = vllm_ascend:register_model", + "ascend_kv_connector = vllm_ascend:register_connector" + ], }, ) diff --git a/vllm_ascend/__init__.py b/vllm_ascend/__init__.py index 90aede78d..f519417d1 100644 --- a/vllm_ascend/__init__.py +++ b/vllm_ascend/__init__.py @@ -27,3 +27,8 @@ def register_model(): from .models import register_model register_model() + + +def register_connector(): + from vllm_ascend.distributed import register_connector + register_connector() diff --git a/vllm_ascend/distributed/__init__.py b/vllm_ascend/distributed/__init__.py index b73f13553..0915b38a5 100644 --- a/vllm_ascend/distributed/__init__.py +++ b/vllm_ascend/distributed/__init__.py @@ -18,21 +18,23 @@ from vllm.distributed.kv_transfer.kv_connector.factory import \ KVConnectorFactory -KVConnectorFactory.register_connector( - "LLMDataDistCMgrConnector", - "vllm_ascend.distributed.llmdatadist_c_mgr_connector", - "LLMDataDistCMgrConnector") -KVConnectorFactory.register_connector( - "MooncakeConnectorV1", "vllm_ascend.distributed.mooncake_connector", - "MooncakeConnector") +def register_connector(): + KVConnectorFactory.register_connector( + "LLMDataDistCMgrConnector", + "vllm_ascend.distributed.llmdatadist_c_mgr_connector", + "LLMDataDistCMgrConnector") -KVConnectorFactory.register_connector( - "MooncakeConnectorStoreV1", - "vllm_ascend.distributed.mooncake.mooncake_store_connector_v1", - "MooncakeConnectorV1") + KVConnectorFactory.register_connector( + "MooncakeConnectorV1", "vllm_ascend.distributed.mooncake_connector", + "MooncakeConnector") -KVConnectorFactory.register_connector( - "MooncakeLayerwiseConnector", - "vllm_ascend.distributed.mooncake_layerwise_connector", - "MooncakeLayerwiseConnector") + KVConnectorFactory.register_connector( + "MooncakeConnectorStoreV1", + "vllm_ascend.distributed.mooncake.mooncake_store_connector_v1", + "MooncakeConnectorV1") + + KVConnectorFactory.register_connector( + "MooncakeLayerwiseConnector", + "vllm_ascend.distributed.mooncake_layerwise_connector", + "MooncakeLayerwiseConnector")