[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>
Nick Hill committed
Jun 23, 2026 at 19:27 UTC
6af0559ddbc9ffa30a8eb3399dd1f27d8e46db56
2 files changed
+42
-20
tests/v1/core/test_scheduler.py
+39
-16
@@ -292,34 +292,32 @@ def test_schedule_prefills_gating(has_running: bool):
292
assert any(r.req_id == "new0" for r in output.scheduled_new_reqs)
293
294
295
-def test_throttle_prefills_excludes_remote_kv_resume():
296
- """A request resuming after a completed async KV load (num_computed_tokens
297
- > 0, e.g. the decode side of P/D disaggregation) must NOT be throttled by
298
- the DP prefill cadence: only fresh prefills are deferred. Otherwise the
299
- resumed request's first (single-token) step would be needlessly delayed.
295
+def _setup_remote_kv_resume(num_prompt_tokens: int, matched_tokens: int):
296
+ """Drive a remote-KV request `r2` to the resume point (async load complete)
297
+ while another request `r1` is already decoding, so the step is throttle-
298
+ eligible. Returns the scheduler. The connector matches `matched_tokens` of
299
+ `r2`'s prompt; the rest (if any) is local prefill.
300
"""
301
from tests.v1.kv_connector.unit.utils import create_model_runner_output
302
303
BLOCK_SIZE = 16
304
- NUM_MATCHED = BLOCK_SIZE * 2
304
scheduler = create_scheduler(
305
enable_prefix_caching=True,
307
- use_kv_connector=mock_kv(matched_tokens=NUM_MATCHED, is_async=True),
306
+ use_kv_connector=mock_kv(matched_tokens=matched_tokens, is_async=True),
307
block_size=BLOCK_SIZE,
308
)
310
-
311
- # Two remote-KV requests with distinct prompts (so r2 gets no local prefix
312
- # cache hit from r1, only the connector's external async load).
309
+ # Distinct prompts so r2 gets no local prefix cache hit from r1, only the
310
+ # connector's external async load.
311
r1, r2 = create_requests(
312
num_requests=2,
315
- num_tokens=NUM_MATCHED * 2,
313
+ num_tokens=num_prompt_tokens,
314
max_tokens=20,
315
block_size=BLOCK_SIZE,
316
req_ids=["r1", "r2"],
317
)
318
321
- # r1: drive through its async KV load and into the running (decode) state,
322
- # so that self.running is non-empty for the assertion below.
319
+ # r1: drive through its async KV load into the running (decode) state, so
320
+ # self.running is non-empty (which makes the next step throttle-eligible).
321
scheduler.add_request(r1)
322
_step_until_kv_transfer_finished(scheduler, ["r1"])
323
output = scheduler.schedule() # promote + schedule r1
@@ -337,15 +335,40 @@ def test_throttle_prefills_excludes_remote_kv_resume():
335
output, create_model_runner_output([r1], finished_recving={"r2"})
336
)
337
assert "r2" in scheduler.finished_recving_kv_req_ids
338
+ return scheduler
339
+
340
+
341
+def test_throttle_prefills_excludes_fully_transferred_remote_kv():
342
+ """A remote-KV resume whose whole prompt was transferred (no local prefill
343
+ left, e.g. the decode side of P/D disaggregation) must NOT be throttled by
344
+ the DP prefill cadence -- its single-token step has no prefill compute to
345
+ defer, so delaying it would be pointless.
346
+ """
347
+ block_size = 16
348
+ num_prompt = block_size * 2
349
+ # Fully matched: the whole prompt is loaded remotely.
350
+ scheduler = _setup_remote_kv_resume(num_prompt, matched_tokens=num_prompt)
351
341
- # Throttle prefills. r2's load is complete, so it must be promoted and
342
- # scheduled (a resume, not a fresh prefill) even though the running decode
343
- # (r1) would otherwise make this a throttled step.
352
output = scheduler.schedule(throttle_prefills=True)
353
assert "r2" in output.num_scheduled_tokens
354
assert "r1" in output.num_scheduled_tokens
355
356
357
+def test_throttle_prefills_defers_remote_kv_resume_with_local_prefill():
358
+ """A remote-KV resume with local prefill still to compute (the connector
359
+ only matched part of the prompt) IS throttled by the DP prefill cadence,
360
+ like any other request doing local prefill compute this step.
361
+ """
362
+ block_size = 16
363
+ num_prompt = block_size * 4
364
+ # Half matched: the remaining half is local prefill compute.
365
+ scheduler = _setup_remote_kv_resume(num_prompt, matched_tokens=num_prompt // 2)
366
+
367
+ output = scheduler.schedule(throttle_prefills=True)
368
+ assert "r2" not in output.num_scheduled_tokens # deferred (has local prefill)
369
+ assert "r1" in output.num_scheduled_tokens
370
+
371
+
372
def test_throttle_defers_inflight_prefill_chunk():
373
"""DP prefill balancing throttles ALL prefill compute on a throttled step,
374
not just new admissions: an in-progress (chunked) prefill already in the
vllm/v1/core/sched/scheduler.py
+3
-4
@@ -783,10 +783,9 @@ class Scheduler(SchedulerInterface):
783
# KVTransfer: loading remote KV, do not allocate for new work.
784
assert num_external_computed_tokens > 0
785
num_new_tokens = 0
786
- elif defer_prefills and request.num_computed_tokens == 0:
787
- # DP prefill balancing: async KV loads (the branch above) are
788
- # allowed to start even on throttled steps, but committing new
789
- # prefill compute is deferred to a cadence-aligned step.
786
+ elif defer_prefills and num_computed_tokens < request.num_tokens - 1:
787
+ # DP prefill balancing: defer this step's local prefill
788
+ # compute to a cadence-aligned step.
789
break
790
else:
791
# Number of tokens to be scheduled.