@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 84747489de
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, 09:15 UTC

Tier offload followup (#42529)

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

Ronen Schaffer committed May 18, 2026 at 22:41 UTC 84747489ded65265ee7d43815bfa3373b0d42279
8 files changed +234 -456
tests/v1/kv_offload/test_tiering_offloading.py
+23 -147
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """
4 -Unit tests for TieringOffloadingManager and ExampleSecondaryTier.
4 +Unit tests for TieringOffloadingManager and ExampleSecondaryTierManager.
5
6 These tests verify:
7 1. Basic tiered offloading operations (store, load, lookup)
@@ -14,7 +14,6 @@ These tests verify:
14 from collections.abc import Iterable
15 from unittest.mock import MagicMock
16
17 -import numpy as np
17 import pytest
18 import torch
19
@@ -23,8 +22,7 @@ from vllm.v1.kv_offload.base import (
22 ReqContext,
23 make_offload_key,
24 )
26 -from vllm.v1.kv_offload.tiering.base import JobMetadata
27 -from vllm.v1.kv_offload.tiering.example import ExampleSecondaryTier
25 +from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
26 from vllm.v1.kv_offload.tiering.manager import (
27 CPUPrimaryTierOffloadingManager,
28 TieringOffloadingManager,
@@ -63,14 +61,17 @@ def count_hits(manager, keys: list[OffloadKey]) -> int | None:
61 return count
62
63
66 -class TestExampleSecondaryTier:
67 - """Tests for ExampleSecondaryTier implementation."""
64 +class TestExampleSecondaryTierManager:
65 + """Tests for ExampleSecondaryTierManager implementation."""
66
67 def test_basic_store_and_lookup(self):
68 """Test basic store and lookup operations."""
69 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
70 + tier = ExampleSecondaryTierManager(
71 + vllm_config=_MOCK_VLLM_CONFIG,
72 + primary_kv_view=mock_view,
73 + tier_type="example",
74 + custom_param=67,
75 )
76
77 # Initially empty
@@ -88,81 +89,6 @@ class TestExampleSecondaryTier:
89 # Third block not present
90 assert tier.lookup(blocks[2], _CTX) is False
91
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 -
92
93 class TestTieringOffloadingManager:
94 """Tests for TieringOffloadingManager."""
@@ -178,11 +104,15 @@ class TestTieringOffloadingManager:
104 mock_view = mock_region.create_kv_memoryview()
105
106 # 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
107 + self.secondary_tier1 = ExampleSecondaryTierManager(
108 + vllm_config=_MOCK_VLLM_CONFIG,
109 + primary_kv_view=mock_view,
110 + tier_type="example",
111 )
184 - self.secondary_tier2 = ExampleSecondaryTier(
185 - vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10
112 + self.secondary_tier2 = ExampleSecondaryTierManager(
113 + vllm_config=_MOCK_VLLM_CONFIG,
114 + primary_kv_view=mock_view,
115 + tier_type="example",
116 )
117
118 # Create tiered manager
@@ -359,20 +289,19 @@ class TestTieringOffloadingManager:
289 self.manager.complete_store(blocks, _CTX, success=True)
290 list(self.manager.take_events())
291
292 + self.secondary_tier1.touch = MagicMock(wraps=self.secondary_tier1.touch)
293 + self.secondary_tier2.touch = MagicMock(wraps=self.secondary_tier2.touch)
294 +
295 # Touch blocks
296 self.manager.touch(blocks, _CTX)
297
298 # Verify touch was called on primary tier (check LRU order)
366 - # In LRU, touched blocks should be at the end
299 primary_keys = list(self.primary_tier._policy.blocks.keys())
300 assert primary_keys[-3:] == list(reversed(blocks))
301
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))
302 + # Verify touch was propagated to all secondary tiers
303 + self.secondary_tier1.touch.assert_called_once_with(blocks, _CTX)
304 + self.secondary_tier2.touch.assert_called_once_with(blocks, _CTX)
305
306 def test_failed_store_no_cascade(self, manager_setup):
307 """Test that failed GPU→primary store doesn't cascade."""
@@ -396,59 +325,6 @@ class TestTieringOffloadingManager:
325 self.secondary_tier1.submit_store.assert_not_called()
326 self.secondary_tier2.submit_store.assert_not_called()
327
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 -
328 def test_lookup_batches_submit_load_per_request(self, manager_setup):
329 """lookup() defers submit_load until take_events(), one call per request.
330
vllm/v1/kv_offload/cpu/spec.py
+2 -4
@@ -81,9 +81,7 @@ class CPUOffloadingSpec(OffloadingSpec):
81 )
82 return self._manager
83
84 - def _create_handlers(
85 - self, kv_caches: CanonicalKVCaches
86 - ) -> CpuGpuOffloadingHandlers:
84 + def create_handlers(self, kv_caches: CanonicalKVCaches) -> CpuGpuOffloadingHandlers:
85 return CpuGpuOffloadingHandlers(
86 kv_caches=kv_caches,
87 block_size_factor=self.block_size_factor,
@@ -98,7 +96,7 @@ class CPUOffloadingSpec(OffloadingSpec):
96 raise Exception(
97 "CPU Offloading is currently only supported on CUDA-alike GPUs"
98 )
101 - self._handlers = self._create_handlers(kv_caches)
99 + self._handlers = self.create_handlers(kv_caches)
100
101 assert self._handlers is not None
102 yield GPULoadStoreSpec, CPULoadStoreSpec, self._handlers.gpu_to_cpu_handler
vllm/v1/kv_offload/tiering/base.py
+14 -14
@@ -53,9 +53,22 @@ class SecondaryTierManager(ABC):
53 async jobs; get_finished() polls for completion.
54 """
55
56 - def __init__(self, vllm_config: "VllmConfig", primary_kv_view: memoryview) -> None:
56 + def __init__(
57 + self,
58 + vllm_config: "VllmConfig",
59 + primary_kv_view: memoryview,
60 + tier_type: str,
61 + ) -> None:
62 + """
63 + Args:
64 + vllm_config: Global vLLM configuration.
65 + primary_kv_view: Memoryview of the primary tier's CPU KV cache.
66 + tier_type: Tier type identifier, set by SecondaryTierFactory
67 + from the registered tier type.
68 + """
69 self._vllm_config = vllm_config
70 self._primary_kv_view: memoryview = primary_kv_view
71 + self.tier_type = tier_type
72
73 @abstractmethod
74 def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
@@ -153,16 +166,3 @@ class SecondaryTierManager(ABC):
166 def shutdown(self) -> None:
167 """Release resources held by this tier (threads, connections, etc.)."""
168 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
-235
@@ -1,235 +0,0 @@
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/example/manager.py new
+137
@@ -0,0 +1,137 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +ExampleSecondaryTierManager: A simple in-memory secondary tier.
5 +
6 +This implementation provides a minimal secondary tier that stores blocks
7 +in memory (using a dictionary) with immediate completion. It serves as a
8 +reference for writing new tiers and is useful for testing the
9 +TieringOffloadingManager without requiring actual storage or network backends.
10 +"""
11 +
12 +import logging
13 +from collections.abc import Iterable
14 +from typing import TYPE_CHECKING
15 +
16 +from vllm.v1.kv_offload.base import OffloadKey, ReqContext
17 +from vllm.v1.kv_offload.tiering.base import (
18 + JobMetadata,
19 + JobResult,
20 + SecondaryTierManager,
21 +)
22 +
23 +logger = logging.getLogger(__name__)
24 +
25 +if TYPE_CHECKING:
26 + from vllm.config import VllmConfig
27 +
28 +
29 +class ExampleSecondaryTierManager(SecondaryTierManager):
30 + """
31 + A simple in-memory secondary tier.
32 +
33 + This implementation:
34 + - Stores blocks in a dictionary (key -> True)
35 + - Completes transfers immediately (synchronous)
36 + """
37 +
38 + def __init__(
39 + self,
40 + vllm_config: "VllmConfig",
41 + primary_kv_view: memoryview,
42 + tier_type: str,
43 + custom_param: int = 0,
44 + ):
45 + """
46 + Initialize the example secondary tier.
47 +
48 + Args:
49 + custom_param: Dummy parameter demonstrating custom args.
50 + """
51 + super().__init__(
52 + vllm_config=vllm_config,
53 + primary_kv_view=primary_kv_view,
54 + tier_type=tier_type,
55 + )
56 +
57 + logger.info(
58 + "ExampleSecondaryTierManager initialized with custom_param=%d", custom_param
59 + )
60 +
61 + # key -> True (only care about presence)
62 + self.blocks: dict[OffloadKey, bool] = {}
63 +
64 + # Completed jobs waiting to be retrieved by get_finished()
65 + self.completed_jobs: list[JobResult] = []
66 +
67 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
68 + """
69 + Check whether a block exists in this secondary tier.
70 +
71 + Args:
72 + key: Offload key to look up.
73 + req_context: Per-request context.
74 +
75 + Returns:
76 + True if the block is present, False if not found.
77 + """
78 + return key in self.blocks
79 +
80 + def submit_store(self, job_metadata: JobMetadata) -> None:
81 + """
82 + Submit a job to store blocks from primary tier to this tier.
83 +
84 + Args:
85 + job_metadata: Job metadata including job_id, keys, and
86 + spec for reading blocks from the primary tier.
87 + """
88 + keys = job_metadata.keys
89 + block_ids = job_metadata.block_ids
90 +
91 + assert len(keys) == len(block_ids), (
92 + f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids"
93 + )
94 +
95 + for key in keys:
96 + self.blocks[key] = True
97 + self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))
98 +
99 + def submit_load(self, job_metadata: JobMetadata) -> None:
100 + """
101 + Submit a job to load blocks from this tier to primary tier.
102 +
103 + Args:
104 + job_metadata: Job metadata including job_id, keys, and
105 + spec for writing blocks into the primary tier.
106 + """
107 + keys = job_metadata.keys
108 + block_ids = job_metadata.block_ids
109 +
110 + assert len(keys) == len(block_ids), (
111 + f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids"
112 + )
113 +
114 + for key in keys:
115 + if key not in self.blocks:
116 + self.completed_jobs.append(
117 + JobResult(job_id=job_metadata.job_id, success=False)
118 + )
119 + return
120 +
121 + self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))
122 +
123 + def get_finished(self) -> Iterable[JobResult]:
124 + """
125 + Poll for finished jobs.
126 +
127 + Returns:
128 + Iterable of JobResult objects for all jobs that have
129 + finished since the last call.
130 + """
131 + result = self.completed_jobs
132 + self.completed_jobs = []
133 + return result
134 +
135 + def get_num_blocks(self) -> int:
136 + """Get the number of blocks currently stored in this tier."""
137 + return len(self.blocks)
vllm/v1/kv_offload/tiering/factory.py
+49 -46
@@ -1,56 +1,59 @@
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 -
3 +import importlib
4 +from collections.abc import Callable
5 from typing import TYPE_CHECKING
6
7 from vllm.v1.kv_offload.tiering.base import SecondaryTierManager
10 -from vllm.v1.kv_offload.tiering.example import ExampleSecondaryTier
8
9 if TYPE_CHECKING:
10 from vllm.config import VllmConfig
11
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)}"
12 +
13 +class SecondaryTierFactory:
14 + _registry: dict[str, Callable[[], type[SecondaryTierManager]]] = {}
15 +
16 + @classmethod
17 + def register_tier(cls, tier_type: str, module_path: str, class_name: str) -> None:
18 + if tier_type in cls._registry:
19 + raise ValueError(f"Tier '{tier_type}' is already registered.")
20 +
21 + def loader() -> type[SecondaryTierManager]:
22 + module = importlib.import_module(module_path)
23 + return getattr(module, class_name)
24 +
25 + cls._registry[tier_type] = loader
26 +
27 + @classmethod
28 + def create_secondary_tier(
29 + cls,
30 + tier_config: dict,
31 + primary_kv_view: memoryview,
32 + vllm_config: "VllmConfig",
33 + ) -> SecondaryTierManager:
34 + config = tier_config.copy()
35 +
36 + tier_type = config.pop("type", None)
37 + if not tier_type:
38 + raise ValueError("Secondary tier configuration must include 'type'")
39 +
40 + if tier_type not in cls._registry:
41 + raise ValueError(
42 + f"Unknown secondary tier type: {tier_type!r}. "
43 + f"Supported types: {list(cls._registry)}"
44 + )
45 +
46 + tier_cls = cls._registry[tier_type]()
47 + return tier_cls(
48 + vllm_config=vllm_config,
49 + primary_kv_view=primary_kv_view,
50 + tier_type=tier_type,
51 + **config,
52 )
56 - return cls(vllm_config=vllm_config, primary_kv_view=primary_kv_view, **config)
53 +
54 +
55 +SecondaryTierFactory.register_tier(
56 + "example",
57 + "vllm.v1.kv_offload.tiering.example.manager",
58 + "ExampleSecondaryTierManager",
59 +)
vllm/v1/kv_offload/tiering/manager.py
+1 -1
@@ -194,7 +194,7 @@ class TieringOffloadingManager(OffloadingManager):
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"
197 + f" ({tier.tier_type}) not in _transfer_jobs"
198 )
199
200 if job_metadata.is_promotion:
vllm/v1/kv_offload/tiering/spec.py
+8 -9
@@ -25,15 +25,14 @@ Example configuration:
25 "secondary_tiers": [
26 {
27 "type": "example",
28 - # Tier-specific parameters (for ExampleSecondaryTier):
29 - "max_blocks": 10000,
30 - "simulate_async": False
28 + "custom_param": 67
29 }
30 ]
31 }
32 """
33
34 import torch
35 +from typing_extensions import override
36
37 from vllm.config import VllmConfig
38 from vllm.logger import init_logger
@@ -42,7 +41,7 @@ from vllm.v1.kv_offload.base import CanonicalKVCaches, OffloadingManager
41 from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers
42 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
43 from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
45 -from vllm.v1.kv_offload.tiering.factory import create_secondary_tier
44 +from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory
45 from vllm.v1.kv_offload.tiering.manager import (
46 CPUPrimaryTierOffloadingManager,
47 TieringOffloadingManager,
@@ -77,6 +76,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
76 # Scheduler-side mmap (rank=None); kept for cleanup
77 self._scheduler_mmap: SharedOffloadRegion | None = None
78
79 + @override
80 def get_manager(self) -> OffloadingManager:
81 """
82 Get the TieringOffloadingManager.
@@ -123,14 +123,14 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
123 secondary_tiers = []
124 for i, tier_config in enumerate(self.secondary_tier_configs):
125 try:
126 - tier = create_secondary_tier(
126 + tier = SecondaryTierFactory.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(),
133 + tier.tier_type,
134 )
135 except Exception as e:
136 logger.error(
@@ -164,9 +164,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
164
165 return self._manager
166
167 - def _create_handlers(
168 - self, kv_caches: CanonicalKVCaches
169 - ) -> CpuGpuOffloadingHandlers:
167 + @override
168 + def create_handlers(self, kv_caches: CanonicalKVCaches) -> CpuGpuOffloadingHandlers:
169 world_size = self.vllm_config.parallel_config.world_size
170 rank = torch.accelerator.current_device_index()
171 worker_mmap = SharedOffloadRegion(