diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 5911859c9d7..6ce9dcc0756 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -218,6 +218,11 @@ class Scheduler(SchedulerInterface): self.num_lookahead_tokens = self.num_spec_tokens if speculative_config.uses_draft_model(): self.num_lookahead_tokens = self.num_spec_tokens + if speculative_config.use_dflash(): + # DFlash requires an extra lookahead slot since it uses in-fill-style + # decoding instead of standard next-token sampling, so it has a query + # for the last sampled token plus queries for each draft token. + self.num_lookahead_tokens = self.num_spec_tokens + 1 # Create the KV cache manager. if hash_block_size is None: @@ -702,8 +707,9 @@ class Scheduler(SchedulerInterface): # extra block gets allocated which # creates a mismatch between the number # of local and remote blocks. + limit_lookahead_tokens = load_kv_async and self.use_eagle effective_lookahead_tokens = ( - 0 if request.num_computed_tokens == 0 else self.num_lookahead_tokens + 0 if limit_lookahead_tokens else self.num_lookahead_tokens ) # Determine if we need to allocate cross-attention blocks.