iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf] Adaptive budget for spec scheduled input tokens, ~60% better Kimi K3 DSpark TTFT (#51725)

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Wentao Ye
2026-08-10 20:52:21 -07:00
committed by GitHub
parent 0f2ea973be
commit 0914ed2e81
4 changed files with 54 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/core/test_scheduler.py
+24
View File
@@ -1280,6 +1280,30 @@ def test_reset_connector_cache_no_connector_is_no_op_success():
assert scheduler.reset_prefix_cache(reset_connector=True) is True
def test_draft_slots_budgeted_per_scheduled_request(tmp_path, monkeypatch):
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
(tmp_path / "config.json").write_text(
'{"architectures": ["OPTForCausalLM"], "model_type": "opt"}'
)
scheduler = create_scheduler(
model=str(tmp_path),
max_num_seqs=16,
max_num_batched_tokens=20,
num_speculative_tokens=4,
parallel_drafting=True,
skip_tokenizer_init=True,
)
speculative_config = scheduler.vllm_config.speculative_config
assert speculative_config is not None
assert scheduler.max_num_scheduled_tokens == 20
assert speculative_config.max_num_new_slots_for_drafting == 3
for request in create_requests(num_requests=2, num_tokens=10):
scheduler.add_request(request)
assert scheduler.schedule().num_scheduled_tokens == {"0": 10, "1": 4}
# Note - these test cases mirror some of those in test_rejection_sampler.py
@pytest.mark.parametrize(
"spec_tokens,output_tokens,expected",
tests/v1/core/utils.py
+2
View File
@@ -61,6 +61,7 @@ def create_scheduler(
max_model_len: int | None = None,
num_speculative_tokens: int | None = None,
speculative_method: str | None = None,
parallel_drafting: bool = False,
skip_tokenizer_init: bool = False,
async_scheduling: bool = False,
pipeline_parallel_size: int = 1,
@@ -151,6 +152,7 @@ def create_scheduler(
spec_kwargs["prompt_lookup_max"] = num_speculative_tokens
spec_kwargs["prompt_lookup_min"] = 1
speculative_config = SpeculativeConfig(**spec_kwargs)
speculative_config.parallel_drafting = parallel_drafting
ec_transfer_config = (
ECTransferConfig(
vllm/config/vllm.py
+8 -19
View File
@@ -1782,21 +1782,15 @@ class VllmConfig:
def _set_max_num_scheduled_tokens(self):
"""
In most cases, the scheduler may schedule a batch with as many tokens as the
worker is configured to handle. However for some speculative decoding methods,
the drafter model may insert additional slots into the batch when drafting.
To account for this, we need to decrease the max_num_scheduled_tokens by an
upper bound on the number of slots that can be added.
worker is configured to handle.
"""
if self.speculative_config is not None:
scheduled_token_delta = (
self.speculative_config.max_num_new_slots_for_drafting
* self.scheduler_config.max_num_seqs
)
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 - scheduled_token_delta
)
self.scheduler_config.max_num_scheduled_tokens = max_num_batched_tokens
if self.scheduler_config.max_num_scheduled_tokens <= 0:
raise ValueError(
@@ -1805,7 +1799,7 @@ class VllmConfig:
" the speculative decoding settings, which does not allow"
" any tokens to be scheduled. Increase max_num_batched_tokens"
" to accommodate the additional draft token slots, or decrease"
" num_speculative_tokens or max_num_seqs."
" num_speculative_tokens."
)
if self.scheduler_config.max_num_scheduled_tokens < 8192:
logger.warning_once(
@@ -1814,19 +1808,14 @@ class VllmConfig:
" the speculative decoding settings. This may lead to suboptimal"
" performance. Consider increasing max_num_batched_tokens to"
" accommodate the additional draft token slots, or decrease"
" num_speculative_tokens or max_num_seqs.",
" num_speculative_tokens.",
)
max_num_scheduled_tokens = self.scheduler_config.max_num_scheduled_tokens
if max_num_batched_tokens < max_num_scheduled_tokens + (
self.speculative_config.max_num_new_slots_for_drafting
* self.scheduler_config.max_num_seqs
):
if max_num_batched_tokens <= scheduled_token_delta:
raise ValueError(
f"VllmConfig received max_num_scheduled_tokens but it does not have"
" enough slots to support the speculative decoding settings."
f" It should be greater by at least {scheduled_token_delta}, but"
f" got {max_num_batched_tokens=} and {max_num_scheduled_tokens=}."
"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=}."
)
def _set_cudagraph_sizes(self):
vllm/v1/core/sched/scheduler.py
+20 -5
View File
@@ -456,6 +456,9 @@ class Scheduler(SchedulerInterface):
req_to_new_blocks: dict[str, KVCacheBlocks] = {}
num_scheduled_tokens: dict[str, int] = {}
token_budget = self.max_num_scheduled_tokens
spec = self.vllm_config.speculative_config
draft_slots = spec.max_num_new_slots_for_drafting if spec is not None else 0
input_budget = self.scheduler_config.max_num_batched_tokens
if self._pause_state == PauseState.PAUSED_ALL:
# Do not schedule any requests when paused.
token_budget = 0
@@ -483,6 +486,8 @@ class Scheduler(SchedulerInterface):
req_index = 0
while req_index < len(self.running) and token_budget > 0:
request = self.running[req_index]
if input_budget <= draft_slots:
break
if (
request.num_output_placeholders > 0
@@ -519,7 +524,9 @@ class Scheduler(SchedulerInterface):
)
if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens:
num_new_tokens = self.scheduler_config.long_prefill_token_threshold
num_new_tokens = min(num_new_tokens, token_budget)
num_new_tokens = min(
num_new_tokens, token_budget, input_budget - draft_slots
)
# Make sure the input position does not exceed the max model len.
# This is necessary when using spec decoding.
@@ -605,7 +612,9 @@ class Scheduler(SchedulerInterface):
if preempted_req in scheduled_running_reqs:
preempted_req_id = preempted_req.request_id
scheduled_running_reqs.remove(preempted_req)
token_budget += num_scheduled_tokens.pop(preempted_req_id)
restored = num_scheduled_tokens.pop(preempted_req_id)
token_budget += restored
input_budget += restored + draft_slots
req_to_new_blocks.pop(preempted_req_id)
scheduled_spec_decode_tokens.pop(preempted_req_id, None)
preempted_encoder_inputs = scheduled_encoder_inputs.pop(
@@ -643,6 +652,7 @@ class Scheduler(SchedulerInterface):
req_to_new_blocks[request_id] = new_blocks
num_scheduled_tokens[request_id] = num_new_tokens
token_budget -= num_new_tokens
input_budget -= num_new_tokens + draft_slots
req_index += 1
# Speculative decode related.
@@ -693,6 +703,8 @@ class Scheduler(SchedulerInterface):
step_skipped_waiting = create_request_queue(self.policy)
while (self.waiting or self.skipped_waiting) and token_budget > 0:
if input_budget <= draft_slots:
break
# Paused streaming sessions (WAITING_FOR_STREAMING_REQ) are not
# in `running` but still hold a model-runner request slot.
num_running = len(self.running) + self.num_waiting_for_streaming_input
@@ -866,6 +878,7 @@ class Scheduler(SchedulerInterface):
# compute to a cadence-aligned step.
break
else:
request_token_budget = min(token_budget, input_budget - draft_slots)
# Number of tokens to be scheduled.
# We use `request.num_tokens` instead of
# `request.num_prompt_tokens` to consider the resumed
@@ -883,7 +896,7 @@ class Scheduler(SchedulerInterface):
):
num_new_tokens = 1 + self.num_spec_tokens
if (
num_new_tokens > token_budget
num_new_tokens > request_token_budget
or num_computed_tokens + num_new_tokens > self.max_model_len
):
# Prefer to not schedule than schedule un-padded here.
@@ -898,13 +911,13 @@ class Scheduler(SchedulerInterface):
# pooling requests to be chunked
if (
not self.scheduler_config.enable_chunked_prefill
and num_new_tokens > token_budget
and num_new_tokens > request_token_budget
):
# If chunked_prefill is disabled,
# we can stop the scheduling here.
break
num_new_tokens = min(num_new_tokens, token_budget)
num_new_tokens = min(num_new_tokens, request_token_budget)
assert num_new_tokens > 0
# Apply Mamba alignment before encoder caps.
@@ -1065,6 +1078,7 @@ class Scheduler(SchedulerInterface):
)
num_scheduled_tokens[request_id] = num_new_tokens
token_budget -= num_new_tokens
input_budget -= num_new_tokens + draft_slots
request.status = RequestStatus.RUNNING
request.num_computed_tokens = num_computed_tokens
if pad_spec_decode:
@@ -1104,6 +1118,7 @@ class Scheduler(SchedulerInterface):
assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens
assert token_budget >= 0
assert input_budget >= 0
assert len(self.running) <= self.max_num_running_reqs
# Since some requests in the RUNNING queue may not be scheduled in
# this step, the total number of scheduled requests can be smaller than