@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7e612a0f06

Fork 0

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

[KV Offloading] Implement `reset_cache` for `TieringOffloadingManager` (#44541)

Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Ronen Schaffer committed Jun 15, 2026 at 21:42 UTC 7e612a0f06ad9e31b4609726266fea3cfb0883fe

10 files changed +255 -22

tests/v1/kv_connector/unit/test_offloading_connector.py

+5 -11

@@ -111,7 +111,7 @@ class MockSubscriber:

111 self.sub.close()

112

113

114 -def _wait_for_prefix_cache_reset(llm: LLM, reset_connector: bool = False) -> None:

114 +def _wait_for_prefix_cache_reset(llm: LLM) -> None:

115 """Wait for async offload transfers to finish so prefix cache can reset.

116

117 The GPU-to-CPU offload runs on a CUDA stream asynchronously. While blocks

@@ -119,14 +119,10 @@ def _wait_for_prefix_cache_reset(llm: LLM, reset_connector: bool = False) -> Non

119 ``False``. Between retries we send a dummy single-token prefill to force

120 the engine to step, which polls the worker for completed transfers and

121 frees GPU blocks.

122 -

123 - Args:

124 - llm: The LLM instance to reset.

125 - reset_connector: If True, also reset the KV connector state.

122 """

123 _dummy_params = SamplingParams(max_tokens=1)

124 deadline = time.monotonic() + _RESET_CACHE_TIMEOUT

129 - while not llm.reset_prefix_cache(reset_connector=reset_connector):

125 + while not llm.reset_prefix_cache():

126 if time.monotonic() > deadline:

127 raise TimeoutError(

128 "reset_prefix_cache did not succeed within "

@@ -141,9 +137,7 @@ def _wait_for_prefix_cache_reset(llm: LLM, reset_connector: bool = False) -> Non

137 )

138

139

144 -def _latency_test(

145 - llm: LLM, subscriber: MockSubscriber | None, reset_connector: bool = False

146 -):

140 +def _latency_test(llm: LLM, subscriber: MockSubscriber | None):

141 sampling_params = SamplingParams(max_tokens=1)

142

143 num_times_cpu_better_than_cold = 0

@@ -173,7 +167,7 @@ def _latency_test(

167

168 # Wait for the async CPU offload to finish, then reset prefix cache

169 # so the next generate() must reload from CPU rather than GPU.

176 - _wait_for_prefix_cache_reset(llm, reset_connector=reset_connector)

170 + _wait_for_prefix_cache_reset(llm)

171

172 # Verify CPU stored events arrived (offload is done before we

173 # attempt to load from CPU).

@@ -549,7 +543,7 @@ def test_fs_tiering_offloading(tmp_path) -> None:

543 topic=kv_events_config.topic,

544 )

545 try:

552 - _latency_test(llm, subscriber, reset_connector=True)

546 + _latency_test(llm, subscriber)

547 _accuracy_test(llm, subscriber)

548 finally:

549 subscriber.close()

tests/v1/kv_offload/tiering/test_fs_tier.py

+22

@@ -10,6 +10,7 @@ data integrity throughout the process.

11 import mmap

12 import os

13 +import threading

14 import time

15 from unittest.mock import MagicMock

@@ -22,6 +23,7 @@ from vllm.v1.kv_offload.tiering.base import JobMetadata

23 from vllm.v1.kv_offload.tiering.fs.manager import (

24 FileSystemTierManager,

25 )

26 +from vllm.v1.kv_offload.tiering.fs.thread_pool import DualQueueThreadPool

28 # ---------------------------------------------------------------------------

29 # Helpers

@@ -296,3 +298,23 @@ def test_store_load_data_integrity(fs_tier):

298 assert torch.allclose(tensor[bid], expected[i]), (

299 f"Block {bid} data mismatch after store+load"

300 )

301 +

302 +

303 +def test_wait_idle_blocks_until_tasks_complete():

304 + """wait_idle must not return while a task is still in flight."""

305 + pool = DualQueueThreadPool(n_read_threads=1, n_write_threads=1)

306 + gate = threading.Event()

307 + pool.enqueue_store(job_id=1, n_tasks=1, tasks=[lambda: gate.wait(timeout=5.0)])

308 +

309 + waiter = threading.Thread(target=pool.wait_idle)

310 + waiter.start()

311 + try:

312 + waiter.join(timeout=0.2)

313 + assert waiter.is_alive(), "wait_idle returned before task completed"

314 + gate.set()

315 + waiter.join(timeout=5.0)

316 + assert not waiter.is_alive(), "wait_idle did not unblock"

317 + finally:

318 + gate.set()

319 + pool.shutdown(wait=True)

320 + waiter.join(timeout=5.0)

tests/v1/kv_offload/tiering/test_obj_tier.py

+27

@@ -290,6 +290,33 @@ class TestMockObjTierBasic:

290 assert len(results) == 1

291 assert results[0].success

292

293 + def test_drain_jobs_polls_until_transfers_complete(self):

294 + """drain_jobs must keep polling check_xfer_state until every

295 + in-flight transfer finishes. A buggy implementation that only

296 + polled once would return with _transfers still populated.

297 + """

298 + call_count = [0]

299 + original = self.agent.check_xfer_state

300 +

301 + def delayed(h):

302 + call_count[0] += 1

303 + # Stay in PROC for the first 2 polls, then DONE.

304 + return "PROC" if call_count[0] < 3 else original(h)

305 +

306 + self.agent.check_xfer_state = delayed

307 +

308 + self.tier.submit_store(make_job(1, [key(1)], [0]))

309 + assert self.tier._transfers # in flight

310 +

311 + self.tier.drain_jobs()

312 +

313 + assert not self.tier._transfers # fully drained

314 + assert call_count[0] >= 3 # polled past the initial PROC responses

315 + # Result is buffered for the next get_finished_jobs() call.

316 + results = list(self.tier.get_finished_jobs())

317 + assert len(results) == 1

318 + assert results[0].success

319 +

320

321 class TestMockObjTierMultiBlock:

322 def test_store_multiple_blocks(self):

tests/v1/kv_offload/tiering/test_tiering_offloading.py

+79

@@ -490,6 +490,85 @@ class TestTieringOffloadingManager:

490 # tier2 (block-level) does not get existing blocks here.

491 self.secondary_tier2.submit_store.assert_not_called()

492

493 + def test_reset_cache_clears_all_state(self, manager_setup):

494 + """reset_cache wipes every kind of orchestrator state and resets

495 + primary tier; pending submissions are dropped without being sent

496 + to the secondary tier."""

497 + # Cascade — populates primary blocks and leaves cascade jobs

498 + # in _transfer_jobs (the synchronous example tier has already

499 + # queued completions); reset_cache's drain loop will pick them up.

500 + blocks = to_keys(range(3))

501 + self.manager.prepare_store(blocks, _CTX)

502 + self.manager.complete_store(blocks, _CTX, success=True)

503 + assert self.manager._transfer_jobs

504 +

505 + # Pending promotion submission (deferred — no on_schedule_end after

506 + # the lookup that staged it).

507 + promo_block = to_keys([99])[0]

508 + self.secondary_tier1.blocks[promo_block] = True

509 + assert self.manager.lookup(promo_block, ReqContext(req_id="pending")) is None

510 + assert self.manager._pending_load_submissions

511 +

512 + # Request-level tier registration.

513 + self.secondary_tier1.on_new_request = (

514 + lambda req_context: RequestOffloadingContext(

515 + policy=OffloadPolicy.REQUEST_LEVEL

516 + )

517 + )

518 + self.manager.on_new_request(ReqContext(req_id="rl"))

519 + assert self.manager._request_level_tiers

520 +

521 + # Mark this step as already polled (reset_cache must clear it).

522 + self.manager._processed_jobs_this_step = True

523 +

524 + # Spy: pending submission must NOT reach the tier.

525 + self.secondary_tier1.submit_load = MagicMock(

526 + wraps=self.secondary_tier1.submit_load

527 + )

528 +

529 + self.manager.reset_cache()

530 +

531 + # Orchestrator state cleared.

532 + assert self.manager._transfer_jobs == {}

533 + assert self.manager._pending_load_submissions == {}

534 + assert self.manager._request_level_tiers == {}

535 + assert self.manager._processed_jobs_this_step is False

536 +

537 + # Primary tier reset to a fresh state.

538 + assert self.primary_tier._num_allocated_blocks == 0

539 + assert self.primary_tier._free_list == []

540 + for block in blocks:

541 + assert self.primary_tier.lookup(block, _CTX) is False

542 +

543 + # Pending submission was dropped, not submitted.

544 + self.secondary_tier1.submit_load.assert_not_called()

545 +

546 + def test_reset_cache_drains_all_tiers(self, manager_setup):

547 + """reset_cache must drain each secondary tier before resetting

548 + the primary tier so no tier I/O is touching primary memory.

549 + Without the drain, an in-flight transfer could write into, or

550 + read junk from, a primary slot that the post-reset path has

551 + reallocated.

552 + """

553 + self.secondary_tier1.drain_jobs = MagicMock(

554 + wraps=self.secondary_tier1.drain_jobs

555 + )

556 + self.secondary_tier2.drain_jobs = MagicMock(

557 + wraps=self.secondary_tier2.drain_jobs

558 + )

559 +

560 + # Drive a cascade so a job lands in _transfer_jobs.

561 + blocks = to_keys(range(3))

562 + self.manager.prepare_store(blocks, _CTX)

563 + self.manager.complete_store(blocks, _CTX, success=True)

564 + assert self.manager._transfer_jobs

565 +

566 + self.manager.reset_cache()

567 +

568 + self.secondary_tier1.drain_jobs.assert_called_once()

569 + self.secondary_tier2.drain_jobs.assert_called_once()

570 + assert self.manager._transfer_jobs == {}

571 +

572

573 class TestTieringOffloadingWithoutSecondaryTiers:

574 """Test TieringOffloadingManager with no secondary tiers (backward compat)."""

vllm/v1/kv_offload/tiering/base.py

+17

@@ -193,6 +193,23 @@ class SecondaryTierManager(ABC):

193 """

194 return

195

196 + @abstractmethod

197 + def drain_jobs(self) -> None:

198 + """Block until every submitted load/store job has completed or failed.

199 +

200 + After this returns, no tier I/O is touching the primary memoryview,

201 + and every submitted job's result is available from `get_finished_jobs()`

202 + (yielded by a prior call or queued for the next one). Used by

203 + `TieringOffloadingManager.reset_cache` to release primary slots

204 + without racing with in-flight transfers.

205 +

206 + Implementations must not abort a mid-flight transfer: a partial copy

207 + would corrupt either the primary memoryview or the secondary backing

208 + store. Queued (not-yet-started) transfers may be cancelled, but their

209 + failure result must still appear in `get_finished_jobs()`.

210 + """

211 + pass

212 +

213 def shutdown(self) -> None:

214 """Release resources held by this tier (threads, connections, etc.)."""

215 return

vllm/v1/kv_offload/tiering/example/manager.py

@@ -142,6 +142,12 @@ class ExampleSecondaryTierManager(SecondaryTierManager):

142 def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:

143 return RequestOffloadingContext()

144

145 + @override

146 + def drain_jobs(self) -> None:

147 + """Synchronous tier — submit_*() returns only after the operation

148 + completes, so there is nothing to wait for."""

149 + return

150 +

151 def get_num_blocks(self) -> int:

152 """Get the number of blocks currently stored in this tier."""

153 return len(self.blocks)

vllm/v1/kv_offload/tiering/fs/manager.py

@@ -179,6 +179,10 @@ class FileSystemTierManager(SecondaryTierManager):

179 )

180

181 @override

182 + def drain_jobs(self) -> None:

183 + """Block until all in-flight transfers in the threadpool finish."""

184 + self._pool.wait_idle()

185 +

186 def on_request_finished(self, req_context: ReqContext) -> None:

187 self._lookup_manager.cleanup(req_context.req_id)

188

vllm/v1/kv_offload/tiering/fs/thread_pool.py

+23 -1

@@ -68,6 +68,7 @@ class DualQueueThreadPool:

68 self._stop = False

69 self._threads: list[threading.Thread] = []

70 self._finished_q: deque[tuple[JobId, bool]] = deque()

71 + self._inflight_jobs = 0 # guarded by _condition

73 for i in range(n_read_threads):

74 t = threading.Thread(

@@ -98,6 +99,7 @@ class DualQueueThreadPool:

99 """Enqueue load tasks for a job (high-priority for load-priority threads)."""

100 state = JobState(job_id, n_tasks)

101 with self._condition:

102 + self._inflight_jobs += 1

103 for fn in tasks:

104 self._load_q.append((fn, state))

105 self._condition.notify(n_tasks)

@@ -111,21 +113,38 @@ class DualQueueThreadPool:

113 """Enqueue store tasks for a job (high-priority for store-priority threads)."""

114 state = JobState(job_id, n_tasks)

115 with self._condition:

116 + self._inflight_jobs += 1

117 for fn in tasks:

118 self._store_q.append((fn, state))

119 self._condition.notify(n_tasks)

120

121 def get_finished(self) -> list[tuple[JobId, bool]]:

122 + # No lock needed: deque is thread-safe for concurrent append/popleft,

123 + # and the manager is the sole popper.

124 jobs = []

125 while self._finished_q:

126 jobs.append(self._finished_q.popleft())

127 return jobs

128

129 + def wait_idle(self) -> None:

130 + """Block until there are no in-flight jobs.

131 +

132 + After this returns, every submitted job has had its last task

133 + finish, so no worker thread is still copying data. Note:

134 + completed jobs may still be sitting in ``_finished_q`` waiting

135 + for ``get_finished()`` to drain them.

136 + """

137 + with self._condition:

138 + self._condition.wait_for(lambda: self._inflight_jobs == 0)

139 +

140 def shutdown(self, wait: bool = True) -> None:

141 with self._condition:

142 self._stop = True

143 self._load_q.clear()

144 self._store_q.clear()

145 + # Cancelled tasks will not decrement _inflight_jobs; reset it so a

146 + # subsequent wait_idle() returns instead of hanging.

147 + self._inflight_jobs = 0

148 self._condition.notify_all()

149 if wait:

150 for t in self._threads:

@@ -155,4 +174,7 @@ class DualQueueThreadPool:

174 job_finished, success = state.task_done(False)

175

176 if job_finished:

158 - self._finished_q.append((state.job_id, success))

177 + with self._condition:

178 + self._finished_q.append((state.job_id, success))

179 + self._inflight_jobs -= 1

180 + self._condition.notify_all()

vllm/v1/kv_offload/tiering/manager.py

+29

@@ -590,6 +590,35 @@ class TieringOffloadingManager(OffloadingManager):

590

591 yield from self.primary_tier.take_events()

592

593 + @override

594 + def reset_cache(self) -> None:

595 + """Drop all tracked state in the orchestrator and primary tier.

596 +

597 + Called during sleep, weight update, or resume. Each secondary tier

598 + drains its in-flight transfers via drain_jobs() so no tier I/O is

599 + touching primary memory before the primary tier is reset. A stuck

600 + tier will block here visibly — preferable to silent corruption

601 + from reusing primary slots while a transfer is mid-copy.

602 +

603 + Secondary tiers are intentionally not reset: persistent stores

604 + (FS, network) keep their data across resets.

605 + """

606 + for tier in self.secondary_tiers:

607 + tier.drain_jobs()

608 + # All tier I/O has stopped; consume their completion notifications

609 + # so manager bookkeeping is consistent before the primary reset.

610 + self._process_finished_jobs()

611 +

612 + # Deferred promotion submissions reserve primary slots that the

613 + # reset below invalidates; their submit_load() has not yet been

614 + # called so no tier I/O is touching that memory.

615 + self._pending_load_submissions.clear()

616 +

617 + self.primary_tier.reset_cache()

618 +

619 + self._request_level_tiers.clear()

620 + self._processed_jobs_this_step = False

621 +

622 @override

623 def shutdown(self) -> None:

624 """Shutdown all tiers and release resources."""

vllm/v1/kv_offload/tiering/obj/manager.py

+43 -10

@@ -3,6 +3,7 @@

3 """Object store secondary tier implementation."""

5 import ctypes

6 +import time

7 from collections.abc import Iterable

8 from typing import TYPE_CHECKING, NamedTuple

@@ -104,7 +105,10 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

105 params = {**obj_config.to_nixl_params(), "num_threads": str(io_threads)}

106 self._agent.create_backend("OBJ", params)

107 self._transfers: dict[int, TransferEntry] = {}

107 - self._failed_jobs: list[JobResult] = []

108 + # Buffered results awaiting the next get_finished_jobs() call:

109 + # submission-time failures + poll-time completions accumulated

110 + # during drain_jobs().

111 + self._pending_results: list[JobResult] = []

112 self._primary_reg = None

113 self._block_size_bytes: int = 0

114 root_dir = f"{prefix}/" if prefix else ""

@@ -182,14 +186,14 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

186 files_desc = self._agent.register_memory(nixl_files, "OBJ")

187 if files_desc is None:

188 logger.warning("register_memory (OBJ) failed for job %d", job_id)

185 - self._failed_jobs.append(JobResult(job_id=job_id, success=False))

189 + self._pending_results.append(JobResult(job_id=job_id, success=False))

190 return

191

192 obj_handle = self._agent.prep_xfer_dlist("ObjAgent", files_desc.trim())

193 if not obj_handle:

194 logger.warning("prep_xfer_dlist (OBJ) failed for job %d", job_id)

195 self._agent.deregister_memory(files_desc)

192 - self._failed_jobs.append(JobResult(job_id=job_id, success=False))

196 + self._pending_results.append(JobResult(job_id=job_id, success=False))

197 return

198

199 xfer_handle = self._agent.make_prepped_xfer(

@@ -203,7 +207,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

207 logger.warning("make_prepped_xfer failed for job %d", job_id)

208 self._agent.release_dlist_handle(obj_handle)

209 self._agent.deregister_memory(files_desc)

206 - self._failed_jobs.append(JobResult(job_id=job_id, success=False))

210 + self._pending_results.append(JobResult(job_id=job_id, success=False))

211 return

212

213 state = self._agent.transfer(xfer_handle)

@@ -212,7 +216,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

216 self._agent.release_dlist_handle(obj_handle)

217 self._agent.deregister_memory(files_desc)

218 self._agent.release_xfer_handle(xfer_handle)

215 - self._failed_jobs.append(JobResult(job_id=job_id, success=False))

219 + self._pending_results.append(JobResult(job_id=job_id, success=False))

220 return

221

222 self._transfers[job_id] = TransferEntry(xfer_handle, files_desc, obj_handle)

@@ -241,10 +245,9 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

245 def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:

246 return RequestOffloadingContext()

247

244 - def get_finished_jobs(self) -> Iterable[JobResult]:

245 - """Poll in-flight transfers; return completed (job_id, success) pairs."""

246 - results: list[JobResult] = self._failed_jobs

247 - self._failed_jobs = []

248 + def _poll_active_transfers(self) -> None:

249 + """Poll all in-flight transfers once; move newly-completed (success or

250 + failure) into ``_pending_results`` and release their NIXL handles."""

251 for job_id, entry in list(self._transfers.items()):

252 try:

253 state = self._agent.check_xfer_state(entry.xfer_handle)

@@ -263,9 +266,39 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

266 self._agent.release_xfer_handle(entry.xfer_handle)

267 self._agent.release_dlist_handle(entry.obj_handle)

268 self._agent.deregister_memory(entry.files_desc)

266 - results.append(JobResult(job_id=job_id, success=success))

269 + self._pending_results.append(JobResult(job_id=job_id, success=success))

270 +

271 + def get_finished_jobs(self) -> Iterable[JobResult]:

272 + """Poll in-flight transfers; return completed (job_id, success) pairs."""

273 + self._poll_active_transfers()

274 + results = self._pending_results

275 + self._pending_results = []

276 return results

277

278 + def drain_jobs(self) -> None:

279 + """Block until every submitted transfer has completed or failed.

280 +

281 + nixl exposes only ``check_xfer_state`` (poll-based), so this loops

282 + until ``_transfers`` is empty. Results accumulate in

283 + ``_pending_results`` and are surfaced by the next

284 + ``get_finished_jobs()`` call.

285 + """

286 + start = time.monotonic()

287 + warned = False

288 + while self._transfers:

289 + self._poll_active_transfers()

290 + if not self._transfers:

291 + break

292 + if not warned and time.monotonic() - start > 5.0:

293 + logger.warning(

294 + "ObjectStoreSecondaryTierManager.drain_jobs: still "

295 + "draining after 5s (%d transfers in flight); a stuck "

296 + "transfer will block the engine.",

297 + len(self._transfers),

298 + )

299 + warned = True

300 + time.sleep(0.001)

301 +

302 def shutdown(self) -> None:

303 self._lookup_manager.shutdown()

304 for job_id, entry in self._transfers.items():