xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-24 14:56:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

[BugFix] Avoid calling KV connector layer APIs when metadata is unset (#28253)

Browse Source 
Signed-off-by: David Ben-David <davidb@pliops.com>
Co-authored-by: David Ben-David <davidb@pliops.com>
Co-authored-by: Mark McLoughlin <markmc@redhat.com>
This commit is contained in:
David Ben-David 2025-11-11 09:39:36 +02:00 committed by GitHub
parent a7adbc6c6b
commit cc079763c5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 18 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/attention/layer.py
View File

@ -837,6 +837,8 @@ def wait_for_kv_layer_from_connector(layer_name: str):
return
connector = get_kv_transfer_group()
if not connector.has_connector_metadata():
return
forward_context: ForwardContext = get_forward_context()
attn_metadata = forward_context.attn_metadata
@ -854,6 +856,8 @@ def maybe_save_kv_layer_to_connector(
return
connector = get_kv_transfer_group()
if not connector.has_connector_metadata():
return
forward_context: ForwardContext = get_forward_context()
attn_metadata = forward_context.attn_metadata

9
vllm/distributed/kv_transfer/kv_connector/v1/base.py
View File

@ -204,11 +204,18 @@ class KVConnectorBase_V1(ABC):
Returns:
ConnectorMetadata: the connector metadata.
"""
# Should only be called while set to valid metadata.
assert self._connector_metadata is not None
return self._connector_metadata
def has_connector_metadata(self) -> bool:
"""Check whether the connector metadata is currently set.
Returns:
bool: True if connector metadata exists, False otherwise.
"""
return self._connector_metadata is not None
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
"""
Initialize with the KV caches. Useful for pre-registering the

6
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
View File

@ -171,16 +171,22 @@ class MultiConnector(KVConnectorBase_V1):
# We must override the base class method here because we need to bind
# the metadata to each connector in the order of the connectors in the
# MultiKVConnectorMetadata.
#
# Note: Call the base class method to ensure metadata is also set on the
# MultiConnector instance itself; otherwise, `has_connector_metadata()` will
# always return False.
def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None:
assert isinstance(connector_metadata, MultiKVConnectorMetadata)
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)
super().bind_connector_metadata(connector_metadata)
def clear_connector_metadata(self) -> None:
for c in self._connectors:
c.clear_connector_metadata()
super().clear_connector_metadata()
def shutdown(self):
exception: Exception | None = None