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, 02:15 UTC
[KVConnector] Always call connector `clear_metadata()` at end of step (#20756)
Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: David Ben-David <sdavidbd@gmail.com>
Nick Hill committed
Jul 10, 2025 at 22:37 UTC
574ad60db9de900d350e6b7b904c5524c01e96be
4 files changed
+25
-26
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+6
-3
@@ -57,7 +57,7 @@ class KVConnectorRole(enum.Enum):
57
WORKER = 1
58
59
60
-class KVConnectorMetadata:
60
+class KVConnectorMetadata(ABC): # noqa: B024
61
"""
62
Abstract Metadata used to communicate between the
63
Scheduler KVConnector and Worker KVConnector.
@@ -71,7 +71,7 @@ class KVConnectorBase_V1(ABC):
71
logger.warning(
72
"Initializing KVConnectorBase_V1. This API is experimental and "
73
"subject to change in the future as we iterate the design.")
74
- self._connector_metadata = KVConnectorMetadata()
74
+ self._connector_metadata: Optional[KVConnectorMetadata] = None
75
self._vllm_config = vllm_config
76
self._role = role
77
@@ -102,7 +102,7 @@ class KVConnectorBase_V1(ABC):
102
This function should be called by the model runner every time
103
after the model execution.
104
"""
105
- self._connector_metadata = KVConnectorMetadata()
105
+ self._connector_metadata = None
106
107
def _get_connector_metadata(self) -> KVConnectorMetadata:
108
"""Get the connector metadata.
@@ -112,6 +112,9 @@ class KVConnectorBase_V1(ABC):
112
Returns:
113
ConnectorMetadata: the connector metadata.
114
"""
115
+
116
+ # Should only be called while set to valid metadata.
117
+ assert self._connector_metadata is not None
118
return self._connector_metadata
119
120
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
vllm/v1/executor/multiproc_executor.py
+15
-19
@@ -250,28 +250,24 @@ class MultiprocExecutor(Executor):
250
self, outputs: list[ModelRunnerOutput]) -> ModelRunnerOutput:
251
# aggregate finished_sending, finished_recving from all workers
252
253
- finished_sending = set[str]()
254
- finished_recving = set[str]()
255
- for output in outputs:
256
- # update finished_sending
257
- for req_id in output.finished_sending or []:
258
- new_count = self._send_remaining_count[req_id] - 1
253
+ def update_finished_set(req_ids: Optional[set[str]],
254
+ remaining_count_dict: dict[str, int],
255
+ finished_set: set[str]) -> None:
256
+ for req_id in req_ids or ():
257
+ new_count = remaining_count_dict[req_id] - 1
258
if new_count == 0:
260
- # got response from all workers, report back to scheduler
261
- finished_sending.add(req_id)
262
- del self._send_remaining_count[req_id]
259
+ finished_set.add(req_id)
260
+ del remaining_count_dict[req_id]
261
else:
264
- self._send_remaining_count[req_id] = new_count
262
+ remaining_count_dict[req_id] = new_count
263
266
- # update finished_recving
267
- for req_id in output.finished_recving or []:
268
- new_count = self._recv_remaining_count[req_id] - 1
269
- if new_count == 0:
270
- # got response from all workers, report back to scheduler
271
- finished_recving.add(req_id)
272
- del self._recv_remaining_count[req_id]
273
- else:
274
- self._recv_remaining_count[req_id] = new_count
264
+ finished_sending = set[str]()
265
+ finished_recving = set[str]()
266
+ for output in outputs:
267
+ update_finished_set(output.finished_sending,
268
+ self._send_remaining_count, finished_sending)
269
+ update_finished_set(output.finished_recving,
270
+ self._recv_remaining_count, finished_recving)
271
272
# select output of the worker specified by output_rank
273
output = outputs[self.output_rank]
vllm/v1/worker/gpu_model_runner.py
-4
@@ -1539,10 +1539,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1539
attn_metadata,
1540
)
1541
1542
- # Clear KVConnector state after all KVs are generated.
1543
- if has_kv_transfer_group():
1544
- get_kv_transfer_group().clear_connector_metadata()
1545
-
1542
self.eplb_step()
1543
1544
return ModelRunnerOutput(
vllm/v1/worker/gpu_worker.py
+4
@@ -338,6 +338,10 @@ class Worker(WorkerBase):
338
output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
339
output.finished_sending = finished_sending
340
output.finished_recving = finished_recving
341
+
342
+ # Clear KVConnector state for this step.
343
+ get_kv_transfer_group().clear_connector_metadata()
344
+
345
# with a connector, the scheduler expects output from all workers
346
return output
347