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
[KV Connector] More async support for `get_num_new_matched_tokens` (#23620)
Signed-off-by: ApostaC <yihua98@uchicago.edu>
Yihua Cheng committed
Sep 9, 2025 at 21:23 UTC
b4a01aaf95f54bba90cb0b072e9254ddb998af8f
6 files changed
+23
-8
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+6
-3
@@ -243,7 +243,7 @@ class KVConnectorBase_V1(ABC):
243
self,
244
request: "Request",
245
num_computed_tokens: int,
246
- ) -> tuple[int, bool]:
246
+ ) -> tuple[Optional[int], bool]:
247
"""
248
Get number of new tokens that can be loaded from the
249
external KV cache beyond the num_computed_tokens.
@@ -255,8 +255,11 @@ class KVConnectorBase_V1(ABC):
255
256
Returns:
257
A tuple with the following elements:
258
- - The number of tokens that can be loaded from the
259
- external KV cache beyond what is already computed.
258
+ - An optional number of tokens that can be loaded from the
259
+ external KV cache beyond what is already computed.
260
+ If None, it means that the connector needs more time to
261
+ determine the number of matched tokens, and the scheduler
262
+ should query for this request again later.
263
- `True` if external KV cache tokens will be loaded
264
asynchronously (between scheduler steps). Must be
265
'False' if the first element is 0.
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py
+1
-1
@@ -110,7 +110,7 @@ class LMCacheConnectorV1(KVConnectorBase_V1):
110
self,
111
request: "Request",
112
num_computed_tokens: int,
113
- ) -> tuple[int, bool]:
113
+ ) -> tuple[Optional[int], bool]:
114
"""
115
Get number of new tokens that can be loaded from the
116
external KV cache beyond the num_computed_tokens.
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+5
-1
@@ -143,11 +143,15 @@ class MultiConnector(KVConnectorBase_V1):
143
self,
144
request: "Request",
145
num_computed_tokens: int,
146
- ) -> tuple[int, bool]:
146
+ ) -> tuple[Optional[int], bool]:
147
to_return = (0, False)
148
for i, c in enumerate(self._connectors):
149
toks, load_async = c.get_num_new_matched_tokens(
150
request, num_computed_tokens)
151
+ # If there is a connector still looking up the matches,
152
+ # we return None to indicate that we are not done yet.
153
+ if toks is None:
154
+ return (None, False)
155
# The first connector that has new matched tokens will be assigned
156
# to this request.
157
if to_return[0] == 0 and toks > 0:
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+1
-1
@@ -162,7 +162,7 @@ class NixlConnector(KVConnectorBase_V1):
162
163
def get_num_new_matched_tokens(
164
self, request: "Request",
165
- num_computed_tokens: int) -> tuple[int, bool]:
165
+ num_computed_tokens: int) -> tuple[Optional[int], bool]:
166
assert self.connector_scheduler is not None
167
return self.connector_scheduler.get_num_new_matched_tokens(
168
request, num_computed_tokens)
vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py
+2
-2
@@ -3,7 +3,7 @@
3
import hashlib
4
import os
5
from dataclasses import dataclass
6
-from typing import TYPE_CHECKING
6
+from typing import TYPE_CHECKING, Optional
7
8
import safetensors
9
import torch
@@ -238,7 +238,7 @@ class SharedStorageConnector(KVConnectorBase_V1):
238
self,
239
request: "Request",
240
num_computed_tokens: int,
241
- ) -> tuple[int, bool]:
241
+ ) -> tuple[Optional[int], bool]:
242
"""
243
Get number of new tokens that can be loaded from the
244
external KV cache beyond the num_computed_tokens.
vllm/v1/core/sched/scheduler.py
+8
@@ -387,6 +387,14 @@ class Scheduler(SchedulerInterface):
387
self.connector.get_num_new_matched_tokens(
388
request, num_new_local_computed_tokens))
389
390
+ if num_external_computed_tokens is None:
391
+ # The request cannot be scheduled because
392
+ # the KVConnector couldn't determine
393
+ # the number of matched tokens.
394
+ self.waiting.pop_request()
395
+ skipped_waiting_requests.prepend_request(request)
396
+ continue
397
+
398
# Total computed tokens (local + external).
399
num_computed_tokens = (num_new_local_computed_tokens +
400
num_external_computed_tokens)