b3a87b2d04 · tduncan/Nvidia-Optimization-Stack-VLLM

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

+1

@@ -517,6 +517,7 @@ class RequestRunner:

517 if complete_transfers:

518 self.offloading_spec.complete_transfers()

519

520 + self.worker_connector.wait_for_save()

521 finished_sending, finished_recving = self.worker_connector.get_finished(

522 scheduler_output.finished_req_ids

523 )

tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py

+1

@@ -71,6 +71,7 @@ def test_kv_connector_mixin_clears_metadata():

71 assert connector._connector_metadata is None

72 # Test connector wrapper records method calls

73 assert connector.call_record.get("bind_connector_metadata", 0) == 1

74 + assert connector.call_record.get("wait_for_save", 0) == 1

75 assert connector.call_record.get("clear_connector_metadata", 0) == 1

76 finally:

77 # Ensure we clean up the global connector between tests

tests/v1/kv_connector/unit/test_lmcache_connector.py

+31

@@ -1,10 +1,12 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from collections.abc import Iterator

4 from unittest.mock import MagicMock

5

6 import pytest

7

8 from vllm.distributed.kv_events import BlockStored

9 +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata

10 from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector import (

11 LMCacheConnectorV1,

12 LMCacheKVEvents,

@@ -783,3 +785,32 @@ class TestIntegrationScenarios:

785 assert aggregated_events[0].block_hashes == ["hash_common"]

786 assert aggregated_events[0].parent_block_hash == "parent_common"

787 assert aggregated_events[0].token_ids == [1, 2, 3]

788 +

789 +

790 +def test_legacy_layerwise_engine_does_not_reuse_storers_on_empty_steps():

791 + """Empty steps must not advance generators exhausted by the last forward."""

792 +

793 + class LegacyLayerwiseEngine:

794 + layerwise_storers: list[Iterator[None]]

795 +

796 + def save_kv_layer(self, *args, **kwargs):

797 + self.layerwise_storers.append(iter([None]))

798 +

799 + def wait_for_save(self):

800 + for storer in self.layerwise_storers:

801 + next(storer)

802 +

803 + connector = LMCacheConnectorV1.__new__(LMCacheConnectorV1)

804 + connector._lmcache_engine = LegacyLayerwiseEngine()

805 + metadata = KVConnectorMetadata()

806 + connector.bind_connector_metadata(metadata)

807 + connector.save_kv_layer("layer", MagicMock(), MagicMock())

808 + connector.wait_for_save()

809 + connector.clear_connector_metadata()

810 +

811 + for _ in range(2):

812 + connector.bind_connector_metadata(metadata)

813 + connector.wait_for_save()

814 + connector.clear_connector_metadata()

815 +

816 + assert connector._lmcache_engine.layerwise_storers == []

tests/v1/worker/test_gpu_kv_connector.py

+1 -1

@@ -96,4 +96,4 @@ def test_no_forward_starts_deferred_load_once(monkeypatch: pytest.MonkeyPatch):

96

97 connector.no_forward(_scheduler_output(False)) # type: ignore[arg-type]

98

99 - assert events == ["handle", "bind", "start", "clear"]

99 + assert events == ["handle", "bind", "start", "wait", "clear"]

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

+8 -5

@@ -31,7 +31,7 @@ The class provides the following primitives:

31 wait_for_layer_load() - blocks until layer i load is done

32

33 save_kv_layer() - starts saving KV for layer i (maybe async)

34 - wait_for_save() - blocks until all saves are done

34 + wait_for_save() - submits saves and performs required synchronization

35

36 get_transfer_results() - returns async send/receive completions and

37 receive failures in one snapshot.

@@ -356,11 +356,14 @@ class KVConnectorBase_V1(ABC):

356

357 @abstractmethod

358 def wait_for_save(self):

359 - """Block until all the save operations is done. This is called

360 - as the forward context exits to ensure that the async saving

361 - from save_kv_layer is complete before finishing the forward.

359 + """Submit saves and perform required synchronization.

360

363 - This prevents overwrites of paged KV buffer before saving done.

361 + Called once per step, including empty steps, after any target and draft

362 + forwards. Metadata remains bound during this call. Do nothing if there

363 + is no save work.

364 +

365 + Async saves must retain their source blocks until completion is reported

366 + to the scheduler. Otherwise, wait for saves to finish before returning.

367 """

368 pass

369

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

+1

@@ -142,6 +142,7 @@ class LMCacheConnectorV1(KVConnectorBase_V1):

142 engine.current_layer = 0

143 if hasattr(engine, "layerwise_retrievers"):

144 engine.layerwise_retrievers = []

145 + engine.layerwise_storers = []

146

147 def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:

148 """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

+2

@@ -606,6 +606,7 @@ class LMCacheConnectorV1Impl:

606

607 self.async_loading = config.enable_async_loading

608 self.layerwise_retrievers: list[Generator[torch.Tensor | None, None, None]] = []

609 + self.layerwise_storers: list[Generator[Any, None, None]] = []

610 self._stats_monitor = LMCStatsMonitor.GetOrCreate()

611 if role == KVConnectorRole.SCHEDULER:

612 # Create lookup client using factory

@@ -798,6 +799,7 @@ class LMCacheConnectorV1Impl:

799 """

800 self.current_layer = 0

801 self.layerwise_retrievers = []

802 + self.layerwise_storers = []

803

804 @_lmcache_nvtx_annotate

805 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):

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

+4 -9

@@ -131,19 +131,14 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):

131 pass

132

133 def wait_for_save(self):

134 - # Store deferral is handled in get_finished(), which always runs even

135 - # when wait_for_save() is skipped (e.g. kv_connector_no_forward).

136 - pass

137 -

138 - def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:

134 assert self.connector_worker is not None

135 assert isinstance(self._connector_metadata, OffloadingConnectorMetadata)

141 -

142 - # Defer store jobs to the next step's start_kv_transfers. Done here

143 - # (rather than wait_for_save) so stores are queued even on steps where

144 - # wait_for_save is skipped.

136 + # Defer store jobs to the next step's start_kv_transfers.

137 self.connector_worker.prepare_store_kv(self._connector_metadata)

138

139 + def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:

140 + assert self.connector_worker is not None

141 + assert isinstance(self._connector_metadata, OffloadingConnectorMetadata)

142 return self.connector_worker.get_finished(finished_req_ids)

143

144 def build_connector_worker_meta(self) -> OffloadingWorkerMetadata | None:

vllm/v1/simple_kv_offload/worker.py

+1 -10

@@ -67,7 +67,6 @@ class SimpleCPUOffloadWorker:

67

68 # Metadata for the current step

69 self._connector_metadata: SimpleCPUOffloadMetadata | None = None

70 - self._store_submitted = False

70

71 # Compute-done event recorded before each store; reused across steps

72 # (get_finished runs once per step, copy queue is FIFO).

@@ -246,15 +245,12 @@ class SimpleCPUOffloadWorker:

245

246 def bind_connector_metadata(self, metadata: SimpleCPUOffloadMetadata) -> None:

247 self._connector_metadata = metadata

249 - self._store_submitted = False

248 if metadata.load_event >= 0:

249 self._pending_load_event_indices.add(metadata.load_event)

250 if metadata.store_event >= 0:

251 self._pending_store_event_indices.add(metadata.store_event)

252

253 def clear_connector_metadata(self) -> None:

256 - # No-forward steps skip the normal wait_for_save hook.

257 - self.wait_for_save()

254 self._connector_metadata = None

255

256 def start_load_kv(self) -> None:

@@ -283,11 +279,7 @@ class SimpleCPUOffloadWorker:

279 #45704 for the bug and #39306 for the srcAccessOrder rationale.

280 """

281 metadata = self._connector_metadata

286 - if (

287 - metadata is not None

288 - and metadata.store_gpu_blocks

289 - and not self._store_submitted

290 - ):

282 + if metadata is not None and metadata.store_gpu_blocks:

283 backend = self._backend

284 assert backend is not None

285 if self._store_compute_done is None:

@@ -301,7 +293,6 @@ class SimpleCPUOffloadWorker:

293 events_list=self._store_events,

294 wait_event=self._store_compute_done,

295 )

304 - self._store_submitted = True

296

297 def get_finished(

298 self, finished_req_ids: set[str]

vllm/v1/worker/gpu/kv_connector.py

+4 -9

@@ -35,9 +35,7 @@ class KVConnector:

35 def finish_forward(self) -> None:

36 pass

37

38 - def post_forward(

39 - self, finished_req_ids: set[str], wait_for_save: bool = True

40 - ) -> KVConnectorOutput | None:

38 + def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None:

39 return None

40

41 def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:

@@ -97,9 +95,7 @@ class ActiveKVConnector(KVConnector):

95 def reset_capture_state(self) -> None:

96 self.kv_connector.reset_capture_state()

97

100 - def post_forward(

101 - self, finished_req_ids: set[str], wait_for_save: bool = True

102 - ) -> KVConnectorOutput | None:

98 + def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None:

99 if self._disabled:

100 return None

101

@@ -107,8 +103,7 @@ class ActiveKVConnector(KVConnector):

103 self._start_load_kv()

104

105 output = KVConnectorOutput()

110 - if wait_for_save:

111 - self.kv_connector.wait_for_save()

106 + self.kv_connector.wait_for_save()

107 transfer_results = self.kv_connector.get_transfer_results(finished_req_ids)

108 output.finished_sending = transfer_results.finished_sending or None

109 output.finished_recving = transfer_results.finished_recving or None

@@ -129,7 +124,7 @@ class ActiveKVConnector(KVConnector):

124 self.pre_forward(scheduler_output)

125 self.finish_forward()

126 finished_req_ids = scheduler_output.finished_req_ids

132 - kv_connector_output = self.post_forward(finished_req_ids, wait_for_save=False)

127 + kv_connector_output = self.post_forward(finished_req_ids)

128 return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)

129

130 def set_disabled(self, disabled: bool) -> None:

vllm/v1/worker/kv_connector_model_runner_mixin.py

+2 -3

@@ -29,7 +29,7 @@ class KVConnectorModelRunnerMixin:

29 with (

30 set_forward_context(None, vllm_config),

31 KVConnectorModelRunnerMixin._get_kv_connector_output(

32 - scheduler_output, wait_for_save=False

32 + scheduler_output

33 ) as kv_connector_output,

34 ):

35 pass

@@ -66,7 +66,6 @@ class KVConnectorModelRunnerMixin:

66 @contextmanager

67 def _get_kv_connector_output(

68 scheduler_output: "SchedulerOutput",

69 - wait_for_save: bool = True,

69 defer_finalize: bool = False,

70 ) -> Generator[KVConnectorOutput, None, None]:

71 output = KVConnectorOutput()

@@ -89,7 +88,7 @@ class KVConnectorModelRunnerMixin:

88 finally:

89 if start_after_forward:

90 kv_connector.start_load_kv(get_forward_context())

92 - if wait_for_save and not defer_finalize:

91 + if not defer_finalize:

92 kv_connector.wait_for_save()

93

94 transfer_results = kv_connector.get_transfer_results(

[Bugfix][KVConnector] Finalize saves on steps without a forward (#57775)