@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5bd8c71e79

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 3, 16:30 UTC

[kv_offload] Implement `reset_cache()` for the offloading connector (#41956)

Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Or Ozeri <or@ozery.com>

Martin Hickey committed May 14, 2026 at 14:00 UTC 5bd8c71e792a6d1632fd92fec8c4a1b79b02027e

8 files changed +153

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

+77

@@ -919,3 +919,80 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):

919 # Finish: no store pending -> no further call.

920 runner.run(decoded_tokens=[EOS_TOKEN_ID])

921 assert runner.manager.complete_store.call_count == 0

922 +

923 +

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

925 +def test_reset_cache(request_runner, async_scheduling: bool):

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

927 + next_stored_block_idx for active requests and clears job tracking."""

928 + block_size = 4

929 + block_size_factor = 3

930 + offloaded_block_size = block_size * block_size_factor

931 + num_gpu_blocks = 100

932 +

933 + runner = request_runner(

934 + block_size=block_size,

935 + num_gpu_blocks=num_gpu_blocks,

936 + async_scheduling=async_scheduling,

937 + block_size_factor=block_size_factor,

938 + )

939 +

940 + # Store 1 offloaded block (3 GPU blocks) to CPU.

941 + runner.new_request(token_ids=[0] * offloaded_block_size)

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

943 + generate_store_output(keys)

944 + )

945 + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0, 1, 2))

946 +

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

948 + # Leave the load in-flight so that reset_cache must flush it.

949 + runner.scheduler.reset_prefix_cache()

950 + runner.new_request(token_ids=[0] * offloaded_block_size)

951 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1

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

953 + generate_store_output([])

954 + )

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

956 +

957 + # Capture in-flight load job IDs before reset.

958 + load_job_ids = {

959 + jid

960 + for jid, status in runner.connector_scheduler._jobs.items()

961 + if not status.is_store

962 + }

963 + assert load_job_ids, "expected in-flight load jobs before reset"

964 +

965 + # Record job counter to verify the reset counter is set correctly.

966 + job_counter_before_reset = runner.connector_scheduler._job_counter

967 +

968 + # After update_state_after_alloc, next_stored_block_idx is advanced to

969 + # skip the loaded prefix; reset_cache must bring it back to 0.

970 + for req_status in runner.connector_scheduler._req_status.values():

971 + for group_state in req_status.group_states:

972 + assert group_state.next_stored_block_idx > 0

973 +

974 + # Reset the cache

975 + runner.connector_scheduler.reset_cache()

976 +

977 + # manager.reset_cache() must be called exactly once.

978 + runner.manager.reset_cache.assert_called_once()

979 +

980 + # In-flight load jobs must be queued for flushing to prevent CUDA stream

981 + # races between old loads and new post-reset stores.

982 + assert load_job_ids <= runner.connector_scheduler._current_batch_jobs_to_flush

983 +

984 + # All internal job tracking must be cleared.

985 + assert not runner.connector_scheduler._jobs

986 + assert not runner.connector_scheduler._block_id_to_pending_jobs

987 + if runner.connector_scheduler._blocks_being_loaded is not None:

988 + assert not runner.connector_scheduler._blocks_being_loaded

989 +

990 + # Job reset counter must equal the job counter so that completions for

991 + # pre-reset jobs arriving from workers are silently discarded.

992 + assert runner.connector_scheduler._stale_job_threshold == job_counter_before_reset

993 +

994 + # next_stored_block_idx must be reset to 0 for every active request so

995 + # that post-reset stores restart from block 0.

996 + for req_status in runner.connector_scheduler._req_status.values():

997 + for group_state in req_status.group_states:

998 + assert group_state.next_stored_block_idx == 0

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

+38

@@ -223,6 +223,9 @@ class OffloadingConnectorScheduler:

223

224 # Job ID counter shared by loads and stores.

225 self._job_counter: int = 0

226 + # Threshold value for stale jobs. All job ids >= _stale_job_threshold are

227 + # active jobs.

228 + self._stale_job_threshold: int = 0

229 self._jobs: dict[int, TransferJobStatus] = {}

230

231 # block_id -> pending store job_ids. Used to track jobs that needs

@@ -802,6 +805,13 @@ class OffloadingConnectorScheduler:

805 meta = OffloadingWorkerMetadata()

806 for job_id, count in meta.completed_jobs.items():

807 assert count > 0

808 + if job_id < self._stale_job_threshold:

809 + logger.debug(

810 + "Skipping stale completed job %d (pre-reset counter: %d)",

811 + job_id,

812 + self._stale_job_threshold,

813 + )

814 + continue

815 job_status = self._jobs[job_id]

816 job_status.pending_count -= count

817 if job_status.pending_count > 0:

@@ -882,5 +892,33 @@ class OffloadingConnectorScheduler:

892 lora_name=None,

893 )

894

895 + def reset_cache(self) -> None:

896 + """Reset the offloading manager cache, evicting all stored blocks."""

897 +

898 + # reset_cache cannot be called in the middle of a schedule step

899 + assert not self._current_batch_load_jobs

900 + assert not self._current_batch_jobs_to_flush

901 +

902 + # Flush all in-flight jobs

903 + self._current_batch_jobs_to_flush.update(self._jobs.keys())

904 +

905 + # Reset offloading manager cache

906 + self.manager.reset_cache()

907 +

908 + # Reset store progress so active requests re-offload from block 0

909 + for status in self._req_status.values():

910 + for group_state in status.group_states:

911 + group_state.next_stored_block_idx = 0

912 +

913 + # Discard jobs and save job_counter to be able to discard worker responses

914 + self._stale_job_threshold = self._job_counter

915 + self._jobs.clear()

916 + self._block_id_to_pending_jobs.clear()

917 +

918 + # Note: _current_batch_jobs_to_flush is intentionally NOT cleared.

919 + # The load flush IDs collected above must be delivered to workers.

920 + if self._blocks_being_loaded is not None:

921 + self._blocks_being_loaded.clear()

922 +

923 def shutdown(self) -> None:

924 self.manager.shutdown()

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

@@ -174,6 +174,11 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):

174 def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str | None:

175 return "HND"

176

177 + def reset_cache(self) -> bool | None:

178 + assert self.connector_scheduler is not None

179 + self.connector_scheduler.reset_cache()

180 + return True

181 +

182 def get_kv_connector_stats(self) -> KVConnectorStats | None:

183 if self.connector_worker is None:

184 return None # We only emit stats from the worker-side

vllm/v1/kv_offload/base.py

@@ -218,6 +218,10 @@ class OffloadingManager(ABC):

218 """

219 return ()

220

221 + def reset_cache(self) -> None:

222 + """Evict all tracked blocks and reset internal state."""

223 + return

224 +

225 def shutdown(self) -> None:

226 """Shutdown the manager and release any resources."""

227 return

vllm/v1/kv_offload/cpu/manager.py

+11

@@ -223,6 +223,17 @@ class CPUOffloadingManager(OffloadingManager):

223 )

224 )

225

226 + def reset_cache(self) -> None:

227 + # Clear ALL blocks unconditionally. The scheduler's _stale_job_threshold

228 + # guarantees that complete_load / complete_store are never called for

229 + # pre-reset jobs, so no lazy cleanup is needed. The scheduler also

230 + # flushes in-flight load job IDs to the workers before any new stores

231 + # can begin, preventing a cross-direction data race on reused offload block IDs.

232 + self._policy.clear()

233 +

234 + self._free_list.clear()

235 + self._num_allocated_blocks = 0

236 +

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

238 if self.events is not None:

239 yield from self.events

vllm/v1/kv_offload/cpu/policies/arc.py

@@ -94,6 +94,13 @@ class ARCCachePolicy(CachePolicy):

94 # move to MRU position (end) to keep it fresh in the ghost list

95 self.b2.move_to_end(key)

97 + def clear(self) -> None:

98 + self.t1.clear()

99 + self.t2.clear()

100 + self.b1.clear()

101 + self.b2.clear()

102 + self.target_t1_size = 0.0

103 +

104 def evict(

105 self, n: int, protected: set[OffloadKey]

106 ) -> list[tuple[OffloadKey, BlockStatus]] | None:

vllm/v1/kv_offload/cpu/policies/base.py

@@ -74,3 +74,11 @@ class CachePolicy(ABC):

74 For ARC: ghost list cleanup (trimming to cache_capacity) is performed

75 at the end of a successful eviction.

76 """

77 +

78 + @abstractmethod

79 + def clear(self) -> None:

80 + """

81 + Remove ALL blocks regardless of ref_cnt.

82 +

83 + Ghost lists and adaptive state are also reset.

84 + """

vllm/v1/kv_offload/cpu/policies/lru.py

@@ -28,6 +28,9 @@ class LRUCachePolicy(CachePolicy):

28 if key in self.blocks:

29 self.blocks.move_to_end(key)

31 + def clear(self) -> None:

32 + self.blocks.clear()

33 +

34 def evict(

35 self, n: int, protected: set[OffloadKey]

36 ) -> list[tuple[OffloadKey, BlockStatus]] | None: