[Scheduler] Soften Long Prefill Tokens Threshhold (#57951)
Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Sep 21, 2026 at 20:19 UTC
c64b15cde583c7dd3cbb906b30164131a030d6b4
5 files changed
+52
-11
tests/v1/core/test_deferred_block_free.py
+1
-1
@@ -417,7 +417,7 @@ def test_abort_mid_prefill_defers_free():
417
chunk is still in flight must withhold its blocks.
418
"""
419
scheduler = create_scheduler(
420
- model=MODEL, async_scheduling=True, long_prefill_token_threshold=16
420
+ model=MODEL, async_scheduling=True, max_num_batched_tokens=16
421
)
422
scheduler.defer_block_free = True
423
pool = scheduler.kv_cache_manager.block_pool
tests/v1/core/test_scheduler.py
+30
@@ -762,6 +762,36 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
762
assert output2.num_scheduled_tokens[requests[2].request_id] == 800 - 224 - 224
763
764
765
+def test_long_prefill_threshold_ignored_when_alone():
766
+ """A lone long prefill is not capped by the threshold: it has no other
767
+ request to starve."""
768
+ scheduler = create_scheduler(
769
+ max_num_batched_tokens=1024,
770
+ long_prefill_token_threshold=400,
771
+ )
772
+ request = create_requests(num_requests=1, num_tokens=2000)[0]
773
+ scheduler.add_request(request)
774
+
775
+ output = scheduler.schedule()
776
+ assert output.num_scheduled_tokens[request.request_id] == 1024
777
+
778
+
779
+def test_long_prefill_threshold_applies_with_other_requests():
780
+ """The threshold caps the prefill as soon as another request is queued."""
781
+ scheduler = create_scheduler(
782
+ max_num_batched_tokens=1024,
783
+ long_prefill_token_threshold=400,
784
+ )
785
+ long_req = create_requests(num_requests=1, num_tokens=2000)[0]
786
+ short_req = create_requests(num_requests=1, num_tokens=10, req_ids=["short"])[0]
787
+ for request in [long_req, short_req]:
788
+ scheduler.add_request(request)
789
+
790
+ output = scheduler.schedule()
791
+ assert output.num_scheduled_tokens[long_req.request_id] == 400
792
+ assert output.num_scheduled_tokens[short_req.request_id] == 10
793
+
794
+
795
def test_update_from_output_routes_sampling_masks_by_request():
796
"""Each request receives the sampler row at its own batch index."""
797
scheduler = create_scheduler()
tests/v1/e2e/general/test_pooling_chunked_prefill.py
+4
-4
@@ -77,13 +77,13 @@ def test_pooling_chunked_prefill(vllm_runner, monkeypatch):
77
78
chunk_size = 10
79
80
- # Set chunking parameters to force chunked prefill
81
- # Note: Chunked prefill is automatically handled by vLLM
82
- # internally based on the model size and prompt
80
+ # Cap the token budget to force chunked prefill; the
81
+ # long_prefill_token_threshold is not applied to a lone request.
82
with vllm_runner(
83
model_id,
84
runner="pooling",
86
- long_prefill_token_threshold=chunk_size,
85
+ max_num_batched_tokens=chunk_size,
86
+ max_num_seqs=1,
87
tensor_parallel_size=1,
88
enforce_eager=True,
89
enable_chunked_prefill=True,
vllm/config/scheduler.py
+4
-1
@@ -79,7 +79,10 @@ class SchedulerConfig:
79
80
long_prefill_token_threshold: int = Field(default=0, ge=0)
81
"""For chunked prefill, a request is considered long if the prompt is
82
- longer than this number of tokens. 0 disables the cap (default)."""
82
+ longer than this number of tokens. 0 disables the cap (default).
83
+
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
max_num_queued_reqs: int | None = Field(default=None, ge=0)
88
"""Maximum number of requests that can be in-flight (waiting or running)
vllm/v1/core/sched/scheduler.py
+13
-5
@@ -599,6 +599,15 @@ class Scheduler(SchedulerInterface):
599
throttle_prefills and not self.prefill_capacity_bound
600
) and any(not r.is_prefill_chunk for r in self.running)
601
602
+ # `long_prefill_token_threshold` exists to stop a long prefill from
603
+ # starving other requests of the token budget. When it is the only
604
+ # request there is nobody to starve, so let it use the whole budget.
605
+ long_prefill_token_threshold = (
606
+ self.scheduler_config.long_prefill_token_threshold
607
+ if len(self.running) + len(self.waiting) + len(self.skipped_waiting) > 1
608
+ else 0
609
+ )
610
+
611
# First, schedule the RUNNING requests.
612
req_index = 0
613
while req_index < len(self.running) and token_budget > 0:
@@ -650,8 +659,8 @@ class Scheduler(SchedulerInterface):
659
+ request.num_output_placeholders
660
- request.num_computed_tokens
661
)
653
- if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens:
654
- num_new_tokens = self.scheduler_config.long_prefill_token_threshold
662
+ if 0 < long_prefill_token_threshold < num_new_tokens:
663
+ num_new_tokens = long_prefill_token_threshold
664
num_new_tokens = min(
665
num_new_tokens, token_budget, input_budget - draft_slots
666
)
@@ -1090,9 +1099,8 @@ class Scheduler(SchedulerInterface):
1099
num_new_tokens = padded_num_tokens
1100
pad_spec_decode = True
1101
1093
- threshold = self.scheduler_config.long_prefill_token_threshold
1094
- if 0 < threshold < num_new_tokens:
1095
- num_new_tokens = threshold
1102
+ if 0 < long_prefill_token_threshold < num_new_tokens:
1103
+ num_new_tokens = long_prefill_token_threshold
1104
1105
# chunked prefill has to be enabled explicitly to allow
1106
# pooling requests to be chunked