[BugFix] Reserve the bonus query slot in DFlash scheduling budget (#51256)
Signed-off-by: HF-001 <1670186653@qq.com>
kx committed
Aug 13, 2026 at 11:50 UTC
2ac1f683f11cf0f526c896e5fa5e0e22ebc6ffa2
2 files changed
+59
-10
tests/test_config.py
+24
@@ -490,6 +490,30 @@ def test_draft_model_enables_async_scheduling_by_default():
490
assert cfg.scheduler_config.async_scheduling is True
491
492
493
+@pytest.mark.parametrize(
494
+ ("method", "parallel_drafting", "expected_slots"),
495
+ [
496
+ pytest.param("eagle3", False, 0, id="eagle3"),
497
+ pytest.param("eagle3", True, 7, id="p-eagle"),
498
+ pytest.param("dflash", True, 8, id="dflash"),
499
+ pytest.param("dspark", True, 7, id="dspark"),
500
+ pytest.param("mtp", False, 0, id="mtp"),
501
+ pytest.param("ngram", False, 0, id="ngram"),
502
+ pytest.param("draft_model", False, 1, id="draft-model"),
503
+ pytest.param("draft_model", True, 8, id="pard"),
504
+ ],
505
+)
506
+def test_max_num_new_slots_for_drafting(method, parallel_drafting, expected_slots):
507
+ speculative_config = SpeculativeConfig(
508
+ model="ngram",
509
+ num_speculative_tokens=8,
510
+ )
511
+ speculative_config.method = method
512
+ speculative_config.parallel_drafting = parallel_drafting
513
+
514
+ assert speculative_config.max_num_new_slots_for_drafting == expected_slots
515
+
516
+
517
@dataclass
518
class _TestConfigFields:
519
a: int
vllm/config/speculative.py
+35
-10
@@ -1448,19 +1448,44 @@ class SpeculativeConfig:
1448
1449
@property
1450
def max_num_new_slots_for_drafting(self) -> int:
1451
+ """Return the maximum additional drafting slots per request.
1452
+
1453
+ The scheduler budget already includes one query slot per decoding request.
1454
+ Let K be ``num_speculative_tokens``. Standard configurations require:
1455
+
1456
+ ==================== ============= ======== ================
1457
+ Algorithm Method Parallel Additional slots
1458
+ ==================== ============= ======== ================
1459
+ EAGLE3 eagle3 No 0
1460
+ P-EAGLE eagle3 Yes K - 1
1461
+ DFlash dflash Yes K
1462
+ DSpark dspark Yes K - 1
1463
+ MTP mtp No 0
1464
+ N-gram ngram No 0
1465
+ Draft model draft_model No 1
1466
+ PARD draft_model Yes K
1467
+ ==================== ============= ======== ================
1468
"""
1452
- Calculate the maximum number of new slots that might be added to the batch
1453
- when drafting.
1454
- """
1455
- slots_per_req = 0 # for serial non-draft-model methods, no change needed
1469
+ num_draft_tokens = self.num_speculative_tokens
1470
+
1471
+ if self.use_dflash():
1472
+ # DFlash uses one bonus query followed by K mask queries.
1473
+ return num_draft_tokens
1474
+
1475
if self.parallel_drafting:
1457
- # For parallel drafting, we need one new slot per 'masked' token
1458
- slots_per_req = self.num_speculative_tokens - 1
1476
+ if self.uses_draft_model():
1477
+ # PARD does not shift the existing input, so all K query
1478
+ # positions require additional slots.
1479
+ return num_draft_tokens
1480
+
1481
+ # The existing query is reused; only masked queries need new slots.
1482
+ return num_draft_tokens - 1
1483
+
1484
if self.uses_draft_model():
1460
- # For draft model-based speculation, we need one new slot per request
1461
- # Since we do not slice the draft tokens
1462
- slots_per_req += 1
1463
- return slots_per_req
1485
+ # The autoregressive draft-model input retains one unsliced token.
1486
+ return 1
1487
+
1488
+ return 0
1489
1490
def use_gemma4_mtp(self) -> bool:
1491
return (