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 3, 10:30 UTC
[KV Connector][Offloading] Flush all pending jobs on last step (#42611)
Signed-off-by: Liran Schour <lirans@il.ibm.com> Signed-off-by: liranschour <liranschour@users.noreply.github.com> Co-authored-by: Or Ozeri <or@ozery.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
liranschour committed
May 18, 2026 at 15:59 UTC
e5417657e55ec2f42809816e4aa5c9753f390cdd
3 files changed
+71
-14
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+55
-10
@@ -258,9 +258,12 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
258
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
259
generate_store_output(keys)
260
)
261
+ # With sync scheduling, all-finished flush fires within this run.
262
+ # With async scheduling, the finish is delayed so flush fires later.
263
runner.run(
264
decoded_tokens=[EOS_TOKEN_ID],
265
expected_stored=(0, 1, 2),
266
+ expected_flushed=(0, 1, 2) if not async_scheduling else (),
267
)
268
269
# start a request to load the first block, but don't complete
@@ -325,6 +328,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
328
runner.run(
329
decoded_tokens=[EOS_TOKEN_ID],
330
expected_stored=(0, 1, 2),
331
+ expected_flushed=(0, 1, 2) if not async_scheduling else (),
332
)
333
334
# start a request to load the first block, but don't complete
@@ -351,6 +355,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
355
runner.run(
356
decoded_tokens=[],
357
expected_loaded=(0, 1, 2),
358
+ expected_flushed=(0, 1, 2),
359
)
360
361
# assert request is deleted
@@ -766,7 +771,11 @@ def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bo
771
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
772
generate_store_output(keys)
773
)
769
- runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0, 1, 2))
774
+ runner.run(
775
+ decoded_tokens=[EOS_TOKEN_ID],
776
+ expected_stored=(0, 1, 2),
777
+ expected_flushed=(0, 1, 2) if not async_scheduling else (),
778
+ )
779
780
# Reset GPU prefix cache so the next request must load from CPU.
781
runner.scheduler.reset_prefix_cache()
@@ -831,8 +840,13 @@ def test_fence_at_update_state_after_alloc(request_runner):
840
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
841
generate_store_output(keys)
842
)
834
- runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)
835
- assert runner.connector_scheduler._block_id_to_pending_jobs
843
+ runner.run(
844
+ decoded_tokens=[EOS_TOKEN_ID],
845
+ complete_transfers=False,
846
+ expected_stored=(0,),
847
+ expected_flushed=(0,),
848
+ )
849
+ assert runner.connector_scheduler._block_id_to_pending_jobs == {}
850
851
runner.scheduler.reset_prefix_cache()
852
runner.new_request(token_ids=[0] * 4)
@@ -843,8 +857,6 @@ def test_fence_at_update_state_after_alloc(request_runner):
857
runner.run(
858
decoded_tokens=[],
859
complete_transfers=False,
846
- expected_stored=(0,),
847
- expected_flushed=(0,),
860
)
861
assert runner.connector_scheduler._block_id_to_pending_jobs == {}
862
@@ -864,8 +876,13 @@ def test_fence_at_build_store_jobs(request_runner):
876
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
877
generate_store_output(keys)
878
)
867
- runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)
868
- assert runner.connector_scheduler._block_id_to_pending_jobs
879
+ runner.run(
880
+ decoded_tokens=[EOS_TOKEN_ID],
881
+ complete_transfers=False,
882
+ expected_stored=(0,),
883
+ expected_flushed=(0,),
884
+ )
885
+ assert runner.connector_scheduler._block_id_to_pending_jobs == {}
886
887
runner.scheduler.reset_prefix_cache()
888
runner.new_request(token_ids=[1] * 4)
@@ -875,8 +892,6 @@ def test_fence_at_build_store_jobs(request_runner):
892
)
893
runner.run(
894
decoded_tokens=[EOS_TOKEN_ID],
878
- expected_stored=(0,),
879
- expected_flushed=(0,),
895
)
896
assert runner.connector_scheduler._block_id_to_pending_jobs == {}
897
@@ -921,6 +936,32 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
936
assert runner.manager.complete_store.call_count == 0
937
938
939
+def test_flush_all_jobs_when_no_requests_remain(request_runner):
940
+ """When all tracked requests are finished, build_connector_meta flushes
941
+ all pending jobs since there will be no future step to complete them."""
942
+ block_size = 4
943
+ block_size_factor = 1
944
+ offloaded_block_size = block_size * block_size_factor
945
+
946
+ runner = request_runner(
947
+ block_size=block_size,
948
+ num_gpu_blocks=100,
949
+ async_scheduling=False,
950
+ block_size_factor=block_size_factor,
951
+ )
952
+
953
+ runner.new_request(token_ids=[0] * offloaded_block_size)
954
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
955
+ generate_store_output(keys)
956
+ )
957
+ runner.run(
958
+ decoded_tokens=[EOS_TOKEN_ID],
959
+ complete_transfers=False,
960
+ expected_stored=(0,),
961
+ expected_flushed=(0,),
962
+ )
963
+
964
+
965
@pytest.mark.parametrize("async_scheduling", [True, False])
966
def test_reset_cache(request_runner, async_scheduling: bool):
967
"""reset_cache flushes in-flight loads, calls manager.reset_cache(), resets
@@ -942,7 +983,11 @@ def test_reset_cache(request_runner, async_scheduling: bool):
983
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
984
generate_store_output(keys)
985
)
945
- runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0, 1, 2))
986
+ runner.run(
987
+ decoded_tokens=[EOS_TOKEN_ID],
988
+ expected_stored=(0, 1, 2),
989
+ expected_flushed=(0, 1, 2) if not async_scheduling else (),
990
+ )
991
992
# Reset GPU prefix cache then start a request that loads from CPU.
993
# Leave the load in-flight so that reset_cache must flush it.
tests/v1/kv_connector/unit/offloading_connector/utils.py
+8
-4
@@ -348,10 +348,14 @@ class RequestRunner:
348
def _parse_transfers(self):
349
for transfer_spec in self.offloading_spec.get_flushed_transfers():
350
src_spec, dst_spec = transfer_spec
351
- assert isinstance(src_spec, GPULoadStoreSpec)
352
-
353
- for block_id in src_spec.block_ids:
354
- self.flushed_gpu_blocks.add(self.gpu_blocks[block_id.item()])
351
+ if isinstance(src_spec, GPULoadStoreSpec):
352
+ # store flush
353
+ for block_id in src_spec.block_ids:
354
+ self.flushed_gpu_blocks.add(self.gpu_blocks[block_id.item()])
355
+ else:
356
+ # load flush
357
+ for block_id in dst_spec.block_ids:
358
+ self.flushed_gpu_blocks.add(self.gpu_blocks[block_id.item()])
359
360
block_size_factor = self.block_size_factor
361
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+8
@@ -782,6 +782,14 @@ class OffloadingConnectorScheduler:
782
assert self._jobs[any_jid].is_store
783
self._current_batch_jobs_to_flush.update(req_status.transfer_jobs)
784
785
+ # If all tracked requests are finished, flush all pending jobs
786
+ # (both store and load) - there might not be a future scheduler
787
+ # step to trigger their completion.
788
+ if self._req_status and all(
789
+ rs.req.is_finished() for rs in self._req_status.values()
790
+ ):
791
+ self._current_batch_jobs_to_flush.update(self._jobs.keys())
792
+
793
meta = OffloadingConnectorMetadata(
794
load_jobs=self._current_batch_load_jobs,
795
store_jobs=self._build_store_jobs(scheduler_output),