diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index a5c8d8673d..b50e712a8e 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Unit tests for TieringOffloadingManager and ExampleSecondaryTier. +Unit tests for TieringOffloadingManager and ExampleSecondaryTierManager. These tests verify: 1. Basic tiered offloading operations (store, load, lookup) @@ -14,7 +14,6 @@ These tests verify: from collections.abc import Iterable from unittest.mock import MagicMock -import numpy as np import pytest import torch @@ -23,8 +22,7 @@ from vllm.v1.kv_offload.base import ( ReqContext, make_offload_key, ) -from vllm.v1.kv_offload.tiering.base import JobMetadata -from vllm.v1.kv_offload.tiering.example import ExampleSecondaryTier +from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager from vllm.v1.kv_offload.tiering.manager import ( CPUPrimaryTierOffloadingManager, TieringOffloadingManager, @@ -63,14 +61,17 @@ def count_hits(manager, keys: list[OffloadKey]) -> int | None: return count -class TestExampleSecondaryTier: - """Tests for ExampleSecondaryTier implementation.""" +class TestExampleSecondaryTierManager: + """Tests for ExampleSecondaryTierManager implementation.""" def test_basic_store_and_lookup(self): """Test basic store and lookup operations.""" mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy()) - tier = ExampleSecondaryTier( - vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10 + tier = ExampleSecondaryTierManager( + vllm_config=_MOCK_VLLM_CONFIG, + primary_kv_view=mock_view, + tier_type="example", + custom_param=67, ) # Initially empty @@ -88,81 +89,6 @@ class TestExampleSecondaryTier: # Third block not present assert tier.lookup(blocks[2], _CTX) is False - def test_lru_eviction(self): - """Test LRU eviction policy.""" - mock_view = memoryview(torch.zeros((4, 16), dtype=torch.int8).numpy()) - tier = ExampleSecondaryTier( - vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=3 - ) - - # Fill tier to capacity - blocks = to_keys(range(3)) - for block in blocks: - tier.blocks[block] = True - - assert tier.get_num_blocks() == 3 - - # Touch first block (make it most recently used) - tier.touch([blocks[0]], _CTX) - - # Store new block should evict blocks[1] (least recently used) - new_block = to_keys([3])[0] - - tier.submit_store( - JobMetadata( - job_id=1, - keys=[new_block], - block_ids=np.array([0], dtype=np.int64), - is_promotion=False, - req_context=_CTX, - ) - ) - - # Complete the job - tier.get_finished() - - # Verify new block is stored and blocks[1] was evicted (LRU) - assert new_block in tier.blocks - assert blocks[1] not in tier.blocks - # blocks[0] and blocks[2] should still be present - assert blocks[0] in tier.blocks - assert blocks[2] in tier.blocks - - def test_async_simulation(self): - """Test simulated async behavior.""" - mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy()) - tier = ExampleSecondaryTier( - vllm_config=_MOCK_VLLM_CONFIG, - primary_kv_view=mock_view, - max_blocks=10, - simulate_async=True, - ) - - blocks = to_keys(range(2)) - - # Submit store job - tier.submit_store( - JobMetadata( - job_id=1, - keys=blocks, - block_ids=np.array([0, 1], dtype=np.int64), - is_promotion=False, - req_context=_CTX, - ) - ) - - # Blocks should not yet be stored (pending async completion) - assert tier.get_num_blocks() == 0 - - # First get_finished() should complete the job - completed = list(tier.get_finished()) - assert len(completed) == 1 - assert completed[0].job_id == 1 - assert completed[0].success is True - - # Blocks should now be stored - assert tier.get_num_blocks() == 2 - class TestTieringOffloadingManager: """Tests for TieringOffloadingManager.""" @@ -178,11 +104,15 @@ class TestTieringOffloadingManager: mock_view = mock_region.create_kv_memoryview() # Create secondary tiers with the primary view - self.secondary_tier1 = ExampleSecondaryTier( - vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10 + self.secondary_tier1 = ExampleSecondaryTierManager( + vllm_config=_MOCK_VLLM_CONFIG, + primary_kv_view=mock_view, + tier_type="example", ) - self.secondary_tier2 = ExampleSecondaryTier( - vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10 + self.secondary_tier2 = ExampleSecondaryTierManager( + vllm_config=_MOCK_VLLM_CONFIG, + primary_kv_view=mock_view, + tier_type="example", ) # Create tiered manager @@ -359,20 +289,19 @@ class TestTieringOffloadingManager: self.manager.complete_store(blocks, _CTX, success=True) list(self.manager.take_events()) + self.secondary_tier1.touch = MagicMock(wraps=self.secondary_tier1.touch) + self.secondary_tier2.touch = MagicMock(wraps=self.secondary_tier2.touch) + # Touch blocks self.manager.touch(blocks, _CTX) # Verify touch was called on primary tier (check LRU order) - # In LRU, touched blocks should be at the end primary_keys = list(self.primary_tier._policy.blocks.keys()) assert primary_keys[-3:] == list(reversed(blocks)) - # Verify touch was called on all secondary tiers - secondary1_keys = list(self.secondary_tier1.blocks.keys()) - assert secondary1_keys[-3:] == list(reversed(blocks)) - - secondary2_keys = list(self.secondary_tier2.blocks.keys()) - assert secondary2_keys[-3:] == list(reversed(blocks)) + # Verify touch was propagated to all secondary tiers + self.secondary_tier1.touch.assert_called_once_with(blocks, _CTX) + self.secondary_tier2.touch.assert_called_once_with(blocks, _CTX) def test_failed_store_no_cascade(self, manager_setup): """Test that failed GPU→primary store doesn't cascade.""" @@ -396,59 +325,6 @@ class TestTieringOffloadingManager: self.secondary_tier1.submit_store.assert_not_called() self.secondary_tier2.submit_store.assert_not_called() - def test_multiple_secondary_tiers_independent_eviction(self): - """Test that secondary tiers manage their own evictions.""" - mock_region = _mock_mmap_region(10) - mock_view = mock_region.create_kv_memoryview() - - # Create tier with small capacity - small_tier = ExampleSecondaryTier( - vllm_config=_MOCK_VLLM_CONFIG, - primary_kv_view=mock_view, - max_blocks=5, - simulate_async=False, - ) - large_tier = ExampleSecondaryTier( - vllm_config=_MOCK_VLLM_CONFIG, - primary_kv_view=mock_view, - max_blocks=10, - simulate_async=False, - ) - - # Create a fresh primary tier for this test - primary_tier = CPUPrimaryTierOffloadingManager( - num_blocks=10, mmap_region=mock_region - ) - - manager = TieringOffloadingManager( - primary_tier=primary_tier, - secondary_tiers=[small_tier, large_tier], - ) - - # First, store 5 blocks to fill the small tier - blocks1 = to_keys(range(5)) - result = manager.prepare_store(blocks1, _CTX) - assert result is not None - manager.complete_store(blocks1, _CTX, success=True) - list(manager.take_events()) - - # Both tiers should have 5 blocks - assert small_tier.get_num_blocks() == 5 - assert large_tier.get_num_blocks() == 5 - - # Now store 3 more blocks - small tier should evict 3 blocks - blocks2 = to_keys(range(5, 8)) - result = manager.prepare_store(blocks2, _CTX) - assert result is not None - manager.complete_store(blocks2, _CTX, success=True) - list(manager.take_events()) - - # Small tier should still have 5 blocks (evicted 3, added 3) - assert small_tier.get_num_blocks() == 5 - - # Large tier should have all 8 blocks - assert large_tier.get_num_blocks() == 8 - def test_lookup_batches_submit_load_per_request(self, manager_setup): """lookup() defers submit_load until take_events(), one call per request. diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 99265e09c0..6d17d5317f 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -81,9 +81,7 @@ class CPUOffloadingSpec(OffloadingSpec): ) return self._manager - def _create_handlers( - self, kv_caches: CanonicalKVCaches - ) -> CpuGpuOffloadingHandlers: + def create_handlers(self, kv_caches: CanonicalKVCaches) -> CpuGpuOffloadingHandlers: return CpuGpuOffloadingHandlers( kv_caches=kv_caches, block_size_factor=self.block_size_factor, @@ -98,7 +96,7 @@ class CPUOffloadingSpec(OffloadingSpec): raise Exception( "CPU Offloading is currently only supported on CUDA-alike GPUs" ) - self._handlers = self._create_handlers(kv_caches) + self._handlers = self.create_handlers(kv_caches) assert self._handlers is not None yield GPULoadStoreSpec, CPULoadStoreSpec, self._handlers.gpu_to_cpu_handler diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index ace325185a..d9bd238f3c 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -53,9 +53,22 @@ class SecondaryTierManager(ABC): async jobs; get_finished() polls for completion. """ - def __init__(self, vllm_config: "VllmConfig", primary_kv_view: memoryview) -> None: + def __init__( + self, + vllm_config: "VllmConfig", + primary_kv_view: memoryview, + tier_type: str, + ) -> None: + """ + Args: + vllm_config: Global vLLM configuration. + primary_kv_view: Memoryview of the primary tier's CPU KV cache. + tier_type: Tier type identifier, set by SecondaryTierFactory + from the registered tier type. + """ self._vllm_config = vllm_config self._primary_kv_view: memoryview = primary_kv_view + self.tier_type = tier_type @abstractmethod def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: @@ -153,16 +166,3 @@ class SecondaryTierManager(ABC): def shutdown(self) -> None: """Release resources held by this tier (threads, connections, etc.).""" return - - @staticmethod - @abstractmethod - def get_tier_type() -> str: - """ - Get the type identifier of this tier (e.g., "example", "storage"). - - Must match the "type" field in the tier config dict. - - Returns: - Tier type string. - """ - pass diff --git a/vllm/v1/kv_offload/tiering/example/__init__.py b/vllm/v1/kv_offload/tiering/example/__init__.py index 8c3edb693f..e69de29bb2 100644 --- a/vllm/v1/kv_offload/tiering/example/__init__.py +++ b/vllm/v1/kv_offload/tiering/example/__init__.py @@ -1,235 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -ExampleSecondaryTier: A simple in-memory secondary tier for testing. - -This implementation provides a minimal secondary tier that stores blocks -in memory (using a dictionary) and simulates async transfers with immediate -completion. It's useful for testing the TieringOffloadingManager without -requiring actual storage or network backends. -""" - -from collections import OrderedDict -from collections.abc import Collection, Iterable -from dataclasses import dataclass -from typing import TYPE_CHECKING - -from vllm.v1.kv_offload.base import OffloadKey, ReqContext -from vllm.v1.kv_offload.tiering.base import ( - JobId, - JobMetadata, - JobResult, - SecondaryTierManager, -) - -if TYPE_CHECKING: - from vllm.config import VllmConfig - - -@dataclass -class _JobMetadata: - """Internal metadata for tracking job details.""" - - job_id: JobId - keys: Collection[OffloadKey] - is_store: bool # True for store jobs, False for load jobs - - -class ExampleSecondaryTier(SecondaryTierManager): - """ - A simple in-memory secondary tier for testing. - - This implementation: - - Stores blocks in a dictionary (key -> True) - - Simulates async transfers with immediate completion - - Uses LRU eviction policy - """ - - def __init__( - self, - vllm_config: "VllmConfig", - primary_kv_view: memoryview, - max_blocks: int = 1000, - simulate_async: bool = False, - ): - """ - Initialize the example secondary tier. - - Args: - vllm_config: Global vLLM configuration. - primary_kv_view: Memoryview of the primary tier's CPU KV cache. - max_blocks: Maximum number of blocks this tier can store - simulate_async: If True, jobs complete on next get_finished() call. - If False, jobs complete immediately. - """ - super().__init__(vllm_config, primary_kv_view) - self.max_blocks = max_blocks - self.simulate_async = simulate_async - - # key -> True (only care about presence) - self.blocks: OrderedDict[OffloadKey, bool] = OrderedDict() - - # Completed jobs waiting to be retrieved by get_finished() - self.completed_jobs: list[JobResult] = [] - - # Pending jobs (for simulated async mode) - self.pending_jobs: list[_JobMetadata] = [] - - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: - """ - Check whether a block exists in this secondary tier. - - Args: - key: Offload key to look up. - req_context: Per-request context. - - Returns: - True if the block is present, False if not found. - """ - return key in self.blocks - - def submit_store(self, job_metadata: JobMetadata) -> None: - """ - Submit an async job to store blocks from primary tier to this tier. - - Args: - job_metadata: Job metadata including job_id, keys, and - spec for reading blocks from the primary tier. - """ - job_id = job_metadata.job_id - keys = job_metadata.keys - block_ids = job_metadata.block_ids - - assert len(keys) == len(block_ids), ( - f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids" - ) - - # Filter out blocks already present - blocks_to_store = [bh for bh in keys if bh not in self.blocks] - - if not blocks_to_store: - # All blocks already present - return - - # Evict blocks if needed (LRU policy) - num_blocks_to_evict = len(blocks_to_store) - ( - self.max_blocks - len(self.blocks) - ) - - evicted = [] - if num_blocks_to_evict > 0: - # Collect eviction candidates first (LRU order), then delete atomically - protected = set(keys) - for key in self.blocks: - if key not in protected: - evicted.append(key) - if len(evicted) == num_blocks_to_evict: - break - else: - # Could not collect enough eviction candidates - return - for key in evicted: - del self.blocks[key] - - # Create internal job metadata - internal_job_metadata = _JobMetadata( - job_id=job_id, keys=blocks_to_store, is_store=True - ) - - if self.simulate_async: - # Job will complete on next get_finished() call - self.pending_jobs.append(internal_job_metadata) - else: - # Job completes immediately - self._complete_store_job(internal_job_metadata) - - def submit_load(self, job_metadata: JobMetadata) -> None: - """ - Submit an async job to load blocks from this tier to primary tier. - - Args: - job_metadata: Job metadata including job_id, keys, and - spec for writing blocks into the primary tier. - """ - job_id = job_metadata.job_id - keys = job_metadata.keys - block_ids = job_metadata.block_ids - - assert len(keys) == len(block_ids), ( - f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids" - ) - - # Verify all blocks exist - for key in keys: - if key not in self.blocks: - return - - # Create internal job metadata - internal_job_metadata = _JobMetadata(job_id=job_id, keys=keys, is_store=False) - - if self.simulate_async: - # Job will complete on next get_finished() call - self.pending_jobs.append(internal_job_metadata) - else: - # Job completes immediately - self._complete_load_job(internal_job_metadata) - - def get_finished(self) -> Iterable[JobResult]: - """ - Poll for finished async jobs. - - Returns: - Iterable of JobResult objects for all jobs that have - finished since the last call. - """ - # Move pending jobs to completed - if self.simulate_async and self.pending_jobs: - for job_metadata in self.pending_jobs: - if job_metadata.is_store: - self._complete_store_job(job_metadata) - else: - self._complete_load_job(job_metadata) - self.pending_jobs.clear() - - # Return completed jobs - result = self.completed_jobs - self.completed_jobs = [] - return result - - def _complete_store_job(self, job_metadata: _JobMetadata): - """Complete a store job by adding blocks to storage.""" - for key in job_metadata.keys: - self.blocks[key] = True - # Return simplified JobResult (only job_id and success) - self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) - - def _complete_load_job(self, job_metadata: _JobMetadata): - """Complete a load job.""" - # Return simplified JobResult (only job_id and success) - self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) - - def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): - """ - Mark blocks as recently used (move to end of LRU list). - - Args: - keys: Blocks to mark as recently used. - req_context: Per-request context. - """ - for key in reversed(list(keys)): - if key in self.blocks: - self.blocks.move_to_end(key) - - @staticmethod - def get_tier_type() -> str: - return "example" - - def get_num_blocks(self) -> int: - """Get the number of blocks currently stored in this tier.""" - return len(self.blocks) - - def clear(self): - """Clear all blocks and pending jobs (for testing).""" - self.blocks.clear() - self.completed_jobs.clear() - self.pending_jobs.clear() diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py new file mode 100644 index 0000000000..b7a59f6738 --- /dev/null +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +ExampleSecondaryTierManager: A simple in-memory secondary tier. + +This implementation provides a minimal secondary tier that stores blocks +in memory (using a dictionary) with immediate completion. It serves as a +reference for writing new tiers and is useful for testing the +TieringOffloadingManager without requiring actual storage or network backends. +""" + +import logging +from collections.abc import Iterable +from typing import TYPE_CHECKING + +from vllm.v1.kv_offload.base import OffloadKey, ReqContext +from vllm.v1.kv_offload.tiering.base import ( + JobMetadata, + JobResult, + SecondaryTierManager, +) + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from vllm.config import VllmConfig + + +class ExampleSecondaryTierManager(SecondaryTierManager): + """ + A simple in-memory secondary tier. + + This implementation: + - Stores blocks in a dictionary (key -> True) + - Completes transfers immediately (synchronous) + """ + + def __init__( + self, + vllm_config: "VllmConfig", + primary_kv_view: memoryview, + tier_type: str, + custom_param: int = 0, + ): + """ + Initialize the example secondary tier. + + Args: + custom_param: Dummy parameter demonstrating custom args. + """ + super().__init__( + vllm_config=vllm_config, + primary_kv_view=primary_kv_view, + tier_type=tier_type, + ) + + logger.info( + "ExampleSecondaryTierManager initialized with custom_param=%d", custom_param + ) + + # key -> True (only care about presence) + self.blocks: dict[OffloadKey, bool] = {} + + # Completed jobs waiting to be retrieved by get_finished() + self.completed_jobs: list[JobResult] = [] + + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + """ + Check whether a block exists in this secondary tier. + + Args: + key: Offload key to look up. + req_context: Per-request context. + + Returns: + True if the block is present, False if not found. + """ + return key in self.blocks + + def submit_store(self, job_metadata: JobMetadata) -> None: + """ + Submit a job to store blocks from primary tier to this tier. + + Args: + job_metadata: Job metadata including job_id, keys, and + spec for reading blocks from the primary tier. + """ + keys = job_metadata.keys + block_ids = job_metadata.block_ids + + assert len(keys) == len(block_ids), ( + f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids" + ) + + for key in keys: + self.blocks[key] = True + self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) + + def submit_load(self, job_metadata: JobMetadata) -> None: + """ + Submit a job to load blocks from this tier to primary tier. + + Args: + job_metadata: Job metadata including job_id, keys, and + spec for writing blocks into the primary tier. + """ + keys = job_metadata.keys + block_ids = job_metadata.block_ids + + assert len(keys) == len(block_ids), ( + f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids" + ) + + for key in keys: + if key not in self.blocks: + self.completed_jobs.append( + JobResult(job_id=job_metadata.job_id, success=False) + ) + return + + self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) + + def get_finished(self) -> Iterable[JobResult]: + """ + Poll for finished jobs. + + Returns: + Iterable of JobResult objects for all jobs that have + finished since the last call. + """ + result = self.completed_jobs + self.completed_jobs = [] + return result + + def get_num_blocks(self) -> int: + """Get the number of blocks currently stored in this tier.""" + return len(self.blocks) diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index e8f4a9b32d..8980841bb2 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -1,56 +1,59 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Factory for creating secondary tier implementations. -""" - +import importlib +from collections.abc import Callable from typing import TYPE_CHECKING from vllm.v1.kv_offload.tiering.base import SecondaryTierManager -from vllm.v1.kv_offload.tiering.example import ExampleSecondaryTier if TYPE_CHECKING: from vllm.config import VllmConfig -SUPPORTED_TIERS: tuple[type[SecondaryTierManager], ...] = (ExampleSecondaryTier,) -_TIER_REGISTRY: dict[str, type[SecondaryTierManager]] = { - cls.get_tier_type(): cls for cls in SUPPORTED_TIERS -} +class SecondaryTierFactory: + _registry: dict[str, Callable[[], type[SecondaryTierManager]]] = {} + @classmethod + def register_tier(cls, tier_type: str, module_path: str, class_name: str) -> None: + if tier_type in cls._registry: + raise ValueError(f"Tier '{tier_type}' is already registered.") -def create_secondary_tier( - tier_config: dict, - primary_kv_view: memoryview, - vllm_config: "VllmConfig", -) -> SecondaryTierManager: - """ - Create a secondary tier from configuration. + def loader() -> type[SecondaryTierManager]: + module = importlib.import_module(module_path) + return getattr(module, class_name) - Args: - tier_config: Dictionary with tier configuration containing: - - type (required): Type of secondary tier (e.g., "example") - - Additional tier-specific parameters are passed directly - to the tier constructor - primary_kv_view: Memoryview of the primary tier's CPU KV cache. - vllm_config: Global vLLM configuration. + cls._registry[tier_type] = loader - Returns: - SecondaryTierManager instance + @classmethod + def create_secondary_tier( + cls, + tier_config: dict, + primary_kv_view: memoryview, + vllm_config: "VllmConfig", + ) -> SecondaryTierManager: + config = tier_config.copy() - Raises: - ValueError: If tier type is unknown or configuration is invalid - """ - config = tier_config.copy() + tier_type = config.pop("type", None) + if not tier_type: + raise ValueError("Secondary tier configuration must include 'type'") - tier_type = config.pop("type", None) - if not tier_type: - raise ValueError("Secondary tier configuration must include 'type'") + if tier_type not in cls._registry: + raise ValueError( + f"Unknown secondary tier type: {tier_type!r}. " + f"Supported types: {list(cls._registry)}" + ) - cls = _TIER_REGISTRY.get(tier_type) - if cls is None: - raise ValueError( - f"Unknown secondary tier type: {tier_type!r}. " - f"Supported types: {list(_TIER_REGISTRY)}" + tier_cls = cls._registry[tier_type]() + return tier_cls( + vllm_config=vllm_config, + primary_kv_view=primary_kv_view, + tier_type=tier_type, + **config, ) - return cls(vllm_config=vllm_config, primary_kv_view=primary_kv_view, **config) + + +SecondaryTierFactory.register_tier( + "example", + "vllm.v1.kv_offload.tiering.example.manager", + "ExampleSecondaryTierManager", +) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 578992fdd7..4a2d79254e 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -194,7 +194,7 @@ class TieringOffloadingManager(OffloadingManager): job_metadata = self._transfer_jobs.pop(job_id, None) assert job_metadata is not None, ( f"Finished job_id {job_id} from tier #{i}" - f" ({tier.get_tier_type()}) not in _transfer_jobs" + f" ({tier.tier_type}) not in _transfer_jobs" ) if job_metadata.is_promotion: diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 0a372ae8ef..515dae75d3 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -25,15 +25,14 @@ Example configuration: "secondary_tiers": [ { "type": "example", - # Tier-specific parameters (for ExampleSecondaryTier): - "max_blocks": 10000, - "simulate_async": False + "custom_param": 67 } ] } """ import torch +from typing_extensions import override from vllm.config import VllmConfig from vllm.logger import init_logger @@ -42,7 +41,7 @@ from vllm.v1.kv_offload.base import CanonicalKVCaches, OffloadingManager from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec -from vllm.v1.kv_offload.tiering.factory import create_secondary_tier +from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory from vllm.v1.kv_offload.tiering.manager import ( CPUPrimaryTierOffloadingManager, TieringOffloadingManager, @@ -77,6 +76,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec): # Scheduler-side mmap (rank=None); kept for cleanup self._scheduler_mmap: SharedOffloadRegion | None = None + @override def get_manager(self) -> OffloadingManager: """ Get the TieringOffloadingManager. @@ -123,14 +123,14 @@ class TieringOffloadingSpec(CPUOffloadingSpec): secondary_tiers = [] for i, tier_config in enumerate(self.secondary_tier_configs): try: - tier = create_secondary_tier( + tier = SecondaryTierFactory.create_secondary_tier( tier_config, primary_kv_view, self.vllm_config ) secondary_tiers.append(tier) logger.info( "Created secondary tier #%d (%s)", i, - tier.get_tier_type(), + tier.tier_type, ) except Exception as e: logger.error( @@ -164,9 +164,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec): return self._manager - def _create_handlers( - self, kv_caches: CanonicalKVCaches - ) -> CpuGpuOffloadingHandlers: + @override + def create_handlers(self, kv_caches: CanonicalKVCaches) -> CpuGpuOffloadingHandlers: world_size = self.vllm_config.parallel_config.world_size rank = torch.accelerator.current_device_index() worker_mmap = SharedOffloadRegion(