From 7fb9c0197a3173f2a2edcc9d64f6c0e73ef20717 Mon Sep 17 00:00:00 2001 From: Benjamin Chislett Date: Wed, 27 May 2026 17:45:34 -0400 Subject: [PATCH] [Bugfix][DFlash]allocate the proper number of lookahead slots (#43733) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Benjamin Chislett Signed-off-by: Benjamin Chislett Co-authored-by: Nicolò Lucchesi --- vllm/v1/core/sched/scheduler.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) 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.