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
[BugFix] [P/D] Handle lookahead token count edge-case with Eagle Spec Decoding and P/D (#22317)
Signed-off-by: Pradyun Ramadorai <pradyunr@amazon.com> Signed-off-by: Pradyun92 <142861237+Pradyun92@users.noreply.github.com> Co-authored-by: Pradyun Ramadorai <pradyunr@amazon.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com>
Pradyun92 committed
Aug 8, 2025 at 20:04 UTC
35afe1b30b154114dc2ee8329e12f8cf3fe9f576
1 file changed
+11
-1
vllm/v1/core/sched/scheduler.py
+11
-1
@@ -437,14 +437,24 @@ class Scheduler(SchedulerInterface):
437
# The request cannot be scheduled.
438
break
439
440
+ # Handles an edge case when P/D Disaggregation
441
+ # is used with Spec Decoding where an
442
+ # extra block gets allocated which
443
+ # creates a mismatch between the number
444
+ # of local and remote blocks.
445
+ effective_lookahead_tokens = (0 if request.num_computed_tokens
446
+ == 0 else
447
+ self.num_lookahead_tokens)
448
+
449
new_blocks = self.kv_cache_manager.allocate_slots(
450
request,
451
num_new_tokens + num_external_computed_tokens,
452
num_new_local_computed_tokens,
453
new_computed_blocks,
445
- num_lookahead_tokens=self.num_lookahead_tokens,
454
+ num_lookahead_tokens=effective_lookahead_tokens,
455
delay_cache_blocks=load_kv_async,
456
)
457
+
458
if new_blocks is None:
459
# The request cannot be scheduled.
460
break