[Scheduler] Tune --long-prefill-token-threshold adaptiveness (#58459)
Signed-off-by: Robert Shaw <robertgshaw2@gmail.com>
Robert Shaw committed
Sep 23, 2026 at 22:23 UTC
6697a7dd30ceebbe69a9d37e4012d11d40f79685
5 files changed
+54
-1
tests/v1/core/test_scheduler.py
+20
@@ -796,6 +796,26 @@ def test_long_prefill_threshold_applies_with_other_requests():
796
assert output.num_scheduled_tokens[short_req.request_id] == 10
797
798
799
+def test_long_prefill_threshold_floored_by_fair_share():
800
+ """With the adaptive flag, the effective threshold never falls below the
801
+ fair share of the token budget: max_num_batched_tokens / num queued +
802
+ running requests."""
803
+ scheduler = create_scheduler(
804
+ max_num_batched_tokens=1024,
805
+ long_prefill_token_threshold=100,
806
+ long_prefill_token_threshold_adaptive=True,
807
+ )
808
+ long_req = create_requests(num_requests=1, num_tokens=2000)[0]
809
+ short_req = create_requests(num_requests=1, num_tokens=10, req_ids=["short"])[0]
810
+ for request in [long_req, short_req]:
811
+ scheduler.add_request(request)
812
+
813
+ output = scheduler.schedule()
814
+ # 100 is below the fair share (1024 // 2 = 512), so the floor binds.
815
+ assert output.num_scheduled_tokens[long_req.request_id] == 512
816
+ assert output.num_scheduled_tokens[short_req.request_id] == 10
817
+
818
+
819
def test_update_from_output_routes_sampling_masks_by_request():
820
"""Each request receives the sampler row at its own batch index."""
821
scheduler = create_scheduler()
tests/v1/core/utils.py
+5
@@ -58,6 +58,7 @@ def create_scheduler(
58
enable_chunked_prefill: bool = True,
59
enable_prefix_caching: bool = False,
60
long_prefill_token_threshold: int = 0,
61
+ long_prefill_token_threshold_adaptive: bool = False,
62
disable_chunked_mm_input: bool = False,
63
use_kv_connector: None | bool | str | MockKVConfig = None,
64
kv_role: str = "kv_both",
@@ -90,6 +91,9 @@ def create_scheduler(
91
enable_prefix_caching: optionally force APC config
92
(True/False) or use default
93
(False)
94
+ long_prefill_token_threshold: cap on prefill chunk size
95
+ long_prefill_token_threshold_adaptive: floor the cap at a
96
+ fair share of the token budget
97
98
Returns:
99
{class}`Scheduler` instance
@@ -115,6 +119,7 @@ def create_scheduler(
119
max_num_batched_tokens=max_num_batched_tokens,
120
max_model_len=max_model_len,
121
long_prefill_token_threshold=long_prefill_token_threshold,
122
+ long_prefill_token_threshold_adaptive=(long_prefill_token_threshold_adaptive),
123
disable_chunked_mm_input=disable_chunked_mm_input,
124
enable_chunked_prefill=enable_chunked_prefill,
125
async_scheduling=async_scheduling,
vllm/config/scheduler.py
+6
@@ -84,6 +84,12 @@ class SchedulerConfig:
84
The cap is not applied when the request is the only one in the batch,
85
since there is no other request for it to starve."""
86
87
+ long_prefill_token_threshold_adaptive: bool = Field(default=False)
88
+ """Floor the effective long prefill token threshold at a fair share of
89
+ the token budget: max_num_batched_tokens divided by the number of
90
+ queued and running requests. Only applies when
91
+ long_prefill_token_threshold is nonzero."""
92
+
93
max_num_queued_reqs: int | None = Field(default=None, ge=0)
94
"""Maximum number of requests that can be in-flight (waiting or running)
95
at the same time, or None for no limit. When the limit is reached, new
vllm/engine/arg_utils.py
+10
@@ -573,6 +573,9 @@ class EngineArgs:
573
max_num_batched_tokens: int | None = None
574
max_num_scheduled_tokens: int | None = None
575
long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold
576
+ long_prefill_token_threshold_adaptive: bool = (
577
+ SchedulerConfig.long_prefill_token_threshold_adaptive
578
+ )
579
max_num_seqs: int | None = None
580
max_num_active_seqs: int | None = SchedulerConfig.max_num_active_seqs
581
max_num_queued_reqs: int | None = None
@@ -1643,6 +1646,10 @@ class EngineArgs:
1646
"--long-prefill-token-threshold",
1647
**scheduler_kwargs["long_prefill_token_threshold"],
1648
)
1649
+ scheduler_group.add_argument(
1650
+ "--long-prefill-token-threshold-adaptive",
1651
+ **scheduler_kwargs["long_prefill_token_threshold_adaptive"],
1652
+ )
1653
# multi-step scheduling has been removed; corresponding arguments
1654
# are no longer supported.
1655
scheduler_group.add_argument(
@@ -2489,6 +2496,9 @@ class EngineArgs:
2496
policy=self.scheduling_policy,
2497
scheduler_cls=self.scheduler_cls,
2498
long_prefill_token_threshold=self.long_prefill_token_threshold,
2499
+ long_prefill_token_threshold_adaptive=(
2500
+ self.long_prefill_token_threshold_adaptive
2501
+ ),
2502
scheduler_reserve_full_isl=self.scheduler_reserve_full_isl,
2503
watermark=self.watermark,
2504
prefill_schedule_interval=self.prefill_schedule_interval,
vllm/v1/core/sched/scheduler.py
+13
-1
@@ -134,6 +134,9 @@ class Scheduler(SchedulerInterface):
134
if self.scheduler_config.max_num_scheduled_tokens is not None
135
else self.scheduler_config.max_num_batched_tokens
136
)
137
+ self.adaptive_long_prefill_threshold = (
138
+ self.scheduler_config.long_prefill_token_threshold_adaptive
139
+ )
140
self.max_model_len = vllm_config.model_config.max_model_len
141
self.enable_kv_cache_events = (
142
self.kv_events_config is not None
@@ -602,11 +605,20 @@ class Scheduler(SchedulerInterface):
605
# `long_prefill_token_threshold` exists to stop a long prefill from
606
# starving other requests of the token budget. When it is the only
607
# request there is nobody to starve, so let it use the whole budget.
608
+ num_eligible_reqs = (
609
+ len(self.running) + len(self.waiting) + len(self.skipped_waiting)
610
+ )
611
long_prefill_token_threshold = (
612
self.scheduler_config.long_prefill_token_threshold
607
- if len(self.running) + len(self.waiting) + len(self.skipped_waiting) > 1
613
+ if num_eligible_reqs > 1
614
else 0
615
)
616
+ if long_prefill_token_threshold > 0 and self.adaptive_long_prefill_threshold:
617
+ # Floor the cap at a fair share of the input budget so it never
618
+ # cuts a request below max_num_batched_tokens / num requests.
619
+ long_prefill_token_threshold = max(
620
+ long_prefill_token_threshold, input_budget // num_eligible_reqs
621
+ )
622
623
# First, schedule the RUNNING requests.
624
req_index = 0