[Spec Decode] Don't reserve DFlash/DSpark draft slots in the scheduler token budget (#59468)

Signed-off-by: zixi-qi <zixi@inferact.ai>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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"),

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",

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():

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