A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[BugFix][Nixl][PD] Fix heterogenous TP (#22663)
Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com>
Nicolò Lucchesi committed
Aug 12, 2025 at 14:37 UTC
d030b01548d52a5e3afe56fdb8ce7a367b9799e3
2 files changed
+31
-17
vllm/distributed/kv_transfer/kv_connector/factory.py
+24
-13
@@ -4,13 +4,17 @@
4
import importlib
5
from typing import TYPE_CHECKING, Callable
6
7
+# yapf: disable
8
import vllm.envs as envs
8
-from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
9
+from vllm.distributed.kv_transfer.kv_connector.base import (
10
+ KVConnectorBase, KVConnectorBaseType)
11
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
12
from vllm.logger import init_logger
13
14
+# yapf: enable
15
+
16
if TYPE_CHECKING:
13
- from vllm.config import VllmConfig
17
+ from vllm.config import KVTransferConfig, VllmConfig
18
19
logger = init_logger(__name__)
20
@@ -42,17 +46,7 @@ class KVConnectorFactory:
46
f"but found {envs.VLLM_USE_V1=}")
47
48
kv_transfer_config = config.kv_transfer_config
45
- connector_name = kv_transfer_config.kv_connector
46
- if connector_name in cls._registry:
47
- connector_cls = cls._registry[connector_name]()
48
- else:
49
- connector_module_path = kv_transfer_config.kv_connector_module_path
50
- if connector_module_path is None:
51
- raise ValueError(
52
- f"Unsupported connector type: {connector_name}")
53
- connector_module = importlib.import_module(connector_module_path)
54
- connector_cls = getattr(connector_module, connector_name)
55
- assert issubclass(connector_cls, KVConnectorBase)
49
+ connector_cls = cls.get_connector_class(kv_transfer_config)
50
logger.info("Creating v1 connector with name: %s and engine_id: %s",
51
connector_cls.__name__, kv_transfer_config.engine_id)
52
# NOTE(Kuntai): v1 connector is explicitly separated into two roles.
@@ -65,6 +59,23 @@ class KVConnectorFactory:
59
# We build separately to enforce strict separation
60
return connector_cls(config, role)
61
62
+ @classmethod
63
+ def get_connector_class(
64
+ cls, kv_transfer_config: "KVTransferConfig"
65
+ ) -> type[KVConnectorBaseType]:
66
+ """Get the connector class by name."""
67
+ connector_name = kv_transfer_config.kv_connector
68
+ if connector_name in cls._registry:
69
+ connector_cls = cls._registry[connector_name]()
70
+ else:
71
+ connector_module_path = kv_transfer_config.kv_connector_module_path
72
+ if connector_module_path is None:
73
+ raise ValueError(
74
+ f"Unsupported connector type: {connector_name}")
75
+ connector_module = importlib.import_module(connector_module_path)
76
+ connector_cls = getattr(connector_module, connector_name)
77
+ return connector_cls
78
+
79
80
# Register various connectors here.
81
# The registration should not be done in each individual file, as we want to
vllm/distributed/kv_transfer/kv_connector/utils.py
+7
-4
@@ -13,8 +13,8 @@ import torch
13
import vllm.envs as envs
14
from vllm import _custom_ops as ops
15
from vllm.config import VllmConfig, get_current_vllm_config
16
-from vllm.distributed.kv_transfer.kv_connector.v1.base import (
17
- KVConnectorBase_V1)
16
+from vllm.distributed.kv_transfer.kv_connector.factory import (
17
+ KVConnectorFactory)
18
from vllm.logger import init_logger
19
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
20
@@ -106,8 +106,9 @@ def get_kv_connector_cache_layout():
106
vllm_config = get_current_vllm_config()
107
kv_config = vllm_config.kv_transfer_config
108
if kv_config is not None:
109
- required_kvcache_layout = (
110
- KVConnectorBase_V1.get_required_kvcache_layout(vllm_config))
109
+ connector_cls = KVConnectorFactory.get_connector_class(kv_config)
110
+ required_kvcache_layout = connector_cls.get_required_kvcache_layout(
111
+ vllm_config)
112
if required_kvcache_layout is not None:
113
return required_kvcache_layout
114
logger.info_once("Connectors do not specify a " \
@@ -143,6 +144,8 @@ class KVOutputAggregator:
144
finished_recving = set[str]()
145
for output in outputs:
146
output = output.kv_connector_output
147
+ if not output:
148
+ continue
149
update_finished_set(output.finished_sending,
150
self._send_remaining_count, finished_sending)
151
update_finished_set(output.finished_recving,