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 8, 22:45 UTC
[KV Connector][BugFix] scheduler: Delay freeing blocks of aborted async loads (#32255)
Fixes a not-yet-reported case where it was possible for blocks to be freed by an abort before an async transfer completed, resulting in corrupted KV data. Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Feb 4, 2026 at 13:16 UTC
8e3269086916d81b010a2d6209784a106ac2994a
3 files changed
+123
-6
tests/v1/core/test_scheduler.py
+49
@@ -3473,3 +3473,52 @@ def test_prepend_skipped_requests_order():
3473
3474
# verify waiting order is preserved
3475
assert list(scheduler.waiting) == expected_waiting_reqs
3476
+
3477
+
3478
+def test_abort_request_waiting_for_remote_kvs():
3479
+ scheduler = create_scheduler(use_kv_connector=True)
3480
+
3481
+ # add a single request
3482
+ request = create_requests(num_requests=1)[0]
3483
+ scheduler.add_request(request)
3484
+
3485
+ # set request to waiting for remote KVs, and abort it
3486
+ request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
3487
+ scheduler.finish_requests((request.request_id,), RequestStatus.FINISHED_ABORTED)
3488
+ assert request.status == RequestStatus.FINISHED_ABORTED
3489
+
3490
+ # verify request is not deleted
3491
+ assert request.request_id in scheduler.requests
3492
+
3493
+ # finish recving request
3494
+ scheduler_output = scheduler.schedule()
3495
+ model_runner_output = ModelRunnerOutput(
3496
+ req_ids=[],
3497
+ req_id_to_index={},
3498
+ kv_connector_output=KVConnectorOutput(finished_recving={request.request_id}),
3499
+ )
3500
+ scheduler.update_from_output(scheduler_output, model_runner_output)
3501
+
3502
+ # assert request is deleted
3503
+ assert request.request_id not in scheduler.requests
3504
+ assert not scheduler.finished_recving_kv_req_ids
3505
+
3506
+
3507
+def test_abort_request_finished_recving():
3508
+ scheduler = create_scheduler(use_kv_connector=True)
3509
+
3510
+ # add a single request
3511
+ request = create_requests(num_requests=1)[0]
3512
+ scheduler.add_request(request)
3513
+
3514
+ # set request to waiting for remote KVs, finished but not yet updated
3515
+ request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
3516
+ scheduler.finished_recving_kv_req_ids.add(request.request_id)
3517
+
3518
+ # abort request
3519
+ scheduler.finish_requests((request.request_id,), RequestStatus.FINISHED_ABORTED)
3520
+ assert request.status == RequestStatus.FINISHED_ABORTED
3521
+
3522
+ # verify request is deleted
3523
+ assert request.request_id not in scheduler.requests
3524
+ assert not scheduler.finished_recving_kv_req_ids
tests/v1/kv_connector/unit/test_offloading_connector.py
+53
-2
@@ -42,7 +42,7 @@ from vllm.v1.kv_offload.worker.worker import (
42
TransferSpec,
43
)
44
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput
45
-from vllm.v1.request import Request
45
+from vllm.v1.request import Request, RequestStatus
46
47
from .utils import (
48
EOS_TOKEN_ID,
@@ -355,7 +355,7 @@ class RequestRunner:
355
self.scheduler.update_from_output(scheduler_output, model_runner_output)
356
357
if (
358
- prev_token_id is EOS_TOKEN_ID
358
+ prev_token_id == EOS_TOKEN_ID
359
and prev_token_id != token_id
360
and self.scheduler.requests
361
):
@@ -730,6 +730,57 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner):
730
assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
731
732
733
+def test_abort_loading_requests(request_runner):
734
+ offloaded_block_size = 12
735
+ gpu_block_size = 4
736
+ num_gpu_blocks = 100
737
+
738
+ runner = request_runner(
739
+ offloaded_block_size=offloaded_block_size,
740
+ gpu_block_size=gpu_block_size,
741
+ num_gpu_blocks=num_gpu_blocks,
742
+ )
743
+
744
+ # store 1 blocks
745
+ runner.new_request(token_ids=[0] * offloaded_block_size)
746
+ runner.manager.prepare_store.side_effect = (
747
+ lambda block_hashes: generate_store_output(block_hashes)
748
+ )
749
+ runner.run(
750
+ decoded_tokens=[EOS_TOKEN_ID],
751
+ expected_stored_gpu_block_indexes=(0, 1, 2),
752
+ )
753
+
754
+ # start a request to load the first block, but don't complete
755
+ runner.scheduler.reset_prefix_cache()
756
+ runner.new_request(token_ids=[0] * offloaded_block_size)
757
+ runner.manager.lookup.return_value = 1
758
+ runner.run(
759
+ decoded_tokens=[],
760
+ complete_transfers=False,
761
+ )
762
+
763
+ # request triggered a load
764
+ transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
765
+ assert transfer_jobs
766
+
767
+ # abort request
768
+ req_id = str(runner.req_id)
769
+ runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED)
770
+
771
+ # verify request is not deleted
772
+ assert req_id in runner.scheduler.requests
773
+
774
+ # complete loading request
775
+ runner.run(
776
+ decoded_tokens=[],
777
+ expected_loaded_gpu_block_indexes=(0, 1, 2),
778
+ )
779
+
780
+ # assert request is deleted
781
+ assert req_id not in runner.scheduler.requests
782
+
783
+
784
class TestOffloadingConnectorStats:
785
"""Tests for OffloadingConnector stats reconstruction and operations."""
786
vllm/v1/core/sched/scheduler.py
+21
-4
@@ -1670,19 +1670,30 @@ class Scheduler(SchedulerInterface):
1670
1671
# Second pass: set status and free requests
1672
for request in valid_requests:
1673
+ delay_free_blocks = False
1674
+ if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
1675
+ delay_free_blocks = (
1676
+ request.request_id not in self.finished_recving_kv_req_ids
1677
+ )
1678
+ self.finished_recving_kv_req_ids.discard(request.request_id)
1679
+ self.failed_recving_kv_req_ids.discard(request.request_id)
1680
+
1681
request.status = finished_status
1674
- self._free_request(request)
1682
+ self._free_request(request, delay_free_blocks=delay_free_blocks)
1683
1676
- def _free_request(self, request: Request) -> dict[str, Any] | None:
1684
+ def _free_request(
1685
+ self, request: Request, delay_free_blocks: bool = False
1686
+ ) -> dict[str, Any] | None:
1687
assert request.is_finished()
1688
1679
- delay_free_blocks, kv_xfer_params = self._connector_finished(request)
1689
+ connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
1690
self.encoder_cache_manager.free(request)
1691
request_id = request.request_id
1692
self.finished_req_ids.add(request_id)
1693
if self.finished_req_ids_dict is not None:
1694
self.finished_req_ids_dict[request.client_index].add(request_id)
1695
1696
+ delay_free_blocks |= connector_delay_free_blocks
1697
if not delay_free_blocks:
1698
self._free_blocks(request)
1699
@@ -1954,7 +1965,13 @@ class Scheduler(SchedulerInterface):
1965
# KV Connector:: update recv and send status from last step.
1966
for req_id in kv_connector_output.finished_recving or ():
1967
logger.debug("Finished recving KV transfer for request %s", req_id)
1957
- self.finished_recving_kv_req_ids.add(req_id)
1968
+ assert req_id in self.requests
1969
+ req = self.requests[req_id]
1970
+ if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
1971
+ self.finished_recving_kv_req_ids.add(req_id)
1972
+ else:
1973
+ assert RequestStatus.is_finished(req.status)
1974
+ self._free_blocks(self.requests[req_id])
1975
for req_id in kv_connector_output.finished_sending or ():
1976
logger.debug("Finished sending KV transfer for request %s", req_id)
1977
assert req_id in self.requests