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
[BugFix] Fix multi async save in MultiConnector (#18246)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
May 16, 2025 at 08:13 UTC
1db4f47f8128f36f43ca1b61905075ebfb4ddacb
1 file changed
+17
-6
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+17
-6
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
import copy
3
+from dataclasses import dataclass
4
from typing import TYPE_CHECKING, Any, Optional
5
6
import torch
@@ -21,9 +22,10 @@ if TYPE_CHECKING:
22
logger = init_logger(__name__)
23
24
24
-class MultiKVConnectorMetadata(tuple[KVConnectorMetadata, ...],
25
- KVConnectorMetadata):
26
- pass
25
+@dataclass
26
+class MultiKVConnectorMetadata(KVConnectorMetadata):
27
+ metadata: tuple[KVConnectorMetadata, ...]
28
+ extra_async_saves: Optional[dict[str, int]] = None
29
30
31
class MultiConnector(KVConnectorBase_V1):
@@ -54,6 +56,7 @@ class MultiConnector(KVConnectorBase_V1):
56
# Keeps track of *additional* remaining async saves (beyond 1) to be
57
# finished per request. Not needed for async loads since we only allow
58
# a single connector to load.
59
+ # Propagated from scheduler to worker side via the connector metadata.
60
self._extra_async_saves: dict[str, int] = {}
61
62
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
@@ -66,7 +69,10 @@ class MultiConnector(KVConnectorBase_V1):
69
def bind_connector_metadata(
70
self, connector_metadata: KVConnectorMetadata) -> None:
71
assert isinstance(connector_metadata, MultiKVConnectorMetadata)
69
- for c, cm in zip(self._connectors, connector_metadata):
72
+ if connector_metadata.extra_async_saves:
73
+ self._extra_async_saves.update(
74
+ connector_metadata.extra_async_saves)
75
+ for c, cm in zip(self._connectors, connector_metadata.metadata):
76
c.bind_connector_metadata(cm)
77
78
def clear_connector_metadata(self) -> None:
@@ -152,8 +158,13 @@ class MultiConnector(KVConnectorBase_V1):
158
def build_connector_meta(
159
self,
160
scheduler_output: SchedulerOutput) -> MultiKVConnectorMetadata:
155
- return MultiKVConnectorMetadata(
156
- c.build_connector_meta(scheduler_output) for c in self._connectors)
161
+ metadata = MultiKVConnectorMetadata(metadata=tuple(
162
+ c.build_connector_meta(scheduler_output)
163
+ for c in self._connectors))
164
+ if self._extra_async_saves:
165
+ metadata.extra_async_saves = self._extra_async_saves
166
+ self._extra_async_saves = {}
167
+ return metadata
168
169
def request_finished(
170
self,