Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
OpenDAS
vllm_cscc
Commits
1db4f47f
Unverified
Commit
1db4f47f
authored
May 16, 2025
by
Nick Hill
Committed by
GitHub
May 16, 2025
Browse files
[BugFix] Fix multi async save in MultiConnector (#18246)
Signed-off-by:
Nick Hill
<
nhill@redhat.com
>
parent
d3d91b6f
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
17 additions
and
6 deletions
+17
-6
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
...istributed/kv_transfer/kv_connector/v1/multi_connector.py
+17
-6
No files found.
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
View file @
1db4f47f
# SPDX-License-Identifier: Apache-2.0
import
copy
from
dataclasses
import
dataclass
from
typing
import
TYPE_CHECKING
,
Any
,
Optional
import
torch
...
...
@@ -21,9 +22,10 @@ if TYPE_CHECKING:
logger
=
init_logger
(
__name__
)
class
MultiKVConnectorMetadata
(
tuple
[
KVConnectorMetadata
,
...],
KVConnectorMetadata
):
pass
@
dataclass
class
MultiKVConnectorMetadata
(
KVConnectorMetadata
):
metadata
:
tuple
[
KVConnectorMetadata
,
...]
extra_async_saves
:
Optional
[
dict
[
str
,
int
]]
=
None
class
MultiConnector
(
KVConnectorBase_V1
):
...
...
@@ -54,6 +56,7 @@ class MultiConnector(KVConnectorBase_V1):
# Keeps track of *additional* remaining async saves (beyond 1) to be
# finished per request. Not needed for async loads since we only allow
# a single connector to load.
# Propagated from scheduler to worker side via the connector metadata.
self
.
_extra_async_saves
:
dict
[
str
,
int
]
=
{}
def
register_kv_caches
(
self
,
kv_caches
:
dict
[
str
,
torch
.
Tensor
]):
...
...
@@ -66,7 +69,10 @@ class MultiConnector(KVConnectorBase_V1):
def
bind_connector_metadata
(
self
,
connector_metadata
:
KVConnectorMetadata
)
->
None
:
assert
isinstance
(
connector_metadata
,
MultiKVConnectorMetadata
)
for
c
,
cm
in
zip
(
self
.
_connectors
,
connector_metadata
):
if
connector_metadata
.
extra_async_saves
:
self
.
_extra_async_saves
.
update
(
connector_metadata
.
extra_async_saves
)
for
c
,
cm
in
zip
(
self
.
_connectors
,
connector_metadata
.
metadata
):
c
.
bind_connector_metadata
(
cm
)
def
clear_connector_metadata
(
self
)
->
None
:
...
...
@@ -152,8 +158,13 @@ class MultiConnector(KVConnectorBase_V1):
def
build_connector_meta
(
self
,
scheduler_output
:
SchedulerOutput
)
->
MultiKVConnectorMetadata
:
return
MultiKVConnectorMetadata
(
c
.
build_connector_meta
(
scheduler_output
)
for
c
in
self
.
_connectors
)
metadata
=
MultiKVConnectorMetadata
(
metadata
=
tuple
(
c
.
build_connector_meta
(
scheduler_output
)
for
c
in
self
.
_connectors
))
if
self
.
_extra_async_saves
:
metadata
.
extra_async_saves
=
self
.
_extra_async_saves
self
.
_extra_async_saves
=
{}
return
metadata
def
request_finished
(
self
,
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment