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 14, 05:30 UTC
v1: Pass KVConnectorOutput to scheduler-side (#22157)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Aug 9, 2025 at 09:09 UTC
7ad7adb67f1350b6e9f7cfdd7aacf38eed093bb1
3 files changed
+22
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+13
@@ -12,6 +12,8 @@ The class provides the following primitives:
12
times for a given request and should be side-effect free.
13
update_state_after_alloc() - update KVConnector state after
14
temporary buffer alloc by the CacheManager.
15
+ update_connector_output() - update KVConnector state after
16
+ output is received from worker-side connectors.
17
request_finished() - called when a request is finished, with
18
the computed kv cache blocks for the request.
19
Returns whether KV cache should be freed now or will be
@@ -38,6 +40,7 @@ import torch
40
41
from vllm.logger import init_logger
42
from vllm.v1.core.sched.output import SchedulerOutput
43
+from vllm.v1.outputs import KVConnectorOutput
44
45
if TYPE_CHECKING:
46
from vllm.attention.backends.abstract import AttentionMetadata
@@ -283,6 +286,16 @@ class KVConnectorBase_V1(ABC):
286
"""
287
pass
288
289
+ def update_connector_output(self, connector_output: KVConnectorOutput):
290
+ """
291
+ Update KVConnector state from worker-side connectors output.
292
+
293
+ Args:
294
+ connector_output (KVConnectorOutput): the worker-side
295
+ connectors output.
296
+ """
297
+ return
298
+
299
def request_finished(
300
self,
301
request: "Request",
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+5
@@ -14,6 +14,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
14
from vllm.logger import init_logger
15
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
16
from vllm.v1.core.sched.output import SchedulerOutput
17
+from vllm.v1.outputs import KVConnectorOutput
18
19
if TYPE_CHECKING:
20
from vllm.attention.backends.abstract import AttentionMetadata
@@ -177,6 +178,10 @@ class MultiConnector(KVConnectorBase_V1):
178
self._extra_async_saves = {}
179
return metadata
180
181
+ def update_connector_output(self, connector_output: KVConnectorOutput):
182
+ for c in self._connectors:
183
+ c.update_connector_output(connector_output)
184
+
185
def request_finished(
186
self,
187
request: "Request",
vllm/v1/core/sched/scheduler.py
+4
@@ -1150,6 +1150,10 @@ class Scheduler(SchedulerInterface):
1150
# if finished_recving: add to state so we can
1151
scheduler the request during the next step.
1152
"""
1153
+
1154
+ assert self.connector is not None
1155
+ self.connector.update_connector_output(kv_connector_output)
1156
+
1157
# KV Connector:: update recv and send status from last step.
1158
for req_id in (kv_connector_output.finished_recving or ()):
1159
logger.debug("Finished recving KV transfer for request %s", req_id)