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

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 10, 09:00 UTC

[BugFix] scheduler: Fix resuming of preempted requests after async load (#31583)

Signed-off-by: Or Ozeri <oro@il.ibm.com>

Or Ozeri committed Jan 10, 2026 at 22:39 UTC 028599739d617b76c55179a9dabddb458e1fe035

3 files changed +28 -4

tests/v1/core/test_scheduler.py

+21 -2

@@ -1261,10 +1261,11 @@ def test_kv_connector_unable_to_allocate(use_ec_connector, ec_role):

1261 assert len(scheduler.waiting) == 0

1262

1263

1264 +@pytest.mark.parametrize("is_async", [False, True])

1265 @pytest.mark.parametrize(

1266 "use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")]

1267 )

1267 -def test_kv_connector_handles_preemption(use_ec_connector, ec_role):

1268 +def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role):

1269 """

1270 Test whether scheduler with KVConnector is able to handle

1271 unable to allocate (run out of blocks in allocate_slots().

@@ -1277,7 +1278,9 @@ def test_kv_connector_handles_preemption(use_ec_connector, ec_role):

1278 NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE

1279 scheduler = create_scheduler(

1280 enable_prefix_caching=True,

1280 - use_kv_connector=mock_kv(matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=False),

1281 + use_kv_connector=mock_kv(

1282 + matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=is_async

1283 + ),

1284 block_size=BLOCK_SIZE,

1285 num_blocks=NUM_BLOCKS,

1286 # encoder connector should not affect test results

@@ -1315,6 +1318,12 @@ def test_kv_connector_handles_preemption(use_ec_connector, ec_role):

1318

1319 # All can be scheduled - 1st token.

1320 output = scheduler.schedule()

1321 + if is_async:

1322 + assert len(scheduler.waiting) == 2

1323 + assert scheduler.running == []

1324 + _step_until_kv_transfer_finished(scheduler, req_ids)

1325 + output = scheduler.schedule()

1326 +

1327 _assert_right_scheduler_output(

1328 output,

1329 # 2 remote kv cache hits.

@@ -1367,6 +1376,12 @@ def test_kv_connector_handles_preemption(use_ec_connector, ec_role):

1376 # Restarts the preempted request - generate 3rd token.

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

1378 output = scheduler.schedule()

1379 + if is_async:

1380 + waiting_req_ids = [req.request_id for req in scheduler.waiting]

1381 + assert len(waiting_req_ids) == 1

1382 + _step_until_kv_transfer_finished(scheduler, waiting_req_ids)

1383 + output = scheduler.schedule()

1384 +

1385 _assert_right_scheduler_output(

1386 output,

1387 # 1 remote kv_cache hit!

@@ -1377,6 +1392,8 @@ def test_kv_connector_handles_preemption(use_ec_connector, ec_role):

1392 )

1393 assert len(scheduler.running) == 1

1394 assert len(scheduler.waiting) == 0

1395 + assert output.scheduled_cached_reqs.num_reqs == 1

1396 + assert output.scheduled_new_reqs == []

1397 _ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)

1398 assert len(scheduler.running) == 1

1399 assert len(scheduler.waiting) == 0

@@ -1389,6 +1406,8 @@ def test_kv_connector_handles_preemption(use_ec_connector, ec_role):

1406 num_requests=0,

1407 expected_num_scheduled_tokens=1,

1408 )

1409 + assert output.scheduled_cached_reqs.num_reqs == 1

1410 + assert output.scheduled_new_reqs == []

1411 assert len(scheduler.running) == 1

1412 _ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)

1413 assert len(scheduler.running) == 0

vllm/v1/core/sched/scheduler.py

+6 -1

@@ -445,7 +445,12 @@ class Scheduler(SchedulerInterface):

445 if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:

446 is_ready = self._update_waiting_for_remote_kv(request)

447 if is_ready:

448 - request.status = RequestStatus.WAITING

448 + if request.num_preemptions:

449 + # We must be loading for a resumed preemption

450 + # rather than a new request.

451 + request.status = RequestStatus.PREEMPTED

452 + else:

453 + request.status = RequestStatus.WAITING

454 else:

455 logger.debug(

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

vllm/v1/request.py

+1 -1

@@ -123,7 +123,7 @@ class Request:

123 # indicates that the output is corrupted

124 self.num_nans_in_logits = 0

125

126 - # The number of requests being preempted by the scheduler

126 + # The number of times this request has been preempted by the scheduler.

127 self.num_preemptions = 0

128

129 # The number of tokens that have been computed remotely.