[Perf] Adaptive budget for spec scheduled input tokens, ~60% better Kimi K3 DSpark TTFT (#51725)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 10, 2026 at 23:52 UTC
0914ed2e816fa1987951f22afba1c00b7786f110
4 files changed
+54
-24
tests/v1/core/test_scheduler.py
+24
@@ -1280,6 +1280,30 @@ def test_reset_connector_cache_no_connector_is_no_op_success():
1280
assert scheduler.reset_prefix_cache(reset_connector=True) is True
1281
1282
1283
+def test_draft_slots_budgeted_per_scheduled_request(tmp_path, monkeypatch):
1284
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
1285
+ (tmp_path / "config.json").write_text(
1286
+ '{"architectures": ["OPTForCausalLM"], "model_type": "opt"}'
1287
+ )
1288
+ scheduler = create_scheduler(
1289
+ model=str(tmp_path),
1290
+ max_num_seqs=16,
1291
+ max_num_batched_tokens=20,
1292
+ num_speculative_tokens=4,
1293
+ parallel_drafting=True,
1294
+ skip_tokenizer_init=True,
1295
+ )
1296
+ speculative_config = scheduler.vllm_config.speculative_config
1297
+ assert speculative_config is not None
1298
+ assert scheduler.max_num_scheduled_tokens == 20
1299
+ assert speculative_config.max_num_new_slots_for_drafting == 3
1300
+
1301
+ for request in create_requests(num_requests=2, num_tokens=10):
1302
+ scheduler.add_request(request)
1303
+
1304
+ assert scheduler.schedule().num_scheduled_tokens == {"0": 10, "1": 4}
1305
+
1306
+
1307
# Note - these test cases mirror some of those in test_rejection_sampler.py
1308
@pytest.mark.parametrize(
1309
"spec_tokens,output_tokens,expected",
tests/v1/core/utils.py
+2
@@ -61,6 +61,7 @@ def create_scheduler(
61
max_model_len: int | None = None,
62
num_speculative_tokens: int | None = None,
63
speculative_method: str | None = None,
64
+ parallel_drafting: bool = False,
65
skip_tokenizer_init: bool = False,
66
async_scheduling: bool = False,
67
pipeline_parallel_size: int = 1,
@@ -151,6 +152,7 @@ def create_scheduler(
152
spec_kwargs["prompt_lookup_max"] = num_speculative_tokens
153
spec_kwargs["prompt_lookup_min"] = 1
154
speculative_config = SpeculativeConfig(**spec_kwargs)
155
+ speculative_config.parallel_drafting = parallel_drafting
156
157
ec_transfer_config = (
158
ECTransferConfig(
vllm/config/vllm.py
+8
-19
@@ -1782,21 +1782,15 @@ class VllmConfig:
1782
def _set_max_num_scheduled_tokens(self):
1783
"""
1784
In most cases, the scheduler may schedule a batch with as many tokens as the
1785
- worker is configured to handle. However for some speculative decoding methods,
1786
- the drafter model may insert additional slots into the batch when drafting.
1787
- To account for this, we need to decrease the max_num_scheduled_tokens by an
1788
- upper bound on the number of slots that can be added.
1785
+ worker is configured to handle.
1786
"""
1787
if self.speculative_config is not None:
1788
scheduled_token_delta = (
1789
self.speculative_config.max_num_new_slots_for_drafting
1793
- * self.scheduler_config.max_num_seqs
1790
)
1791
max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
1792
if self.scheduler_config.max_num_scheduled_tokens is None:
1797
- self.scheduler_config.max_num_scheduled_tokens = (
1798
- max_num_batched_tokens - scheduled_token_delta
1799
- )
1793
+ self.scheduler_config.max_num_scheduled_tokens = max_num_batched_tokens
1794
1795
if self.scheduler_config.max_num_scheduled_tokens <= 0:
1796
raise ValueError(
@@ -1805,7 +1799,7 @@ class VllmConfig:
1799
" the speculative decoding settings, which does not allow"
1800
" any tokens to be scheduled. Increase max_num_batched_tokens"
1801
" to accommodate the additional draft token slots, or decrease"
1808
- " num_speculative_tokens or max_num_seqs."
1802
+ " num_speculative_tokens."
1803
)
1804
if self.scheduler_config.max_num_scheduled_tokens < 8192:
1805
logger.warning_once(
@@ -1814,19 +1808,14 @@ class VllmConfig:
1808
" the speculative decoding settings. This may lead to suboptimal"
1809
" performance. Consider increasing max_num_batched_tokens to"
1810
" accommodate the additional draft token slots, or decrease"
1817
- " num_speculative_tokens or max_num_seqs.",
1811
+ " num_speculative_tokens.",
1812
)
1813
1820
- max_num_scheduled_tokens = self.scheduler_config.max_num_scheduled_tokens
1821
- if max_num_batched_tokens < max_num_scheduled_tokens + (
1822
- self.speculative_config.max_num_new_slots_for_drafting
1823
- * self.scheduler_config.max_num_seqs
1824
- ):
1814
+ if max_num_batched_tokens <= scheduled_token_delta:
1815
raise ValueError(
1826
- f"VllmConfig received max_num_scheduled_tokens but it does not have"
1827
- " enough slots to support the speculative decoding settings."
1828
- f" It should be greater by at least {scheduled_token_delta}, but"
1829
- f" got {max_num_batched_tokens=} and {max_num_scheduled_tokens=}."
1816
+ "VllmConfig does not have enough slots to schedule a token and"
1817
+ " support the speculative decoding settings."
1818
+ f" Got {max_num_batched_tokens=} and {scheduled_token_delta=}."
1819
)
1820
1821
def _set_cudagraph_sizes(self):
vllm/v1/core/sched/scheduler.py
+20
-5
@@ -456,6 +456,9 @@ class Scheduler(SchedulerInterface):
456
req_to_new_blocks: dict[str, KVCacheBlocks] = {}
457
num_scheduled_tokens: dict[str, int] = {}
458
token_budget = self.max_num_scheduled_tokens
459
+ spec = self.vllm_config.speculative_config
460
+ draft_slots = spec.max_num_new_slots_for_drafting if spec is not None else 0
461
+ input_budget = self.scheduler_config.max_num_batched_tokens
462
if self._pause_state == PauseState.PAUSED_ALL:
463
# Do not schedule any requests when paused.
464
token_budget = 0
@@ -483,6 +486,8 @@ class Scheduler(SchedulerInterface):
486
req_index = 0
487
while req_index < len(self.running) and token_budget > 0:
488
request = self.running[req_index]
489
+ if input_budget <= draft_slots:
490
+ break
491
492
if (
493
request.num_output_placeholders > 0
@@ -519,7 +524,9 @@ class Scheduler(SchedulerInterface):
524
)
525
if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens:
526
num_new_tokens = self.scheduler_config.long_prefill_token_threshold
522
- num_new_tokens = min(num_new_tokens, token_budget)
527
+ num_new_tokens = min(
528
+ num_new_tokens, token_budget, input_budget - draft_slots
529
+ )
530
531
# Make sure the input position does not exceed the max model len.
532
# This is necessary when using spec decoding.
@@ -605,7 +612,9 @@ class Scheduler(SchedulerInterface):
612
if preempted_req in scheduled_running_reqs:
613
preempted_req_id = preempted_req.request_id
614
scheduled_running_reqs.remove(preempted_req)
608
- token_budget += num_scheduled_tokens.pop(preempted_req_id)
615
+ restored = num_scheduled_tokens.pop(preempted_req_id)
616
+ token_budget += restored
617
+ input_budget += restored + draft_slots
618
req_to_new_blocks.pop(preempted_req_id)
619
scheduled_spec_decode_tokens.pop(preempted_req_id, None)
620
preempted_encoder_inputs = scheduled_encoder_inputs.pop(
@@ -643,6 +652,7 @@ class Scheduler(SchedulerInterface):
652
req_to_new_blocks[request_id] = new_blocks
653
num_scheduled_tokens[request_id] = num_new_tokens
654
token_budget -= num_new_tokens
655
+ input_budget -= num_new_tokens + draft_slots
656
req_index += 1
657
658
# Speculative decode related.
@@ -693,6 +703,8 @@ class Scheduler(SchedulerInterface):
703
step_skipped_waiting = create_request_queue(self.policy)
704
705
while (self.waiting or self.skipped_waiting) and token_budget > 0:
706
+ if input_budget <= draft_slots:
707
+ break
708
# Paused streaming sessions (WAITING_FOR_STREAMING_REQ) are not
709
# in `running` but still hold a model-runner request slot.
710
num_running = len(self.running) + self.num_waiting_for_streaming_input
@@ -866,6 +878,7 @@ class Scheduler(SchedulerInterface):
878
# compute to a cadence-aligned step.
879
break
880
else:
881
+ request_token_budget = min(token_budget, input_budget - draft_slots)
882
# Number of tokens to be scheduled.
883
# We use `request.num_tokens` instead of
884
# `request.num_prompt_tokens` to consider the resumed
@@ -883,7 +896,7 @@ class Scheduler(SchedulerInterface):
896
):
897
num_new_tokens = 1 + self.num_spec_tokens
898
if (
886
- num_new_tokens > token_budget
899
+ num_new_tokens > request_token_budget
900
or num_computed_tokens + num_new_tokens > self.max_model_len
901
):
902
# Prefer to not schedule than schedule un-padded here.
@@ -898,13 +911,13 @@ class Scheduler(SchedulerInterface):
911
# pooling requests to be chunked
912
if (
913
not self.scheduler_config.enable_chunked_prefill
901
- and num_new_tokens > token_budget
914
+ and num_new_tokens > request_token_budget
915
):
916
# If chunked_prefill is disabled,
917
# we can stop the scheduling here.
918
break
919
907
- num_new_tokens = min(num_new_tokens, token_budget)
920
+ num_new_tokens = min(num_new_tokens, request_token_budget)
921
assert num_new_tokens > 0
922
923
# Apply Mamba alignment before encoder caps.
@@ -1065,6 +1078,7 @@ class Scheduler(SchedulerInterface):
1078
)
1079
num_scheduled_tokens[request_id] = num_new_tokens
1080
token_budget -= num_new_tokens
1081
+ input_budget -= num_new_tokens + draft_slots
1082
request.status = RequestStatus.RUNNING
1083
request.num_computed_tokens = num_computed_tokens
1084
if pad_spec_decode:
@@ -1104,6 +1118,7 @@ class Scheduler(SchedulerInterface):
1118
assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens
1119
1120
assert token_budget >= 0
1121
+ assert input_budget >= 0
1122
assert len(self.running) <= self.max_num_running_reqs
1123
# Since some requests in the RUNNING queue may not be scheduled in
1124
# this step, the total number of scheduled requests can be smaller than