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

Fork 0

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

[KV Connector][Offloading] Avoid blocking the engine to flush offloads on idle (#45595)

Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Signed-off-by: Or Ozeri <oro@il.ibm.com> Signed-off-by: Itay Etelis <Itay.etelis@gmail.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: Itay Etelis <Itay.etelis@gmail.com>

Itay Etelis committed Jun 17, 2026 at 11:35 UTC 6d8fff5698aea7225792b37c3b075cc64c4a0178

6 files changed +49 -81

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

+12 -73

@@ -262,12 +262,9 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:

262 runner.manager.prepare_store.side_effect = lambda keys, req_context: (

263 generate_store_output(keys)

264 )

265 - # With sync scheduling, all-finished flush fires within this run.

266 - # With async scheduling, the finish is delayed so flush fires later.

265 runner.run(

266 decoded_tokens=[EOS_TOKEN_ID],

267 expected_stored=(0, 1, 2),

270 - expected_flushed=(0, 1, 2) if not async_scheduling else (),

268 )

269

270 # start a request to load the first block, but don't complete

@@ -332,7 +329,6 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):

329 runner.run(

330 decoded_tokens=[EOS_TOKEN_ID],

331 expected_stored=(0, 1, 2),

335 - 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

@@ -359,7 +355,6 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):

355 runner.run(

356 decoded_tokens=[],

357 expected_loaded=(0, 1, 2),

362 - expected_flushed=(0, 1, 2),

358 )

359

360 # assert request is deleted

@@ -774,7 +769,6 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling

769 runner.run(

770 decoded_tokens=[EOS_TOKEN_ID],

771 expected_stored=(0, 1, 2),

777 - expected_flushed=(0, 1, 2) if not async_scheduling else (),

772 )

773

774 # Reset GPU prefix cache so the next request must load from CPU.

@@ -841,13 +835,8 @@ def test_fence_at_update_state_after_alloc(request_runner):

835 runner.manager.prepare_store.side_effect = lambda keys, req_context: (

836 generate_store_output(keys)

837 )

844 - runner.run(

845 - decoded_tokens=[EOS_TOKEN_ID],

846 - complete_transfers=False,

847 - expected_stored=(0,),

848 - expected_flushed=(0,),

849 - )

850 - assert runner.connector_scheduler._block_id_to_pending_jobs == {}

838 + runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)

839 + assert runner.connector_scheduler._block_id_to_pending_jobs

840

841 runner.scheduler.reset_prefix_cache()

842 runner.new_request(token_ids=[0] * 4)

@@ -858,6 +847,8 @@ def test_fence_at_update_state_after_alloc(request_runner):

847 runner.run(

848 decoded_tokens=[],

849 complete_transfers=False,

850 + expected_stored=(0,),

851 + expected_flushed=(0,),

852 )

853 assert runner.connector_scheduler._block_id_to_pending_jobs == {}

854

@@ -877,13 +868,8 @@ def test_fence_at_build_store_jobs(request_runner):

868 runner.manager.prepare_store.side_effect = lambda keys, req_context: (

869 generate_store_output(keys)

870 )

880 - runner.run(

881 - decoded_tokens=[EOS_TOKEN_ID],

882 - complete_transfers=False,

883 - expected_stored=(0,),

884 - expected_flushed=(0,),

885 - )

886 - assert runner.connector_scheduler._block_id_to_pending_jobs == {}

871 + runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)

872 + assert runner.connector_scheduler._block_id_to_pending_jobs

873

874 runner.scheduler.reset_prefix_cache()

875 runner.new_request(token_ids=[1] * 4)

@@ -893,6 +879,8 @@ def test_fence_at_build_store_jobs(request_runner):

879 )

880 runner.run(

881 decoded_tokens=[EOS_TOKEN_ID],

882 + expected_stored=(0,),

883 + expected_flushed=(0,),

884 )

885 assert runner.connector_scheduler._block_id_to_pending_jobs == {}

886

@@ -966,10 +954,10 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):

954 lambda keys, req_context: generate_store_output(keys)

955 )

956

969 - # With sync scheduling, the connector flushes completed stores when the

970 - # request finishes; async scheduling defers the flush to the next step.

971 - flushed_all = all_offsets if not async_scheduling else ()

972 - flushed_two = (0, 1, 2, 3, 4, 5) if not async_scheduling else ()

957 + # Pending offloads drain via non-blocking stepping, not a flush, so no

958 + # blocks are flushed when the request finishes.

959 + flushed_all: tuple[int, ...] = ()

960 + flushed_two: tuple[int, ...] = ()

961

962 # None -> no cap, all 9 offsets stored

963 r = make_runner()

@@ -1081,32 +1069,6 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool):

1069 assert len(offered_keys) == num_prompt_blocks

1070

1071

1084 -def test_flush_all_jobs_when_no_requests_remain(request_runner):

1085 - """When all tracked requests are finished, build_connector_meta flushes

1086 - all pending jobs since there will be no future step to complete them."""

1087 - block_size = 4

1088 - block_size_factor = 1

1089 - offloaded_block_size = block_size * block_size_factor

1090 -

1091 - runner = request_runner(

1092 - block_size=block_size,

1093 - num_gpu_blocks=100,

1094 - async_scheduling=False,

1095 - block_size_factor=block_size_factor,

1096 - )

1097 -

1098 - runner.new_request(token_ids=[0] * offloaded_block_size)

1099 - runner.manager.prepare_store.side_effect = lambda keys, req_context: (

1100 - generate_store_output(keys)

1101 - )

1102 - runner.run(

1103 - decoded_tokens=[EOS_TOKEN_ID],

1104 - complete_transfers=False,

1105 - expected_stored=(0,),

1106 - expected_flushed=(0,),

1107 - )

1108 -

1109 -

1072 @pytest.mark.parametrize("async_scheduling", [True, False])

1073 def test_reset_cache(request_runner, async_scheduling: bool):

1074 """reset_cache flushes in-flight loads, calls manager.reset_cache(), resets

@@ -1131,7 +1093,6 @@ def test_reset_cache(request_runner, async_scheduling: bool):

1093 runner.run(

1094 decoded_tokens=[EOS_TOKEN_ID],

1095 expected_stored=(0, 1, 2),

1134 - expected_flushed=(0, 1, 2) if not async_scheduling else (),

1096 )

1097

1098 # Reset GPU prefix cache then start a request that loads from CPU.

@@ -1381,7 +1342,6 @@ def test_stale_sliding_window_block_after_prepare_store_failure(

1342 runner.run(

1343 decoded_tokens=[EOS_TOKEN_ID],

1344 expected_stored=(2, 3),

1384 - expected_flushed=(2, 3) if not async_scheduling else (),

1345 )

1346

1347

@@ -1409,7 +1369,6 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool):

1369 runner.run(

1370 decoded_tokens=[EOS_TOKEN_ID],

1371 expected_stored=(0, 1, 2),

1412 - expected_flushed=(0, 1, 2) if not async_scheduling else (),

1372 )

1373

1374 # Reset GPU prefix cache so the next request cannot hit locally.

@@ -1429,7 +1388,6 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool):

1388 decoded_tokens=[EOS_TOKEN_ID],

1389 expected_loaded=(), # no CPU loads must happen

1390 expected_stored=(0, 1, 2), # tokens still offloaded to CPU

1432 - expected_flushed=(0, 1, 2) if not async_scheduling else (),

1391 )

1392

1393 # The external lookup must have been completely skipped.

@@ -1936,15 +1894,6 @@ class TestEagle:

1894 (1, 0),

1895 (1, 1),

1896 ),

1939 - expected_flushed=(

1940 - (0, 0),

1941 - (0, 1),

1942 - (0, 2),

1943 - (1, 0),

1944 - (1, 1),

1945 - )

1946 - if not async_scheduling

1947 - else (),

1897 )

1898

1899 @pytest.mark.parametrize("async_scheduling", [True, False])

@@ -1990,7 +1939,6 @@ class TestEagle:

1939 runner.run(

1940 decoded_tokens=[EOS_TOKEN_ID],

1941 expected_stored=((0, 0), (0, 1)),

1993 - expected_flushed=((0, 0), (0, 1)) if not async_scheduling else (),

1942 )

1943

1944 @pytest.mark.parametrize("async_scheduling", [True, False])

@@ -2087,15 +2035,6 @@ class TestEagle:

2035 (1, 0),

2036 (1, 1),

2037 ),

2090 - expected_flushed=(

2091 - (0, 0),

2092 - (0, 1),

2093 - (0, 2),

2094 - (1, 0),

2095 - (1, 1),

2096 - )

2097 - if not async_scheduling

2098 - else (),

2038 )

2039

2040 runner.scheduler.reset_prefix_cache()

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+8 -8

@@ -1016,14 +1016,6 @@ class OffloadingConnectorScheduler:

1016 for jid in self._block_id_to_pending_jobs[bid]

1017 )

1018

1019 - # If all tracked requests are finished, flush all pending jobs

1020 - # (both store and load) - there might not be a future scheduler

1021 - # step to trigger their completion.

1022 - if self._req_status and all(

1023 - rs.req.is_finished() for rs in self._req_status.values()

1024 - ):

1025 - self._current_batch_jobs_to_flush.update(self._jobs.keys())

1026 -

1019 meta = OffloadingConnectorMetadata(

1020 load_jobs=self._current_batch_load_jobs,

1021 store_jobs=self._build_store_jobs(scheduler_output),

@@ -1034,6 +1026,14 @@ class OffloadingConnectorScheduler:

1026 self._current_batch_allocated_block_ids = set()

1027 return meta

1028

1029 + def has_pending_push_work(self) -> bool:

1030 + """Whether the engine must keep stepping.

1031 +

1032 + While True, build_connector_meta() and update_connector_output()

1033 + continue to be called even when no requests are scheduled.

1034 + """

1035 + return bool(self._jobs) or self.manager.has_pending_work()

1036 +

1037 def update_connector_output(self, connector_output: KVConnectorOutput):

1038 """

1039 Update KVConnector state from worker-side connectors output.

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

@@ -150,6 +150,10 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):

150 assert self.connector_scheduler is not None

151 return self.connector_scheduler.build_connector_meta(scheduler_output)

152

153 + def has_pending_push_work(self) -> bool:

154 + assert self.connector_scheduler is not None

155 + return self.connector_scheduler.has_pending_push_work()

156 +

157 def update_connector_output(self, connector_output: KVConnectorOutput):

158 assert self.connector_scheduler is not None

159 self.connector_scheduler.update_connector_output(connector_output)

vllm/v1/kv_offload/base.py

@@ -288,6 +288,14 @@ class OffloadingManager(ABC):

288 """

289 return

290

291 + def has_pending_work(self) -> bool:

292 + """Whether this manager needs the engine to keep stepping.

293 +

294 + While True, on_schedule_end() and get_finished_jobs() continue

295 + to be called even when no requests are scheduled.

296 + """

297 + return False

298 +

299 def reset_cache(self) -> None:

300 """Evict all tracked blocks and reset internal state."""

301 return

vllm/v1/kv_offload/tiering/base.py

@@ -153,6 +153,14 @@ class SecondaryTierManager(ABC):

153 """

154 pass

155

156 + def has_pending_work(self) -> bool:

157 + """Whether this tier needs the engine to keep stepping.

158 +

159 + While True, on_schedule_end() and get_finished_jobs() continue

160 + to be called even when no requests are scheduled.

161 + """

162 + return False

163 +

164 def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):

165 """

166 Mark blocks as recently used for eviction policy.

vllm/v1/kv_offload/tiering/manager.py

@@ -577,6 +577,15 @@ class TieringOffloadingManager(OffloadingManager):

577 for tier in self.secondary_tiers:

578 tier.on_schedule_end()

579

580 + @override

581 + def has_pending_work(self) -> bool:

582 + # In-flight primary<->secondary transfers (pending promotions are

583 + # translated to transfer jobs in on_schedule_end), plus any work the

584 + # secondary tiers themselves still have outstanding.

585 + return bool(self._transfer_jobs) or any(

586 + tier.has_pending_work() for tier in self.secondary_tiers

587 + )

588 +

589 @override

590 def take_events(self) -> Iterable[OffloadingEvent]:

591 """Yield offloading events collected since the last call.