diff --git a/tests/test_config.py b/tests/test_config.py index e0978bac7c..93bcdbc258 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -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 diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index be70da78b3..d8f049755c 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1448,19 +1448,44 @@ class SpeculativeConfig: @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 (