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, 09:45 UTC
[BugFix] Fix handling of num_computed_tokens with connector (#18232)
Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com>
Nick Hill committed
May 19, 2025 at 09:03 UTC
1b15df2546e97c409668da92954d8802c48d13af
2 files changed
+29
-16
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+11
-5
@@ -209,7 +209,17 @@ class NixlConnectorScheduler:
209
rounded_num_prompt_tokens = round_down(
210
len(request.prompt_token_ids), self.block_size)
211
count = max(rounded_num_prompt_tokens - num_computed_tokens, 0)
212
- return count, count > 0
212
+ if count > 0:
213
+ return count, True
214
+
215
+ # NOTE: if count is 0 here, we have less than block_size
216
+ # tokens to pull after subtracting the local prefix cache hit.
217
+ # The remote only sends fully computed blocks, so there is
218
+ # nothing to transfer but we still need to notify the
219
+ # prefill worker so that the remote blocks are freed.
220
+ if all(p in params for p in ("remote_engine_id", "remote_host",
221
+ "remote_port")):
222
+ self._reqs_need_recv[request.request_id] = (request, [])
223
224
# No remote prefill for this request.
225
return 0, False
@@ -225,10 +235,6 @@ class NixlConnectorScheduler:
235
num_external_tokens, params)
236
237
if params is not None and params.get("do_remote_prefill"):
228
- # NOTE(rob): if prompt < block_size, no remote blocks
229
- # since the remote only sends fully computed blocks, so
230
- # skip recving for this request. num_external_tokens
231
- # should be 0 if there are no remote blocks.
238
if params.get("remote_block_ids"):
239
if all(p in params for p in ("remote_engine_id", "remote_host",
240
"remote_port")):
vllm/v1/core/sched/scheduler.py
+18
-11
@@ -345,32 +345,38 @@ class Scheduler(SchedulerInterface):
345
skipped_waiting_requests.appendleft(request)
346
continue
347
348
+ num_external_computed_tokens = 0
349
+ load_kv_async = False
350
+
351
# Get already-cached tokens.
352
if num_prealloc_computed_tokens == 0:
353
new_computed_blocks, num_native_computed_tokens = \
354
self.kv_cache_manager.get_computed_blocks(
355
request)
356
+
357
+ # Get externally-cached tokens if using a KVConnector.
358
+ if self.connector is not None:
359
+ num_external_computed_tokens, load_kv_async = (
360
+ self.connector.get_num_new_matched_tokens(
361
+ request, num_native_computed_tokens))
362
+
363
+ # Total computed tokens (local + external).
364
+ num_computed_tokens = (num_native_computed_tokens +
365
+ num_external_computed_tokens)
366
else:
367
# P/D: skip checking prefix cache if loaded from remote kvs.
368
new_computed_blocks = KVCacheBlocks.create_empty()
369
num_native_computed_tokens = 0
370
358
- # Get externally-cached tokens if using a KVConnector.
359
- num_external_computed_tokens, load_kv_async = (
360
- (0, False) if self.connector is None else
361
- self.connector.get_num_new_matched_tokens(
362
- request, num_native_computed_tokens))
363
-
364
- # Total computed tokens (local + external).
365
- num_computed_tokens = (num_native_computed_tokens +
366
- num_external_computed_tokens +
367
- num_prealloc_computed_tokens)
371
+ # Total computed tokens (allocated in prior step).
372
+ num_computed_tokens = num_prealloc_computed_tokens
373
374
encoder_inputs_to_schedule = None
375
new_encoder_budget = encoder_budget
376
377
# P/D: loading remote KV, do not allocate for new work.
378
if load_kv_async:
379
+ assert num_external_computed_tokens > 0
380
num_new_tokens = 0
381
# Number of tokens to be scheduled.
382
else:
@@ -411,7 +417,8 @@ class Scheduler(SchedulerInterface):
417
# KVConnector: update internal state after allocation.
418
# This information is used to determine if a load is
419
# needed for this request.
414
- if self.connector is not None:
420
+ if num_external_computed_tokens:
421
+ assert self.connector is not None
422
self.connector.update_state_after_alloc(
423
request,
424
new_computed_blocks + new_blocks,