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 6, 00:45 UTC
[Core][KV Connector] Remove use of num_cached_tokens in error handling (#38096)
Signed-off-by: Mark McLoughlin <markmc@redhat.com>
Mark McLoughlin committed
Mar 25, 2026 at 18:20 UTC
e38817fadbe10393ec8ef2fd52fe9a086d7365c1
1 file changed
+17
-12
vllm/v1/core/sched/scheduler.py
+17
-12
@@ -1326,7 +1326,8 @@ class Scheduler(SchedulerInterface):
1326
# load. Identify affected requests and adjust their computed token
1327
# count to trigger recomputation of the invalid blocks.
1328
failed_kv_load_req_ids = self._handle_invalid_blocks(
1329
- kv_connector_output.invalid_block_ids
1329
+ kv_connector_output.invalid_block_ids,
1330
+ num_scheduled_tokens,
1331
)
1332
1333
# NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
@@ -2133,6 +2134,7 @@ class Scheduler(SchedulerInterface):
2134
self,
2135
requests: Iterable[Request],
2136
invalid_block_ids: set[int],
2137
+ num_scheduled_tokens: dict[str, int],
2138
evict_blocks: bool = True,
2139
) -> tuple[set[str], int, set[int]]:
2140
"""
@@ -2146,6 +2148,7 @@ class Scheduler(SchedulerInterface):
2148
Args:
2149
requests: The set of requests to scan for invalid blocks.
2150
invalid_block_ids: IDs of invalid blocks.
2151
+ num_scheduled_tokens: req_id -> number of scheduled tokens.
2152
evict_blocks: Whether to collect blocks for eviction (False for
2153
async requests which aren't cached yet).
2154
@@ -2173,12 +2176,9 @@ class Scheduler(SchedulerInterface):
2176
(req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id)
2177
# We iterate only over blocks that may contain externally computed
2178
# tokens
2176
- if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
2177
- # Async loading. num_computed_tokens does not include new tokens
2178
- req_num_computed_tokens = request.num_computed_tokens
2179
- else:
2180
- # Sync loading. num_computed_tokens includes new tokens
2181
- req_num_computed_tokens = request.num_cached_tokens
2179
+ req_num_computed_tokens = (
2180
+ request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0)
2181
+ )
2182
2183
req_num_computed_blocks = (
2184
req_num_computed_tokens + self.block_size - 1
@@ -2225,15 +2225,17 @@ class Scheduler(SchedulerInterface):
2225
# Currently this only applies to sync loading; Async
2226
# loading does not yet support block sharing
2227
total_affected_tokens += (
2228
- request.num_computed_tokens - request.num_cached_tokens
2228
+ request.num_computed_tokens - req_num_computed_tokens
2229
)
2230
- request.num_computed_tokens = request.num_cached_tokens
2230
+ request.num_computed_tokens = req_num_computed_tokens
2231
2232
affected_req_ids.add(request.request_id)
2233
2234
return affected_req_ids, total_affected_tokens, blocks_to_evict
2235
2236
- def _handle_invalid_blocks(self, invalid_block_ids: set[int]) -> set[str]:
2236
+ def _handle_invalid_blocks(
2237
+ self, invalid_block_ids: set[int], num_scheduled_tokens: dict[str, int]
2238
+ ) -> set[str]:
2239
"""
2240
Handle requests affected by invalid KV cache blocks.
2241
@@ -2250,7 +2252,10 @@ class Scheduler(SchedulerInterface):
2252
)
2253
async_failed_req_ids, num_failed_tokens, _ = (
2254
self._update_requests_with_invalid_blocks(
2253
- async_load_reqs, invalid_block_ids, evict_blocks=False
2255
+ async_load_reqs,
2256
+ invalid_block_ids,
2257
+ num_scheduled_tokens,
2258
+ evict_blocks=False,
2259
)
2260
)
2261
@@ -2260,7 +2265,7 @@ class Scheduler(SchedulerInterface):
2265
# handle sync loads (may be cached, collect blocks for eviction)
2266
sync_failed_req_ids, num_failed_tokens, sync_blocks_to_evict = (
2267
self._update_requests_with_invalid_blocks(
2263
- self.running, invalid_block_ids, evict_blocks=True
2268
+ self.running, invalid_block_ids, num_scheduled_tokens, evict_blocks=True
2269
)
2270
)
2271