[Core][DP] Throttle prefills based on local prefill work (#46532)

Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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

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.