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

Fork 0

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

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 23:00 UTC

[Perf] Optimize scheduler overhead for PD disaggregation, around 5% E2E perf improvement (#35781)

Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

Wentao Ye committed Mar 11, 2026 at 00:25 UTC a8ff2cca92807f1b15b9b8d21135784298ad7814

6 files changed +243 -105

tests/v1/core/test_scheduler.py

+98 -6

@@ -1115,12 +1115,16 @@ def _step_until_done(

1115 all_finished = all_done

1116

1117

1118 +def _num_waiting_requests(scheduler: Scheduler) -> int:

1119 + return len(scheduler.waiting) + len(scheduler.skipped_waiting)

1120 +

1121 +

1122 def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):

1123 """Cycle requests through a KV transfer cycle."""

1124

1125 # Requests should first transition to WAITING_FOR_REMOTE_KVS

1126 output = scheduler.schedule()

1123 - assert len(scheduler.waiting) == len(req_ids)

1127 + assert _num_waiting_requests(scheduler) == len(req_ids)

1128 assert len(scheduler.running) == 0

1129 assert len(output.scheduled_new_reqs) == 0

1130 for req in scheduler.requests.values():

@@ -1139,7 +1143,7 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):

1143

1144 # Simulate KV transfer completion using KVConnectorOutput.finished_recving

1145 output = scheduler.schedule()

1142 - assert len(scheduler.waiting) == len(req_ids)

1146 + assert _num_waiting_requests(scheduler) == len(req_ids)

1147 assert len(scheduler.running) == 0

1148

1149 MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

@@ -1546,7 +1550,7 @@ def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role):

1550 # All can be scheduled - 1st token.

1551 output = scheduler.schedule()

1552 if is_async:

1549 - assert len(scheduler.waiting) == 2

1553 + assert _num_waiting_requests(scheduler) == 2

1554 assert scheduler.running == []

1555 _step_until_kv_transfer_finished(scheduler, req_ids)

1556 output = scheduler.schedule()

@@ -1604,7 +1608,11 @@ def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role):

1608 # This will have a local and remote cache hit.

1609 output = scheduler.schedule()

1610 if is_async:

1607 - waiting_req_ids = [req.request_id for req in scheduler.waiting]

1611 + waiting_req_ids = [

1612 + req.request_id

1613 + for req in scheduler.skipped_waiting

1614 + if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS

1615 + ]

1616 assert len(waiting_req_ids) == 1

1617 _step_until_kv_transfer_finished(scheduler, waiting_req_ids)

1618 output = scheduler.schedule()

@@ -2439,7 +2447,8 @@ def test_schedule_skip_tokenizer_init_structured_output_request():

2447 output = scheduler.schedule()

2448 assert len(output.scheduled_new_reqs) == 0

2449 assert len(scheduler.running) == 0

2442 - assert len(scheduler.waiting) == 1

2450 + assert len(scheduler.waiting) == 0

2451 + assert len(scheduler.skipped_waiting) == 1

2452

2453

2454 @pytest.mark.parametrize(

@@ -3626,6 +3635,9 @@ def test_prepend_skipped_requests_order():

3635 # simulate first 2 waiting requests are waiting for remote KVs

3636 for req in expected_waiting_reqs[:2]:

3637 req.status = RequestStatus.WAITING_FOR_REMOTE_KVS

3638 + scheduler.waiting.remove_requests(expected_waiting_reqs[:2])

3639 + for req in expected_waiting_reqs[:2]:

3640 + scheduler.skipped_waiting.add_request(req)

3641

3642 # schedule step

3643 # expect the first 2 waiting to be skipped, the third running,

@@ -3636,7 +3648,87 @@ def test_prepend_skipped_requests_order():

3648 expected_waiting_reqs.pop(2)

3649

3650 # verify waiting order is preserved

3639 - assert list(scheduler.waiting) == expected_waiting_reqs

3651 + waiting_reqs = list(scheduler.skipped_waiting) + list(scheduler.waiting)

3652 + assert waiting_reqs == expected_waiting_reqs

3653 +

3654 +

3655 +def test_remote_kv_promotion_keeps_fcfs_with_fsm_prefix():

3656 + scheduler = create_scheduler(max_num_seqs=1)

3657 + scheduler.connector = Mock()

3658 + scheduler.connector.get_num_new_matched_tokens.return_value = (0, False)

3659 +

3660 + requests = create_requests(num_requests=4)

3661 + for request in requests:

3662 + scheduler.add_request(request)

3663 +

3664 + req_fsm_1, req_fsm_2, req_remote, req_tail = list(scheduler.waiting)

3665 +

3666 + # simulate two FSM requests at the waiting head that become ready now.

3667 + req_fsm_1.status = RequestStatus.WAITING_FOR_FSM

3668 + req_fsm_1.structured_output_request = Mock(grammar=object())

3669 + req_fsm_2.status = RequestStatus.WAITING_FOR_FSM

3670 + req_fsm_2.structured_output_request = Mock(grammar=object())

3671 +

3672 + # simulate a remote-KV request that is ready to be promoted now.

3673 + req_remote.status = RequestStatus.WAITING_FOR_REMOTE_KVS

3674 + scheduler.waiting.remove_requests([req_fsm_1, req_fsm_2, req_remote])

3675 + scheduler.skipped_waiting.add_request(req_fsm_1)

3676 + scheduler.skipped_waiting.add_request(req_fsm_2)

3677 + scheduler.skipped_waiting.add_request(req_remote)

3678 + scheduler.finished_recving_kv_req_ids.add(req_remote.request_id)

3679 + scheduler._update_waiting_for_remote_kv = Mock()

3680 +

3681 + output = scheduler.schedule()

3682 +

3683 + assert output.scheduled_new_reqs

3684 + assert output.scheduled_new_reqs[0].req_id == req_fsm_1.request_id

3685 + waiting_req_ids = [

3686 + req.request_id

3687 + for req in list(scheduler.skipped_waiting) + list(scheduler.waiting)

3688 + ]

3689 + assert waiting_req_ids == [

3690 + req_fsm_2.request_id,

3691 + req_remote.request_id,

3692 + req_tail.request_id,

3693 + ]

3694 +

3695 +

3696 +def test_fcfs_mixed_skipped_waiting_types_keep_order():

3697 + scheduler = create_scheduler(max_num_batched_tokens=20)

3698 + scheduler._update_waiting_for_remote_kv = Mock()

3699 +

3700 + mk_req = lambda req_id, num_tokens=1: create_requests( # noqa: E731

3701 + num_requests=1, num_tokens=num_tokens, req_ids=[req_id]

3702 + )[0]

3703 + req_fsm, req_remote, req_stream = mk_req("fsm"), mk_req("remote"), mk_req("stream")

3704 + req_regular, req_tail = mk_req("regular", 20), mk_req("tail")

3705 + req_fsm.status = RequestStatus.WAITING_FOR_FSM

3706 + req_fsm.structured_output_request = Mock(grammar=None)

3707 + req_remote.status = RequestStatus.WAITING_FOR_REMOTE_KVS

3708 + req_stream.status = RequestStatus.WAITING_FOR_STREAMING_REQ

3709 +

3710 + for req in (req_fsm, req_remote, req_stream, req_regular, req_tail):

3711 + scheduler.add_request(req)

3712 + scheduler.schedule()

3713 + assert list(scheduler.skipped_waiting) == [req_fsm, req_remote, req_stream]

3714 +

3715 + scheduler.finish_requests(req_regular.request_id, RequestStatus.FINISHED_ABORTED)

3716 + assert not scheduler.running

3717 +

3718 + req_fsm.structured_output_request = Mock(grammar=object())

3719 + scheduler.finished_recving_kv_req_ids.add(req_remote.request_id)

3720 + req_stream.status = RequestStatus.WAITING

3721 +

3722 + second_output = scheduler.schedule()

3723 + expected_order = [

3724 + req_fsm.request_id,

3725 + req_remote.request_id,

3726 + req_stream.request_id,

3727 + req_tail.request_id,

3728 + ]

3729 + assert [req.req_id for req in second_output.scheduled_new_reqs] == expected_order

3730 + assert [req.request_id for req in scheduler.running] == expected_order

3731 + scheduler._update_waiting_for_remote_kv.assert_called_once_with(req_remote)

3732

3733

3734 def test_abort_request_waiting_for_remote_kvs():

tests/v1/kv_connector/unit/test_error_propagation.py

+2 -1

@@ -119,7 +119,7 @@ def test_error_propagation_async_load(fail_scheduler: Scheduler):

119

120 scheduler_output = fail_scheduler.schedule()

121

122 - assert len(fail_scheduler.waiting) == 1

122 + assert len(fail_scheduler.skipped_waiting) == 1

123 assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS

124 assert request.num_computed_tokens == num_external_computed_tokens

125

@@ -145,3 +145,4 @@ def test_error_propagation_async_load(fail_scheduler: Scheduler):

145 assert output.finish_reason == FinishReason.ERROR

146

147 assert len(fail_scheduler.waiting) == 0

148 + assert len(fail_scheduler.skipped_waiting) == 0

tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py

+1 -1

@@ -337,7 +337,7 @@ def test_async_recompute_blocks_not_cached_when_invalid(

337 scheduler_output = recompute_scheduler.schedule()

338

339 # request should be waiting for remote KVs

340 - assert len(recompute_scheduler.waiting) == 1

340 + assert len(recompute_scheduler.skipped_waiting) == 1

341 assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS

342 assert request.num_computed_tokens == num_external_computed_tokens

343

tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py

+12 -8

@@ -76,8 +76,9 @@ def test_async_load_failure(

77 scheduler_output = scheduler.schedule()

79 - assert len(scheduler.waiting) == 3

80 - for request in scheduler.waiting:

79 + assert len(scheduler.waiting) == 0

80 + assert len(scheduler.skipped_waiting) == 3

81 + for request in scheduler.skipped_waiting:

82 assert request.num_computed_tokens == num_external_computed_tokens

83 assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS

84 assert scheduler.connector.get_num_new_matched_tokens.call_count == 3

@@ -96,8 +97,9 @@ def test_async_load_failure(

98 min_invalid_block_idx = min(invalid_block_idxs)

99 - assert len(scheduler.waiting) == 3

100 - for request in scheduler.waiting:

100 + assert len(scheduler.waiting) == 0

101 + assert len(scheduler.skipped_waiting) == 3

102 + for request in scheduler.skipped_waiting:

103 if request.request_id == request2.request_id:

104 assert request.num_computed_tokens == (

105 min_invalid_block_idx * scheduler.block_size

@@ -303,8 +305,9 @@ def test_async_progressive_load_failure(

305

306 scheduler_output = scheduler.schedule()

307

306 - assert len(scheduler.waiting) == 1

307 - assert scheduler.waiting.peek_request().request_id == request.request_id

308 + assert len(scheduler.waiting) == 0

309 + assert len(scheduler.skipped_waiting) == 1

310 + assert scheduler.skipped_waiting.peek_request().request_id == request.request_id

311 assert request.num_computed_tokens == num_external_computed_tokens

312 assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS

313 assert scheduler.connector.get_num_new_matched_tokens.call_count == 1

@@ -325,8 +328,9 @@ def test_async_progressive_load_failure(

328

329 min_invalid_block_idx = min(min_invalid_block_idx, invalid_block_idx)

330

328 - assert len(scheduler.waiting) == 1

329 - assert scheduler.waiting.peek_request().request_id == request.request_id

331 + assert len(scheduler.waiting) == 0

332 + assert len(scheduler.skipped_waiting) == 1

333 + assert scheduler.skipped_waiting.peek_request().request_id == request.request_id

334 assert request.num_computed_tokens == (

335 min_invalid_block_idx * scheduler.block_size

336 )

tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py

+29 -25

@@ -18,6 +18,10 @@ from .utils import (

18 pytestmark = pytest.mark.cpu_test

21 +def _num_waiting_requests(scheduler) -> int:

22 + return len(scheduler.waiting) + len(scheduler.skipped_waiting)

23 +

24 +

25 def test_basic_lifecycle():

26 """Test lifecycle of a remote prefill."""

@@ -54,8 +58,8 @@ def test_basic_lifecycle():

58 assert scheduler_output.total_num_scheduled_tokens == 0

60 # Req waiting for KVs with no computed/scheduled toks ...

57 - assert len(scheduler.waiting) == 1

58 - assert request in scheduler.waiting

61 + assert _num_waiting_requests(scheduler) == 1

62 + assert request in scheduler.skipped_waiting

63 assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS

64 assert request.num_computed_tokens == NUM_TOKENS

@@ -81,7 +85,7 @@ def test_basic_lifecycle():

85 # STEP (2):

86 # (2a): schedule(): nothing happens!

87 scheduler_output = scheduler.schedule()

84 - assert len(scheduler.waiting) == 1

88 + assert _num_waiting_requests(scheduler) == 1

89 assert len(scheduler.running) == 0

91 # (2b): forward(): request finishes recv.

@@ -94,7 +98,7 @@ def test_basic_lifecycle():

98 engine_core_outputs = scheduler.update_from_output(

99 scheduler_output, model_runner_output

100 )

97 - assert len(scheduler.waiting) == 1

101 + assert _num_waiting_requests(scheduler) == 1

102 assert request_id in scheduler.finished_recving_kv_req_ids

103

104 # STEP (3):

@@ -180,7 +184,7 @@ def test_interleaved_lifecycle():

184 scheduler.add_request(request_remote)

185 scheduler_output = scheduler.schedule()

186 assert len(scheduler.running) == 2

183 - assert len(scheduler.waiting) == 1

187 + assert _num_waiting_requests(scheduler) == 1

188 assert len(scheduler_output.scheduled_new_reqs) == 1

189 assert scheduler_output.scheduled_cached_reqs.num_reqs == 1

190

@@ -190,7 +194,7 @@ def test_interleaved_lifecycle():

194 # STEP 3: continue running, KVs not arrived yet.

195 scheduler_output = scheduler.schedule()

196 assert len(scheduler.running) == 2

193 - assert len(scheduler.waiting) == 1

197 + assert _num_waiting_requests(scheduler) == 1

198 assert len(scheduler_output.scheduled_new_reqs) == 0

199 assert scheduler_output.scheduled_cached_reqs.num_reqs == 2

200

@@ -199,14 +203,14 @@ def test_interleaved_lifecycle():

203 )

204 scheduler.update_from_output(scheduler_output, model_runner_output)

205 assert len(scheduler.running) == 2

202 - assert len(scheduler.waiting) == 1

206 + assert _num_waiting_requests(scheduler) == 1

207 assert len(scheduler_output.scheduled_new_reqs) == 0

208 assert scheduler_output.scheduled_cached_reqs.num_reqs == 2

209

210 # STEP 4: KVs arrive.

211 scheduler_output = scheduler.schedule()

212 assert len(scheduler.running) == 2

209 - assert len(scheduler.waiting) == 1

213 + assert _num_waiting_requests(scheduler) == 1

214 assert len(scheduler_output.scheduled_new_reqs) == 0

215 assert scheduler_output.scheduled_cached_reqs.num_reqs == 2

216

@@ -218,7 +222,7 @@ def test_interleaved_lifecycle():

222 # STEP 5: RECVed KVs are sent to ModelRunner.

223 scheduler_output = scheduler.schedule()

224 assert len(scheduler.running) == 3

221 - assert len(scheduler.waiting) == 0

225 + assert _num_waiting_requests(scheduler) == 0

226 assert len(scheduler_output.scheduled_new_reqs) == 1

227 assert scheduler_output.scheduled_cached_reqs.num_reqs == 2

228

@@ -279,14 +283,14 @@ def test_no_spurious_prefix_caching():

283 scheduler.add_request(request_remote)

284 scheduler_output = scheduler.schedule()

285 scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)

282 - assert len(scheduler.waiting) == 1

286 + assert _num_waiting_requests(scheduler) == 1

287

288 # Schedule the local prefill request. This should

289 # cause blocks to be cached, but separately from

290 scheduler.add_request(request_local)

291 scheduler_output = scheduler.schedule()

292 assert len(scheduler.running) == 1

289 - assert len(scheduler.waiting) == 1

293 + assert _num_waiting_requests(scheduler) == 1

294

295 local_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[

296 0

@@ -348,7 +352,7 @@ def test_full_block_prompt():

352 finished_recving={request_id}

353 )

354 scheduler.update_from_output(scheduler_output, model_runner_output)

351 - assert len(scheduler.waiting) == 1

355 + assert _num_waiting_requests(scheduler) == 1

356 assert request_id in scheduler.finished_recving_kv_req_ids

357

358 # # STEP (3): Run as usual.

@@ -418,7 +422,7 @@ def test_cannot_schedule_after_recv():

422 model_runner_output = create_model_runner_output(reqs=[request_normal])

423 scheduler.update_from_output(scheduler_output, model_runner_output)

424 assert len(scheduler.running) == 1

421 - assert len(scheduler.waiting) == 0

425 + assert _num_waiting_requests(scheduler) == 0

426

427 # Step 2: 5 blocks are in use (2 new for remote blocks).

428 scheduler.add_request(request_remote)

@@ -426,7 +430,7 @@ def test_cannot_schedule_after_recv():

430 model_runner_output = create_model_runner_output(reqs=[request_normal])

431 scheduler.update_from_output(scheduler_output, model_runner_output)

432 assert len(scheduler.running) == 1

429 - assert len(scheduler.waiting) == 1

433 + assert _num_waiting_requests(scheduler) == 1

434

435 # Step 3: finish recving (5 blocks in use)

436 scheduler_output = scheduler.schedule()

@@ -435,7 +439,7 @@ def test_cannot_schedule_after_recv():

439 )

440 scheduler.update_from_output(scheduler_output, model_runner_output)

441 assert len(scheduler.running) == 1

438 - assert len(scheduler.waiting) == 1

442 + assert _num_waiting_requests(scheduler) == 1

443

444 # Step 4: try to schedule, remote request is put to running list

445 # because the transfer is completed.

@@ -445,7 +449,7 @@ def test_cannot_schedule_after_recv():

449 )

450 scheduler.update_from_output(scheduler_output, model_runner_output)

451 assert len(scheduler.running) == 2

448 - assert len(scheduler.waiting) == 0

452 + assert _num_waiting_requests(scheduler) == 0

453

454 # Step 5: Remote request will be put back to waiting list

455 # because it needs new block to hold generated token.

@@ -453,7 +457,7 @@ def test_cannot_schedule_after_recv():

457 model_runner_output = create_model_runner_output(reqs=[request_normal])

458 scheduler.update_from_output(scheduler_output, model_runner_output)

459 assert len(scheduler.running) == 1

456 - assert len(scheduler.waiting) == 1

460 + assert _num_waiting_requests(scheduler) == 1

461

462 # Step 6: finish the request, free it.

463 scheduler_output = scheduler.schedule()

@@ -462,7 +466,7 @@ def test_cannot_schedule_after_recv():

466 )

467 scheduler.update_from_output(scheduler_output, model_runner_output)

468 assert len(scheduler.running) == 0

465 - assert len(scheduler.waiting) == 1

469 + assert _num_waiting_requests(scheduler) == 1

470

471 # Step 7: now we can schedule (with 2 blocks computed),

472 # request is retrieved from preempted list.

@@ -474,7 +478,7 @@ def test_cannot_schedule_after_recv():

478 )

479 scheduler.update_from_output(scheduler_output, model_runner_output)

480 assert len(scheduler.running) == 1

477 - assert len(scheduler.waiting) == 0

481 + assert _num_waiting_requests(scheduler) == 0

482

483 # Step 8: free everything.

484 scheduler_output = scheduler.schedule()

@@ -521,7 +525,7 @@ def test_cannot_recv():

525 model_runner_output = create_model_runner_output(reqs=[request_normal])

526 scheduler.update_from_output(scheduler_output, model_runner_output)

527 assert len(scheduler.running) == 1

524 - assert len(scheduler.waiting) == 0

528 + assert _num_waiting_requests(scheduler) == 0

529

530 # Step 2: 3 blocks are in use,

531 # need 3 new for remote blocks but only 2 are available.

@@ -530,7 +534,7 @@ def test_cannot_recv():

534 model_runner_output = create_model_runner_output(reqs=[request_normal])

535 scheduler.update_from_output(scheduler_output, model_runner_output)

536 assert len(scheduler.running) == 1

533 - assert len(scheduler.waiting) == 1

537 + assert _num_waiting_requests(scheduler) == 1

538 # Should not have KV transfer in progress.

539 assert request_remote.status != RequestStatus.WAITING_FOR_REMOTE_KVS

540

@@ -541,14 +545,14 @@ def test_cannot_recv():

545 )

546 scheduler.update_from_output(scheduler_output, model_runner_output)

547 assert len(scheduler.running) == 0

544 - assert len(scheduler.waiting) == 1

548 + assert _num_waiting_requests(scheduler) == 1

549

550 # Step 4: now we can initiate KV transfer (with 2 blocks computed).

551 scheduler_output = scheduler.schedule()

552 model_runner_output = create_model_runner_output(reqs=[])

553 scheduler.update_from_output(scheduler_output, model_runner_output)

554 assert len(scheduler.running) == 0

551 - assert len(scheduler.waiting) == 1

555 + assert _num_waiting_requests(scheduler) == 1

556 assert request_remote.status == RequestStatus.WAITING_FOR_REMOTE_KVS

557

558 # Step 5: finish recving (5 blocks in use)

@@ -558,14 +562,14 @@ def test_cannot_recv():

562 )

563 scheduler.update_from_output(scheduler_output, model_runner_output)

564 assert len(scheduler.running) == 0

561 - assert len(scheduler.waiting) == 1

565 + assert _num_waiting_requests(scheduler) == 1

566

567 # Step 6: schedule remote request

568 scheduler_output = scheduler.schedule()

569 model_runner_output = create_model_runner_output(reqs=[request_remote])

570 scheduler.update_from_output(scheduler_output, model_runner_output)

571 assert len(scheduler.running) == 1

568 - assert len(scheduler.waiting) == 0

572 + assert _num_waiting_requests(scheduler) == 0

573

574 # Step 7: free everything.

575 scheduler_output = scheduler.schedule()

vllm/v1/core/sched/scheduler.py

+101 -64

@@ -45,7 +45,11 @@ from vllm.v1.core.sched.output import (

45 NewRequestData,

46 SchedulerOutput,

47 )

48 -from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue

48 +from vllm.v1.core.sched.request_queue import (

49 + RequestQueue,

50 + SchedulingPolicy,

51 + create_request_queue,

52 +)

53 from vllm.v1.core.sched.utils import check_stop, remove_all

54 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs

55 from vllm.v1.kv_cache_interface import KVCacheConfig

@@ -160,6 +164,8 @@ class Scheduler(SchedulerInterface):

164 ) from e

165 # Priority queues for requests.

166 self.waiting = create_request_queue(self.policy)

167 + # requests skipped in waiting flow due async deps or constraints.

168 + self.skipped_waiting = create_request_queue(self.policy)

169 self.running: list[Request] = []

170

171 # The request IDs that are finished in between the previous and the

@@ -531,52 +537,29 @@ class Scheduler(SchedulerInterface):

537

538 # Next, schedule the WAITING requests.

539 if not preempted_reqs and self._pause_state == PauseState.UNPAUSED:

534 - # Use a temporary RequestQueue to collect requests that need to be

535 - # skipped and put back at the head of the waiting queue later

536 - skipped_waiting_requests = create_request_queue(self.policy)

540 + step_skipped_waiting = create_request_queue(self.policy)

541

538 - while self.waiting and token_budget > 0:

542 + while (self.waiting or self.skipped_waiting) and token_budget > 0:

543 if len(self.running) == self.max_num_running_reqs:

544 break

545

542 - request = self.waiting.peek_request()

546 + request_queue = self._select_waiting_queue_for_scheduling()

547 + assert request_queue is not None

548 +

549 + request = request_queue.peek_request()

550 request_id = request.request_id

551

545 - # KVTransfer: skip request if still waiting for remote kvs.

546 - if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:

547 - is_ready = self._update_waiting_for_remote_kv(request)

548 - if is_ready:

549 - if request.num_preemptions:

550 - # We must be loading for a resumed preemption

551 - # rather than a new request.

552 - request.status = RequestStatus.PREEMPTED

553 - else:

554 - request.status = RequestStatus.WAITING

555 - else:

552 + # try to promote blocked statuses while traversing skipped queue.

553 + if self._is_blocked_waiting_status(

554 + request.status

555 + ) and not self._try_promote_blocked_waiting_request(request):

556 + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:

557 logger.debug(

558 "%s is still in WAITING_FOR_REMOTE_KVS state.",

559 request_id,

560 )

560 - self.waiting.pop_request()

561 - skipped_waiting_requests.prepend_request(request)

562 - continue

563 -

564 - # Skip request if the structured output request is still waiting

565 - # for FSM compilation.

566 - if request.status == RequestStatus.WAITING_FOR_FSM:

567 - structured_output_req = request.structured_output_request

568 - if structured_output_req and structured_output_req.grammar:

569 - request.status = RequestStatus.WAITING

570 - else:

571 - self.waiting.pop_request()

572 - skipped_waiting_requests.prepend_request(request)

573 - continue

574 -

575 - # Streaming: skip request if still waiting for next streaming req.

576 - if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ:

577 - assert not request.streaming_queue

578 - self.waiting.pop_request()

579 - skipped_waiting_requests.prepend_request(request)

561 + request_queue.pop_request()

562 + step_skipped_waiting.prepend_request(request)

563 continue

564

565 # Check that adding the request still respects the max_loras

@@ -590,8 +573,8 @@ class Scheduler(SchedulerInterface):

573 )

574 ):

575 # Scheduling would exceed max_loras, skip.

593 - self.waiting.pop_request()

594 - skipped_waiting_requests.prepend_request(request)

576 + request_queue.pop_request()

577 + step_skipped_waiting.prepend_request(request)

578 continue

579

580 num_external_computed_tokens = 0

@@ -617,8 +600,8 @@ class Scheduler(SchedulerInterface):

600 # The request cannot be scheduled because

601 # the KVConnector couldn't determine

602 # the number of matched tokens.

620 - self.waiting.pop_request()

621 - skipped_waiting_requests.prepend_request(request)

603 + request_queue.pop_request()

604 + step_skipped_waiting.prepend_request(request)

605 continue

606

607 request.num_external_computed_tokens = ext_tokens

@@ -761,14 +744,12 @@ class Scheduler(SchedulerInterface):

744 preempted=request.num_preemptions > 0,

745 )

746

764 - # Request was already popped from self.waiting

765 - # unless it was re-added above due to new_blocks being None.

766 - request = self.waiting.pop_request()

747 + request = request_queue.pop_request()

748 if load_kv_async:

749 # If loading async, allocate memory and put request

750 # into the WAITING_FOR_REMOTE_KV state.

770 - skipped_waiting_requests.prepend_request(request)

751 request.status = RequestStatus.WAITING_FOR_REMOTE_KVS

752 + step_skipped_waiting.prepend_request(request)

753 # Set num_computed_tokens even though KVs are not yet loaded.

754 # request.num_computed_tokens will not be used anywhere until

755 # the request finished the KV transfer.

@@ -825,9 +806,9 @@ class Scheduler(SchedulerInterface):

806 if self.ec_connector is not None:

807 self.ec_connector.update_state_after_alloc(request, i)

808

828 - # Put back any skipped requests at the head of the waiting queue

829 - if skipped_waiting_requests:

830 - self.waiting.prepend_requests(skipped_waiting_requests)

809 + # re-queue requests skipped in this pass ahead of older skipped items.

810 + if step_skipped_waiting:

811 + self.skipped_waiting.prepend_requests(step_skipped_waiting)

812

813 # Check if the scheduling constraints are satisfied.

814 total_num_scheduled_tokens = sum(num_scheduled_tokens.values())

@@ -1531,6 +1512,32 @@ class Scheduler(SchedulerInterface):

1512

1513 return engine_core_outputs

1514

1515 + @staticmethod

1516 + def _is_blocked_waiting_status(status: RequestStatus) -> bool:

1517 + return status in (

1518 + RequestStatus.WAITING_FOR_FSM,

1519 + RequestStatus.WAITING_FOR_REMOTE_KVS,

1520 + RequestStatus.WAITING_FOR_STREAMING_REQ,

1521 + )

1522 +

1523 + def _enqueue_waiting_request(self, request: Request) -> None:

1524 + if self._is_blocked_waiting_status(request.status):

1525 + self.skipped_waiting.add_request(request)

1526 + else:

1527 + self.waiting.add_request(request)

1528 +

1529 + def _select_waiting_queue_for_scheduling(self) -> RequestQueue | None:

1530 + if self.policy == SchedulingPolicy.FCFS:

1531 + return self.skipped_waiting or self.waiting or None

1532 +

1533 + # PRIORITY mode: compare queue heads when both queues are non-empty.

1534 + if self.waiting and self.skipped_waiting:

1535 + waiting_req = self.waiting.peek_request()

1536 + skipped_req = self.skipped_waiting.peek_request()

1537 + return self.waiting if waiting_req < skipped_req else self.skipped_waiting

1538 +

1539 + return self.waiting or self.skipped_waiting or None

1540 +

1541 def _handle_stopped_request(self, request: Request) -> bool:

1542 """Return True if finished (can be False for resumable requests)."""

1543 if not request.resumable:

@@ -1546,7 +1553,7 @@ class Scheduler(SchedulerInterface):

1553 request.status = RequestStatus.WAITING_FOR_STREAMING_REQ

1554 self.num_waiting_for_streaming_input += 1

1555

1549 - self.waiting.add_request(request)

1556 + self._enqueue_waiting_request(request)

1557 return False

1558

1559 def _get_routed_experts(self, request: Request) -> np.ndarray | None:

@@ -1677,7 +1684,7 @@ class Scheduler(SchedulerInterface):

1684

1685 def get_request_counts(self) -> tuple[int, int]:

1686 """Returns (num_running_reqs, num_waiting_reqs)."""

1680 - return len(self.running), len(self.waiting)

1687 + return len(self.running), len(self.waiting) + len(self.skipped_waiting)

1688

1689 def add_request(self, request: Request) -> None:

1690 existing = self.requests.get(request.request_id)

@@ -1696,7 +1703,7 @@ class Scheduler(SchedulerInterface):

1703 else:

1704 if request.resumable:

1705 request.streaming_queue = deque()

1699 - self.waiting.add_request(request)

1706 + self._enqueue_waiting_request(request)

1707 self.requests[request.request_id] = request

1708 if self.log_stats:

1709 request.record_event(EngineCoreEventType.QUEUED)

@@ -1747,6 +1754,7 @@ class Scheduler(SchedulerInterface):

1754 self.running = remove_all(self.running, running_requests_to_remove)

1755 if waiting_requests_to_remove:

1756 self.waiting.remove_requests(waiting_requests_to_remove)

1757 + self.skipped_waiting.remove_requests(waiting_requests_to_remove)

1758

1759 # Second pass: set status and free requests

1760 for request in valid_requests:

@@ -1798,7 +1806,11 @@ class Scheduler(SchedulerInterface):

1806 return 0

1807 if self._pause_state == PauseState.PAUSED_NEW:

1808 return len(self.running)

1801 - num_waiting = len(self.waiting) - self.num_waiting_for_streaming_input

1809 + num_waiting = (

1810 + len(self.waiting)

1811 + + len(self.skipped_waiting)

1812 + - self.num_waiting_for_streaming_input

1813 + )

1814 return num_waiting + len(self.running)

1815

1816 def has_finished_requests(self) -> bool:

@@ -1898,7 +1910,7 @@ class Scheduler(SchedulerInterface):

1910 )

1911 return SchedulerStats(

1912 num_running_reqs=len(self.running),

1901 - num_waiting_reqs=len(self.waiting),

1913 + num_waiting_reqs=len(self.waiting) + len(self.skipped_waiting),

1914 kv_cache_usage=self.kv_cache_manager.usage,

1915 encoder_cache_usage=self._get_encoder_cache_usage(),

1916 prefix_cache_stats=prefix_cache_stats,

@@ -1981,21 +1993,15 @@ class Scheduler(SchedulerInterface):

1993

1994 return self.connector.request_finished_all_groups(request, block_ids)

1995

1984 - def _update_waiting_for_remote_kv(self, request: Request) -> bool:

1996 + def _update_waiting_for_remote_kv(self, request: Request) -> None:

1997 """

1986 - KV Connector: check if the request_id is finished_recving.

1987 -

1988 - The finished_recving_kv_req_ids list is populated

1989 - on the previous steps()'s update_from_output based

1990 - on the worker side connector.

1998 + KV Connector: update request state after async recv is finished.

1999

2000 When the kv transfer is ready, we cache the blocks

2001 and the request state will be moved back to WAITING from

2002 WAITING_FOR_REMOTE_KV.

2003 """

2004 assert self.connector is not None

1997 - if request.request_id not in self.finished_recving_kv_req_ids:

1998 - return False

2005

2006 if request.request_id in self.failed_recving_kv_req_ids:

2007 # Request had KV load failures; num_computed_tokens was already

@@ -2023,9 +2029,40 @@ class Scheduler(SchedulerInterface):

2029 if request.num_cached_tokens < 0:

2030 request.num_cached_tokens = request.num_computed_tokens

2031

2026 - # Return that we are ready.

2032 self.finished_recving_kv_req_ids.remove(request.request_id)

2028 - return True

2033 +

2034 + def _try_promote_blocked_waiting_request(self, request: Request) -> bool:

2035 + """

2036 + Try to promote a blocked waiting request back to schedulable states.

2037 + """

2038 + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:

2039 + # finished_recving_kv_req_ids is populated during

2040 + # update_from_output(), based on worker-side connector signals

2041 + # in KVConnectorOutput.finished_recving

2042 + if request.request_id not in self.finished_recving_kv_req_ids:

2043 + return False

2044 + self._update_waiting_for_remote_kv(request)

2045 + if request.num_preemptions:

2046 + request.status = RequestStatus.PREEMPTED

2047 + else:

2048 + request.status = RequestStatus.WAITING

2049 + return True

2050 +

2051 + if request.status == RequestStatus.WAITING_FOR_FSM:

2052 + structured_output_req = request.structured_output_request

2053 + if not (structured_output_req and structured_output_req.grammar):

2054 + return False

2055 + request.status = RequestStatus.WAITING

2056 + return True

2057 +

2058 + if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ:

2059 + assert not request.streaming_queue

2060 + return False

2061 +

2062 + raise AssertionError(

2063 + "Unexpected blocked waiting status in promotion: "

2064 + f"{request.status.name} for request {request.request_id}"

2065 + )

2066

2067 def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput):

2068 """

@@ -2172,7 +2209,7 @@ class Scheduler(SchedulerInterface):

2209 # handle async KV loads (not cached yet, evict_blocks=False)

2210 async_load_reqs = (

2211 req

2175 - for req in self.waiting

2212 + for req in self.skipped_waiting

2213 if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS

2214 )

2215 async_failed_req_ids, num_failed_tokens, _ = (