From 7cd7982dfda0d29c68f73c54cd343b674b290a47 Mon Sep 17 00:00:00 2001 From: qizixi <22851944+zixi-qi@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:23:58 -0700 Subject: [PATCH] [Spec Decode] Don't reserve DFlash/DSpark draft slots in the scheduler token budget (#59468) Signed-off-by: zixi-qi Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/test_config.py | 4 ++-- tests/v1/core/test_scheduler.py | 40 +++++++++++++++++++++++++++++++++ vllm/config/speculative.py | 12 +++++----- vllm/config/vllm.py | 29 +++++++++++++++++++----- 4 files changed, 73 insertions(+), 12 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 0e65d87150..6558400552 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -2095,8 +2095,8 @@ def test_draft_parallel_config_preserves_ep_without_model( [ 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("dflash", True, 0, id="dflash"), + pytest.param("dspark", True, 0, 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"), diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 0d7897edc1..7b2b5b13b6 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -2159,6 +2159,46 @@ def test_draft_slots_budgeted_per_scheduled_request(tmp_path, monkeypatch): assert scheduler.schedule().num_scheduled_tokens == {"0": 10, "1": 4} +@pytest.mark.parametrize(("method", "max_num_reqs"), [("dflash", 4), ("dspark", 5)]) +def test_block_drafters_cap_requests_instead_of_draft_slots( + tmp_path, method, max_num_reqs +): + (tmp_path / "config.json").write_text( + '{"architectures": ["OPTForCausalLM"], "model_type": "opt"}' + ) + base = create_scheduler( + model=str(tmp_path), + max_num_batched_tokens=20, + num_speculative_tokens=4, + parallel_drafting=True, + skip_tokenizer_init=True, + ) + # DFlash/DSpark need a resolvable draft model, so switch the ngram config in + # place and redo the config step. DSpark then has 4 lookahead tokens, not 5. + vllm_config = base.vllm_config + object.__setattr__(vllm_config.speculative_config, "method", method) + object.__setattr__( + vllm_config.speculative_config, + "draft_model_config", + SimpleNamespace(hf_config=SimpleNamespace()), + ) + vllm_config._set_max_num_scheduled_tokens() + scheduler = Scheduler( + vllm_config=vllm_config, + kv_cache_config=base.kv_cache_config, + structured_output_manager=base.structured_output_manager, + block_size=base.block_size, + ) + for request in create_requests(num_requests=8, num_tokens=2): + scheduler.add_request(request) + + # Admission stops at 20 // num_lookahead_tokens, and no draft slots are + # charged: each admitted prompt is scheduled in full. + assert scheduler.schedule().num_scheduled_tokens == { + str(i): 2 for i in range(max_num_reqs) + } + + # Note - these test cases mirror some of those in test_rejection_sampler.py @pytest.mark.parametrize( "spec_tokens,output_tokens,expected,expected_per_req", diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index c80016fd36..6b7e661165 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1853,19 +1853,21 @@ class SpeculativeConfig: ==================== ============= ======== ================ EAGLE3 eagle3 No 0 P-EAGLE eagle3 Yes K - 1 - DFlash dflash Yes K - DSpark dspark Yes K - 1 + DFlash dflash Yes 0 + DSpark dspark Yes 0 MTP mtp No 0 N-gram ngram No 0 Draft model draft_model No 1 PARD draft_model Yes K ==================== ============= ======== ================ + + DFlash and DSpark draft in a separate query pass rather than adding + rows to the target batch. """ 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.use_dflash() or self.use_dspark(): + return 0 if self.parallel_drafting: if self.uses_draft_model(): diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 837bf1e80a..608652e08c 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2300,9 +2300,6 @@ class VllmConfig: worker is configured to handle. """ if self.speculative_config is not None: - scheduled_token_delta = ( - self.speculative_config.max_num_new_slots_for_drafting - ) max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens if self.scheduler_config.max_num_scheduled_tokens is None: self.scheduler_config.max_num_scheduled_tokens = max_num_batched_tokens @@ -2326,12 +2323,34 @@ class VllmConfig: " num_speculative_tokens.", ) - if max_num_batched_tokens <= scheduled_token_delta: + # num_lookahead_tokens is at least every method's per-request draft + # slots, and covers the DFlash/DSpark draft pass's query rows. + if max_num_batched_tokens <= self.num_lookahead_tokens: raise ValueError( "VllmConfig does not have enough slots to schedule a token and" " support the speculative decoding settings." - f" Got {max_num_batched_tokens=} and {scheduled_token_delta=}." + f" Got {max_num_batched_tokens=} and" + f" num_lookahead_tokens={self.num_lookahead_tokens}." ) + spec = self.speculative_config + if spec.use_dflash() or spec.use_dspark(): + max_num_active_seqs = ( + self.scheduler_config.max_num_active_seqs + or self.scheduler_config.max_num_seqs + ) + max_draft_seqs = max_num_batched_tokens // self.num_lookahead_tokens + if max_draft_seqs < max_num_active_seqs: + logger.warning_once( + "The %s draft pass needs %d query tokens per request within " + "max_num_batched_tokens=%d. Overriding max_num_active_seqs " + "from %d to %d.", + spec.method, + self.num_lookahead_tokens, + max_num_batched_tokens, + max_num_active_seqs, + max_draft_seqs, + ) + self.scheduler_config.max_num_active_seqs = max_draft_seqs def _set_cudagraph_sizes(self): """VLLM defines the default candidate list of batch sizes for CUDA graph