Skip to content
Merged
Show file tree
Hide file tree
Changes from all 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
24 changes: 24 additions & 0 deletions tests/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -490,6 +490,30 @@ def test_draft_model_enables_async_scheduling_by_default():
assert cfg.scheduler_config.async_scheduling is True


@pytest.mark.parametrize(
("method", "parallel_drafting", "expected_slots"),
[
pytest.param("eagle3", False, 0, id="eagle3"),
pytest.param("eagle3", True, 7, id="p-eagle"),
pytest.param("dflash", True, 8, id="dflash"),
pytest.param("dspark", True, 7, id="dspark"),
pytest.param("mtp", False, 0, id="mtp"),
pytest.param("ngram", False, 0, id="ngram"),
pytest.param("draft_model", False, 1, id="draft-model"),
pytest.param("draft_model", True, 8, id="pard"),
],
)
def test_max_num_new_slots_for_drafting(method, parallel_drafting, expected_slots):
speculative_config = SpeculativeConfig(
model="ngram",
num_speculative_tokens=8,
)
speculative_config.method = method
speculative_config.parallel_drafting = parallel_drafting

assert speculative_config.max_num_new_slots_for_drafting == expected_slots


@dataclass
class _TestConfigFields:
a: int
Expand Down
45 changes: 35 additions & 10 deletions vllm/config/speculative.py
Original file line number Diff line number Diff line change
Expand Up @@ -1437,19 +1437,44 @@ def verify_equal_vocab_size_if_draft_model(self):

@property
def max_num_new_slots_for_drafting(self) -> int:
"""Return the maximum additional drafting slots per request.

The scheduler budget already includes one query slot per decoding request.
Let K be ``num_speculative_tokens``. Standard configurations require:

==================== ============= ======== ================
Algorithm Method Parallel Additional slots
==================== ============= ======== ================
EAGLE3 eagle3 No 0
P-EAGLE eagle3 Yes K - 1
DFlash dflash Yes K
DSpark dspark Yes K - 1
MTP mtp No 0
N-gram ngram No 0
Draft model draft_model No 1
PARD draft_model Yes K
==================== ============= ======== ================
"""
Calculate the maximum number of new slots that might be added to the batch
when drafting.
"""
slots_per_req = 0 # for serial non-draft-model methods, no change needed
num_draft_tokens = self.num_speculative_tokens

if self.use_dflash():
# DFlash uses one bonus query followed by K mask queries.
return num_draft_tokens

if self.parallel_drafting:
# For parallel drafting, we need one new slot per 'masked' token
slots_per_req = self.num_speculative_tokens - 1
if self.uses_draft_model():
# PARD does not shift the existing input, so all K query
# positions require additional slots.
return num_draft_tokens

# The existing query is reused; only masked queries need new slots.
return num_draft_tokens - 1

if self.uses_draft_model():
# For draft model-based speculation, we need one new slot per request
# Since we do not slice the draft tokens
slots_per_req += 1
return slots_per_req
# The autoregressive draft-model input retains one unsliced token.
return 1

return 0

def use_gemma4_mtp(self) -> bool:
return (
Expand Down
Loading