diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index dcfbfd5b1b..004adf4a67 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -292,34 +292,32 @@ def test_schedule_prefills_gating(has_running: bool): assert any(r.req_id == "new0" for r in output.scheduled_new_reqs) -def test_throttle_prefills_excludes_remote_kv_resume(): - """A request resuming after a completed async KV load (num_computed_tokens - > 0, e.g. the decode side of P/D disaggregation) must NOT be throttled by - the DP prefill cadence: only fresh prefills are deferred. Otherwise the - resumed request's first (single-token) step would be needlessly delayed. +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- + eligible. Returns the scheduler. The connector matches `matched_tokens` of + `r2`'s prompt; the rest (if any) is local prefill. """ from tests.v1.kv_connector.unit.utils import create_model_runner_output BLOCK_SIZE = 16 - NUM_MATCHED = BLOCK_SIZE * 2 scheduler = create_scheduler( enable_prefix_caching=True, - use_kv_connector=mock_kv(matched_tokens=NUM_MATCHED, is_async=True), + use_kv_connector=mock_kv(matched_tokens=matched_tokens, is_async=True), block_size=BLOCK_SIZE, ) - - # Two remote-KV requests with distinct prompts (so r2 gets no local prefix - # cache hit from r1, only the connector's external async load). + # Distinct prompts so r2 gets no local prefix cache hit from r1, only the + # connector's external async load. r1, r2 = create_requests( num_requests=2, - num_tokens=NUM_MATCHED * 2, + num_tokens=num_prompt_tokens, max_tokens=20, block_size=BLOCK_SIZE, req_ids=["r1", "r2"], ) - # r1: drive through its async KV load and into the running (decode) state, - # so that self.running is non-empty for the assertion below. + # r1: drive through its async KV load into the running (decode) state, so + # self.running is non-empty (which makes the next step throttle-eligible). scheduler.add_request(r1) _step_until_kv_transfer_finished(scheduler, ["r1"]) output = scheduler.schedule() # promote + schedule r1 @@ -337,15 +335,40 @@ def test_throttle_prefills_excludes_remote_kv_resume(): output, create_model_runner_output([r1], finished_recving={"r2"}) ) assert "r2" in scheduler.finished_recving_kv_req_ids + return scheduler + + +def test_throttle_prefills_excludes_fully_transferred_remote_kv(): + """A remote-KV resume whose whole prompt was transferred (no local prefill + left, e.g. the decode side of P/D disaggregation) must NOT be throttled by + the DP prefill cadence -- its single-token step has no prefill compute to + defer, so delaying it would be pointless. + """ + block_size = 16 + num_prompt = block_size * 2 + # Fully matched: the whole prompt is loaded remotely. + scheduler = _setup_remote_kv_resume(num_prompt, matched_tokens=num_prompt) - # Throttle prefills. r2's load is complete, so it must be promoted and - # scheduled (a resume, not a fresh prefill) even though the running decode - # (r1) would otherwise make this a throttled step. output = scheduler.schedule(throttle_prefills=True) assert "r2" in output.num_scheduled_tokens assert "r1" in output.num_scheduled_tokens +def test_throttle_prefills_defers_remote_kv_resume_with_local_prefill(): + """A remote-KV resume with local prefill still to compute (the connector + only matched part of the prompt) IS throttled by the DP prefill cadence, + like any other request doing local prefill compute this step. + """ + block_size = 16 + num_prompt = block_size * 4 + # Half matched: the remaining half is local prefill compute. + scheduler = _setup_remote_kv_resume(num_prompt, matched_tokens=num_prompt // 2) + + output = scheduler.schedule(throttle_prefills=True) + assert "r2" not in output.num_scheduled_tokens # deferred (has local prefill) + assert "r1" in output.num_scheduled_tokens + + def test_throttle_defers_inflight_prefill_chunk(): """DP prefill balancing throttles ALL prefill compute on a throttled step, not just new admissions: an in-progress (chunked) prefill already in the diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 90d93a110c..55e40b2043 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -783,10 +783,9 @@ class Scheduler(SchedulerInterface): # KVTransfer: loading remote KV, do not allocate for new work. assert num_external_computed_tokens > 0 num_new_tokens = 0 - elif defer_prefills and request.num_computed_tokens == 0: - # DP prefill balancing: async KV loads (the branch above) are - # allowed to start even on throttled steps, but committing new - # prefill compute is deferred to a cadence-aligned step. + elif defer_prefills and num_computed_tokens < request.num_tokens - 1: + # DP prefill balancing: defer this step's local prefill + # compute to a cadence-aligned step. break else: # Number of tokens to be scheduled.