From a348d25ce4c3ebad90130153cf46323a792f9ba8 Mon Sep 17 00:00:00 2001 From: HF-001 <1670186653@qq.com> Date: Thu, 6 Aug 2026 18:39:02 +0800 Subject: [PATCH 1/4] [BugFix] fix dflash max_num_scheduled_tokens Signed-off-by: HF-001 <1670186653@qq.com> --- tests/test_config.py | 11 +++++++++++ vllm/config/speculative.py | 3 +++ 2 files changed, 14 insertions(+) diff --git a/tests/test_config.py b/tests/test_config.py index febc9dab2822..8377ffc4ea14 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -430,6 +430,17 @@ def test_draft_model_enables_async_scheduling_by_default(): assert cfg.scheduler_config.async_scheduling is True +def test_dflash_max_num_new_slots_for_drafting(): + speculative_config = SpeculativeConfig( + model="ngram", + num_speculative_tokens=8, + ) + speculative_config.method = "dflash" + speculative_config.parallel_drafting = True + + assert speculative_config.max_num_new_slots_for_drafting == 8 + + @dataclass class _TestConfigFields: a: int diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 3a084acfc099..85496cfae704 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1408,6 +1408,9 @@ def max_num_new_slots_for_drafting(self) -> int: 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.use_dflash(): + # DFlash adds a bonus query before the masked tokens. + slots_per_req += 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 From 94a97a0e5d801f339dd051705bed569b105bdeb4 Mon Sep 17 00:00:00 2001 From: HF-001 <1670186653@qq.com> Date: Fri, 7 Aug 2026 09:09:33 +0800 Subject: [PATCH 2/4] enrich test config cases for dflash slots Signed-off-by: HF-001 <1670186653@qq.com> --- tests/test_config.py | 13 ++++++++++--- vllm/config/speculative.py | 9 +++------ 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 8377ffc4ea14..ab21be62d1b7 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -430,15 +430,22 @@ def test_draft_model_enables_async_scheduling_by_default(): assert cfg.scheduler_config.async_scheduling is True -def test_dflash_max_num_new_slots_for_drafting(): +@pytest.mark.parametrize( + ("method", "expected_slots"), + [ + ("eagle3", 7), + ("dflash", 8), + ], +) +def test_parallel_drafting_max_num_new_slots(method, expected_slots): speculative_config = SpeculativeConfig( model="ngram", num_speculative_tokens=8, ) - speculative_config.method = "dflash" + speculative_config.method = method speculative_config.parallel_drafting = True - assert speculative_config.max_num_new_slots_for_drafting == 8 + assert speculative_config.max_num_new_slots_for_drafting == expected_slots @dataclass diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index eeadb829dede..983ac65e2171 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1421,12 +1421,9 @@ def max_num_new_slots_for_drafting(self) -> int: 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.use_dflash(): - # DFlash adds a bonus query before the masked tokens. - slots_per_req += 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 + if self.uses_draft_model() or self.use_dflash(): + # Draft models do not slice the draft tokens, while DFlash adds a + # bonus query before the masked tokens. slots_per_req += 1 return slots_per_req From 79c0215f4882ee29020f909bdb2a5505889a55b6 Mon Sep 17 00:00:00 2001 From: HF-001 <1670186653@qq.com> Date: Wed, 12 Aug 2026 14:20:05 +0800 Subject: [PATCH 3/4] optimize max_num_new_slots_for_drafting and test cases Signed-off-by: HF-001 <1670186653@qq.com> --- tests/test_config.py | 16 ++++++++----- vllm/config/speculative.py | 46 +++++++++++++++++++++++++++++--------- 2 files changed, 46 insertions(+), 16 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index ab21be62d1b7..1f06287f2c37 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -431,19 +431,25 @@ def test_draft_model_enables_async_scheduling_by_default(): @pytest.mark.parametrize( - ("method", "expected_slots"), + ("method", "parallel_drafting", "expected_slots"), [ - ("eagle3", 7), - ("dflash", 8), + 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_parallel_drafting_max_num_new_slots(method, expected_slots): +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 = True + speculative_config.parallel_drafting = parallel_drafting assert speculative_config.max_num_new_slots_for_drafting == expected_slots diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 983ac65e2171..fe219c354473 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1413,19 +1413,43 @@ 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 already reserves one slot for the request's next token. + 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() or self.use_dflash(): - # Draft models do not slice the draft tokens, while DFlash adds a - # bonus query before the masked tokens. - slots_per_req += 1 - return slots_per_req + if self.uses_draft_model(): + # PARD retains the input query instead of reusing its slot. + 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(): + # The autoregressive draft-model input retains one unsliced token. + return 1 + + return 0 def use_gemma4_mtp(self) -> bool: return ( From 96d40d056443248f791d9c402ce41c6d2bc753bd Mon Sep 17 00:00:00 2001 From: HF-001 <1670186653@qq.com> Date: Wed, 12 Aug 2026 14:34:29 +0800 Subject: [PATCH 4/4] optmize annotation Signed-off-by: HF-001 <1670186653@qq.com> --- vllm/config/speculative.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index fe219c354473..bb65f38d45bf 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1415,7 +1415,7 @@ def verify_equal_vocab_size_if_draft_model(self): def max_num_new_slots_for_drafting(self) -> int: """Return the maximum additional drafting slots per request. - The scheduler already reserves one slot for the request's next token. + The scheduler budget already includes one query slot per decoding request. Let K be ``num_speculative_tokens``. Standard configurations require: ==================== ============= ======== ================ @@ -1439,7 +1439,8 @@ def max_num_new_slots_for_drafting(self) -> int: if self.parallel_drafting: if self.uses_draft_model(): - # PARD retains the input query instead of reusing its slot. + # 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.