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

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 2, 23:15 UTC

Keep scheduler alive for delayed KV connector frees (#43433)

Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>

Gabriel Wu committed May 23, 2026 at 14:23 UTC 82536acc54fc8d828b0d8add9d093aee0da1370e

4 files changed +69 -12

tests/v1/core/test_scheduler.py

+39

@@ -3928,6 +3928,45 @@ def test_abort_request_finished_recving():

3928 assert not scheduler.finished_recving_kv_req_ids

3929

3930

3931 +def test_delayed_kv_connector_free_keeps_scheduler_active():

3932 + scheduler = create_scheduler(use_kv_connector=True)

3933 + queued_request, request = create_requests(

3934 + num_requests=2, req_ids=["queued", "finished"]

3935 + )

3936 + scheduler.add_request(queued_request)

3937 +

3938 + assert not scheduler.has_finished_requests()

3939 +

3940 + request.status = RequestStatus.FINISHED_STOPPED

3941 + scheduler.requests[request.request_id] = request

3942 + scheduler.finished_req_ids = set()

3943 +

3944 + assert scheduler.has_finished_requests()

3945 + assert scheduler.has_requests()

3946 +

3947 + scheduler_output = SchedulerOutput(

3948 + scheduled_new_reqs=[],

3949 + scheduled_cached_reqs=CachedRequestData.make_empty(),

3950 + num_scheduled_tokens={},

3951 + total_num_scheduled_tokens=0,

3952 + scheduled_encoder_inputs={},

3953 + scheduled_spec_decode_tokens={},

3954 + num_common_prefix_blocks=[],

3955 + finished_req_ids=set(),

3956 + free_encoder_mm_hashes=[],

3957 + )

3958 + model_runner_output = ModelRunnerOutput(

3959 + req_ids=[],

3960 + req_id_to_index={},

3961 + kv_connector_output=KVConnectorOutput(finished_sending={request.request_id}),

3962 + )

3963 +

3964 + scheduler.update_from_output(scheduler_output, model_runner_output)

3965 +

3966 + assert request.request_id not in scheduler.requests

3967 + assert not scheduler.has_finished_requests()

3968 +

3969 +

3970 # ==============================================================================

3971 # Variable-length encoder cross-attention block allocation tests

3972 # ==============================================================================

tests/v1/kv_connector/unit/test_multi_connector.py

+16 -6

@@ -258,10 +258,12 @@ def test_multi_example_connector_consistency():

258 )

259

260 events = get_connector_events()

261 + storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])

262 + storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])

263 # First event is bind_gpu_block_pool from initialization, then

264 # set_xfer_handshake_metadata, then on_new_request when the request is enqueued,

265 # then get_num_new_matched_tokens and update_state_after_alloc from generate().

264 - assert events["storage1-SCHEDULER"][:6] == [

266 + assert storage1_scheduler_events[:6] == [

267 "bind_gpu_block_pool",

268 "set_xfer_handshake_metadata",

269 "on_new_request",

@@ -281,7 +283,7 @@ def test_multi_example_connector_consistency():

283 "wait_for_layer_load",

284 "save_kv_layer",

285 ]

284 - assert events["storage2-SCHEDULER"][:6] == [

286 + assert storage2_scheduler_events[:6] == [

287 "bind_gpu_block_pool",

288 "set_xfer_handshake_metadata",

289 "on_new_request",

@@ -312,13 +314,15 @@ def test_multi_example_connector_consistency():

314 # connector so update_state_after_alloc will be with allocated blocks

315 # on that one but with zero blocks for others (first nonzero match is

316 # chosen).

315 - assert events["storage1-SCHEDULER"][:4] == [

317 + storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])

318 + storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])

319 + assert storage1_scheduler_events[:4] == [

320 "on_new_request",

321 "get_num_new_matched_tokens 0",

322 "update_state_after_alloc num_blocks=[7] 96",

323 "build_connector_meta",

324 ]

321 - assert events["storage2-SCHEDULER"][:4] == [

325 + assert storage2_scheduler_events[:4] == [

326 "on_new_request",

327 "get_num_new_matched_tokens 0",

328 "update_state_after_alloc num_blocks=[0] 0",

@@ -340,13 +344,15 @@ def test_multi_example_connector_consistency():

344 # return 0 from the first connector, but the second connector should have

345 # a hit, so update_state_after_alloc will only be called with allocated

346 # blocks for the second connector.

343 - assert events["storage1-SCHEDULER"][:4] == [

347 + storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])

348 + storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])

349 + assert storage1_scheduler_events[:4] == [

350 "on_new_request",

351 "get_num_new_matched_tokens 0",

352 "update_state_after_alloc num_blocks=[0] 0",

353 "build_connector_meta",

354 ]

349 - assert events["storage2-SCHEDULER"][:4] == [

355 + assert storage2_scheduler_events[:4] == [

356 "on_new_request",

357 "get_num_new_matched_tokens 0",

358 "update_state_after_alloc num_blocks=[7] 96",

@@ -358,6 +364,10 @@ def test_multi_example_connector_consistency():

364 shutil.rmtree(storage_2_path)

365

366

367 +def _ignore_event_collection(events: list[str]) -> list[str]:

368 + return [event for event in events if event != "take_events"]

369 +

370 +

371 def get_connector_events() -> dict[str, list[str]]:

372 # Read in connector events and reset the files.

373 import glob

vllm/v1/core/sched/scheduler.py

+10 -1

@@ -1882,7 +1882,16 @@ class Scheduler(SchedulerInterface):

1882 return num_waiting + len(self.running)

1883

1884 def has_finished_requests(self) -> bool:

1885 - return len(self.finished_req_ids) > 0

1885 + if self.finished_req_ids:

1886 + return True

1887 + if self.connector is None:

1888 + return False

1889 + # Finished requests waiting on delayed connector cleanup remain in

1890 + # self.requests after they have been removed from scheduling queues.

1891 + num_in_queues = (

1892 + len(self.waiting) + len(self.skipped_waiting) + len(self.running)

1893 + )

1894 + return len(self.requests) > num_in_queues

1895

1896 def reset_prefix_cache(

1897 self, reset_running_requests: bool = False, reset_connector: bool = False

vllm/v1/engine/core.py

+4 -5

@@ -1242,11 +1242,10 @@ class EngineCoreProc(EngineCore):

1242 # Post-step hook.

1243 self.post_step(model_executed)

1244

1245 - # If no model execution happened but there are waiting requests

1246 - # (e.g., WAITING_FOR_REMOTE_KVS), yield the GIL briefly to allow

1247 - # background threads (like NIXL handshake) to make progress.

1248 - # Without this, the tight polling loop can starve background threads.

1249 - if not model_executed and self.scheduler.has_unfinished_requests():

1245 + # If no model execution happened but there is still scheduler work

1246 + # (e.g. WAITING_FOR_REMOTE_KVS or delayed KV connector frees), yield

1247 + # the GIL briefly to allow background transfer threads to make progress.

1248 + if not model_executed and self.scheduler.has_requests():

1249 time.sleep(0.001)

1250

1251 return model_executed