Skip to content
Closed
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions vllm/v1/core/sched/scheduler.py
Original file line number Diff line number Diff line change
Expand Up @@ -210,6 +210,7 @@ def __init__(

speculative_config = vllm_config.speculative_config
self.use_eagle = False
self.use_dflash = False
self.num_spec_tokens = self.num_lookahead_tokens = 0
if speculative_config:
self.num_spec_tokens = speculative_config.num_speculative_tokens
Expand All @@ -219,6 +220,7 @@ def __init__(
if speculative_config.uses_draft_model():
self.num_lookahead_tokens = self.num_spec_tokens
if speculative_config.use_dflash():
self.use_dflash = True
# 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.
Expand Down Expand Up @@ -707,9 +709,14 @@ def schedule(self) -> SchedulerOutput:
# 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 limit_lookahead_tokens else self.num_lookahead_tokens
0
if (
self.use_eagle
and not self.use_dflash
and request.num_computed_tokens == 0
)
else self.num_lookahead_tokens
)

# Determine if we need to allocate cross-attention blocks.
Expand Down
Loading