From ed3050add4abd770b118ad6fa337bdde0ba2ae08 Mon Sep 17 00:00:00 2001 From: fululi12 Date: Thu, 15 Oct 2026 09:03:13 -0700 Subject: [PATCH] [Core] Apply prefill_schedule_interval outside data parallelism (#54627) Signed-off-by: lifulu Signed-off-by: lifulu Co-authored-by: Cursor Agent Co-authored-by: lifulu Co-authored-by: Claude Opus 5 (1M context) --- tests/v1/core/test_scheduler.py | 243 ++++++++++++++++++++++++++++++++ tests/v1/core/utils.py | 2 + vllm/config/scheduler.py | 9 +- vllm/v1/core/sched/scheduler.py | 78 ++++++++++ 4 files changed, 329 insertions(+), 3 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 99eb02ff33..0d7897edc1 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -558,6 +558,249 @@ def test_schedule_prefills_gating(has_running: bool): assert any(r.req_id == "new0" for r in output.scheduled_new_reqs) +def _decode_one(scheduler, output, req_ids: list[str]) -> None: + """Feed back one sampled token per request, so they keep decoding.""" + scheduler.update_from_output( + output, + ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, + sampled_token_ids=[[0] for _ in req_ids], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ), + ) + + +def _start_decoding(scheduler, req_id: str = "run0"): + """Prefill one request and take it into the decode state.""" + (request,) = create_requests( + num_requests=1, num_tokens=8, max_tokens=100, req_ids=[req_id] + ) + scheduler.add_request(request) + output = scheduler.schedule() + assert req_id in output.num_scheduled_tokens + _decode_one(scheduler, output, [req_id]) + assert len(scheduler.running) == 1 + return output + + +@pytest.mark.parametrize("interval", [2, 4]) +def test_local_prefill_interval_holds_prefill_between_bursts(interval: int): + """A single engine applies `prefill_schedule_interval` in the scheduler, + timed from the last prefill-carrying step: a waiting prefill is held for the + remainder of the interval while running requests keep decoding, then + admitted once the interval has elapsed. + """ + scheduler = create_scheduler(prefill_schedule_interval=interval) + # This prefill starts the interval. + _start_decoding(scheduler) + assert scheduler.last_prefill_step == scheduler.num_executed_steps + + (new_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["new0"]) + scheduler.add_request(new_req) + + # Steps inside the interval decode alone; the waiting prefill is held. + for _ in range(interval - 1): + output = scheduler.schedule() + assert "new0" not in output.num_scheduled_tokens + assert "run0" in output.num_scheduled_tokens + assert new_req.status == RequestStatus.WAITING + _decode_one(scheduler, output, ["run0"]) + + # The interval has now elapsed, so prefill is admitted again. + output = scheduler.schedule() + assert "new0" in output.num_scheduled_tokens + assert scheduler.last_prefill_step == scheduler.num_executed_steps + + +def test_empty_schedule_does_not_advance_the_prefill_cadence(): + """An empty schedule() must not consume interval budget. + + `current_step` counts every schedule() call. With async scheduling, + step_with_batch_queue() calls schedule() again whenever the batch queue has + room and has_requests() is true -- so a call lands while the previous batch + is still in flight and nothing new can be scheduled. An async KV connector + makes this common, since requests sit in lookup or WAITING_FOR_REMOTE_KVS. + Timing the cadence on current_step lets those empty calls expire the hold + early. Reported on #54627 against the LMCache MP connector. + """ + scheduler = create_scheduler(prefill_schedule_interval=4) + _start_decoding(scheduler) + + # One real decode step, left in flight (no update_from_output yet). + output = scheduler.schedule() + assert output.total_num_scheduled_tokens > 0 + executed_before = scheduler.num_executed_steps + step_before = scheduler.current_step + + # Re-entering schedule() while that batch is still unaccounted for + # schedules nothing -- this is the async-scheduling shape. + for _ in range(3): + empty = scheduler.schedule() + assert empty.total_num_scheduled_tokens == 0 + + assert scheduler.current_step == step_before + 3, "current_step should advance" + assert scheduler.num_executed_steps == executed_before, ( + "an empty schedule() must not count as an executed step" + ) + + +def test_prefill_cadence_counts_executed_steps_not_schedule_calls(): + """The hold spans `interval` EXECUTED steps, however many empty schedule() + calls are interleaved. + + Counting schedule() calls instead lets each empty call eat interval budget, + so with an async KV connector prefill lands on nearly every step and the + feature stops doing anything. + """ + interval = 4 + scheduler = create_scheduler(prefill_schedule_interval=interval) + _start_decoding(scheduler) + prefill_at = scheduler.num_executed_steps + + (new_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["new0"]) + scheduler.add_request(new_req) + + empty_calls = 0 + admitted_at = None + for _ in range(50): + output = scheduler.schedule() + if "new0" in output.num_scheduled_tokens: + admitted_at = scheduler.num_executed_steps + break + if output.total_num_scheduled_tokens == 0: + empty_calls += 1 + continue + # A real decode step. Leave it in flight and re-enter schedule() twice, + # as async scheduling does while the batch queue has room. + for _ in range(2): + extra = scheduler.schedule() + if "new0" in extra.num_scheduled_tokens: + admitted_at = scheduler.num_executed_steps + break + empty_calls += 1 + if admitted_at is not None: + break + _decode_one(scheduler, output, ["run0"]) + + assert admitted_at is not None, "prefill was never admitted" + assert empty_calls > 0, "test never exercised an empty schedule() call" + # The hold is measured in executed steps. Counting schedule() calls would + # admit after far fewer, because the empty re-entries above would count. + assert admitted_at - prefill_at == interval, ( + f"hold lasted {admitted_at - prefill_at} executed steps, expected " + f"{interval}; empty schedule() calls were counted against the cadence" + ) + assert scheduler.current_step > scheduler.num_executed_steps + + +def test_local_prefill_interval_defaults_to_no_gating(): + """The default interval of 1 leaves admission exactly as it was: a waiting + prefill is admitted on the very next step. + """ + scheduler = create_scheduler() + assert scheduler.local_prefill_interval == 1 + _start_decoding(scheduler) + + (new_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["new0"]) + scheduler.add_request(new_req) + output = scheduler.schedule() + assert "new0" in output.num_scheduled_tokens + + +def test_local_prefill_interval_left_to_engine_under_dp(): + """Under data parallelism the cadence must stay aligned across ranks, so the + DP engine core drives it and the scheduler must not gate on its own. + """ + scheduler = create_scheduler(prefill_schedule_interval=4, data_parallel_size=2) + assert scheduler.local_prefill_interval == 1 + _start_decoding(scheduler) + + (new_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["new0"]) + scheduler.add_request(new_req) + output = scheduler.schedule() + assert "new0" in output.num_scheduled_tokens + + +def test_local_prefill_interval_never_idles_the_model(): + """With nothing running to decode in prefill's place, the interval must not + hold prefill back, or the step would be wasted entirely. + """ + scheduler = create_scheduler(prefill_schedule_interval=8) + _start_decoding(scheduler) + + # Finish the only running request, leaving nothing to decode. + scheduler.finish_requests("run0", RequestStatus.FINISHED_ABORTED) + assert not scheduler.running + + (new_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["new0"]) + scheduler.add_request(new_req) + output = scheduler.schedule() + assert "new0" in output.num_scheduled_tokens + + +@pytest.mark.parametrize("interval", [2, 3]) +def test_local_prefill_interval_defers_inflight_prefill_chunk(interval: int): + """The single-engine interval gates prefill WORK, not just new admissions. + + A request that is already partially prefilled is held for the remainder of + the interval too, exactly as `throttle_prefills` does for DP. That is the + point of the cadence -- a held step must carry no prefill at all, or decode + does not get the uninterrupted stretch the interval exists to give it. + """ + scheduler = create_scheduler( + max_num_seqs=16, + max_num_batched_tokens=50, + enable_chunked_prefill=True, + prefill_schedule_interval=interval, + ) + + # A short request that finishes prefill in one step -> a running decode. + (decode_req,) = create_requests(num_requests=1, num_tokens=4, req_ids=["dec0"]) + scheduler.add_request(decode_req) + output = scheduler.schedule() + _decode_one(scheduler, output, ["dec0"]) + assert decode_req in scheduler.running and not decode_req.is_prefill_chunk + + # A long request (80 tokens, budget 50) -> prefilled in chunks. That first + # prefill closed the interval, so walk to the next step that opens it. + (chunk_req,) = create_requests(num_requests=1, num_tokens=80, req_ids=["chk0"]) + scheduler.add_request(chunk_req) + for _ in range(interval - 1): + output = scheduler.schedule() + assert "chk0" not in output.num_scheduled_tokens + _decode_one(scheduler, output, ["dec0"]) + + output = scheduler.schedule() + assert output.num_scheduled_tokens["chk0"] > 0 + scheduler.update_from_output( + output, + ModelRunnerOutput( + req_ids=["dec0", "chk0"], + req_id_to_index={"dec0": 0, "chk0": 1}, + sampled_token_ids=[[0], []], # no token sampled for partial prefill + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ), + ) + assert chunk_req.is_prefill_chunk # still mid-prefill, in running + + # Held steps: the in-flight chunk is deferred along with new admissions, + # while the decode keeps running. + for _ in range(interval - 1): + output = scheduler.schedule() + assert "chk0" not in output.num_scheduled_tokens + assert "dec0" in output.num_scheduled_tokens + _decode_one(scheduler, output, ["dec0"]) + + # Once the interval has elapsed the chunk resumes. + output = scheduler.schedule() + assert "chk0" in output.num_scheduled_tokens + + def _setup_remote_kv_resume(num_prompt_tokens: int, matched_tokens: int): """Drive a remote-KV request `r2` to the resume point (async load complete) while another request `r1` is already decoding, so the step is throttle- diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 541e55f5b6..25e9373888 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -85,6 +85,7 @@ def create_scheduler( kv_cache_spec: KVCacheSpec | None = None, runner: RunnerOption = "auto", per_request_spec_decode_metrics: str = "none", + prefill_schedule_interval: int = 1, scheduling_policy: SchedulerPolicy = "fcfs", diffusion_canvas_length: int | None = None, scheduler_cls: type[Scheduler] | None = None, @@ -134,6 +135,7 @@ def create_scheduler( enable_chunked_prefill=enable_chunked_prefill, async_scheduling=async_scheduling, is_encoder_decoder=model_config.is_encoder_decoder, + prefill_schedule_interval=prefill_schedule_interval, # Ensure admission/preemption mechanics are deterministic watermark=0.0, policy=scheduling_policy, diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 1dbe4b72eb..5ad38c0a46 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -201,9 +201,12 @@ class SchedulerConfig: (the default) disables the watermark.""" prefill_schedule_interval: int = Field(default=1, ge=1) - """For data-parallel deployments, only admit new prefill requests - once every N engine steps, aligned across DP ranks, to better balance - per-step forward-pass times.""" + """Only admit prefill once every N engine steps, leaving the steps in + between to decode alone. 1 (the default) admits prefill on every step. + + A deferred step carries no prefill at all, including chunks of an + already-admitted prompt, so a long prompt can take up to N times as many + steps to finish prefilling.""" async_scheduling: bool | None = None """If set to False, disable async scheduling. Async scheduling helps to diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index a6ca6074df..c7ece2a1f0 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -368,6 +368,36 @@ class Scheduler(SchedulerInterface): # prefill batch fully drained the waiting queue. Prefill throttling # is disabled in this case. self.prefill_capacity_bound = False + # The DP engine core drives this cadence from a counter shared across + # ranks; an engine without one never throttles, so apply the interval + # locally, timed from the last prefill-carrying step. Three regimes: + # 1. no DP -- local cadence applies; + # 2. dense DP rank -- `reconfigure_for_independent_dp_rank` already + # set data_parallel_size=1, so each rank gets its own cadence. + # Intended, but a behaviour change: this used to be a no-op; + # 3. MoE DP rank (size > 1) -- the engine core drives it, no-op here. + self.local_prefill_interval = ( + self.scheduler_config.prefill_schedule_interval + if self.parallel_config.data_parallel_size == 1 + else 1 + ) + if self.local_prefill_interval > 1: + # Neither the hold nor its release is otherwise visible in logs or + # stats, which makes a misapplied cadence hard to diagnose. + logger.info_once( + "Prefill cadence: applying prefill_schedule_interval=%d " + "locally (this engine is not DP-coordinated); prefill is held " + "until %d executed steps have passed since the last " + "prefill-carrying step.", + self.local_prefill_interval, + self.local_prefill_interval, + ) + # A full interval back, so the first step is never gated. + self.last_prefill_step = -self.local_prefill_interval + # Steps that actually scheduled tokens. `current_step` counts every + # schedule() call including empty ones, so it would expire the hold + # early; it keeps its own meaning (next_decode_eligible_step). + self.num_executed_steps = 0 self.scheduler_reserve_full_isl = ( self.scheduler_config.scheduler_reserve_full_isl ) @@ -623,6 +653,9 @@ class Scheduler(SchedulerInterface): scheduled_spec_decode_tokens: dict[str, list[int]] = {} # Whether the running batch contains any prefill requests. prefill_scheduled = False + # Whether a waiting request was admitted with prompt left to compute. + # Separate from `prefill_scheduled`, which covers only the running loop. + prefill_admitted = False # Whether any scheduled request has a synchronous connector KV load. has_sync_kv_loads = False @@ -654,6 +687,28 @@ class Scheduler(SchedulerInterface): long_prefill_token_threshold, input_budget // num_eligible_reqs ) + # Local prefill cadence: hold prefill until the interval since the last + # prefill-carrying step has elapsed, so decode runs in longer + # uninterrupted stretches. Mirrors the DP gate above, including its + # capacity escape: once a release step cannot drain the waiting queue, + # throttling stops (#44558, test_throttle_capacity_bound_guard_admits). + # Without that escape a backed-up engine holds both admissions and + # in-flight chunks for interval-1 of every interval steps, capping + # prefill at roughly one budget per interval and growing the waiting + # queue without bound. + if ( + not defer_prefills + and self.local_prefill_interval > 1 + and not self.prefill_capacity_bound + # This call becomes executed step num_executed_steps + 1 if it + # schedules anything, and last_prefill_step is stamped on the same + # convention below. + and (self.num_executed_steps + 1) - self.last_prefill_step + < self.local_prefill_interval + and any(not r.is_prefill_chunk for r in self.running) + ): + defer_prefills = True + # First, schedule the RUNNING requests. req_index = 0 while req_index < len(self.running) and token_budget > 0: @@ -1318,6 +1373,12 @@ class Scheduler(SchedulerInterface): ) if load_kv_async: + # This path deliberately does not reach `prefill_admitted` + # below, so it does not restart the prefill cadence: + # `num_new_tokens` was set to 0 above, so the step carries + # no prefill compute for this request. The interval is + # timed from steps that actually prefilled; refreshing it + # here would start the cadence from a step that did none. # If loading async, allocate memory and put request # into the WAITING_FOR_REMOTE_KV state. request.status = RequestStatus.WAITING_FOR_REMOTE_KVS @@ -1373,6 +1434,14 @@ class Scheduler(SchedulerInterface): scheduled_spec_decode_tokens[request_id] = [ -1 ] * self.num_spec_tokens + # Must use the SAME predicate as the deferral gate in the + # waiting loop (`num_computed_tokens < num_tokens - 1`). + # Stamping on `< num_prompt_tokens` instead let a request with + # exactly one token left pass the gate on a held step and still + # restart the interval -- which is every full prefix-cache hit, + # and every remote-KV resume, since those set + # num_computed_tokens = num_tokens - 1. + prefill_admitted |= num_computed_tokens < request.num_tokens - 1 # Only track requests that will still be prefilling after this chunk. if num_computed_tokens + num_new_tokens < request.num_tokens: self._inflight_prefills.add(request) @@ -1403,6 +1472,15 @@ class Scheduler(SchedulerInterface): # Check if the scheduling constraints are satisfied. total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens + # Count this step before stamping the cadence below, so that + # last_prefill_step includes the step that carried the prefill. + if total_num_scheduled_tokens > 0: + self.num_executed_steps += 1 + + # Restart the cadence interval from any step that carried prefill, + # whether a new admission or a running chunk. + if self.local_prefill_interval > 1 and (prefill_scheduled or prefill_admitted): + self.last_prefill_step = self.num_executed_steps assert token_budget >= 0 assert input_budget >= 0