@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d8d95998dc

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Core] Add prefill step cadence for better non-PD DP balancing (#44558)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Jun 16, 2026 at 13:17 UTC d8d95998dcb261229369eaae28c4413268ea803d

7 files changed +306 -4

tests/v1/core/test_scheduler.py

+181

@@ -207,6 +207,187 @@ def test_schedule_partial_requests():

207 assert requests[2].request_id not in output.num_scheduled_tokens

208

209

210 +@pytest.mark.parametrize("has_running", [True, False])

211 +def test_schedule_prefills_gating(has_running: bool):

212 + """DP prefill-balancing gate: when `throttle_prefills` is True, a new

213 + WAITING (prefill) request is deferred ONLY if this rank has running work to

214 + protect. With no running requests, the prefill is admitted regardless (so a

215 + throttled step is never wasted as a dummy), and running/decode requests are

216 + unaffected. Once the cadence allows prefills again, the request is admitted.

217 + """

218 + scheduler = create_scheduler(max_num_seqs=16, max_num_batched_tokens=8192)

219 +

220 + if has_running:

221 + # Establish a running (decode) request via a prefill + output step.

222 + (running_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["run0"])

223 + scheduler.add_request(running_req)

224 + output = scheduler.schedule()

225 + assert len(output.scheduled_new_reqs) == 1

226 + scheduler.update_from_output(

227 + output,

228 + ModelRunnerOutput(

229 + req_ids=["run0"],

230 + req_id_to_index={"run0": 0},

231 + sampled_token_ids=[[0]],

232 + logprobs=None,

233 + prompt_logprobs_dict={},

234 + pooler_output=[],

235 + ),

236 + )

237 + assert len(scheduler.running) == 1

238 +

239 + # Add a new WAITING (prefill) request, with prefills gated off.

240 + (new_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["new0"])

241 + scheduler.add_request(new_req)

242 + output = scheduler.schedule(throttle_prefills=True)

243 +

244 + if has_running:

245 + # There is running work to protect, so the new prefill is deferred...

246 + assert "new0" not in output.num_scheduled_tokens

247 + assert new_req.status == RequestStatus.WAITING

248 + # ...while the running/decode request keeps being scheduled.

249 + assert "run0" in output.num_scheduled_tokens

250 + # When the cadence allows prefills again, the request is admitted.

251 + output = scheduler.schedule()

252 +

253 + # No running work to protect (or cadence now open): the prefill is admitted.

254 + assert "new0" in output.num_scheduled_tokens

255 + assert any(r.req_id == "new0" for r in output.scheduled_new_reqs)

256 +

257 +

258 +def test_throttle_prefills_excludes_remote_kv_resume():

259 + """A request resuming after a completed async KV load (num_computed_tokens

260 + > 0, e.g. the decode side of P/D disaggregation) must NOT be throttled by

261 + the DP prefill cadence: only fresh prefills are deferred. Otherwise the

262 + resumed request's first (single-token) step would be needlessly delayed.

263 + """

264 + from tests.v1.kv_connector.unit.utils import create_model_runner_output

265 +

266 + BLOCK_SIZE = 16

267 + NUM_MATCHED = BLOCK_SIZE * 2

268 + scheduler = create_scheduler(

269 + enable_prefix_caching=True,

270 + use_kv_connector=mock_kv(matched_tokens=NUM_MATCHED, is_async=True),

271 + block_size=BLOCK_SIZE,

272 + )

273 +

274 + # Two remote-KV requests with distinct prompts (so r2 gets no local prefix

275 + # cache hit from r1, only the connector's external async load).

276 + r1, r2 = create_requests(

277 + num_requests=2,

278 + num_tokens=NUM_MATCHED * 2,

279 + max_tokens=20,

280 + block_size=BLOCK_SIZE,

281 + req_ids=["r1", "r2"],

282 + )

283 +

284 + # r1: drive through its async KV load and into the running (decode) state,

285 + # so that self.running is non-empty for the assertion below.

286 + scheduler.add_request(r1)

287 + _step_until_kv_transfer_finished(scheduler, ["r1"])

288 + output = scheduler.schedule() # promote + schedule r1

289 + assert "r1" in output.num_scheduled_tokens

290 + scheduler.update_from_output(

291 + output, create_model_runner_output([r1], token_id=1000)

292 + )

293 + assert scheduler.running # r1 now decoding

294 +

295 + # r2: a second remote-KV request; complete its async load while r1 decodes.

296 + scheduler.add_request(r2)

297 + output = scheduler.schedule() # r1 decodes; r2 -> WAITING_FOR_REMOTE_KVS

298 + assert r2.status == RequestStatus.WAITING_FOR_REMOTE_KVS

299 + scheduler.update_from_output(

300 + output, create_model_runner_output([r1], finished_recving={"r2"})

301 + )

302 + assert "r2" in scheduler.finished_recving_kv_req_ids

303 +

304 + # Throttle prefills. r2's load is complete, so it must be promoted and

305 + # scheduled (a resume, not a fresh prefill) even though the running decode

306 + # (r1) would otherwise make this a throttled step.

307 + output = scheduler.schedule(throttle_prefills=True)

308 + assert "r2" in output.num_scheduled_tokens

309 + assert "r1" in output.num_scheduled_tokens

310 +

311 +

312 +def test_throttle_defers_inflight_prefill_chunk():

313 + """DP prefill balancing throttles ALL prefill compute on a throttled step,

314 + not just new admissions: an in-progress (chunked) prefill already in the

315 + running queue is also deferred, so the step runs decode-only, while a

316 + separate decode keeps being scheduled."""

317 + scheduler = create_scheduler(

318 + max_num_seqs=16, max_num_batched_tokens=50, enable_chunked_prefill=True

319 + )

320 +

321 + # A short request that finishes prefill in one step -> a running decode.

322 + (decode_req,) = create_requests(num_requests=1, num_tokens=4, req_ids=["dec0"])

323 + scheduler.add_request(decode_req)

324 + output = scheduler.schedule()

325 + scheduler.update_from_output(

326 + output,

327 + ModelRunnerOutput(

328 + req_ids=["dec0"],

329 + req_id_to_index={"dec0": 0},

330 + sampled_token_ids=[[0]],

331 + logprobs=None,

332 + prompt_logprobs_dict={},

333 + pooler_output=[],

334 + ),

335 + )

336 + assert decode_req in scheduler.running and not decode_req.is_prefill_chunk

337 +

338 + # A long request (80 tokens, budget 50) -> prefilled in chunks.

339 + (chunk_req,) = create_requests(num_requests=1, num_tokens=80, req_ids=["chk0"])

340 + scheduler.add_request(chunk_req)

341 + output = scheduler.schedule() # first chunk of chk0 + decode of dec0

342 + assert output.num_scheduled_tokens["chk0"] > 0

343 + scheduler.update_from_output(

344 + output,

345 + ModelRunnerOutput(

346 + req_ids=["dec0", "chk0"],

347 + req_id_to_index={"dec0": 0, "chk0": 1},

348 + sampled_token_ids=[[0], []], # no token sampled for partial prefill

349 + logprobs=None,

350 + prompt_logprobs_dict={},

351 + pooler_output=[],

352 + ),

353 + )

354 + assert chunk_req.is_prefill_chunk # still mid-prefill, in running

355 +

356 + # Throttled step: the in-flight prefill chunk is deferred, the decode runs.

357 + output = scheduler.schedule(throttle_prefills=True)

358 + assert "chk0" not in output.num_scheduled_tokens

359 + assert "dec0" in output.num_scheduled_tokens

360 +

361 + # When the cadence opens again, the prefill chunk resumes.

362 + output = scheduler.schedule()

363 + assert "chk0" in output.num_scheduled_tokens

364 +

365 +

366 +def test_throttle_capacity_bound_guard_admits():

367 + """Saturation guard: if a cadence-aligned release step cannot drain the

368 + waiting prefill queue (it ran out of token budget), the throttle backs off on

369 + the next step so the backlog cannot grow into a TTFT avalanche -- prefills are

370 + admitted even though throttle_prefills is set."""

371 + scheduler = create_scheduler(

372 + max_num_seqs=16, max_num_batched_tokens=200, enable_chunked_prefill=True

373 + )

374 + a, b = create_requests(num_requests=2, num_tokens=200, req_ids=["a", "b"])

375 + scheduler.add_request(a)

376 + scheduler.add_request(b)

377 +

378 + # Release step (throttle off): `a` fills the 200-token budget; `b` cannot be

379 + # reached, so the waiting queue is not drained -> capacity-bound.

380 + output = scheduler.schedule()

381 + assert "a" in output.num_scheduled_tokens

382 + assert "b" not in output.num_scheduled_tokens

383 + assert scheduler.prefill_capacity_bound

384 +

385 + # Throttle. Because the previous release was capacity-bound, the guard backs

386 + # off and `b` is admitted rather than stalling the backlog.

387 + output = scheduler.schedule(throttle_prefills=True)

388 + assert "b" in output.num_scheduled_tokens

389 +

390 +

391 def test_no_mm_input_chunking():

392 # Disable multimodal input chunking.

393 scheduler = create_scheduler(

tests/v1/distributed/test_async_llm_dp.py

+61

@@ -186,6 +186,67 @@ async def test_load(

186 )

187

188

189 +@pytest.mark.parametrize("prefill_schedule_interval", [1, 4])

190 +@pytest.mark.asyncio

191 +async def test_dp_prefill_schedule_interval(prefill_schedule_interval: int):

192 + """Throttling new prefills to every Nth step (DP balancing) must not break

193 + generation: a stream of staggered requests should still all complete with

194 + the expected number of tokens.

195 +

196 + The throttle only engages in the DP MoE/EP engine-core path

197 + (`DPEngineCoreProc`), so this uses an MoE model with expert parallel.

198 + """

199 + with ExitStack() as after:

200 + prompt = "This is a test of data parallel"

201 +

202 + engine_args = AsyncEngineArgs(

203 + model="ibm-research/PowerMoE-3b",

204 + enforce_eager=True,

205 + tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),

206 + data_parallel_size=DP_SIZE,

207 + data_parallel_backend="mp",

208 + enable_expert_parallel=True,

209 + prefill_schedule_interval=prefill_schedule_interval,

210 + )

211 + engine = AsyncLLM.from_engine_args(engine_args)

212 + after.callback(engine.shutdown)

213 +

214 + NUM_REQUESTS = 50

215 + NUM_EXPECTED_TOKENS = 10

216 +

217 + request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]

218 +

219 + # Create requests with a small stagger so they arrive across many

220 + # steps and (with interval > 1) accumulate in the waiting queue

221 + # before being admitted together on cadence-aligned steps.

222 + tasks = []

223 + for request_id in request_ids:

224 + tasks.append(

225 + asyncio.create_task(

226 + generate(

227 + engine,

228 + request_id,

229 + prompt,

230 + RequestOutputKind.DELTA,

231 + NUM_EXPECTED_TOKENS,

232 + )

233 + )

234 + )

235 + await asyncio.sleep(0.01)

236 +

237 + done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)

238 + for task in pending:

239 + task.cancel()

240 + for task in done:

241 + num_generated_tokens, request_id = await task

242 + assert num_generated_tokens == NUM_EXPECTED_TOKENS, (

243 + f"{request_id} generated {num_generated_tokens} but "

244 + f"expected {NUM_EXPECTED_TOKENS}"

245 + )

246 +

247 + assert not engine.output_processor.has_unfinished_requests()

248 +

249 +

250 # =============================================================================

251 # DP Pause/Resume Tests

252 # =============================================================================

vllm/config/scheduler.py

@@ -150,6 +150,11 @@ class SchedulerConfig:

150 of requests when GPU memory is scarce. Must be in the range [0.0, 1.0); 0.0

151 (the default) disables the watermark."""

152

153 + prefill_schedule_interval: int = Field(default=1, ge=1)

154 + """For data-parallel deployments, only admit new prefill requests

155 + once every N engine steps, aligned across DP ranks, to better balance

156 + per-step forward-pass times."""

157 +

158 async_scheduling: bool | None = None

159 """If set to False, disable async scheduling. Async scheduling helps to

160 avoid gaps in GPU utilization, leading to better latency and throughput.

vllm/engine/arg_utils.py

@@ -599,6 +599,7 @@ class EngineArgs:

599 disable_chunked_mm_input: bool = SchedulerConfig.disable_chunked_mm_input

600

601 scheduler_reserve_full_isl: bool = SchedulerConfig.scheduler_reserve_full_isl

602 + prefill_schedule_interval: int = SchedulerConfig.prefill_schedule_interval

603

604 watermark: float = SchedulerConfig.watermark

605

@@ -1417,6 +1418,10 @@ class EngineArgs:

1418 **scheduler_kwargs["scheduler_reserve_full_isl"],

1419 )

1420 scheduler_group.add_argument("--watermark", **scheduler_kwargs["watermark"])

1421 + scheduler_group.add_argument(

1422 + "--prefill-schedule-interval",

1423 + **scheduler_kwargs["prefill_schedule_interval"],

1424 + )

1425 scheduler_group.add_argument(

1426 "--disable-hybrid-kv-cache-manager",

1427 **scheduler_kwargs["disable_hybrid_kv_cache_manager"],

@@ -2064,6 +2069,7 @@ class EngineArgs:

2069 long_prefill_token_threshold=self.long_prefill_token_threshold,

2070 scheduler_reserve_full_isl=self.scheduler_reserve_full_isl,

2071 watermark=self.watermark,

2072 + prefill_schedule_interval=self.prefill_schedule_interval,

2073 disable_hybrid_kv_cache_manager=self.disable_hybrid_kv_cache_manager,

2074 async_scheduling=self.async_scheduling,

2075 stream_interval=self.stream_interval,

vllm/v1/core/sched/interface.py

+7 -1

@@ -49,7 +49,7 @@ class SchedulerInterface(ABC):

49 raise NotImplementedError

51 @abstractmethod

52 - def schedule(self) -> "SchedulerOutput":

52 + def schedule(self, throttle_prefills: bool = False) -> "SchedulerOutput":

53 """Schedule the requests to process in this scheduling step.

55 The scheduling decision is made at the iteration level. Each scheduling

@@ -68,6 +68,12 @@ class SchedulerInterface(ABC):

68 or the batch as a whole. The model runner will use this information in

69 preparing inputs to the model.

71 + Args:

72 + throttle_prefills: DP prefill balancing. When True (set by the DP

73 + engine core on non-cadence-aligned steps), new prefill compute is

74 + deferred to a later step so prefills stay aligned across DP ranks;

75 + automatically overridden when the rank is saturated.

76 +

77 Returns:

78 A SchedulerOutput object containing information about the scheduled

79 requests.

vllm/v1/core/sched/scheduler.py

+27 -1

@@ -275,6 +275,10 @@ class Scheduler(SchedulerInterface):

275 # Scheduler iteration counter. Drives the V2+PP+async decode-throttle

276 # cadence (`next_decode_eligible_step`).

277 self.current_step = 0

278 + # DP prefill balancing: Flag to track whether the last cadence-aligned

279 + # prefill batch fully drained the waiting queue. Prefill throttling

280 + # is disabled in this case.

281 + self.prefill_capacity_bound = False

282 self.scheduler_reserve_full_isl = (

283 self.scheduler_config.scheduler_reserve_full_isl

284 )

@@ -380,7 +384,7 @@ class Scheduler(SchedulerInterface):

384 num_new_tokens = num_new_tokens // block_size * block_size

385 return num_new_tokens

386

383 - def schedule(self) -> SchedulerOutput:

387 + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:

388 self.current_step += 1

389 # NOTE(woosuk) on the scheduling algorithm:

390 # There's no "decoding phase" nor "prefill phase" in the scheduler.

@@ -416,6 +420,12 @@ class Scheduler(SchedulerInterface):

420

421 self.kv_cache_manager.new_step_starts()

422

423 + # DP prefill balancing: on a throttled (non-cadence-aligned) step, defer

424 + # all prefill compute unless saturated.

425 + defer_prefills = (

426 + throttle_prefills and not self.prefill_capacity_bound

427 + ) and any(not r.is_prefill_chunk for r in self.running)

428 +

429 # First, schedule the RUNNING requests.

430 req_index = 0

431 while req_index < len(self.running) and token_budget > 0:

@@ -443,6 +453,12 @@ class Scheduler(SchedulerInterface):

453 req_index += 1

454 continue

455

456 + if defer_prefills and request.is_prefill_chunk:

457 + # DP prefill balancing: defer this in-progress prefill chunk to a

458 + # cadence-aligned step; decodes still run to fill this step.

459 + req_index += 1

460 + continue

461 +

462 num_new_tokens = (

463 request.num_tokens_with_spec

464 + request.num_output_placeholders

@@ -766,6 +782,11 @@ class Scheduler(SchedulerInterface):

782 # KVTransfer: loading remote KV, do not allocate for new work.

783 assert num_external_computed_tokens > 0

784 num_new_tokens = 0

785 + elif defer_prefills and request.num_computed_tokens == 0:

786 + # DP prefill balancing: async KV loads (the branch above) are

787 + # allowed to start even on throttled steps, but committing new

788 + # prefill compute is deferred to a cadence-aligned step.

789 + break

790 else:

791 # Number of tokens to be scheduled.

792 # We use `request.num_tokens` instead of

@@ -959,6 +980,11 @@ class Scheduler(SchedulerInterface):

980 if step_skipped_waiting:

981 self.skipped_waiting.prepend_requests(step_skipped_waiting)

982

983 + # DP prefill balancing: on a step that admitted prefills (release),

984 + # record whether it was capacity-bound.

985 + if not defer_prefills:

986 + self.prefill_capacity_bound = bool(self.waiting)

987 +

988 # Check if the scheduling constraints are satisfied.

989 total_num_scheduled_tokens = sum(num_scheduled_tokens.values())

990 assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens

vllm/v1/engine/core.py

+19 -2

@@ -471,6 +471,11 @@ class EngineCore:

471 )

472 self._iteration_index += 1

473

474 + def _should_throttle_prefills(self) -> bool:

475 + """Whether to defer new prefills this step (DP prefill balancing).

476 + Overridden by the DP engine core; never throttles otherwise."""

477 + return False

478 +

479 def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]:

480 """Schedule, execute, and make output.

481

@@ -482,7 +487,7 @@ class EngineCore:

487 # or finished and not yet removed from the batch.

488 if not self.scheduler.has_requests():

489 return {}, False

485 - scheduler_output = self.scheduler.schedule()

490 + scheduler_output = self.scheduler.schedule(self._should_throttle_prefills())

491 future = self.model_executor.execute_model(scheduler_output, non_block=True)

492 grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output)

493 with (

@@ -539,7 +544,7 @@ class EngineCore:

544 model_executed = False

545 deferred_scheduler_output = None

546 if self.scheduler.has_requests():

542 - scheduler_output = self.scheduler.schedule()

547 + scheduler_output = self.scheduler.schedule(self._should_throttle_prefills())

548 with self.log_error_detail(scheduler_output):

549 exec_future = self.model_executor.execute_model(

550 scheduler_output, non_block=True

@@ -1751,6 +1756,9 @@ class DPEngineCoreProc(EngineCoreProc):

1756 "DPEngineCoreProc should only be used for MoE models"

1757 )

1758

1759 + scheduler_config = vllm_config.scheduler_config

1760 + self.prefill_schedule_interval = scheduler_config.prefill_schedule_interval

1761 +

1762 # Counts forward-passes of the model so that we can synchronize

1763 # finished with DP peers every N steps.

1764 self.step_counter = 0

@@ -1901,6 +1909,15 @@ class DPEngineCoreProc(EngineCoreProc):

1909 )

1910 self.output_queue.put_nowait((-1, EngineCoreOutputs(scheduler_stats=stats)))

1911

1912 + def _should_throttle_prefills(self) -> bool:

1913 + # Throttle new prefills to cadence-aligned steps for DP balancing.

1914 + # step_counter is identical across DP ranks. On a fresh wave the

1915 + # counter is 0, so prefills are admitted immediately after idle.

1916 + return (

1917 + self.prefill_schedule_interval > 1

1918 + and self.step_counter % self.prefill_schedule_interval != 0

1919 + )

1920 +

1921 def run_busy_loop(self):

1922 """Core busy loop of the EngineCore for data parallel case."""

1923