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 2, 17:01 UTC
[Bugfix][DFlash]allocate the proper number of lookahead slots (#43733)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com>
Benjamin Chislett committed
May 27, 2026 at 17:45 UTC
7fb9c0197a3173f2a2edcc9d64f6c0e73ef20717
1 file changed
+7
-1
vllm/v1/core/sched/scheduler.py
+7
-1
@@ -218,6 +218,11 @@ class Scheduler(SchedulerInterface):
218
self.num_lookahead_tokens = self.num_spec_tokens
219
if speculative_config.uses_draft_model():
220
self.num_lookahead_tokens = self.num_spec_tokens
221
+ if speculative_config.use_dflash():
222
+ # DFlash requires an extra lookahead slot since it uses in-fill-style
223
+ # decoding instead of standard next-token sampling, so it has a query
224
+ # for the last sampled token plus queries for each draft token.
225
+ self.num_lookahead_tokens = self.num_spec_tokens + 1
226
227
# Create the KV cache manager.
228
if hash_block_size is None:
@@ -702,8 +707,9 @@ class Scheduler(SchedulerInterface):
707
# extra block gets allocated which
708
# creates a mismatch between the number
709
# of local and remote blocks.
710
+ limit_lookahead_tokens = load_kv_async and self.use_eagle
711
effective_lookahead_tokens = (
706
- 0 if request.num_computed_tokens == 0 else self.num_lookahead_tokens
712
+ 0 if limit_lookahead_tokens else self.num_lookahead_tokens
713
)
714
715
# Determine if we need to allocate cross-attention blocks.