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 15, 09:45 UTC
[V1] Support multiple kv connectors (#17564)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com>
Michael Goin committed
May 14, 2025 at 19:28 UTC
2142035b51795d69d7434ce74e7aad746a2bab7a
4 files changed
+424
-2
tests/v1/kv_connector/unit/test_multi_connector.py
new
+241
@@ -0,0 +1,241 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import filecmp
3
+import shutil
4
+import tempfile
5
+from collections import defaultdict
6
+from pathlib import Path
7
+
8
+from vllm import LLM, SamplingParams
9
+from vllm.config import KVTransferConfig, VllmConfig
10
+from vllm.distributed.kv_transfer.kv_connector.factory import (
11
+ KVConnectorFactory)
12
+from vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector import ( # noqa
13
+ SharedStorageConnector)
14
+
15
+MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
16
+
17
+PROMPT_CONTEXT = "Hi " * 100
18
+PROMPTS = [
19
+ PROMPT_CONTEXT + "Hello, my name is",
20
+ PROMPT_CONTEXT + "The capital of France is",
21
+]
22
+
23
+SAMPLING_PARAMS = SamplingParams(temperature=0, max_tokens=20)
24
+
25
+
26
+class TestSharedStorageConnector(SharedStorageConnector):
27
+
28
+ def __init__(self, config: VllmConfig, role):
29
+ self.name = config.kv_transfer_config.kv_connector_extra_config["name"]
30
+ self._connector = SharedStorageConnector(config, role)
31
+ self.call_record: dict[str, int] = defaultdict(int)
32
+ # Use a unique temp file per connector
33
+ self._event_file = tempfile.gettempdir(
34
+ ) + f"/connector_{self.name}_events.log"
35
+ # Start with an empty file
36
+ with open(self._event_file, "w") as _:
37
+ pass
38
+
39
+ def __getattribute__(self, name):
40
+ if name in ("_connector", "call_record", "name", "_event_file",
41
+ "__class__", "__dict__", "__getattribute__",
42
+ "__init__"): # avoid recursion
43
+ return object.__getattribute__(self, name)
44
+ if not hasattr(self._connector, name):
45
+ return object.__getattribute__(self, name)
46
+ attr = getattr(self._connector, name)
47
+
48
+ # Intercept calls to the connector interface and write an event
49
+ # for each one to a file, which can be read back in the main test proc.
50
+ if callable(attr):
51
+
52
+ def wrapper(*args, **kwargs):
53
+ self.call_record[name] += 1
54
+ # Log the event as a line to the file
55
+ try:
56
+ with open(self._event_file, "a") as f:
57
+ f.write(name + "\n")
58
+ except Exception as e:
59
+ print(f"[ERROR] Could not log event {name} "
60
+ f"for {self.name}: {e}")
61
+ return attr(*args, **kwargs)
62
+
63
+ return wrapper
64
+ return attr
65
+
66
+
67
+KVConnectorFactory.register_connector("TestSharedStorageConnector",
68
+ TestSharedStorageConnector.__module__,
69
+ TestSharedStorageConnector.__name__)
70
+
71
+
72
+# Helper function to compare directories recursively
73
+def _compare_directories(dir1: Path, dir2: Path) -> bool:
74
+ """Compares two directories recursively for identical content."""
75
+ dcmp = filecmp.dircmp(dir1, dir2)
76
+ if dcmp.left_only or dcmp.right_only or dcmp.diff_files:
77
+ print(f"Differences found between {dir1} and {dir2}:")
78
+ print(f" Left only: {dcmp.left_only}")
79
+ print(f" Right only: {dcmp.right_only}")
80
+ print(f" Different files: {dcmp.diff_files}")
81
+ return False
82
+ for sub_dir in dcmp.common_dirs:
83
+ if not _compare_directories(dir1 / sub_dir, dir2 / sub_dir):
84
+ return False
85
+ return True
86
+
87
+
88
+def test_multi_shared_storage_connector_consistency():
89
+ """
90
+ Tests that MultiConnector with two SharedStorageConnectors saves
91
+ identical KV cache data to separate storage locations.
92
+ """
93
+ storage_1_path = Path("storage_1/")
94
+ storage_2_path = Path("storage_2/")
95
+ shutil.rmtree(storage_1_path, ignore_errors=True)
96
+ shutil.rmtree(storage_2_path, ignore_errors=True)
97
+ storage_1_path.mkdir()
98
+ storage_2_path.mkdir()
99
+
100
+ # Configure MultiConnector with two SharedStorageConnectors
101
+ kv_transfer_config = KVTransferConfig(
102
+ kv_connector="MultiConnector",
103
+ kv_role="kv_both",
104
+ kv_connector_extra_config={
105
+ "connectors": [{
106
+ "kv_connector": "TestSharedStorageConnector",
107
+ "kv_role": "kv_both",
108
+ "kv_connector_extra_config": {
109
+ "shared_storage_path": str(storage_1_path),
110
+ "name": "storage1",
111
+ }
112
+ }, {
113
+ "kv_connector": "TestSharedStorageConnector",
114
+ "kv_role": "kv_both",
115
+ "kv_connector_extra_config": {
116
+ "shared_storage_path": str(storage_2_path),
117
+ "name": "storage2",
118
+ }
119
+ }]
120
+ },
121
+ )
122
+
123
+ llm = LLM(
124
+ model=MODEL_NAME,
125
+ enforce_eager=True,
126
+ gpu_memory_utilization=0.5,
127
+ kv_transfer_config=kv_transfer_config,
128
+ )
129
+ # Run generation - this should trigger saving KV cache
130
+ _ = llm.generate(PROMPTS, SAMPLING_PARAMS)
131
+
132
+ # --- Verification ---
133
+
134
+ # Check that both storage directories were populated
135
+ local_subdirs = list(storage_1_path.iterdir())
136
+ external_subdirs = list(storage_2_path.iterdir())
137
+
138
+ assert len(
139
+ local_subdirs
140
+ ) > 0, f"Local storage path {storage_1_path} is empty after generation."
141
+ assert len(external_subdirs) > 0, (
142
+ f"External storage path {storage_2_path} is empty after generation.")
143
+ assert len(local_subdirs) == len(external_subdirs), (
144
+ f"Mismatch in number of cache entries: "
145
+ f"Local={len(local_subdirs)}, External={len(external_subdirs)}")
146
+
147
+ # The subdirectories should correspond to the prompt hashes
148
+ # Since prompts are the same, the hash directories should be the same name
149
+ local_subdir_names = sorted([d.name for d in local_subdirs])
150
+ external_subdir_names = sorted([d.name for d in external_subdirs])
151
+ assert local_subdir_names == external_subdir_names, (
152
+ "Cache directory names do not match between local and external storage"
153
+ )
154
+
155
+ # Compare the contents of each corresponding cache directory
156
+ for subdir_name in local_subdir_names:
157
+ print(f"Comparing contents of cache directory: {subdir_name}")
158
+ assert _compare_directories(storage_1_path / subdir_name,
159
+ storage_2_path / subdir_name), \
160
+ (f"Contents differ for cache directory '{subdir_name}' between "
161
+ f"{storage_1_path} and {storage_2_path}")
162
+
163
+ events = get_connector_events()
164
+ # get_num_new_matched_tokens will be called on each connector in turn.
165
+ # neither of them have hits so update_state_after_alloc won't be called.
166
+ assert events["storage1"][:3] == [
167
+ 'get_num_new_matched_tokens', 'build_connector_meta',
168
+ 'bind_connector_metadata'
169
+ ]
170
+ assert events["storage2"][:3] == [
171
+ 'get_num_new_matched_tokens', 'build_connector_meta',
172
+ 'bind_connector_metadata'
173
+ ]
174
+
175
+ # Reset prefix cache or else we'll just get the tokens back from there.
176
+ llm.reset_prefix_cache()
177
+
178
+ # Run generation again - this should trigger loading from the first
179
+ # connector.
180
+ _ = llm.generate(PROMPTS, SAMPLING_PARAMS)
181
+
182
+ events = get_connector_events()
183
+ # get_num_new_matched_tokens will return new tokens from the first
184
+ # connector so update_state_after_alloc will be called once blocks
185
+ # are allocated for the first connector.
186
+ # get_num_new_matched_tokens *won't* be called on the second connector
187
+ # in this case.
188
+ assert events["storage1"][:4] == [
189
+ 'get_num_new_matched_tokens', 'update_state_after_alloc',
190
+ 'build_connector_meta', 'bind_connector_metadata'
191
+ ]
192
+ assert events["storage2"][:2] == [
193
+ 'build_connector_meta', 'bind_connector_metadata'
194
+ ]
195
+
196
+ # Delete storage1 connector state
197
+ shutil.rmtree(storage_1_path)
198
+
199
+ # Reset prefix cache or else we'll just get the tokens back from there.
200
+ llm.reset_prefix_cache()
201
+
202
+ # Run generation again - this should trigger loading from the first
203
+ # connector.
204
+ _ = llm.generate(PROMPTS, SAMPLING_PARAMS)
205
+
206
+ events = get_connector_events()
207
+ # get_num_new_matched_tokens will be called for the first connector but it
208
+ # won't have a hit so update_state_after_alloc won't be called.
209
+ # get_num_new_matched_tokens will also be called on the second connector,
210
+ # but it should have a hit so update_state_after_alloc will be called.
211
+ assert events["storage1"][:3] == [
212
+ 'get_num_new_matched_tokens', 'build_connector_meta',
213
+ 'bind_connector_metadata'
214
+ ]
215
+ assert events["storage2"][:4] == [
216
+ 'get_num_new_matched_tokens', 'update_state_after_alloc',
217
+ 'build_connector_meta', 'bind_connector_metadata'
218
+ ]
219
+
220
+ # Clean up
221
+ shutil.rmtree(storage_1_path)
222
+ shutil.rmtree(storage_2_path)
223
+
224
+
225
+def get_connector_events() -> dict[str, list[str]]:
226
+ # Read in connector events and reset the files.
227
+ import glob
228
+ event_files = glob.glob(tempfile.gettempdir() + "/connector_*_events.log")
229
+ connector_events = {}
230
+ for fname in event_files:
231
+ name = fname.split("connector_")[1].split("_events.log")[0]
232
+ try:
233
+ with open(fname, "r+") as f:
234
+ connector_events[name] = [
235
+ line.strip() for line in f if line.strip()
236
+ ]
237
+ f.truncate(0)
238
+ except Exception as e:
239
+ print(f"[ERROR] Could not read connector events for {name}: {e}")
240
+
241
+ return connector_events
vllm/distributed/kv_transfer/kv_connector/factory.py
+5
@@ -110,3 +110,8 @@ KVConnectorFactory.register_connector(
110
"NixlConnector",
111
"vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector",
112
"NixlConnector")
113
+
114
+KVConnectorFactory.register_connector(
115
+ "MultiConnector",
116
+ "vllm.distributed.kv_transfer.kv_connector.v1.multi_connector",
117
+ "MultiConnector")
vllm/distributed/kv_transfer/kv_connector/v1/base.py
-2
@@ -22,7 +22,6 @@ The class provides the following primitives:
22
23
import enum
24
from abc import ABC, abstractmethod
25
-from dataclasses import dataclass
25
from typing import TYPE_CHECKING, Any, Optional
26
27
import torch
@@ -48,7 +47,6 @@ class KVConnectorRole(enum.Enum):
47
WORKER = 1
48
49
51
-@dataclass
50
class KVConnectorMetadata:
51
"""
52
Abstract Metadata used to communicate between the
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
new
+178
@@ -0,0 +1,178 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import copy
3
+from typing import TYPE_CHECKING, Any, Optional
4
+
5
+import torch
6
+
7
+from vllm.config import KVTransferConfig, VllmConfig
8
+from vllm.distributed.kv_transfer.kv_connector.factory import (
9
+ KVConnectorFactory)
10
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
11
+ KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)
12
+from vllm.logger import init_logger
13
+from vllm.v1.core.sched.output import SchedulerOutput
14
+
15
+if TYPE_CHECKING:
16
+ from vllm.attention.backends.abstract import AttentionMetadata
17
+ from vllm.forward_context import ForwardContext
18
+ from vllm.v1.core.kv_cache_manager import KVCacheBlocks
19
+ from vllm.v1.request import Request
20
+
21
+logger = init_logger(__name__)
22
+
23
+
24
+class MultiKVConnectorMetadata(tuple[KVConnectorMetadata, ...],
25
+ KVConnectorMetadata):
26
+ pass
27
+
28
+
29
+class MultiConnector(KVConnectorBase_V1):
30
+ """
31
+ A wrapper for using multiple KVConnectors at the same time.
32
+
33
+ The current logic is:
34
+ - Load KV from the first connector that advertises available tokens from
35
+ get_num_new_matched_tokens(), based on the order in the config.
36
+ - Save to all connectors.
37
+ """
38
+
39
+ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
40
+ super().__init__(vllm_config=vllm_config, role=role)
41
+ self._connectors = []
42
+ ktcs = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
43
+ "connectors")
44
+ assert ktcs is not None
45
+ for ktc in ktcs:
46
+ temp_config = copy.copy(vllm_config)
47
+ temp_config.kv_transfer_config = KVTransferConfig(**ktc)
48
+ self._connectors.append(
49
+ KVConnectorFactory.create_connector_v1(temp_config, role))
50
+
51
+ # A mapping from request id to the connector that is assigned to it.
52
+ self._requests_to_connector: dict[str, KVConnectorBase_V1] = {}
53
+
54
+ # Keeps track of *additional* remaining async saves (beyond 1) to be
55
+ # finished per request. Not needed for async loads since we only allow
56
+ # a single connector to load.
57
+ self._extra_async_saves: dict[str, int] = {}
58
+
59
+ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
60
+ for c in self._connectors:
61
+ c.register_kv_caches(kv_caches)
62
+
63
+ # We must override the base class method here because we need to bind
64
+ # the metadata to each connector in the order of the connectors in the
65
+ # MultiKVConnectorMetadata.
66
+ def bind_connector_metadata(
67
+ self, connector_metadata: KVConnectorMetadata) -> None:
68
+ assert isinstance(connector_metadata, MultiKVConnectorMetadata)
69
+ for c, cm in zip(self._connectors, connector_metadata):
70
+ c.bind_connector_metadata(cm)
71
+
72
+ def clear_connector_metadata(self) -> None:
73
+ for c in self._connectors:
74
+ c.clear_connector_metadata()
75
+
76
+ # ==============================
77
+ # Worker-side methods
78
+ # ==============================
79
+ def start_load_kv(self, forward_context: "ForwardContext",
80
+ **kwargs) -> None:
81
+ for c in self._connectors:
82
+ c.start_load_kv(forward_context, **kwargs)
83
+
84
+ def wait_for_layer_load(self, layer_name: str) -> None:
85
+ for c in self._connectors:
86
+ c.wait_for_layer_load(layer_name)
87
+
88
+ def save_kv_layer(self, layer_name: str, kv_layer: torch.Tensor,
89
+ attn_metadata: "AttentionMetadata", **kwargs) -> None:
90
+ for c in self._connectors:
91
+ c.save_kv_layer(layer_name, kv_layer, attn_metadata, **kwargs)
92
+
93
+ def wait_for_save(self):
94
+ for c in self._connectors:
95
+ c.wait_for_save()
96
+
97
+ def get_finished(
98
+ self, finished_req_ids: set[str]
99
+ ) -> tuple[Optional[set[str]], Optional[set[str]]]:
100
+ finished_recving: set[str] = set()
101
+ finished_sending: set[str] = set()
102
+ for c in self._connectors:
103
+ recving, sending = c.get_finished(finished_req_ids)
104
+ if not recving and not sending:
105
+ continue
106
+ # Aggregate finished recving request ids.
107
+ finished_recving.update(recving or ())
108
+ # Aggregate finished sending request ids - only include
109
+ # once we've drained the "extra" count (for cases where
110
+ # more than one connector is async-saving the same request).
111
+ for req_id in sending or ():
112
+ extra_pending = self._extra_async_saves.get(req_id)
113
+ if extra_pending is None:
114
+ finished_sending.add(req_id)
115
+ continue
116
+ assert extra_pending > 0
117
+ if extra_pending == 1:
118
+ del self._extra_async_saves[req_id]
119
+ else:
120
+ self._extra_async_saves[req_id] = extra_pending - 1
121
+
122
+ return finished_recving or None, finished_sending or None
123
+
124
+ # ==============================
125
+ # Scheduler-side methods
126
+ # ==============================
127
+ def get_num_new_matched_tokens(
128
+ self,
129
+ request: "Request",
130
+ num_computed_tokens: int,
131
+ ) -> tuple[int, bool]:
132
+ for c in self._connectors:
133
+ toks, load_async = c.get_num_new_matched_tokens(
134
+ request, num_computed_tokens)
135
+ # The first connector that has new matched tokens will be assigned
136
+ # to this request.
137
+ if toks > 0:
138
+ self._requests_to_connector[request.request_id] = c
139
+ return toks, load_async
140
+ return 0, False
141
+
142
+ def update_state_after_alloc(self, request: "Request",
143
+ blocks: "KVCacheBlocks",
144
+ num_external_tokens: int):
145
+ # If the request is not assigned to any connector, we do nothing.
146
+ if request.request_id not in self._requests_to_connector:
147
+ return
148
+ # We assume that the request is assigned to only one connector.
149
+ c = self._requests_to_connector.pop(request.request_id)
150
+ c.update_state_after_alloc(request, blocks, num_external_tokens)
151
+
152
+ def build_connector_meta(
153
+ self,
154
+ scheduler_output: SchedulerOutput) -> MultiKVConnectorMetadata:
155
+ return MultiKVConnectorMetadata(
156
+ c.build_connector_meta(scheduler_output) for c in self._connectors)
157
+
158
+ def request_finished(
159
+ self,
160
+ request: "Request",
161
+ blocks: "KVCacheBlocks",
162
+ ) -> tuple[bool, Optional[dict[str, Any]]]:
163
+ async_saves = 0
164
+ kv_txfer_params = None
165
+ for c in self._connectors:
166
+ async_save, txfer_params = c.request_finished(request, blocks)
167
+ if async_save:
168
+ async_saves += 1
169
+ if txfer_params is not None:
170
+ if kv_txfer_params is not None:
171
+ #TODO we can probably change this to merge the dicts here,
172
+ # checking for key clashes.
173
+ raise RuntimeError(
174
+ "Only one connector can produce KV transfer params")
175
+ kv_txfer_params = txfer_params
176
+ if async_saves > 1:
177
+ self._extra_async_saves[request.request_id] = async_saves - 1
178
+ return async_saves > 0, kv_txfer_params