[Core] Apply prefill_schedule_interval outside data parallelism (#54627)

Signed-off-by: lifulu <fululi12@amd.com>
Signed-off-by: lifulu <lifulu@users.noreply.github.com>
Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: lifulu <lifulu@users.noreply.github.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

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-

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,

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

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