@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 11f6b545d4
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, 19:00 UTC

[kv_offload] Add multi-tier KV cache offloading framework (#40020)

Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>

Ronen Schaffer committed May 13, 2026 at 17:21 UTC 11f6b545d41a42422410bb529466463f466cafd9
14 files changed +1841 -11
tests/v1/kv_connector/unit/test_offloading_connector.py
+45
@@ -292,3 +292,48 @@ def test_cpu_offloading(
292 if subscriber is not None:
293 subscriber.close()
294 del llm
295 +
296 +
297 +def test_tiering_offloading() -> None:
298 + """Tests OffloadingConnector with TieringOffloadingSpec."""
299 + extra_config: dict = {
300 + "cpu_bytes_to_use": 500 << 20,
301 + "block_size": 48,
302 + "spec_name": "TieringOffloadingSpec",
303 + "secondary_tiers": [{"type": "example"}],
304 + }
305 + kv_transfer_config = KVTransferConfig(
306 + kv_connector="OffloadingConnector",
307 + kv_role="kv_both",
308 + kv_connector_extra_config=extra_config,
309 + )
310 +
311 + port: int
312 + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
313 + s.bind(("0.0.0.0", 0))
314 + port = s.getsockname()[1]
315 + events_endpoint = f"tcp://*:{port}"
316 + kv_events_config = KVEventsConfig(
317 + enable_kv_cache_events=True,
318 + publisher="zmq",
319 + endpoint=events_endpoint,
320 + topic="test",
321 + )
322 +
323 + llm = LLM(
324 + model="meta-llama/Llama-3.2-1B-Instruct",
325 + max_model_len=4096,
326 + gpu_memory_utilization=0.5,
327 + kv_events_config=kv_events_config,
328 + kv_transfer_config=kv_transfer_config,
329 + )
330 + subscriber = MockSubscriber(
331 + events_endpoint.replace("*", "127.0.0.1"),
332 + topic=kv_events_config.topic,
333 + )
334 + try:
335 + _latency_test(llm, subscriber)
336 + _accuracy_test(llm, subscriber)
337 + finally:
338 + subscriber.close()
339 + del llm
tests/v1/kv_offload/cpu/test_manager.py
+37
@@ -262,6 +262,43 @@ def test_cpu_manager():
262 )
263
264
265 +def test_prepare_load_preserves_key_order():
266 + """block_ids[i] must correspond to keys[i] (co-indexed invariant)."""
267 + manager = CPUOffloadingManager(num_blocks=4, cache_policy="lru")
268 +
269 + key_a, key_b, key_c = to_key(0), to_key(1), to_key(2)
270 +
271 + # Store all three keys and learn their block ID assignments
272 + store_output = manager.prepare_store([key_a, key_b, key_c], _EMPTY_REQ_CTX)
273 + assert store_output is not None
274 + assert isinstance(store_output.store_spec, CPULoadStoreSpec)
275 + key_to_block_id = {
276 + k: int(bid)
277 + for k, bid in zip(store_output.keys_to_store, store_output.store_spec.block_ids)
278 + }
279 + manager.complete_store([key_a, key_b, key_c], _EMPTY_REQ_CTX)
280 +
281 + # Forward order: [a, b, c]
282 + spec_fwd = manager.prepare_load([key_a, key_b, key_c], _EMPTY_REQ_CTX)
283 + assert isinstance(spec_fwd, CPULoadStoreSpec)
284 + assert [int(x) for x in spec_fwd.block_ids] == [
285 + key_to_block_id[key_a],
286 + key_to_block_id[key_b],
287 + key_to_block_id[key_c],
288 + ]
289 + manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant
290 +
291 + # Arbitrary permutation: [b, c, a]
292 + spec_perm = manager.prepare_load([key_b, key_c, key_a], _EMPTY_REQ_CTX)
293 + assert isinstance(spec_perm, CPULoadStoreSpec)
294 + assert [int(x) for x in spec_perm.block_ids] == [
295 + key_to_block_id[key_b],
296 + key_to_block_id[key_c],
297 + key_to_block_id[key_a],
298 + ]
299 + manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant
300 +
301 +
302 class TestARCPolicy:
303 """Unit tests for CPUOffloadingManager with ARC eviction policy."""
304
tests/v1/kv_offload/test_tiering_offloading.py new
+565
@@ -0,0 +1,565 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Unit tests for TieringOffloadingManager and ExampleSecondaryTier.
5 +
6 +These tests verify:
7 +1. Basic tiered offloading operations (store, load, lookup)
8 +2. Cascade behavior (blocks stored to all secondary tiers)
9 +3. Promotion behavior (blocks loaded from secondary to primary to GPU)
10 +4. ref_cnt management (blocks protected during async transfers)
11 +5. Eviction coordination between tiers
12 +"""
13 +
14 +from collections.abc import Iterable
15 +from unittest.mock import MagicMock
16 +
17 +import numpy as np
18 +import pytest
19 +import torch
20 +
21 +from vllm.v1.kv_offload.base import (
22 + OffloadKey,
23 + ReqContext,
24 + make_offload_key,
25 +)
26 +from vllm.v1.kv_offload.tiering.base import JobMetadata
27 +from vllm.v1.kv_offload.tiering.example import ExampleSecondaryTier
28 +from vllm.v1.kv_offload.tiering.manager import (
29 + CPUPrimaryTierOffloadingManager,
30 + TieringOffloadingManager,
31 +)
32 +
33 +_CTX = ReqContext(req_id="test")
34 +_MOCK_VLLM_CONFIG = MagicMock()
35 +
36 +
37 +def _mock_mmap_region(num_blocks: int, row_bytes: int = 16):
38 + """Create a mock SharedOffloadRegion for testing."""
39 + mock = MagicMock()
40 + view = memoryview(torch.zeros((num_blocks, row_bytes), dtype=torch.int8).numpy())
41 + mock.create_kv_memoryview.return_value = view
42 + return mock
43 +
44 +
45 +def to_keys(int_ids: Iterable[int]) -> list[OffloadKey]:
46 + return [make_offload_key(str(i).encode(), 0) for i in int_ids]
47 +
48 +
49 +def count_hits(manager, keys: list[OffloadKey]) -> int | None:
50 + """Count consecutive lookup hits from the start of keys.
51 +
52 + Returns the count of leading True results, or None if any lookup
53 + returns None (retry-later signal).
54 + """
55 + count = 0
56 + for key in keys:
57 + result = manager.lookup(key, _CTX)
58 + if result is None:
59 + return None
60 + if not result:
61 + break
62 + count += 1
63 + return count
64 +
65 +
66 +class TestExampleSecondaryTier:
67 + """Tests for ExampleSecondaryTier implementation."""
68 +
69 + def test_basic_store_and_lookup(self):
70 + """Test basic store and lookup operations."""
71 + mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy())
72 + tier = ExampleSecondaryTier(
73 + vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10
74 + )
75 +
76 + # Initially empty
77 + blocks = to_keys(range(3))
78 + assert tier.lookup(blocks[0], _CTX) is False
79 +
80 + # Store blocks (simulate with direct insertion for testing)
81 + tier.blocks[blocks[0]] = True
82 + tier.blocks[blocks[1]] = True
83 +
84 + # Lookup should find first two blocks
85 + assert tier.lookup(blocks[0], _CTX) is True
86 + assert tier.lookup(blocks[1], _CTX) is True
87 +
88 + # Third block not present
89 + assert tier.lookup(blocks[2], _CTX) is False
90 +
91 + def test_lru_eviction(self):
92 + """Test LRU eviction policy."""
93 + mock_view = memoryview(torch.zeros((4, 16), dtype=torch.int8).numpy())
94 + tier = ExampleSecondaryTier(
95 + vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=3
96 + )
97 +
98 + # Fill tier to capacity
99 + blocks = to_keys(range(3))
100 + for block in blocks:
101 + tier.blocks[block] = True
102 +
103 + assert tier.get_num_blocks() == 3
104 +
105 + # Touch first block (make it most recently used)
106 + tier.touch([blocks[0]], _CTX)
107 +
108 + # Store new block should evict blocks[1] (least recently used)
109 + new_block = to_keys([3])[0]
110 +
111 + tier.submit_store(
112 + JobMetadata(
113 + job_id=1,
114 + keys=[new_block],
115 + block_ids=np.array([0], dtype=np.int64),
116 + is_promotion=False,
117 + req_context=_CTX,
118 + )
119 + )
120 +
121 + # Complete the job
122 + tier.get_finished()
123 +
124 + # Verify new block is stored and blocks[1] was evicted (LRU)
125 + assert new_block in tier.blocks
126 + assert blocks[1] not in tier.blocks
127 + # blocks[0] and blocks[2] should still be present
128 + assert blocks[0] in tier.blocks
129 + assert blocks[2] in tier.blocks
130 +
131 + def test_async_simulation(self):
132 + """Test simulated async behavior."""
133 + mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy())
134 + tier = ExampleSecondaryTier(
135 + vllm_config=_MOCK_VLLM_CONFIG,
136 + primary_kv_view=mock_view,
137 + max_blocks=10,
138 + simulate_async=True,
139 + )
140 +
141 + blocks = to_keys(range(2))
142 +
143 + # Submit store job
144 + tier.submit_store(
145 + JobMetadata(
146 + job_id=1,
147 + keys=blocks,
148 + block_ids=np.array([0, 1], dtype=np.int64),
149 + is_promotion=False,
150 + req_context=_CTX,
151 + )
152 + )
153 +
154 + # Blocks should not yet be stored (pending async completion)
155 + assert tier.get_num_blocks() == 0
156 +
157 + # First get_finished() should complete the job
158 + completed = list(tier.get_finished())
159 + assert len(completed) == 1
160 + assert completed[0].job_id == 1
161 + assert completed[0].success is True
162 +
163 + # Blocks should now be stored
164 + assert tier.get_num_blocks() == 2
165 +
166 +
167 +class TestTieringOffloadingManager:
168 + """Tests for TieringOffloadingManager."""
169 +
170 + @pytest.fixture
171 + def manager_setup(self):
172 + # Create primary tier (CPU-based)
173 + mock_region = _mock_mmap_region(5)
174 + self.primary_tier = CPUPrimaryTierOffloadingManager(
175 + num_blocks=5, mmap_region=mock_region
176 + )
177 +
178 + mock_view = mock_region.create_kv_memoryview()
179 +
180 + # Create secondary tiers with the primary view
181 + self.secondary_tier1 = ExampleSecondaryTier(
182 + vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10
183 + )
184 + self.secondary_tier2 = ExampleSecondaryTier(
185 + vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10
186 + )
187 +
188 + # Create tiered manager
189 + self.manager = TieringOffloadingManager(
190 + primary_tier=self.primary_tier,
191 + secondary_tiers=[self.secondary_tier1, self.secondary_tier2],
192 + )
193 +
194 + def test_basic_store_to_primary(self, manager_setup):
195 + """Test basic store operation to primary tier."""
196 + blocks = to_keys(range(3))
197 +
198 + # Prepare store
199 + result = self.manager.prepare_store(blocks, _CTX)
200 + assert result is not None
201 + assert len(result.keys_to_store) == 3
202 +
203 + # Complete store
204 + self.manager.complete_store(blocks, _CTX, success=True)
205 +
206 + # Blocks should be in primary tier
207 + assert count_hits(self.primary_tier, blocks) == 3
208 +
209 + def test_cascade_to_all_secondary_tiers(self, manager_setup):
210 + """Test that blocks are cascaded to ALL secondary tiers."""
211 + blocks = to_keys(range(3))
212 +
213 + self.secondary_tier1.submit_store = MagicMock(
214 + wraps=self.secondary_tier1.submit_store
215 + )
216 + self.secondary_tier2.submit_store = MagicMock(
217 + wraps=self.secondary_tier2.submit_store
218 + )
219 +
220 + # Store to primary
221 + result = self.manager.prepare_store(blocks, _CTX)
222 + assert result is not None
223 +
224 + # Complete store (triggers cascade via submit_store on each tier)
225 + self.manager.complete_store(blocks, _CTX, success=True)
226 +
227 + # submit_store was called once per secondary tier
228 + self.secondary_tier1.submit_store.assert_called_once()
229 + self.secondary_tier2.submit_store.assert_called_once()
230 +
231 + # Blocks should be in both secondary tiers
232 + assert self.secondary_tier1.get_num_blocks() == 3
233 + assert self.secondary_tier2.get_num_blocks() == 3
234 +
235 + # Verify blocks are present
236 + assert all(self.secondary_tier1.lookup(b, _CTX) for b in blocks)
237 + assert all(self.secondary_tier2.lookup(b, _CTX) for b in blocks)
238 +
239 + def test_ref_cnt_protection_during_cascade(self, manager_setup):
240 + """Test that ref_cnt protects blocks during cascade."""
241 + blocks = to_keys(range(3))
242 +
243 + # Store to primary
244 + result = self.manager.prepare_store(blocks, _CTX)
245 + assert result is not None
246 + self.manager.complete_store(blocks, _CTX, success=True)
247 +
248 + # After complete_store, blocks should have ref_cnt > 0
249 + # (one for each secondary tier)
250 + for block_hash in blocks:
251 + block = self.primary_tier._policy.get(block_hash)
252 + # ref_cnt should be 2 (one for each secondary tier)
253 + assert block.ref_cnt == 2
254 +
255 + # End of step 1: _maybe_process_finished_jobs() was already called by
256 + # prepare_store() above (setting the per-step flag), so take_events()
257 + # does NOT poll get_finished() again — cascade completions remain
258 + # unprocessed until the next step.
259 + list(self.manager.take_events())
260 +
261 + # ref_cnt still held: cascade jobs finished (sync tier) but haven't
262 + # been polled yet because the per-step guard skipped the second call.
263 + for block_hash in blocks:
264 + block = self.primary_tier._policy.get(block_hash)
265 + assert block.ref_cnt == 2
266 +
267 + # Secondary tiers have completed jobs waiting to be drained
268 + assert len(self.secondary_tier1.completed_jobs) > 0
269 + assert len(self.secondary_tier2.completed_jobs) > 0
270 +
271 + # End of step 2: flag was reset, so _maybe_process_finished_jobs()
272 + # runs and processes the cascade completions (complete_read → ref_cnt--)
273 + list(self.manager.take_events())
274 +
275 + # After cascade completes, ref_cnt should be 0
276 + for block_hash in blocks:
277 + block = self.primary_tier._policy.get(block_hash)
278 + assert block.ref_cnt == 0
279 +
280 + # All completed jobs have been drained
281 + assert len(self.secondary_tier1.completed_jobs) == 0
282 + assert len(self.secondary_tier2.completed_jobs) == 0
283 +
284 + def test_lookup_from_primary(self, manager_setup):
285 + """Test lookup when blocks are in primary tier."""
286 + blocks = to_keys(range(3))
287 +
288 + # Store blocks
289 + self.manager.prepare_store(blocks, _CTX)
290 + self.manager.complete_store(blocks, _CTX, success=True)
291 +
292 + # Lookup should find all blocks in primary
293 + assert count_hits(self.manager, blocks) == 3
294 +
295 + def test_promotion_from_secondary(self, manager_setup):
296 + """Test promotion of blocks from secondary to primary tier."""
297 + blocks = to_keys(range(3))
298 +
299 + # Manually add blocks to secondary tier (simulate previous cascade)
300 + for block in blocks:
301 + self.secondary_tier1.blocks[block] = True
302 +
303 + # Lookup each block to initiate promotion for all of them
304 + for block in blocks:
305 + result = self.manager.lookup(block, _CTX)
306 + assert result is None # Retry later (promotion initiated)
307 +
308 + # End of step 1: flushes deferred submit_load() calls
309 + list(self.manager.take_events())
310 +
311 + # End of step 2: processes the completed promotion jobs
312 + list(self.manager.take_events())
313 +
314 + # Now blocks should be in primary tier
315 + assert count_hits(self.primary_tier, blocks) == 3
316 +
317 + # Next lookup should succeed
318 + assert count_hits(self.manager, blocks) == 3
319 +
320 + def test_partial_lookup(self, manager_setup):
321 + """Test lookup with partial hits."""
322 + blocks = to_keys(range(5))
323 +
324 + # Store first 3 blocks to primary
325 + self.manager.prepare_store(blocks[:3], _CTX)
326 + self.manager.complete_store(blocks[:3], _CTX, success=True)
327 +
328 + # Lookup all 5 blocks should return 3 (first 3 found)
329 + assert count_hits(self.manager, blocks) == 3
330 +
331 + def test_eviction_in_primary_tier(self, manager_setup):
332 + """Test eviction in primary tier when capacity is exceeded."""
333 + # Primary tier has capacity of 5 blocks
334 + # First, fill the primary tier
335 + blocks = to_keys(range(5))
336 + result = self.manager.prepare_store(blocks, _CTX)
337 + assert result is not None
338 + assert len(result.keys_to_store) == 5
339 + self.manager.complete_store(blocks, _CTX, success=True)
340 +
341 + # End of step: release ref_cnt from cascade
342 + list(self.manager.take_events())
343 +
344 + # Now try to store 2 more blocks (should trigger eviction)
345 + more_blocks = to_keys(range(5, 7))
346 + result = self.manager.prepare_store(more_blocks, _CTX)
347 +
348 + # Should evict 2 blocks from primary tier
349 + assert result is not None
350 + assert len(result.evicted_keys) == 2
351 + assert len(result.keys_to_store) == 2
352 +
353 + def test_touch_propagates_to_all_tiers(self, manager_setup):
354 + """Test that touch() propagates to all tiers."""
355 + blocks = to_keys(range(3))
356 +
357 + # Store blocks
358 + self.manager.prepare_store(blocks, _CTX)
359 + self.manager.complete_store(blocks, _CTX, success=True)
360 + list(self.manager.take_events())
361 +
362 + # Touch blocks
363 + self.manager.touch(blocks, _CTX)
364 +
365 + # Verify touch was called on primary tier (check LRU order)
366 + # In LRU, touched blocks should be at the end
367 + primary_keys = list(self.primary_tier._policy.blocks.keys())
368 + assert primary_keys[-3:] == list(reversed(blocks))
369 +
370 + # Verify touch was called on all secondary tiers
371 + secondary1_keys = list(self.secondary_tier1.blocks.keys())
372 + assert secondary1_keys[-3:] == list(reversed(blocks))
373 +
374 + secondary2_keys = list(self.secondary_tier2.blocks.keys())
375 + assert secondary2_keys[-3:] == list(reversed(blocks))
376 +
377 + def test_failed_store_no_cascade(self, manager_setup):
378 + """Test that failed GPU→primary store doesn't cascade."""
379 + blocks = to_keys(range(3))
380 +
381 + self.secondary_tier1.submit_store = MagicMock(
382 + wraps=self.secondary_tier1.submit_store
383 + )
384 + self.secondary_tier2.submit_store = MagicMock(
385 + wraps=self.secondary_tier2.submit_store
386 + )
387 +
388 + # Prepare store
389 + result = self.manager.prepare_store(blocks, _CTX)
390 + assert result is not None
391 +
392 + # Complete store with failure — cascade must not happen
393 + self.manager.complete_store(blocks, _CTX, success=False)
394 +
395 + # submit_store was never called on either secondary tier
396 + self.secondary_tier1.submit_store.assert_not_called()
397 + self.secondary_tier2.submit_store.assert_not_called()
398 +
399 + def test_multiple_secondary_tiers_independent_eviction(self):
400 + """Test that secondary tiers manage their own evictions."""
401 + mock_region = _mock_mmap_region(10)
402 + mock_view = mock_region.create_kv_memoryview()
403 +
404 + # Create tier with small capacity
405 + small_tier = ExampleSecondaryTier(
406 + vllm_config=_MOCK_VLLM_CONFIG,
407 + primary_kv_view=mock_view,
408 + max_blocks=5,
409 + simulate_async=False,
410 + )
411 + large_tier = ExampleSecondaryTier(
412 + vllm_config=_MOCK_VLLM_CONFIG,
413 + primary_kv_view=mock_view,
414 + max_blocks=10,
415 + simulate_async=False,
416 + )
417 +
418 + # Create a fresh primary tier for this test
419 + primary_tier = CPUPrimaryTierOffloadingManager(
420 + num_blocks=10, mmap_region=mock_region
421 + )
422 +
423 + manager = TieringOffloadingManager(
424 + primary_tier=primary_tier,
425 + secondary_tiers=[small_tier, large_tier],
426 + )
427 +
428 + # First, store 5 blocks to fill the small tier
429 + blocks1 = to_keys(range(5))
430 + result = manager.prepare_store(blocks1, _CTX)
431 + assert result is not None
432 + manager.complete_store(blocks1, _CTX, success=True)
433 + list(manager.take_events())
434 +
435 + # Both tiers should have 5 blocks
436 + assert small_tier.get_num_blocks() == 5
437 + assert large_tier.get_num_blocks() == 5
438 +
439 + # Now store 3 more blocks - small tier should evict 3 blocks
440 + blocks2 = to_keys(range(5, 8))
441 + result = manager.prepare_store(blocks2, _CTX)
442 + assert result is not None
443 + manager.complete_store(blocks2, _CTX, success=True)
444 + list(manager.take_events())
445 +
446 + # Small tier should still have 5 blocks (evicted 3, added 3)
447 + assert small_tier.get_num_blocks() == 5
448 +
449 + # Large tier should have all 8 blocks
450 + assert large_tier.get_num_blocks() == 8
451 +
452 + def test_lookup_batches_submit_load_per_request(self, manager_setup):
453 + """lookup() defers submit_load until take_events(), one call per request.
454 +
455 + Blocks from different requests each get their own submit_load call, each
456 + carrying the correct req_context.
457 + """
458 + blocks = to_keys(range(4))
459 + for block in blocks:
460 + self.secondary_tier1.blocks[block] = True
461 +
462 + self.secondary_tier1.submit_load = MagicMock(
463 + wraps=self.secondary_tier1.submit_load
464 + )
465 +
466 + ctx_a = ReqContext(req_id="req_a")
467 + ctx_b = ReqContext(req_id="req_b")
468 +
469 + # All lookups return None: secondary hit triggers promotion (in-flight)
470 + assert self.manager.lookup(blocks[0], ctx_a) is None
471 + assert self.manager.lookup(blocks[1], ctx_a) is None
472 + assert self.manager.lookup(blocks[2], ctx_b) is None
473 + assert self.manager.lookup(blocks[3], ctx_b) is None
474 +
475 + # submit_load must not fire during lookup - only at end of step
476 + self.secondary_tier1.submit_load.assert_not_called()
477 +
478 + # simulate end of step
479 + list(self.manager.take_events())
480 +
481 + assert self.secondary_tier1.submit_load.call_count == 2
482 + calls = self.secondary_tier1.submit_load.call_args_list
483 + jm_a = calls[0].args[0]
484 + jm_b = calls[1].args[0]
485 + assert set(jm_a.keys) == {blocks[0], blocks[1]}
486 + assert jm_a.req_context is ctx_a
487 + assert set(jm_b.keys) == {blocks[2], blocks[3]}
488 + assert jm_b.req_context is ctx_b
489 +
490 + def test_lookup_shared_block_no_duplicate_promotion(self, manager_setup):
491 + """A block looked up by two requests in the same step is promoted once.
492 +
493 + The first lookup initiates promotion (returns None via secondary hit).
494 + The second lookup sees ref_cnt=-1 on the primary slot and returns None
495 + via the primary in-flight path — without triggering a second promotion.
496 + """
497 + shared_block = to_keys([0])[0]
498 + self.secondary_tier1.blocks[shared_block] = True
499 +
500 + self.secondary_tier1.submit_load = MagicMock(
501 + wraps=self.secondary_tier1.submit_load
502 + )
503 +
504 + ctx_a = ReqContext(req_id="req_a")
505 + ctx_b = ReqContext(req_id="req_b")
506 +
507 + result_a = self.manager.lookup(shared_block, ctx_a)
508 + result_b = self.manager.lookup(shared_block, ctx_b)
509 +
510 + # Both see None (in-flight), but promotion is only queued once
511 + assert result_a is None
512 + assert result_b is None
513 +
514 + list(self.manager.take_events())
515 +
516 + # Only one submit_load call despite two lookups
517 + self.secondary_tier1.submit_load.assert_called_once()
518 + job_metadata = self.secondary_tier1.submit_load.call_args.args[0]
519 + assert list(job_metadata.keys) == [shared_block]
520 + assert job_metadata.req_context is ctx_a
521 +
522 + def test_complete_store_forwards_req_context_to_submit_store(self, manager_setup):
523 + """complete_store cascades to secondary tiers with the correct req_context."""
524 + blocks = to_keys(range(2))
525 +
526 + self.secondary_tier1.submit_store = MagicMock(
527 + wraps=self.secondary_tier1.submit_store
528 + )
529 +
530 + ctx = ReqContext(req_id="req_ctx", kv_transfer_params={"key": "value"})
531 +
532 + self.manager.prepare_store(blocks, ctx)
533 + self.manager.complete_store(blocks, ctx, success=True)
534 +
535 + assert self.secondary_tier1.submit_store.call_count == 1
536 + job_metadata = self.secondary_tier1.submit_store.call_args.args[0]
537 + assert job_metadata.req_context is ctx
538 +
539 +
540 +class TestTieringOffloadingWithoutSecondaryTiers:
541 + """Test TieringOffloadingManager with no secondary tiers (backward compat)."""
542 +
543 + def test_works_without_secondary_tiers(self):
544 + """Test that manager works with empty secondary_tiers list."""
545 + primary_tier = CPUPrimaryTierOffloadingManager(
546 + num_blocks=5, mmap_region=_mock_mmap_region(5)
547 + )
548 +
549 + # Create manager with no secondary tiers
550 + manager = TieringOffloadingManager(
551 + primary_tier=primary_tier, secondary_tiers=[]
552 + )
553 +
554 + blocks = to_keys(range(3))
555 +
556 + # Should work like a regular OffloadingManager
557 + result = manager.prepare_store(blocks, _CTX)
558 + assert result is not None
559 + manager.complete_store(blocks, _CTX, success=True)
560 +
561 + assert count_hits(manager, blocks) == 3
562 +
563 +
564 +if __name__ == "__main__":
565 + pytest.main([__file__, "-v"])
vllm/v1/kv_offload/base.py
+3 -5
@@ -4,8 +4,6 @@
4 Core abstractions for KV cache offloading in vLLM v1.
5 """
6
7 -from __future__ import annotations
8 -
7 from abc import ABC, abstractmethod
8 from collections.abc import Collection, Iterable, Iterator, Sequence
9 from dataclasses import dataclass
@@ -201,7 +199,7 @@ class OffloadingManager(ABC):
199 """
200 Marks blocks which were previously prepared to be stored, as stored.
201 Following this call, the blocks become loadable.
204 - If if_success is False, blocks that were not marked as stored will be
202 + If success is False, blocks that were not marked as stored will be
203 removed.
204
205 Args:
@@ -328,7 +326,7 @@ class CanonicalKVCaches:
326 class OffloadingSpec(ABC):
327 """Spec for an offloading connector"""
328
331 - def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
329 + def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"):
330 logger.warning(
331 "Initializing OffloadingSpec. This API is experimental and "
332 "subject to change in the future as we iterate the design."
@@ -394,7 +392,7 @@ class OffloadingSpec(ABC):
392 @abstractmethod
393 def get_handlers(
394 self, kv_caches: CanonicalKVCaches
397 - ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
395 + ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], "OffloadingHandler"]]:
396 """
397 Get offloading handlers along with their respective src and dst types.
398
vllm/v1/kv_offload/cpu/gpu_worker.py
+7
@@ -124,6 +124,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
124 block_size_factor: int,
125 kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]],
126 gpu_to_cpu: bool,
127 + mmap_region: SharedOffloadRegion | None = None,
128 ):
129 """
130 Initialize a SingleDirectionOffloadingHandler.
@@ -167,6 +168,8 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
168 self.dst_block_size_factor = block_size_factor if self.gpu_to_cpu else 1
169
170 self.transfer_type = ("GPU", "CPU") if self.gpu_to_cpu else ("CPU", "GPU")
171 + # mmap_region to clean up on shutdown (gpu_to_cpu handler owns it)
172 + self._mmap_region = mmap_region
173 # job_id -> event
174 self._transfer_events: dict[int, torch.Event] = {}
175 # queue of transfers (job_id, stream, event)
@@ -382,6 +385,9 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
385 self._event_pool.clear()
386 self.src_tensors.clear()
387 self.dst_tensors.clear()
388 + if self._mmap_region is not None:
389 + self._mmap_region.cleanup()
390 + self._mmap_region = None
391
392
393 class CpuGpuOffloadingHandlers:
@@ -434,6 +440,7 @@ class CpuGpuOffloadingHandlers:
440 block_size_factor=block_size_factor,
441 kv_cache_groups_data_refs=kv_caches.group_data_refs,
442 gpu_to_cpu=True,
443 + mmap_region=mmap_region,
444 )
445
446 self.cpu_to_gpu_handler = SingleDirectionOffloadingHandler(
vllm/v1/kv_offload/cpu/shared_offload_region.py
+14
@@ -153,6 +153,20 @@ class SharedOffloadRegion:
153 self._views.append(worker_layer_view)
154 return worker_layer_view
155
156 + def create_kv_memoryview(self) -> memoryview:
157 + """Return a zero-copy memoryview over the entire KV buffer.
158 +
159 + Shape: (num_blocks, row_stride_bytes). Secondary tiers address
160 + block *b* as ``view[b]``.
161 + """
162 + kv_tensor = self._base.view(self.num_blocks, self._row_stride)
163 + np_arr = kv_tensor.numpy()
164 + assert np_arr.ctypes.data == self._base.data_ptr(), (
165 + "view()/numpy() created a copy instead of sharing the mmap buffer; "
166 + "secondary tiers require zero-copy access to primary KV data"
167 + )
168 + return memoryview(np_arr)
169 +
170 def cleanup(self) -> None:
171 if self.is_pinned and self._base is not None:
172 base_ptr = self._base.data_ptr()
vllm/v1/kv_offload/cpu/spec.py
+14 -6
@@ -44,6 +44,10 @@ class CPUOffloadingSpec(OffloadingSpec):
44 if kv_bytes_per_offloaded_block > 0
45 else 0
46 )
47 + world_size = vllm_config.parallel_config.world_size
48 + self.cpu_page_size_per_worker: int = (
49 + kv_bytes_per_offloaded_block // world_size if world_size > 0 else 0
50 + )
51
52 # scheduler-side
53 self._manager: OffloadingManager | None = None
@@ -77,6 +81,15 @@ class CPUOffloadingSpec(OffloadingSpec):
81 )
82 return self._manager
83
84 + def _create_handlers(
85 + self, kv_caches: CanonicalKVCaches
86 + ) -> CpuGpuOffloadingHandlers:
87 + return CpuGpuOffloadingHandlers(
88 + kv_caches=kv_caches,
89 + block_size_factor=self.block_size_factor,
90 + num_cpu_blocks=self.num_blocks,
91 + )
92 +
93 def get_handlers(
94 self, kv_caches: CanonicalKVCaches
95 ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
@@ -85,12 +98,7 @@ class CPUOffloadingSpec(OffloadingSpec):
98 raise Exception(
99 "CPU Offloading is currently only supported on CUDA-alike GPUs"
100 )
88 -
89 - self._handlers = CpuGpuOffloadingHandlers(
90 - kv_caches=kv_caches,
91 - block_size_factor=self.block_size_factor,
92 - num_cpu_blocks=self.num_blocks,
93 - )
101 + self._handlers = self._create_handlers(kv_caches)
102
103 assert self._handlers is not None
104 yield GPULoadStoreSpec, CPULoadStoreSpec, self._handlers.gpu_to_cpu_handler
vllm/v1/kv_offload/factory.py
+5
@@ -56,3 +56,8 @@ class OffloadingSpecFactory:
56 OffloadingSpecFactory.register_spec(
57 "CPUOffloadingSpec", "vllm.v1.kv_offload.cpu.spec", "CPUOffloadingSpec"
58 )
59 +OffloadingSpecFactory.register_spec(
60 + "TieringOffloadingSpec",
61 + "vllm.v1.kv_offload.tiering.spec",
62 + "TieringOffloadingSpec",
63 +)
vllm/v1/kv_offload/tiering/__init__.py
vllm/v1/kv_offload/tiering/base.py new
+168
@@ -0,0 +1,168 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Abstract interfaces and data types for the secondary tiering layer.
5 +"""
6 +
7 +from abc import ABC, abstractmethod
8 +from collections.abc import Collection, Iterable
9 +from dataclasses import dataclass
10 +from typing import TYPE_CHECKING
11 +
12 +import numpy as np
13 +
14 +from vllm.v1.kv_offload.base import OffloadKey, ReqContext
15 +
16 +if TYPE_CHECKING:
17 + from vllm.config import VllmConfig
18 +
19 +# Type alias for job IDs used in async transfer tracking
20 +JobId = int
21 +
22 +
23 +@dataclass
24 +class JobMetadata:
25 + """Metadata for an in-flight async transfer job."""
26 +
27 + job_id: JobId
28 + keys: Collection[OffloadKey]
29 + block_ids: np.ndarray
30 + is_promotion: bool
31 + req_context: ReqContext
32 +
33 +
34 +@dataclass
35 +class JobResult:
36 + """Result of an async transfer job (successful or failed)."""
37 +
38 + job_id: JobId
39 + success: bool
40 +
41 +
42 +class SecondaryTierManager(ABC):
43 + """
44 + Abstract interface for managing a single non-primary offloading tier.
45 +
46 + Secondary tiers cannot directly access GPU memory. All data transfers
47 + must go through the CPU (primary) tier:
48 + - Store: GPU → CPU (primary) → secondary (cascade)
49 + - Load: secondary → CPU (primary) → GPU (promotion)
50 +
51 + IMPORTANT: All methods run in the Scheduler process and must be
52 + lightweight and non-blocking. submit_load() and submit_store() submit
53 + async jobs; get_finished() polls for completion.
54 + """
55 +
56 + def __init__(self, vllm_config: "VllmConfig", primary_kv_view: memoryview) -> None:
57 + self._vllm_config = vllm_config
58 + self._primary_kv_view: memoryview = primary_kv_view
59 +
60 + @abstractmethod
61 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
62 + """
63 + Check whether a block exists in this secondary tier.
64 +
65 + Args:
66 + key: Offload key to look up.
67 + req_context: per-request context (e.g. kv_transfer_params).
68 +
69 + Returns:
70 + True if the block is present and ready,
71 + False if not found,
72 + or None if the block is being transferred (retry later).
73 + """
74 + pass
75 +
76 + @abstractmethod
77 + def submit_store(self, job_metadata: JobMetadata) -> None:
78 + """
79 + Submit an async job to store blocks from the primary tier to this
80 + secondary tier.
81 +
82 + This method must be lightweight and non-blocking: allocate metadata
83 + and submit the transfer, but do NOT perform the data copy on the
84 + calling thread.
85 +
86 + Preconditions (guaranteed by the framework):
87 + - ``job_metadata.block_ids`` are valid primary-tier slots, pinned
88 + (ref-counted) for the duration of the transfer.
89 +
90 + The implementation is responsible for:
91 + 1. Filtering out blocks already present in this tier
92 + 2. Evicting blocks if capacity is needed
93 + 3. Allocating space in this tier
94 + 4. Submitting the async transfer (read from primary via block_ids)
95 +
96 + Report completion via ``get_finished()``.
97 +
98 + Args:
99 + job_metadata: Job metadata including job_id, keys, and block_ids
100 + identifying the primary-tier slots to read from.
101 + """
102 + pass
103 +
104 + @abstractmethod
105 + def submit_load(self, job_metadata: JobMetadata) -> None:
106 + """
107 + Submit an async job to load blocks from this secondary tier to the
108 + primary tier.
109 +
110 + This method must be lightweight and non-blocking: mark blocks as
111 + in-flight and submit the transfer, but do NOT perform the data copy
112 + on the calling thread.
113 +
114 + Preconditions (guaranteed by the framework):
115 + - ``job_metadata.block_ids`` are allocated primary-tier slots
116 + ready to receive data.
117 +
118 + The implementation must copy data from this tier into the
119 + primary-tier slots identified by ``block_ids``.
120 +
121 + Report completion via ``get_finished()``.
122 +
123 + Args:
124 + job_metadata: Job metadata including job_id, keys, and block_ids
125 + identifying the primary-tier slots to write into.
126 + """
127 + pass
128 +
129 + @abstractmethod
130 + def get_finished(self) -> Iterable[JobResult]:
131 + """
132 + Return all jobs (loads and stores) that completed since the last call.
133 +
134 + The framework uses these results to release resources and finalize
135 + transfers.
136 +
137 + Returns:
138 + Iterable of JobResult objects for jobs finished since the
139 + last call.
140 + """
141 + pass
142 +
143 + def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
144 + """
145 + Mark blocks as recently used for eviction policy.
146 +
147 + Args:
148 + keys: Offload keys to mark as recently used.
149 + req_context: Per-request context.
150 + """
151 + return
152 +
153 + def shutdown(self) -> None:
154 + """Release resources held by this tier (threads, connections, etc.)."""
155 + return
156 +
157 + @staticmethod
158 + @abstractmethod
159 + def get_tier_type() -> str:
160 + """
161 + Get the type identifier of this tier (e.g., "example", "storage").
162 +
163 + Must match the "type" field in the tier config dict.
164 +
165 + Returns:
166 + Tier type string.
167 + """
168 + pass
vllm/v1/kv_offload/tiering/example/__init__.py new
+235
@@ -0,0 +1,235 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +ExampleSecondaryTier: A simple in-memory secondary tier for testing.
5 +
6 +This implementation provides a minimal secondary tier that stores blocks
7 +in memory (using a dictionary) and simulates async transfers with immediate
8 +completion. It's useful for testing the TieringOffloadingManager without
9 +requiring actual storage or network backends.
10 +"""
11 +
12 +from collections import OrderedDict
13 +from collections.abc import Collection, Iterable
14 +from dataclasses import dataclass
15 +from typing import TYPE_CHECKING
16 +
17 +from vllm.v1.kv_offload.base import OffloadKey, ReqContext
18 +from vllm.v1.kv_offload.tiering.base import (
19 + JobId,
20 + JobMetadata,
21 + JobResult,
22 + SecondaryTierManager,
23 +)
24 +
25 +if TYPE_CHECKING:
26 + from vllm.config import VllmConfig
27 +
28 +
29 +@dataclass
30 +class _JobMetadata:
31 + """Internal metadata for tracking job details."""
32 +
33 + job_id: JobId
34 + keys: Collection[OffloadKey]
35 + is_store: bool # True for store jobs, False for load jobs
36 +
37 +
38 +class ExampleSecondaryTier(SecondaryTierManager):
39 + """
40 + A simple in-memory secondary tier for testing.
41 +
42 + This implementation:
43 + - Stores blocks in a dictionary (key -> True)
44 + - Simulates async transfers with immediate completion
45 + - Uses LRU eviction policy
46 + """
47 +
48 + def __init__(
49 + self,
50 + vllm_config: "VllmConfig",
51 + primary_kv_view: memoryview,
52 + max_blocks: int = 1000,
53 + simulate_async: bool = False,
54 + ):
55 + """
56 + Initialize the example secondary tier.
57 +
58 + Args:
59 + vllm_config: Global vLLM configuration.
60 + primary_kv_view: Memoryview of the primary tier's CPU KV cache.
61 + max_blocks: Maximum number of blocks this tier can store
62 + simulate_async: If True, jobs complete on next get_finished() call.
63 + If False, jobs complete immediately.
64 + """
65 + super().__init__(vllm_config, primary_kv_view)
66 + self.max_blocks = max_blocks
67 + self.simulate_async = simulate_async
68 +
69 + # key -> True (only care about presence)
70 + self.blocks: OrderedDict[OffloadKey, bool] = OrderedDict()
71 +
72 + # Completed jobs waiting to be retrieved by get_finished()
73 + self.completed_jobs: list[JobResult] = []
74 +
75 + # Pending jobs (for simulated async mode)
76 + self.pending_jobs: list[_JobMetadata] = []
77 +
78 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
79 + """
80 + Check whether a block exists in this secondary tier.
81 +
82 + Args:
83 + key: Offload key to look up.
84 + req_context: Per-request context.
85 +
86 + Returns:
87 + True if the block is present, False if not found.
88 + """
89 + return key in self.blocks
90 +
91 + def submit_store(self, job_metadata: JobMetadata) -> None:
92 + """
93 + Submit an async job to store blocks from primary tier to this tier.
94 +
95 + Args:
96 + job_metadata: Job metadata including job_id, keys, and
97 + spec for reading blocks from the primary tier.
98 + """
99 + job_id = job_metadata.job_id
100 + keys = job_metadata.keys
101 + block_ids = job_metadata.block_ids
102 +
103 + assert len(keys) == len(block_ids), (
104 + f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids"
105 + )
106 +
107 + # Filter out blocks already present
108 + blocks_to_store = [bh for bh in keys if bh not in self.blocks]
109 +
110 + if not blocks_to_store:
111 + # All blocks already present
112 + return
113 +
114 + # Evict blocks if needed (LRU policy)
115 + num_blocks_to_evict = len(blocks_to_store) - (
116 + self.max_blocks - len(self.blocks)
117 + )
118 +
119 + evicted = []
120 + if num_blocks_to_evict > 0:
121 + # Collect eviction candidates first (LRU order), then delete atomically
122 + protected = set(keys)
123 + for key in self.blocks:
124 + if key not in protected:
125 + evicted.append(key)
126 + if len(evicted) == num_blocks_to_evict:
127 + break
128 + else:
129 + # Could not collect enough eviction candidates
130 + return
131 + for key in evicted:
132 + del self.blocks[key]
133 +
134 + # Create internal job metadata
135 + internal_job_metadata = _JobMetadata(
136 + job_id=job_id, keys=blocks_to_store, is_store=True
137 + )
138 +
139 + if self.simulate_async:
140 + # Job will complete on next get_finished() call
141 + self.pending_jobs.append(internal_job_metadata)
142 + else:
143 + # Job completes immediately
144 + self._complete_store_job(internal_job_metadata)
145 +
146 + def submit_load(self, job_metadata: JobMetadata) -> None:
147 + """
148 + Submit an async job to load blocks from this tier to primary tier.
149 +
150 + Args:
151 + job_metadata: Job metadata including job_id, keys, and
152 + spec for writing blocks into the primary tier.
153 + """
154 + job_id = job_metadata.job_id
155 + keys = job_metadata.keys
156 + block_ids = job_metadata.block_ids
157 +
158 + assert len(keys) == len(block_ids), (
159 + f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids"
160 + )
161 +
162 + # Verify all blocks exist
163 + for key in keys:
164 + if key not in self.blocks:
165 + return
166 +
167 + # Create internal job metadata
168 + internal_job_metadata = _JobMetadata(job_id=job_id, keys=keys, is_store=False)
169 +
170 + if self.simulate_async:
171 + # Job will complete on next get_finished() call
172 + self.pending_jobs.append(internal_job_metadata)
173 + else:
174 + # Job completes immediately
175 + self._complete_load_job(internal_job_metadata)
176 +
177 + def get_finished(self) -> Iterable[JobResult]:
178 + """
179 + Poll for finished async jobs.
180 +
181 + Returns:
182 + Iterable of JobResult objects for all jobs that have
183 + finished since the last call.
184 + """
185 + # Move pending jobs to completed
186 + if self.simulate_async and self.pending_jobs:
187 + for job_metadata in self.pending_jobs:
188 + if job_metadata.is_store:
189 + self._complete_store_job(job_metadata)
190 + else:
191 + self._complete_load_job(job_metadata)
192 + self.pending_jobs.clear()
193 +
194 + # Return completed jobs
195 + result = self.completed_jobs
196 + self.completed_jobs = []
197 + return result
198 +
199 + def _complete_store_job(self, job_metadata: _JobMetadata):
200 + """Complete a store job by adding blocks to storage."""
201 + for key in job_metadata.keys:
202 + self.blocks[key] = True
203 + # Return simplified JobResult (only job_id and success)
204 + self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))
205 +
206 + def _complete_load_job(self, job_metadata: _JobMetadata):
207 + """Complete a load job."""
208 + # Return simplified JobResult (only job_id and success)
209 + self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))
210 +
211 + def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
212 + """
213 + Mark blocks as recently used (move to end of LRU list).
214 +
215 + Args:
216 + keys: Blocks to mark as recently used.
217 + req_context: Per-request context.
218 + """
219 + for key in reversed(list(keys)):
220 + if key in self.blocks:
221 + self.blocks.move_to_end(key)
222 +
223 + @staticmethod
224 + def get_tier_type() -> str:
225 + return "example"
226 +
227 + def get_num_blocks(self) -> int:
228 + """Get the number of blocks currently stored in this tier."""
229 + return len(self.blocks)
230 +
231 + def clear(self):
232 + """Clear all blocks and pending jobs (for testing)."""
233 + self.blocks.clear()
234 + self.completed_jobs.clear()
235 + self.pending_jobs.clear()
vllm/v1/kv_offload/tiering/factory.py new
+56
@@ -0,0 +1,56 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Factory for creating secondary tier implementations.
5 +"""
6 +
7 +from typing import TYPE_CHECKING
8 +
9 +from vllm.v1.kv_offload.tiering.base import SecondaryTierManager
10 +from vllm.v1.kv_offload.tiering.example import ExampleSecondaryTier
11 +
12 +if TYPE_CHECKING:
13 + from vllm.config import VllmConfig
14 +
15 +SUPPORTED_TIERS: tuple[type[SecondaryTierManager], ...] = (ExampleSecondaryTier,)
16 +
17 +_TIER_REGISTRY: dict[str, type[SecondaryTierManager]] = {
18 + cls.get_tier_type(): cls for cls in SUPPORTED_TIERS
19 +}
20 +
21 +
22 +def create_secondary_tier(
23 + tier_config: dict,
24 + primary_kv_view: memoryview,
25 + vllm_config: "VllmConfig",
26 +) -> SecondaryTierManager:
27 + """
28 + Create a secondary tier from configuration.
29 +
30 + Args:
31 + tier_config: Dictionary with tier configuration containing:
32 + - type (required): Type of secondary tier (e.g., "example")
33 + - Additional tier-specific parameters are passed directly
34 + to the tier constructor
35 + primary_kv_view: Memoryview of the primary tier's CPU KV cache.
36 + vllm_config: Global vLLM configuration.
37 +
38 + Returns:
39 + SecondaryTierManager instance
40 +
41 + Raises:
42 + ValueError: If tier type is unknown or configuration is invalid
43 + """
44 + config = tier_config.copy()
45 +
46 + tier_type = config.pop("type", None)
47 + if not tier_type:
48 + raise ValueError("Secondary tier configuration must include 'type'")
49 +
50 + cls = _TIER_REGISTRY.get(tier_type)
51 + if cls is None:
52 + raise ValueError(
53 + f"Unknown secondary tier type: {tier_type!r}. "
54 + f"Supported types: {list(_TIER_REGISTRY)}"
55 + )
56 + return cls(vllm_config=vllm_config, primary_kv_view=primary_kv_view, **config)
vllm/v1/kv_offload/tiering/manager.py new
+505
@@ -0,0 +1,505 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +TieringOffloadingManager: Multi-tier KV cache offloading orchestrator.
5 +
6 +This manager coordinates between a CPU primary tier (with direct GPU access)
7 +and zero or more secondary tiers (Storage, Network, etc.) to provide
8 +hierarchical KV cache offloading.
9 +
10 +Key Design Principles:
11 +1. Always offload to all tiers — When a block is stored to the primary tier,
12 + it is cascaded to ALL secondary tiers
13 +2. Primary tier is the gateway — Secondary tiers cannot access GPU memory
14 + directly; all data flows through the CPU primary tier
15 +3. Staged promotion — Blocks in secondary tiers must be promoted to the
16 + primary tier before GPU can access them
17 +4. Transparent retry mechanism — Return None from lookup() to signal
18 + "data is being promoted, try later"
19 +5. ref_cnt as eviction protection — primary.prepare_read() increments ref_cnt,
20 + protecting blocks from eviction until complete_read() is called
21 +"""
22 +
23 +from collections.abc import Collection, Iterable
24 +from dataclasses import dataclass, field
25 +
26 +import numpy as np
27 +
28 +from vllm.logger import init_logger
29 +from vllm.v1.kv_offload.base import (
30 + LoadStoreSpec,
31 + OffloadingEvent,
32 + OffloadingManager,
33 + OffloadKey,
34 + PrepareStoreOutput,
35 + ReqContext,
36 +)
37 +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
38 +from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
39 +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
40 +from vllm.v1.kv_offload.tiering.base import (
41 + JobId,
42 + JobMetadata,
43 + SecondaryTierManager,
44 +)
45 +
46 +logger = init_logger(__name__)
47 +
48 +
49 +@dataclass
50 +class PendingPromotion:
51 + """Accumulator for blocks awaiting submit_load() for one (tier, request)."""
52 +
53 + req_context: ReqContext
54 + keys: list[OffloadKey] = field(default_factory=list)
55 + block_ids: list[int] = field(default_factory=list)
56 +
57 +
58 +class CPUPrimaryTierOffloadingManager(CPUOffloadingManager):
59 + """CPUOffloadingManager with a primary/secondary transfer interface.
60 +
61 + The inherited prepare_store/complete_store/prepare_load/complete_load are the
62 + GPU-facing OffloadingManager interface. These aliases expose the same operations
63 + from the secondary tier perspective, where read/write refers to secondary
64 + accessing primary. This avoids confusion when reading TieringOffloadingManager
65 + code (e.g. calling prepare_load inside a cascade/store path would be misleading).
66 + """
67 +
68 + def __init__(
69 + self,
70 + num_blocks: int,
71 + mmap_region: SharedOffloadRegion,
72 + cache_policy: str = "lru",
73 + enable_events: bool = False,
74 + ):
75 + super().__init__(
76 + num_blocks=num_blocks,
77 + cache_policy=cache_policy, # type: ignore[arg-type]
78 + enable_events=enable_events,
79 + )
80 + self._mmap_region = mmap_region
81 + # read/write is for CPU<->secondary transfers,
82 + # load/store is for CPU<->GPU transfers.
83 + # These aliases avoid calling prepare_load inside a store path.
84 + self.prepare_read = self.prepare_load
85 + self.complete_read = self.complete_load
86 + self.prepare_write = self.prepare_store
87 + self.complete_write = self.complete_store
88 +
89 + self._kv_memoryview = mmap_region.create_kv_memoryview()
90 +
91 + def get_kv_memoryview(self) -> memoryview:
92 + """Return the memoryview over the primary tier's KV cache buffer.
93 +
94 + The view has shape (num_blocks, row_stride_bytes) and is backed by the
95 + SharedOffloadRegion mmap. Secondary tiers address block *b* as
96 + ``view[b]``.
97 + """
98 + return self._kv_memoryview
99 +
100 + def shutdown(self) -> None:
101 + super().shutdown()
102 + self._kv_memoryview.release()
103 + self._mmap_region.cleanup()
104 +
105 +
106 +class TieringOffloadingManager(OffloadingManager):
107 + """
108 + Orchestrates multi-tier KV cache offloading.
109 +
110 + This manager coordinates between a CPU primary tier (with direct GPU access)
111 + and zero or more secondary tiers (Storage, Network, etc.) to provide
112 + hierarchical KV cache offloading.
113 +
114 + Key internal state:
115 + - Minimal state tracking; relies on secondary tiers to report completion
116 + via get_finished()
117 + - Secondary tiers return JobResult objects containing all necessary
118 + information
119 + - job_id_counter: monotonically increasing counter for job IDs
120 + """
121 +
122 + def __init__(
123 + self,
124 + primary_tier: CPUPrimaryTierOffloadingManager,
125 + secondary_tiers: list[SecondaryTierManager] | None = None,
126 + enable_events: bool = False,
127 + ):
128 + """
129 + Initialize the TieringOffloadingManager.
130 +
131 + Args:
132 + primary_tier: The primary tier manager (CPU-based).
133 + secondary_tiers: List of secondary tier managers (e.g., Storage,
134 + Network). Can be None or empty list.
135 + enable_events: Whether to track offloading events
136 + """
137 + self.primary_tier: CPUPrimaryTierOffloadingManager = primary_tier
138 + self.secondary_tiers = secondary_tiers or []
139 +
140 + self._job_id_counter: int = 0
141 + self.events: list[OffloadingEvent] | None = [] if enable_events else None
142 +
143 + # Job tracking: maps job_id to metadata for all in-flight transfers.
144 + # JobMetadata.is_promotion distinguishes direction:
145 + # True: secondary → primary (promotion)
146 + # False: primary → secondary (cascade)
147 + self._transfer_jobs: dict[JobId, JobMetadata] = {}
148 +
149 + # Pending promotion requests accumulated during lookup() calls; flushed
150 + # as one batched submit_load() per (tier, request) in take_events().
151 + # Outer key: tier. Inner key: req_context.req_id — the same ReqContext
152 + # object is reused for all block lookups of a given request per engine step.
153 + self._pending_load_submissions: dict[
154 + SecondaryTierManager, dict[str, PendingPromotion]
155 + ] = {}
156 +
157 + # Gate for once-per-step execution of _maybe_process_finished_jobs().
158 + # Reset at the end of each step in take_events().
159 + self._processed_jobs_this_step: bool = False
160 +
161 + def _next_job_id(self) -> JobId:
162 + """Generate a unique job ID for async transfer tracking."""
163 + job_id = self._job_id_counter
164 + self._job_id_counter += 1
165 + return job_id
166 +
167 + def _maybe_process_finished_jobs(self):
168 + """
169 + Poll secondary tiers for completed jobs (at most once per step).
170 +
171 + Guarded by _processed_jobs_this_step: the first call in an engine step
172 + does the actual polling; subsequent calls are no-ops. The flag is reset
173 + in take_events() at the end of each step.
174 + """
175 + if self._processed_jobs_this_step:
176 + return
177 + self._processed_jobs_this_step = True
178 + self._process_finished_jobs()
179 +
180 + def _process_finished_jobs(self):
181 + """
182 + Unconditionally poll all secondary tiers for completed jobs.
183 +
184 + This method:
185 + 1. Calls get_finished() on each secondary tier
186 + 2. For completed stores (primary→secondary): calls primary.complete_read()
187 + to decrement ref_cnt
188 + 3. For completed loads (secondary→primary): calls primary.complete_write()
189 + to make blocks available
190 + """
191 + for i, tier in enumerate(self.secondary_tiers):
192 + for completed_job in tier.get_finished():
193 + job_id = completed_job.job_id
194 + job_metadata = self._transfer_jobs.pop(job_id, None)
195 + assert job_metadata is not None, (
196 + f"Finished job_id {job_id} from tier #{i}"
197 + f" ({tier.get_tier_type()}) not in _transfer_jobs"
198 + )
199 +
200 + if job_metadata.is_promotion:
201 + # secondary→primary transfer (promotion) completed.
202 + # Make blocks available in primary tier.
203 + self.primary_tier.complete_write(
204 + job_metadata.keys,
205 + job_metadata.req_context,
206 + completed_job.success,
207 + )
208 + else:
209 + # primary→secondary transfer completed.
210 + # Decrement ref_cnt on primary blocks.
211 + self.primary_tier.complete_read(
212 + job_metadata.keys, job_metadata.req_context
213 + )
214 +
215 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
216 + """
217 + Check whether a single block is offloaded and ready.
218 +
219 + Algorithm:
220 + 1. Process any completed async jobs first.
221 + 2. Query primary tier — short-circuit on hit or in-flight.
222 + 3. On primary miss, query secondary tiers — stop on first
223 + hit and initiate promotion.
224 +
225 + Args:
226 + key: Block hash to look up.
227 + req_context: Per-request context.
228 +
229 + Returns:
230 + True — block is ready in the primary tier.
231 + None — block found but not yet ready (primary in-flight,
232 + promotion started, or a secondary tier is busy).
233 + False — block not found in any tier, or primary is full
234 + and cannot accept a promotion.
235 + """
236 + self._maybe_process_finished_jobs()
237 +
238 + primary_hit = self.primary_tier.lookup(key, req_context)
239 + if primary_hit is True:
240 + return True
241 + if primary_hit is None:
242 + return None
243 +
244 + any_none = False
245 + for tier in self.secondary_tiers:
246 + result = tier.lookup(key, req_context)
247 + if result is True:
248 + if not self._initiate_promotion(tier, key, req_context):
249 + return False # primary full, block unavailable
250 + return None # promotion started, retry later
251 + if result is None:
252 + any_none = True
253 +
254 + if any_none:
255 + return None
256 + return False
257 +
258 + def _initiate_promotion(
259 + self,
260 + tier: SecondaryTierManager,
261 + key: OffloadKey,
262 + req_context: ReqContext,
263 + ) -> bool:
264 + """
265 + Queue a block for promotion from a secondary tier to the primary tier.
266 +
267 + Allocates space in the primary tier immediately (sets ref_cnt=-1 so
268 + subsequent lookups within the same step see the slot as in-flight),
269 + then defers the actual submit_load() call to _flush_pending_promotions()
270 + so all blocks queued during one engine step are submitted as a single
271 + batched job.
272 +
273 + Args:
274 + tier: The secondary tier to promote from
275 + key: Block to promote
276 + req_context: Per-request context forwarded to primary.prepare_write().
277 +
278 + Returns:
279 + True if promotion was initiated, False if primary tier is full.
280 + """
281 + # Allocate space in primary tier for promoted block.
282 + # Must happen immediately so primary.lookup() returns None (in-flight)
283 + # for this key on any subsequent lookup() call within the same step,
284 + # preventing duplicate promotion attempts.
285 + primary_write_result = self.primary_tier.prepare_write([key], req_context)
286 +
287 + if primary_write_result is None:
288 + # Primary tier is full; caller should treat the block as unavailable
289 + # rather than retrying indefinitely.
290 + return False
291 +
292 + store_spec = primary_write_result.store_spec
293 + assert isinstance(store_spec, CPULoadStoreSpec)
294 + # Defer submit_load to take_events(). Group by (tier, request) so each
295 + # request's blocks are submitted as one batched job per tier.
296 + tier_pending = self._pending_load_submissions.setdefault(tier, {})
297 + ctx_id = req_context.req_id
298 + if ctx_id not in tier_pending:
299 + tier_pending[ctx_id] = PendingPromotion(
300 + keys=[], block_ids=[], req_context=req_context
301 + )
302 + entry = tier_pending[ctx_id]
303 + entry.keys.extend(primary_write_result.keys_to_store)
304 + entry.block_ids.extend(store_spec.block_ids)
305 + return True
306 +
307 + def _flush_pending_promotions(self) -> None:
308 + """Submit one batched submit_load() per (tier, request).
309 +
310 + Called from take_events() at the end of each engine step, flushing
311 + all promotion requests deferred during lookup().
312 + """
313 + if not self._pending_load_submissions:
314 + return
315 +
316 + for tier, pending_by_ctx in self._pending_load_submissions.items():
317 + for entry in pending_by_ctx.values():
318 + job_id = self._next_job_id()
319 + job_metadata = JobMetadata(
320 + job_id=job_id,
321 + keys=entry.keys,
322 + block_ids=np.array(entry.block_ids, dtype=np.int64),
323 + is_promotion=True,
324 + req_context=entry.req_context,
325 + )
326 + self._transfer_jobs[job_id] = job_metadata
327 + tier.submit_load(job_metadata)
328 +
329 + self._pending_load_submissions.clear()
330 +
331 + def prepare_load(
332 + self, keys: Collection[OffloadKey], req_context: ReqContext
333 + ) -> LoadStoreSpec:
334 + """
335 + Prepare blocks to be loaded from primary tier to GPU.
336 +
337 + CRITICAL: This method calls _maybe_process_finished_jobs() FIRST to ensure
338 + that any completed promotions have been finalized and blocks are ready.
339 +
340 + This increments ref_cnt on the blocks in the primary tier, protecting
341 + them from eviction during the transfer.
342 +
343 + Args:
344 + keys: Blocks to prepare for loading.
345 + req_context: Per-request context.
346 +
347 + Returns:
348 + LoadStoreSpec for reading from primary tier.
349 + """
350 + # Process completed promotions to ensure blocks are ready
351 + self._maybe_process_finished_jobs()
352 +
353 + return self.primary_tier.prepare_load(keys, req_context)
354 +
355 + def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
356 + """
357 + Mark blocks as recently used in all tiers.
358 +
359 + Args:
360 + keys: Blocks to mark as recently used.
361 + req_context: Per-request context.
362 + """
363 + self.primary_tier.touch(keys, req_context)
364 + for tier in self.secondary_tiers:
365 + tier.touch(keys, req_context)
366 +
367 + def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext):
368 + """
369 + Mark blocks as done loading from primary tier to GPU.
370 +
371 + This decrements ref_cnt on the blocks in the primary tier, allowing
372 + them to be evicted again.
373 +
374 + Args:
375 + keys: Blocks that finished loading.
376 + req_context: Per-request context.
377 + """
378 + self.primary_tier.complete_load(keys, req_context)
379 +
380 + def prepare_store(
381 + self, keys: Collection[OffloadKey], req_context: ReqContext
382 + ) -> PrepareStoreOutput | None:
383 + """
384 + Prepare blocks to be stored from GPU to primary tier.
385 +
386 + CRITICAL: This method calls _maybe_process_finished_jobs() FIRST to ensure
387 + that any completed async transfers have their ref_cnt decremented
388 + before the primary tier makes eviction decisions.
389 +
390 + Args:
391 + keys: Blocks to prepare for storing.
392 + req_context: Per-request context.
393 +
394 + Returns:
395 + PrepareStoreOutput describing where to store blocks and what was
396 + evicted, or None if store cannot proceed.
397 + """
398 + # Step 1: Poll for completed async jobs FIRST
399 + # This decrements ref_cnt on primary blocks that have been
400 + # successfully transferred to secondary tiers.
401 + self._maybe_process_finished_jobs()
402 +
403 + # Step 2: Store to primary tier
404 + primary_result = self.primary_tier.prepare_store(keys, req_context)
405 +
406 + # Note: Secondary tier cascading will happen in complete_store()
407 + # after the GPU→Primary transfer completes and blocks are ready.
408 +
409 + return primary_result
410 +
411 + def complete_store(
412 + self,
413 + keys: Collection[OffloadKey],
414 + req_context: ReqContext,
415 + success: bool = True,
416 + ):
417 + """
418 + Mark blocks as done storing from GPU to primary tier.
419 +
420 + This is where secondary tier cascading happens — after blocks are
421 + confirmed to be in the primary tier, they are cascaded to ALL
422 + secondary tiers.
423 +
424 + For each secondary tier:
425 + 1. Call primary.prepare_read() to get LoadStoreSpec AND increment
426 + ref_cnt (protecting blocks during async transfer)
427 + 2. Call tier.submit_store() to start async transfer: primary→secondary
428 + 3. Track the job in _store_jobs dictionary
429 +
430 + Args:
431 + keys: Blocks that finished storing.
432 + success: Whether the GPU→primary transfer succeeded.
433 + req_context: Per-request context forwarded to primary.prepare_read().
434 + """
435 + # Step 1: Complete store in primary tier (makes blocks loadable)
436 + self.primary_tier.complete_store(keys, req_context, success)
437 +
438 + if not success:
439 + # If GPU→Primary transfer failed, don't cascade to secondary tiers
440 + return
441 +
442 + # Step 2: Cascade to ALL secondary tiers
443 + # For each secondary tier, call primary.prepare_read() to get the
444 + # LoadStoreSpec AND to increment ref_cnt (protecting blocks from
445 + # eviction during the async transfer). One prepare_read() call per
446 + # secondary tier.
447 + for tier in self.secondary_tiers:
448 + primary_blocks_spec = self.primary_tier.prepare_read(keys, req_context)
449 +
450 + # Submit async store job: primary→secondary
451 + job_id = self._next_job_id()
452 +
453 + # Track this store job
454 + assert isinstance(primary_blocks_spec, CPULoadStoreSpec)
455 + job_metadata = JobMetadata(
456 + job_id=job_id,
457 + keys=keys,
458 + block_ids=primary_blocks_spec.block_ids,
459 + is_promotion=False,
460 + req_context=req_context,
461 + )
462 + self._transfer_jobs[job_id] = job_metadata
463 +
464 + tier.submit_store(job_metadata)
465 +
466 + # Note: The async transfers are now in flight. Their completion is
467 + # tracked via get_finished() / _maybe_process_finished_jobs().
468 +
469 + def take_events(self) -> Iterable[OffloadingEvent]:
470 + """
471 + End-of-step hook: flush deferred work, yield events, reset per-step state.
472 +
473 + Called once per engine step from Scheduler.update_from_output() →
474 + connector.take_events(). Ensures _maybe_process_finished_jobs() has run
475 + at least once this step, flushes pending promotions, yields collected
476 + events, and resets the per-step flag.
477 +
478 + Yields:
479 + New OffloadingEvents collected since the last call.
480 + """
481 + # TODO: Move _flush_pending_promotions() to a dedicated end_of_batch()
482 + # hook once one exists. For now, take_events() serves as the flush
483 + # point under the assumption that it is called at the end of each
484 + # engine step (Scheduler.update_from_output() → connector.take_events()).
485 + # When the dedicated hook is added, update tests that rely on
486 + # take_events() to signal end of step.
487 +
488 + self._maybe_process_finished_jobs()
489 +
490 + self._flush_pending_promotions()
491 +
492 + # Reset the per-step gate so next step's first call does real work.
493 + self._processed_jobs_this_step = False
494 +
495 + if self.events is not None:
496 + yield from self.events
497 + self.events.clear()
498 +
499 + yield from self.primary_tier.take_events()
500 +
501 + def shutdown(self) -> None:
502 + """Shutdown all tiers and release resources."""
503 + for tier in self.secondary_tiers:
504 + tier.shutdown()
505 + self.primary_tier.shutdown()
vllm/v1/kv_offload/tiering/spec.py new
+187
@@ -0,0 +1,187 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +TieringOffloadingSpec: Spec for multi-tier KV cache offloading.
5 +
6 +This spec creates a TieringOffloadingManager with a CPU primary tier
7 +and configurable secondary tiers (e.g., Storage, Network).
8 +
9 +Configuration via kv_connector_extra_config:
10 + - cpu_bytes_to_use: (required) Bytes to allocate for CPU primary tier
11 + - block_size: (optional) Block size for offloaded blocks (default: GPU block size)
12 + - eviction_policy: (optional) Primary tier eviction policy: "lru" or
13 + "arc" (default: "lru")
14 + - secondary_tiers: (optional) List of secondary tier configurations
15 + Each secondary tier config is a dict with:
16 + - type: (required) Type of secondary tier (e.g., "example", "storage", "network")
17 + - Additional tier-specific parameters are passed directly to the tier
18 + constructor. See each tier's documentation for supported parameters.
19 +
20 +Example configuration:
21 +{
22 + "cpu_bytes_to_use": 10737418240, # 10 GB
23 + "block_size": 16,
24 + "eviction_policy": "lru",
25 + "secondary_tiers": [
26 + {
27 + "type": "example",
28 + # Tier-specific parameters (for ExampleSecondaryTier):
29 + "max_blocks": 10000,
30 + "simulate_async": False
31 + }
32 + ]
33 +}
34 +"""
35 +
36 +import torch
37 +
38 +from vllm.config import VllmConfig
39 +from vllm.logger import init_logger
40 +from vllm.v1.kv_cache_interface import KVCacheConfig
41 +from vllm.v1.kv_offload.base import CanonicalKVCaches, OffloadingManager
42 +from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers
43 +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
44 +from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
45 +from vllm.v1.kv_offload.tiering.factory import create_secondary_tier
46 +from vllm.v1.kv_offload.tiering.manager import (
47 + CPUPrimaryTierOffloadingManager,
48 + TieringOffloadingManager,
49 +)
50 +
51 +logger = init_logger(__name__)
52 +
53 +
54 +class TieringOffloadingSpec(CPUOffloadingSpec):
55 + """
56 + Spec for multi-tier KV cache offloading.
57 +
58 + Creates a TieringOffloadingManager with:
59 + - Primary tier: CPU (LRU or ARC eviction policy)
60 + - Secondary tiers: Configurable via extra_config
61 +
62 + The CPU primary tier has direct GPU access and serves as the gateway for
63 + all GPU↔offload operations. Secondary tiers cannot directly access GPU
64 + memory and must transfer data through the primary tier.
65 + """
66 +
67 + def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
68 + super().__init__(vllm_config, kv_cache_config)
69 + # Redeclare for mypy: parent sets this but `--follow-imports skip` hides it
70 + self._manager: OffloadingManager | None = None
71 +
72 + # Parse secondary tier configurations
73 + self.secondary_tier_configs = self.extra_config.get("secondary_tiers", [])
74 + if not isinstance(self.secondary_tier_configs, list):
75 + raise ValueError("secondary_tiers must be a list of tier configurations")
76 +
77 + # Scheduler-side mmap (rank=None); kept for cleanup
78 + self._scheduler_mmap: SharedOffloadRegion | None = None
79 +
80 + def get_manager(self) -> OffloadingManager:
81 + """
82 + Get the TieringOffloadingManager.
83 +
84 + Creates a TieringOffloadingManager with:
85 + - Primary tier: CPU (LRU or ARC)
86 + - Secondary tiers: As configured in extra_config
87 +
88 + Returns:
89 + TieringOffloadingManager instance
90 + """
91 + if not self._manager:
92 + kv_events_config = self.vllm_config.kv_events_config
93 + enable_events = (
94 + kv_events_config is not None and kv_events_config.enable_kv_cache_events
95 + )
96 +
97 + # Create scheduler-side SharedOffloadRegion (rank=None) so the
98 + # primary tier can eagerly create a memoryview over _base.
99 + world_size = self.vllm_config.parallel_config.world_size
100 + scheduler_mmap = SharedOffloadRegion(
101 + instance_id=self.vllm_config.instance_id,
102 + total_size_bytes=self.cpu_page_size_per_worker
103 + * world_size
104 + * self.num_blocks,
105 + num_blocks=self.num_blocks,
106 + rank=None,
107 + num_workers=world_size,
108 + cpu_page_size=self.cpu_page_size_per_worker,
109 + )
110 + self._scheduler_mmap = scheduler_mmap
111 +
112 + # Create primary tier (CPU-based)
113 + assert len(self.gpu_block_size) == 1
114 + primary_tier = CPUPrimaryTierOffloadingManager(
115 + num_blocks=self.num_blocks,
116 + cache_policy=self.eviction_policy, # type: ignore[arg-type]
117 + enable_events=enable_events,
118 + mmap_region=scheduler_mmap,
119 + )
120 +
121 + # Create secondary tiers
122 + primary_kv_view = primary_tier.get_kv_memoryview()
123 + secondary_tiers = []
124 + for i, tier_config in enumerate(self.secondary_tier_configs):
125 + try:
126 + tier = create_secondary_tier(
127 + tier_config, primary_kv_view, self.vllm_config
128 + )
129 + secondary_tiers.append(tier)
130 + logger.info(
131 + "Created secondary tier #%d (%s)",
132 + i,
133 + tier.get_tier_type(),
134 + )
135 + except Exception as e:
136 + logger.error(
137 + "Failed to create secondary tier from config %s: %s",
138 + tier_config,
139 + e,
140 + )
141 + raise
142 +
143 + # Create TieringOffloadingManager. GPU↔CPU transfers use the inherited
144 + # get_handlers(); secondary tier transfers are handled by the
145 + # secondary tier managers and need no additional handlers here.
146 + tiering_manager = TieringOffloadingManager(
147 + primary_tier=primary_tier,
148 + secondary_tiers=secondary_tiers,
149 + enable_events=enable_events,
150 + )
151 + if int(self.extra_config.get("store_threshold", 0)) >= 2:
152 + raise ValueError(
153 + "store_threshold is not supported for TieringOffloadingSpec"
154 + )
155 + self._manager = tiering_manager
156 +
157 + logger.info(
158 + "Created TieringOffloadingManager with primary tier "
159 + "(%s, %s blocks) and %s secondary tier(s)",
160 + self.eviction_policy,
161 + self.num_blocks,
162 + len(secondary_tiers),
163 + )
164 +
165 + return self._manager
166 +
167 + def _create_handlers(
168 + self, kv_caches: CanonicalKVCaches
169 + ) -> CpuGpuOffloadingHandlers:
170 + world_size = self.vllm_config.parallel_config.world_size
171 + rank = torch.accelerator.current_device_index()
172 + worker_mmap = SharedOffloadRegion(
173 + instance_id=self.vllm_config.instance_id,
174 + total_size_bytes=self.cpu_page_size_per_worker
175 + * world_size
176 + * self.num_blocks,
177 + num_blocks=self.num_blocks,
178 + rank=rank,
179 + num_workers=world_size,
180 + cpu_page_size=self.cpu_page_size_per_worker,
181 + )
182 + return CpuGpuOffloadingHandlers(
183 + kv_caches=kv_caches,
184 + block_size_factor=self.block_size_factor,
185 + num_cpu_blocks=self.num_blocks,
186 + mmap_region=worker_mmap,
187 + )