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, 07:15 UTC
[Docs] Improve V1 KVConnector interface documentation (#19172)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Jun 6, 2025 at 01:22 UTC
65c69444b1b7d5da9ec5482ca27e32a752cd0a29
2 files changed
+32
-9
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+26
-3
@@ -8,9 +8,15 @@ The class provides the following primitives:
8
Scheduler-side: runs in the scheduler, binds metadata, which
9
is used by the worker-side to load/save KV cache.
10
get_num_new_matched_tokens() - get number of new tokens
11
- that exist in the remote KV cache
11
+ that exist in the remote KV cache. Might be called multiple
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
+ request_finished() - called when a request is finished, with
16
+ the computed kv cache blocks for the request.
17
+ Returns whether KV cache should be freed now or will be
18
+ freed asynchronously and optionally returns KV transfer
19
+ params.
20
21
Worker-side: runs in each worker, loads/saves KV cache to/from
22
the Connector based on the metadata.
@@ -19,6 +25,9 @@ The class provides the following primitives:
25
26
save_kv_layer() - starts saving KV for layer i (maybe async)
27
wait_for_save() - blocks until all saves are done
28
+
29
+ get_finished() - called with ids of finished requests, returns
30
+ ids of requests that have completed async sending/recving.
31
"""
32
33
import enum
@@ -184,7 +193,8 @@ class KVConnectorBase_V1(ABC):
193
finished generating tokens.
194
195
Returns:
187
- ids of requests that have finished asynchronous transfer,
196
+ ids of requests that have finished asynchronous transfer
197
+ (requests that previously returned True from request_finished()),
198
tuple of (sending/saving ids, recving/loading ids).
199
The finished saves/sends req ids must belong to a set provided in a
200
call to this method (this call or a prior one).
@@ -215,7 +225,8 @@ class KVConnectorBase_V1(ABC):
225
- The number of tokens that can be loaded from the
226
external KV cache beyond what is already computed.
227
- `True` if external KV cache tokens will be loaded
218
- asynchronously (between scheduler steps).
228
+ asynchronously (between scheduler steps). Must be
229
+ 'False' if the first element is 0.
230
"""
231
pass
232
@@ -225,6 +236,18 @@ class KVConnectorBase_V1(ABC):
236
num_external_tokens: int):
237
"""
238
Update KVConnector state after block allocation.
239
+
240
+ If get_num_new_matched_tokens previously returned True for a
241
+ request, this function may be called twice for that same request -
242
+ first when blocks are allocated for the connector tokens to be
243
+ asynchronously loaded into, and second when any additional blocks
244
+ are allocated, after the load/transfer is complete.
245
+
246
+ Args:
247
+ request (Request): the request object.
248
+ blocks (KVCacheBlocks): the blocks allocated for the request.
249
+ num_external_tokens (int): the number of tokens that will be
250
+ loaded from the external KV cache.
251
"""
252
pass
253
vllm/v1/core/sched/scheduler.py
+6
-6
@@ -101,7 +101,7 @@ class Scheduler(SchedulerInterface):
101
# This is flushed at the end of each scheduling step.
102
self.finished_req_ids: set[str] = set()
103
104
- # P/D: requests in process of recving KV transfers
104
+ # KV Connector: requests in process of async KV loading or recving
105
self.finished_recving_kv_req_ids: set[str] = set()
106
107
# OPTIMIZATION: Cache the CachedRequestData objects to avoid creating
@@ -822,7 +822,7 @@ class Scheduler(SchedulerInterface):
822
if not stopped:
823
new_running.append(request)
824
825
- # P/D: update state for finished KV Transfers.
825
+ # KV Connector: update state for finished KV Transfers.
826
self._update_from_kv_xfer_finished(model_runner_output)
827
828
# Return the cached request data to the queue so they can be reused.
@@ -969,7 +969,7 @@ class Scheduler(SchedulerInterface):
969
self.kv_event_publisher.shutdown()
970
971
########################################################################
972
- # P/D Related Methods
972
+ # KV Connector Related Methods
973
########################################################################
974
975
def get_kv_connector(self) -> Optional[KVConnectorBase_V1]:
@@ -992,7 +992,7 @@ class Scheduler(SchedulerInterface):
992
993
def _update_waiting_for_remote_kv(self, request: Request) -> bool:
994
"""
995
- P/D: check if the request_id is finished_recving.
995
+ KV Connector: check if the request_id is finished_recving.
996
997
The finished_recving_kv_req_ids list is populated
998
on the previous steps()'s update_from_output based
@@ -1029,7 +1029,7 @@ class Scheduler(SchedulerInterface):
1029
def _update_from_kv_xfer_finished(self,
1030
model_runner_output: ModelRunnerOutput):
1031
"""
1032
- P/D: update the scheduler state based on the output.
1032
+ KV Connector: update the scheduler state based on the output.
1033
1034
The Worker side connectors add finished_recving and
1035
finished_sending reqs to the output.
@@ -1037,7 +1037,7 @@ class Scheduler(SchedulerInterface):
1037
# if finished_recving: add to state so we can
1038
scheduler the request during the next step.
1039
"""
1040
- # P/D: update recv and send status from last step.
1040
+ # KV Connector:: update recv and send status from last step.
1041
for req_id in (model_runner_output.finished_recving or ()):
1042
logger.debug("Finished recving KV transfer for request %s", req_id)
1043
self.finished_recving_kv_req_ids.add(req_id)