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, 13:15 UTC
[BugFix] Avoid calling KV connector layer APIs when metadata is unset (#28253)
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>
David Ben-David committed
Nov 11, 2025 at 09:39 UTC
cc079763c59adb8c03305663a5b8857ba85deb1b
3 files changed
+18
-1
vllm/attention/layer.py
+4
@@ -837,6 +837,8 @@ def wait_for_kv_layer_from_connector(layer_name: str):
837
return
838
839
connector = get_kv_transfer_group()
840
+ if not connector.has_connector_metadata():
841
+ return
842
843
forward_context: ForwardContext = get_forward_context()
844
attn_metadata = forward_context.attn_metadata
@@ -854,6 +856,8 @@ def maybe_save_kv_layer_to_connector(
856
return
857
858
connector = get_kv_transfer_group()
859
+ if not connector.has_connector_metadata():
860
+ return
861
862
forward_context: ForwardContext = get_forward_context()
863
attn_metadata = forward_context.attn_metadata
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+8
-1
@@ -204,11 +204,18 @@ class KVConnectorBase_V1(ABC):
204
Returns:
205
ConnectorMetadata: the connector metadata.
206
"""
207
-
207
# Should only be called while set to valid metadata.
208
assert self._connector_metadata is not None
209
return self._connector_metadata
210
211
+ def has_connector_metadata(self) -> bool:
212
+ """Check whether the connector metadata is currently set.
213
+
214
+ Returns:
215
+ bool: True if connector metadata exists, False otherwise.
216
+ """
217
+ return self._connector_metadata is not None
218
+
219
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
220
"""
221
Initialize with the KV caches. Useful for pre-registering the
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+6
@@ -171,16 +171,22 @@ class MultiConnector(KVConnectorBase_V1):
171
# We must override the base class method here because we need to bind
172
# the metadata to each connector in the order of the connectors in the
173
# MultiKVConnectorMetadata.
174
+ #
175
+ # Note: Call the base class method to ensure metadata is also set on the
176
+ # MultiConnector instance itself; otherwise, `has_connector_metadata()` will
177
+ # always return False.
178
def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None:
179
assert isinstance(connector_metadata, MultiKVConnectorMetadata)
180
if connector_metadata.extra_async_saves:
181
self._extra_async_saves.update(connector_metadata.extra_async_saves)
182
for c, cm in zip(self._connectors, connector_metadata.metadata):
183
c.bind_connector_metadata(cm)
184
+ super().bind_connector_metadata(connector_metadata)
185
186
def clear_connector_metadata(self) -> None:
187
for c in self._connectors:
188
c.clear_connector_metadata()
189
+ super().clear_connector_metadata()
190
191
def shutdown(self):
192
exception: Exception | None = None