A high-throughput and memory-efficient inference and serving engine for LLMs
KV-Cache multi-tier offloading async batched lookup (#44193)
Signed-off-by: Effi Ofer <effi.ofer@gmail.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Effi Ofer committed
Jun 10, 2026 at 17:59 UTC
af65e08fc5e42a909e66be34ff9f35e5307fc6ab
6 files changed
+518
-31
tests/v1/kv_offload/tiering/test_async_lookup.py
new
+158
@@ -0,0 +1,158 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for AsyncLookupManager."""
4
+
5
+import threading
6
+from collections.abc import Iterable
7
+
8
+from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key
9
+from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager
10
+
11
+
12
+def _key(i: int) -> OffloadKey:
13
+ return make_offload_key(str(i).encode(), 0)
14
+
15
+
16
+def _ctx(req_id: str = "r1") -> ReqContext:
17
+ return ReqContext(req_id=req_id)
18
+
19
+
20
+class InMemoryLookupManager(AsyncLookupManager):
21
+ """Test subclass backed by an in-memory set."""
22
+
23
+ def __init__(self, existing_keys: set[OffloadKey] | None = None):
24
+ super().__init__(tier_type="test")
25
+ self._existing = existing_keys or set()
26
+ self._results_ready = threading.Event()
27
+
28
+ def batch_lookup(
29
+ self, keys: list[OffloadKey], req_context: ReqContext
30
+ ) -> Iterable[bool]:
31
+ results = [k in self._existing for k in keys]
32
+ self._results_ready.set()
33
+ return results
34
+
35
+
36
+class TestAsyncLookupManager:
37
+ def test_new_key_returns_none(self):
38
+ mgr = InMemoryLookupManager()
39
+ assert mgr.lookup(_key(1), _ctx()) is None
40
+ mgr.shutdown()
41
+
42
+ def test_found_key_returns_true(self):
43
+ mgr = InMemoryLookupManager(existing_keys={_key(1)})
44
+ assert mgr.lookup(_key(1), _ctx()) is None
45
+ mgr.flush()
46
+ mgr._results_ready.wait()
47
+ mgr._results_ready.clear()
48
+ assert mgr.lookup(_key(1), _ctx()) is True
49
+ mgr.shutdown()
50
+
51
+ def test_not_found_key_returns_false(self):
52
+ mgr = InMemoryLookupManager(existing_keys=set())
53
+ assert mgr.lookup(_key(1), _ctx()) is None
54
+ mgr.flush()
55
+ mgr._results_ready.wait()
56
+ mgr._results_ready.clear()
57
+ assert mgr.lookup(_key(1), _ctx()) is False
58
+ mgr.shutdown()
59
+
60
+ def test_multiple_keys_single_step(self):
61
+ existing = {_key(1), _key(3)}
62
+ mgr = InMemoryLookupManager(existing_keys=existing)
63
+ ctx = _ctx()
64
+ for i in range(1, 5):
65
+ assert mgr.lookup(_key(i), ctx) is None
66
+ mgr.flush()
67
+ mgr._results_ready.wait()
68
+ mgr._results_ready.clear()
69
+ assert mgr.lookup(_key(1), ctx) is True
70
+ assert mgr.lookup(_key(2), ctx) is False
71
+ assert mgr.lookup(_key(3), ctx) is True
72
+ assert mgr.lookup(_key(4), ctx) is False
73
+ mgr.shutdown()
74
+
75
+ def test_cleanup_removes_entries(self):
76
+ mgr = InMemoryLookupManager(existing_keys={_key(1)})
77
+ ctx = _ctx("req_a")
78
+ mgr.lookup(_key(1), ctx)
79
+ mgr.flush()
80
+ mgr._results_ready.wait()
81
+ mgr._results_ready.clear()
82
+ assert mgr.lookup(_key(1), ctx) is True
83
+ mgr.cleanup("req_a")
84
+ assert _key(1) not in mgr._lookup_state
85
+ mgr.shutdown()
86
+
87
+ def test_cleanup_preserves_shared_entries(self):
88
+ mgr = InMemoryLookupManager(existing_keys={_key(1)})
89
+ ctx_a = _ctx("req_a")
90
+ ctx_b = _ctx("req_b")
91
+ mgr.lookup(_key(1), ctx_a)
92
+ mgr.lookup(_key(1), ctx_b)
93
+ mgr.flush()
94
+ mgr._results_ready.wait()
95
+ mgr._results_ready.clear()
96
+ # Drain so result is applied
97
+ mgr.lookup(_key(1), ctx_a)
98
+ mgr.cleanup("req_a")
99
+ # Key still present because req_b still references it
100
+ assert _key(1) in mgr._lookup_state
101
+ mgr.cleanup("req_b")
102
+ assert _key(1) not in mgr._lookup_state
103
+ mgr.shutdown()
104
+
105
+ def test_flush_no_queue_post_when_empty(self):
106
+ mgr = InMemoryLookupManager()
107
+ mgr.flush()
108
+ assert mgr._lookup_queue.empty()
109
+ mgr.shutdown()
110
+
111
+ def test_repeated_lookup_same_key_no_duplicate_batch(self):
112
+ mgr = InMemoryLookupManager(existing_keys={_key(1)})
113
+ ctx = _ctx()
114
+ mgr.lookup(_key(1), ctx)
115
+ mgr.lookup(_key(1), ctx)
116
+ assert len(mgr._lookup_batch) == 1
117
+ mgr.shutdown()
118
+
119
+ def test_cleanup_unknown_req_id_is_noop(self):
120
+ mgr = InMemoryLookupManager(existing_keys={_key(1)})
121
+ ctx = _ctx("req_a")
122
+ mgr.lookup(_key(1), ctx)
123
+ mgr.flush()
124
+ mgr._results_ready.wait()
125
+ mgr._results_ready.clear()
126
+ mgr.lookup(_key(1), ctx)
127
+ mgr.cleanup("nonexistent")
128
+ assert _key(1) in mgr._lookup_state
129
+ mgr.shutdown()
130
+
131
+ def test_multiple_flushes_across_steps(self):
132
+ existing = {_key(1), _key(2), _key(3)}
133
+ mgr = InMemoryLookupManager(existing_keys=existing)
134
+ ctx = _ctx()
135
+
136
+ # Step 1: lookup key 1, flush
137
+ mgr.lookup(_key(1), ctx)
138
+ mgr.flush()
139
+ mgr._results_ready.wait()
140
+ mgr._results_ready.clear()
141
+
142
+ # Step 2: lookup keys 2 and 3, flush
143
+ mgr.lookup(_key(2), ctx)
144
+ mgr.lookup(_key(3), ctx)
145
+ mgr.flush()
146
+ mgr._results_ready.wait()
147
+ mgr._results_ready.clear()
148
+
149
+ # All results should be available
150
+ assert mgr.lookup(_key(1), ctx) is True
151
+ assert mgr.lookup(_key(2), ctx) is True
152
+ assert mgr.lookup(_key(3), ctx) is True
153
+ mgr.shutdown()
154
+
155
+ def test_shutdown_unblocks_worker(self):
156
+ mgr = InMemoryLookupManager()
157
+ mgr.shutdown()
158
+ assert not mgr._thread.is_alive()
tests/v1/kv_offload/tiering/test_fs_tier.py
+27
-12
@@ -71,9 +71,9 @@ def make_job(
71
)
72
73
74
-def drain(tier: FileSystemTierManager, max_rounds: int = 40) -> list:
74
+def drain(tier: FileSystemTierManager, max_rounds: int = 100) -> list:
75
"""
76
- Call get_finished_jobs() repeatedly until no new results arrive for 5
76
+ Call get_finished_jobs() repeatedly until no new results arrive for 20
77
consecutive rounds or max_rounds is reached.
78
"""
79
results = []
@@ -86,11 +86,29 @@ def drain(tier: FileSystemTierManager, max_rounds: int = 40) -> list:
86
idle = 0
87
else:
88
idle += 1
89
- if idle >= 5:
89
+ if idle >= 20:
90
break
91
return results
92
93
94
+def lookup_and_wait(
95
+ tier: FileSystemTierManager,
96
+ keys: list[OffloadKey],
97
+ ctx: ReqContext = _CTX,
98
+ timeout: float = 1.0,
99
+) -> list[bool]:
100
+ """Perform a full async lookup cycle and return resolved results."""
101
+ for k in keys:
102
+ tier.lookup(k, ctx)
103
+ tier.on_schedule_end()
104
+ deadline = time.monotonic() + timeout
105
+ while time.monotonic() < deadline:
106
+ if not tier._lookup_manager._pending_results.empty():
107
+ break
108
+ time.sleep(0.01)
109
+ return [tier.lookup(k, ctx) for k in keys]
110
+
111
+
112
def _page_aligned_zero_tensor(
113
num_blocks: int, block_elements: int, dtype: torch.dtype = _DTYPE
114
) -> torch.Tensor:
@@ -145,8 +163,8 @@ def fs_tier(tmp_path):
163
164
def test_lookup_empty_tier(fs_tier):
165
tier, _ = fs_tier
148
- assert tier.lookup(key(1), _CTX) is False
149
- assert tier.lookup(key(2), _CTX) is False
166
+ results = lookup_and_wait(tier, [key(1), key(2)])
167
+ assert results == [False, False]
168
169
170
def test_store_creates_file_and_lookup_succeeds(fs_tier):
@@ -156,7 +174,7 @@ def test_store_creates_file_and_lookup_succeeds(fs_tier):
174
results = drain(tier)
175
assert len(results) == 1
176
assert results[0].success
159
- assert tier.lookup(key(1), _CTX) is True
177
+ assert lookup_and_wait(tier, [key(1)]) == [True]
178
dest = tier.file_mapper.get_file_name(key(1))
179
assert os.path.exists(dest), f"Expected file at {dest}"
180
@@ -168,16 +186,14 @@ def test_store_then_load_roundtrip(fs_tier):
186
store_results = drain(tier)
187
assert all(r.success for r in store_results)
188
171
- assert tier.lookup(key(1), _CTX) is True
172
- assert tier.lookup(key(2), _CTX) is True
189
+ assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True]
190
191
job_l = make_job(2, [key(1), key(2)], [2, 3], is_promotion=True)
192
tier.submit_load(job_l)
193
load_results = drain(tier)
194
assert all(r.success for r in load_results)
195
# Blocks stay on disk after load
179
- assert tier.lookup(key(1), _CTX) is True
180
- assert tier.lookup(key(2), _CTX) is True
196
+ assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True]
197
198
199
def test_invalid_path_raises_at_construction():
@@ -213,8 +229,7 @@ def test_multiple_jobs_tracked_independently(fs_tier):
229
results = drain(tier)
230
job_ids = {r.job_id for r in results}
231
assert job_ids == {1, 2}
216
- assert tier.lookup(key(1), _CTX) is True
217
- assert tier.lookup(key(2), _CTX) is True
232
+ assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True]
233
234
235
def test_multi_block_job_partial_failure(fs_tier):
tests/v1/kv_offload/tiering/test_obj_tier.py
+25
-8
@@ -8,6 +8,7 @@ without S3 credentials or a live object store. They verify the manager's
8
state machine: job submission, transfer completion polling, and lookup.
9
"""
10
11
+import time
12
import uuid
13
from collections.abc import Callable
14
from types import SimpleNamespace
@@ -206,6 +207,24 @@ def drain(
207
return results
208
209
210
+def lookup_and_wait(
211
+ tier: ObjectStoreSecondaryTierManager,
212
+ keys: list[OffloadKey],
213
+ ctx: ReqContext = _CTX,
214
+ timeout: float = 1.0,
215
+) -> list[bool]:
216
+ """Perform a full async lookup cycle and return resolved results."""
217
+ for k in keys:
218
+ tier.lookup(k, ctx)
219
+ tier.on_schedule_end()
220
+ deadline = time.monotonic() + timeout
221
+ while time.monotonic() < deadline:
222
+ if not tier._lookup_manager._pending_results.empty():
223
+ break
224
+ time.sleep(0.01)
225
+ return [tier.lookup(k, ctx) for k in keys]
226
+
227
+
228
# ---------------------------------------------------------------------------
229
# Tests
230
# ---------------------------------------------------------------------------
@@ -216,19 +235,19 @@ class TestMockObjTierBasic:
235
self.tier, self.agent = _make_tier(num_blocks=4)
236
237
def test_lookup_empty_tier(self):
219
- assert self.tier.lookup(key(1), _CTX) is False
238
+ assert lookup_and_wait(self.tier, [key(1)]) == [False]
239
240
def test_store_and_lookup(self):
241
self.tier.submit_store(make_job(1, [key(1)], [0]))
242
results = drain(self.tier)
243
assert len(results) == 1
244
assert results[0].success
226
- assert self.tier.lookup(key(1), _CTX) is True
245
+ assert lookup_and_wait(self.tier, [key(1)]) == [True]
246
247
def test_lookup_unrelated_key_returns_false(self):
248
self.tier.submit_store(make_job(1, [key(1)], [0]))
249
drain(self.tier)
231
- assert self.tier.lookup(key(999), _CTX) is False
250
+ assert lookup_and_wait(self.tier, [key(999)]) == [False]
251
252
def test_store_then_load_roundtrip(self):
253
self.tier.submit_store(make_job(1, [key(1), key(2)], [0, 1]))
@@ -280,15 +299,13 @@ class TestMockObjTierMultiBlock:
299
results = drain(tier)
300
assert len(results) == 1
301
assert results[0].success
283
- assert all(tier.lookup(k, _CTX) for k in keys)
302
+ assert lookup_and_wait(tier, keys) == [True] * 8
303
304
def test_partial_block_lookup(self):
305
tier, _ = _make_tier(num_blocks=4)
306
tier.submit_store(make_job(1, [key(0), key(1)], [0, 1]))
307
drain(tier)
289
- assert tier.lookup(key(0), _CTX) is True
290
- assert tier.lookup(key(1), _CTX) is True
291
- assert tier.lookup(key(2), _CTX) is False
308
+ assert lookup_and_wait(tier, [key(0), key(1), key(2)]) == [True, True, False]
309
310
311
class TestMockObjTierFailures:
@@ -297,7 +314,7 @@ class TestMockObjTierFailures:
314
agent.query_memory = lambda *a, **k: (_ for _ in ()).throw(
315
RuntimeError("backend error")
316
)
300
- assert tier.lookup(key(1), _CTX) is False
317
+ assert lookup_and_wait(tier, [key(1)]) == [False]
318
319
def test_submit_store_register_memory_failure_reported_in_get_finished(self):
320
tier, agent = _make_tier(num_blocks=4)
vllm/v1/kv_offload/tiering/async_lookup.py
new
+231
@@ -0,0 +1,231 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+AsyncLookupManager: per-tier async lookup manager for secondary tier
5
+existence checks.
6
+
7
+Each secondary tier that wants non-blocking lookups composes its own
8
+AsyncLookupManager instance internally. The manager maintains lookup
9
+state and uses a background thread to execute batch_lookup() calls.
10
+
11
+Locking design
12
+--------------
13
+There is no explicit lock. Thread safety is achieved by ownership:
14
+
15
+* _lookup_state and _lookup_batch are owned exclusively by the scheduler
16
+ thread. lookup(), flush(), and cleanup() read and write them directly.
17
+
18
+* _lookup_queue is written by the scheduler (flush → put_nowait, one item
19
+ per step) and read by the background thread (get). queue.Queue is
20
+ thread-safe.
21
+
22
+* _pending_results is written by the background thread (put) and read by
23
+ the scheduler (get_nowait inside drain_results). queue.SimpleQueue is
24
+ thread-safe by design.
25
+
26
+lookup() accumulates new keys in _lookup_batch without touching the queue.
27
+flush() is called once per step from the tier's on_schedule_end(), posting
28
+the entire batch as a single queue item so the background thread sees one
29
+batch per step.
30
+drain_results() is called before any lookup() calls in the same step, so
31
+lookup() is a pure OrderedDict operation.
32
+"""
33
+
34
+import queue
35
+import threading
36
+from abc import ABC, abstractmethod
37
+from collections.abc import Iterable
38
+from dataclasses import dataclass, field
39
+
40
+from vllm.logger import init_logger
41
+from vllm.v1.kv_offload.base import OffloadKey, ReqContext
42
+
43
+logger = init_logger(__name__)
44
+
45
+
46
+@dataclass(slots=True)
47
+class LookupState:
48
+ result: bool | None = None # True (found), False (not found), None
49
+ request_ids: set[str] = field(default_factory=set) # requests asking for the lookup
50
+
51
+
52
+class AsyncLookupManager(ABC):
53
+ """
54
+ Per-tier async lookup manager for secondary tier existence checks.
55
+
56
+ Each secondary tier that wants non-blocking lookups composes its own
57
+ AsyncLookupManager instance internally. The manager maintains lookup
58
+ state (cache, queue) and uses a background thread to execute the actual
59
+ batch_lookup() calls.
60
+
61
+ Subclasses implement only batch_lookup() — all queue management,
62
+ state tracking, and result delivery is provided by this base class.
63
+
64
+ The owning tier delegates its lookup(), on_schedule_end(), and
65
+ on_request_finished() to this manager:
66
+ - lookup() → drain_results() + lookup state check
67
+ - on_schedule_end() → flush()
68
+ - on_request_finished() → cleanup()
69
+ """
70
+
71
+ def __init__(
72
+ self,
73
+ tier_type: str,
74
+ ) -> None:
75
+ self._tier_type = tier_type
76
+
77
+ # key → LookupState; scheduler-owned, no lock needed.
78
+ self._lookup_state: dict[OffloadKey, LookupState] = {}
79
+ # req_id → keys looked up by that request (reverse index for cleanup).
80
+ self._req_keys: dict[str, set[OffloadKey]] = {}
81
+
82
+ # Accumulates (key, req_context) pairs during lookup() calls.
83
+ # Flushed as one queue item per step by flush().
84
+ self._lookup_batch: list[tuple[OffloadKey, ReqContext]] = []
85
+
86
+ # Scheduler → worker: one full step's batch per item.
87
+ # None is used as a shutdown sentinel.
88
+ self._lookup_queue: queue.SimpleQueue[
89
+ list[tuple[OffloadKey, ReqContext]] | None
90
+ ] = queue.SimpleQueue()
91
+
92
+ # Worker → scheduler: completed result batches.
93
+ # Each item is a list of (key, found) pairs.
94
+ # SimpleQueue is explicitly thread-safe for one writer / one reader.
95
+ self._pending_results: queue.SimpleQueue[list[tuple[OffloadKey, bool]]] = (
96
+ queue.SimpleQueue()
97
+ )
98
+ self._need_to_drain: bool = False
99
+
100
+ self._thread = threading.Thread(
101
+ target=self._worker,
102
+ name=f"vllm_offloading_lookup_{tier_type}",
103
+ daemon=True,
104
+ )
105
+ self._thread.start()
106
+
107
+ @abstractmethod
108
+ def batch_lookup(
109
+ self, keys: list[OffloadKey], req_context: ReqContext
110
+ ) -> Iterable[bool]:
111
+ """
112
+ Check whether a batch of blocks exist in this tier.
113
+
114
+ Called from the worker thread — must be synchronous and must not
115
+ touch the primary tier or scheduler state.
116
+
117
+ Returns a list parallel to keys: True if present, False if not.
118
+ """
119
+ ...
120
+
121
+ # ------------------------------------------------------------------
122
+ # Scheduler-thread API
123
+ # ------------------------------------------------------------------
124
+
125
+ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
126
+ """
127
+ Non-blocking lookup called from the scheduler thread.
128
+
129
+ Returns:
130
+ True — block is present in this tier.
131
+ False — block is not present in this tier.
132
+ None — result not yet available; retry next step.
133
+ """
134
+ if self._need_to_drain:
135
+ self.drain_results()
136
+ self._need_to_drain = False
137
+ req_id = req_context.req_id
138
+ state = self._lookup_state.get(key)
139
+ if state is None:
140
+ state = LookupState()
141
+ self._lookup_state[key] = state
142
+ self._lookup_batch.append((key, req_context))
143
+ state.request_ids.add(req_id)
144
+ self._req_keys.setdefault(req_id, set()).add(key)
145
+ return state.result
146
+
147
+ def flush(self) -> None:
148
+ """Post this step's accumulated keys to the worker thread.
149
+
150
+ Called once per step from on_schedule_end() after all lookup() calls
151
+ are done. The worker receives the full batch and processes it during
152
+ the model-execution window, maximising time available before the next
153
+ step's drain_results(). Safe to call with an empty batch (no-op).
154
+ """
155
+ self._need_to_drain = True
156
+ if self._lookup_batch:
157
+ self._lookup_queue.put(self._lookup_batch)
158
+ self._lookup_batch = []
159
+
160
+ def drain_results(self) -> None:
161
+ """Apply pending worker results to _lookup_state.
162
+
163
+ Called from lookup() before checking state.
164
+ """
165
+ while True:
166
+ try:
167
+ batch = self._pending_results.get_nowait()
168
+ except queue.Empty:
169
+ break
170
+ for key, result in batch:
171
+ state = self._lookup_state.get(key)
172
+ if state is not None:
173
+ state.result = result
174
+
175
+ def cleanup(self, req_id: str) -> None:
176
+ """Remove entries no longer needed by any active request.
177
+
178
+ Called from the tier's on_request_finished(). Uses the reverse
179
+ index to visit only keys associated with this request.
180
+ """
181
+ for key in self._req_keys.pop(req_id, ()):
182
+ state = self._lookup_state[key]
183
+ state.request_ids.discard(req_id)
184
+ if not state.request_ids:
185
+ del self._lookup_state[key]
186
+
187
+ def shutdown(self) -> None:
188
+ """Stop the worker thread."""
189
+ self._lookup_queue.put(None) # unblock _worker from _lookup_queue.get()
190
+ self._thread.join()
191
+
192
+ # ------------------------------------------------------------------
193
+ # Internal helpers
194
+ # ------------------------------------------------------------------
195
+
196
+ def _worker(self) -> None:
197
+ while True:
198
+ pending = self._lookup_queue.get()
199
+ if pending is None:
200
+ break
201
+
202
+ # Group by req_id.
203
+ batches: dict[str, tuple[ReqContext, list[OffloadKey]]] = {}
204
+ for key, req_context in pending:
205
+ req_id = req_context.req_id
206
+ if req_id not in batches:
207
+ batches[req_id] = (req_context, [])
208
+ batches[req_id][1].append(key)
209
+
210
+ if not batches:
211
+ continue
212
+
213
+ results: list[tuple[OffloadKey, bool]] = []
214
+ for req_context, keys in batches.values():
215
+ try:
216
+ hits = self.batch_lookup(keys, req_context)
217
+ except Exception as exc:
218
+ logger.warning(
219
+ "batch_lookup failed on tier %s for %d keys: %s",
220
+ self._tier_type,
221
+ len(keys),
222
+ exc,
223
+ )
224
+ hits = (False for _ in keys)
225
+
226
+ for key, hit in zip(keys, hits):
227
+ results.append((key, hit))
228
+
229
+ # Post the entire batch as one item — no lock needed.
230
+ if results:
231
+ self._pending_results.put(results)
vllm/v1/kv_offload/tiering/fs/manager.py
+33
-6
@@ -26,6 +26,7 @@ from typing_extensions import override
26
from vllm.logger import init_logger
27
from vllm.v1.kv_offload.base import OffloadKey, ReqContext
28
from vllm.v1.kv_offload.file_mapper import FileMapper
29
+from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager
30
from vllm.v1.kv_offload.tiering.base import (
31
JobMetadata,
32
JobResult,
@@ -41,6 +42,23 @@ if TYPE_CHECKING:
42
logger = init_logger(__name__)
43
44
45
+class FsAsyncLookupManager(AsyncLookupManager):
46
+ """Async lookup manager for FileSystemTierManager."""
47
+
48
+ def __init__(
49
+ self,
50
+ tier: "FileSystemTierManager",
51
+ tier_type: str,
52
+ ) -> None:
53
+ super().__init__(tier_type=tier_type)
54
+ self._tier = tier
55
+
56
+ def batch_lookup(
57
+ self, keys: list[OffloadKey], req_context: ReqContext
58
+ ) -> Iterable[bool]:
59
+ return (os.path.exists(self._tier.file_mapper.get_file_name(k)) for k in keys)
60
+
61
+
62
class FileSystemTierManager(SecondaryTierManager):
63
"""
64
Pure-Python disk-backed secondary tier.
@@ -111,15 +129,15 @@ class FileSystemTierManager(SecondaryTierManager):
129
thread_name_prefix="vllm_kv_py_fs",
130
)
131
132
+ self._lookup_manager = FsAsyncLookupManager(tier=self, tier_type=self.tier_type)
133
+
134
@override
135
def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
136
return RequestOffloadingContext()
137
138
@override
119
- def lookup(
120
- self, key: OffloadKey, req_context: ReqContext | None = None
121
- ) -> bool | None:
122
- return os.path.exists(self.file_mapper.get_file_name(key))
139
+ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
140
+ return self._lookup_manager.lookup(key, req_context)
141
142
@override
143
def submit_store(self, job_metadata: JobMetadata) -> None:
@@ -159,12 +177,21 @@ class FileSystemTierManager(SecondaryTierManager):
177
for job_id, success in self._pool.get_finished()
178
)
179
180
+ @override
181
+ def on_request_finished(self, req_context: ReqContext) -> None:
182
+ self._lookup_manager.cleanup(req_context.req_id)
183
+
184
+ @override
185
+ def on_schedule_end(self) -> None:
186
+ self._lookup_manager.flush()
187
+
188
@override
189
def shutdown(self) -> None:
190
"""
191
Release resources held by this tier.
192
167
- Shuts down the thread pool, clearing pending tasks and waiting for
168
- active threads to complete.
193
+ Shuts down the lookup manager and the thread pool,
194
+ clearing pending tasks and waiting for active threads to complete.
195
"""
196
+ self._lookup_manager.shutdown()
197
self._pool.shutdown(wait=True)
vllm/v1/kv_offload/tiering/obj/manager.py
+44
-5
@@ -11,6 +11,7 @@ from vllm.distributed.nixl_utils import nixl_agent_config
11
from vllm.logger import init_logger
12
from vllm.v1.kv_offload.base import OffloadKey, ReqContext
13
from vllm.v1.kv_offload.file_mapper import FileMapper
14
+from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager
15
from vllm.v1.kv_offload.tiering.base import (
16
JobMetadata,
17
JobResult,
@@ -49,6 +50,37 @@ class TransferEntry(NamedTuple):
50
obj_handle: "nixl_prepped_dlist_handle"
51
52
53
+class ObjAsyncLookupManager(AsyncLookupManager):
54
+ """Async lookup manager for ObjectStoreSecondaryTierManager.
55
+
56
+ Batches existence probes into a single query_memory() call so the
57
+ background thread issues one round-trip per step instead of one per key.
58
+ """
59
+
60
+ def __init__(
61
+ self,
62
+ tier: "ObjectStoreSecondaryTierManager",
63
+ tier_type: str,
64
+ ) -> None:
65
+ super().__init__(tier_type=tier_type)
66
+ self._tier = tier
67
+
68
+ def batch_lookup(
69
+ self, keys: list[OffloadKey], req_context: ReqContext
70
+ ) -> Iterable[bool]:
71
+ descriptors = [
72
+ (
73
+ _PROBE_ADDR,
74
+ _PROBE_LEN,
75
+ _PROBE_DEV_ID,
76
+ self._tier._file_mapper.get_file_name(k),
77
+ )
78
+ for k in keys
79
+ ]
80
+ results = self._tier._agent.query_memory(descriptors, "OBJ", "OBJ")
81
+ return (r is not None for r in results)
82
+
83
+
84
class ObjectStoreSecondaryTierManager(SecondaryTierManager):
85
"""Secondary tier that offloads KV cache blocks to an S3-compatible store.
86
@@ -99,6 +131,10 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
131
self._agent.prep_xfer_dlist("NIXL_INIT_AGENT", all_blocks, "DRAM")
132
)
133
134
+ self._lookup_manager = ObjAsyncLookupManager(
135
+ tier=self, tier_type=self.tier_type
136
+ )
137
+
138
def _probe_connectivity(self) -> None:
139
"""Verify object store connectivity at startup via a NIXL lookup probe.
140
@@ -179,11 +215,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
215
self._transfers[job_id] = TransferEntry(xfer_handle, files_desc, obj_handle)
216
217
def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
182
- try:
183
- return self._exists(self._file_mapper.get_file_name(key))
184
- except Exception as e:
185
- logger.warning("lookup failed for key %s: %s", key, e)
186
- return False
218
+ return self._lookup_manager.lookup(key, req_context)
219
220
def submit_store(self, job_metadata: JobMetadata) -> None:
221
obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys)
@@ -197,6 +229,12 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
229
job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_READ
230
)
231
232
+ def on_request_finished(self, req_context: ReqContext) -> None:
233
+ self._lookup_manager.cleanup(req_context.req_id)
234
+
235
+ def on_schedule_end(self) -> None:
236
+ self._lookup_manager.flush()
237
+
238
def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
239
return RequestOffloadingContext()
240
@@ -226,6 +264,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
264
return results
265
266
def shutdown(self) -> None:
267
+ self._lookup_manager.shutdown()
268
for job_id, entry in self._transfers.items():
269
try:
270
self._agent.release_xfer_handle(entry.xfer_handle)