diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index c96e401962..062b38c2ef 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -796,6 +796,26 @@ def test_long_prefill_threshold_applies_with_other_requests(): assert output.num_scheduled_tokens[short_req.request_id] == 10 +def test_long_prefill_threshold_floored_by_fair_share(): + """With the adaptive flag, the effective threshold never falls below the + fair share of the token budget: max_num_batched_tokens / num queued + + running requests.""" + scheduler = create_scheduler( + max_num_batched_tokens=1024, + long_prefill_token_threshold=100, + long_prefill_token_threshold_adaptive=True, + ) + long_req = create_requests(num_requests=1, num_tokens=2000)[0] + short_req = create_requests(num_requests=1, num_tokens=10, req_ids=["short"])[0] + for request in [long_req, short_req]: + scheduler.add_request(request) + + output = scheduler.schedule() + # 100 is below the fair share (1024 // 2 = 512), so the floor binds. + assert output.num_scheduled_tokens[long_req.request_id] == 512 + assert output.num_scheduled_tokens[short_req.request_id] == 10 + + def test_update_from_output_routes_sampling_masks_by_request(): """Each request receives the sampler row at its own batch index.""" scheduler = create_scheduler() diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index ebda4f092c..46ca29f67b 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -58,6 +58,7 @@ def create_scheduler( enable_chunked_prefill: bool = True, enable_prefix_caching: bool = False, long_prefill_token_threshold: int = 0, + long_prefill_token_threshold_adaptive: bool = False, disable_chunked_mm_input: bool = False, use_kv_connector: None | bool | str | MockKVConfig = None, kv_role: str = "kv_both", @@ -90,6 +91,9 @@ def create_scheduler( enable_prefix_caching: optionally force APC config (True/False) or use default (False) + long_prefill_token_threshold: cap on prefill chunk size + long_prefill_token_threshold_adaptive: floor the cap at a + fair share of the token budget Returns: {class}`Scheduler` instance @@ -115,6 +119,7 @@ def create_scheduler( max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len, long_prefill_token_threshold=long_prefill_token_threshold, + long_prefill_token_threshold_adaptive=(long_prefill_token_threshold_adaptive), disable_chunked_mm_input=disable_chunked_mm_input, enable_chunked_prefill=enable_chunked_prefill, async_scheduling=async_scheduling, diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 86f863ab3a..1fb94d3efd 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -84,6 +84,12 @@ class SchedulerConfig: The cap is not applied when the request is the only one in the batch, since there is no other request for it to starve.""" + long_prefill_token_threshold_adaptive: bool = Field(default=False) + """Floor the effective long prefill token threshold at a fair share of + the token budget: max_num_batched_tokens divided by the number of + queued and running requests. Only applies when + long_prefill_token_threshold is nonzero.""" + max_num_queued_reqs: int | None = Field(default=None, ge=0) """Maximum number of requests that can be in-flight (waiting or running) at the same time, or None for no limit. When the limit is reached, new diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 0a5472bce7..f676fc9f18 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -573,6 +573,9 @@ class EngineArgs: max_num_batched_tokens: int | None = None max_num_scheduled_tokens: int | None = None long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold + long_prefill_token_threshold_adaptive: bool = ( + SchedulerConfig.long_prefill_token_threshold_adaptive + ) max_num_seqs: int | None = None max_num_active_seqs: int | None = SchedulerConfig.max_num_active_seqs max_num_queued_reqs: int | None = None @@ -1643,6 +1646,10 @@ class EngineArgs: "--long-prefill-token-threshold", **scheduler_kwargs["long_prefill_token_threshold"], ) + scheduler_group.add_argument( + "--long-prefill-token-threshold-adaptive", + **scheduler_kwargs["long_prefill_token_threshold_adaptive"], + ) # multi-step scheduling has been removed; corresponding arguments # are no longer supported. scheduler_group.add_argument( @@ -2489,6 +2496,9 @@ class EngineArgs: policy=self.scheduling_policy, scheduler_cls=self.scheduler_cls, long_prefill_token_threshold=self.long_prefill_token_threshold, + long_prefill_token_threshold_adaptive=( + self.long_prefill_token_threshold_adaptive + ), scheduler_reserve_full_isl=self.scheduler_reserve_full_isl, watermark=self.watermark, prefill_schedule_interval=self.prefill_schedule_interval, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index d77b63b4aa..59ba590d9d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -134,6 +134,9 @@ class Scheduler(SchedulerInterface): if self.scheduler_config.max_num_scheduled_tokens is not None else self.scheduler_config.max_num_batched_tokens ) + self.adaptive_long_prefill_threshold = ( + self.scheduler_config.long_prefill_token_threshold_adaptive + ) self.max_model_len = vllm_config.model_config.max_model_len self.enable_kv_cache_events = ( self.kv_events_config is not None @@ -602,11 +605,20 @@ class Scheduler(SchedulerInterface): # `long_prefill_token_threshold` exists to stop a long prefill from # starving other requests of the token budget. When it is the only # request there is nobody to starve, so let it use the whole budget. + num_eligible_reqs = ( + len(self.running) + len(self.waiting) + len(self.skipped_waiting) + ) long_prefill_token_threshold = ( self.scheduler_config.long_prefill_token_threshold - if len(self.running) + len(self.waiting) + len(self.skipped_waiting) > 1 + if num_eligible_reqs > 1 else 0 ) + if long_prefill_token_threshold > 0 and self.adaptive_long_prefill_threshold: + # Floor the cap at a fair share of the input budget so it never + # cuts a request below max_num_batched_tokens / num requests. + long_prefill_token_threshold = max( + long_prefill_token_threshold, input_budget // num_eligible_reqs + ) # First, schedule the RUNNING requests. req_index = 0