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 7, 05:15 UTC
[KVConnector] Scheduler: Fix num_computed_tokens after async KV load (#34616)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Mar 5, 2026 at 16:25 UTC
612e7729c2a548a7b6c9baa1821f419909777ffa
5 files changed
+32
-22
tests/v1/kv_connector/unit/test_error_propagation.py
+1
-1
@@ -121,7 +121,7 @@ def test_error_propagation_async_load(fail_scheduler: Scheduler):
121
122
assert len(fail_scheduler.waiting) == 1
123
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
124
- assert request.num_computed_tokens == 0
124
+ assert request.num_computed_tokens == num_external_computed_tokens
125
126
(req_block_ids,) = fail_scheduler.kv_cache_manager.get_block_ids(request.request_id)
127
invalid_block_ids = {req_block_ids[invalid_block_idx]}
tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py
+1
-1
@@ -339,7 +339,7 @@ def test_async_recompute_blocks_not_cached_when_invalid(
339
# request should be waiting for remote KVs
340
assert len(recompute_scheduler.waiting) == 1
341
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
342
- assert request.num_computed_tokens == 0
342
+ assert request.num_computed_tokens == num_external_computed_tokens
343
344
# get the allocated block IDs
345
(req_block_ids,) = recompute_scheduler.kv_cache_manager.get_block_ids(
tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py
+3
-3
@@ -78,7 +78,7 @@ def test_async_load_failure(
78
79
assert len(scheduler.waiting) == 3
80
for request in scheduler.waiting:
81
- assert request.num_computed_tokens == 0
81
+ assert request.num_computed_tokens == num_external_computed_tokens
82
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
83
assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
84
@@ -103,7 +103,7 @@ def test_async_load_failure(
103
min_invalid_block_idx * scheduler.block_size
104
)
105
else:
106
- assert request.num_computed_tokens == 0
106
+ assert request.num_computed_tokens == num_external_computed_tokens
107
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
108
assert scheduler.failed_recving_kv_req_ids == {request2.request_id}
109
assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
@@ -305,7 +305,7 @@ def test_async_progressive_load_failure(
305
306
assert len(scheduler.waiting) == 1
307
assert scheduler.waiting.peek_request().request_id == request.request_id
308
- assert request.num_computed_tokens == 0
308
+ assert request.num_computed_tokens == num_external_computed_tokens
309
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
310
assert scheduler.connector.get_num_new_matched_tokens.call_count == 1
311
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py
+1
-1
@@ -57,7 +57,7 @@ def test_basic_lifecycle():
57
assert len(scheduler.waiting) == 1
58
assert request in scheduler.waiting
59
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
60
- assert request.num_computed_tokens == 0
60
+ assert request.num_computed_tokens == NUM_TOKENS
61
62
# ... but should have (uncached) blocks allocated to it.
63
block_pool = scheduler.kv_cache_manager.block_pool
vllm/v1/core/sched/scheduler.py
+26
-16
@@ -638,6 +638,7 @@ class Scheduler(SchedulerInterface):
638
num_computed_tokens = (
639
num_new_local_computed_tokens + num_external_computed_tokens
640
)
641
+ assert num_computed_tokens <= request.num_tokens
642
else:
643
# KVTransfer: WAITING reqs have num_computed_tokens > 0
644
# after async KV recvs are completed.
@@ -773,6 +774,20 @@ class Scheduler(SchedulerInterface):
774
# into the WAITING_FOR_REMOTE_KV state.
775
skipped_waiting_requests.prepend_request(request)
776
request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
777
+ # Set num_computed_tokens even though KVs are not yet loaded.
778
+ # request.num_computed_tokens will not be used anywhere until
779
+ # the request finished the KV transfer.
780
+ #
781
+ # If a transfer error is reported by the connector,
782
+ # request.num_computed_tokens will be re-set accordingly in
783
+ # _update_requests_with_invalid_blocks.
784
+ #
785
+ # When the transfer is finished, either successfully or not,
786
+ # request.num_computed_tokens will correctly reflect the number
787
+ # of computed tokens.
788
+ # _update_waiting_for_remote_kv will then cache
789
+ # only the successfully loaded tokens.
790
+ request.num_computed_tokens = num_computed_tokens
791
continue
792
793
self.running.append(request)
@@ -1994,17 +2009,17 @@ class Scheduler(SchedulerInterface):
2009
self.failed_recving_kv_req_ids.remove(request.request_id)
2010
else:
2011
# Now that the blocks are ready, actually cache them.
1997
- (block_ids,) = self.kv_cache_manager.get_block_ids(request.request_id)
1998
- num_computed_tokens = len(block_ids) * self.block_size
1999
- # Handle the case where num request tokens less than one block.
2000
- num_computed_tokens = min(num_computed_tokens, request.num_tokens)
2001
- if num_computed_tokens == request.num_tokens:
2002
- num_computed_tokens -= 1
2012
# This will cache the blocks iff caching is enabled.
2004
- self.kv_cache_manager.cache_blocks(request, num_computed_tokens)
2013
+ self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens)
2014
2006
- # Update the request state for scheduling.
2007
- request.num_computed_tokens = num_computed_tokens
2015
+ # on a full prompt hit, we need to re-compute the last token
2016
+ # in order to be able to sample the next token
2017
+ if request.num_computed_tokens == request.num_tokens:
2018
+ request.num_computed_tokens = request.num_tokens - 1
2019
+
2020
+ # Count the number of prefix cached tokens.
2021
+ if request.num_cached_tokens < 0:
2022
+ request.num_cached_tokens = request.num_computed_tokens
2023
2024
# Return that we are ready.
2025
self.finished_recving_kv_req_ids.remove(request.request_id)
@@ -2084,13 +2099,8 @@ class Scheduler(SchedulerInterface):
2099
# We iterate only over blocks that may contain externally computed
2100
# tokens
2101
if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
2087
- # Async loading. If num_computed_tokens is set it implies we
2088
- # already processed some block failures for it in a prior step
2089
- req_num_computed_tokens = (
2090
- request.num_computed_tokens
2091
- if req_id in self.failed_recving_kv_req_ids
2092
- else len(req_block_ids) * self.block_size
2093
- )
2102
+ # Async loading. num_computed_tokens does not include new tokens
2103
+ req_num_computed_tokens = request.num_computed_tokens
2104
else:
2105
# Sync loading. num_computed_tokens includes new tokens
2106
req_num_computed_tokens = request.num_cached_tokens