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, 22:35 UTC
feat(distributed): add `get_required_kvcache_layout` class method to kv connector api (#20433)
Signed-off-by: wxsm <wxsms@foxmail.com>
wxsm committed
Jul 31, 2025 at 00:41 UTC
f4135232b9a8c4845f8961fb1cd17581c56ae2ce
7 files changed
+186
-28
tests/distributed/test_kvlayout.py
new
+72
@@ -0,0 +1,72 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.config import (DeviceConfig, KVTransferConfig, ModelConfig,
5
+ VllmConfig, set_current_vllm_config)
6
+from vllm.distributed.kv_transfer.kv_connector.utils import (
7
+ get_kv_connector_cache_layout)
8
+from vllm.logger import init_logger
9
+
10
+logger = init_logger("test_expert_parallel")
11
+
12
+
13
+def test_get_kv_connector_cache_layout_without_kv_connector():
14
+ vllm_config = VllmConfig(device_config=DeviceConfig("cpu"))
15
+ with set_current_vllm_config(vllm_config):
16
+ # Test with default settings
17
+ layout = get_kv_connector_cache_layout()
18
+ assert layout == "NHD"
19
+
20
+
21
+def test_get_kv_connector_cache_layout_with_lmcache_connector():
22
+ kv_transfer_config = KVTransferConfig(
23
+ kv_connector="LMCacheConnectorV1",
24
+ kv_role="kv_both",
25
+ )
26
+ vllm_config = VllmConfig(device_config=DeviceConfig("cpu"),
27
+ kv_transfer_config=kv_transfer_config)
28
+ with set_current_vllm_config(vllm_config):
29
+ # Test with default settings
30
+ layout = get_kv_connector_cache_layout()
31
+ assert layout == "NHD"
32
+
33
+
34
+def test_get_kv_connector_cache_layout_with_nixl_connector():
35
+ kv_transfer_config = KVTransferConfig(
36
+ kv_connector="NixlConnector",
37
+ kv_role="kv_both",
38
+ )
39
+ model_config = ModelConfig()
40
+ vllm_config = VllmConfig(device_config=DeviceConfig("cpu"),
41
+ model_config=model_config,
42
+ kv_transfer_config=kv_transfer_config)
43
+ with set_current_vllm_config(vllm_config):
44
+ # Test with default settings
45
+ layout = get_kv_connector_cache_layout()
46
+ assert layout == "HND"
47
+
48
+
49
+def test_get_kv_connector_cache_layout_with_multi_connector():
50
+ kv_transfer_config = KVTransferConfig(kv_connector="MultiConnector",
51
+ kv_role="kv_both",
52
+ kv_connector_extra_config={
53
+ "connectors": [{
54
+ "kv_connector":
55
+ "SharedStorageConnector",
56
+ "kv_role":
57
+ "kv_both"
58
+ }, {
59
+ "kv_connector":
60
+ "NixlConnector",
61
+ "kv_role":
62
+ "kv_both"
63
+ }]
64
+ })
65
+ model_config = ModelConfig()
66
+ vllm_config = VllmConfig(device_config=DeviceConfig("cpu"),
67
+ model_config=model_config,
68
+ kv_transfer_config=kv_transfer_config)
69
+ with set_current_vllm_config(vllm_config):
70
+ # Test with default settings
71
+ layout = get_kv_connector_cache_layout()
72
+ assert layout == "HND"
vllm/distributed/kv_transfer/kv_connector/base.py
+15
-1
@@ -9,7 +9,7 @@ The class provides two primary abstract methods:
9
"""
10
11
from abc import ABC, abstractmethod
12
-from typing import TYPE_CHECKING, Union
12
+from typing import TYPE_CHECKING, Optional, Union
13
14
import torch
15
@@ -124,5 +124,19 @@ class KVConnectorBase(ABC):
124
125
raise NotImplementedError
126
127
+ @classmethod
128
+ def get_required_kvcache_layout(
129
+ cls, vllm_config: "VllmConfig") -> Optional[str]:
130
+ """
131
+ Get the required KV cache layout for this connector.
132
+ Args:
133
+ vllm_config (VllmConfig): the vllm config.
134
+
135
+ Returns:
136
+ str: the required KV cache layout. e.g. HND, or NHD.
137
+ None if the connector does not require a specific layout.
138
+ """
139
+ return None
140
+
141
142
KVConnectorBaseType = Union[KVConnectorBase, KVConnectorBase_V1]
vllm/distributed/kv_transfer/kv_connector/factory.py
+21
-16
@@ -5,6 +5,7 @@ import importlib
5
from typing import TYPE_CHECKING, Callable
6
7
import vllm.envs as envs
8
+from vllm.config import KVTransferConfig
9
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType
10
from vllm.distributed.kv_transfer.kv_connector.v1 import (KVConnectorBase_V1,
11
KVConnectorRole)
@@ -41,25 +42,15 @@ class KVConnectorFactory:
42
raise ValueError("Attempting to initialize a V0 Connector, "
43
f"but found {envs.VLLM_USE_V1=}")
44
44
- connector_name = config.kv_transfer_config.kv_connector
45
- if connector_name not in cls._registry:
46
- raise ValueError(f"Unsupported connector type: {connector_name}")
47
-
48
- connector_cls = cls._registry[connector_name]()
45
+ connector_cls = cls.get_connector_class(config.kv_transfer_config)
46
assert issubclass(connector_cls, KVConnectorBase)
47
return connector_cls(rank, local_rank, config)
48
49
@classmethod
53
- def create_connector_v1(
54
- cls,
55
- config: "VllmConfig",
56
- role: KVConnectorRole,
57
- ) -> KVConnectorBase_V1:
58
- if not envs.VLLM_USE_V1:
59
- raise ValueError("Attempting to initialize a V1 Connector, "
60
- f"but found {envs.VLLM_USE_V1=}")
61
-
62
- kv_transfer_config = config.kv_transfer_config
50
+ def get_connector_class(
51
+ cls, kv_transfer_config: "KVTransferConfig"
52
+ ) -> type[KVConnectorBaseType]:
53
+ """Get the connector class by name."""
54
connector_name = kv_transfer_config.kv_connector
55
if connector_name in cls._registry:
56
connector_cls = cls._registry[connector_name]()
@@ -70,9 +61,23 @@ class KVConnectorFactory:
61
f"Unsupported connector type: {connector_name}")
62
connector_module = importlib.import_module(connector_module_path)
63
connector_cls = getattr(connector_module, connector_name)
64
+ return connector_cls
65
+
66
+ @classmethod
67
+ def create_connector_v1(
68
+ cls,
69
+ config: "VllmConfig",
70
+ role: KVConnectorRole,
71
+ ) -> KVConnectorBase_V1:
72
+ if not envs.VLLM_USE_V1:
73
+ raise ValueError("Attempting to initialize a V1 Connector, "
74
+ f"but found {envs.VLLM_USE_V1=}")
75
+
76
+ kv_transfer_config = config.kv_transfer_config
77
+ connector_cls = cls.get_connector_class(kv_transfer_config)
78
assert issubclass(connector_cls, KVConnectorBase_V1)
79
logger.info("Creating v1 connector with name: %s and engine_id: %s",
75
- connector_name, kv_transfer_config.engine_id)
80
+ connector_cls.__name__, kv_transfer_config.engine_id)
81
# NOTE(Kuntai): v1 connector is explicitly separated into two roles.
82
# Scheduler connector:
83
# - Co-locate with scheduler process
vllm/distributed/kv_transfer/kv_connector/utils.py
+10
-9
@@ -13,6 +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.factory import (
17
+ KVConnectorFactory)
18
from vllm.logger import init_logger
19
from vllm.v1.outputs import ModelRunnerOutput
20
@@ -103,15 +105,14 @@ def get_kv_connector_cache_layout():
105
# used for faster transfer.
106
vllm_config = get_current_vllm_config()
107
kv_config = vllm_config.kv_transfer_config
106
- if kv_config is not None and vllm_config.model_config is None:
107
- logger.warning_once("Unable to detect current VLLM config. " \
108
- "Defaulting to NHD kv cache layout.")
109
- elif kv_config is not None:
110
- use_mla = vllm_config.model_config.use_mla
111
- if not use_mla and kv_config.kv_connector == "NixlConnector":
112
- logger.info_once("NixlConnector detected. Setting KV cache " \
113
- "layout to HND for better xfer performance.")
114
- return "HND"
108
+ if kv_config is not None:
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 " \
115
+ "kv cache layout, defaulting to NHD.")
116
return "NHD"
117
118
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+14
@@ -299,3 +299,17 @@ class KVConnectorBase_V1(ABC):
299
returned by the engine.
300
"""
301
return False, None
302
+
303
+ @classmethod
304
+ def get_required_kvcache_layout(
305
+ cls, vllm_config: "VllmConfig") -> Optional[str]:
306
+ """
307
+ Get the required KV cache layout for this connector.
308
+ Args:
309
+ vllm_config (VllmConfig): the vllm config.
310
+
311
+ Returns:
312
+ str: the required KV cache layout. e.g. HND, or NHD.
313
+ None if the connector does not require a specific layout.
314
+ """
315
+ return None
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+33
@@ -202,3 +202,36 @@ class MultiConnector(KVConnectorBase_V1):
202
self._requests_to_connector.pop(request.request_id, None)
203
204
return async_saves > 0, kv_txfer_params
205
+
206
+ @classmethod
207
+ def get_required_kvcache_layout(
208
+ cls, vllm_config: "VllmConfig") -> Optional[str]:
209
+ """
210
+ Get the required KV cache layout for this connector.
211
+ Args:
212
+ vllm_config (VllmConfig): the vllm config.
213
+
214
+ Returns:
215
+ str: the required KV cache layout. e.g. HND, or NHD.
216
+ None if the connector does not require a specific layout.
217
+ """
218
+ ktcs = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
219
+ "connectors")
220
+ assert ktcs is not None
221
+ layouts: set[str] = set()
222
+ temp_vllm_config = copy.copy(vllm_config)
223
+ for ktc in ktcs:
224
+ kv_transfer_config = KVTransferConfig(**ktc)
225
+ temp_vllm_config.kv_transfer_config = kv_transfer_config
226
+ required_kvcache_layout = KVConnectorFactory.get_connector_class(
227
+ kv_transfer_config).get_required_kvcache_layout(
228
+ temp_vllm_config)
229
+ if required_kvcache_layout is not None:
230
+ layouts.add(required_kvcache_layout)
231
+
232
+ if len(layouts) > 1:
233
+ raise ValueError(f"KV cache layout mismatch: "
234
+ f"found {len(layouts)} different layouts "
235
+ f"({', '.join(layouts) })."
236
+ f"All connectors must use the same layout.")
237
+ return next(iter(layouts), None)
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+21
-2
@@ -133,6 +133,25 @@ class NixlConnector(KVConnectorBase_V1):
133
self.connector_worker = NixlConnectorWorker(
134
vllm_config, self.engine_id)
135
136
+ ############################################################
137
+ # Class Methods
138
+ ############################################################
139
+ @classmethod
140
+ def get_required_kvcache_layout(cls, vllm_config: VllmConfig):
141
+ if vllm_config.model_config is None:
142
+ logger.warning_once("Unable to detect current VLLM config. "
143
+ "Fallback to default kv cache layout.")
144
+ return None
145
+ use_mla = vllm_config.model_config.use_mla
146
+ if use_mla:
147
+ # return None when we have mla
148
+ # as the layout should not matter in that case,
149
+ # which fallback to the default behavior.
150
+ return None
151
+ logger.info_once("NixlConnector setting KV cache "
152
+ "layout to HND for better xfer performance.")
153
+ return "HND"
154
+
155
############################################################
156
# Scheduler Side Methods
157
############################################################
@@ -236,13 +255,13 @@ class NixlConnectorScheduler:
255
"""
256
For remote prefill, pull all prompt blocks from remote
257
asynchronously relative to engine execution.
239
-
258
+
259
Args:
260
request (Request): the request object.
261
num_computed_tokens (int): the number of locally
262
computed tokens for this request
263
Returns:
245
- * the number of tokens that can be loaded from the
264
+ * the number of tokens that can be loaded from the
265
external KV cache beyond what is already computed.
266
* true if the external KV cache tokens will be loaded
267
asynchronously (between scheduler steps).