@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 58279c60b5
Fork 0

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 13, 20:45 UTC

[KV Connector] Make KVCacheConfig an explicit constructor argument (#27887)

Signed-off-by: Mark McLoughlin <markmc@redhat.com>

Mark McLoughlin committed Nov 4, 2025 at 07:00 UTC 58279c60b52c7e6e286799a313416949f43aeefe
14 files changed +410 -43
tests/v1/kv_connector/unit/test_backwards_compatibility.py new
+275
@@ -0,0 +1,275 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Unit tests for backwards compatibility with external KV connector implementations.
5 +
6 +This test ensures that external connectors (loaded via kv_connector_module_path)
7 +implemented with the old signature continue to work:
8 +- Old signature: __init__(self, vllm_config, role)
9 +- New signature: __init__(self, vllm_config, role, kv_cache_config)
10 +"""
11 +
12 +from typing import TYPE_CHECKING
13 +from unittest.mock import patch
14 +
15 +import pytest
16 +
17 +from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
18 +from vllm.distributed.kv_transfer.kv_connector.v1 import (
19 + KVConnectorBase_V1,
20 + KVConnectorRole,
21 +)
22 +from vllm.v1.core.sched.output import SchedulerOutput
23 +
24 +from .utils import create_scheduler, create_vllm_config
25 +
26 +if TYPE_CHECKING:
27 + from vllm.attention.backends.abstract import AttentionMetadata
28 + from vllm.config import VllmConfig
29 + from vllm.forward_context import ForwardContext
30 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
31 + from vllm.v1.kv_cache_interface import KVCacheConfig
32 + from vllm.v1.request import Request
33 +
34 +
35 +class OldStyleTestConnector(KVConnectorBase_V1):
36 + """
37 + Test connector using the old signature with 2 required arguments.
38 + This simulates external connectors that haven't been updated yet.
39 + """
40 +
41 + def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
42 + # Old-style call to super().__init__ with only 2 arguments
43 + super().__init__(vllm_config=vllm_config, role=role)
44 +
45 + def get_num_new_matched_tokens(
46 + self, request: "Request", num_computed_tokens: int
47 + ) -> tuple[int | None, bool]:
48 + return 0, False
49 +
50 + def update_state_after_alloc(
51 + self,
52 + request: "Request",
53 + blocks: "KVCacheBlocks",
54 + num_external_tokens: int,
55 + ):
56 + pass
57 +
58 + def build_connector_meta(self, scheduler_output: SchedulerOutput):
59 + return None
60 +
61 + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
62 + pass
63 +
64 + def wait_for_layer_load(self, layer_name: str) -> None:
65 + pass
66 +
67 + def save_kv_layer(
68 + self,
69 + layer_name: str,
70 + kv_layer,
71 + attn_metadata: "AttentionMetadata",
72 + **kwargs,
73 + ) -> None:
74 + pass
75 +
76 + def wait_for_save(self):
77 + pass
78 +
79 +
80 +class NewStyleTestConnector(KVConnectorBase_V1):
81 + """
82 + Test connector using the new signature with 3 required arguments.
83 + """
84 +
85 + def __init__(
86 + self,
87 + vllm_config: "VllmConfig",
88 + role: KVConnectorRole,
89 + kv_cache_config: "KVCacheConfig",
90 + ):
91 + # New-style call to super().__init__ with all 3 arguments
92 + super().__init__(
93 + vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config
94 + )
95 +
96 + def get_num_new_matched_tokens(
97 + self, request: "Request", num_computed_tokens: int
98 + ) -> tuple[int | None, bool]:
99 + return 0, False
100 +
101 + def update_state_after_alloc(
102 + self,
103 + request: "Request",
104 + blocks: "KVCacheBlocks",
105 + num_external_tokens: int,
106 + ):
107 + pass
108 +
109 + def build_connector_meta(self, scheduler_output: SchedulerOutput):
110 + return None
111 +
112 + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
113 + pass
114 +
115 + def wait_for_layer_load(self, layer_name: str) -> None:
116 + pass
117 +
118 + def save_kv_layer(
119 + self,
120 + layer_name: str,
121 + kv_layer,
122 + attn_metadata: "AttentionMetadata",
123 + **kwargs,
124 + ) -> None:
125 + pass
126 +
127 + def wait_for_save(self):
128 + pass
129 +
130 +
131 +@pytest.mark.parametrize("role", [KVConnectorRole.SCHEDULER, KVConnectorRole.WORKER])
132 +def test_external_old_signature_factory_instantiation(role):
133 + """
134 + Test that external connectors with old signature (2 required args) loaded
135 + via kv_connector_module_path are correctly instantiated with backwards
136 + compatibility support.
137 + """
138 + vllm_config = create_vllm_config()
139 + vllm_config.kv_transfer_config.kv_connector = "OldStyleTestConnector"
140 + vllm_config.kv_transfer_config.kv_connector_module_path = (
141 + "tests.v1.kv_connector.unit.test_backwards_compatibility"
142 + )
143 +
144 + scheduler = create_scheduler(vllm_config)
145 + kv_cache_config = scheduler.kv_cache_config
146 +
147 + connector = KVConnectorFactory.create_connector(vllm_config, role, kv_cache_config)
148 +
149 + assert connector is not None
150 + assert isinstance(connector, OldStyleTestConnector)
151 + assert connector.role == role
152 + assert connector._kv_cache_config is None
153 +
154 +
155 +@pytest.mark.parametrize("role", [KVConnectorRole.SCHEDULER, KVConnectorRole.WORKER])
156 +def test_external_new_signature_factory_instantiation(role):
157 + """
158 + Test that external connectors with new signature (3 required args) loaded
159 + via kv_connector_module_path are correctly instantiated.
160 + """
161 + vllm_config = create_vllm_config()
162 + vllm_config.kv_transfer_config.kv_connector = "NewStyleTestConnector"
163 + vllm_config.kv_transfer_config.kv_connector_module_path = (
164 + "tests.v1.kv_connector.unit.test_backwards_compatibility"
165 + )
166 +
167 + scheduler = create_scheduler(vllm_config)
168 + kv_cache_config = scheduler.kv_cache_config
169 +
170 + connector = KVConnectorFactory.create_connector(vllm_config, role, kv_cache_config)
171 +
172 + assert connector is not None
173 + assert isinstance(connector, NewStyleTestConnector)
174 + assert connector.role == role
175 + assert connector._kv_cache_config is not None
176 + assert connector._kv_cache_config == kv_cache_config
177 +
178 +
179 +@pytest.mark.parametrize("role", [KVConnectorRole.SCHEDULER, KVConnectorRole.WORKER])
180 +def test_old_signature_super_init(role):
181 + """
182 + Test that old-style connectors can call super().__init__() without
183 + kv_cache_config parameter.
184 + """
185 + vllm_config = create_vllm_config()
186 +
187 + connector = OldStyleTestConnector(vllm_config, role)
188 +
189 + assert connector is not None
190 + assert connector.role == role
191 + assert connector._kv_cache_config is None
192 +
193 +
194 +def test_old_signature_super_init_with_kwargs():
195 + """
196 + Test that old-style connectors can call super().__init__() with keyword
197 + arguments in different orders.
198 + """
199 + vllm_config = create_vllm_config()
200 +
201 + # Test with vllm_config= and role= kwargs
202 + connector1 = OldStyleTestConnector(
203 + vllm_config=vllm_config, role=KVConnectorRole.SCHEDULER
204 + )
205 + assert connector1 is not None
206 + assert connector1._kv_cache_config is None
207 +
208 + # Test with role= and vllm_config= in reversed order
209 + connector2 = OldStyleTestConnector(
210 + role=KVConnectorRole.WORKER, vllm_config=vllm_config
211 + )
212 + assert connector2 is not None
213 + assert connector2._kv_cache_config is None
214 +
215 +
216 +def test_internal_connector_uses_new_signature():
217 + """
218 + Test that internal connectors (registered in factory) always use the new
219 + signature and get kv_cache_config.
220 + """
221 + from vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector import (
222 + SharedStorageConnector,
223 + )
224 +
225 + vllm_config = create_vllm_config()
226 + vllm_config.kv_transfer_config.kv_connector = "SharedStorageConnector"
227 +
228 + scheduler = create_scheduler(vllm_config)
229 + kv_cache_config = scheduler.kv_cache_config
230 +
231 + connector = KVConnectorFactory.create_connector(
232 + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
233 + )
234 +
235 + assert connector is not None
236 + assert isinstance(connector, SharedStorageConnector)
237 + assert connector._kv_cache_config is not None
238 + assert connector._kv_cache_config == kv_cache_config
239 +
240 +
241 +def test_signature_detection_with_mocking():
242 + """
243 + Test that the factory correctly applies compat_sig flag returned from
244 + _get_connector_class_with_compat.
245 + """
246 + vllm_config = create_vllm_config()
247 + scheduler = create_scheduler(vllm_config)
248 + kv_cache_config = scheduler.kv_cache_config
249 +
250 + # Mock _get_connector_class_with_compat to return old-style connector
251 + with patch.object(
252 + KVConnectorFactory,
253 + "_get_connector_class_with_compat",
254 + return_value=(OldStyleTestConnector, True),
255 + ):
256 + old_connector = KVConnectorFactory.create_connector(
257 + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
258 + )
259 + assert old_connector is not None
260 + assert isinstance(old_connector, OldStyleTestConnector)
261 + assert old_connector._kv_cache_config is None
262 +
263 + # Mock _get_connector_class_with_compat to return new-style connector
264 + with patch.object(
265 + KVConnectorFactory,
266 + "_get_connector_class_with_compat",
267 + return_value=(NewStyleTestConnector, False),
268 + ):
269 + new_connector = KVConnectorFactory.create_connector(
270 + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
271 + )
272 + assert new_connector is not None
273 + assert isinstance(new_connector, NewStyleTestConnector)
274 + assert new_connector._kv_cache_config is not None
275 + assert new_connector._kv_cache_config == kv_cache_config
tests/v1/kv_connector/unit/utils.py
+1 -1
@@ -254,7 +254,7 @@ def create_model_runner_output(
254
255
256 class TestSharedStorageConnector(SharedStorageConnector):
257 - def __init__(self, config: VllmConfig, role):
257 + def __init__(self, config: VllmConfig, role, kv_cache_config):
258 self.name = config.kv_transfer_config.kv_connector_extra_config["name"]
259 self._connector = SharedStorageConnector(config, role)
260 self.call_record: dict[str, int] = defaultdict(int)
vllm/distributed/kv_transfer/kv_connector/factory.py
+33 -8
@@ -3,10 +3,9 @@
3
4 import importlib
5 from collections.abc import Callable
6 -from typing import TYPE_CHECKING, cast
6 +from typing import TYPE_CHECKING, Optional, cast
7
8 import vllm.envs as envs
9 -from vllm.config import VllmConfig
9 from vllm.distributed.kv_transfer.kv_connector.base import (
10 KVConnectorBase,
11 KVConnectorBaseType,
@@ -16,9 +15,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
15 supports_hma,
16 )
17 from vllm.logger import init_logger
18 +from vllm.utils.func_utils import supports_kw
19
20 if TYPE_CHECKING:
21 + from vllm.config import VllmConfig
22 from vllm.config.kv_transfer import KVTransferConfig
23 + from vllm.v1.kv_cache_interface import KVCacheConfig
24
25 logger = init_logger(__name__)
26
@@ -41,8 +43,9 @@ class KVConnectorFactory:
43 @classmethod
44 def create_connector(
45 cls,
44 - config: VllmConfig,
46 + config: "VllmConfig",
47 role: KVConnectorRole,
48 + kv_cache_config: Optional["KVCacheConfig"] = None,
49 ) -> KVConnectorBase:
50 if not envs.VLLM_USE_V1:
51 raise ValueError(
@@ -53,7 +56,9 @@ class KVConnectorFactory:
56 kv_transfer_config = config.kv_transfer_config
57 if kv_transfer_config is None:
58 raise ValueError("kv_transfer_config must be set to create a connector")
56 - connector_cls = cls.get_connector_class(kv_transfer_config)
59 + connector_cls, compat_sig = cls._get_connector_class_with_compat(
60 + kv_transfer_config
61 + )
62
63 # check if the connector supports HMA
64 hma_enabled = not config.scheduler_config.disable_hybrid_kv_cache_manager
@@ -76,7 +81,12 @@ class KVConnectorFactory:
81 # - Co-locate with worker process
82 # - Should only be used inside the forward context & attention layer
83 # We build separately to enforce strict separation
79 - return connector_cls(config, role)
84 + if compat_sig:
85 + # Old signature: __init__(self, vllm_config, role)
86 + return connector_cls(config, role)
87 + else:
88 + # New signature: __init__(self, vllm_config, role, kv_cache_config)
89 + return connector_cls(config, role, kv_cache_config)
90
91 @classmethod
92 def get_connector_class_by_name(
@@ -97,13 +107,13 @@ class KVConnectorFactory:
107 return cls._registry[connector_name]()
108
109 @classmethod
100 - def get_connector_class(
110 + def _get_connector_class_with_compat(
111 cls, kv_transfer_config: "KVTransferConfig"
102 - ) -> type[KVConnectorBaseType]:
103 - """Get the connector class by name."""
112 + ) -> tuple[type[KVConnectorBaseType], bool]:
113 connector_name = kv_transfer_config.kv_connector
114 if connector_name is None:
115 raise ValueError("Connector name is not set in KVTransferConfig")
116 + compat_sig = False
117 if connector_name in cls._registry:
118 connector_cls = cls._registry[connector_name]()
119 else:
@@ -118,6 +128,21 @@ class KVConnectorFactory:
128 f"Class {connector_name} not found in {connector_module_path}"
129 ) from e
130 connector_cls = cast(type[KVConnectorBaseType], connector_cls)
131 + if not supports_kw(connector_cls, "kv_cache_config"):
132 + compat_sig = True
133 + logger.warning(
134 + "Connector %s uses deprecated signature with 2 required arguments. "
135 + "Please update to include kv_cache_config as the second argument.",
136 + connector_cls.__name__,
137 + )
138 + return connector_cls, compat_sig
139 +
140 + @classmethod
141 + def get_connector_class(
142 + cls, kv_transfer_config: "KVTransferConfig"
143 + ) -> type[KVConnectorBaseType]:
144 + """Get the connector class by name."""
145 + connector_cls, _ = cls._get_connector_class_with_compat(kv_transfer_config)
146 return connector_cls
147
148
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+15 -1
@@ -58,6 +58,7 @@ if TYPE_CHECKING:
58 )
59 from vllm.forward_context import ForwardContext
60 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
61 + from vllm.v1.kv_cache_interface import KVCacheConfig
62 from vllm.v1.request import Request
63
64 # s_tensor_list, d_tensor_list, s_indices, d_indices, direction
@@ -141,7 +142,12 @@ class KVConnectorMetadata(ABC): # noqa: B024
142
143
144 class KVConnectorBase_V1(ABC):
144 - def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
145 + def __init__(
146 + self,
147 + vllm_config: "VllmConfig",
148 + role: KVConnectorRole,
149 + kv_cache_config: Optional["KVCacheConfig"] = None,
150 + ):
151 logger.warning(
152 "Initializing KVConnectorBase_V1. This API is experimental and "
153 "subject to change in the future as we iterate the design."
@@ -152,6 +158,14 @@ class KVConnectorBase_V1(ABC):
158 self._kv_transfer_config = vllm_config.kv_transfer_config
159 else:
160 raise ValueError("kv_transfer_config must be set for KVConnectorBase_V1")
161 + self._kv_cache_config = kv_cache_config
162 + if self._kv_cache_config is None:
163 + logger.warning(
164 + "KVConnectorBase_V1 initialized without kv_cache_config. "
165 + "This is deprecated - please update your connector to accept "
166 + "kv_cache_config as the third constructor argument and pass it "
167 + "to super().__init__()."
168 + )
169 self._role = role
170
171 @property
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+9 -3
@@ -32,7 +32,7 @@ Usage:
32 """
33
34 from dataclasses import dataclass
35 -from typing import TYPE_CHECKING, Any
35 +from typing import TYPE_CHECKING, Any, Optional
36
37 import torch
38
@@ -50,6 +50,7 @@ if TYPE_CHECKING:
50 from vllm.forward_context import ForwardContext
51 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
52 from vllm.v1.core.sched.output import SchedulerOutput
53 + from vllm.v1.kv_cache_interface import KVCacheConfig
54 from vllm.v1.request import Request
55
56 logger = init_logger(__name__)
@@ -79,8 +80,13 @@ class DecodeBenchConnector(KVConnectorBase_V1):
80 testing of the decoder with larger input sequence lengths.
81 """
82
82 - def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
83 - super().__init__(vllm_config, role)
83 + def __init__(
84 + self,
85 + vllm_config: "VllmConfig",
86 + role: KVConnectorRole,
87 + kv_cache_config: Optional["KVCacheConfig"] = None,
88 + ):
89 + super().__init__(vllm_config, role, kv_cache_config)
90
91 self.connector_scheduler: DecodeBenchConnectorScheduler | None = None
92 self.connector_worker: DecodeBenchConnectorWorker | None = None
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py
+10 -2
@@ -20,14 +20,22 @@ if TYPE_CHECKING:
20 from vllm.attention.backends.abstract import AttentionMetadata
21 from vllm.forward_context import ForwardContext
22 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
23 + from vllm.v1.kv_cache_interface import KVCacheConfig
24 from vllm.v1.request import Request
25
26 logger = init_logger(__name__)
27
28
29 class LMCacheConnectorV1(KVConnectorBase_V1):
29 - def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
30 - super().__init__(vllm_config=vllm_config, role=role)
30 + def __init__(
31 + self,
32 + vllm_config: "VllmConfig",
33 + role: KVConnectorRole,
34 + kv_cache_config: "KVCacheConfig",
35 + ):
36 + super().__init__(
37 + vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config
38 + )
39 assert vllm_config.kv_transfer_config is not None
40 use_native = vllm_config.kv_transfer_config.get_from_extra_config(
41 "use_native", False
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+11 -3
@@ -31,6 +31,7 @@ if TYPE_CHECKING:
31 from vllm.distributed.kv_events import KVCacheEvent
32 from vllm.forward_context import ForwardContext
33 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
34 + from vllm.v1.kv_cache_interface import KVCacheConfig
35 from vllm.v1.request import Request
36
37 logger = init_logger(__name__)
@@ -109,15 +110,22 @@ class MultiConnector(KVConnectorBase_V1):
110 - Save to all connectors.
111 """
112
112 - def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
113 - super().__init__(vllm_config=vllm_config, role=role)
113 + def __init__(
114 + self,
115 + vllm_config: "VllmConfig",
116 + role: KVConnectorRole,
117 + kv_cache_config: "KVCacheConfig",
118 + ):
119 + super().__init__(
120 + vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config
121 + )
122
123 self._connectors: list[KVConnectorBase_V1] = []
124 self._ktc_kv_transfer_config = []
125 for connector_cls, temp_config in self._get_connector_classes_and_configs(
126 vllm_config
127 ):
120 - self._connectors.append(connector_cls(temp_config, role))
128 + self._connectors.append(connector_cls(temp_config, role, kv_cache_config))
129 self._ktc_kv_transfer_config.append(temp_config.kv_transfer_config)
130
131 # A mapping from request id to the index of the connector chosen to
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+10 -2
@@ -13,7 +13,7 @@ from collections import defaultdict
13 from collections.abc import Iterator
14 from concurrent.futures import Future, ThreadPoolExecutor
15 from dataclasses import dataclass
16 -from typing import TYPE_CHECKING, Any
16 +from typing import TYPE_CHECKING, Any, Optional
17
18 import msgspec
19 import numpy as np
@@ -52,6 +52,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
52 if TYPE_CHECKING:
53 from vllm.attention.backends.abstract import AttentionMetadata
54 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
55 + from vllm.v1.kv_cache_interface import KVCacheConfig
56 from vllm.v1.request import Request
57
58 Transfer = tuple[int, float] # (xfer_handle, start_time)
@@ -150,7 +151,14 @@ class NixlConnectorMetadata(KVConnectorMetadata):
151
152
153 class NixlConnector(KVConnectorBase_V1):
153 - def __init__(self, vllm_config: VllmConfig, role: KVConnectorRole):
154 + def __init__(
155 + self,
156 + vllm_config: VllmConfig,
157 + role: KVConnectorRole,
158 + kv_cache_config: Optional["KVCacheConfig"] = None,
159 + ):
160 + super().__init__(vllm_config, role, kv_cache_config)
161 +
162 assert vllm_config.kv_transfer_config is not None
163 assert vllm_config.kv_transfer_config.engine_id is not None
164 self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+8 -2
@@ -21,6 +21,7 @@ from vllm.logger import init_logger
21 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
22 from vllm.v1.core.kv_cache_utils import BlockHash
23 from vllm.v1.core.sched.output import SchedulerOutput
24 +from vllm.v1.kv_cache_interface import KVCacheConfig
25 from vllm.v1.kv_offload.abstract import OffloadingManager
26 from vllm.v1.kv_offload.factory import OffloadingSpecFactory
27 from vllm.v1.kv_offload.mediums import GPULoadStoreSpec
@@ -41,8 +42,13 @@ class OffloadingConnectorMetadata(KVConnectorMetadata):
42
43
44 class OffloadingConnector(KVConnectorBase_V1):
44 - def __init__(self, vllm_config: VllmConfig, role: KVConnectorRole):
45 - super().__init__(vllm_config, role)
45 + def __init__(
46 + self,
47 + vllm_config: VllmConfig,
48 + role: KVConnectorRole,
49 + kv_cache_config: KVCacheConfig | None = None,
50 + ):
51 + super().__init__(vllm_config, role, kv_cache_config)
52
53 spec = OffloadingSpecFactory.create_spec(vllm_config)
54
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py
+13 -3
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5 -from typing import TYPE_CHECKING, Any
5 +from typing import TYPE_CHECKING, Any, Optional
6
7 import regex as re
8 import torch
@@ -25,6 +25,7 @@ if TYPE_CHECKING:
25 from vllm.attention.backends.abstract import AttentionMetadata
26 from vllm.forward_context import ForwardContext
27 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
28 + from vllm.v1.kv_cache_interface import KVCacheConfig
29 from vllm.v1.request import Request
30
31 logger = init_logger(__name__)
@@ -71,8 +72,17 @@ class P2pNcclConnectorMetadata(KVConnectorMetadata):
72
73
74 class P2pNcclConnector(KVConnectorBase_V1):
74 - def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
75 - super().__init__(vllm_config=vllm_config, role=role)
75 + def __init__(
76 + self,
77 + vllm_config: "VllmConfig",
78 + role: KVConnectorRole,
79 + kv_cache_config: Optional["KVCacheConfig"] = None,
80 + ):
81 + super().__init__(
82 + vllm_config=vllm_config,
83 + role=role,
84 + kv_cache_config=kv_cache_config,
85 + )
86 self._block_size = vllm_config.cache_config.block_size
87 self._requests_need_load: dict[str, Any] = {}
88 self.is_producer = self._kv_transfer_config.is_kv_producer
vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py
+13 -3
@@ -3,7 +3,7 @@
3 import hashlib
4 import os
5 from dataclasses import dataclass, field
6 -from typing import TYPE_CHECKING, Any
6 +from typing import TYPE_CHECKING, Any, Optional
7
8 import safetensors
9 import torch
@@ -22,6 +22,7 @@ if TYPE_CHECKING:
22 from vllm.attention.backends.abstract import AttentionMetadata
23 from vllm.forward_context import ForwardContext
24 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
25 + from vllm.v1.kv_cache_interface import KVCacheConfig
26 from vllm.v1.request import Request
27
28 logger = init_logger(__name__)
@@ -86,8 +87,17 @@ class SharedStorageConnector(KVConnectorBase_V1):
87 # It does extra work which will overwrite the existing prefix-cache in GPU
88 # - to remove the overhead, need to add some "mask" in the ReqMeta class
89
89 - def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
90 - super().__init__(vllm_config=vllm_config, role=role)
90 + def __init__(
91 + self,
92 + vllm_config: "VllmConfig",
93 + role: KVConnectorRole,
94 + kv_cache_config: Optional["KVCacheConfig"] = None,
95 + ):
96 + super().__init__(
97 + vllm_config=vllm_config,
98 + role=role,
99 + kv_cache_config=kv_cache_config,
100 + )
101 self._block_size = vllm_config.cache_config.block_size
102 self._requests_need_load: dict[str, Request] = {}
103 self._storage_path = self._kv_transfer_config.get_from_extra_config(
vllm/distributed/kv_transfer/kv_transfer_state.py
+8 -3
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import TYPE_CHECKING
3 +from typing import TYPE_CHECKING, Optional
4
5 from vllm import envs
6 from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType
@@ -12,6 +12,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
12
13 if TYPE_CHECKING:
14 from vllm.config import VllmConfig
15 + from vllm.v1.kv_cache_interface import KVCacheConfig
16
17 _KV_CONNECTOR_AGENT: KVConnectorBaseType | None = None
18
@@ -48,7 +49,9 @@ def is_v1_kv_transfer_group(connector: KVConnectorBaseType | None = None) -> boo
49 return isinstance(connector, KVConnectorBase_V1)
50
51
51 -def ensure_kv_transfer_initialized(vllm_config: "VllmConfig") -> None:
52 +def ensure_kv_transfer_initialized(
53 + vllm_config: "VllmConfig", kv_cache_config: Optional["KVCacheConfig"] = None
54 +) -> None:
55 """
56 Initialize KV cache transfer parallel group.
57 """
@@ -64,7 +67,9 @@ def ensure_kv_transfer_initialized(vllm_config: "VllmConfig") -> None:
67 ):
68 if envs.VLLM_USE_V1:
69 _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector(
67 - config=vllm_config, role=KVConnectorRole.WORKER
70 + config=vllm_config,
71 + role=KVConnectorRole.WORKER,
72 + kv_cache_config=kv_cache_config,
73 )
74 else:
75 raise ValueError("V0 is no longer supported")
vllm/v1/core/sched/scheduler.py
+3 -9
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import copy
3 import itertools
4 import time
5 from collections import defaultdict
@@ -92,15 +91,10 @@ class Scheduler(SchedulerInterface):
91 assert not self.is_encoder_decoder, (
92 "Encoder-decoder models are not currently supported with KV connectors"
93 )
95 -
96 - connector_vllm_config = copy.copy(self.vllm_config)
97 -
98 - # We're dynamically inserting a kv_cache_config variable into the
99 - # connector_vllm_config. This is distinct from the cache_config
100 - # that is already in there.
101 - connector_vllm_config.kv_cache_config = copy.copy(kv_cache_config) # type: ignore[attr-defined]
94 self.connector = KVConnectorFactory.create_connector(
103 - config=connector_vllm_config, role=KVConnectorRole.SCHEDULER
95 + config=self.vllm_config,
96 + role=KVConnectorRole.SCHEDULER,
97 + kv_cache_config=self.kv_cache_config,
98 )
99 if self.log_stats:
100 self.connector_prefix_cache_stats = PrefixCacheStats()
vllm/v1/worker/gpu_worker.py
+1 -3
@@ -380,9 +380,7 @@ class Worker(WorkerBase):
380 # NOTE(Kuntai): This need to be done before `initialize_kv_cache`,
381 # because `initialize_kv_cache` will inject kv cache groups not
382 # related to kv cache connector (e.g. kv cache sharing layers).
383 - connector_vllm_config = copy.copy(self.vllm_config)
384 - connector_vllm_config.kv_cache_config = copy.copy(kv_cache_config)
385 - ensure_kv_transfer_initialized(connector_vllm_config)
383 + ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
384
385 if self.vllm_config.model_config.enable_sleep_mode:
386 from vllm.device_allocator.cumem import CuMemAllocator