@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2aac565cae

[Core][KV Connector] Start async KV loads after the forward launch when no sync loads are scheduled (#53333)

Signed-off-by: Summer Yang <girasoleyang@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Summer Yang committed Aug 28, 2026 at 14:58 UTC 2aac565cae880087d752e90f1a08dcd9b369f9a0
17 files changed +293 -118
tests/v1/core/test_scheduler.py
+41
@@ -1975,6 +1975,47 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):
1975 return initial_ecos
1976
1977
1978 +@pytest.mark.parametrize(
1979 + ("load_modes", "expected_has_sync_loads"),
1980 + [
1981 + ((True,), False),
1982 + ((False,), True),
1983 + ((True, False), True),
1984 + ],
1985 +)
1986 +@pytest.mark.skip_global_cleanup
1987 +def test_has_sync_kv_loads(
1988 + load_modes: tuple[bool, ...],
1989 + expected_has_sync_loads: bool,
1990 + tmp_path,
1991 +):
1992 + (tmp_path / "config.json").write_text(
1993 + '{"architectures": ["OPTForCausalLM"], "model_type": "opt"}'
1994 + )
1995 + block_size = 16
1996 + scheduler = create_scheduler(
1997 + model=str(tmp_path),
1998 + skip_tokenizer_init=True,
1999 + use_kv_connector=mock_kv(matched_tokens=block_size, is_async=False),
2000 + block_size=block_size,
2001 + )
2002 + requests = create_requests(
2003 + num_requests=len(load_modes),
2004 + num_tokens=block_size * 2,
2005 + block_size=block_size,
2006 + )
2007 + for request in requests:
2008 + scheduler.add_request(request)
2009 +
2010 + scheduler.connector.get_num_new_matched_tokens = Mock(
2011 + side_effect=[(block_size, is_async) for is_async in load_modes]
2012 + )
2013 +
2014 + output = scheduler.schedule()
2015 +
2016 + assert output.has_sync_kv_loads is expected_has_sync_loads
2017 +
2018 +
2019 @pytest.mark.parametrize("is_async", [False, True])
2020 def test_kv_connector_basic(is_async: bool):
2021 """
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+1 -1
@@ -2960,7 +2960,7 @@ def test_putting_consumer_queues_decode_save():
2960
2961 with patch.object(torch.cuda, "Event") as event_cls:
2962 event = event_cls.return_value
2963 - w.get_finished(set(), meta)
2963 + w.wait_for_save(meta)
2964
2965 event.record.assert_called_once_with()
2966 send_thread.add_request.assert_called_once_with(req)
tests/v1/kv_connector/unit/test_multi_connector.py
+9 -15
@@ -294,18 +294,23 @@ def test_multi_example_connector_consistency():
294 "update_state_after_alloc num_blocks=[7] 0",
295 "build_connector_meta",
296 ]
297 - # First three events are from initialization (register_kv_caches,
298 - # set_host_xfer_buffer_ops, get_handshake_metadata), then generate() events.
299 - assert events["storage1-WORKER"][:8] == [
297 + # First three events are from initialization. During generate(), layer hooks
298 + # run before the deferred load is started after the forward pass.
299 + expected_worker_prefix = [
300 "register_kv_caches",
301 "set_host_xfer_buffer_ops",
302 "get_handshake_metadata",
303 "handle_preemptions",
304 "bind_connector_metadata",
305 - "start_load_kv",
305 "wait_for_layer_load",
306 "save_kv_layer",
307 ]
308 + for connector_name in ("storage1-WORKER", "storage2-WORKER"):
309 + worker_events = events[connector_name]
310 + assert worker_events[:7] == expected_worker_prefix
311 + assert worker_events.index("start_load_kv") > worker_events.index(
312 + "save_kv_layer"
313 + )
314 assert storage2_scheduler_events[:7] == [
315 "bind_gpu_block_pool",
316 "get_finished_count",
@@ -315,17 +320,6 @@ def test_multi_example_connector_consistency():
320 "update_state_after_alloc num_blocks=[7] 0",
321 "build_connector_meta",
322 ]
318 - assert events["storage2-WORKER"][:8] == [
319 - "register_kv_caches",
320 - "set_host_xfer_buffer_ops",
321 - "get_handshake_metadata",
322 - "handle_preemptions",
323 - "bind_connector_metadata",
324 - "start_load_kv",
325 - "wait_for_layer_load",
326 - "save_kv_layer",
327 - ]
328 -
323 # Reset prefix cache or else we'll just get the tokens back from there.
324 llm.reset_prefix_cache()
325
tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py
+4 -1
@@ -213,7 +213,7 @@ def _multi_connector_config(swa_enabled: bool = False):
213 def test_nixl_wins_load_over_cpu_offload():
214 """When NixlConnector (index 0) has matched tokens from a remote prefill, it should
215 win the load: Nixl metadata tracks the recv while CPU offload metadata has no load
216 - scheduled."""
216 + scheduled."""
217 vllm_config, kv_cache_config = _multi_connector_config()
218 scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config)
219 mc = scheduler.connector
@@ -227,6 +227,9 @@ def test_nixl_wins_load_over_cpu_offload():
227 )
228 scheduler.add_request(request)
229 sched_out = scheduler.schedule()
230 + # The remote-prefill load is async, so the step has no sync KV loads
231 + # and the worker will start the load after the forward.
232 + assert not sched_out.has_sync_kv_loads
233
234 assert mc._requests_to_connector[request.request_id] == 0
235
tests/v1/simple_kv_offload/test_worker.py
+4 -3
@@ -152,8 +152,8 @@ class _RecordingBackend:
152 self.calls.append({"is_store": is_store, "wait_event": wait_event})
153
154
155 -def test_get_finished_passes_wait_event_for_store_only():
156 - """get_finished gates stores on a compute-done event but not loads."""
155 +def test_transfer_hooks_pass_wait_event_for_store_only():
156 + """wait_for_save gates stores on a compute-done event; start_load_kv does not."""
157 worker = SimpleCPUOffloadWorker(
158 vllm_config=None, kv_cache_config=None, cpu_capacity_bytes=0
159 )
@@ -168,7 +168,8 @@ def test_get_finished_passes_wait_event_for_store_only():
168 store_cpu_blocks=[1],
169 )
170
171 - worker.get_finished(set())
171 + worker.start_load_kv()
172 + worker.wait_for_save()
173
174 store_calls = [c for c in recording.calls if c["is_store"]]
175 load_calls = [c for c in recording.calls if not c["is_store"]]
tests/v1/worker/test_gpu_kv_connector.py new
+79
@@ -0,0 +1,79 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +from unittest.mock import Mock
6 +
7 +import pytest
8 +
9 +import vllm.v1.worker.gpu.kv_connector as kv_connector_module
10 +from vllm.config import KVTransferConfig
11 +from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
12 +
13 +
14 +def _make_connector(
15 + monkeypatch: pytest.MonkeyPatch,
16 + events: list[str],
17 +) -> ActiveKVConnector:
18 + backend = Mock()
19 + backend.handle_preemptions.side_effect = lambda _: events.append("handle")
20 + backend.bind_connector_metadata.side_effect = lambda _: events.append("bind")
21 + backend.start_load_kv.side_effect = lambda _: events.append("start")
22 + backend.wait_for_save.side_effect = lambda: events.append("wait")
23 + backend.get_finished.side_effect = lambda _: (set(), set())
24 + backend.get_block_ids_with_load_errors.return_value = set()
25 + backend.get_kv_connector_stats.return_value = None
26 + backend.get_kv_connector_kv_cache_events.return_value = None
27 + backend.build_connector_worker_meta.return_value = None
28 + backend.clear_connector_metadata.side_effect = lambda: events.append("clear")
29 + monkeypatch.setattr(kv_connector_module, "get_kv_transfer_group", lambda: backend)
30 + monkeypatch.setattr(
31 + kv_connector_module, "is_forward_context_available", lambda: True
32 + )
33 + monkeypatch.setattr(kv_connector_module, "get_forward_context", object)
34 +
35 + kv_config = KVTransferConfig(
36 + kv_connector="NixlConnector",
37 + kv_role="kv_consumer",
38 + kv_buffer_device="cpu",
39 + )
40 + connector = ActiveKVConnector( # type: ignore[arg-type]
41 + SimpleNamespace(kv_transfer_config=kv_config), {}
42 + )
43 + events.clear()
44 + return connector
45 +
46 +
47 +def _scheduler_output(has_sync_kv_loads: bool) -> SimpleNamespace:
48 + return SimpleNamespace(
49 + kv_connector_metadata=object(),
50 + finished_req_ids=set(),
51 + has_sync_kv_loads=has_sync_kv_loads,
52 + )
53 +
54 +
55 +@pytest.mark.parametrize("has_sync_kv_loads", [False, True])
56 +def test_load_start_phase(
57 + monkeypatch: pytest.MonkeyPatch,
58 + has_sync_kv_loads: bool,
59 +):
60 + events: list[str] = []
61 + connector = _make_connector(monkeypatch, events)
62 + output = _scheduler_output(has_sync_kv_loads)
63 +
64 + connector.pre_forward(output) # type: ignore[arg-type]
65 + assert events == (
66 + ["handle", "bind", "start"] if has_sync_kv_loads else ["handle", "bind"]
67 + )
68 +
69 + connector.post_forward(set())
70 + assert events == ["handle", "bind", "start", "wait", "clear"]
71 +
72 +
73 +def test_no_forward_starts_deferred_load_once(monkeypatch: pytest.MonkeyPatch):
74 + events: list[str] = []
75 + connector = _make_connector(monkeypatch, events)
76 +
77 + connector.no_forward(_scheduler_output(False)) # type: ignore[arg-type]
78 +
79 + assert events == ["handle", "bind", "start", "clear"]
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+2 -2
@@ -289,8 +289,8 @@ class KVConnectorBase_V1(ABC):
289 def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
290 """
291 Start loading the KV cache from the connector to vLLM's paged
292 - KV buffer. This is called from the forward context before the
293 - forward pass to enable async loading during model execution.
292 + KV buffer. Loads required by the current forward start before it;
293 + independent asynchronous loads may start after it is submitted.
294
295 Args:
296 forward_context (ForwardContext): the forward context.
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py
+15
@@ -133,6 +133,21 @@ class LMCacheConnectorV1(KVConnectorBase_V1):
133 "please check and use the latest version"
134 )
135
136 + def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None:
137 + super().bind_connector_metadata(connector_metadata)
138 + engine = self._lmcache_engine
139 + if hasattr(engine, "bind_connector_metadata"):
140 + engine.bind_connector_metadata(connector_metadata)
141 + else:
142 + # Compat for engines predating the bind hook: start_load_kv may
143 + # run after the forward (SchedulerOutput.has_sync_kv_loads), so
144 + # the per-step layerwise state it used to reset must be reset
145 + # here. Remove once LMCache ships bind_connector_metadata.
146 + if hasattr(engine, "current_layer"):
147 + engine.current_layer = 0
148 + if hasattr(engine, "layerwise_retrievers"):
149 + engine.layerwise_retrievers = []
150 +
151 def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
152 """
153 Start loading the KV cache from the connector to vLLM's paged
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py
+11 -4
@@ -783,6 +783,17 @@ class LMCacheConnectorV1Impl:
783 ####################
784 # Worker side APIs
785 ####################
786 + def bind_connector_metadata(self, metadata: "KVConnectorMetadata") -> None:
787 + """Per-step init, called when the runner binds this step's metadata.
788 +
789 + The layerwise hooks fire during every forward once metadata is
790 + bound, while start_load_kv may run after the forward launch on
791 + steps without sync loads (SchedulerOutput.has_sync_kv_loads), so
792 + the per-step state they consume must be reset here.
793 + """
794 + self.current_layer = 0
795 + self.layerwise_retrievers = []
796 +
797 @_lmcache_nvtx_annotate
798 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
799 logger.info("Registering KV caches")
@@ -806,8 +817,6 @@ class LMCacheConnectorV1Impl:
817 The number of elements in kv_caches and layer_names should be
818 the same.
819 """
809 - self.current_layer = 0
810 -
820 if len(self.kv_caches) == 0:
821 self._init_kv_caches_from_forward_context(forward_context)
822
@@ -826,8 +835,6 @@ class LMCacheConnectorV1Impl:
835
836 self.lmcache_engine.post_init(kvcaches=kvcaches)
837
829 - self.layerwise_retrievers = []
830 -
838 for idx, request in enumerate(metadata.requests):
839 if request.load_spec is None:
840 continue
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
+8 -4
@@ -283,8 +283,10 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
283 self.connector_worker.register_kv_caches(kv_caches)
284
285 def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
286 - # No-op: loads are issued in get_finished() for compute overlap.
287 - pass
286 + assert self.connector_worker is not None
287 + metadata = self._get_connector_metadata()
288 + assert isinstance(metadata, MooncakeStoreConnectorMetadata)
289 + self.connector_worker.start_load_kv(metadata)
290
291 def wait_for_layer_load(self, layer_name: str) -> None:
292 # No layerwise support - no-op
@@ -301,8 +303,10 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
303 return
304
305 def wait_for_save(self):
304 - # No-op: stores are issued in get_finished() for compute overlap.
305 - pass
306 + assert self.connector_worker is not None
307 + metadata = self._get_connector_metadata()
308 + assert isinstance(metadata, MooncakeStoreConnectorMetadata)
309 + self.connector_worker.wait_for_save(metadata)
310
311 def get_finished(
312 self, finished_req_ids: set[str]
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+40 -40
@@ -1743,36 +1743,17 @@ class MooncakeStoreWorker:
1743 "Started %d Mooncake KV-load receive thread(s)", self.num_recv_threads
1744 )
1745
1746 - def start_load_kv(
1747 - self,
1748 - metadata: MooncakeStoreConnectorMetadata,
1749 - ):
1750 - """No-op: loads are issued in get_finished() for overlap."""
1751 - pass
1746 + def start_load_kv(self, metadata: MooncakeStoreConnectorMetadata):
1747 + """Issue async loads.
1748
1753 - def wait_for_save(
1754 - self,
1755 - metadata: MooncakeStoreConnectorMetadata,
1756 - ):
1757 - """No-op: stores are issued in get_finished() for overlap."""
1758 - pass
1759 -
1760 - def get_finished(
1761 - self,
1762 - finished_req_ids: set[str],
1763 - meta: MooncakeStoreConnectorMetadata,
1764 - ) -> tuple[set[str], set[str]]:
1765 - """Issue all I/O and get completed send/recv request IDs.
1766 -
1767 - All load and store I/O requests are issued here (after model
1768 - compute is launched on the compute stream) for better
1769 - compute-I/O overlap.
1749 + Runs after the forward launch on steps without sync loads
1750 + (SchedulerOutput.has_sync_kv_loads), keeping load submission off
1751 + the critical path while preserving compute-I/O overlap.
1752 """
1753 if self._capacity_only:
1772 - return set(), set()
1754 + return
1755
1774 - # Issue async loads
1775 - for request in meta.requests:
1756 + for request in metadata.requests:
1757 load_spec = request.load_spec
1758 if load_spec is None or not load_spec.can_load:
1759 continue
@@ -1781,21 +1762,40 @@ class MooncakeStoreWorker:
1762 self.recv_request_queue.put(request)
1763
1764 assert self.load_async, "load_async must be True for better performance."
1784 - # Issue stores with CUDA event synchronization.
1785 - if self.can_put:
1786 - current_event = None
1787 - for request in meta.requests:
1788 - if request.can_save:
1789 - current_event = torch.cuda.Event()
1790 - current_event.record()
1791 - break
1765
1793 - for request in meta.requests:
1794 - if not request.can_save:
1795 - continue
1796 - request.current_event = current_event
1797 - assert self.kv_send_thread is not None
1798 - self.kv_send_thread.add_request(request)
1766 + def wait_for_save(self, metadata: MooncakeStoreConnectorMetadata):
1767 + """Issue async stores with CUDA event synchronization.
1768 +
1769 + Runs after the forward launch for compute-I/O overlap.
1770 + """
1771 + if self._capacity_only or not self.can_put:
1772 + return
1773 +
1774 + current_event = None
1775 + for request in metadata.requests:
1776 + if request.can_save:
1777 + current_event = torch.cuda.Event()
1778 + current_event.record()
1779 + break
1780 +
1781 + for request in metadata.requests:
1782 + if not request.can_save:
1783 + continue
1784 + request.current_event = current_event
1785 + assert self.kv_send_thread is not None
1786 + self.kv_send_thread.add_request(request)
1787 +
1788 + def get_finished(
1789 + self, finished_req_ids: set[str], meta: MooncakeStoreConnectorMetadata
1790 + ) -> tuple[set[str], set[str]]:
1791 + """Get completed send/recv request IDs.
1792 +
1793 + Loads are issued in start_load_kv() and stores in wait_for_save().
1794 + """
1795 + if self._capacity_only:
1796 + return set(), set()
1797 +
1798 + if self.can_put:
1799 self._close_ended_store_requests(finished_req_ids, meta)
1800
1801 # Blocks read by a store job are released by the scheduler when the job
vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
+6 -4
@@ -196,10 +196,11 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
196 self.worker_handler.handle_preemptions(kv_connector_metadata)
197
198 def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
199 - pass # Launch loads ops in get_finished() after launching model execution
199 + if self.worker_handler is not None:
200 + self.worker_handler.start_load_kv()
201
202 def wait_for_layer_load(self, layer_name: str) -> None:
202 - pass # Always load asynchronously and deferred to get_finished()
203 + pass # Always load asynchronously, issued in start_load_kv()
204
205 def save_kv_layer(
206 self,
@@ -208,10 +209,11 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
209 attn_metadata: "AttentionMetadata",
210 **kwargs: Any,
211 ) -> None:
211 - pass # Always save asynchronously and deferred to get_finished()
212 + pass # Always save asynchronously, issued in wait_for_save()
213
214 def wait_for_save(self) -> None:
214 - pass # All stores are driven by get_finished() and no wait needed
215 + if self.worker_handler is not None:
216 + self.worker_handler.wait_for_save()
217
218 def get_finished(
219 self,
vllm/v1/core/sched/output.py
+4
@@ -262,6 +262,10 @@ class SchedulerOutput:
262 # KV Cache Connector metadata.
263 kv_connector_metadata: KVConnectorMetadata | None = None
264
265 + # Whether any scheduled request consumes KV that the connector loads
266 + # synchronously during this step (load_async=False).
267 + has_sync_kv_loads: bool = False
268 +
269 # EC Cache Connector metadata
270 ec_connector_metadata: ECConnectorMetadata | None = None
271 # EC Cache Manager metadata
vllm/v1/core/sched/scheduler.py
+6
@@ -532,6 +532,8 @@ class Scheduler(SchedulerInterface):
532 scheduled_spec_decode_tokens: dict[str, list[int]] = {}
533 # Whether the running batch contains any prefill requests.
534 prefill_scheduled = False
535 + # Whether any scheduled request has a synchronous connector KV load.
536 + has_sync_kv_loads = False
537
538 # For logging.
539 scheduled_timestamp = time.monotonic()
@@ -1142,6 +1144,9 @@ class Scheduler(SchedulerInterface):
1144 continue
1145
1146 self.running.append(request)
1147 + if num_external_computed_tokens > 0:
1148 + # load_kv_async is False here
1149 + has_sync_kv_loads = True
1150 if self.log_stats:
1151 request.record_event(
1152 EngineCoreEventType.SCHEDULED, scheduled_timestamp
@@ -1318,6 +1323,7 @@ class Scheduler(SchedulerInterface):
1323 finished_req_ids=self.finished_req_ids,
1324 free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
1325 new_block_ids_to_zero=self._get_new_block_ids_to_zero(),
1326 + has_sync_kv_loads=has_sync_kv_loads,
1327 kv_cache_block_copies=pending_kv_cache_block_copies,
1328 partial_tail_offloads=pending_partial_tail_offloads,
1329 num_spec_tokens_to_schedule=num_spec_tokens_to_schedule,
vllm/v1/simple_kv_offload/worker.py
+40 -39
@@ -231,58 +231,59 @@ class SimpleCPUOffloadWorker:
231 self._connector_metadata = None
232
233 def start_load_kv(self) -> None:
234 - # NOTE: we defer launching both load and store to get_finished(),
235 - # which runs after model execution. This hides the CPU-side
236 - # block copy op overhead (~5ms) behind GPU compute.
237 - pass
234 + # Invoked after the forward launch on steps without sync loads
235 + # (SchedulerOutput.has_sync_kv_loads), so the CPU-side block copy
236 + # op overhead (~5ms) stays hidden behind GPU compute. Stores are
237 + # issued in wait_for_save().
238 + metadata = self._connector_metadata
239 + if metadata is not None and metadata.load_cpu_blocks:
240 + backend = self._backend
241 + assert backend is not None
242 + backend.launch_copy(
243 + metadata.load_cpu_blocks,
244 + metadata.load_gpu_blocks,
245 + is_store=False,
246 + event_idx=metadata.load_event,
247 + events_list=self._load_events,
248 + )
249
250 def wait_for_save(self) -> None:
240 - pass
241 -
242 - def get_finished(
243 - self,
244 - finished_req_ids: set[str],
245 - ) -> tuple[set[str] | None, set[str] | None]:
246 - """Submit transfers and report completed events to the scheduler.
251 + """Submit async stores.
252
253 Stores (GPU->CPU) read the live KV cache, which the compute stream may
254 still be writing under v1 overlapped execution, so they are ordered
250 - after a compute-done event recorded on the current stream. Loads
251 - (CPU->GPU) read stable pinned host memory and launch immediately. See
255 + after a compute-done event recorded on the current stream. See
256 #45704 for the bug and #39306 for the srcAccessOrder rationale.
257 + """
258 + metadata = self._connector_metadata
259 + if metadata is not None and metadata.store_gpu_blocks:
260 + backend = self._backend
261 + assert backend is not None
262 + if self._store_compute_done is None:
263 + self._store_compute_done = torch.Event()
264 + self._store_compute_done.record(torch.cuda.current_stream())
265 + backend.launch_copy(
266 + metadata.store_gpu_blocks,
267 + metadata.store_cpu_blocks,
268 + is_store=True,
269 + event_idx=metadata.store_event,
270 + events_list=self._store_events,
271 + wait_event=self._store_compute_done,
272 + )
273 +
274 + def get_finished(
275 + self, finished_req_ids: set[str]
276 + ) -> tuple[set[str] | None, set[str] | None]:
277 + """Report completed transfer events to the scheduler.
278 +
279 + Loads are submitted in start_load_kv() and stores in wait_for_save().
280
281 Returns:
282 tuple of (finished_sending, finished_recving).
283 - finished_sending: always None (stores use worker metadata).
284 - finished_recving: req_ids whose loads have completed.
285 """
259 - # (1) Submit transfers
286 metadata = self._connector_metadata
261 - if metadata is not None:
262 - backend = self._backend
263 - assert backend is not None
264 - if metadata.load_cpu_blocks:
265 - backend.launch_copy(
266 - metadata.load_cpu_blocks,
267 - metadata.load_gpu_blocks,
268 - is_store=False,
269 - event_idx=metadata.load_event,
270 - events_list=self._load_events,
271 - )
272 - if metadata.store_gpu_blocks:
273 - if self._store_compute_done is None:
274 - self._store_compute_done = torch.Event()
275 - self._store_compute_done.record(torch.cuda.current_stream())
276 - backend.launch_copy(
277 - metadata.store_gpu_blocks,
278 - metadata.store_cpu_blocks,
279 - is_store=True,
280 - event_idx=metadata.store_event,
281 - events_list=self._store_events,
282 - wait_event=self._store_compute_done,
283 - )
284 -
285 - # (2) Track completed transfer events
287 finished_recving: set[str] = set()
288
289 if self._pending_load_event_indices:
vllm/v1/worker/gpu/kv_connector.py
+14
@@ -54,6 +54,7 @@ class ActiveKVConnector(KVConnector):
54 self.kv_connector.register_kv_caches(kv_caches_dict)
55 self.kv_connector.set_host_xfer_buffer_ops(copy_kv_blocks)
56
57 + self._pending_load_start = False
58 self._disabled = False
59
60 def pre_forward(self, scheduler_output: "SchedulerOutput") -> None:
@@ -65,6 +66,16 @@ class ActiveKVConnector(KVConnector):
66 self.kv_connector.handle_preemptions(kv_connector_metadata)
67 self.kv_connector.bind_connector_metadata(kv_connector_metadata)
68
69 + if scheduler_output.has_sync_kv_loads:
70 + # Sync loads need to run before this step's forward.
71 + self._start_load_kv()
72 + else:
73 + # Start any async loads in post-forward instead, keeping
74 + # their host-side submission cost off the critical path.
75 + self._pending_load_start = True
76 +
77 + def _start_load_kv(self) -> None:
78 + self._pending_load_start = False
79 # TODO: sort out KV Connectors' use of forward_context
80 if is_forward_context_available():
81 self.kv_connector.start_load_kv(get_forward_context())
@@ -78,6 +89,9 @@ class ActiveKVConnector(KVConnector):
89 if self._disabled:
90 return None
91
92 + if self._pending_load_start:
93 + self._start_load_kv()
94 +
95 output = KVConnectorOutput()
96 if wait_for_save:
97 self.kv_connector.wait_for_save()
vllm/v1/worker/kv_connector_model_runner_mixin.py
+9 -5
@@ -79,14 +79,18 @@ class KVConnectorModelRunnerMixin:
79 assert scheduler_output.kv_connector_metadata is not None
80 kv_connector.bind_connector_metadata(scheduler_output.kv_connector_metadata)
81
82 - # Background KV cache transfers happen here.
83 - # These transfers are designed to be async and the requests
84 - # involved may be disjoint from the running requests.
85 - # Do this here to save a collective_rpc.
86 - kv_connector.start_load_kv(get_forward_context())
82 + # Start this step's KV loads, ordered after any in-flight KV block
83 + # zeroing. Sync loads feed this step's forward so must precede it;
84 + # otherwise start (async) loads after the forward launch, keeping
85 + # their host-side submission cost off the critical path.
86 + start_after_forward = not scheduler_output.has_sync_kv_loads
87 + if not start_after_forward:
88 + kv_connector.start_load_kv(get_forward_context())
89 try:
90 yield output
91 finally:
92 + if start_after_forward:
93 + kv_connector.start_load_kv(get_forward_context())
94 if wait_for_save and not defer_finalize:
95 kv_connector.wait_for_save()
96