From 84c6059c55bf43854212609e740075c3717a7c68 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 16 Apr 2026 16:15:33 +0300 Subject: [PATCH 01/55] [kv_offload] Add tiered KV offloading framework Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 423 ++++++++++++++++ vllm/v1/kv_offload/abstract.py | 155 +++++- vllm/v1/kv_offload/cpu/spec.py | 51 +- vllm/v1/kv_offload/factory.py | 5 + .../v1/kv_offload/secondary_tiers/__init__.py | 18 + vllm/v1/kv_offload/secondary_tiers/dummy.py | 299 ++++++++++++ vllm/v1/kv_offload/tiering/__init__.py | 0 vllm/v1/kv_offload/tiering/manager.py | 455 ++++++++++++++++++ vllm/v1/kv_offload/tiering/spec.py | 231 +++++++++ vllm/v1/kv_offload/worker/cpu_gpu.py | 7 + 10 files changed, 1623 insertions(+), 21 deletions(-) create mode 100644 tests/v1/kv_offload/test_tiering_offloading.py create mode 100644 vllm/v1/kv_offload/secondary_tiers/__init__.py create mode 100644 vllm/v1/kv_offload/secondary_tiers/dummy.py create mode 100644 vllm/v1/kv_offload/tiering/__init__.py create mode 100644 vllm/v1/kv_offload/tiering/manager.py create mode 100644 vllm/v1/kv_offload/tiering/spec.py diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py new file mode 100644 index 000000000000..ded6ce8b004b --- /dev/null +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -0,0 +1,423 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Unit tests for TieringOffloadingManager and DummySecondaryTier. + +These tests verify: +1. Basic tiered offloading operations (store, load, lookup) +2. Cascade behavior (blocks stored to all secondary tiers) +3. Promotion behavior (blocks loaded from secondary to primary to GPU) +4. ref_cnt management (blocks protected during async transfers) +5. Eviction coordination between tiers +""" + +from collections.abc import Iterable + +import pytest +import torch + +from vllm.v1.kv_offload.abstract import JobMetadata, OffloadKey, make_offload_key +from vllm.v1.kv_offload.mediums import CPULoadStoreSpec +from vllm.v1.kv_offload.secondary_tiers.dummy import DummySecondaryTier +from vllm.v1.kv_offload.tiering.manager import ( + CPUPrimaryTierOffloadingManager, + TieringOffloadingManager, +) + + +def to_keys(int_ids: Iterable[int]) -> list[OffloadKey]: + return [make_offload_key(str(i).encode(), 0) for i in int_ids] + + +class TestDummySecondaryTier: + """Tests for DummySecondaryTier implementation.""" + + def test_basic_store_and_lookup(self): + """Test basic store and lookup operations.""" + tier = DummySecondaryTier(tier_name="Test", max_blocks=10) + + # Initially empty + blocks = to_keys(range(3)) + assert tier.lookup(blocks) == 0 + + # Store blocks (simulate with direct insertion for testing) + tier.blocks[blocks[0]] = True + tier.blocks[blocks[1]] = True + + # Lookup should find 2 blocks + assert tier.lookup(blocks) == 2 + + # Third block not present + assert tier.lookup([blocks[2]]) == 0 + + def test_in_flight_blocks_return_none(self): + """Test that in-flight blocks cause lookup to return None.""" + tier = DummySecondaryTier(tier_name="Test", max_blocks=10) + + blocks = to_keys(range(3)) + + # Mark first block as in-flight + tier.in_flight[blocks[0]] = 1 + + # Lookup should return None (retry later) + assert tier.lookup(blocks) is None + + def test_lru_eviction(self): + """Test LRU eviction policy.""" + tier = DummySecondaryTier(tier_name="Test", 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]]) + + # Store new block should evict blocks[1] (least recently used) + new_block = to_keys([3])[0] + + mock_tensor = torch.zeros((4, 16), dtype=torch.float32) + tier.set_primary_view(memoryview(mock_tensor.numpy())) + + tier.submit_store( + JobMetadata( + job_id=1, + keys=[new_block], + spec=CPULoadStoreSpec([0]), + ) + ) + + # 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.""" + tier = DummySecondaryTier(tier_name="Test", max_blocks=10, simulate_async=True) + + blocks = to_keys(range(2)) + + mock_tensor = torch.zeros((10, 16), dtype=torch.float32) + tier.set_primary_view(memoryview(mock_tensor.numpy())) + + # Submit store job + tier.submit_store( + JobMetadata( + job_id=1, + keys=blocks, + spec=CPULoadStoreSpec([0, 1]), + ) + ) + + # Blocks should be in-flight + assert tier.get_num_in_flight() == 2 + 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 + assert tier.get_num_in_flight() == 0 + + +class TestTieringOffloadingManager: + """Tests for TieringOffloadingManager.""" + + @pytest.fixture + def manager_setup(self): + # Create primary tier (CPU-based) + self.primary_tier = CPUPrimaryTierOffloadingManager(block_size=16, num_blocks=5) + + # Mock get_primary_kv_tensor to return test tensor + # Create mock CPU tensor (5 blocks, 16 bytes per block) + mock_cpu_tensor = torch.zeros((5, 16), dtype=torch.float32) + self.primary_tier.get_primary_kv_tensor = lambda: mock_cpu_tensor + + # Create secondary tiers + self.secondary_tier1 = DummySecondaryTier(tier_name="Storage", max_blocks=10) + self.secondary_tier2 = DummySecondaryTier(tier_name="Network", max_blocks=10) + + # Create tiered manager + self.manager = TieringOffloadingManager( + primary_tier=self.primary_tier, + secondary_tiers=[self.secondary_tier1, self.secondary_tier2], + ) + + def test_basic_store_to_primary(self, manager_setup): + """Test basic store operation to primary tier.""" + blocks = to_keys(range(3)) + + # Prepare store + result = self.manager.prepare_store(blocks) + assert result is not None + assert len(result.keys_to_store) == 3 + + # Complete store + self.manager.complete_store(blocks, success=True) + + # Blocks should be in primary tier + assert self.primary_tier.lookup(blocks) == 3 + + def test_cascade_to_all_secondary_tiers(self, manager_setup): + """Test that blocks are cascaded to ALL secondary tiers.""" + blocks = to_keys(range(3)) + + # Store to primary + result = self.manager.prepare_store(blocks) + assert result is not None + + # Complete store (triggers cascade) + self.manager.complete_store(blocks, success=True) + + # Process finished jobs to complete cascade + self.manager._process_finished_jobs() + + # Blocks should be in both secondary tiers + assert self.secondary_tier1.get_num_blocks() == 3 + assert self.secondary_tier2.get_num_blocks() == 3 + + # Verify blocks are present + assert self.secondary_tier1.lookup(blocks) == 3 + assert self.secondary_tier2.lookup(blocks) == 3 + + def test_ref_cnt_protection_during_cascade(self, manager_setup): + """Test that ref_cnt protects blocks during cascade.""" + blocks = to_keys(range(3)) + + # Store to primary + result = self.manager.prepare_store(blocks) + assert result is not None + self.manager.complete_store(blocks, success=True) + + # After complete_store, blocks should have ref_cnt > 0 + # (one for each secondary tier) + for block_hash in blocks: + block = self.primary_tier._policy.get(block_hash) + # ref_cnt should be 2 (one for each secondary tier) + assert block.ref_cnt == 2 + + # Process finished jobs to complete cascade + self.manager._process_finished_jobs() + + # After cascade completes, ref_cnt should be 0 + for block_hash in blocks: + block = self.primary_tier._policy.get(block_hash) + assert block.ref_cnt == 0 + + def test_lookup_from_primary(self, manager_setup): + """Test lookup when blocks are in primary tier.""" + blocks = to_keys(range(3)) + + # Store blocks + self.manager.prepare_store(blocks) + self.manager.complete_store(blocks, success=True) + + # Lookup should find all blocks in primary + assert self.manager.lookup(blocks) == 3 + + def test_promotion_from_secondary(self, manager_setup): + """Test promotion of blocks from secondary to primary tier.""" + blocks = to_keys(range(3)) + + # Manually add blocks to secondary tier (simulate previous cascade) + for block in blocks: + self.secondary_tier1.blocks[block] = True + + # Lookup should initiate promotion + result = self.manager.lookup(blocks) + assert result is None # Retry later + + # Process finished jobs to complete promotion + self.manager._process_finished_jobs() + + # Now blocks should be in primary tier + assert self.primary_tier.lookup(blocks) == 3 + + # Next lookup should succeed + assert self.manager.lookup(blocks) == 3 + + def test_partial_lookup(self, manager_setup): + """Test lookup with partial hits.""" + blocks = to_keys(range(5)) + + # Store first 3 blocks to primary + self.manager.prepare_store(blocks[:3]) + self.manager.complete_store(blocks[:3], success=True) + + # Lookup all 5 blocks should return 3 (first 3 found) + assert self.manager.lookup(blocks) == 3 + + def test_eviction_in_primary_tier(self, manager_setup): + """Test eviction in primary tier when capacity is exceeded.""" + # Primary tier has capacity of 5 blocks + # First, fill the primary tier + blocks = to_keys(range(5)) + result = self.manager.prepare_store(blocks) + assert result is not None + assert len(result.keys_to_store) == 5 + self.manager.complete_store(blocks, success=True) + + # Process finished jobs to release ref_cnt from cascade + self.manager._process_finished_jobs() + + # Now try to store 2 more blocks (should trigger eviction) + more_blocks = to_keys(range(5, 7)) + result = self.manager.prepare_store(more_blocks) + + # Should evict 2 blocks from primary tier + assert result is not None + assert len(result.evicted_keys) == 2 + assert len(result.keys_to_store) == 2 + + def test_touch_propagates_to_all_tiers(self, manager_setup): + """Test that touch() propagates to all tiers.""" + blocks = to_keys(range(3)) + + # Store blocks + self.manager.prepare_store(blocks) + self.manager.complete_store(blocks, success=True) + self.manager._process_finished_jobs() + + # Touch blocks + self.manager.touch(blocks) + + # 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)) + + def test_failed_store_no_cascade(self, manager_setup): + """Test that failed GPU→primary store doesn't cascade.""" + blocks = to_keys(range(3)) + + # Prepare store + result = self.manager.prepare_store(blocks) + assert result is not None + + # Complete store with failure + self.manager.complete_store(blocks, success=False) + + # Process finished jobs + self.manager._process_finished_jobs() + + # Blocks should NOT be in secondary tiers + assert self.secondary_tier1.get_num_blocks() == 0 + assert self.secondary_tier2.get_num_blocks() == 0 + + def test_multiple_secondary_tiers_independent_eviction(self): + """Test that secondary tiers manage their own evictions.""" + # Create tier with small capacity + small_tier = DummySecondaryTier( + tier_name="SmallStorage", max_blocks=5, simulate_async=False + ) + large_tier = DummySecondaryTier( + tier_name="LargeStorage", max_blocks=10, simulate_async=False + ) + + # Create a fresh primary tier for this test + primary_tier = CPUPrimaryTierOffloadingManager(block_size=16, num_blocks=10) + + # Mock get_primary_kv_tensor to return test tensor + mock_cpu_tensor = torch.zeros((10, 16), dtype=torch.float32) + primary_tier.get_primary_kv_tensor = lambda: mock_cpu_tensor + + 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) + assert result is not None + manager.complete_store(blocks1, success=True) + manager._process_finished_jobs() + + # 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) + assert result is not None + manager.complete_store(blocks2, success=True) + manager._process_finished_jobs() + + # 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_prepare_store_processes_finished_jobs_first(self, manager_setup): + """Test that prepare_store() calls _process_finished_jobs() first.""" + blocks = to_keys(range(3)) + + # Store blocks + self.manager.prepare_store(blocks) + self.manager.complete_store(blocks, success=True) + + # Blocks should have ref_cnt = 2 (one for each secondary tier) + for block_hash in blocks: + block = self.primary_tier._policy.get(block_hash) + assert block.ref_cnt == 2 + + # Call prepare_store again (should process finished jobs first) + more_blocks = to_keys(range(3, 5)) + self.manager.prepare_store(more_blocks) + + # Original blocks should now have ref_cnt = 0 + for block_hash in blocks: + block = self.primary_tier._policy.get(block_hash) + assert block.ref_cnt == 0 + + +class TestTieringOffloadingWithoutSecondaryTiers: + """Test TieringOffloadingManager with no secondary tiers (backward compat).""" + + def test_works_without_secondary_tiers(self): + """Test that manager works with empty secondary_tiers list.""" + primary_tier = CPUPrimaryTierOffloadingManager(block_size=16, num_blocks=5) + + # Mock get_primary_kv_tensor to return test tensor + mock_cpu_tensor = torch.zeros((5, 16), dtype=torch.float32) + primary_tier.get_primary_kv_tensor = lambda: mock_cpu_tensor + + # Create manager with no secondary tiers + manager = TieringOffloadingManager( + primary_tier=primary_tier, secondary_tiers=[] + ) + + blocks = to_keys(range(3)) + + # Should work like a regular OffloadingManager + result = manager.prepare_store(blocks) + assert result is not None + manager.complete_store(blocks, success=True) + + assert manager.lookup(blocks) == 3 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/vllm/v1/kv_offload/abstract.py b/vllm/v1/kv_offload/abstract.py index 8f809ceaa08a..8dacd2c5aa71 100644 --- a/vllm/v1/kv_offload/abstract.py +++ b/vllm/v1/kv_offload/abstract.py @@ -57,6 +57,10 @@ class ReqContext: kv_transfer_params: dict[str, Any] | None = None +# Type alias for job IDs used in async transfer tracking +JobId = int + + class LoadStoreSpec(ABC): """ Abstract metadata that encapsulates information allowing a worker @@ -88,6 +92,23 @@ class OffloadingEvent: removed: bool +@dataclass +class JobMetadata: + """Metadata for an in-flight async transfer job.""" + + job_id: JobId + keys: list[OffloadKey] + spec: LoadStoreSpec + + +@dataclass +class JobResult: + """Result of an async transfer job (successful or failed).""" + + job_id: JobId + success: bool + + class OffloadingManager(ABC): @abstractmethod def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: @@ -174,7 +195,7 @@ def complete_store(self, keys: Iterable[OffloadKey], success: bool = True): """ Marks blocks which were previously prepared to be stored, as stored. Following this call, the blocks become loadable. - If if_success is False, blocks that were not marked as stored will be + If success is False, blocks that were not marked as stored will be removed. Args: @@ -195,3 +216,135 @@ def take_events(self) -> Iterable[OffloadingEvent]: def shutdown(self) -> None: """Shutdown the manager and release any resources.""" return + + +class SecondaryTierManager(ABC): + """ + Abstract interface for managing a single non-primary offloading tier. + + Secondary tiers cannot directly access GPU memory. All data transfers + must go through the primary tier (implemented as CPU in current version): + - Store: GPU → primary → secondary (cascade) + - Load: secondary → primary → GPU (promotion) + + IMPORTANT: All methods run in the Scheduler process and must be + lightweight and non-blocking. submit_load() and submit_store() submit + async jobs; get_finished() polls for completion. + """ + + @abstractmethod + def lookup(self, keys: Iterable[OffloadKey]) -> int | None: + """ + Check which blocks exist in this secondary tier. + + Args: + keys: Offload keys to look up. + + Returns: + Number of consecutive blocks (from start) that are present and ready, + or None if blocks are being transferred (retry later). + """ + pass + + @abstractmethod + def submit_store(self, job_metadata: JobMetadata) -> None: + """ + Submit an async job to store blocks from the primary tier to this + secondary tier. + + This method is lightweight: it allocates metadata and submits the + transfer job, but does NOT perform the actual data transfer on the + calling thread. + + The caller (TieringOffloadingManager) must have already called + primary.prepare_read(keys) to obtain job_metadata.spec and + to increment ref_cnt on those blocks. ref_cnt will be decremented + when get_finished() reports this job_id as complete and + primary.unprepare_read() is called. + + This method is responsible for: + 1. Filtering out blocks already present in this secondary tier + 2. Evicting blocks from this secondary tier if needed (secondary + tiers are responsible for their own evictions) + 3. Allocating space in this secondary tier + 4. Submitting the async transfer: primary → secondary + + Args: + job_metadata: Job metadata including job_id, keys, and + spec for reading blocks from the primary tier + (obtained via primary.prepare_read()). + spec is a CPULoadStoreSpec with block_ids. + """ + pass + + @abstractmethod + def submit_load(self, job_metadata: JobMetadata) -> None: + """ + Submit an async job to load blocks from this secondary tier to the + primary tier. + + This method is lightweight: it marks blocks as in-flight and submits + the transfer job, but does NOT perform the actual data transfer on + the calling thread. + + The caller (TieringOffloadingManager) must have already called + primary.prepare_write(keys) to obtain job_metadata.spec and + to allocate space in the primary tier. When get_finished() reports + this job_id as complete, primary.complete_write() is called to make + the blocks available for GPU loads. + + Args: + job_metadata: Job metadata including job_id, keys, and + spec for writing blocks into the primary tier + (obtained via primary.prepare_write()). + spec is a CPULoadStoreSpec with block_ids. + """ + pass + + @abstractmethod + def get_finished(self) -> Iterable[JobResult]: + """ + Poll for finished async jobs (both loads and stores). + + This is the mechanism by which the TieringOffloadingManager learns + that a transfer has finished and can: + - Call primary.unprepare_read() to decrement ref_cnt (for stores) + - Call primary.complete_write() to make blocks loadable (for loads) + + Returns: + Iterable of JobResult objects for all jobs that have + finished since the last call. + """ + pass + + def set_primary_view(self, view: memoryview) -> None: + """ + Provide a long-lived memoryview of the primary-tier CPU tensor. + + Called once by TieringOffloadingManager during initialisation. + Override to store the view for use in `submit_store` and `submit_load`. + Use `view.strides[0]` to obtain the byte stride between block slots. + + Args: + view: Memoryview of the primary tier's CPU KV cache tensor. + """ + return + + def touch(self, keys: Iterable[OffloadKey]): + """ + Mark blocks as recently used for eviction policy. + + Args: + keys: Offload keys to mark as recently used. + """ + return + + @abstractmethod + def get_tier_name(self) -> str: + """ + Get the name of this tier (e.g., "Storage", "Network"). + + Returns: + Tier name string. + """ + pass diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 8d4b744a0b4f..f3b606b2020c 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -40,6 +40,10 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): if kv_bytes_per_offloaded_block > 0 else 0 ) + world_size = vllm_config.parallel_config.world_size + self.cpu_page_size_per_worker: int = ( + kv_bytes_per_offloaded_block // world_size if world_size > 0 else 0 + ) # scheduler-side self._manager: OffloadingManager | None = None @@ -49,6 +53,22 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): self.eviction_policy: str = self.extra_config.get("eviction_policy", "lru") + def _maybe_apply_store_filter( + self, manager: OffloadingManager + ) -> OffloadingManager: + # store_threshold: how many times a block must appear in lookup() + # before it is eligible for CPU offloading. Values < 2 disable + # filtering (a threshold of 1 equals no filter; 0 is the default). + store_threshold = int(self.extra_config.get("store_threshold", 0)) + if store_threshold >= 2: + max_tracker_size = int(self.extra_config.get("max_tracker_size", 64_000)) + return FilterReusedOffloadingManager( + backing=manager, + store_threshold=store_threshold, + max_tracker_size=max_tracker_size, + ) + return manager + def get_manager(self) -> OffloadingManager: if not self._manager: kv_events_config = self.vllm_config.kv_events_config @@ -61,22 +81,18 @@ def get_manager(self) -> OffloadingManager: cache_policy=self.eviction_policy, # type: ignore[arg-type] enable_events=enable_events, ) - - # store_threshold: how many times a block must appear in lookup() - # before it is eligible for CPU offloading. Values < 2 disable - # filtering (a threshold of 1 equals no filter; 0 is the default). - store_threshold = int(self.extra_config.get("store_threshold", 0)) - if store_threshold >= 2: - max_tracker_size = int( - self.extra_config.get("max_tracker_size", 64_000) - ) - self._manager = FilterReusedOffloadingManager( - backing=self._manager, - store_threshold=store_threshold, - max_tracker_size=max_tracker_size, - ) + self._manager = self._maybe_apply_store_filter(self._manager) return self._manager + def _create_handlers( + self, kv_caches: CanonicalKVCaches + ) -> CpuGpuOffloadingHandlers: + return CpuGpuOffloadingHandlers( + kv_caches=kv_caches, + block_size_factor=self.block_size_factor, + num_cpu_blocks=self.num_blocks, + ) + def get_handlers( self, kv_caches: CanonicalKVCaches ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]: @@ -85,12 +101,7 @@ def get_handlers( raise Exception( "CPU Offloading is currently only supported on CUDA-alike GPUs" ) - - self._handlers = CpuGpuOffloadingHandlers( - kv_caches=kv_caches, - block_size_factor=self.block_size_factor, - num_cpu_blocks=self.num_blocks, - ) + 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/factory.py b/vllm/v1/kv_offload/factory.py index ecbaebb0d967..f3be7b2d59a8 100644 --- a/vllm/v1/kv_offload/factory.py +++ b/vllm/v1/kv_offload/factory.py @@ -56,3 +56,8 @@ def create_spec( OffloadingSpecFactory.register_spec( "CPUOffloadingSpec", "vllm.v1.kv_offload.cpu.spec", "CPUOffloadingSpec" ) +OffloadingSpecFactory.register_spec( + "TieringOffloadingSpec", + "vllm.v1.kv_offload.tiering.spec", + "TieringOffloadingSpec", +) diff --git a/vllm/v1/kv_offload/secondary_tiers/__init__.py b/vllm/v1/kv_offload/secondary_tiers/__init__.py new file mode 100644 index 000000000000..c4b5a473aa7b --- /dev/null +++ b/vllm/v1/kv_offload/secondary_tiers/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Secondary tier implementations for KV cache offloading. + +This package contains various secondary tier storage backends that can be used +with the TieringOffloadingManager for multi-tier KV cache management. +""" + +from vllm.v1.kv_offload.secondary_tiers.dummy import ( + DummyLoadStoreSpec, + DummySecondaryTier, +) + +__all__ = [ + "DummyLoadStoreSpec", + "DummySecondaryTier", +] diff --git a/vllm/v1/kv_offload/secondary_tiers/dummy.py b/vllm/v1/kv_offload/secondary_tiers/dummy.py new file mode 100644 index 000000000000..f75edafd6187 --- /dev/null +++ b/vllm/v1/kv_offload/secondary_tiers/dummy.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +DummySecondaryTier: 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 Iterable +from dataclasses import dataclass + +from vllm.v1.kv_offload.abstract import ( + JobId, + JobMetadata, + JobResult, + LoadStoreSpec, + OffloadKey, + SecondaryTierManager, +) +from vllm.v1.kv_offload.mediums import CPULoadStoreSpec + + +@dataclass +class _JobMetadata: + """Internal metadata for tracking job details.""" + + job_id: JobId + keys: list[OffloadKey] + is_store: bool # True for store jobs, False for load jobs + + +class DummyLoadStoreSpec(LoadStoreSpec): + """ + Minimal LoadStoreSpec for DummySecondaryTier testing. + + This spec is never actually used for I/O since DummySecondaryTier + stores blocks in memory. It exists to provide proper type semantics + and serve as a template for real secondary tier implementations. + """ + + @staticmethod + def medium() -> str: + return "Dummy" + + +class DummySecondaryTier(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 + - Tracks in-flight transfers to return None from lookup() + """ + + def __init__( + self, + tier_name: str = "DummyStorage", + max_blocks: int = 1000, + simulate_async: bool = False, + ): + """ + Initialize the dummy secondary tier. + + Args: + tier_name: Name of this tier (for identification) + 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. + """ + self.tier_name = tier_name + self.max_blocks = max_blocks + self.simulate_async = simulate_async + + self._primary_view: memoryview | None = None + + # key -> True (only care about presence) + self.blocks: OrderedDict[OffloadKey, bool] = OrderedDict() + + # Tracks in-flight transfers: key -> job_id + self.in_flight: dict[OffloadKey, JobId] = {} + + # 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 set_primary_view(self, view: memoryview) -> None: + self._primary_view = view + + def lookup(self, keys: Iterable[OffloadKey]) -> int | None: + """ + Check which blocks exist in this secondary tier. + + Args: + keys: Block hashes to look up. + + Returns: + Number of consecutive blocks (from start) that are present and ready, + or None if blocks are being transferred (retry later). + """ + hit_count = 0 + for key in keys: + # Check if block is in-flight + if key in self.in_flight: + # Block is being transferred, return None (retry later) + return None + + # Check if block exists in this tier + if key not in self.blocks: + break + + hit_count += 1 + + return hit_count + + 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_list = list(job_metadata.keys) + primary_read_spec = job_metadata.spec + + # Validate spec type and consistency + assert isinstance(primary_read_spec, CPULoadStoreSpec), ( + f"Expected CPULoadStoreSpec, got {type(primary_read_spec)}" + ) + assert len(keys_list) == len(primary_read_spec.block_ids), ( + f"Length mismatch: {len(keys_list)} keys but " + f"{len(primary_read_spec.block_ids)} block_ids in spec" + ) + + # Filter out blocks already present + blocks_to_store = [bh for bh in keys_list 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_list) + for key in self.blocks: + if key not in protected and key not in self.in_flight: + 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] + + # Mark blocks as in-flight + for key in blocks_to_store: + self.in_flight[key] = job_id + + # 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_list = list(job_metadata.keys) + primary_write_spec = job_metadata.spec + + # Validate spec type and consistency + assert isinstance(primary_write_spec, CPULoadStoreSpec), ( + f"Expected CPULoadStoreSpec, got {type(primary_write_spec)}" + ) + assert len(keys_list) == len(primary_write_spec.block_ids), ( + f"Length mismatch: {len(keys_list)} keys but " + f"{len(primary_write_spec.block_ids)} block_ids in spec" + ) + + # Verify all blocks exist + for key in keys_list: + if key not in self.blocks: + return + + # Mark blocks as in-flight + for key in keys_list: + self.in_flight[key] = job_id + + # Create internal job metadata + internal_job_metadata = _JobMetadata( + job_id=job_id, keys=keys_list, 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 + del self.in_flight[key] + # 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 by removing in-flight markers.""" + for key in job_metadata.keys: + del self.in_flight[key] + # 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: Iterable[OffloadKey]): + """ + Mark blocks as recently used (move to end of LRU list). + + Args: + keys: Blocks to mark as recently used. + """ + for key in reversed(list(keys)): + if key in self.blocks: + self.blocks.move_to_end(key) + + def get_tier_name(self) -> str: + """ + Get the name of this tier. + + Returns: + Tier name string. + """ + return self.tier_name + + def get_num_blocks(self) -> int: + """Get the number of blocks currently stored in this tier.""" + return len(self.blocks) + + def get_num_in_flight(self) -> int: + """Get the number of blocks currently in-flight.""" + return len(self.in_flight) + + def clear(self): + """Clear all blocks and in-flight transfers (for testing).""" + self.blocks.clear() + self.in_flight.clear() + self.completed_jobs.clear() + self.pending_jobs.clear() diff --git a/vllm/v1/kv_offload/tiering/__init__.py b/vllm/v1/kv_offload/tiering/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py new file mode 100644 index 000000000000..91c62bbd01b8 --- /dev/null +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -0,0 +1,455 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +TieringOffloadingManager: Multi-tier KV cache offloading orchestrator. + +This manager coordinates between a primary tier (with GPU access, currently +CPU-based) and zero or more secondary tiers (Storage, Network, etc.) to +provide hierarchical KV cache offloading. + +Key Design Principles: +1. Always offload to all tiers — When a block is stored to the primary tier, + it is cascaded to ALL secondary tiers +2. Primary tier is the gateway — Only the primary tier can directly access + GPU memory (currently implemented using CPU memory) +3. Staged promotion — Blocks in secondary tiers must be promoted to the + primary tier before GPU can access them +4. Transparent retry mechanism — Return None from lookup() to signal + "data is being promoted, try later" +5. ref_cnt as eviction protection — primary.prepare_read() increments ref_cnt, + protecting blocks from eviction until complete_read() is called +""" + +from collections.abc import Iterable + +import torch + +from vllm.logger import init_logger +from vllm.v1.kv_offload.abstract import ( + JobId, + JobMetadata, + LoadStoreSpec, + OffloadingEvent, + OffloadingManager, + OffloadKey, + PrepareStoreOutput, + SecondaryTierManager, +) +from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion + +logger = init_logger(__name__) + + +class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): + """CPUOffloadingManager with a primary/secondary transfer interface. + + The inherited prepare_store/complete_store/prepare_load/complete_load are the + GPU-facing OffloadingManager interface. These aliases expose the same operations + from the secondary tier perspective, where read/write refers to secondary + accessing primary. This avoids confusion when reading TieringOffloadingManager + code (e.g. calling prepare_load inside a cascade/store path would be misleading). + """ + + def __init__( + self, + block_size: int, + num_blocks: int, + cache_policy: str = "lru", + enable_events: bool = False, + mmap_region: SharedOffloadRegion | None = None, + ): + super().__init__( + block_size=block_size, + num_blocks=num_blocks, + cache_policy=cache_policy, # type: ignore[arg-type] + enable_events=enable_events, + ) + self._mmap_region = mmap_region + + def prepare_write(self, keys) -> PrepareStoreOutput | None: + """Allocate space in primary for a secondary->primary write (promotion).""" + return self.prepare_store(keys) + + def complete_write(self, keys, success: bool = True) -> None: + """Finalize secondary->primary write, making blocks available.""" + self.complete_store(keys, success) + + def prepare_read(self, keys) -> LoadStoreSpec: + """Protect primary blocks for a primary->secondary read (cascade), + incrementing ref_cnt.""" + return self.prepare_load(keys) + + def complete_read(self, keys) -> None: + """Release protection after primary->secondary read completes, + decrementing ref_cnt.""" + self.complete_load(keys) + + def get_primary_kv_tensor(self) -> torch.Tensor: + """ + Get the primary tier's KV cache tensor. + + Returns a 2-D int8 tensor of shape (num_blocks, row_stride_bytes) + backed by the SharedOffloadRegion mmap, where row_stride_bytes = + cpu_page_size * world_size. Secondary tiers address block b as + view[b], and view.strides[0] gives the per-block byte stride. + + Returns: + 2-D int8 CPU tensor of shape (num_blocks, row_stride_bytes). + """ + assert self._mmap_region is not None, ( + "mmap_region must be provided to CPUPrimaryTierOffloadingManager" + ) + return self._mmap_region._base.view( + self._mmap_region.num_blocks, self._mmap_region._row_stride + ) + + +class TieringOffloadingManager(OffloadingManager): + """ + Orchestrates multi-tier KV cache offloading. + + This manager coordinates between a primary tier (with GPU access, currently + CPU-based) and zero or more secondary tiers (Storage, Network, etc.) to + provide hierarchical KV cache offloading. + + Key internal state: + - Minimal state tracking; relies on secondary tiers to report completion + via get_finished() + - Secondary tiers return JobResult objects containing all necessary + information + - job_id_counter: monotonically increasing counter for job IDs + """ + + def __init__( + self, + primary_tier: CPUPrimaryTierOffloadingManager, + secondary_tiers: list[SecondaryTierManager] | None = None, + enable_events: bool = False, + ): + """ + Initialize the TieringOffloadingManager. + + Args: + primary_tier: The primary tier manager (CPU-based). + secondary_tiers: List of secondary tier managers (e.g., Storage, + Network). Can be None or empty list. + enable_events: Whether to track offloading events + """ + self.primary_tier: CPUPrimaryTierOffloadingManager = primary_tier + self.secondary_tiers = secondary_tiers or [] + + self._job_id_counter: int = 0 + self.events: list[OffloadingEvent] | None = [] if enable_events else None + + # Job tracking: maps job_id to metadata for each transfer direction + # Store jobs: primary → secondary transfers + self._store_jobs: dict[JobId, JobMetadata] = {} + # Load jobs: secondary → primary transfers (promotions) + self._load_jobs: dict[JobId, JobMetadata] = {} + + # Wire each secondary tier with a long-lived memoryview of the primary + # CPU tensor (one independent view per tier). Views are stored so they + # can be released on shutdown(). + self._secondary_views: list[memoryview] = [] + cpu_tensor = primary_tier.get_primary_kv_tensor() + for tier in self.secondary_tiers: + view = memoryview(cpu_tensor.numpy()) + self._secondary_views.append(view) + tier.set_primary_view(view) + + def _next_job_id(self) -> JobId: + """Generate a unique job ID for async transfer tracking.""" + job_id = self._job_id_counter + self._job_id_counter += 1 + return job_id + + def _process_finished_jobs(self): + """ + Poll all secondary tiers for completed jobs and update state accordingly. + + This method: + 1. Calls get_finished() on each secondary tier + 2. For completed stores (primary→secondary): calls primary.complete_read() + to decrement ref_cnt + 3. For completed loads (secondary→primary): calls primary.complete_write() + to make blocks available + """ + for tier in self.secondary_tiers: + for completed_job in tier.get_finished(): + job_id = completed_job.job_id + + # Determine job type by checking which dictionary contains the job_id + if job_id in self._store_jobs: + # primary→secondary transfer completed. + # Decrement ref_cnt on primary blocks. + job_metadata = self._store_jobs.pop(job_id) + self.primary_tier.complete_read(job_metadata.keys) + elif job_id in self._load_jobs: + # secondary→primary transfer (promotion) completed. + # Make blocks available in primary tier. + job_metadata = self._load_jobs.pop(job_id) + self.primary_tier.complete_write( + job_metadata.keys, completed_job.success + ) + else: + # Job ID not found in either dictionary - this shouldn't happen + logger.error( + "Received finished job for unknown job_id %d from tier %s", + job_id, + tier.get_tier_name(), + ) + + def lookup(self, keys: Iterable[OffloadKey]) -> int | None: + """ + Find the length of the maximal series of blocks that are offloaded. + + Algorithm: + 1. Check primary tier first + 2. If not all blocks found, check all secondary tiers sequentially, + promoting blocks from each tier that has hits and updating the + remaining blocks to search for + 3. Return None to signal "retry later" if any promotions were initiated + + Args: + keys: Block hashes to look up. + + Returns: + Number of consecutive blocks (from start) that are present, + or None if blocks are being transferred (retry later). + """ + # Process any completed async jobs first to ensure promoted blocks + # are finalized and available in the primary tier + self._process_finished_jobs() + + keys_list = list(keys) + + # Step 1: Check primary tier + primary_hits = self.primary_tier.lookup(keys_list) + + if primary_hits is None: + # Primary tier is busy (blocks being transferred) + return None + + if primary_hits == len(keys_list): + # All blocks in primary tier + return primary_hits + + # Step 2: Check all secondary tiers for remaining blocks + remaining_keys = keys_list[primary_hits:] + + # Track whether any promotions were initiated + has_promotions = False + + for tier in self.secondary_tiers: + if not remaining_keys: + # All blocks have been found + break + + secondary_hits = tier.lookup(remaining_keys) + + # Skip if tier is busy (None) or has no hits (0) + if not secondary_hits: + continue + + # Found blocks in this secondary tier, initiate promotion + blocks_to_promote = remaining_keys[:secondary_hits] + self._initiate_promotion(tier, blocks_to_promote) + has_promotions = True + + # Update remaining_keys to continue searching for the rest + remaining_keys = remaining_keys[secondary_hits:] + + # Step 3: If any promotions were initiated, return None to signal retry + if has_promotions: + return None + + # No more blocks found in any tier + return primary_hits + + def _initiate_promotion(self, tier: SecondaryTierManager, keys: list[OffloadKey]): + """ + Initiate promotion of blocks from a secondary tier to the primary tier. + + This method: + 1. Calls primary.prepare_write() to allocate space in primary tier + 2. Calls tier.submit_load() to start async transfer: secondary→primary + 3. Tracks the job in _load_jobs dictionary + + Args: + tier: The secondary tier to promote from + keys: Blocks to promote + """ + # Allocate space in primary tier for promoted blocks + primary_store_result = self.primary_tier.prepare_write(keys) + + if primary_store_result is None: + # Cannot allocate space in primary tier (full) + # The next lookup() will retry + return + + # Submit async load job: secondary→primary + job_id = self._next_job_id() + + # Track this load job + job_metadata = JobMetadata( + job_id=job_id, + keys=keys, + spec=primary_store_result.store_spec, + ) + self._load_jobs[job_id] = job_metadata + + tier.submit_load(job_metadata) + + def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec: + """ + Prepare blocks to be loaded from primary tier to GPU. + + CRITICAL: This method calls _process_finished_jobs() FIRST to ensure + that any completed promotions have been finalized and blocks are ready. + + This increments ref_cnt on the blocks in the primary tier, protecting + them from eviction during the transfer. + + Args: + keys: Blocks to prepare for loading. + + Returns: + LoadStoreSpec for reading from primary tier. + """ + # Process completed promotions to ensure blocks are ready + self._process_finished_jobs() + + return self.primary_tier.prepare_load(keys) + + def touch(self, keys: Iterable[OffloadKey]): + """ + Mark blocks as recently used in all tiers. + + Args: + keys: Blocks to mark as recently used. + """ + keys = list(keys) + self.primary_tier.touch(keys) + for tier in self.secondary_tiers: + tier.touch(keys) + + def complete_load(self, keys: Iterable[OffloadKey]): + """ + Mark blocks as done loading from primary tier to GPU. + + This decrements ref_cnt on the blocks in the primary tier, allowing + them to be evicted again. + + Args: + keys: Blocks that finished loading. + """ + self.primary_tier.complete_load(keys) + + def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None: + """ + Prepare blocks to be stored from GPU to primary tier. + + CRITICAL: This method calls _process_finished_jobs() FIRST to ensure + that any completed async transfers have their ref_cnt decremented + before the primary tier makes eviction decisions. + + Args: + keys: Blocks to prepare for storing. + + Returns: + PrepareStoreOutput describing where to store blocks and what was + evicted, or None if store cannot proceed. + """ + # Step 1: Poll for completed async jobs FIRST + # This decrements ref_cnt on primary blocks that have been + # successfully transferred to secondary tiers. + self._process_finished_jobs() + + # Step 2: Store to primary tier + primary_result = self.primary_tier.prepare_store(keys) + + # Note: Secondary tier cascading will happen in complete_store() + # after the GPU→Primary transfer completes and blocks are ready. + + return primary_result + + def complete_store(self, keys: Iterable[OffloadKey], success: bool = True): + """ + Mark blocks as done storing from GPU to primary tier. + + This is where secondary tier cascading happens — after blocks are + confirmed to be in the primary tier, they are cascaded to ALL + secondary tiers. + + For each secondary tier: + 1. Call primary.prepare_read() to get LoadStoreSpec AND increment + ref_cnt (protecting blocks during async transfer) + 2. Call tier.submit_store() to start async transfer: primary→secondary + 3. Track the job in _store_jobs dictionary + + Args: + keys: Blocks that finished storing. + success: Whether the GPU→primary transfer succeeded. + """ + # Materialize only if success=True (needed for cascading to secondary tiers) + keys_list = list(keys) if success else keys + + # Step 1: Complete store in primary tier (makes blocks loadable) + self.primary_tier.complete_store(keys_list, success) + + if not success: + # If GPU→Primary transfer failed, don't cascade to secondary tiers + return + + # At this point, success=True is guaranteed, so keys_list + # is list[OffloadKey] + assert isinstance(keys_list, list) + + # Step 2: Cascade to ALL secondary tiers + # For each secondary tier, call primary.prepare_read() to get the + # LoadStoreSpec AND to increment ref_cnt (protecting blocks from + # eviction during the async transfer). One prepare_read() call per + # secondary tier. + for tier in self.secondary_tiers: + # Get spec for reading from primary tier AND increment ref_cnt + primary_blocks_spec = self.primary_tier.prepare_read(keys_list) + + # Submit async store job: primary→secondary + job_id = self._next_job_id() + + # Track this store job + job_metadata = JobMetadata( + job_id=job_id, keys=keys_list, spec=primary_blocks_spec + ) + self._store_jobs[job_id] = job_metadata + + tier.submit_store(job_metadata) + + # Note: The async transfers are now in flight. + # Their completion is tracked via get_finished() / _process_finished_jobs(). + + def take_events(self) -> Iterable[OffloadingEvent]: + """ + Take offloading events from the primary tier. + + Note: Currently only primary tier events are tracked. Secondary tier + events could be added in the future if needed. + + Yields: + New OffloadingEvents collected since the last call. + """ + if self.events is not None: + yield from self.events + self.events.clear() + + # Also yield events from primary tier + yield from self.primary_tier.take_events() + + def shutdown(self) -> None: + """Release memoryviews and scheduler-side mmap.""" + for view in self._secondary_views: + view.release() + if self.primary_tier._mmap_region is not None: + self.primary_tier._mmap_region.cleanup() + self.primary_tier._mmap_region = None diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py new file mode 100644 index 000000000000..7b11b5bbcfb3 --- /dev/null +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +TieringOffloadingSpec: Spec for multi-tier KV cache offloading. + +This spec creates a TieringOffloadingManager with a CPU-based primary tier +and configurable secondary tiers (e.g., Storage, Network). + +Configuration via kv_connector_extra_config: + - cpu_bytes_to_use: (required) Bytes to allocate for CPU primary tier + - block_size: (optional) Block size for offloaded blocks (default: GPU block size) + - eviction_policy: (optional) Primary tier eviction policy: "lru" or + "arc" (default: "lru") + - store_threshold: (optional) How many times a block must appear in lookup() + before it is eligible for CPU offloading. Values < 2 disable filtering + (default: 0) + - max_tracker_size: (optional) Maximum number of blocks tracked for + store_threshold filtering (default: 64000) + - secondary_tiers: (optional) List of secondary tier configurations + Each secondary tier config is a dict with: + - type: (required) Type of secondary tier (e.g., "dummy", "storage", "network") + - tier_name: (required) Name for this tier (used for logging and identification) + - Additional tier-specific parameters are passed directly to the tier + constructor. See each tier's documentation for supported parameters. + +Example configuration: +{ + "cpu_bytes_to_use": 10737418240, # 10 GB + "block_size": 16, + "eviction_policy": "lru", + "secondary_tiers": [ + { + "type": "dummy", + "tier_name": "TestStorage", + # Tier-specific parameters (for DummySecondaryTier): + "max_blocks": 10000, + "simulate_async": False + } + ] +} +""" + +import torch + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.kv_offload.abstract import OffloadingManager +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.secondary_tiers.dummy import DummySecondaryTier +from vllm.v1.kv_offload.spec import CanonicalKVCaches +from vllm.v1.kv_offload.tiering.manager import ( + CPUPrimaryTierOffloadingManager, + TieringOffloadingManager, +) +from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers + +logger = init_logger(__name__) + + +class TieringOffloadingSpec(CPUOffloadingSpec): + """ + Spec for multi-tier KV cache offloading. + + Creates a TieringOffloadingManager with: + - Primary tier: CPU-based (LRU or ARC eviction policy) + - Secondary tiers: Configurable via extra_config + + The primary tier has direct GPU access and serves as the gateway for all + GPU↔offload operations. Secondary tiers cannot directly access GPU memory + and must coordinate with the primary tier for data transfers. + """ + + def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): + super().__init__(vllm_config, kv_cache_config) + + # Parse secondary tier configurations + self.secondary_tier_configs = self.extra_config.get("secondary_tiers", []) + if not isinstance(self.secondary_tier_configs, list): + raise ValueError("secondary_tiers must be a list of tier configurations") + + # Scheduler-side mmap (rank=None); kept for cleanup + self._scheduler_mmap: SharedOffloadRegion | None = None + + def _create_secondary_tier(self, tier_config: dict): + """ + Create a secondary tier from configuration. + + Args: + tier_config: Dictionary with tier configuration containing: + - type (required): Type of secondary tier (e.g., "dummy") + - tier_name (required): Name for this tier + - Additional tier-specific parameters are passed directly + to the tier constructor + + Returns: + SecondaryTierManager instance + + Raises: + ValueError: If tier type is unknown or configuration is invalid + """ + # Make a copy to avoid modifying the original config + config = tier_config.copy() + + # Extract common parameters + tier_type = config.pop("type", None) + if not tier_type: + raise ValueError("Secondary tier configuration must include 'type'") + + tier_name = config.pop("tier_name", None) + if not tier_name: + raise ValueError("Secondary tier configuration must include 'tier_name'") + + # Remaining parameters in config are tier-specific + if tier_type == "dummy": + # DummySecondaryTier for testing + # Pass tier_name and tier-specific params to constructor + return DummySecondaryTier(tier_name=tier_name, **config) + else: + raise ValueError( + f"Unknown secondary tier type: {tier_type}. Supported types: dummy" + ) + + def get_manager(self) -> OffloadingManager: + """ + Get the TieringOffloadingManager. + + Creates a TieringOffloadingManager with: + - Primary tier: CPU-based (LRU or ARC) + - Secondary tiers: As configured in extra_config + + Returns: + TieringOffloadingManager instance + """ + if not self._manager: + kv_events_config = self.vllm_config.kv_events_config + enable_events = ( + kv_events_config is not None and kv_events_config.enable_kv_cache_events + ) + + # Create scheduler-side SharedOffloadRegion (rank=None) first so + # CPUPrimaryTierOffloadingManager can return _base from + # get_primary_kv_tensor(), which TieringOffloadingManager.__init__ + # calls immediately to wire secondary tier memoryviews. + world_size = self.vllm_config.parallel_config.world_size + scheduler_mmap = SharedOffloadRegion( + instance_id=self.vllm_config.instance_id, + total_size_bytes=self.cpu_page_size_per_worker + * world_size + * self.num_blocks, + num_blocks=self.num_blocks, + rank=None, + num_workers=world_size, + cpu_page_size=self.cpu_page_size_per_worker, + ) + self._scheduler_mmap = scheduler_mmap + + # Create primary tier (CPU-based) + assert len(self.gpu_block_size) == 1 + offloaded_block_size = self.gpu_block_size[0] * self.block_size_factor + primary_tier = CPUPrimaryTierOffloadingManager( + block_size=offloaded_block_size, + num_blocks=self.num_blocks, + cache_policy=self.eviction_policy, # type: ignore[arg-type] + enable_events=enable_events, + mmap_region=scheduler_mmap, + ) + + # Create secondary tiers + secondary_tiers = [] + for tier_config in self.secondary_tier_configs: + try: + tier = self._create_secondary_tier(tier_config) + secondary_tiers.append(tier) + logger.info( + "Created secondary tier: %s (type: %s)", + tier.get_tier_name(), + tier_config.get("type"), + ) + except Exception as e: + logger.error( + "Failed to create secondary tier from config %s: %s", + tier_config, + e, + ) + raise + + # Create TieringOffloadingManager. GPU↔CPU transfers use the inherited + # get_handlers(); secondary tier transfers are handled by the + # secondary tier managers and need no additional handlers here. + tiering_manager = TieringOffloadingManager( + primary_tier=primary_tier, + secondary_tiers=secondary_tiers, + enable_events=enable_events, + ) + # PRNOTE: should the store_filter apply to the TieringOffloadingManager or + # to the primary CPU manager? + self._manager = self._maybe_apply_store_filter(tiering_manager) + + logger.info( + "Created TieringOffloadingManager with primary tier " + "(%s, %s blocks) and %s secondary tier(s)", + self.eviction_policy, + self.num_blocks, + len(secondary_tiers), + ) + + return self._manager + + 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( + instance_id=self.vllm_config.instance_id, + total_size_bytes=self.cpu_page_size_per_worker + * world_size + * self.num_blocks, + num_blocks=self.num_blocks, + rank=rank, + num_workers=world_size, + cpu_page_size=self.cpu_page_size_per_worker, + ) + return CpuGpuOffloadingHandlers( + kv_caches=kv_caches, + block_size_factor=self.block_size_factor, + num_cpu_blocks=self.num_blocks, + mmap_region=worker_mmap, + ) diff --git a/vllm/v1/kv_offload/worker/cpu_gpu.py b/vllm/v1/kv_offload/worker/cpu_gpu.py index dd12a533ede9..d886f4721de3 100644 --- a/vllm/v1/kv_offload/worker/cpu_gpu.py +++ b/vllm/v1/kv_offload/worker/cpu_gpu.py @@ -119,6 +119,7 @@ def __init__( block_size_factor: int, kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]], gpu_to_cpu: bool, + mmap_region: SharedOffloadRegion | None = None, ): """ Initialize a SingleDirectionOffloadingHandler. @@ -181,6 +182,8 @@ def __init__( self.group_block_size_in_bytes.append(group_block_size_in_bytes) self.transfer_type = ("GPU", "CPU") if self.gpu_to_cpu else ("CPU", "GPU") + # mmap_region to clean up on shutdown (gpu_to_cpu handler owns it) + self._mmap_region = mmap_region # job_id -> event self._transfer_events: dict[int, torch.Event] = {} # queue of transfers (job_id, stream, event) @@ -318,6 +321,9 @@ def shutdown(self) -> None: self._event_pool.clear() self.src_tensors.clear() self.dst_tensors.clear() + if self._mmap_region is not None: + self._mmap_region.cleanup() + self._mmap_region = None class CpuGpuOffloadingHandlers: @@ -370,6 +376,7 @@ def __init__( block_size_factor=block_size_factor, kv_cache_groups_data_refs=kv_caches.group_data_refs, gpu_to_cpu=True, + mmap_region=mmap_region, ) self.cpu_to_gpu_handler = SingleDirectionOffloadingHandler( From cc1ae5948979a746f53ab2b99b223aab2fd23c8e Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 20 Apr 2026 16:36:23 +0300 Subject: [PATCH 02/55] Fix rebase issues: remove block_size Signed-off-by: Ronen Schaffer --- tests/v1/kv_offload/test_tiering_offloading.py | 6 +++--- vllm/v1/kv_offload/tiering/manager.py | 2 -- vllm/v1/kv_offload/tiering/spec.py | 1 - 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index ded6ce8b004b..cb83cee63b77 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -139,7 +139,7 @@ class TestTieringOffloadingManager: @pytest.fixture def manager_setup(self): # Create primary tier (CPU-based) - self.primary_tier = CPUPrimaryTierOffloadingManager(block_size=16, num_blocks=5) + self.primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=5) # Mock get_primary_kv_tensor to return test tensor # Create mock CPU tensor (5 blocks, 16 bytes per block) @@ -335,7 +335,7 @@ def test_multiple_secondary_tiers_independent_eviction(self): ) # Create a fresh primary tier for this test - primary_tier = CPUPrimaryTierOffloadingManager(block_size=16, num_blocks=10) + primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=10) # Mock get_primary_kv_tensor to return test tensor mock_cpu_tensor = torch.zeros((10, 16), dtype=torch.float32) @@ -398,7 +398,7 @@ class TestTieringOffloadingWithoutSecondaryTiers: def test_works_without_secondary_tiers(self): """Test that manager works with empty secondary_tiers list.""" - primary_tier = CPUPrimaryTierOffloadingManager(block_size=16, num_blocks=5) + primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=5) # Mock get_primary_kv_tensor to return test tensor mock_cpu_tensor = torch.zeros((5, 16), dtype=torch.float32) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 91c62bbd01b8..8c1777506e0e 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -53,14 +53,12 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): def __init__( self, - block_size: int, num_blocks: int, cache_policy: str = "lru", enable_events: bool = False, mmap_region: SharedOffloadRegion | None = None, ): super().__init__( - block_size=block_size, num_blocks=num_blocks, cache_policy=cache_policy, # type: ignore[arg-type] enable_events=enable_events, diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 7b11b5bbcfb3..0b0b6ecabebc 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -160,7 +160,6 @@ def get_manager(self) -> OffloadingManager: assert len(self.gpu_block_size) == 1 offloaded_block_size = self.gpu_block_size[0] * self.block_size_factor primary_tier = CPUPrimaryTierOffloadingManager( - block_size=offloaded_block_size, num_blocks=self.num_blocks, cache_policy=self.eviction_policy, # type: ignore[arg-type] enable_events=enable_events, From c3d6cb68d5d091f484df62a6c729f36d7d54ab86 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 20 Apr 2026 17:26:45 +0300 Subject: [PATCH 03/55] Fix rebase issues: add ReqContext and update lookup() Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 73 +++++++---- vllm/v1/kv_offload/tiering/manager.py | 123 +++++++++--------- 2 files changed, 107 insertions(+), 89 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index cb83cee63b77..4fd9c8b08a8f 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -16,7 +16,12 @@ import pytest import torch -from vllm.v1.kv_offload.abstract import JobMetadata, OffloadKey, make_offload_key +from vllm.v1.kv_offload.abstract import ( + JobMetadata, + OffloadKey, + ReqContext, + make_offload_key, +) from vllm.v1.kv_offload.mediums import CPULoadStoreSpec from vllm.v1.kv_offload.secondary_tiers.dummy import DummySecondaryTier from vllm.v1.kv_offload.tiering.manager import ( @@ -24,11 +29,30 @@ TieringOffloadingManager, ) +_CTX = ReqContext() + def to_keys(int_ids: Iterable[int]) -> list[OffloadKey]: return [make_offload_key(str(i).encode(), 0) for i in int_ids] +def count_hits(manager, keys: list[OffloadKey]) -> int | None: + """Count consecutive lookup hits from the start of keys. + + Returns the count of leading True results, or None if any lookup + returns None (retry-later signal). + """ + count = 0 + for key in keys: + result = manager.lookup(key, _CTX) + if result is None: + return None + if not result: + break + count += 1 + return count + + class TestDummySecondaryTier: """Tests for DummySecondaryTier implementation.""" @@ -161,7 +185,7 @@ def test_basic_store_to_primary(self, manager_setup): blocks = to_keys(range(3)) # Prepare store - result = self.manager.prepare_store(blocks) + result = self.manager.prepare_store(blocks, _CTX) assert result is not None assert len(result.keys_to_store) == 3 @@ -169,14 +193,14 @@ def test_basic_store_to_primary(self, manager_setup): self.manager.complete_store(blocks, success=True) # Blocks should be in primary tier - assert self.primary_tier.lookup(blocks) == 3 + assert count_hits(self.primary_tier, blocks) == 3 def test_cascade_to_all_secondary_tiers(self, manager_setup): """Test that blocks are cascaded to ALL secondary tiers.""" blocks = to_keys(range(3)) # Store to primary - result = self.manager.prepare_store(blocks) + result = self.manager.prepare_store(blocks, _CTX) assert result is not None # Complete store (triggers cascade) @@ -198,7 +222,7 @@ def test_ref_cnt_protection_during_cascade(self, manager_setup): blocks = to_keys(range(3)) # Store to primary - result = self.manager.prepare_store(blocks) + result = self.manager.prepare_store(blocks, _CTX) assert result is not None self.manager.complete_store(blocks, success=True) @@ -222,11 +246,11 @@ def test_lookup_from_primary(self, manager_setup): blocks = to_keys(range(3)) # Store blocks - self.manager.prepare_store(blocks) + self.manager.prepare_store(blocks, _CTX) self.manager.complete_store(blocks, success=True) # Lookup should find all blocks in primary - assert self.manager.lookup(blocks) == 3 + assert count_hits(self.manager, blocks) == 3 def test_promotion_from_secondary(self, manager_setup): """Test promotion of blocks from secondary to primary tier.""" @@ -236,36 +260,37 @@ def test_promotion_from_secondary(self, manager_setup): for block in blocks: self.secondary_tier1.blocks[block] = True - # Lookup should initiate promotion - result = self.manager.lookup(blocks) - assert result is None # Retry later + # Lookup each block to initiate promotion for all of them + for block in blocks: + result = self.manager.lookup(block, _CTX) + assert result is None # Retry later (promotion initiated) # Process finished jobs to complete promotion self.manager._process_finished_jobs() # Now blocks should be in primary tier - assert self.primary_tier.lookup(blocks) == 3 + assert count_hits(self.primary_tier, blocks) == 3 # Next lookup should succeed - assert self.manager.lookup(blocks) == 3 + assert count_hits(self.manager, blocks) == 3 def test_partial_lookup(self, manager_setup): """Test lookup with partial hits.""" blocks = to_keys(range(5)) # Store first 3 blocks to primary - self.manager.prepare_store(blocks[:3]) + self.manager.prepare_store(blocks[:3], _CTX) self.manager.complete_store(blocks[:3], success=True) # Lookup all 5 blocks should return 3 (first 3 found) - assert self.manager.lookup(blocks) == 3 + assert count_hits(self.manager, blocks) == 3 def test_eviction_in_primary_tier(self, manager_setup): """Test eviction in primary tier when capacity is exceeded.""" # Primary tier has capacity of 5 blocks # First, fill the primary tier blocks = to_keys(range(5)) - result = self.manager.prepare_store(blocks) + result = self.manager.prepare_store(blocks, _CTX) assert result is not None assert len(result.keys_to_store) == 5 self.manager.complete_store(blocks, success=True) @@ -275,7 +300,7 @@ def test_eviction_in_primary_tier(self, manager_setup): # Now try to store 2 more blocks (should trigger eviction) more_blocks = to_keys(range(5, 7)) - result = self.manager.prepare_store(more_blocks) + result = self.manager.prepare_store(more_blocks, _CTX) # Should evict 2 blocks from primary tier assert result is not None @@ -287,7 +312,7 @@ def test_touch_propagates_to_all_tiers(self, manager_setup): blocks = to_keys(range(3)) # Store blocks - self.manager.prepare_store(blocks) + self.manager.prepare_store(blocks, _CTX) self.manager.complete_store(blocks, success=True) self.manager._process_finished_jobs() @@ -311,7 +336,7 @@ def test_failed_store_no_cascade(self, manager_setup): blocks = to_keys(range(3)) # Prepare store - result = self.manager.prepare_store(blocks) + result = self.manager.prepare_store(blocks, _CTX) assert result is not None # Complete store with failure @@ -348,7 +373,7 @@ def test_multiple_secondary_tiers_independent_eviction(self): # First, store 5 blocks to fill the small tier blocks1 = to_keys(range(5)) - result = manager.prepare_store(blocks1) + result = manager.prepare_store(blocks1, _CTX) assert result is not None manager.complete_store(blocks1, success=True) manager._process_finished_jobs() @@ -359,7 +384,7 @@ def test_multiple_secondary_tiers_independent_eviction(self): # Now store 3 more blocks - small tier should evict 3 blocks blocks2 = to_keys(range(5, 8)) - result = manager.prepare_store(blocks2) + result = manager.prepare_store(blocks2, _CTX) assert result is not None manager.complete_store(blocks2, success=True) manager._process_finished_jobs() @@ -375,7 +400,7 @@ def test_prepare_store_processes_finished_jobs_first(self, manager_setup): blocks = to_keys(range(3)) # Store blocks - self.manager.prepare_store(blocks) + self.manager.prepare_store(blocks, _CTX) self.manager.complete_store(blocks, success=True) # Blocks should have ref_cnt = 2 (one for each secondary tier) @@ -385,7 +410,7 @@ def test_prepare_store_processes_finished_jobs_first(self, manager_setup): # Call prepare_store again (should process finished jobs first) more_blocks = to_keys(range(3, 5)) - self.manager.prepare_store(more_blocks) + self.manager.prepare_store(more_blocks, _CTX) # Original blocks should now have ref_cnt = 0 for block_hash in blocks: @@ -412,11 +437,11 @@ def test_works_without_secondary_tiers(self): blocks = to_keys(range(3)) # Should work like a regular OffloadingManager - result = manager.prepare_store(blocks) + result = manager.prepare_store(blocks, _CTX) assert result is not None manager.complete_store(blocks, success=True) - assert manager.lookup(blocks) == 3 + assert count_hits(manager, blocks) == 3 if __name__ == "__main__": diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 8c1777506e0e..ca63ca5ac54f 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -33,6 +33,7 @@ OffloadingManager, OffloadKey, PrepareStoreOutput, + ReqContext, SecondaryTierManager, ) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager @@ -65,18 +66,18 @@ def __init__( ) self._mmap_region = mmap_region - def prepare_write(self, keys) -> PrepareStoreOutput | None: + def prepare_write(self, keys, req_context: ReqContext) -> PrepareStoreOutput | None: """Allocate space in primary for a secondary->primary write (promotion).""" - return self.prepare_store(keys) + return self.prepare_store(keys, req_context) def complete_write(self, keys, success: bool = True) -> None: """Finalize secondary->primary write, making blocks available.""" self.complete_store(keys, success) - def prepare_read(self, keys) -> LoadStoreSpec: + def prepare_read(self, keys, req_context: ReqContext) -> LoadStoreSpec: """Protect primary blocks for a primary->secondary read (cascade), incrementing ref_cnt.""" - return self.prepare_load(keys) + return self.prepare_load(keys, req_context) def complete_read(self, keys) -> None: """Release protection after primary->secondary read completes, @@ -198,74 +199,53 @@ def _process_finished_jobs(self): tier.get_tier_name(), ) - def lookup(self, keys: Iterable[OffloadKey]) -> int | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: """ - Find the length of the maximal series of blocks that are offloaded. + Check whether a single block is offloaded and ready. Algorithm: - 1. Check primary tier first - 2. If not all blocks found, check all secondary tiers sequentially, - promoting blocks from each tier that has hits and updating the - remaining blocks to search for - 3. Return None to signal "retry later" if any promotions were initiated + 1. Process any completed async jobs first + 2. Check primary tier + 3. If not in primary, check secondary tiers and initiate promotion + if the block is found there Args: - keys: Block hashes to look up. + key: Block hash to look up. + req_context: Per-request context. Returns: - Number of consecutive blocks (from start) that are present, - or None if blocks are being transferred (retry later). + True if the block is in the primary tier and ready, + False if not found in any tier, + None if the block is being transferred (retry later). """ - # Process any completed async jobs first to ensure promoted blocks - # are finalized and available in the primary tier self._process_finished_jobs() - keys_list = list(keys) - # Step 1: Check primary tier - primary_hits = self.primary_tier.lookup(keys_list) - - if primary_hits is None: - # Primary tier is busy (blocks being transferred) + primary_hit = self.primary_tier.lookup(key, req_context) + if primary_hit is None: return None + if primary_hit: + return True - if primary_hits == len(keys_list): - # All blocks in primary tier - return primary_hits - - # Step 2: Check all secondary tiers for remaining blocks - remaining_keys = keys_list[primary_hits:] - - # Track whether any promotions were initiated - has_promotions = False - + # Step 2: Check secondary tiers for tier in self.secondary_tiers: - if not remaining_keys: - # All blocks have been found - break - - secondary_hits = tier.lookup(remaining_keys) - - # Skip if tier is busy (None) or has no hits (0) - if not secondary_hits: - continue - - # Found blocks in this secondary tier, initiate promotion - blocks_to_promote = remaining_keys[:secondary_hits] - self._initiate_promotion(tier, blocks_to_promote) - has_promotions = True - - # Update remaining_keys to continue searching for the rest - remaining_keys = remaining_keys[secondary_hits:] - - # Step 3: If any promotions were initiated, return None to signal retry - if has_promotions: - return None - - # No more blocks found in any tier - return primary_hits - - def _initiate_promotion(self, tier: SecondaryTierManager, keys: list[OffloadKey]): + secondary_hits = tier.lookup([key]) + if secondary_hits is None: + # Tier is busy with this block + return None + if secondary_hits > 0: + # Found in secondary — initiate promotion and signal retry + self._initiate_promotion(tier, [key], req_context) + return None + + return False + + def _initiate_promotion( + self, + tier: SecondaryTierManager, + keys: list[OffloadKey], + req_context: ReqContext, + ): """ Initiate promotion of blocks from a secondary tier to the primary tier. @@ -277,9 +257,10 @@ def _initiate_promotion(self, tier: SecondaryTierManager, keys: list[OffloadKey] Args: tier: The secondary tier to promote from keys: Blocks to promote + req_context: Per-request context forwarded to primary.prepare_write(). """ # Allocate space in primary tier for promoted blocks - primary_store_result = self.primary_tier.prepare_write(keys) + primary_store_result = self.primary_tier.prepare_write(keys, req_context) if primary_store_result is None: # Cannot allocate space in primary tier (full) @@ -299,7 +280,9 @@ def _initiate_promotion(self, tier: SecondaryTierManager, keys: list[OffloadKey] tier.submit_load(job_metadata) - def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec: + def prepare_load( + self, keys: Iterable[OffloadKey], req_context: ReqContext + ) -> LoadStoreSpec: """ Prepare blocks to be loaded from primary tier to GPU. @@ -311,6 +294,7 @@ def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec: Args: keys: Blocks to prepare for loading. + req_context: Per-request context. Returns: LoadStoreSpec for reading from primary tier. @@ -318,7 +302,7 @@ def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec: # Process completed promotions to ensure blocks are ready self._process_finished_jobs() - return self.primary_tier.prepare_load(keys) + return self.primary_tier.prepare_load(keys, req_context) def touch(self, keys: Iterable[OffloadKey]): """ @@ -344,7 +328,9 @@ def complete_load(self, keys: Iterable[OffloadKey]): """ self.primary_tier.complete_load(keys) - def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None: + def prepare_store( + self, keys: Iterable[OffloadKey], req_context: ReqContext + ) -> PrepareStoreOutput | None: """ Prepare blocks to be stored from GPU to primary tier. @@ -354,6 +340,7 @@ def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None Args: keys: Blocks to prepare for storing. + req_context: Per-request context. Returns: PrepareStoreOutput describing where to store blocks and what was @@ -365,14 +352,18 @@ def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None self._process_finished_jobs() # Step 2: Store to primary tier - primary_result = self.primary_tier.prepare_store(keys) + primary_result = self.primary_tier.prepare_store(keys, req_context) # Note: Secondary tier cascading will happen in complete_store() # after the GPU→Primary transfer completes and blocks are ready. return primary_result - def complete_store(self, keys: Iterable[OffloadKey], success: bool = True): + def complete_store( + self, + keys: Iterable[OffloadKey], + success: bool = True, + ): """ Mark blocks as done storing from GPU to primary tier. @@ -389,6 +380,7 @@ def complete_store(self, keys: Iterable[OffloadKey], success: bool = True): Args: keys: Blocks that finished storing. success: Whether the GPU→primary transfer succeeded. + req_context: Per-request context forwarded to primary.prepare_read(). """ # Materialize only if success=True (needed for cascading to secondary tiers) keys_list = list(keys) if success else keys @@ -411,7 +403,8 @@ def complete_store(self, keys: Iterable[OffloadKey], success: bool = True): # secondary tier. for tier in self.secondary_tiers: # Get spec for reading from primary tier AND increment ref_cnt - primary_blocks_spec = self.primary_tier.prepare_read(keys_list) + # TODO: pass the actual req_context instead of None + primary_blocks_spec = self.primary_tier.prepare_read(keys_list, None) # Submit async store job: primary→secondary job_id = self._next_job_id() From f05638b20c7ddea1d179fae938349971d8b573d4 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 20 Apr 2026 17:39:00 +0300 Subject: [PATCH 04/55] Fix rebase issues: make pre-commit happy Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 4 +++- vllm/v1/kv_offload/tiering/spec.py | 1 - 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index ca63ca5ac54f..19843a847746 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -404,7 +404,9 @@ def complete_store( for tier in self.secondary_tiers: # Get spec for reading from primary tier AND increment ref_cnt # TODO: pass the actual req_context instead of None - primary_blocks_spec = self.primary_tier.prepare_read(keys_list, None) + primary_blocks_spec = self.primary_tier.prepare_read( + keys_list, ReqContext() + ) # Submit async store job: primary→secondary job_id = self._next_job_id() diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 0b0b6ecabebc..ef776f3cc6ab 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -158,7 +158,6 @@ def get_manager(self) -> OffloadingManager: # Create primary tier (CPU-based) assert len(self.gpu_block_size) == 1 - offloaded_block_size = self.gpu_block_size[0] * self.block_size_factor primary_tier = CPUPrimaryTierOffloadingManager( num_blocks=self.num_blocks, cache_policy=self.eviction_policy, # type: ignore[arg-type] From 7d3798d158da3678461bb2061b7b3c4e3d03fb25 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 21 Apr 2026 09:01:33 +0300 Subject: [PATCH 05/55] Address review: pass filtered keys to submit_load Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 19843a847746..01954586c245 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -273,7 +273,7 @@ def _initiate_promotion( # Track this load job job_metadata = JobMetadata( job_id=job_id, - keys=keys, + keys=primary_store_result.keys_to_store, spec=primary_store_result.store_spec, ) self._load_jobs[job_id] = job_metadata From 7374bf22a8b07bd8f694f4f2b2b01369aa7f2b7b Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 21 Apr 2026 09:45:39 +0300 Subject: [PATCH 06/55] propagate request context to submit_load() Signed-off-by: Ronen Schaffer --- tests/v1/kv_offload/test_tiering_offloading.py | 16 ++++++++++++++++ vllm/v1/kv_offload/abstract.py | 3 ++- vllm/v1/kv_offload/tiering/manager.py | 1 + 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index 4fd9c8b08a8f..40f45b8cae3a 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -12,6 +12,7 @@ """ from collections.abc import Iterable +from unittest.mock import MagicMock import pytest import torch @@ -417,6 +418,21 @@ def test_prepare_store_processes_finished_jobs_first(self, manager_setup): block = self.primary_tier._policy.get(block_hash) assert block.ref_cnt == 0 + def test_req_context_propagated_to_submit_load(self, manager_setup): + """Test that req_context from lookup() is forwarded to submit_load.""" + block = to_keys([0])[0] + self.secondary_tier1.blocks[block] = True # simulate prior cascade + + self.secondary_tier1.submit_load = MagicMock( + wraps=self.secondary_tier1.submit_load + ) + ctx = ReqContext(kv_transfer_params={"priority": "high"}) + self.manager.lookup(block, ctx) + + self.secondary_tier1.submit_load.assert_called_once() + job_metadata = self.secondary_tier1.submit_load.call_args[0][0] + assert job_metadata.req_context.kv_transfer_params == {"priority": "high"} + class TestTieringOffloadingWithoutSecondaryTiers: """Test TieringOffloadingManager with no secondary tiers (backward compat).""" diff --git a/vllm/v1/kv_offload/abstract.py b/vllm/v1/kv_offload/abstract.py index 8dacd2c5aa71..a4cbabe5bef9 100644 --- a/vllm/v1/kv_offload/abstract.py +++ b/vllm/v1/kv_offload/abstract.py @@ -28,7 +28,7 @@ from abc import ABC, abstractmethod from collections.abc import Iterable -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, NewType # `OffloadKey` identifies an offloaded block. It combines a block hash with @@ -99,6 +99,7 @@ class JobMetadata: job_id: JobId keys: list[OffloadKey] spec: LoadStoreSpec + req_context: ReqContext = field(default_factory=ReqContext) @dataclass diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 01954586c245..225d5514ff6d 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -275,6 +275,7 @@ def _initiate_promotion( job_id=job_id, keys=primary_store_result.keys_to_store, spec=primary_store_result.store_spec, + req_context=req_context, ) self._load_jobs[job_id] = job_metadata From 51654d9b7a0fa4f5f3bec574a7d51b4295e52b84 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 21 Apr 2026 10:50:54 +0300 Subject: [PATCH 07/55] Addredd review: Refactor get_primary_kv_tensor() -> create_kv_memoryview() Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 16 ++++----- vllm/v1/kv_offload/tiering/manager.py | 33 +++++++------------ vllm/v1/kv_offload/tiering/spec.py | 6 ++-- 3 files changed, 20 insertions(+), 35 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index 40f45b8cae3a..febe46bdf44f 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -166,10 +166,8 @@ def manager_setup(self): # Create primary tier (CPU-based) self.primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=5) - # Mock get_primary_kv_tensor to return test tensor - # Create mock CPU tensor (5 blocks, 16 bytes per block) - mock_cpu_tensor = torch.zeros((5, 16), dtype=torch.float32) - self.primary_tier.get_primary_kv_tensor = lambda: mock_cpu_tensor + mock_arr = torch.zeros((5, 16), dtype=torch.int8).numpy() + self.primary_tier.create_kv_memoryview = lambda: memoryview(mock_arr) # Create secondary tiers self.secondary_tier1 = DummySecondaryTier(tier_name="Storage", max_blocks=10) @@ -363,9 +361,8 @@ def test_multiple_secondary_tiers_independent_eviction(self): # Create a fresh primary tier for this test primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=10) - # Mock get_primary_kv_tensor to return test tensor - mock_cpu_tensor = torch.zeros((10, 16), dtype=torch.float32) - primary_tier.get_primary_kv_tensor = lambda: mock_cpu_tensor + mock_arr = torch.zeros((10, 16), dtype=torch.int8).numpy() + primary_tier.create_kv_memoryview = lambda: memoryview(mock_arr) manager = TieringOffloadingManager( primary_tier=primary_tier, @@ -441,9 +438,8 @@ def test_works_without_secondary_tiers(self): """Test that manager works with empty secondary_tiers list.""" primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=5) - # Mock get_primary_kv_tensor to return test tensor - mock_cpu_tensor = torch.zeros((5, 16), dtype=torch.float32) - primary_tier.get_primary_kv_tensor = lambda: mock_cpu_tensor + mock_arr = torch.zeros((5, 16), dtype=torch.int8).numpy() + primary_tier.create_kv_memoryview = lambda: memoryview(mock_arr) # Create manager with no secondary tiers manager = TieringOffloadingManager( diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 225d5514ff6d..56815f04d9b2 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -22,8 +22,6 @@ from collections.abc import Iterable -import torch - from vllm.logger import init_logger from vllm.v1.kv_offload.abstract import ( JobId, @@ -84,24 +82,20 @@ def complete_read(self, keys) -> None: decrementing ref_cnt.""" self.complete_load(keys) - def get_primary_kv_tensor(self) -> torch.Tensor: - """ - Get the primary tier's KV cache tensor. - - Returns a 2-D int8 tensor of shape (num_blocks, row_stride_bytes) - backed by the SharedOffloadRegion mmap, where row_stride_bytes = - cpu_page_size * world_size. Secondary tiers address block b as - view[b], and view.strides[0] gives the per-block byte stride. + def create_kv_memoryview(self) -> memoryview: + """Create a memoryview over the primary tier's KV cache buffer. - Returns: - 2-D int8 CPU tensor of shape (num_blocks, row_stride_bytes). + Returns a 2-D memoryview of shape (num_blocks, row_stride_bytes) + backed by the SharedOffloadRegion mmap. Secondary tiers address + block b as view[b]. Caller must call release() when done. """ assert self._mmap_region is not None, ( "mmap_region must be provided to CPUPrimaryTierOffloadingManager" ) - return self._mmap_region._base.view( + kv_tensor = self._mmap_region._base.view( self._mmap_region.num_blocks, self._mmap_region._row_stride ) + return memoryview(kv_tensor.numpy()) class TieringOffloadingManager(OffloadingManager): @@ -148,14 +142,10 @@ def __init__( self._load_jobs: dict[JobId, JobMetadata] = {} # Wire each secondary tier with a long-lived memoryview of the primary - # CPU tensor (one independent view per tier). Views are stored so they - # can be released on shutdown(). - self._secondary_views: list[memoryview] = [] - cpu_tensor = primary_tier.get_primary_kv_tensor() + # CPU buffer. One view is shared across all tiers; released in shutdown(). + self._primary_kv_view = primary_tier.create_kv_memoryview() for tier in self.secondary_tiers: - view = memoryview(cpu_tensor.numpy()) - self._secondary_views.append(view) - tier.set_primary_view(view) + tier.set_primary_view(self._primary_kv_view) def _next_job_id(self) -> JobId: """Generate a unique job ID for async transfer tracking.""" @@ -442,8 +432,7 @@ def take_events(self) -> Iterable[OffloadingEvent]: def shutdown(self) -> None: """Release memoryviews and scheduler-side mmap.""" - for view in self._secondary_views: - view.release() + self._primary_kv_view.release() if self.primary_tier._mmap_region is not None: self.primary_tier._mmap_region.cleanup() self.primary_tier._mmap_region = None diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index ef776f3cc6ab..bd6e0282525f 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -140,9 +140,9 @@ def get_manager(self) -> OffloadingManager: ) # Create scheduler-side SharedOffloadRegion (rank=None) first so - # CPUPrimaryTierOffloadingManager can return _base from - # get_primary_kv_tensor(), which TieringOffloadingManager.__init__ - # calls immediately to wire secondary tier memoryviews. + # CPUPrimaryTierOffloadingManager.create_kv_memoryview() has _base + # available when TieringOffloadingManager.__init__ wires secondary + # tier memoryviews. world_size = self.vllm_config.parallel_config.world_size scheduler_mmap = SharedOffloadRegion( instance_id=self.vllm_config.instance_id, From 9dc1757b79caa3772557bfb32e25e4f109ef767f Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 21 Apr 2026 11:38:12 +0300 Subject: [PATCH 08/55] Address review: remove support for store_threshold Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/spec.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index bd6e0282525f..bfc0235267ef 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -11,11 +11,6 @@ - block_size: (optional) Block size for offloaded blocks (default: GPU block size) - eviction_policy: (optional) Primary tier eviction policy: "lru" or "arc" (default: "lru") - - store_threshold: (optional) How many times a block must appear in lookup() - before it is eligible for CPU offloading. Values < 2 disable filtering - (default: 0) - - max_tracker_size: (optional) Maximum number of blocks tracked for - store_threshold filtering (default: 64000) - secondary_tiers: (optional) List of secondary tier configurations Each secondary tier config is a dict with: - type: (required) Type of secondary tier (e.g., "dummy", "storage", "network") @@ -74,6 +69,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec): def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): super().__init__(vllm_config, kv_cache_config) + # Redeclare for mypy: parent sets this but `--follow-imports skip` hides it + self._manager: OffloadingManager | None = None # Parse secondary tier configurations self.secondary_tier_configs = self.extra_config.get("secondary_tiers", []) @@ -192,9 +189,11 @@ def get_manager(self) -> OffloadingManager: secondary_tiers=secondary_tiers, enable_events=enable_events, ) - # PRNOTE: should the store_filter apply to the TieringOffloadingManager or - # to the primary CPU manager? - self._manager = self._maybe_apply_store_filter(tiering_manager) + if int(self.extra_config.get("store_threshold", 0)) >= 2: + raise ValueError( + "store_threshold is not supported for TieringOffloadingSpec" + ) + self._manager = tiering_manager logger.info( "Created TieringOffloadingManager with primary tier " From 93a2a67b0672065e937816bf11cc6c4ccafc1957 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 23 Apr 2026 11:22:16 +0300 Subject: [PATCH 09/55] Change SecondaryTierManager.lookup() to accept single key instead of multiple keys Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 16 ++++++----- vllm/v1/kv_offload/abstract.py | 11 ++++---- vllm/v1/kv_offload/secondary_tiers/dummy.py | 28 ++++++------------- vllm/v1/kv_offload/tiering/manager.py | 6 ++-- 4 files changed, 27 insertions(+), 34 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index febe46bdf44f..ea1572e793ee 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -63,17 +63,18 @@ def test_basic_store_and_lookup(self): # Initially empty blocks = to_keys(range(3)) - assert tier.lookup(blocks) == 0 + assert tier.lookup(blocks[0]) is False # Store blocks (simulate with direct insertion for testing) tier.blocks[blocks[0]] = True tier.blocks[blocks[1]] = True - # Lookup should find 2 blocks - assert tier.lookup(blocks) == 2 + # Lookup should find first two blocks + assert tier.lookup(blocks[0]) is True + assert tier.lookup(blocks[1]) is True # Third block not present - assert tier.lookup([blocks[2]]) == 0 + assert tier.lookup(blocks[2]) is False def test_in_flight_blocks_return_none(self): """Test that in-flight blocks cause lookup to return None.""" @@ -85,7 +86,8 @@ def test_in_flight_blocks_return_none(self): tier.in_flight[blocks[0]] = 1 # Lookup should return None (retry later) - assert tier.lookup(blocks) is None + assert tier.lookup(blocks[0]) is None + assert tier.lookup(blocks[1]) is False # not in-flight, just absent def test_lru_eviction(self): """Test LRU eviction policy.""" @@ -213,8 +215,8 @@ def test_cascade_to_all_secondary_tiers(self, manager_setup): assert self.secondary_tier2.get_num_blocks() == 3 # Verify blocks are present - assert self.secondary_tier1.lookup(blocks) == 3 - assert self.secondary_tier2.lookup(blocks) == 3 + assert all(self.secondary_tier1.lookup(b) for b in blocks) + assert all(self.secondary_tier2.lookup(b) for b in blocks) def test_ref_cnt_protection_during_cascade(self, manager_setup): """Test that ref_cnt protects blocks during cascade.""" diff --git a/vllm/v1/kv_offload/abstract.py b/vllm/v1/kv_offload/abstract.py index a4cbabe5bef9..9ba172392561 100644 --- a/vllm/v1/kv_offload/abstract.py +++ b/vllm/v1/kv_offload/abstract.py @@ -234,16 +234,17 @@ class SecondaryTierManager(ABC): """ @abstractmethod - def lookup(self, keys: Iterable[OffloadKey]) -> int | None: + def lookup(self, key: OffloadKey) -> bool | None: """ - Check which blocks exist in this secondary tier. + Check whether a block exists in this secondary tier. Args: - keys: Offload keys to look up. + key: Offload key to look up. Returns: - Number of consecutive blocks (from start) that are present and ready, - or None if blocks are being transferred (retry later). + True if the block is present and ready, + False if not found, + or None if the block is being transferred (retry later). """ pass diff --git a/vllm/v1/kv_offload/secondary_tiers/dummy.py b/vllm/v1/kv_offload/secondary_tiers/dummy.py index f75edafd6187..b0475d03cef4 100644 --- a/vllm/v1/kv_offload/secondary_tiers/dummy.py +++ b/vllm/v1/kv_offload/secondary_tiers/dummy.py @@ -94,31 +94,21 @@ def __init__( def set_primary_view(self, view: memoryview) -> None: self._primary_view = view - def lookup(self, keys: Iterable[OffloadKey]) -> int | None: + def lookup(self, key: OffloadKey) -> bool | None: """ - Check which blocks exist in this secondary tier. + Check whether a block exists in this secondary tier. Args: - keys: Block hashes to look up. + key: Offload key to look up. Returns: - Number of consecutive blocks (from start) that are present and ready, - or None if blocks are being transferred (retry later). + True if the block is present and ready, + False if not found, + or None if the block is being transferred (retry later). """ - hit_count = 0 - for key in keys: - # Check if block is in-flight - if key in self.in_flight: - # Block is being transferred, return None (retry later) - return None - - # Check if block exists in this tier - if key not in self.blocks: - break - - hit_count += 1 - - return hit_count + if key in self.in_flight: + return None + return key in self.blocks def submit_store(self, job_metadata: JobMetadata) -> None: """ diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 56815f04d9b2..f4332d150938 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -219,11 +219,11 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: # Step 2: Check secondary tiers for tier in self.secondary_tiers: - secondary_hits = tier.lookup([key]) - if secondary_hits is None: + secondary_hit = tier.lookup(key) + if secondary_hit is None: # Tier is busy with this block return None - if secondary_hits > 0: + if secondary_hit: # Found in secondary — initiate promotion and signal retry self._initiate_promotion(tier, [key], req_context) return None From f1975be17d1d83261c66f53e1f9e1ccd8e128c7e Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 23 Apr 2026 11:32:33 +0300 Subject: [PATCH 10/55] Add ReqContext to SecondaryTierManager.lookup() Signed-off-by: Ronen Schaffer --- tests/v1/kv_offload/test_tiering_offloading.py | 16 ++++++++-------- vllm/v1/kv_offload/abstract.py | 3 ++- vllm/v1/kv_offload/secondary_tiers/dummy.py | 3 ++- vllm/v1/kv_offload/tiering/manager.py | 2 +- 4 files changed, 13 insertions(+), 11 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index ea1572e793ee..634c0ed07e29 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -63,18 +63,18 @@ def test_basic_store_and_lookup(self): # Initially empty blocks = to_keys(range(3)) - assert tier.lookup(blocks[0]) is False + assert tier.lookup(blocks[0], _CTX) is False # Store blocks (simulate with direct insertion for testing) tier.blocks[blocks[0]] = True tier.blocks[blocks[1]] = True # Lookup should find first two blocks - assert tier.lookup(blocks[0]) is True - assert tier.lookup(blocks[1]) is True + assert tier.lookup(blocks[0], _CTX) is True + assert tier.lookup(blocks[1], _CTX) is True # Third block not present - assert tier.lookup(blocks[2]) is False + assert tier.lookup(blocks[2], _CTX) is False def test_in_flight_blocks_return_none(self): """Test that in-flight blocks cause lookup to return None.""" @@ -86,8 +86,8 @@ def test_in_flight_blocks_return_none(self): tier.in_flight[blocks[0]] = 1 # Lookup should return None (retry later) - assert tier.lookup(blocks[0]) is None - assert tier.lookup(blocks[1]) is False # not in-flight, just absent + assert tier.lookup(blocks[0], _CTX) is None + assert tier.lookup(blocks[1], _CTX) is False # not in-flight, just absent def test_lru_eviction(self): """Test LRU eviction policy.""" @@ -215,8 +215,8 @@ def test_cascade_to_all_secondary_tiers(self, manager_setup): assert self.secondary_tier2.get_num_blocks() == 3 # Verify blocks are present - assert all(self.secondary_tier1.lookup(b) for b in blocks) - assert all(self.secondary_tier2.lookup(b) for b in blocks) + assert all(self.secondary_tier1.lookup(b, _CTX) for b in blocks) + assert all(self.secondary_tier2.lookup(b, _CTX) for b in blocks) def test_ref_cnt_protection_during_cascade(self, manager_setup): """Test that ref_cnt protects blocks during cascade.""" diff --git a/vllm/v1/kv_offload/abstract.py b/vllm/v1/kv_offload/abstract.py index 9ba172392561..a17e750c84ba 100644 --- a/vllm/v1/kv_offload/abstract.py +++ b/vllm/v1/kv_offload/abstract.py @@ -234,12 +234,13 @@ class SecondaryTierManager(ABC): """ @abstractmethod - def lookup(self, key: OffloadKey) -> bool | None: + 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 (e.g. kv_transfer_params). Returns: True if the block is present and ready, diff --git a/vllm/v1/kv_offload/secondary_tiers/dummy.py b/vllm/v1/kv_offload/secondary_tiers/dummy.py index b0475d03cef4..c8b67583a04f 100644 --- a/vllm/v1/kv_offload/secondary_tiers/dummy.py +++ b/vllm/v1/kv_offload/secondary_tiers/dummy.py @@ -19,6 +19,7 @@ JobResult, LoadStoreSpec, OffloadKey, + ReqContext, SecondaryTierManager, ) from vllm.v1.kv_offload.mediums import CPULoadStoreSpec @@ -94,7 +95,7 @@ def __init__( def set_primary_view(self, view: memoryview) -> None: self._primary_view = view - def lookup(self, key: OffloadKey) -> bool | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: """ Check whether a block exists in this secondary tier. diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index f4332d150938..051dd3480a7a 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -219,7 +219,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: # Step 2: Check secondary tiers for tier in self.secondary_tiers: - secondary_hit = tier.lookup(key) + secondary_hit = tier.lookup(key, req_context) if secondary_hit is None: # Tier is busy with this block return None From 70bfa596c7b93087504848b3cbdc91d45faf4846 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 23 Apr 2026 14:43:35 +0300 Subject: [PATCH 11/55] Change _initiate_promotion() to handle single key Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 051dd3480a7a..55d8f5dd795c 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -225,7 +225,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: return None if secondary_hit: # Found in secondary — initiate promotion and signal retry - self._initiate_promotion(tier, [key], req_context) + self._initiate_promotion(tier, key, req_context) return None return False @@ -233,11 +233,11 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: def _initiate_promotion( self, tier: SecondaryTierManager, - keys: list[OffloadKey], + key: OffloadKey, req_context: ReqContext, ): """ - Initiate promotion of blocks from a secondary tier to the primary tier. + Initiate promotion of a block from a secondary tier to the primary tier. This method: 1. Calls primary.prepare_write() to allocate space in primary tier @@ -246,11 +246,11 @@ def _initiate_promotion( Args: tier: The secondary tier to promote from - keys: Blocks to promote + key: Block to promote req_context: Per-request context forwarded to primary.prepare_write(). """ - # Allocate space in primary tier for promoted blocks - primary_store_result = self.primary_tier.prepare_write(keys, req_context) + # Allocate space in primary tier for promoted block + primary_store_result = self.primary_tier.prepare_write([key], req_context) if primary_store_result is None: # Cannot allocate space in primary tier (full) From f8a0438bdf8b81ffc50ced51750f834ab411d8f5 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 23 Apr 2026 16:03:22 +0300 Subject: [PATCH 12/55] Address review: re-organize file tree structure Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 4 +- vllm/v1/kv_offload/abstract.py | 158 +---------------- .../v1/kv_offload/secondary_tiers/__init__.py | 18 -- vllm/v1/kv_offload/tiering/base.py | 166 ++++++++++++++++++ .../dummy.py => tiering/dummy/__init__.py} | 8 +- vllm/v1/kv_offload/tiering/factory.py | 44 +++++ vllm/v1/kv_offload/tiering/manager.py | 8 +- vllm/v1/kv_offload/tiering/spec.py | 43 +---- 8 files changed, 223 insertions(+), 226 deletions(-) delete mode 100644 vllm/v1/kv_offload/secondary_tiers/__init__.py create mode 100644 vllm/v1/kv_offload/tiering/base.py rename vllm/v1/kv_offload/{secondary_tiers/dummy.py => tiering/dummy/__init__.py} (98%) create mode 100644 vllm/v1/kv_offload/tiering/factory.py diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index 634c0ed07e29..7b05819f7393 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -18,13 +18,13 @@ import torch from vllm.v1.kv_offload.abstract import ( - JobMetadata, OffloadKey, ReqContext, make_offload_key, ) from vllm.v1.kv_offload.mediums import CPULoadStoreSpec -from vllm.v1.kv_offload.secondary_tiers.dummy import DummySecondaryTier +from vllm.v1.kv_offload.tiering.base import JobMetadata +from vllm.v1.kv_offload.tiering.dummy import DummySecondaryTier from vllm.v1.kv_offload.tiering.manager import ( CPUPrimaryTierOffloadingManager, TieringOffloadingManager, diff --git a/vllm/v1/kv_offload/abstract.py b/vllm/v1/kv_offload/abstract.py index a17e750c84ba..3506e8c0d8d8 100644 --- a/vllm/v1/kv_offload/abstract.py +++ b/vllm/v1/kv_offload/abstract.py @@ -28,7 +28,7 @@ from abc import ABC, abstractmethod from collections.abc import Iterable -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any, NewType # `OffloadKey` identifies an offloaded block. It combines a block hash with @@ -57,10 +57,6 @@ class ReqContext: kv_transfer_params: dict[str, Any] | None = None -# Type alias for job IDs used in async transfer tracking -JobId = int - - class LoadStoreSpec(ABC): """ Abstract metadata that encapsulates information allowing a worker @@ -92,24 +88,6 @@ class OffloadingEvent: removed: bool -@dataclass -class JobMetadata: - """Metadata for an in-flight async transfer job.""" - - job_id: JobId - keys: list[OffloadKey] - spec: LoadStoreSpec - req_context: ReqContext = field(default_factory=ReqContext) - - -@dataclass -class JobResult: - """Result of an async transfer job (successful or failed).""" - - job_id: JobId - success: bool - - class OffloadingManager(ABC): @abstractmethod def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: @@ -217,137 +195,3 @@ def take_events(self) -> Iterable[OffloadingEvent]: def shutdown(self) -> None: """Shutdown the manager and release any resources.""" return - - -class SecondaryTierManager(ABC): - """ - Abstract interface for managing a single non-primary offloading tier. - - Secondary tiers cannot directly access GPU memory. All data transfers - must go through the primary tier (implemented as CPU in current version): - - Store: GPU → primary → secondary (cascade) - - Load: secondary → primary → GPU (promotion) - - IMPORTANT: All methods run in the Scheduler process and must be - lightweight and non-blocking. submit_load() and submit_store() submit - async jobs; get_finished() polls for completion. - """ - - @abstractmethod - 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 (e.g. kv_transfer_params). - - Returns: - True if the block is present and ready, - False if not found, - or None if the block is being transferred (retry later). - """ - pass - - @abstractmethod - def submit_store(self, job_metadata: JobMetadata) -> None: - """ - Submit an async job to store blocks from the primary tier to this - secondary tier. - - This method is lightweight: it allocates metadata and submits the - transfer job, but does NOT perform the actual data transfer on the - calling thread. - - The caller (TieringOffloadingManager) must have already called - primary.prepare_read(keys) to obtain job_metadata.spec and - to increment ref_cnt on those blocks. ref_cnt will be decremented - when get_finished() reports this job_id as complete and - primary.unprepare_read() is called. - - This method is responsible for: - 1. Filtering out blocks already present in this secondary tier - 2. Evicting blocks from this secondary tier if needed (secondary - tiers are responsible for their own evictions) - 3. Allocating space in this secondary tier - 4. Submitting the async transfer: primary → secondary - - Args: - job_metadata: Job metadata including job_id, keys, and - spec for reading blocks from the primary tier - (obtained via primary.prepare_read()). - spec is a CPULoadStoreSpec with block_ids. - """ - pass - - @abstractmethod - def submit_load(self, job_metadata: JobMetadata) -> None: - """ - Submit an async job to load blocks from this secondary tier to the - primary tier. - - This method is lightweight: it marks blocks as in-flight and submits - the transfer job, but does NOT perform the actual data transfer on - the calling thread. - - The caller (TieringOffloadingManager) must have already called - primary.prepare_write(keys) to obtain job_metadata.spec and - to allocate space in the primary tier. When get_finished() reports - this job_id as complete, primary.complete_write() is called to make - the blocks available for GPU loads. - - Args: - job_metadata: Job metadata including job_id, keys, and - spec for writing blocks into the primary tier - (obtained via primary.prepare_write()). - spec is a CPULoadStoreSpec with block_ids. - """ - pass - - @abstractmethod - def get_finished(self) -> Iterable[JobResult]: - """ - Poll for finished async jobs (both loads and stores). - - This is the mechanism by which the TieringOffloadingManager learns - that a transfer has finished and can: - - Call primary.unprepare_read() to decrement ref_cnt (for stores) - - Call primary.complete_write() to make blocks loadable (for loads) - - Returns: - Iterable of JobResult objects for all jobs that have - finished since the last call. - """ - pass - - def set_primary_view(self, view: memoryview) -> None: - """ - Provide a long-lived memoryview of the primary-tier CPU tensor. - - Called once by TieringOffloadingManager during initialisation. - Override to store the view for use in `submit_store` and `submit_load`. - Use `view.strides[0]` to obtain the byte stride between block slots. - - Args: - view: Memoryview of the primary tier's CPU KV cache tensor. - """ - return - - def touch(self, keys: Iterable[OffloadKey]): - """ - Mark blocks as recently used for eviction policy. - - Args: - keys: Offload keys to mark as recently used. - """ - return - - @abstractmethod - def get_tier_name(self) -> str: - """ - Get the name of this tier (e.g., "Storage", "Network"). - - Returns: - Tier name string. - """ - pass diff --git a/vllm/v1/kv_offload/secondary_tiers/__init__.py b/vllm/v1/kv_offload/secondary_tiers/__init__.py deleted file mode 100644 index c4b5a473aa7b..000000000000 --- a/vllm/v1/kv_offload/secondary_tiers/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Secondary tier implementations for KV cache offloading. - -This package contains various secondary tier storage backends that can be used -with the TieringOffloadingManager for multi-tier KV cache management. -""" - -from vllm.v1.kv_offload.secondary_tiers.dummy import ( - DummyLoadStoreSpec, - DummySecondaryTier, -) - -__all__ = [ - "DummyLoadStoreSpec", - "DummySecondaryTier", -] diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py new file mode 100644 index 000000000000..57750fae9141 --- /dev/null +++ b/vllm/v1/kv_offload/tiering/base.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Abstract interfaces and data types for the secondary tiering layer. +""" + +from abc import ABC, abstractmethod +from collections.abc import Iterable +from dataclasses import dataclass, field + +from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadKey, ReqContext + +# Type alias for job IDs used in async transfer tracking +JobId = int + + +@dataclass +class JobMetadata: + """Metadata for an in-flight async transfer job.""" + + job_id: JobId + keys: list[OffloadKey] + spec: LoadStoreSpec + req_context: ReqContext = field(default_factory=ReqContext) + + +@dataclass +class JobResult: + """Result of an async transfer job (successful or failed).""" + + job_id: JobId + success: bool + + +class SecondaryTierManager(ABC): + """ + Abstract interface for managing a single non-primary offloading tier. + + Secondary tiers cannot directly access GPU memory. All data transfers + must go through the primary tier (implemented as CPU in current version): + - Store: GPU → primary → secondary (cascade) + - Load: secondary → primary → GPU (promotion) + + IMPORTANT: All methods run in the Scheduler process and must be + lightweight and non-blocking. submit_load() and submit_store() submit + async jobs; get_finished() polls for completion. + """ + + @abstractmethod + 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 (e.g. kv_transfer_params). + + Returns: + True if the block is present and ready, + False if not found, + or None if the block is being transferred (retry later). + """ + pass + + @abstractmethod + def submit_store(self, job_metadata: JobMetadata) -> None: + """ + Submit an async job to store blocks from the primary tier to this + secondary tier. + + This method is lightweight: it allocates metadata and submits the + transfer job, but does NOT perform the actual data transfer on the + calling thread. + + The caller (TieringOffloadingManager) must have already called + primary.prepare_read(keys) to obtain job_metadata.spec and + to increment ref_cnt on those blocks. ref_cnt will be decremented + when get_finished() reports this job_id as complete and + primary.unprepare_read() is called. + + This method is responsible for: + 1. Filtering out blocks already present in this secondary tier + 2. Evicting blocks from this secondary tier if needed (secondary + tiers are responsible for their own evictions) + 3. Allocating space in this secondary tier + 4. Submitting the async transfer: primary → secondary + + Args: + job_metadata: Job metadata including job_id, keys, and + spec for reading blocks from the primary tier + (obtained via primary.prepare_read()). + spec is a CPULoadStoreSpec with block_ids. + """ + pass + + @abstractmethod + def submit_load(self, job_metadata: JobMetadata) -> None: + """ + Submit an async job to load blocks from this secondary tier to the + primary tier. + + This method is lightweight: it marks blocks as in-flight and submits + the transfer job, but does NOT perform the actual data transfer on + the calling thread. + + The caller (TieringOffloadingManager) must have already called + primary.prepare_write(keys) to obtain job_metadata.spec and + to allocate space in the primary tier. When get_finished() reports + this job_id as complete, primary.complete_write() is called to make + the blocks available for GPU loads. + + Args: + job_metadata: Job metadata including job_id, keys, and + spec for writing blocks into the primary tier + (obtained via primary.prepare_write()). + spec is a CPULoadStoreSpec with block_ids. + """ + pass + + @abstractmethod + def get_finished(self) -> Iterable[JobResult]: + """ + Poll for finished async jobs (both loads and stores). + + This is the mechanism by which the TieringOffloadingManager learns + that a transfer has finished and can: + - Call primary.unprepare_read() to decrement ref_cnt (for stores) + - Call primary.complete_write() to make blocks loadable (for loads) + + Returns: + Iterable of JobResult objects for all jobs that have + finished since the last call. + """ + pass + + def set_primary_view(self, view: memoryview) -> None: + """ + Provide a long-lived memoryview of the primary-tier CPU tensor. + + Called once by TieringOffloadingManager during initialisation. + Override to store the view for use in `submit_store` and `submit_load`. + Use `view.strides[0]` to obtain the byte stride between block slots. + + Args: + view: Memoryview of the primary tier's CPU KV cache tensor. + """ + return + + def touch(self, keys: Iterable[OffloadKey]): + """ + Mark blocks as recently used for eviction policy. + + Args: + keys: Offload keys to mark as recently used. + """ + return + + @abstractmethod + def get_tier_name(self) -> str: + """ + Get the name of this tier (e.g., "Storage", "Network"). + + Returns: + Tier name string. + """ + pass diff --git a/vllm/v1/kv_offload/secondary_tiers/dummy.py b/vllm/v1/kv_offload/tiering/dummy/__init__.py similarity index 98% rename from vllm/v1/kv_offload/secondary_tiers/dummy.py rename to vllm/v1/kv_offload/tiering/dummy/__init__.py index c8b67583a04f..677ac1ee279f 100644 --- a/vllm/v1/kv_offload/secondary_tiers/dummy.py +++ b/vllm/v1/kv_offload/tiering/dummy/__init__.py @@ -13,16 +13,14 @@ from collections.abc import Iterable from dataclasses import dataclass -from vllm.v1.kv_offload.abstract import ( +from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadKey, ReqContext +from vllm.v1.kv_offload.mediums import CPULoadStoreSpec +from vllm.v1.kv_offload.tiering.base import ( JobId, JobMetadata, JobResult, - LoadStoreSpec, - OffloadKey, - ReqContext, SecondaryTierManager, ) -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec @dataclass diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py new file mode 100644 index 000000000000..f5354a7a6655 --- /dev/null +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Factory for creating secondary tier implementations. +""" + +from vllm.v1.kv_offload.tiering.base import SecondaryTierManager + + +def create_secondary_tier(tier_config: dict) -> SecondaryTierManager: + """ + Create a secondary tier from configuration. + + Args: + tier_config: Dictionary with tier configuration containing: + - type (required): Type of secondary tier (e.g., "dummy") + - tier_name (required): Name for this tier + - Additional tier-specific parameters are passed directly + to the tier constructor + + Returns: + SecondaryTierManager instance + + 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_name = config.pop("tier_name", None) + if not tier_name: + raise ValueError("Secondary tier configuration must include 'tier_name'") + + if tier_type == "dummy": + from vllm.v1.kv_offload.tiering.dummy import DummySecondaryTier + + return DummySecondaryTier(tier_name=tier_name, **config) + else: + raise ValueError( + f"Unknown secondary tier type: {tier_type}. Supported types: dummy" + ) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 55d8f5dd795c..12a9a99efb3a 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -24,18 +24,20 @@ from vllm.logger import init_logger from vllm.v1.kv_offload.abstract import ( - JobId, - JobMetadata, LoadStoreSpec, OffloadingEvent, OffloadingManager, OffloadKey, PrepareStoreOutput, ReqContext, - SecondaryTierManager, ) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion +from vllm.v1.kv_offload.tiering.base import ( + JobId, + JobMetadata, + SecondaryTierManager, +) logger = init_logger(__name__) diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index bfc0235267ef..397dd59b8090 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -43,8 +43,8 @@ from vllm.v1.kv_offload.abstract import OffloadingManager 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.secondary_tiers.dummy import DummySecondaryTier from vllm.v1.kv_offload.spec import CanonicalKVCaches +from vllm.v1.kv_offload.tiering.factory import create_secondary_tier from vllm.v1.kv_offload.tiering.manager import ( CPUPrimaryTierOffloadingManager, TieringOffloadingManager, @@ -80,45 +80,6 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): # Scheduler-side mmap (rank=None); kept for cleanup self._scheduler_mmap: SharedOffloadRegion | None = None - def _create_secondary_tier(self, tier_config: dict): - """ - Create a secondary tier from configuration. - - Args: - tier_config: Dictionary with tier configuration containing: - - type (required): Type of secondary tier (e.g., "dummy") - - tier_name (required): Name for this tier - - Additional tier-specific parameters are passed directly - to the tier constructor - - Returns: - SecondaryTierManager instance - - Raises: - ValueError: If tier type is unknown or configuration is invalid - """ - # Make a copy to avoid modifying the original config - config = tier_config.copy() - - # Extract common parameters - tier_type = config.pop("type", None) - if not tier_type: - raise ValueError("Secondary tier configuration must include 'type'") - - tier_name = config.pop("tier_name", None) - if not tier_name: - raise ValueError("Secondary tier configuration must include 'tier_name'") - - # Remaining parameters in config are tier-specific - if tier_type == "dummy": - # DummySecondaryTier for testing - # Pass tier_name and tier-specific params to constructor - return DummySecondaryTier(tier_name=tier_name, **config) - else: - raise ValueError( - f"Unknown secondary tier type: {tier_type}. Supported types: dummy" - ) - def get_manager(self) -> OffloadingManager: """ Get the TieringOffloadingManager. @@ -166,7 +127,7 @@ def get_manager(self) -> OffloadingManager: secondary_tiers = [] for tier_config in self.secondary_tier_configs: try: - tier = self._create_secondary_tier(tier_config) + tier = create_secondary_tier(tier_config) secondary_tiers.append(tier) logger.info( "Created secondary tier: %s (type: %s)", From 98ffa4889c561744e4670e1f54e5d4ad2bc66a06 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 28 Apr 2026 10:20:40 +0300 Subject: [PATCH 13/55] Address review: remove tier_name from config Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 20 +++++++--------- vllm/v1/kv_offload/tiering/base.py | 9 +++++--- vllm/v1/kv_offload/tiering/dummy/__init__.py | 14 +++-------- vllm/v1/kv_offload/tiering/factory.py | 23 ++++++++++--------- vllm/v1/kv_offload/tiering/manager.py | 8 ++++--- vllm/v1/kv_offload/tiering/spec.py | 10 ++++---- 6 files changed, 38 insertions(+), 46 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index 7b05819f7393..88372f853c11 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -59,7 +59,7 @@ class TestDummySecondaryTier: def test_basic_store_and_lookup(self): """Test basic store and lookup operations.""" - tier = DummySecondaryTier(tier_name="Test", max_blocks=10) + tier = DummySecondaryTier(max_blocks=10) # Initially empty blocks = to_keys(range(3)) @@ -78,7 +78,7 @@ def test_basic_store_and_lookup(self): def test_in_flight_blocks_return_none(self): """Test that in-flight blocks cause lookup to return None.""" - tier = DummySecondaryTier(tier_name="Test", max_blocks=10) + tier = DummySecondaryTier(max_blocks=10) blocks = to_keys(range(3)) @@ -91,7 +91,7 @@ def test_in_flight_blocks_return_none(self): def test_lru_eviction(self): """Test LRU eviction policy.""" - tier = DummySecondaryTier(tier_name="Test", max_blocks=3) + tier = DummySecondaryTier(max_blocks=3) # Fill tier to capacity blocks = to_keys(range(3)) @@ -129,7 +129,7 @@ def test_lru_eviction(self): def test_async_simulation(self): """Test simulated async behavior.""" - tier = DummySecondaryTier(tier_name="Test", max_blocks=10, simulate_async=True) + tier = DummySecondaryTier(max_blocks=10, simulate_async=True) blocks = to_keys(range(2)) @@ -172,8 +172,8 @@ def manager_setup(self): self.primary_tier.create_kv_memoryview = lambda: memoryview(mock_arr) # Create secondary tiers - self.secondary_tier1 = DummySecondaryTier(tier_name="Storage", max_blocks=10) - self.secondary_tier2 = DummySecondaryTier(tier_name="Network", max_blocks=10) + self.secondary_tier1 = DummySecondaryTier(max_blocks=10) + self.secondary_tier2 = DummySecondaryTier(max_blocks=10) # Create tiered manager self.manager = TieringOffloadingManager( @@ -353,12 +353,8 @@ def test_failed_store_no_cascade(self, manager_setup): def test_multiple_secondary_tiers_independent_eviction(self): """Test that secondary tiers manage their own evictions.""" # Create tier with small capacity - small_tier = DummySecondaryTier( - tier_name="SmallStorage", max_blocks=5, simulate_async=False - ) - large_tier = DummySecondaryTier( - tier_name="LargeStorage", max_blocks=10, simulate_async=False - ) + small_tier = DummySecondaryTier(max_blocks=5, simulate_async=False) + large_tier = DummySecondaryTier(max_blocks=10, simulate_async=False) # Create a fresh primary tier for this test primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=10) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 57750fae9141..c397f597e28a 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -155,12 +155,15 @@ def touch(self, keys: Iterable[OffloadKey]): """ return + @staticmethod @abstractmethod - def get_tier_name(self) -> str: + def get_tier_type() -> str: """ - Get the name of this tier (e.g., "Storage", "Network"). + Get the type identifier of this tier (e.g., "dummy", "storage"). + + Must match the "type" field in the tier config dict. Returns: - Tier name string. + Tier type string. """ pass diff --git a/vllm/v1/kv_offload/tiering/dummy/__init__.py b/vllm/v1/kv_offload/tiering/dummy/__init__.py index 677ac1ee279f..32cffd9fdcc9 100644 --- a/vllm/v1/kv_offload/tiering/dummy/__init__.py +++ b/vllm/v1/kv_offload/tiering/dummy/__init__.py @@ -59,7 +59,6 @@ class DummySecondaryTier(SecondaryTierManager): def __init__( self, - tier_name: str = "DummyStorage", max_blocks: int = 1000, simulate_async: bool = False, ): @@ -67,12 +66,10 @@ def __init__( Initialize the dummy secondary tier. Args: - tier_name: Name of this tier (for identification) 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. """ - self.tier_name = tier_name self.max_blocks = max_blocks self.simulate_async = simulate_async @@ -263,14 +260,9 @@ def touch(self, keys: Iterable[OffloadKey]): if key in self.blocks: self.blocks.move_to_end(key) - def get_tier_name(self) -> str: - """ - Get the name of this tier. - - Returns: - Tier name string. - """ - return self.tier_name + @staticmethod + def get_tier_type() -> str: + return "dummy" def get_num_blocks(self) -> int: """Get the number of blocks currently stored in this tier.""" diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index f5354a7a6655..63f49606ceab 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -5,6 +5,13 @@ """ from vllm.v1.kv_offload.tiering.base import SecondaryTierManager +from vllm.v1.kv_offload.tiering.dummy import DummySecondaryTier + +SUPPORTED_TIERS: tuple[type[SecondaryTierManager], ...] = (DummySecondaryTier,) + +_TIER_REGISTRY: dict[str, type[SecondaryTierManager]] = { + cls.get_tier_type(): cls for cls in SUPPORTED_TIERS +} def create_secondary_tier(tier_config: dict) -> SecondaryTierManager: @@ -14,7 +21,6 @@ def create_secondary_tier(tier_config: dict) -> SecondaryTierManager: Args: tier_config: Dictionary with tier configuration containing: - type (required): Type of secondary tier (e.g., "dummy") - - tier_name (required): Name for this tier - Additional tier-specific parameters are passed directly to the tier constructor @@ -30,15 +36,10 @@ def create_secondary_tier(tier_config: dict) -> SecondaryTierManager: if not tier_type: raise ValueError("Secondary tier configuration must include 'type'") - tier_name = config.pop("tier_name", None) - if not tier_name: - raise ValueError("Secondary tier configuration must include 'tier_name'") - - if tier_type == "dummy": - from vllm.v1.kv_offload.tiering.dummy import DummySecondaryTier - - return DummySecondaryTier(tier_name=tier_name, **config) - else: + cls = _TIER_REGISTRY.get(tier_type) + if cls is None: raise ValueError( - f"Unknown secondary tier type: {tier_type}. Supported types: dummy" + f"Unknown secondary tier type: {tier_type!r}. " + f"Supported types: {list(_TIER_REGISTRY)}" ) + return cls(**config) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 12a9a99efb3a..adbedd2fc419 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -166,7 +166,7 @@ def _process_finished_jobs(self): 3. For completed loads (secondary→primary): calls primary.complete_write() to make blocks available """ - for tier in self.secondary_tiers: + for i, tier in enumerate(self.secondary_tiers): for completed_job in tier.get_finished(): job_id = completed_job.job_id @@ -186,9 +186,11 @@ def _process_finished_jobs(self): else: # Job ID not found in either dictionary - this shouldn't happen logger.error( - "Received finished job for unknown job_id %d from tier %s", + "Received finished job for unknown job_id %d" + " from tier #%d (%s)", job_id, - tier.get_tier_name(), + i, + tier.get_tier_type(), ) def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 397dd59b8090..8e53ced98293 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -14,7 +14,6 @@ - secondary_tiers: (optional) List of secondary tier configurations Each secondary tier config is a dict with: - type: (required) Type of secondary tier (e.g., "dummy", "storage", "network") - - tier_name: (required) Name for this tier (used for logging and identification) - Additional tier-specific parameters are passed directly to the tier constructor. See each tier's documentation for supported parameters. @@ -26,7 +25,6 @@ "secondary_tiers": [ { "type": "dummy", - "tier_name": "TestStorage", # Tier-specific parameters (for DummySecondaryTier): "max_blocks": 10000, "simulate_async": False @@ -125,14 +123,14 @@ def get_manager(self) -> OffloadingManager: # Create secondary tiers secondary_tiers = [] - for tier_config in self.secondary_tier_configs: + for i, tier_config in enumerate(self.secondary_tier_configs): try: tier = create_secondary_tier(tier_config) secondary_tiers.append(tier) logger.info( - "Created secondary tier: %s (type: %s)", - tier.get_tier_name(), - tier_config.get("type"), + "Created secondary tier #%d (%s)", + i, + tier.get_tier_type(), ) except Exception as e: logger.error( From b8897c31a41e860997df8cfb9a374a8ffd0db248 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 28 Apr 2026 10:43:14 +0300 Subject: [PATCH 14/55] Address review: rename dummy -> example Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 24 +++++++++---------- vllm/v1/kv_offload/tiering/base.py | 2 +- .../tiering/{dummy => example}/__init__.py | 16 ++++++------- vllm/v1/kv_offload/tiering/factory.py | 6 ++--- vllm/v1/kv_offload/tiering/spec.py | 6 ++--- 5 files changed, 27 insertions(+), 27 deletions(-) rename vllm/v1/kv_offload/tiering/{dummy => example}/__init__.py (95%) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index 88372f853c11..276beb266b89 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 DummySecondaryTier. +Unit tests for TieringOffloadingManager and ExampleSecondaryTier. These tests verify: 1. Basic tiered offloading operations (store, load, lookup) @@ -24,7 +24,7 @@ ) from vllm.v1.kv_offload.mediums import CPULoadStoreSpec from vllm.v1.kv_offload.tiering.base import JobMetadata -from vllm.v1.kv_offload.tiering.dummy import DummySecondaryTier +from vllm.v1.kv_offload.tiering.example import ExampleSecondaryTier from vllm.v1.kv_offload.tiering.manager import ( CPUPrimaryTierOffloadingManager, TieringOffloadingManager, @@ -54,12 +54,12 @@ def count_hits(manager, keys: list[OffloadKey]) -> int | None: return count -class TestDummySecondaryTier: - """Tests for DummySecondaryTier implementation.""" +class TestExampleSecondaryTier: + """Tests for ExampleSecondaryTier implementation.""" def test_basic_store_and_lookup(self): """Test basic store and lookup operations.""" - tier = DummySecondaryTier(max_blocks=10) + tier = ExampleSecondaryTier(max_blocks=10) # Initially empty blocks = to_keys(range(3)) @@ -78,7 +78,7 @@ def test_basic_store_and_lookup(self): def test_in_flight_blocks_return_none(self): """Test that in-flight blocks cause lookup to return None.""" - tier = DummySecondaryTier(max_blocks=10) + tier = ExampleSecondaryTier(max_blocks=10) blocks = to_keys(range(3)) @@ -91,7 +91,7 @@ def test_in_flight_blocks_return_none(self): def test_lru_eviction(self): """Test LRU eviction policy.""" - tier = DummySecondaryTier(max_blocks=3) + tier = ExampleSecondaryTier(max_blocks=3) # Fill tier to capacity blocks = to_keys(range(3)) @@ -129,7 +129,7 @@ def test_lru_eviction(self): def test_async_simulation(self): """Test simulated async behavior.""" - tier = DummySecondaryTier(max_blocks=10, simulate_async=True) + tier = ExampleSecondaryTier(max_blocks=10, simulate_async=True) blocks = to_keys(range(2)) @@ -172,8 +172,8 @@ def manager_setup(self): self.primary_tier.create_kv_memoryview = lambda: memoryview(mock_arr) # Create secondary tiers - self.secondary_tier1 = DummySecondaryTier(max_blocks=10) - self.secondary_tier2 = DummySecondaryTier(max_blocks=10) + self.secondary_tier1 = ExampleSecondaryTier(max_blocks=10) + self.secondary_tier2 = ExampleSecondaryTier(max_blocks=10) # Create tiered manager self.manager = TieringOffloadingManager( @@ -353,8 +353,8 @@ def test_failed_store_no_cascade(self, manager_setup): def test_multiple_secondary_tiers_independent_eviction(self): """Test that secondary tiers manage their own evictions.""" # Create tier with small capacity - small_tier = DummySecondaryTier(max_blocks=5, simulate_async=False) - large_tier = DummySecondaryTier(max_blocks=10, simulate_async=False) + small_tier = ExampleSecondaryTier(max_blocks=5, simulate_async=False) + large_tier = ExampleSecondaryTier(max_blocks=10, simulate_async=False) # Create a fresh primary tier for this test primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=10) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index c397f597e28a..1c0c460eafdd 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -159,7 +159,7 @@ def touch(self, keys: Iterable[OffloadKey]): @abstractmethod def get_tier_type() -> str: """ - Get the type identifier of this tier (e.g., "dummy", "storage"). + Get the type identifier of this tier (e.g., "example", "storage"). Must match the "type" field in the tier config dict. diff --git a/vllm/v1/kv_offload/tiering/dummy/__init__.py b/vllm/v1/kv_offload/tiering/example/__init__.py similarity index 95% rename from vllm/v1/kv_offload/tiering/dummy/__init__.py rename to vllm/v1/kv_offload/tiering/example/__init__.py index 32cffd9fdcc9..b4d68b3d8ad8 100644 --- a/vllm/v1/kv_offload/tiering/dummy/__init__.py +++ b/vllm/v1/kv_offload/tiering/example/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -DummySecondaryTier: A simple in-memory secondary tier for testing. +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 @@ -32,21 +32,21 @@ class _JobMetadata: is_store: bool # True for store jobs, False for load jobs -class DummyLoadStoreSpec(LoadStoreSpec): +class ExampleLoadStoreSpec(LoadStoreSpec): """ - Minimal LoadStoreSpec for DummySecondaryTier testing. + Minimal LoadStoreSpec for ExampleSecondaryTier testing. - This spec is never actually used for I/O since DummySecondaryTier + This spec is never actually used for I/O since ExampleSecondaryTier stores blocks in memory. It exists to provide proper type semantics and serve as a template for real secondary tier implementations. """ @staticmethod def medium() -> str: - return "Dummy" + return "Example" -class DummySecondaryTier(SecondaryTierManager): +class ExampleSecondaryTier(SecondaryTierManager): """ A simple in-memory secondary tier for testing. @@ -63,7 +63,7 @@ def __init__( simulate_async: bool = False, ): """ - Initialize the dummy secondary tier. + Initialize the example secondary tier. Args: max_blocks: Maximum number of blocks this tier can store @@ -262,7 +262,7 @@ def touch(self, keys: Iterable[OffloadKey]): @staticmethod def get_tier_type() -> str: - return "dummy" + return "example" def get_num_blocks(self) -> int: """Get the number of blocks currently stored in this tier.""" diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index 63f49606ceab..9653175c02c1 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -5,9 +5,9 @@ """ from vllm.v1.kv_offload.tiering.base import SecondaryTierManager -from vllm.v1.kv_offload.tiering.dummy import DummySecondaryTier +from vllm.v1.kv_offload.tiering.example import ExampleSecondaryTier -SUPPORTED_TIERS: tuple[type[SecondaryTierManager], ...] = (DummySecondaryTier,) +SUPPORTED_TIERS: tuple[type[SecondaryTierManager], ...] = (ExampleSecondaryTier,) _TIER_REGISTRY: dict[str, type[SecondaryTierManager]] = { cls.get_tier_type(): cls for cls in SUPPORTED_TIERS @@ -20,7 +20,7 @@ def create_secondary_tier(tier_config: dict) -> SecondaryTierManager: Args: tier_config: Dictionary with tier configuration containing: - - type (required): Type of secondary tier (e.g., "dummy") + - type (required): Type of secondary tier (e.g., "example") - Additional tier-specific parameters are passed directly to the tier constructor diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 8e53ced98293..bd8f133f7793 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -13,7 +13,7 @@ "arc" (default: "lru") - secondary_tiers: (optional) List of secondary tier configurations Each secondary tier config is a dict with: - - type: (required) Type of secondary tier (e.g., "dummy", "storage", "network") + - type: (required) Type of secondary tier (e.g., "example", "storage", "network") - Additional tier-specific parameters are passed directly to the tier constructor. See each tier's documentation for supported parameters. @@ -24,8 +24,8 @@ "eviction_policy": "lru", "secondary_tiers": [ { - "type": "dummy", - # Tier-specific parameters (for DummySecondaryTier): + "type": "example", + # Tier-specific parameters (for ExampleSecondaryTier): "max_blocks": 10000, "simulate_async": False } From 56c3735624e0803cad3f5839c2edfe2d03afb83b Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 28 Apr 2026 10:46:45 +0300 Subject: [PATCH 15/55] Remove unused ExampleLoadStoreSpec class Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/example/__init__.py | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/example/__init__.py b/vllm/v1/kv_offload/tiering/example/__init__.py index b4d68b3d8ad8..8a0aa7c3577c 100644 --- a/vllm/v1/kv_offload/tiering/example/__init__.py +++ b/vllm/v1/kv_offload/tiering/example/__init__.py @@ -13,7 +13,7 @@ from collections.abc import Iterable from dataclasses import dataclass -from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadKey, ReqContext +from vllm.v1.kv_offload.abstract import OffloadKey, ReqContext from vllm.v1.kv_offload.mediums import CPULoadStoreSpec from vllm.v1.kv_offload.tiering.base import ( JobId, @@ -32,20 +32,6 @@ class _JobMetadata: is_store: bool # True for store jobs, False for load jobs -class ExampleLoadStoreSpec(LoadStoreSpec): - """ - Minimal LoadStoreSpec for ExampleSecondaryTier testing. - - This spec is never actually used for I/O since ExampleSecondaryTier - stores blocks in memory. It exists to provide proper type semantics - and serve as a template for real secondary tier implementations. - """ - - @staticmethod - def medium() -> str: - return "Example" - - class ExampleSecondaryTier(SecondaryTierManager): """ A simple in-memory secondary tier for testing. From 675e0565e51a615663b5ff0ca3ea4246777b8be4 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 28 Apr 2026 12:19:49 +0300 Subject: [PATCH 16/55] Address review: widen JobMetadata.keys from list to Sequence Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/base.py | 4 +-- .../v1/kv_offload/tiering/example/__init__.py | 28 +++++++++---------- 2 files changed, 15 insertions(+), 17 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 1c0c460eafdd..eeceb3329023 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -5,7 +5,7 @@ """ from abc import ABC, abstractmethod -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from dataclasses import dataclass, field from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadKey, ReqContext @@ -19,7 +19,7 @@ class JobMetadata: """Metadata for an in-flight async transfer job.""" job_id: JobId - keys: list[OffloadKey] + keys: Sequence[OffloadKey] spec: LoadStoreSpec req_context: ReqContext = field(default_factory=ReqContext) diff --git a/vllm/v1/kv_offload/tiering/example/__init__.py b/vllm/v1/kv_offload/tiering/example/__init__.py index 8a0aa7c3577c..f3f6fc944fae 100644 --- a/vllm/v1/kv_offload/tiering/example/__init__.py +++ b/vllm/v1/kv_offload/tiering/example/__init__.py @@ -10,7 +10,7 @@ """ from collections import OrderedDict -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from dataclasses import dataclass from vllm.v1.kv_offload.abstract import OffloadKey, ReqContext @@ -28,7 +28,7 @@ class _JobMetadata: """Internal metadata for tracking job details.""" job_id: JobId - keys: list[OffloadKey] + keys: Sequence[OffloadKey] is_store: bool # True for store jobs, False for load jobs @@ -101,20 +101,20 @@ def submit_store(self, job_metadata: JobMetadata) -> None: spec for reading blocks from the primary tier. """ job_id = job_metadata.job_id - keys_list = list(job_metadata.keys) + keys = job_metadata.keys primary_read_spec = job_metadata.spec # Validate spec type and consistency assert isinstance(primary_read_spec, CPULoadStoreSpec), ( f"Expected CPULoadStoreSpec, got {type(primary_read_spec)}" ) - assert len(keys_list) == len(primary_read_spec.block_ids), ( - f"Length mismatch: {len(keys_list)} keys but " + assert len(keys) == len(primary_read_spec.block_ids), ( + f"Length mismatch: {len(keys)} keys but " f"{len(primary_read_spec.block_ids)} block_ids in spec" ) # Filter out blocks already present - blocks_to_store = [bh for bh in keys_list if bh not in self.blocks] + blocks_to_store = [bh for bh in keys if bh not in self.blocks] if not blocks_to_store: # All blocks already present @@ -128,7 +128,7 @@ def submit_store(self, job_metadata: JobMetadata) -> None: evicted = [] if num_blocks_to_evict > 0: # Collect eviction candidates first (LRU order), then delete atomically - protected = set(keys_list) + protected = set(keys) for key in self.blocks: if key not in protected and key not in self.in_flight: evicted.append(key) @@ -165,31 +165,29 @@ def submit_load(self, job_metadata: JobMetadata) -> None: spec for writing blocks into the primary tier. """ job_id = job_metadata.job_id - keys_list = list(job_metadata.keys) + keys = job_metadata.keys primary_write_spec = job_metadata.spec # Validate spec type and consistency assert isinstance(primary_write_spec, CPULoadStoreSpec), ( f"Expected CPULoadStoreSpec, got {type(primary_write_spec)}" ) - assert len(keys_list) == len(primary_write_spec.block_ids), ( - f"Length mismatch: {len(keys_list)} keys but " + assert len(keys) == len(primary_write_spec.block_ids), ( + f"Length mismatch: {len(keys)} keys but " f"{len(primary_write_spec.block_ids)} block_ids in spec" ) # Verify all blocks exist - for key in keys_list: + for key in keys: if key not in self.blocks: return # Mark blocks as in-flight - for key in keys_list: + for key in keys: self.in_flight[key] = job_id # Create internal job metadata - internal_job_metadata = _JobMetadata( - job_id=job_id, keys=keys_list, is_store=False - ) + 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 From a9426dafd3046eb617b542043088cdbe8a2889a8 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 29 Apr 2026 10:11:05 +0300 Subject: [PATCH 17/55] Address review: replace `LoadStoreSpec` with `block_ids: np.ndarray` Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 6 ++--- vllm/v1/kv_offload/tiering/base.py | 16 ++++++------- .../v1/kv_offload/tiering/example/__init__.py | 23 +++++-------------- vllm/v1/kv_offload/tiering/manager.py | 6 +++-- 4 files changed, 21 insertions(+), 30 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index 276beb266b89..1b91f53c35a6 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -14,6 +14,7 @@ from collections.abc import Iterable from unittest.mock import MagicMock +import numpy as np import pytest import torch @@ -22,7 +23,6 @@ ReqContext, make_offload_key, ) -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec from vllm.v1.kv_offload.tiering.base import JobMetadata from vllm.v1.kv_offload.tiering.example import ExampleSecondaryTier from vllm.v1.kv_offload.tiering.manager import ( @@ -113,7 +113,7 @@ def test_lru_eviction(self): JobMetadata( job_id=1, keys=[new_block], - spec=CPULoadStoreSpec([0]), + block_ids=np.array([0], dtype=np.int64), ) ) @@ -141,7 +141,7 @@ def test_async_simulation(self): JobMetadata( job_id=1, keys=blocks, - spec=CPULoadStoreSpec([0, 1]), + block_ids=np.array([0, 1], dtype=np.int64), ) ) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index eeceb3329023..06ae763df963 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -8,7 +8,9 @@ from collections.abc import Iterable, Sequence from dataclasses import dataclass, field -from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadKey, ReqContext +import numpy as np + +from vllm.v1.kv_offload.abstract import OffloadKey, ReqContext # Type alias for job IDs used in async transfer tracking JobId = int @@ -20,7 +22,7 @@ class JobMetadata: job_id: JobId keys: Sequence[OffloadKey] - spec: LoadStoreSpec + block_ids: np.ndarray req_context: ReqContext = field(default_factory=ReqContext) @@ -73,7 +75,7 @@ def submit_store(self, job_metadata: JobMetadata) -> None: calling thread. The caller (TieringOffloadingManager) must have already called - primary.prepare_read(keys) to obtain job_metadata.spec and + primary.prepare_read(keys) to obtain job_metadata.block_ids and to increment ref_cnt on those blocks. ref_cnt will be decremented when get_finished() reports this job_id as complete and primary.unprepare_read() is called. @@ -87,9 +89,8 @@ def submit_store(self, job_metadata: JobMetadata) -> None: Args: job_metadata: Job metadata including job_id, keys, and - spec for reading blocks from the primary tier + block_ids for reading blocks from the primary tier (obtained via primary.prepare_read()). - spec is a CPULoadStoreSpec with block_ids. """ pass @@ -104,16 +105,15 @@ def submit_load(self, job_metadata: JobMetadata) -> None: the calling thread. The caller (TieringOffloadingManager) must have already called - primary.prepare_write(keys) to obtain job_metadata.spec and + primary.prepare_write(keys) to obtain job_metadata.block_ids and to allocate space in the primary tier. When get_finished() reports this job_id as complete, primary.complete_write() is called to make the blocks available for GPU loads. Args: job_metadata: Job metadata including job_id, keys, and - spec for writing blocks into the primary tier + block_ids for writing blocks into the primary tier (obtained via primary.prepare_write()). - spec is a CPULoadStoreSpec with block_ids. """ pass diff --git a/vllm/v1/kv_offload/tiering/example/__init__.py b/vllm/v1/kv_offload/tiering/example/__init__.py index f3f6fc944fae..5dc776431b4a 100644 --- a/vllm/v1/kv_offload/tiering/example/__init__.py +++ b/vllm/v1/kv_offload/tiering/example/__init__.py @@ -14,7 +14,6 @@ from dataclasses import dataclass from vllm.v1.kv_offload.abstract import OffloadKey, ReqContext -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec from vllm.v1.kv_offload.tiering.base import ( JobId, JobMetadata, @@ -102,15 +101,10 @@ def submit_store(self, job_metadata: JobMetadata) -> None: """ job_id = job_metadata.job_id keys = job_metadata.keys - primary_read_spec = job_metadata.spec + block_ids = job_metadata.block_ids - # Validate spec type and consistency - assert isinstance(primary_read_spec, CPULoadStoreSpec), ( - f"Expected CPULoadStoreSpec, got {type(primary_read_spec)}" - ) - assert len(keys) == len(primary_read_spec.block_ids), ( - f"Length mismatch: {len(keys)} keys but " - f"{len(primary_read_spec.block_ids)} block_ids in spec" + assert len(keys) == len(block_ids), ( + f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids" ) # Filter out blocks already present @@ -166,15 +160,10 @@ def submit_load(self, job_metadata: JobMetadata) -> None: """ job_id = job_metadata.job_id keys = job_metadata.keys - primary_write_spec = job_metadata.spec + block_ids = job_metadata.block_ids - # Validate spec type and consistency - assert isinstance(primary_write_spec, CPULoadStoreSpec), ( - f"Expected CPULoadStoreSpec, got {type(primary_write_spec)}" - ) - assert len(keys) == len(primary_write_spec.block_ids), ( - f"Length mismatch: {len(keys)} keys but " - f"{len(primary_write_spec.block_ids)} block_ids in spec" + assert len(keys) == len(block_ids), ( + f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids" ) # Verify all blocks exist diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index adbedd2fc419..866573ae5359 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -268,7 +268,7 @@ def _initiate_promotion( job_metadata = JobMetadata( job_id=job_id, keys=primary_store_result.keys_to_store, - spec=primary_store_result.store_spec, + block_ids=primary_store_result.store_spec.block_ids, req_context=req_context, ) self._load_jobs[job_id] = job_metadata @@ -408,7 +408,9 @@ def complete_store( # Track this store job job_metadata = JobMetadata( - job_id=job_id, keys=keys_list, spec=primary_blocks_spec + job_id=job_id, + keys=keys_list, + block_ids=primary_blocks_spec.block_ids, ) self._store_jobs[job_id] = job_metadata From 2ba6523e3d7947b88c3c4abcae18424ef9c145bd Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 29 Apr 2026 11:37:18 +0300 Subject: [PATCH 18/55] Make pre-commit happy Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 866573ae5359..625e59c7f6f0 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -33,6 +33,7 @@ ) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion +from vllm.v1.kv_offload.mediums import CPULoadStoreSpec from vllm.v1.kv_offload.tiering.base import ( JobId, JobMetadata, @@ -265,10 +266,12 @@ def _initiate_promotion( job_id = self._next_job_id() # Track this load job + store_spec = primary_store_result.store_spec + assert isinstance(store_spec, CPULoadStoreSpec) job_metadata = JobMetadata( job_id=job_id, keys=primary_store_result.keys_to_store, - block_ids=primary_store_result.store_spec.block_ids, + block_ids=store_spec.block_ids, req_context=req_context, ) self._load_jobs[job_id] = job_metadata @@ -407,6 +410,7 @@ def complete_store( job_id = self._next_job_id() # Track this store job + assert isinstance(primary_blocks_spec, CPULoadStoreSpec) job_metadata = JobMetadata( job_id=job_id, keys=keys_list, From bafc17e1113015bae5defbae0ba61c9011496f40 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 29 Apr 2026 15:51:40 +0300 Subject: [PATCH 19/55] Fix merge issues: Iterable to Sequence Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/base.py | 2 +- vllm/v1/kv_offload/tiering/example/__init__.py | 4 ++-- vllm/v1/kv_offload/tiering/manager.py | 9 ++++----- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 06ae763df963..c8f03c5f18ae 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -146,7 +146,7 @@ def set_primary_view(self, view: memoryview) -> None: """ return - def touch(self, keys: Iterable[OffloadKey]): + def touch(self, keys: Sequence[OffloadKey]): """ Mark blocks as recently used for eviction policy. diff --git a/vllm/v1/kv_offload/tiering/example/__init__.py b/vllm/v1/kv_offload/tiering/example/__init__.py index 5dc776431b4a..93cb5e79d3af 100644 --- a/vllm/v1/kv_offload/tiering/example/__init__.py +++ b/vllm/v1/kv_offload/tiering/example/__init__.py @@ -222,14 +222,14 @@ def _complete_load_job(self, job_metadata: _JobMetadata): # 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: Iterable[OffloadKey]): + def touch(self, keys: Sequence[OffloadKey]): """ Mark blocks as recently used (move to end of LRU list). Args: keys: Blocks to mark as recently used. """ - for key in reversed(list(keys)): + for key in reversed(keys): if key in self.blocks: self.blocks.move_to_end(key) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 625e59c7f6f0..afc9bbe6e49a 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -20,7 +20,7 @@ protecting blocks from eviction until complete_read() is called """ -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from vllm.logger import init_logger from vllm.v1.kv_offload.abstract import ( @@ -279,7 +279,7 @@ def _initiate_promotion( tier.submit_load(job_metadata) def prepare_load( - self, keys: Iterable[OffloadKey], req_context: ReqContext + self, keys: Sequence[OffloadKey], req_context: ReqContext ) -> LoadStoreSpec: """ Prepare blocks to be loaded from primary tier to GPU. @@ -302,14 +302,13 @@ def prepare_load( return self.primary_tier.prepare_load(keys, req_context) - def touch(self, keys: Iterable[OffloadKey]): + def touch(self, keys: Sequence[OffloadKey]): """ Mark blocks as recently used in all tiers. Args: keys: Blocks to mark as recently used. """ - keys = list(keys) self.primary_tier.touch(keys) for tier in self.secondary_tiers: tier.touch(keys) @@ -327,7 +326,7 @@ def complete_load(self, keys: Iterable[OffloadKey]): self.primary_tier.complete_load(keys) def prepare_store( - self, keys: Iterable[OffloadKey], req_context: ReqContext + self, keys: Sequence[OffloadKey], req_context: ReqContext ) -> PrepareStoreOutput | None: """ Prepare blocks to be stored from GPU to primary tier. From ad714be21dc6bb3931baa81705003382bb222354 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 30 Apr 2026 12:50:33 +0300 Subject: [PATCH 20/55] Fix merge issues: update imports Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 2 +- vllm/v1/kv_offload/tiering/spec.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 996adfa2e5cd..dabe5c8a2cb0 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -31,9 +31,9 @@ PrepareStoreOutput, ReqContext, ) +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec from vllm.v1.kv_offload.tiering.base import ( JobId, JobMetadata, diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index ac2105d19ecf..d4b58202892a 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -38,10 +38,9 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.kv_offload.base import OffloadingManager +from vllm.v1.kv_offload.base import CanonicalKVCaches, OffloadingManager 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.spec import CanonicalKVCaches from vllm.v1.kv_offload.tiering.factory import create_secondary_tier from vllm.v1.kv_offload.tiering.manager import ( CPUPrimaryTierOffloadingManager, From 2029d4616b285219a8dae1631e1e9bdce62e57eb Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Sun, 3 May 2026 23:24:27 +0300 Subject: [PATCH 21/55] Fix merge issues: update imports Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/spec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index d4b58202892a..a31a22fd1575 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -39,6 +39,7 @@ from vllm.logger import init_logger from vllm.v1.kv_cache_interface import KVCacheConfig 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 @@ -46,7 +47,6 @@ CPUPrimaryTierOffloadingManager, TieringOffloadingManager, ) -from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers logger = init_logger(__name__) From 282cfe744ca515a4c93e4354c69d3389980e7e54 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 4 May 2026 10:37:38 +0300 Subject: [PATCH 22/55] Fix merge issues: change keys type Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index dabe5c8a2cb0..3c62f919647b 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -20,7 +20,7 @@ protecting blocks from eviction until complete_read() is called """ -from collections.abc import Iterable, Sequence +from collections.abc import Collection, Iterable from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( @@ -279,7 +279,7 @@ def _initiate_promotion( tier.submit_load(job_metadata) def prepare_load( - self, keys: Sequence[OffloadKey], req_context: ReqContext + self, keys: Collection[OffloadKey], req_context: ReqContext ) -> LoadStoreSpec: """ Prepare blocks to be loaded from primary tier to GPU. @@ -302,7 +302,7 @@ def prepare_load( return self.primary_tier.prepare_load(keys, req_context) - def touch(self, keys: Sequence[OffloadKey]): + def touch(self, keys: Collection[OffloadKey]): """ Mark blocks as recently used in all tiers. @@ -313,7 +313,7 @@ def touch(self, keys: Sequence[OffloadKey]): for tier in self.secondary_tiers: tier.touch(keys) - def complete_load(self, keys: Iterable[OffloadKey]): + def complete_load(self, keys: Collection[OffloadKey]): """ Mark blocks as done loading from primary tier to GPU. @@ -326,7 +326,7 @@ def complete_load(self, keys: Iterable[OffloadKey]): self.primary_tier.complete_load(keys) def prepare_store( - self, keys: Sequence[OffloadKey], req_context: ReqContext + self, keys: Collection[OffloadKey], req_context: ReqContext ) -> PrepareStoreOutput | None: """ Prepare blocks to be stored from GPU to primary tier. @@ -358,7 +358,7 @@ def prepare_store( def complete_store( self, - keys: Iterable[OffloadKey], + keys: Collection[OffloadKey], success: bool = True, ): """ From 6ffec8b3323f4c22dfd455bba4c5c5721d8a83b9 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 4 May 2026 11:01:37 +0300 Subject: [PATCH 23/55] Fix merge issues: sequence -> collection Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/base.py | 6 +++--- vllm/v1/kv_offload/tiering/example/__init__.py | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 47e65de3d785..500fca0e701b 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -5,7 +5,7 @@ """ from abc import ABC, abstractmethod -from collections.abc import Iterable, Sequence +from collections.abc import Collection, Iterable from dataclasses import dataclass, field import numpy as np @@ -21,7 +21,7 @@ class JobMetadata: """Metadata for an in-flight async transfer job.""" job_id: JobId - keys: Sequence[OffloadKey] + keys: Collection[OffloadKey] block_ids: np.ndarray req_context: ReqContext = field(default_factory=ReqContext) @@ -146,7 +146,7 @@ def set_primary_view(self, view: memoryview) -> None: """ return - def touch(self, keys: Sequence[OffloadKey]): + def touch(self, keys: Collection[OffloadKey]): """ Mark blocks as recently used for eviction policy. diff --git a/vllm/v1/kv_offload/tiering/example/__init__.py b/vllm/v1/kv_offload/tiering/example/__init__.py index 7eada1a35339..5903672bbca6 100644 --- a/vllm/v1/kv_offload/tiering/example/__init__.py +++ b/vllm/v1/kv_offload/tiering/example/__init__.py @@ -10,7 +10,7 @@ """ from collections import OrderedDict -from collections.abc import Iterable, Sequence +from collections.abc import Collection, Iterable from dataclasses import dataclass from vllm.v1.kv_offload.base import OffloadKey, ReqContext @@ -27,7 +27,7 @@ class _JobMetadata: """Internal metadata for tracking job details.""" job_id: JobId - keys: Sequence[OffloadKey] + keys: Collection[OffloadKey] is_store: bool # True for store jobs, False for load jobs @@ -222,14 +222,14 @@ def _complete_load_job(self, job_metadata: _JobMetadata): # 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: Sequence[OffloadKey]): + def touch(self, keys: Collection[OffloadKey]): """ Mark blocks as recently used (move to end of LRU list). Args: keys: Blocks to mark as recently used. """ - for key in reversed(keys): + for key in reversed(list(keys)): if key in self.blocks: self.blocks.move_to_end(key) From e80f48bd6a21063f086f62ee6fdc5c73f20b173b Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 4 May 2026 11:46:06 +0300 Subject: [PATCH 24/55] Change SecondaryTierManager.lookup() to accept single key instead of multiple keys Signed-off-by: Ronen Schaffer --- tests/v1/kv_offload/cpu/test_manager.py | 37 +++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index ef5d61e7b3d2..96c766f14bfb 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -261,6 +261,43 @@ def test_cpu_manager(): ) +def test_prepare_load_preserves_key_order(): + """block_ids[i] must correspond to keys[i] (co-indexed invariant).""" + manager = CPUOffloadingManager(num_blocks=4, cache_policy="lru") + + key_a, key_b, key_c = to_key(0), to_key(1), to_key(2) + + # Store all three keys and learn their block ID assignments + store_output = manager.prepare_store([key_a, key_b, key_c], _EMPTY_REQ_CTX) + assert store_output is not None + assert isinstance(store_output.store_spec, CPULoadStoreSpec) + key_to_block_id = { + k: int(bid) + for k, bid in zip(store_output.keys_to_store, store_output.store_spec.block_ids) + } + manager.complete_store([key_a, key_b, key_c]) + + # Forward order: [a, b, c] + spec_fwd = manager.prepare_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) + assert isinstance(spec_fwd, CPULoadStoreSpec) + assert [int(x) for x in spec_fwd.block_ids] == [ + key_to_block_id[key_a], + key_to_block_id[key_b], + key_to_block_id[key_c], + ] + manager.complete_load([key_a, key_b, key_c]) # order irrelevant + + # Arbitrary permutation: [b, c, a] + spec_perm = manager.prepare_load([key_b, key_c, key_a], _EMPTY_REQ_CTX) + assert isinstance(spec_perm, CPULoadStoreSpec) + assert [int(x) for x in spec_perm.block_ids] == [ + key_to_block_id[key_b], + key_to_block_id[key_c], + key_to_block_id[key_a], + ] + manager.complete_load([key_a, key_b, key_c]) # order irrelevant + + class TestARCPolicy: """Unit tests for CPUOffloadingManager with ARC eviction policy.""" From 6875f73412fb7646e27e416dccc302f385e9d79a Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 4 May 2026 12:16:55 +0300 Subject: [PATCH 25/55] Remove unnecessary list materialism Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 3c62f919647b..de109c6bf027 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -379,20 +379,13 @@ def complete_store( success: Whether the GPU→primary transfer succeeded. req_context: Per-request context forwarded to primary.prepare_read(). """ - # Materialize only if success=True (needed for cascading to secondary tiers) - keys_list = list(keys) if success else keys - # Step 1: Complete store in primary tier (makes blocks loadable) - self.primary_tier.complete_store(keys_list, success) + self.primary_tier.complete_store(keys, success) if not success: # If GPU→Primary transfer failed, don't cascade to secondary tiers return - # At this point, success=True is guaranteed, so keys_list - # is list[OffloadKey] - assert isinstance(keys_list, list) - # Step 2: Cascade to ALL secondary tiers # For each secondary tier, call primary.prepare_read() to get the # LoadStoreSpec AND to increment ref_cnt (protecting blocks from @@ -401,9 +394,7 @@ def complete_store( for tier in self.secondary_tiers: # Get spec for reading from primary tier AND increment ref_cnt # TODO: pass the actual req_context instead of None - primary_blocks_spec = self.primary_tier.prepare_read( - keys_list, ReqContext() - ) + primary_blocks_spec = self.primary_tier.prepare_read(keys, ReqContext()) # Submit async store job: primary→secondary job_id = self._next_job_id() @@ -412,7 +403,7 @@ def complete_store( assert isinstance(primary_blocks_spec, CPULoadStoreSpec) job_metadata = JobMetadata( job_id=job_id, - keys=keys_list, + keys=keys, block_ids=primary_blocks_spec.block_ids, ) self._store_jobs[job_id] = job_metadata From 7e77867757160132ea5df341f7d1b3fa38c6d87d Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 4 May 2026 14:59:46 +0300 Subject: [PATCH 26/55] Refactor tiering.lookup() Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 47 +++++++++++++++------------ 1 file changed, 26 insertions(+), 21 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index de109c6bf027..6a681390f57a 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -200,40 +200,45 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: Algorithm: 1. Process any completed async jobs first - 2. Check primary tier - 3. If not in primary, check secondary tiers and initiate promotion - if the block is found there + 2. Query all tiers (primary + secondaries) unconditionally + 3. Decide based on combined results Args: key: Block hash to look up. req_context: Per-request context. Returns: - True if the block is in the primary tier and ready, + True if the block is ready (primary hit, or found in a secondary + tier with promotion initiated), False if not found in any tier, - None if the block is being transferred (retry later). + None if no tier has the block but at least one tier is busy + (retry later). """ self._process_finished_jobs() - # Step 1: Check primary tier + # Always query every tier to warm up caches / prefetch state primary_hit = self.primary_tier.lookup(key, req_context) - if primary_hit is None: - return None - if primary_hit: - return True - # Step 2: Check secondary tiers + hit_tier = None + any_none = False for tier in self.secondary_tiers: - secondary_hit = tier.lookup(key, req_context) - if secondary_hit is None: - # Tier is busy with this block - return None - if secondary_hit: - # Found in secondary — initiate promotion and signal retry - self._initiate_promotion(tier, key, req_context) - return None - - return False + result = tier.lookup(key, req_context) + if result is True and hit_tier is None: + hit_tier = tier + elif result is None: + any_none = True + + if primary_hit: + return True + elif primary_hit is None: + return None + elif hit_tier is not None: + self._initiate_promotion(hit_tier, key, req_context) + return True + elif any_none: + return None + else: + return False def _initiate_promotion( self, From 82af98ada9ea9467674b5289078b92b8b172608b Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 6 May 2026 10:35:37 +0300 Subject: [PATCH 27/55] cherry-pick: Update lookup() to return None when block isn't ready Signed-off-by: Ronen Schaffer --- tests/v1/kv_offload/cpu/test_manager.py | 12 ++++++------ vllm/v1/kv_offload/cpu/manager.py | 6 +++++- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 96c766f14bfb..d0e89feee856 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -163,9 +163,9 @@ def test_cpu_manager(): ), ) - # lookup [1, 2] -> not ready - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False + # lookup [1, 2] -> write in-flight, not yet ready + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is None + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is None # no events so far assert list(cpu_manager.take_events()) == [] @@ -333,9 +333,9 @@ def test_basic(self): ), ) - # lookup [1, 2] -> not ready - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False + # lookup [1, 2] -> write in-flight, not yet ready + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is None + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is None # no events so far assert list(cpu_manager.take_events()) == [] diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 90d56915ccf4..80bcb568f99a 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -86,7 +86,11 @@ def _get_load_store_spec( def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: block = self._policy.get(key) - return block is not None and block.is_ready + if block is None: + return False + if not block.is_ready: + return None # write in-flight; caller should retry + return True def prepare_load( self, From b4df1360c27f7004a88932c5aaf12c66393d0375 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 6 May 2026 16:01:15 +0300 Subject: [PATCH 28/55] batch and defer submit_load per tier and request Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 74 +++++++++++++-- vllm/v1/kv_offload/tiering/manager.py | 93 ++++++++++++++----- 2 files changed, 138 insertions(+), 29 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index bebbf2216d9e..d8523b466eb5 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -266,6 +266,9 @@ def test_promotion_from_secondary(self, manager_setup): result = self.manager.lookup(block, _CTX) assert result is None # Retry later (promotion initiated) + # Flush deferred submit_load() calls (normally done at end of engine step) + list(self.manager.take_events()) + # Process finished jobs to complete promotion self.manager._process_finished_jobs() @@ -413,20 +416,75 @@ def test_prepare_store_processes_finished_jobs_first(self, manager_setup): block = self.primary_tier._policy.get(block_hash) assert block.ref_cnt == 0 - def test_req_context_propagated_to_submit_load(self, manager_setup): - """Test that req_context from lookup() is forwarded to submit_load.""" - block = to_keys([0])[0] - self.secondary_tier1.blocks[block] = True # simulate prior cascade + def test_lookup_batches_submit_load_per_request(self, manager_setup): + """lookup() defers submit_load until take_events(), one call per request. + + Blocks from different requests each get their own submit_load call, each + carrying the correct req_context. + """ + blocks = to_keys(range(4)) + for block in blocks: + self.secondary_tier1.blocks[block] = True self.secondary_tier1.submit_load = MagicMock( wraps=self.secondary_tier1.submit_load ) - ctx = ReqContext(kv_transfer_params={"priority": "high"}) - self.manager.lookup(block, ctx) + ctx_a = ReqContext() + ctx_b = ReqContext() + + # All lookups return None: secondary hit triggers promotion (in-flight) + assert self.manager.lookup(blocks[0], ctx_a) is None + assert self.manager.lookup(blocks[1], ctx_a) is None + assert self.manager.lookup(blocks[2], ctx_b) is None + assert self.manager.lookup(blocks[3], ctx_b) is None + + # submit_load must not fire during lookup - only at end of step + self.secondary_tier1.submit_load.assert_not_called() + + # simulate end of step + list(self.manager.take_events()) + + assert self.secondary_tier1.submit_load.call_count == 2 + calls = self.secondary_tier1.submit_load.call_args_list + jm_a = calls[0].args[0] + jm_b = calls[1].args[0] + assert set(jm_a.keys) == {blocks[0], blocks[1]} + assert jm_a.req_context is ctx_a + assert set(jm_b.keys) == {blocks[2], blocks[3]} + assert jm_b.req_context is ctx_b + + def test_lookup_shared_block_no_duplicate_promotion(self, manager_setup): + """A block looked up by two requests in the same step is promoted once. + + The first lookup initiates promotion (returns None via secondary hit). + The second lookup sees ref_cnt=-1 on the primary slot and returns None + via the primary in-flight path — without triggering a second promotion. + """ + shared_block = to_keys([0])[0] + self.secondary_tier1.blocks[shared_block] = True + + self.secondary_tier1.submit_load = MagicMock( + wraps=self.secondary_tier1.submit_load + ) + + ctx_a = ReqContext() + ctx_b = ReqContext() + + result_a = self.manager.lookup(shared_block, ctx_a) + result_b = self.manager.lookup(shared_block, ctx_b) + + # Both see None (in-flight), but promotion is only queued once + assert result_a is None + assert result_b is None + + list(self.manager.take_events()) + + # Only one submit_load call despite two lookups self.secondary_tier1.submit_load.assert_called_once() - job_metadata = self.secondary_tier1.submit_load.call_args[0][0] - assert job_metadata.req_context.kv_transfer_params == {"priority": "high"} + job_metadata = self.secondary_tier1.submit_load.call_args.args[0] + assert list(job_metadata.keys) == [shared_block] + assert job_metadata.req_context is ctx_a class TestTieringOffloadingWithoutSecondaryTiers: diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 6a681390f57a..7aedc469324f 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -21,6 +21,9 @@ """ from collections.abc import Collection, Iterable +from dataclasses import dataclass, field + +import numpy as np from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( @@ -43,6 +46,15 @@ logger = init_logger(__name__) +@dataclass +class _PendingPromotion: + """Accumulator for blocks awaiting submit_load() for one (tier, request).""" + + keys: list[OffloadKey] = field(default_factory=list) + block_ids: list[int] = field(default_factory=list) + req_context: ReqContext = field(default_factory=ReqContext) + + class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): """CPUOffloadingManager with a primary/secondary transfer interface. @@ -144,6 +156,14 @@ def __init__( # Load jobs: secondary → primary transfers (promotions) self._load_jobs: dict[JobId, JobMetadata] = {} + # Pending promotion requests accumulated during lookup() calls; flushed + # as one batched submit_load() per (tier, request) in take_events(). + # Outer key: tier. Inner key: id(req_context) — the same ReqContext + # object is reused for all block lookups of a given request per engine step. + self._pending_load_submissions: dict[ + SecondaryTierManager, dict[int, _PendingPromotion] + ] = {} + # Wire each secondary tier with a long-lived memoryview of the primary # CPU buffer. One view is shared across all tiers; released in shutdown(). self._primary_kv_view = primary_tier.create_kv_memoryview() @@ -234,7 +254,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: return None elif hit_tier is not None: self._initiate_promotion(hit_tier, key, req_context) - return True + return None elif any_none: return None else: @@ -247,41 +267,65 @@ def _initiate_promotion( req_context: ReqContext, ): """ - Initiate promotion of a block from a secondary tier to the primary tier. + Queue a block for promotion from a secondary tier to the primary tier. - This method: - 1. Calls primary.prepare_write() to allocate space in primary tier - 2. Calls tier.submit_load() to start async transfer: secondary→primary - 3. Tracks the job in _load_jobs dictionary + Allocates space in the primary tier immediately (sets ref_cnt=-1 so + subsequent lookups within the same step see the slot as in-flight), + then defers the actual submit_load() call to _flush_pending_promotions() + so all blocks queued during one engine step are submitted as a single + batched job. Args: tier: The secondary tier to promote from key: Block to promote req_context: Per-request context forwarded to primary.prepare_write(). """ - # Allocate space in primary tier for promoted block + # Allocate space in primary tier for promoted block. + # Must happen immediately so primary.lookup() returns None (in-flight) + # for this key on any subsequent lookup() call within the same step, + # preventing duplicate promotion attempts. primary_store_result = self.primary_tier.prepare_write([key], req_context) if primary_store_result is None: - # Cannot allocate space in primary tier (full) - # The next lookup() will retry + # Cannot allocate space in primary tier (full); retry next step. return - # Submit async load job: secondary→primary - job_id = self._next_job_id() - - # Track this load job store_spec = primary_store_result.store_spec assert isinstance(store_spec, CPULoadStoreSpec) - job_metadata = JobMetadata( - job_id=job_id, - keys=primary_store_result.keys_to_store, - block_ids=store_spec.block_ids, - req_context=req_context, - ) - self._load_jobs[job_id] = job_metadata + # Defer submit_load to take_events(). Group by (tier, request) so each + # request's blocks are submitted as one batched job per tier. + tier_pending = self._pending_load_submissions.setdefault(tier, {}) + ctx_id = id(req_context) + if ctx_id not in tier_pending: + tier_pending[ctx_id] = _PendingPromotion( + keys=[], block_ids=[], req_context=req_context + ) + entry = tier_pending[ctx_id] + entry.keys.extend(primary_store_result.keys_to_store) + entry.block_ids.extend(store_spec.block_ids) + + def _flush_pending_promotions(self) -> None: + """Submit one batched submit_load() per (tier, request). + + Called from take_events() at the end of each engine step, flushing + all promotion requests deferred during lookup(). + """ + if not self._pending_load_submissions: + return - tier.submit_load(job_metadata) + for tier, pending_by_ctx in self._pending_load_submissions.items(): + for entry in pending_by_ctx.values(): + job_id = self._next_job_id() + job_metadata = JobMetadata( + job_id=job_id, + keys=entry.keys, + block_ids=np.array(entry.block_ids, dtype=np.int64), + req_context=entry.req_context, + ) + self._load_jobs[job_id] = job_metadata + tier.submit_load(job_metadata) + + self._pending_load_submissions.clear() def prepare_load( self, keys: Collection[OffloadKey], req_context: ReqContext @@ -428,6 +472,13 @@ def take_events(self) -> Iterable[OffloadingEvent]: Yields: New OffloadingEvents collected since the last call. """ + # TODO: Move _flush_pending_promotions() to a dedicated end_of_batch() + # hook once one exists. For now, take_events() serves as the flush + # point under the assumption that it is called at the end of each + # engine step (Scheduler.update_from_output() → connector.take_events()). + # Update the relevant tests the rely on take_events() to signal end of step. + self._flush_pending_promotions() + if self.events is not None: yield from self.events self.events.clear() From eac6e0f3aa70252188da2ea8c8c0e045f99ea0e4 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 30 Apr 2026 14:36:28 +0300 Subject: [PATCH 29/55] cherry-pick: Add ReqContext to touch, complete_load, complete_store Signed-off-by: Ronen Schaffer --- tests/v1/kv_offload/cpu/test_manager.py | 64 +++++++++---------- .../kv_connector/v1/offloading/scheduler.py | 14 ++-- vllm/v1/kv_offload/base.py | 14 +++- vllm/v1/kv_offload/cpu/manager.py | 11 +++- vllm/v1/kv_offload/reuse_manager.py | 17 +++-- 5 files changed, 70 insertions(+), 50 deletions(-) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index d0e89feee856..e1c0b7e047c5 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -117,10 +117,10 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): # store [1, 2] and complete manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) - manager.complete_store(to_keys([1, 2])) + manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) # touch [1] to make block 2 the LRU candidate - manager.touch(to_keys([1])) + manager.touch(to_keys([1]), _EMPTY_REQ_CTX) # prepare_store([2, 3, 4, 5]): # - block 2 is already stored -> filtered out of keys_to_store @@ -137,7 +137,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): ) # complete_store must not silently drop block 2 - manager.complete_store(to_keys([2, 3, 4, 5])) + manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX) # block 2 must still be present in the cache assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True @@ -171,7 +171,7 @@ def test_cpu_manager(): assert list(cpu_manager.take_events()) == [] # complete store [1, 2] - cpu_manager.complete_store(to_keys([1, 2])) + cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] @@ -199,7 +199,7 @@ def test_cpu_manager(): assert cpu_manager.prepare_store(to_keys([1, 6]), _EMPTY_REQ_CTX) is None # complete store [2, 3, 4, 5] - cpu_manager.complete_store(to_keys([2, 3, 4, 5])) + cpu_manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX) # lookup (now that we have [2, 3, 4, 5]) assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False @@ -217,7 +217,7 @@ def test_cpu_manager(): assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None # complete load [2, 3] - cpu_manager.complete_load(to_keys([2, 3])) + cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX) # prepare store [6, 7, 8] -> evicts [2, 3, 4] (oldest) prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) @@ -231,10 +231,10 @@ def test_cpu_manager(): ) # complete store [6, 7, 8] - cpu_manager.complete_store(to_keys([6, 7, 8])) + cpu_manager.complete_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) # touch [5, 6, 7] (move to end of LRU order) - cpu_manager.touch(to_keys([5, 6, 7])) + cpu_manager.touch(to_keys([5, 6, 7]), _EMPTY_REQ_CTX) # prepare store [7, 9] -> evicts [8] (oldest following previous touch) prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX) @@ -248,7 +248,7 @@ def test_cpu_manager(): ) # complete store [7, 9] with failure - cpu_manager.complete_store(to_keys([7, 9]), success=False) + cpu_manager.complete_store(to_keys([7, 9]), _EMPTY_REQ_CTX, success=False) # assert [7] is still stored, but [9] is not assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is True @@ -341,7 +341,7 @@ def test_basic(self): assert list(cpu_manager.take_events()) == [] # complete store [1, 2] - cpu_manager.complete_store(to_keys([1, 2])) + cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] @@ -362,14 +362,14 @@ def test_t1_to_t2_promotion(self): # store and complete block 1 cpu_manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX) - cpu_manager.complete_store(to_keys([1])) + cpu_manager.complete_store(to_keys([1]), _EMPTY_REQ_CTX) # block 1 starts in T1 (recent) assert to_keys([1])[0] in arc_policy.t1 assert to_keys([1])[0] not in arc_policy.t2 # touch block 1 (simulate second access) - cpu_manager.touch(to_keys([1])) + cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX) # block 1 should now be in T2 (frequent) assert to_keys([1])[0] not in arc_policy.t1 @@ -394,7 +394,7 @@ def test_eviction_with_load(self): evicted_keys=[], ), ) - cpu_manager.complete_store(to_keys([1, 2, 3, 4])) + cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) # prepare load [2, 3] (increases ref_cnt) prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX) @@ -405,7 +405,7 @@ def test_eviction_with_load(self): assert cpu_manager.prepare_store(to_keys([5, 6, 7]), _EMPTY_REQ_CTX) is None # complete load [2, 3] - cpu_manager.complete_load(to_keys([2, 3])) + cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX) # now prepare store [5, 6, 7] should succeed # ARC will evict blocks one at a time from T1 as needed @@ -426,20 +426,20 @@ def test_adaptive_target(self): # store blocks 1, 2 (fills cache) cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) - cpu_manager.complete_store(to_keys([1, 2])) + cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) initial_target = arc_policy.target_t1_size # store block 3, evicting block 1 (moves to B1 ghost list) cpu_manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX) - cpu_manager.complete_store(to_keys([3])) + cpu_manager.complete_store(to_keys([3]), _EMPTY_REQ_CTX) # block 1 should be in B1 (ghost list) assert to_keys([1])[0] in arc_policy.b1 # touch block 1 (cache miss, but in B1) # this should increase target_t1_size (favor recency) - cpu_manager.touch(to_keys([1])) + cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX) # target should have increased assert arc_policy.target_t1_size > initial_target @@ -453,10 +453,10 @@ def test_t1_t2_eviction_policy(self): # store blocks 1, 2, 3, 4 cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) - cpu_manager.complete_store(to_keys([1, 2, 3, 4])) + cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) # promote blocks 3, 4 to T2 by touching them - cpu_manager.touch(to_keys([3, 4])) + cpu_manager.touch(to_keys([3, 4]), _EMPTY_REQ_CTX) # now: T1 = {1, 2}, T2 = {3, 4} assert len(arc_policy.t1) == 2 @@ -471,7 +471,7 @@ def test_t1_t2_eviction_policy(self): assert output is not None assert to_keys([1]) == output.evicted_keys - cpu_manager.complete_store(to_keys([5])) + cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX) # block 1 should be in B1 (ghost list) assert to_keys([1])[0] in arc_policy.b1 @@ -487,12 +487,12 @@ def test_ghost_list_bounds(self): # fill cache with blocks 1, 2 cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) - cpu_manager.complete_store(to_keys([1, 2])) + cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) # store many blocks to fill ghost lists for i in range(3, 20): cpu_manager.prepare_store(to_keys([i]), _EMPTY_REQ_CTX) - cpu_manager.complete_store(to_keys([i])) + cpu_manager.complete_store(to_keys([i]), _EMPTY_REQ_CTX) # ghost lists should not exceed cache_capacity assert len(arc_policy.b1) <= arc_policy.cache_capacity @@ -507,14 +507,14 @@ def test_touch_ordering(self): # store blocks 1, 2, 3, 4 cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) - cpu_manager.complete_store(to_keys([1, 2, 3, 4])) + cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) # promote 3, 4 to T2 - cpu_manager.touch(to_keys([3, 4])) + cpu_manager.touch(to_keys([3, 4]), _EMPTY_REQ_CTX) # T1 = {1, 2}, T2 = {3, 4} # touch [1, 3, 4] - should promote 1 to T2, and move 3,4 to end of T2 - cpu_manager.touch(to_keys([1, 3, 4])) + cpu_manager.touch(to_keys([1, 3, 4]), _EMPTY_REQ_CTX) # T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent) assert len(arc_policy.t1) == 1 @@ -540,7 +540,7 @@ def test_failed_store(self): # store blocks 1, 2, 3, 4 cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) - cpu_manager.complete_store(to_keys([1, 2, 3, 4])) + cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) # prepare store block 5 (will evict block 1) prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX) @@ -548,7 +548,7 @@ def test_failed_store(self): assert len(prepare_store_output.evicted_keys) == 1 # complete store with failure - cpu_manager.complete_store(to_keys([5]), success=False) + cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX, success=False) # block 5 should not be in cache assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is False @@ -569,7 +569,7 @@ def test_full_scenario(self): # store [1, 2] cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) - cpu_manager.complete_store(to_keys([1, 2])) + cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) # store [3, 4, 5] -> evicts [1] prepare_store_output = cpu_manager.prepare_store( @@ -577,10 +577,10 @@ def test_full_scenario(self): ) assert prepare_store_output is not None assert len(prepare_store_output.evicted_keys) == 1 - cpu_manager.complete_store(to_keys([3, 4, 5])) + cpu_manager.complete_store(to_keys([3, 4, 5]), _EMPTY_REQ_CTX) # promote some blocks to T2 - cpu_manager.touch(to_keys([2, 3])) + cpu_manager.touch(to_keys([2, 3]), _EMPTY_REQ_CTX) # T1 has {4, 5}, T2 has {2, 3} assert len(arc_policy.t1) == 2 @@ -589,7 +589,7 @@ def test_full_scenario(self): # store [6] -> should evict from T1 (4 is oldest in T1) prepare_store_output = cpu_manager.prepare_store(to_keys([6]), _EMPTY_REQ_CTX) assert prepare_store_output is not None - cpu_manager.complete_store(to_keys([6])) + cpu_manager.complete_store(to_keys([6]), _EMPTY_REQ_CTX) # verify blocks 2, 3 (in T2) are still present assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True @@ -646,4 +646,4 @@ def test_filter_reused_manager(): assert prepare_store_output is not None assert prepare_store_output.keys_to_store == [] - manager.complete_store(to_keys([1])) + manager.complete_store(to_keys([1]), _EMPTY_REQ_CTX) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 773fe8f056ac..137eaef9788c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -291,7 +291,7 @@ def _touch(self, req_status: RequestOffloadState): self.config.kv_group_configs, req_status.group_states ): if group_config.sliding_window_size_in_blocks is None: - self.manager.touch(group_state.offload_keys) + self.manager.touch(group_state.offload_keys, req_status.req_context) else: # we aim to keep just blocks that are necessary to hit # the original request (+ decoded blocks) @@ -300,7 +300,10 @@ def _touch(self, req_status: RequestOffloadState): group_state.num_hit_blocks - group_config.sliding_window_size_in_blocks, ) - self.manager.touch(group_state.offload_keys[blocks_to_skip:]) + self.manager.touch( + group_state.offload_keys[blocks_to_skip:], + req_status.req_context, + ) def _lookup(self, req_status: RequestOffloadState) -> int | None: """ @@ -802,14 +805,13 @@ def update_connector_output(self, connector_output: KVConnectorOutput): continue assert job_status.pending_count == 0 + req_status = self._req_status[job_status.req_id] if job_status.is_store: - self.manager.complete_store(job_status.keys) + self.manager.complete_store(job_status.keys, req_status.req_context) else: - self.manager.complete_load(job_status.keys) + self.manager.complete_load(job_status.keys, req_status.req_context) if self._blocks_being_loaded: self._blocks_being_loaded.difference_update(job_status.keys) - - req_status = self._req_status[job_status.req_id] if self._block_id_to_pending_jobs: # Sliding window blocks are tracked from store creation # and must be cleaned up unconditionally. diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index b73a66ba2d9b..afa3a06774c0 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -147,22 +147,24 @@ def prepare_load( """ pass - def touch(self, keys: Collection[OffloadKey]): + def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): """ Mark the given blocks as recently used. This could in practice mean moving them to the end of an LRU list. Args: keys: the keys identifying the blocks. + req_context: per-request context (e.g. kv_transfer_params). """ return - def complete_load(self, keys: Collection[OffloadKey]): + def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext): """ Marks previous blocks that were prepared to load as done loading. Args: keys: the keys identifying the blocks. + req_context: per-request context (e.g. kv_transfer_params). """ return @@ -189,7 +191,12 @@ def prepare_store( """ pass - def complete_store(self, keys: Collection[OffloadKey], success: bool = True): + def complete_store( + self, + keys: Collection[OffloadKey], + req_context: ReqContext, + success: bool = True, + ): """ Marks blocks which were previously prepared to be stored, as stored. Following this call, the blocks become loadable. @@ -198,6 +205,7 @@ def complete_store(self, keys: Collection[OffloadKey], success: bool = True): Args: keys: the keys identifying the blocks. + req_context: per-request context (e.g. kv_transfer_params). success: whether the blocks were stored successfully. """ return diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 80bcb568f99a..dc501d745683 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -106,10 +106,12 @@ def prepare_load( blocks.append(block) return self._get_load_store_spec(keys, blocks) - def touch(self, keys: Collection[OffloadKey]) -> None: + def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None: self._policy.touch(keys) - def complete_load(self, keys: Collection[OffloadKey]) -> None: + def complete_load( + self, keys: Collection[OffloadKey], req_context: ReqContext + ) -> None: for key in keys: block = self._policy.get(key) assert block is not None, f"Block {key!r} not found" @@ -172,7 +174,10 @@ def prepare_store( ) def complete_store( - self, keys: Collection[OffloadKey], success: bool = True + self, + keys: Collection[OffloadKey], + req_context: ReqContext, + success: bool = True, ) -> None: stored_keys: list[OffloadKey] = [] diff --git a/vllm/v1/kv_offload/reuse_manager.py b/vllm/v1/kv_offload/reuse_manager.py index 6cb0a5f7591c..d8c20b9dec7b 100644 --- a/vllm/v1/kv_offload/reuse_manager.py +++ b/vllm/v1/kv_offload/reuse_manager.py @@ -105,16 +105,21 @@ def prepare_load( ) -> LoadStoreSpec: return self._backing.prepare_load(keys, req_context) - def touch(self, keys: Collection[OffloadKey]) -> None: - return self._backing.touch(keys) + def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None: + return self._backing.touch(keys, req_context) - def complete_load(self, keys: Collection[OffloadKey]) -> None: - return self._backing.complete_load(keys) + def complete_load( + self, keys: Collection[OffloadKey], req_context: ReqContext + ) -> None: + return self._backing.complete_load(keys, req_context) def complete_store( - self, keys: Collection[OffloadKey], success: bool = True + self, + keys: Collection[OffloadKey], + req_context: ReqContext, + success: bool = True, ) -> None: - return self._backing.complete_store(keys, success) + return self._backing.complete_store(keys, req_context, success) def take_events(self) -> Iterable[OffloadingEvent]: return self._backing.take_events() From 80b671dd77ffd4d8d6e3f53028c133cde74fc1a8 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 7 May 2026 11:39:14 +0300 Subject: [PATCH 30/55] Add context to tiering offloading manager methods Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 45 +++++++++++++------ vllm/v1/kv_offload/tiering/base.py | 3 +- .../v1/kv_offload/tiering/example/__init__.py | 3 +- vllm/v1/kv_offload/tiering/manager.py | 38 +++++++++------- 4 files changed, 58 insertions(+), 31 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index d8523b466eb5..c7dffed8d3ed 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -101,7 +101,7 @@ def test_lru_eviction(self): assert tier.get_num_blocks() == 3 # Touch first block (make it most recently used) - tier.touch([blocks[0]]) + tier.touch([blocks[0]], _CTX) # Store new block should evict blocks[1] (least recently used) new_block = to_keys([3])[0] @@ -191,7 +191,7 @@ def test_basic_store_to_primary(self, manager_setup): assert len(result.keys_to_store) == 3 # Complete store - self.manager.complete_store(blocks, success=True) + self.manager.complete_store(blocks, _CTX, success=True) # Blocks should be in primary tier assert count_hits(self.primary_tier, blocks) == 3 @@ -205,7 +205,7 @@ def test_cascade_to_all_secondary_tiers(self, manager_setup): assert result is not None # Complete store (triggers cascade) - self.manager.complete_store(blocks, success=True) + self.manager.complete_store(blocks, _CTX, success=True) # Process finished jobs to complete cascade self.manager._process_finished_jobs() @@ -225,7 +225,7 @@ def test_ref_cnt_protection_during_cascade(self, manager_setup): # Store to primary result = self.manager.prepare_store(blocks, _CTX) assert result is not None - self.manager.complete_store(blocks, success=True) + self.manager.complete_store(blocks, _CTX, success=True) # After complete_store, blocks should have ref_cnt > 0 # (one for each secondary tier) @@ -248,7 +248,7 @@ def test_lookup_from_primary(self, manager_setup): # Store blocks self.manager.prepare_store(blocks, _CTX) - self.manager.complete_store(blocks, success=True) + self.manager.complete_store(blocks, _CTX, success=True) # Lookup should find all blocks in primary assert count_hits(self.manager, blocks) == 3 @@ -284,7 +284,7 @@ def test_partial_lookup(self, manager_setup): # Store first 3 blocks to primary self.manager.prepare_store(blocks[:3], _CTX) - self.manager.complete_store(blocks[:3], success=True) + self.manager.complete_store(blocks[:3], _CTX, success=True) # Lookup all 5 blocks should return 3 (first 3 found) assert count_hits(self.manager, blocks) == 3 @@ -297,7 +297,7 @@ def test_eviction_in_primary_tier(self, manager_setup): result = self.manager.prepare_store(blocks, _CTX) assert result is not None assert len(result.keys_to_store) == 5 - self.manager.complete_store(blocks, success=True) + self.manager.complete_store(blocks, _CTX, success=True) # Process finished jobs to release ref_cnt from cascade self.manager._process_finished_jobs() @@ -317,11 +317,11 @@ def test_touch_propagates_to_all_tiers(self, manager_setup): # Store blocks self.manager.prepare_store(blocks, _CTX) - self.manager.complete_store(blocks, success=True) + self.manager.complete_store(blocks, _CTX, success=True) self.manager._process_finished_jobs() # Touch blocks - self.manager.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 @@ -344,7 +344,7 @@ def test_failed_store_no_cascade(self, manager_setup): assert result is not None # Complete store with failure - self.manager.complete_store(blocks, success=False) + self.manager.complete_store(blocks, _CTX, success=False) # Process finished jobs self.manager._process_finished_jobs() @@ -374,7 +374,7 @@ def test_multiple_secondary_tiers_independent_eviction(self): blocks1 = to_keys(range(5)) result = manager.prepare_store(blocks1, _CTX) assert result is not None - manager.complete_store(blocks1, success=True) + manager.complete_store(blocks1, _CTX, success=True) manager._process_finished_jobs() # Both tiers should have 5 blocks @@ -385,7 +385,7 @@ def test_multiple_secondary_tiers_independent_eviction(self): blocks2 = to_keys(range(5, 8)) result = manager.prepare_store(blocks2, _CTX) assert result is not None - manager.complete_store(blocks2, success=True) + manager.complete_store(blocks2, _CTX, success=True) manager._process_finished_jobs() # Small tier should still have 5 blocks (evicted 3, added 3) @@ -400,7 +400,7 @@ def test_prepare_store_processes_finished_jobs_first(self, manager_setup): # Store blocks self.manager.prepare_store(blocks, _CTX) - self.manager.complete_store(blocks, success=True) + self.manager.complete_store(blocks, _CTX, success=True) # Blocks should have ref_cnt = 2 (one for each secondary tier) for block_hash in blocks: @@ -486,6 +486,23 @@ def test_lookup_shared_block_no_duplicate_promotion(self, manager_setup): assert list(job_metadata.keys) == [shared_block] assert job_metadata.req_context is ctx_a + def test_complete_store_forwards_req_context_to_submit_store(self, manager_setup): + """complete_store cascades to secondary tiers with the correct req_context.""" + blocks = to_keys(range(2)) + + self.secondary_tier1.submit_store = MagicMock( + wraps=self.secondary_tier1.submit_store + ) + + ctx = ReqContext(kv_transfer_params={"key": "value"}) + + self.manager.prepare_store(blocks, ctx) + self.manager.complete_store(blocks, ctx, success=True) + + assert self.secondary_tier1.submit_store.call_count == 1 + job_metadata = self.secondary_tier1.submit_store.call_args.args[0] + assert job_metadata.req_context is ctx + class TestTieringOffloadingWithoutSecondaryTiers: """Test TieringOffloadingManager with no secondary tiers (backward compat).""" @@ -507,7 +524,7 @@ def test_works_without_secondary_tiers(self): # Should work like a regular OffloadingManager result = manager.prepare_store(blocks, _CTX) assert result is not None - manager.complete_store(blocks, success=True) + manager.complete_store(blocks, _CTX, success=True) assert count_hits(manager, blocks) == 3 diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 500fca0e701b..04f4cdf10966 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -146,12 +146,13 @@ def set_primary_view(self, view: memoryview) -> None: """ return - def touch(self, keys: Collection[OffloadKey]): + def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): """ Mark blocks as recently used for eviction policy. Args: keys: Offload keys to mark as recently used. + req_context: Per-request context. """ return diff --git a/vllm/v1/kv_offload/tiering/example/__init__.py b/vllm/v1/kv_offload/tiering/example/__init__.py index 5903672bbca6..42a030546de5 100644 --- a/vllm/v1/kv_offload/tiering/example/__init__.py +++ b/vllm/v1/kv_offload/tiering/example/__init__.py @@ -222,12 +222,13 @@ def _complete_load_job(self, job_metadata: _JobMetadata): # 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]): + 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: diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 7aedc469324f..6d3f87697305 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -83,19 +83,21 @@ def prepare_write(self, keys, req_context: ReqContext) -> PrepareStoreOutput | N """Allocate space in primary for a secondary->primary write (promotion).""" return self.prepare_store(keys, req_context) - def complete_write(self, keys, success: bool = True) -> None: + def complete_write( + self, keys, req_context: ReqContext, success: bool = True + ) -> None: """Finalize secondary->primary write, making blocks available.""" - self.complete_store(keys, success) + self.complete_store(keys, req_context, success) def prepare_read(self, keys, req_context: ReqContext) -> LoadStoreSpec: """Protect primary blocks for a primary->secondary read (cascade), incrementing ref_cnt.""" return self.prepare_load(keys, req_context) - def complete_read(self, keys) -> None: + def complete_read(self, keys, req_context: ReqContext) -> None: """Release protection after primary->secondary read completes, decrementing ref_cnt.""" - self.complete_load(keys) + self.complete_load(keys, req_context) def create_kv_memoryview(self) -> memoryview: """Create a memoryview over the primary tier's KV cache buffer. @@ -196,13 +198,17 @@ def _process_finished_jobs(self): # primary→secondary transfer completed. # Decrement ref_cnt on primary blocks. job_metadata = self._store_jobs.pop(job_id) - self.primary_tier.complete_read(job_metadata.keys) + self.primary_tier.complete_read( + job_metadata.keys, job_metadata.req_context + ) elif job_id in self._load_jobs: # secondary→primary transfer (promotion) completed. # Make blocks available in primary tier. job_metadata = self._load_jobs.pop(job_id) self.primary_tier.complete_write( - job_metadata.keys, completed_job.success + job_metadata.keys, + job_metadata.req_context, + completed_job.success, ) else: # Job ID not found in either dictionary - this shouldn't happen @@ -351,18 +357,19 @@ def prepare_load( return self.primary_tier.prepare_load(keys, req_context) - def touch(self, keys: Collection[OffloadKey]): + def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): """ Mark blocks as recently used in all tiers. Args: keys: Blocks to mark as recently used. + req_context: Per-request context. """ - self.primary_tier.touch(keys) + self.primary_tier.touch(keys, req_context) for tier in self.secondary_tiers: - tier.touch(keys) + tier.touch(keys, req_context) - def complete_load(self, keys: Collection[OffloadKey]): + def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext): """ Mark blocks as done loading from primary tier to GPU. @@ -371,8 +378,9 @@ def complete_load(self, keys: Collection[OffloadKey]): Args: keys: Blocks that finished loading. + req_context: Per-request context. """ - self.primary_tier.complete_load(keys) + self.primary_tier.complete_load(keys, req_context) def prepare_store( self, keys: Collection[OffloadKey], req_context: ReqContext @@ -408,6 +416,7 @@ def prepare_store( def complete_store( self, keys: Collection[OffloadKey], + req_context: ReqContext, success: bool = True, ): """ @@ -429,7 +438,7 @@ def complete_store( req_context: Per-request context forwarded to primary.prepare_read(). """ # Step 1: Complete store in primary tier (makes blocks loadable) - self.primary_tier.complete_store(keys, success) + self.primary_tier.complete_store(keys, req_context, success) if not success: # If GPU→Primary transfer failed, don't cascade to secondary tiers @@ -441,9 +450,7 @@ def complete_store( # eviction during the async transfer). One prepare_read() call per # secondary tier. for tier in self.secondary_tiers: - # Get spec for reading from primary tier AND increment ref_cnt - # TODO: pass the actual req_context instead of None - primary_blocks_spec = self.primary_tier.prepare_read(keys, ReqContext()) + primary_blocks_spec = self.primary_tier.prepare_read(keys, req_context) # Submit async store job: primary→secondary job_id = self._next_job_id() @@ -454,6 +461,7 @@ def complete_store( job_id=job_id, keys=keys, block_ids=primary_blocks_spec.block_ids, + req_context=req_context, ) self._store_jobs[job_id] = job_metadata From 727d9cf3a91294b6e787a3670ddc795a5901e741 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 7 May 2026 16:19:52 +0300 Subject: [PATCH 31/55] Invoke _process_finished_jobs() once per engine step Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 84 +++++++++++++------ vllm/v1/kv_offload/tiering/manager.py | 51 ++++++++--- 2 files changed, 98 insertions(+), 37 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index c7dffed8d3ed..f8988cc6b1e9 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -200,15 +200,23 @@ def test_cascade_to_all_secondary_tiers(self, manager_setup): """Test that blocks are cascaded to ALL secondary tiers.""" blocks = to_keys(range(3)) + self.secondary_tier1.submit_store = MagicMock( + wraps=self.secondary_tier1.submit_store + ) + self.secondary_tier2.submit_store = MagicMock( + wraps=self.secondary_tier2.submit_store + ) + # Store to primary result = self.manager.prepare_store(blocks, _CTX) assert result is not None - # Complete store (triggers cascade) + # Complete store (triggers cascade via submit_store on each tier) self.manager.complete_store(blocks, _CTX, success=True) - # Process finished jobs to complete cascade - self.manager._process_finished_jobs() + # submit_store was called once per secondary tier + self.secondary_tier1.submit_store.assert_called_once() + self.secondary_tier2.submit_store.assert_called_once() # Blocks should be in both secondary tiers assert self.secondary_tier1.get_num_blocks() == 3 @@ -234,14 +242,35 @@ def test_ref_cnt_protection_during_cascade(self, manager_setup): # ref_cnt should be 2 (one for each secondary tier) assert block.ref_cnt == 2 - # Process finished jobs to complete cascade - self.manager._process_finished_jobs() + # End of step 1: _maybe_process_finished_jobs() was already called by + # prepare_store() above (setting the per-step flag), so take_events() + # does NOT poll get_finished() again — cascade completions remain + # unprocessed until the next step. + list(self.manager.take_events()) + + # ref_cnt still held: cascade jobs finished (sync tier) but haven't + # been polled yet because the per-step guard skipped the second call. + for block_hash in blocks: + block = self.primary_tier._policy.get(block_hash) + assert block.ref_cnt == 2 + + # Secondary tiers have completed jobs waiting to be drained + assert len(self.secondary_tier1.completed_jobs) > 0 + assert len(self.secondary_tier2.completed_jobs) > 0 + + # End of step 2: flag was reset, so _maybe_process_finished_jobs() + # runs and processes the cascade completions (complete_read → ref_cnt--) + list(self.manager.take_events()) # After cascade completes, ref_cnt should be 0 for block_hash in blocks: block = self.primary_tier._policy.get(block_hash) assert block.ref_cnt == 0 + # All completed jobs have been drained + assert len(self.secondary_tier1.completed_jobs) == 0 + assert len(self.secondary_tier2.completed_jobs) == 0 + def test_lookup_from_primary(self, manager_setup): """Test lookup when blocks are in primary tier.""" blocks = to_keys(range(3)) @@ -266,11 +295,11 @@ def test_promotion_from_secondary(self, manager_setup): result = self.manager.lookup(block, _CTX) assert result is None # Retry later (promotion initiated) - # Flush deferred submit_load() calls (normally done at end of engine step) + # End of step 1: flushes deferred submit_load() calls list(self.manager.take_events()) - # Process finished jobs to complete promotion - self.manager._process_finished_jobs() + # End of step 2: processes the completed promotion jobs + list(self.manager.take_events()) # Now blocks should be in primary tier assert count_hits(self.primary_tier, blocks) == 3 @@ -299,8 +328,8 @@ def test_eviction_in_primary_tier(self, manager_setup): assert len(result.keys_to_store) == 5 self.manager.complete_store(blocks, _CTX, success=True) - # Process finished jobs to release ref_cnt from cascade - self.manager._process_finished_jobs() + # End of step: release ref_cnt from cascade + list(self.manager.take_events()) # Now try to store 2 more blocks (should trigger eviction) more_blocks = to_keys(range(5, 7)) @@ -318,7 +347,7 @@ def test_touch_propagates_to_all_tiers(self, manager_setup): # Store blocks self.manager.prepare_store(blocks, _CTX) self.manager.complete_store(blocks, _CTX, success=True) - self.manager._process_finished_jobs() + list(self.manager.take_events()) # Touch blocks self.manager.touch(blocks, _CTX) @@ -339,19 +368,23 @@ def test_failed_store_no_cascade(self, manager_setup): """Test that failed GPU→primary store doesn't cascade.""" blocks = to_keys(range(3)) + self.secondary_tier1.submit_store = MagicMock( + wraps=self.secondary_tier1.submit_store + ) + self.secondary_tier2.submit_store = MagicMock( + wraps=self.secondary_tier2.submit_store + ) + # Prepare store result = self.manager.prepare_store(blocks, _CTX) assert result is not None - # Complete store with failure + # Complete store with failure — cascade must not happen self.manager.complete_store(blocks, _CTX, success=False) - # Process finished jobs - self.manager._process_finished_jobs() - - # Blocks should NOT be in secondary tiers - assert self.secondary_tier1.get_num_blocks() == 0 - assert self.secondary_tier2.get_num_blocks() == 0 + # submit_store was never called on either secondary tier + 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.""" @@ -375,7 +408,7 @@ def test_multiple_secondary_tiers_independent_eviction(self): result = manager.prepare_store(blocks1, _CTX) assert result is not None manager.complete_store(blocks1, _CTX, success=True) - manager._process_finished_jobs() + list(manager.take_events()) # Both tiers should have 5 blocks assert small_tier.get_num_blocks() == 5 @@ -386,7 +419,7 @@ def test_multiple_secondary_tiers_independent_eviction(self): result = manager.prepare_store(blocks2, _CTX) assert result is not None manager.complete_store(blocks2, _CTX, success=True) - manager._process_finished_jobs() + list(manager.take_events()) # Small tier should still have 5 blocks (evicted 3, added 3) assert small_tier.get_num_blocks() == 5 @@ -394,11 +427,11 @@ def test_multiple_secondary_tiers_independent_eviction(self): # Large tier should have all 8 blocks assert large_tier.get_num_blocks() == 8 - def test_prepare_store_processes_finished_jobs_first(self, manager_setup): - """Test that prepare_store() calls _process_finished_jobs() first.""" + def test_prepare_store_processes_finished_jobs_on_new_step(self, manager_setup): + """Test that prepare_store() processes finished jobs on first call of a step.""" blocks = to_keys(range(3)) - # Store blocks + # Step 1: Store blocks and cascade self.manager.prepare_store(blocks, _CTX) self.manager.complete_store(blocks, _CTX, success=True) @@ -407,7 +440,10 @@ def test_prepare_store_processes_finished_jobs_first(self, manager_setup): block = self.primary_tier._policy.get(block_hash) assert block.ref_cnt == 2 - # Call prepare_store again (should process finished jobs first) + # End of step 1 (resets per-step flag; cascade completions pending) + list(self.manager.take_events()) + + # Step 2: prepare_store processes finished cascade jobs on first call more_blocks = to_keys(range(3, 5)) self.manager.prepare_store(more_blocks, _CTX) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 6d3f87697305..83802b58f989 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -166,6 +166,10 @@ def __init__( SecondaryTierManager, dict[int, _PendingPromotion] ] = {} + # Gate for once-per-step execution of _maybe_process_finished_jobs(). + # Reset at the end of each step in take_events(). + self._processed_jobs_this_step: bool = False + # Wire each secondary tier with a long-lived memoryview of the primary # CPU buffer. One view is shared across all tiers; released in shutdown(). self._primary_kv_view = primary_tier.create_kv_memoryview() @@ -178,9 +182,22 @@ def _next_job_id(self) -> JobId: self._job_id_counter += 1 return job_id + def _maybe_process_finished_jobs(self): + """ + Poll secondary tiers for completed jobs (at most once per step). + + Guarded by _processed_jobs_this_step: the first call in an engine step + does the actual polling; subsequent calls are no-ops. The flag is reset + in take_events() at the end of each step. + """ + if self._processed_jobs_this_step: + return + self._processed_jobs_this_step = True + self._process_finished_jobs() + def _process_finished_jobs(self): """ - Poll all secondary tiers for completed jobs and update state accordingly. + Unconditionally poll all secondary tiers for completed jobs. This method: 1. Calls get_finished() on each secondary tier @@ -240,7 +257,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: None if no tier has the block but at least one tier is busy (retry later). """ - self._process_finished_jobs() + self._maybe_process_finished_jobs() # Always query every tier to warm up caches / prefetch state primary_hit = self.primary_tier.lookup(key, req_context) @@ -339,7 +356,7 @@ def prepare_load( """ Prepare blocks to be loaded from primary tier to GPU. - CRITICAL: This method calls _process_finished_jobs() FIRST to ensure + CRITICAL: This method calls _maybe_process_finished_jobs() FIRST to ensure that any completed promotions have been finalized and blocks are ready. This increments ref_cnt on the blocks in the primary tier, protecting @@ -353,7 +370,7 @@ def prepare_load( LoadStoreSpec for reading from primary tier. """ # Process completed promotions to ensure blocks are ready - self._process_finished_jobs() + self._maybe_process_finished_jobs() return self.primary_tier.prepare_load(keys, req_context) @@ -388,7 +405,7 @@ def prepare_store( """ Prepare blocks to be stored from GPU to primary tier. - CRITICAL: This method calls _process_finished_jobs() FIRST to ensure + CRITICAL: This method calls _maybe_process_finished_jobs() FIRST to ensure that any completed async transfers have their ref_cnt decremented before the primary tier makes eviction decisions. @@ -403,7 +420,7 @@ def prepare_store( # Step 1: Poll for completed async jobs FIRST # This decrements ref_cnt on primary blocks that have been # successfully transferred to secondary tiers. - self._process_finished_jobs() + self._maybe_process_finished_jobs() # Step 2: Store to primary tier primary_result = self.primary_tier.prepare_store(keys, req_context) @@ -467,15 +484,17 @@ def complete_store( tier.submit_store(job_metadata) - # Note: The async transfers are now in flight. - # Their completion is tracked via get_finished() / _process_finished_jobs(). + # Note: The async transfers are now in flight. Their completion is + # tracked via get_finished() / _maybe_process_finished_jobs(). def take_events(self) -> Iterable[OffloadingEvent]: """ - Take offloading events from the primary tier. + End-of-step hook: flush deferred work, yield events, reset per-step state. - Note: Currently only primary tier events are tracked. Secondary tier - events could be added in the future if needed. + Called once per engine step from Scheduler.update_from_output() → + connector.take_events(). Ensures _maybe_process_finished_jobs() has run + at least once this step, flushes pending promotions, yields collected + events, and resets the per-step flag. Yields: New OffloadingEvents collected since the last call. @@ -484,14 +503,20 @@ def take_events(self) -> Iterable[OffloadingEvent]: # hook once one exists. For now, take_events() serves as the flush # point under the assumption that it is called at the end of each # engine step (Scheduler.update_from_output() → connector.take_events()). - # Update the relevant tests the rely on take_events() to signal end of step. + # When the dedicated hook is added, update tests that rely on + # take_events() to signal end of step. + + self._maybe_process_finished_jobs() + self._flush_pending_promotions() + # Reset the per-step gate so next step's first call does real work. + self._processed_jobs_this_step = False + if self.events is not None: yield from self.events self.events.clear() - # Also yield events from primary tier yield from self.primary_tier.take_events() def shutdown(self) -> None: From 1c152067d167a85a05595f98f7f380729e588a7a Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 7 May 2026 16:24:25 +0300 Subject: [PATCH 32/55] Remove redundant tests Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 25 ------------------- 1 file changed, 25 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index f8988cc6b1e9..e4ad27eb8381 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -427,31 +427,6 @@ def test_multiple_secondary_tiers_independent_eviction(self): # Large tier should have all 8 blocks assert large_tier.get_num_blocks() == 8 - def test_prepare_store_processes_finished_jobs_on_new_step(self, manager_setup): - """Test that prepare_store() processes finished jobs on first call of a step.""" - blocks = to_keys(range(3)) - - # Step 1: Store blocks and cascade - self.manager.prepare_store(blocks, _CTX) - self.manager.complete_store(blocks, _CTX, success=True) - - # Blocks should have ref_cnt = 2 (one for each secondary tier) - for block_hash in blocks: - block = self.primary_tier._policy.get(block_hash) - assert block.ref_cnt == 2 - - # End of step 1 (resets per-step flag; cascade completions pending) - list(self.manager.take_events()) - - # Step 2: prepare_store processes finished cascade jobs on first call - more_blocks = to_keys(range(3, 5)) - self.manager.prepare_store(more_blocks, _CTX) - - # Original blocks should now have ref_cnt = 0 - for block_hash in blocks: - block = self.primary_tier._policy.get(block_hash) - assert block.ref_cnt == 0 - def test_lookup_batches_submit_load_per_request(self, manager_setup): """lookup() defers submit_load until take_events(), one call per request. From e1cf201010f4f30b9d188ddaa9e728369b135b98 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 11 May 2026 10:47:52 +0300 Subject: [PATCH 33/55] Address review: update docstrings of SecondaryTierManager methods Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/base.py | 68 +++++++++++++++--------------- 1 file changed, 33 insertions(+), 35 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 04f4cdf10966..3a3acaf13069 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -70,27 +70,25 @@ def submit_store(self, job_metadata: JobMetadata) -> None: Submit an async job to store blocks from the primary tier to this secondary tier. - This method is lightweight: it allocates metadata and submits the - transfer job, but does NOT perform the actual data transfer on the + This method must be lightweight and non-blocking: allocate metadata + and submit the transfer, but do NOT perform the data copy on the calling thread. - The caller (TieringOffloadingManager) must have already called - primary.prepare_read(keys) to obtain job_metadata.block_ids and - to increment ref_cnt on those blocks. ref_cnt will be decremented - when get_finished() reports this job_id as complete and - primary.unprepare_read() is called. + Preconditions (guaranteed by the framework): + - ``job_metadata.block_ids`` are valid primary-tier slots, pinned + (ref-counted) for the duration of the transfer. - This method is responsible for: - 1. Filtering out blocks already present in this secondary tier - 2. Evicting blocks from this secondary tier if needed (secondary - tiers are responsible for their own evictions) - 3. Allocating space in this secondary tier - 4. Submitting the async transfer: primary → secondary + The implementation is responsible for: + 1. Filtering out blocks already present in this tier + 2. Evicting blocks if capacity is needed + 3. Allocating space in this tier + 4. Submitting the async transfer (read from primary via block_ids) + + Report completion via ``get_finished()``. Args: - job_metadata: Job metadata including job_id, keys, and - block_ids for reading blocks from the primary tier - (obtained via primary.prepare_read()). + job_metadata: Job metadata including job_id, keys, and block_ids + identifying the primary-tier slots to read from. """ pass @@ -100,36 +98,36 @@ def submit_load(self, job_metadata: JobMetadata) -> None: Submit an async job to load blocks from this secondary tier to the primary tier. - This method is lightweight: it marks blocks as in-flight and submits - the transfer job, but does NOT perform the actual data transfer on - the calling thread. + This method must be lightweight and non-blocking: mark blocks as + in-flight and submit the transfer, but do NOT perform the data copy + on the calling thread. + + Preconditions (guaranteed by the framework): + - ``job_metadata.block_ids`` are allocated primary-tier slots + ready to receive data. + + The implementation must copy data from this tier into the + primary-tier slots identified by ``block_ids``. - The caller (TieringOffloadingManager) must have already called - primary.prepare_write(keys) to obtain job_metadata.block_ids and - to allocate space in the primary tier. When get_finished() reports - this job_id as complete, primary.complete_write() is called to make - the blocks available for GPU loads. + Report completion via ``get_finished()``. Args: - job_metadata: Job metadata including job_id, keys, and - block_ids for writing blocks into the primary tier - (obtained via primary.prepare_write()). + job_metadata: Job metadata including job_id, keys, and block_ids + identifying the primary-tier slots to write into. """ pass @abstractmethod def get_finished(self) -> Iterable[JobResult]: """ - Poll for finished async jobs (both loads and stores). + Return all jobs (loads and stores) that completed since the last call. - This is the mechanism by which the TieringOffloadingManager learns - that a transfer has finished and can: - - Call primary.unprepare_read() to decrement ref_cnt (for stores) - - Call primary.complete_write() to make blocks loadable (for loads) + The framework uses these results to release resources and finalize + transfers. Returns: - Iterable of JobResult objects for all jobs that have - finished since the last call. + Iterable of JobResult objects for jobs finished since the + last call. """ pass @@ -137,7 +135,7 @@ def set_primary_view(self, view: memoryview) -> None: """ Provide a long-lived memoryview of the primary-tier CPU tensor. - Called once by TieringOffloadingManager during initialisation. + Called once during initialisation. Override to store the view for use in `submit_store` and `submit_load`. Use `view.strides[0]` to obtain the byte stride between block slots. From 9a64da17a6aab7347b39dc0e5032a6f2c19d3259 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 11 May 2026 11:00:39 +0300 Subject: [PATCH 34/55] Address review: remove hedging language Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/base.py | 6 +++--- vllm/v1/kv_offload/tiering/manager.py | 14 +++++++------- vllm/v1/kv_offload/tiering/spec.py | 12 ++++++------ 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 3a3acaf13069..1424d26d822b 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -39,9 +39,9 @@ class SecondaryTierManager(ABC): Abstract interface for managing a single non-primary offloading tier. Secondary tiers cannot directly access GPU memory. All data transfers - must go through the primary tier (implemented as CPU in current version): - - Store: GPU → primary → secondary (cascade) - - Load: secondary → primary → GPU (promotion) + must go through the CPU (primary) tier: + - Store: GPU → CPU (primary) → secondary (cascade) + - Load: secondary → CPU (primary) → GPU (promotion) IMPORTANT: All methods run in the Scheduler process and must be lightweight and non-blocking. submit_load() and submit_store() submit diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 83802b58f989..18261435b894 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -3,15 +3,15 @@ """ TieringOffloadingManager: Multi-tier KV cache offloading orchestrator. -This manager coordinates between a primary tier (with GPU access, currently -CPU-based) and zero or more secondary tiers (Storage, Network, etc.) to -provide hierarchical KV cache offloading. +This manager coordinates between a CPU primary tier (with direct GPU access) +and zero or more secondary tiers (Storage, Network, etc.) to provide +hierarchical KV cache offloading. Key Design Principles: 1. Always offload to all tiers — When a block is stored to the primary tier, it is cascaded to ALL secondary tiers 2. Primary tier is the gateway — Only the primary tier can directly access - GPU memory (currently implemented using CPU memory) + GPU memory 3. Staged promotion — Blocks in secondary tiers must be promoted to the primary tier before GPU can access them 4. Transparent retry mechanism — Return None from lookup() to signal @@ -119,9 +119,9 @@ class TieringOffloadingManager(OffloadingManager): """ Orchestrates multi-tier KV cache offloading. - This manager coordinates between a primary tier (with GPU access, currently - CPU-based) and zero or more secondary tiers (Storage, Network, etc.) to - provide hierarchical KV cache offloading. + This manager coordinates between a CPU primary tier (with direct GPU access) + and zero or more secondary tiers (Storage, Network, etc.) to provide + hierarchical KV cache offloading. Key internal state: - Minimal state tracking; relies on secondary tiers to report completion diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index a31a22fd1575..7fc41407e87a 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -3,7 +3,7 @@ """ TieringOffloadingSpec: Spec for multi-tier KV cache offloading. -This spec creates a TieringOffloadingManager with a CPU-based primary tier +This spec creates a TieringOffloadingManager with a CPU primary tier and configurable secondary tiers (e.g., Storage, Network). Configuration via kv_connector_extra_config: @@ -56,12 +56,12 @@ class TieringOffloadingSpec(CPUOffloadingSpec): Spec for multi-tier KV cache offloading. Creates a TieringOffloadingManager with: - - Primary tier: CPU-based (LRU or ARC eviction policy) + - Primary tier: CPU (LRU or ARC eviction policy) - Secondary tiers: Configurable via extra_config - The primary tier has direct GPU access and serves as the gateway for all - GPU↔offload operations. Secondary tiers cannot directly access GPU memory - and must coordinate with the primary tier for data transfers. + The CPU primary tier has direct GPU access and serves as the gateway for + all GPU↔offload operations. Secondary tiers cannot directly access GPU + memory and must transfer data through the primary tier. """ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): @@ -82,7 +82,7 @@ def get_manager(self) -> OffloadingManager: Get the TieringOffloadingManager. Creates a TieringOffloadingManager with: - - Primary tier: CPU-based (LRU or ARC) + - Primary tier: CPU (LRU or ARC) - Secondary tiers: As configured in extra_config Returns: From d72847e235f7bed050e0d39043c42c1e56b0d284 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 11 May 2026 14:09:33 +0300 Subject: [PATCH 35/55] Address feedback: clarify a comment Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 18261435b894..cace99702bc6 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -10,8 +10,8 @@ Key Design Principles: 1. Always offload to all tiers — When a block is stored to the primary tier, it is cascaded to ALL secondary tiers -2. Primary tier is the gateway — Only the primary tier can directly access - GPU memory +2. Primary tier is the gateway — Secondary tiers cannot access GPU memory + directly; all data flows through the CPU primary tier 3. Staged promotion — Blocks in secondary tiers must be promoted to the primary tier before GPU can access them 4. Transparent retry mechanism — Return None from lookup() to signal From 489a960cf884952c07af9acb689e79a84688eb2a Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 11 May 2026 14:51:26 +0300 Subject: [PATCH 36/55] Address review: save the extra function call by alias Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 27 +++++++-------------------- 1 file changed, 7 insertions(+), 20 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index cace99702bc6..f27938b14e89 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -78,26 +78,13 @@ def __init__( enable_events=enable_events, ) self._mmap_region = mmap_region - - def prepare_write(self, keys, req_context: ReqContext) -> PrepareStoreOutput | None: - """Allocate space in primary for a secondary->primary write (promotion).""" - return self.prepare_store(keys, req_context) - - def complete_write( - self, keys, req_context: ReqContext, success: bool = True - ) -> None: - """Finalize secondary->primary write, making blocks available.""" - self.complete_store(keys, req_context, success) - - def prepare_read(self, keys, req_context: ReqContext) -> LoadStoreSpec: - """Protect primary blocks for a primary->secondary read (cascade), - incrementing ref_cnt.""" - return self.prepare_load(keys, req_context) - - def complete_read(self, keys, req_context: ReqContext) -> None: - """Release protection after primary->secondary read completes, - decrementing ref_cnt.""" - self.complete_load(keys, req_context) + # read/write is for CPU<->secondary transfers, + # load/store is for CPU<->GPU transfers. + # These aliases avoid calling prepare_load inside a store path. + self.prepare_read = self.prepare_load + self.complete_read = self.complete_load + self.prepare_write = self.prepare_store + self.complete_write = self.complete_store def create_kv_memoryview(self) -> memoryview: """Create a memoryview over the primary tier's KV cache buffer. From ab564077cee75b11de5314f381986214b0f69dc5 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 11 May 2026 15:27:09 +0300 Subject: [PATCH 37/55] Address review: assert zero-copy access to primary kv tensor Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index f27938b14e89..883834bbfcd7 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -99,7 +99,12 @@ def create_kv_memoryview(self) -> memoryview: kv_tensor = self._mmap_region._base.view( self._mmap_region.num_blocks, self._mmap_region._row_stride ) - return memoryview(kv_tensor.numpy()) + np_arr = kv_tensor.numpy() + assert np_arr.ctypes.data == self._mmap_region._base.data_ptr(), ( + "view()/numpy() created a copy instead of sharing the mmap buffer; " + "secondary tiers require zero-copy access to primary KV data" + ) + return memoryview(np_arr) class TieringOffloadingManager(OffloadingManager): From be11bf9ccabfb86923b41e9a595ffdeb54e2356f Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 11 May 2026 15:31:17 +0300 Subject: [PATCH 38/55] Fix merge issue Signed-off-by: Ronen Schaffer --- tests/v1/kv_offload/cpu/test_manager.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index e1c0b7e047c5..c723f09591fc 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -275,7 +275,7 @@ def test_prepare_load_preserves_key_order(): k: int(bid) for k, bid in zip(store_output.keys_to_store, store_output.store_spec.block_ids) } - manager.complete_store([key_a, key_b, key_c]) + manager.complete_store([key_a, key_b, key_c], _EMPTY_REQ_CTX) # Forward order: [a, b, c] spec_fwd = manager.prepare_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) @@ -285,7 +285,7 @@ def test_prepare_load_preserves_key_order(): key_to_block_id[key_b], key_to_block_id[key_c], ] - manager.complete_load([key_a, key_b, key_c]) # order irrelevant + manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant # Arbitrary permutation: [b, c, a] spec_perm = manager.prepare_load([key_b, key_c, key_a], _EMPTY_REQ_CTX) @@ -295,7 +295,7 @@ def test_prepare_load_preserves_key_order(): key_to_block_id[key_c], key_to_block_id[key_a], ] - manager.complete_load([key_a, key_b, key_c]) # order irrelevant + manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant class TestARCPolicy: From 8aac05ee9ca7e129684487746a78fb949e13ad3e Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 11 May 2026 16:32:51 +0300 Subject: [PATCH 39/55] Address review: unify _store_jobs and _load_jobs into _transfer_jobs Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 2 + vllm/v1/kv_offload/tiering/base.py | 1 + vllm/v1/kv_offload/tiering/manager.py | 45 +++++++++---------- 3 files changed, 25 insertions(+), 23 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index e4ad27eb8381..ed81b7a0bcd8 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -114,6 +114,7 @@ def test_lru_eviction(self): job_id=1, keys=[new_block], block_ids=np.array([0], dtype=np.int64), + is_promotion=False, ) ) @@ -142,6 +143,7 @@ def test_async_simulation(self): job_id=1, keys=blocks, block_ids=np.array([0, 1], dtype=np.int64), + is_promotion=False, ) ) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 1424d26d822b..11a6120a744d 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -23,6 +23,7 @@ class JobMetadata: job_id: JobId keys: Collection[OffloadKey] block_ids: np.ndarray + is_promotion: bool req_context: ReqContext = field(default_factory=ReqContext) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 883834bbfcd7..91f30cdad6b6 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -144,11 +144,11 @@ def __init__( self._job_id_counter: int = 0 self.events: list[OffloadingEvent] | None = [] if enable_events else None - # Job tracking: maps job_id to metadata for each transfer direction - # Store jobs: primary → secondary transfers - self._store_jobs: dict[JobId, JobMetadata] = {} - # Load jobs: secondary → primary transfers (promotions) - self._load_jobs: dict[JobId, JobMetadata] = {} + # Job tracking: maps job_id to metadata for all in-flight transfers. + # JobMetadata.is_promotion distinguishes direction: + # True: secondary → primary (promotion) + # False: primary → secondary (cascade) + self._transfer_jobs: dict[JobId, JobMetadata] = {} # Pending promotion requests accumulated during lookup() calls; flushed # as one batched submit_load() per (tier, request) in take_events(). @@ -201,32 +201,29 @@ def _process_finished_jobs(self): for i, tier in enumerate(self.secondary_tiers): for completed_job in tier.get_finished(): job_id = completed_job.job_id + job_metadata = self._transfer_jobs.pop(job_id, None) - # Determine job type by checking which dictionary contains the job_id - if job_id in self._store_jobs: - # primary→secondary transfer completed. - # Decrement ref_cnt on primary blocks. - job_metadata = self._store_jobs.pop(job_id) - self.primary_tier.complete_read( - job_metadata.keys, job_metadata.req_context + if job_metadata is None: + logger.error( + "Received finished job for unknown job_id %d" + " from tier #%d (%s)", + job_id, + i, + tier.get_tier_type(), ) - elif job_id in self._load_jobs: + elif job_metadata.is_promotion: # secondary→primary transfer (promotion) completed. # Make blocks available in primary tier. - job_metadata = self._load_jobs.pop(job_id) self.primary_tier.complete_write( job_metadata.keys, job_metadata.req_context, completed_job.success, ) else: - # Job ID not found in either dictionary - this shouldn't happen - logger.error( - "Received finished job for unknown job_id %d" - " from tier #%d (%s)", - job_id, - i, - tier.get_tier_type(), + # primary→secondary transfer completed. + # Decrement ref_cnt on primary blocks. + self.primary_tier.complete_read( + job_metadata.keys, job_metadata.req_context ) def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: @@ -335,9 +332,10 @@ def _flush_pending_promotions(self) -> None: job_id=job_id, keys=entry.keys, block_ids=np.array(entry.block_ids, dtype=np.int64), + is_promotion=True, req_context=entry.req_context, ) - self._load_jobs[job_id] = job_metadata + self._transfer_jobs[job_id] = job_metadata tier.submit_load(job_metadata) self._pending_load_submissions.clear() @@ -470,9 +468,10 @@ def complete_store( job_id=job_id, keys=keys, block_ids=primary_blocks_spec.block_ids, + is_promotion=False, req_context=req_context, ) - self._store_jobs[job_id] = job_metadata + self._transfer_jobs[job_id] = job_metadata tier.submit_store(job_metadata) From f19ba044825cfdfd3b62fc700ee6937fb036fd8d Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 11 May 2026 16:38:29 +0300 Subject: [PATCH 40/55] Address review: replace error logging with assert Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 91f30cdad6b6..cc3e0fe66396 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -202,16 +202,12 @@ def _process_finished_jobs(self): for completed_job in tier.get_finished(): job_id = completed_job.job_id 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" + ) - if job_metadata is None: - logger.error( - "Received finished job for unknown job_id %d" - " from tier #%d (%s)", - job_id, - i, - tier.get_tier_type(), - ) - elif job_metadata.is_promotion: + if job_metadata.is_promotion: # secondary→primary transfer (promotion) completed. # Make blocks available in primary tier. self.primary_tier.complete_write( From 3e187e101882029213cdfef2a0c0f005458ceeca Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 11 May 2026 17:31:33 +0300 Subject: [PATCH 41/55] Rename primary_store_result-> primary_write_result Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index cc3e0fe66396..d251aa8a8a97 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -292,13 +292,13 @@ def _initiate_promotion( # Must happen immediately so primary.lookup() returns None (in-flight) # for this key on any subsequent lookup() call within the same step, # preventing duplicate promotion attempts. - primary_store_result = self.primary_tier.prepare_write([key], req_context) + primary_write_result = self.primary_tier.prepare_write([key], req_context) - if primary_store_result is None: + if primary_write_result is None: # Cannot allocate space in primary tier (full); retry next step. return - store_spec = primary_store_result.store_spec + store_spec = primary_write_result.store_spec assert isinstance(store_spec, CPULoadStoreSpec) # Defer submit_load to take_events(). Group by (tier, request) so each # request's blocks are submitted as one batched job per tier. @@ -309,7 +309,7 @@ def _initiate_promotion( keys=[], block_ids=[], req_context=req_context ) entry = tier_pending[ctx_id] - entry.keys.extend(primary_store_result.keys_to_store) + entry.keys.extend(primary_write_result.keys_to_store) entry.block_ids.extend(store_spec.block_ids) def _flush_pending_promotions(self) -> None: From cf2cc5e7eaf3732b28bce5907c6b537feca4011c Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 11 May 2026 18:01:26 +0300 Subject: [PATCH 42/55] Address review: return False in lookup() when primary is full Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index d251aa8a8a97..6bfaea04db4e 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -261,8 +261,9 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: elif primary_hit is None: return None elif hit_tier is not None: - self._initiate_promotion(hit_tier, key, req_context) - return None + if self._initiate_promotion(hit_tier, key, req_context): + return None # promotion started, retry later + return False # primary full, block unavailable elif any_none: return None else: @@ -273,7 +274,7 @@ def _initiate_promotion( tier: SecondaryTierManager, key: OffloadKey, req_context: ReqContext, - ): + ) -> bool: """ Queue a block for promotion from a secondary tier to the primary tier. @@ -287,6 +288,9 @@ def _initiate_promotion( tier: The secondary tier to promote from key: Block to promote req_context: Per-request context forwarded to primary.prepare_write(). + + Returns: + True if promotion was initiated, False if primary tier is full. """ # Allocate space in primary tier for promoted block. # Must happen immediately so primary.lookup() returns None (in-flight) @@ -295,8 +299,9 @@ def _initiate_promotion( primary_write_result = self.primary_tier.prepare_write([key], req_context) if primary_write_result is None: - # Cannot allocate space in primary tier (full); retry next step. - return + # Primary tier is full; caller should treat the block as unavailable + # rather than retrying indefinitely. + return False store_spec = primary_write_result.store_spec assert isinstance(store_spec, CPULoadStoreSpec) @@ -311,6 +316,7 @@ def _initiate_promotion( entry = tier_pending[ctx_id] entry.keys.extend(primary_write_result.keys_to_store) entry.block_ids.extend(store_spec.block_ids) + return True def _flush_pending_promotions(self) -> None: """Submit one batched submit_load() per (tier, request). From c1994924538547dacbdf716ffd54a7766903ca55 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Mon, 11 May 2026 18:15:28 +0300 Subject: [PATCH 43/55] Address review: remove in-flight block tracking from example Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 17 +--------- .../v1/kv_offload/tiering/example/__init__.py | 33 +++---------------- 2 files changed, 6 insertions(+), 44 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index ed81b7a0bcd8..730fd2a1c71c 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -76,19 +76,6 @@ def test_basic_store_and_lookup(self): # Third block not present assert tier.lookup(blocks[2], _CTX) is False - def test_in_flight_blocks_return_none(self): - """Test that in-flight blocks cause lookup to return None.""" - tier = ExampleSecondaryTier(max_blocks=10) - - blocks = to_keys(range(3)) - - # Mark first block as in-flight - tier.in_flight[blocks[0]] = 1 - - # Lookup should return None (retry later) - assert tier.lookup(blocks[0], _CTX) is None - assert tier.lookup(blocks[1], _CTX) is False # not in-flight, just absent - def test_lru_eviction(self): """Test LRU eviction policy.""" tier = ExampleSecondaryTier(max_blocks=3) @@ -147,8 +134,7 @@ def test_async_simulation(self): ) ) - # Blocks should be in-flight - assert tier.get_num_in_flight() == 2 + # Blocks should not yet be stored (pending async completion) assert tier.get_num_blocks() == 0 # First get_finished() should complete the job @@ -159,7 +145,6 @@ def test_async_simulation(self): # Blocks should now be stored assert tier.get_num_blocks() == 2 - assert tier.get_num_in_flight() == 0 class TestTieringOffloadingManager: diff --git a/vllm/v1/kv_offload/tiering/example/__init__.py b/vllm/v1/kv_offload/tiering/example/__init__.py index 42a030546de5..69897570cf93 100644 --- a/vllm/v1/kv_offload/tiering/example/__init__.py +++ b/vllm/v1/kv_offload/tiering/example/__init__.py @@ -39,7 +39,6 @@ class ExampleSecondaryTier(SecondaryTierManager): - Stores blocks in a dictionary (key -> True) - Simulates async transfers with immediate completion - Uses LRU eviction policy - - Tracks in-flight transfers to return None from lookup() """ def __init__( @@ -63,9 +62,6 @@ def __init__( # key -> True (only care about presence) self.blocks: OrderedDict[OffloadKey, bool] = OrderedDict() - # Tracks in-flight transfers: key -> job_id - self.in_flight: dict[OffloadKey, JobId] = {} - # Completed jobs waiting to be retrieved by get_finished() self.completed_jobs: list[JobResult] = [] @@ -81,14 +77,11 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: Args: key: Offload key to look up. + req_context: Per-request context. Returns: - True if the block is present and ready, - False if not found, - or None if the block is being transferred (retry later). + True if the block is present, False if not found. """ - if key in self.in_flight: - return None return key in self.blocks def submit_store(self, job_metadata: JobMetadata) -> None: @@ -124,7 +117,7 @@ def submit_store(self, job_metadata: JobMetadata) -> None: # Collect eviction candidates first (LRU order), then delete atomically protected = set(keys) for key in self.blocks: - if key not in protected and key not in self.in_flight: + if key not in protected: evicted.append(key) if len(evicted) == num_blocks_to_evict: break @@ -134,10 +127,6 @@ def submit_store(self, job_metadata: JobMetadata) -> None: for key in evicted: del self.blocks[key] - # Mark blocks as in-flight - for key in blocks_to_store: - self.in_flight[key] = job_id - # Create internal job metadata internal_job_metadata = _JobMetadata( job_id=job_id, keys=blocks_to_store, is_store=True @@ -171,10 +160,6 @@ def submit_load(self, job_metadata: JobMetadata) -> None: if key not in self.blocks: return - # Mark blocks as in-flight - for key in keys: - self.in_flight[key] = job_id - # Create internal job metadata internal_job_metadata = _JobMetadata(job_id=job_id, keys=keys, is_store=False) @@ -211,14 +196,11 @@ 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 - del self.in_flight[key] # 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 by removing in-flight markers.""" - for key in job_metadata.keys: - del self.in_flight[key] + """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)) @@ -242,13 +224,8 @@ def get_num_blocks(self) -> int: """Get the number of blocks currently stored in this tier.""" return len(self.blocks) - def get_num_in_flight(self) -> int: - """Get the number of blocks currently in-flight.""" - return len(self.in_flight) - def clear(self): - """Clear all blocks and in-flight transfers (for testing).""" + """Clear all blocks and pending jobs (for testing).""" self.blocks.clear() - self.in_flight.clear() self.completed_jobs.clear() self.pending_jobs.clear() From d73e4f52c30a2ae1804b5c536e3ae096582f86d5 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 12 May 2026 11:31:21 +0300 Subject: [PATCH 44/55] Address review: rename _PendingPromotion -> PendingPromotion Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 6bfaea04db4e..b2fc8ef6d756 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -47,7 +47,7 @@ @dataclass -class _PendingPromotion: +class PendingPromotion: """Accumulator for blocks awaiting submit_load() for one (tier, request).""" keys: list[OffloadKey] = field(default_factory=list) @@ -155,7 +155,7 @@ def __init__( # Outer key: tier. Inner key: id(req_context) — the same ReqContext # object is reused for all block lookups of a given request per engine step. self._pending_load_submissions: dict[ - SecondaryTierManager, dict[int, _PendingPromotion] + SecondaryTierManager, dict[int, PendingPromotion] ] = {} # Gate for once-per-step execution of _maybe_process_finished_jobs(). @@ -310,7 +310,7 @@ def _initiate_promotion( tier_pending = self._pending_load_submissions.setdefault(tier, {}) ctx_id = id(req_context) if ctx_id not in tier_pending: - tier_pending[ctx_id] = _PendingPromotion( + tier_pending[ctx_id] = PendingPromotion( keys=[], block_ids=[], req_context=req_context ) entry = tier_pending[ctx_id] From 96440684e29c8afec7a4d12149c32ec17060bca9 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 12 May 2026 11:53:47 +0300 Subject: [PATCH 45/55] Address review: add test_tiering_offloading() Signed-off-by: Ronen Schaffer --- .../unit/test_offloading_connector.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index 555daea50e38..35ebfa1e73b4 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -292,3 +292,48 @@ def test_cpu_offloading( if subscriber is not None: subscriber.close() del llm + + +def test_tiering_offloading() -> None: + """Tests OffloadingConnector with TieringOffloadingSpec.""" + extra_config: dict = { + "cpu_bytes_to_use": 500 << 20, + "block_size": 48, + "spec_name": "TieringOffloadingSpec", + "secondary_tiers": [{"type": "example"}], + } + kv_transfer_config = KVTransferConfig( + kv_connector="OffloadingConnector", + kv_role="kv_both", + kv_connector_extra_config=extra_config, + ) + + port: int + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("0.0.0.0", 0)) + port = s.getsockname()[1] + events_endpoint = f"tcp://*:{port}" + kv_events_config = KVEventsConfig( + enable_kv_cache_events=True, + publisher="zmq", + endpoint=events_endpoint, + topic="test", + ) + + llm = LLM( + model="meta-llama/Llama-3.2-1B-Instruct", + max_model_len=4096, + gpu_memory_utilization=0.5, + kv_events_config=kv_events_config, + kv_transfer_config=kv_transfer_config, + ) + subscriber = MockSubscriber( + events_endpoint.replace("*", "127.0.0.1"), + topic=kv_events_config.topic, + ) + try: + _latency_test(llm, subscriber) + _accuracy_test(llm, subscriber) + finally: + subscriber.close() + del llm From b7cc9526437265954070144443382c291fec0d87 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 12 May 2026 12:14:10 +0300 Subject: [PATCH 46/55] Address review: refactor lookup Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 44 ++++++++++++--------------- 1 file changed, 20 insertions(+), 24 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index b2fc8ef6d756..0deed5b383eb 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -227,47 +227,43 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: Check whether a single block is offloaded and ready. Algorithm: - 1. Process any completed async jobs first - 2. Query all tiers (primary + secondaries) unconditionally - 3. Decide based on combined results + 1. Process any completed async jobs first. + 2. Query primary tier — short-circuit on hit or in-flight. + 3. On primary miss, query secondary tiers — stop on first + hit and initiate promotion. Args: key: Block hash to look up. req_context: Per-request context. Returns: - True if the block is ready (primary hit, or found in a secondary - tier with promotion initiated), - False if not found in any tier, - None if no tier has the block but at least one tier is busy - (retry later). + True — block is ready in the primary tier. + None — block found but not yet ready (primary in-flight, + promotion started, or a secondary tier is busy). + False — block not found in any tier, or primary is full + and cannot accept a promotion. """ self._maybe_process_finished_jobs() - # Always query every tier to warm up caches / prefetch state primary_hit = self.primary_tier.lookup(key, req_context) + if primary_hit is True: + return True + if primary_hit is None: + return None - hit_tier = None any_none = False for tier in self.secondary_tiers: result = tier.lookup(key, req_context) - if result is True and hit_tier is None: - hit_tier = tier - elif result is None: + if result is True: + if not self._initiate_promotion(tier, key, req_context): + return False # primary full, block unavailable + return None # promotion started, retry later + if result is None: any_none = True - if primary_hit: - return True - elif primary_hit is None: + if any_none: return None - elif hit_tier is not None: - if self._initiate_promotion(hit_tier, key, req_context): - return None # promotion started, retry later - return False # primary full, block unavailable - elif any_none: - return None - else: - return False + return False def _initiate_promotion( self, From 7cc73263046bdb51259ccc034579ccf2be4ba281 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 12 May 2026 13:44:50 +0300 Subject: [PATCH 47/55] Address review: wire primary_kv_memoryview during init Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 46 ++++++++++--------- vllm/v1/kv_offload/tiering/base.py | 16 ++----- .../v1/kv_offload/tiering/example/__init__.py | 8 ++-- vllm/v1/kv_offload/tiering/factory.py | 8 +++- vllm/v1/kv_offload/tiering/manager.py | 42 ++++++++++------- vllm/v1/kv_offload/tiering/spec.py | 9 ++-- 6 files changed, 65 insertions(+), 64 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index 730fd2a1c71c..f9356d68489c 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -59,7 +59,8 @@ class TestExampleSecondaryTier: def test_basic_store_and_lookup(self): """Test basic store and lookup operations.""" - tier = ExampleSecondaryTier(max_blocks=10) + mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy()) + tier = ExampleSecondaryTier(primary_kv_view=mock_view, max_blocks=10) # Initially empty blocks = to_keys(range(3)) @@ -78,7 +79,8 @@ def test_basic_store_and_lookup(self): def test_lru_eviction(self): """Test LRU eviction policy.""" - tier = ExampleSecondaryTier(max_blocks=3) + mock_view = memoryview(torch.zeros((4, 16), dtype=torch.int8).numpy()) + tier = ExampleSecondaryTier(primary_kv_view=mock_view, max_blocks=3) # Fill tier to capacity blocks = to_keys(range(3)) @@ -93,9 +95,6 @@ def test_lru_eviction(self): # Store new block should evict blocks[1] (least recently used) new_block = to_keys([3])[0] - mock_tensor = torch.zeros((4, 16), dtype=torch.float32) - tier.set_primary_view(memoryview(mock_tensor.numpy())) - tier.submit_store( JobMetadata( job_id=1, @@ -117,13 +116,13 @@ def test_lru_eviction(self): def test_async_simulation(self): """Test simulated async behavior.""" - tier = ExampleSecondaryTier(max_blocks=10, simulate_async=True) + mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy()) + tier = ExampleSecondaryTier( + primary_kv_view=mock_view, max_blocks=10, simulate_async=True + ) blocks = to_keys(range(2)) - mock_tensor = torch.zeros((10, 16), dtype=torch.float32) - tier.set_primary_view(memoryview(mock_tensor.numpy())) - # Submit store job tier.submit_store( JobMetadata( @@ -155,12 +154,15 @@ def manager_setup(self): # Create primary tier (CPU-based) self.primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=5) - mock_arr = torch.zeros((5, 16), dtype=torch.int8).numpy() - self.primary_tier.create_kv_memoryview = lambda: memoryview(mock_arr) + mock_view = memoryview(torch.zeros((5, 16), dtype=torch.int8).numpy()) - # Create secondary tiers - self.secondary_tier1 = ExampleSecondaryTier(max_blocks=10) - self.secondary_tier2 = ExampleSecondaryTier(max_blocks=10) + # Create secondary tiers with the primary view + self.secondary_tier1 = ExampleSecondaryTier( + primary_kv_view=mock_view, max_blocks=10 + ) + self.secondary_tier2 = ExampleSecondaryTier( + primary_kv_view=mock_view, max_blocks=10 + ) # Create tiered manager self.manager = TieringOffloadingManager( @@ -375,16 +377,19 @@ def test_failed_store_no_cascade(self, manager_setup): def test_multiple_secondary_tiers_independent_eviction(self): """Test that secondary tiers manage their own evictions.""" + mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy()) + # Create tier with small capacity - small_tier = ExampleSecondaryTier(max_blocks=5, simulate_async=False) - large_tier = ExampleSecondaryTier(max_blocks=10, simulate_async=False) + small_tier = ExampleSecondaryTier( + primary_kv_view=mock_view, max_blocks=5, simulate_async=False + ) + large_tier = ExampleSecondaryTier( + 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) - mock_arr = torch.zeros((10, 16), dtype=torch.int8).numpy() - primary_tier.create_kv_memoryview = lambda: memoryview(mock_arr) - manager = TieringOffloadingManager( primary_tier=primary_tier, secondary_tiers=[small_tier, large_tier], @@ -509,9 +514,6 @@ def test_works_without_secondary_tiers(self): """Test that manager works with empty secondary_tiers list.""" primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=5) - mock_arr = torch.zeros((5, 16), dtype=torch.int8).numpy() - primary_tier.create_kv_memoryview = lambda: memoryview(mock_arr) - # Create manager with no secondary tiers manager = TieringOffloadingManager( primary_tier=primary_tier, secondary_tiers=[] diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 11a6120a744d..53350fb7f384 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -49,6 +49,9 @@ class SecondaryTierManager(ABC): async jobs; get_finished() polls for completion. """ + def __init__(self, primary_kv_view: memoryview) -> None: + self._primary_kv_view: memoryview = primary_kv_view + @abstractmethod def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: """ @@ -132,19 +135,6 @@ def get_finished(self) -> Iterable[JobResult]: """ pass - def set_primary_view(self, view: memoryview) -> None: - """ - Provide a long-lived memoryview of the primary-tier CPU tensor. - - Called once during initialisation. - Override to store the view for use in `submit_store` and `submit_load`. - Use `view.strides[0]` to obtain the byte stride between block slots. - - Args: - view: Memoryview of the primary tier's CPU KV cache tensor. - """ - return - def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): """ Mark blocks as recently used for eviction policy. diff --git a/vllm/v1/kv_offload/tiering/example/__init__.py b/vllm/v1/kv_offload/tiering/example/__init__.py index 69897570cf93..b283cc9df36d 100644 --- a/vllm/v1/kv_offload/tiering/example/__init__.py +++ b/vllm/v1/kv_offload/tiering/example/__init__.py @@ -43,6 +43,7 @@ class ExampleSecondaryTier(SecondaryTierManager): def __init__( self, + primary_kv_view: memoryview, max_blocks: int = 1000, simulate_async: bool = False, ): @@ -50,15 +51,15 @@ def __init__( Initialize the example secondary tier. Args: + 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__(primary_kv_view) self.max_blocks = max_blocks self.simulate_async = simulate_async - self._primary_view: memoryview | None = None - # key -> True (only care about presence) self.blocks: OrderedDict[OffloadKey, bool] = OrderedDict() @@ -68,9 +69,6 @@ def __init__( # Pending jobs (for simulated async mode) self.pending_jobs: list[_JobMetadata] = [] - def set_primary_view(self, view: memoryview) -> None: - self._primary_view = view - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: """ Check whether a block exists in this secondary tier. diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index 9653175c02c1..ce4637b7a2be 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -14,7 +14,10 @@ } -def create_secondary_tier(tier_config: dict) -> SecondaryTierManager: +def create_secondary_tier( + tier_config: dict, + primary_kv_view: memoryview, +) -> SecondaryTierManager: """ Create a secondary tier from configuration. @@ -23,6 +26,7 @@ def create_secondary_tier(tier_config: dict) -> SecondaryTierManager: - 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. Returns: SecondaryTierManager instance @@ -42,4 +46,4 @@ def create_secondary_tier(tier_config: dict) -> SecondaryTierManager: f"Unknown secondary tier type: {tier_type!r}. " f"Supported types: {list(_TIER_REGISTRY)}" ) - return cls(**config) + return cls(primary_kv_view=primary_kv_view, **config) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 0deed5b383eb..93115c7e1e94 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -86,16 +86,24 @@ def __init__( self.prepare_write = self.prepare_store self.complete_write = self.complete_store - def create_kv_memoryview(self) -> memoryview: - """Create a memoryview over the primary tier's KV cache buffer. + self._kv_memoryview: memoryview | None = None + if mmap_region is not None: + self._kv_memoryview = self._create_kv_memoryview() - Returns a 2-D memoryview of shape (num_blocks, row_stride_bytes) - backed by the SharedOffloadRegion mmap. Secondary tiers address - block b as view[b]. Caller must call release() when done. + def get_kv_memoryview(self) -> memoryview: + """Return the memoryview over the primary tier's KV cache buffer. + + The view has shape (num_blocks, row_stride_bytes) and is backed by the + SharedOffloadRegion mmap. Secondary tiers address block *b* as + ``view[b]``. """ - assert self._mmap_region is not None, ( + assert self._kv_memoryview is not None, ( "mmap_region must be provided to CPUPrimaryTierOffloadingManager" ) + return self._kv_memoryview + + def _create_kv_memoryview(self) -> memoryview: + assert self._mmap_region is not None kv_tensor = self._mmap_region._base.view( self._mmap_region.num_blocks, self._mmap_region._row_stride ) @@ -106,6 +114,15 @@ def create_kv_memoryview(self) -> memoryview: ) return memoryview(np_arr) + def shutdown(self) -> None: + super().shutdown() + if self._kv_memoryview is not None: + self._kv_memoryview.release() + self._kv_memoryview = None + if self._mmap_region is not None: + self._mmap_region.cleanup() + self._mmap_region = None + class TieringOffloadingManager(OffloadingManager): """ @@ -162,12 +179,6 @@ def __init__( # Reset at the end of each step in take_events(). self._processed_jobs_this_step: bool = False - # Wire each secondary tier with a long-lived memoryview of the primary - # CPU buffer. One view is shared across all tiers; released in shutdown(). - self._primary_kv_view = primary_tier.create_kv_memoryview() - for tier in self.secondary_tiers: - tier.set_primary_view(self._primary_kv_view) - def _next_job_id(self) -> JobId: """Generate a unique job ID for async transfer tracking.""" job_id = self._job_id_counter @@ -509,8 +520,5 @@ def take_events(self) -> Iterable[OffloadingEvent]: yield from self.primary_tier.take_events() def shutdown(self) -> None: - """Release memoryviews and scheduler-side mmap.""" - self._primary_kv_view.release() - if self.primary_tier._mmap_region is not None: - self.primary_tier._mmap_region.cleanup() - self.primary_tier._mmap_region = None + """Shutdown all tiers and release resources.""" + self.primary_tier.shutdown() diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 7fc41407e87a..49d8f9429ff5 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -94,10 +94,8 @@ def get_manager(self) -> OffloadingManager: kv_events_config is not None and kv_events_config.enable_kv_cache_events ) - # Create scheduler-side SharedOffloadRegion (rank=None) first so - # CPUPrimaryTierOffloadingManager.create_kv_memoryview() has _base - # available when TieringOffloadingManager.__init__ wires secondary - # tier memoryviews. + # Create scheduler-side SharedOffloadRegion (rank=None) so the + # primary tier can eagerly create a memoryview over _base. world_size = self.vllm_config.parallel_config.world_size scheduler_mmap = SharedOffloadRegion( instance_id=self.vllm_config.instance_id, @@ -121,10 +119,11 @@ def get_manager(self) -> OffloadingManager: ) # Create secondary tiers + primary_kv_view = primary_tier.get_kv_memoryview() secondary_tiers = [] for i, tier_config in enumerate(self.secondary_tier_configs): try: - tier = create_secondary_tier(tier_config) + tier = create_secondary_tier(tier_config, primary_kv_view) secondary_tiers.append(tier) logger.info( "Created secondary tier #%d (%s)", From 1129c1d45eb633ec8c47e8e8b1174b272d3059cc Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 13 May 2026 10:18:38 +0300 Subject: [PATCH 48/55] Add shutdown to SecondaryTierManager Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/base.py | 4 ++++ vllm/v1/kv_offload/tiering/manager.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 53350fb7f384..842fba5df216 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -145,6 +145,10 @@ def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): """ return + def shutdown(self) -> None: + """Release resources held by this tier (threads, connections, etc.).""" + return + @staticmethod @abstractmethod def get_tier_type() -> str: diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 93115c7e1e94..d50312e013c7 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -521,4 +521,6 @@ def take_events(self) -> Iterable[OffloadingEvent]: def shutdown(self) -> None: """Shutdown all tiers and release resources.""" + for tier in self.secondary_tiers: + tier.shutdown() self.primary_tier.shutdown() From c1e0a2594b9cbf86a5948922ae18ad11c17c440b Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 13 May 2026 10:59:43 +0300 Subject: [PATCH 49/55] Pass vllm_config to SecondaryTierManager Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 28 ++++++++++++++----- vllm/v1/kv_offload/tiering/base.py | 9 +++++- .../v1/kv_offload/tiering/example/__init__.py | 10 ++++++- vllm/v1/kv_offload/tiering/factory.py | 11 +++++++- vllm/v1/kv_offload/tiering/spec.py | 4 ++- 5 files changed, 51 insertions(+), 11 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index f9356d68489c..f33474715253 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -31,6 +31,7 @@ ) _CTX = ReqContext() +_MOCK_VLLM_CONFIG = MagicMock() def to_keys(int_ids: Iterable[int]) -> list[OffloadKey]: @@ -60,7 +61,9 @@ class TestExampleSecondaryTier: 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(primary_kv_view=mock_view, max_blocks=10) + tier = ExampleSecondaryTier( + vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10 + ) # Initially empty blocks = to_keys(range(3)) @@ -80,7 +83,9 @@ def test_basic_store_and_lookup(self): def test_lru_eviction(self): """Test LRU eviction policy.""" mock_view = memoryview(torch.zeros((4, 16), dtype=torch.int8).numpy()) - tier = ExampleSecondaryTier(primary_kv_view=mock_view, max_blocks=3) + tier = ExampleSecondaryTier( + vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=3 + ) # Fill tier to capacity blocks = to_keys(range(3)) @@ -118,7 +123,10 @@ def test_async_simulation(self): """Test simulated async behavior.""" mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy()) tier = ExampleSecondaryTier( - primary_kv_view=mock_view, max_blocks=10, simulate_async=True + vllm_config=_MOCK_VLLM_CONFIG, + primary_kv_view=mock_view, + max_blocks=10, + simulate_async=True, ) blocks = to_keys(range(2)) @@ -158,10 +166,10 @@ def manager_setup(self): # Create secondary tiers with the primary view self.secondary_tier1 = ExampleSecondaryTier( - primary_kv_view=mock_view, max_blocks=10 + vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10 ) self.secondary_tier2 = ExampleSecondaryTier( - primary_kv_view=mock_view, max_blocks=10 + vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10 ) # Create tiered manager @@ -381,10 +389,16 @@ def test_multiple_secondary_tiers_independent_eviction(self): # Create tier with small capacity small_tier = ExampleSecondaryTier( - primary_kv_view=mock_view, max_blocks=5, simulate_async=False + vllm_config=_MOCK_VLLM_CONFIG, + primary_kv_view=mock_view, + max_blocks=5, + simulate_async=False, ) large_tier = ExampleSecondaryTier( - primary_kv_view=mock_view, max_blocks=10, simulate_async=False + vllm_config=_MOCK_VLLM_CONFIG, + primary_kv_view=mock_view, + max_blocks=10, + simulate_async=False, ) # Create a fresh primary tier for this test diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 842fba5df216..c15d3963b577 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -4,14 +4,20 @@ Abstract interfaces and data types for the secondary tiering layer. """ +from __future__ import annotations + from abc import ABC, abstractmethod from collections.abc import Collection, Iterable from dataclasses import dataclass, field +from typing import TYPE_CHECKING import numpy as np from vllm.v1.kv_offload.base import OffloadKey, ReqContext +if TYPE_CHECKING: + from vllm.config import VllmConfig + # Type alias for job IDs used in async transfer tracking JobId = int @@ -49,7 +55,8 @@ class SecondaryTierManager(ABC): async jobs; get_finished() polls for completion. """ - def __init__(self, primary_kv_view: memoryview) -> None: + def __init__(self, vllm_config: VllmConfig, primary_kv_view: memoryview) -> None: + self._vllm_config = vllm_config self._primary_kv_view: memoryview = primary_kv_view @abstractmethod diff --git a/vllm/v1/kv_offload/tiering/example/__init__.py b/vllm/v1/kv_offload/tiering/example/__init__.py index b283cc9df36d..97a9a44d1617 100644 --- a/vllm/v1/kv_offload/tiering/example/__init__.py +++ b/vllm/v1/kv_offload/tiering/example/__init__.py @@ -9,9 +9,12 @@ requiring actual storage or network backends. """ +from __future__ import annotations + 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 ( @@ -21,6 +24,9 @@ SecondaryTierManager, ) +if TYPE_CHECKING: + from vllm.config import VllmConfig + @dataclass class _JobMetadata: @@ -43,6 +49,7 @@ class ExampleSecondaryTier(SecondaryTierManager): def __init__( self, + vllm_config: VllmConfig, primary_kv_view: memoryview, max_blocks: int = 1000, simulate_async: bool = False, @@ -51,12 +58,13 @@ def __init__( 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__(primary_kv_view) + super().__init__(vllm_config, primary_kv_view) self.max_blocks = max_blocks self.simulate_async = simulate_async diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index ce4637b7a2be..e24b1963fbb3 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -4,9 +4,16 @@ Factory for creating secondary tier implementations. """ +from __future__ import annotations + +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]] = { @@ -17,6 +24,7 @@ def create_secondary_tier( tier_config: dict, primary_kv_view: memoryview, + vllm_config: VllmConfig, ) -> SecondaryTierManager: """ Create a secondary tier from configuration. @@ -27,6 +35,7 @@ def create_secondary_tier( - 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. Returns: SecondaryTierManager instance @@ -46,4 +55,4 @@ def create_secondary_tier( f"Unknown secondary tier type: {tier_type!r}. " f"Supported types: {list(_TIER_REGISTRY)}" ) - return cls(primary_kv_view=primary_kv_view, **config) + return cls(vllm_config=vllm_config, primary_kv_view=primary_kv_view, **config) diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 49d8f9429ff5..0a372ae8ef8d 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -123,7 +123,9 @@ def get_manager(self) -> OffloadingManager: secondary_tiers = [] for i, tier_config in enumerate(self.secondary_tier_configs): try: - tier = create_secondary_tier(tier_config, primary_kv_view) + tier = create_secondary_tier( + tier_config, primary_kv_view, self.vllm_config + ) secondary_tiers.append(tier) logger.info( "Created secondary tier #%d (%s)", From d63db1e02ce4c418fcd443b8d723209caf8d4e57 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 13 May 2026 12:48:11 +0300 Subject: [PATCH 50/55] Address review: move create_kv_memoryview to SharedOffloadRegion Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/cpu/shared_offload_region.py | 14 ++++++++++++++ vllm/v1/kv_offload/tiering/manager.py | 14 +------------- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index b2e21d06c9b7..1166b44fc7e3 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -153,6 +153,20 @@ def create_next_view(self, tensor_page_size: int) -> torch.Tensor: self._views.append(worker_layer_view) return worker_layer_view + def create_kv_memoryview(self) -> memoryview: + """Return a zero-copy memoryview over the entire KV buffer. + + Shape: (num_blocks, row_stride_bytes). Secondary tiers address + block *b* as ``view[b]``. + """ + kv_tensor = self._base.view(self.num_blocks, self._row_stride) + np_arr = kv_tensor.numpy() + assert np_arr.ctypes.data == self._base.data_ptr(), ( + "view()/numpy() created a copy instead of sharing the mmap buffer; " + "secondary tiers require zero-copy access to primary KV data" + ) + return memoryview(np_arr) + def cleanup(self) -> None: if self.is_pinned and self._base is not None: base_ptr = self._base.data_ptr() diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index d50312e013c7..28817769c097 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -88,7 +88,7 @@ def __init__( self._kv_memoryview: memoryview | None = None if mmap_region is not None: - self._kv_memoryview = self._create_kv_memoryview() + self._kv_memoryview = mmap_region.create_kv_memoryview() def get_kv_memoryview(self) -> memoryview: """Return the memoryview over the primary tier's KV cache buffer. @@ -102,18 +102,6 @@ def get_kv_memoryview(self) -> memoryview: ) return self._kv_memoryview - def _create_kv_memoryview(self) -> memoryview: - assert self._mmap_region is not None - kv_tensor = self._mmap_region._base.view( - self._mmap_region.num_blocks, self._mmap_region._row_stride - ) - np_arr = kv_tensor.numpy() - assert np_arr.ctypes.data == self._mmap_region._base.data_ptr(), ( - "view()/numpy() created a copy instead of sharing the mmap buffer; " - "secondary tiers require zero-copy access to primary KV data" - ) - return memoryview(np_arr) - def shutdown(self) -> None: super().shutdown() if self._kv_memoryview is not None: From 414def68dba001443cd6ae5c72c09b340c52f035 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 13 May 2026 11:43:43 +0300 Subject: [PATCH 51/55] cherry-pick: Add req_id to ReqContext for per-request tracking Signed-off-by: Ronen Schaffer --- .../unit/offloading_connector/test_scheduler.py | 2 +- tests/v1/kv_offload/cpu/test_manager.py | 6 ++++-- .../kv_transfer/kv_connector/v1/offloading/scheduler.py | 5 ++++- vllm/v1/kv_offload/base.py | 1 + 4 files changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 22eb1a86d4fb..9a8783b084e0 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -627,7 +627,7 @@ def _make_scheduler_with_lookup( return scheduler -_EMPTY_REQ_CTX = ReqContext() +_EMPTY_REQ_CTX = ReqContext(req_id="") class TestMaximalPrefixLookup: diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 95f8f305ddd1..3957294f8b0f 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -19,9 +19,11 @@ from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy -def make_req_context(kv_transfer_params: dict | None = None) -> ReqContext: +def make_req_context( + req_id: str = "", kv_transfer_params: dict | None = None +) -> ReqContext: """Create a ReqContext as production code would, from a request's params.""" - return ReqContext(kv_transfer_params=kv_transfer_params) + return ReqContext(req_id=req_id, kv_transfer_params=kv_transfer_params) _EMPTY_REQ_CTX = make_req_context() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 137eaef9788c..ef016bb02ebb 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -138,7 +138,10 @@ def __post_init__(self) -> None: self.group_states = tuple( RequestGroupState() for _ in self.config.kv_group_configs ) - self.req_context = ReqContext(kv_transfer_params=self.req.kv_transfer_params) + self.req_context = ReqContext( + req_id=self.req.request_id, + kv_transfer_params=self.req.kv_transfer_params, + ) def update_offload_keys(self) -> None: for group_config, group_state in zip( diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index afa3a06774c0..bcf1d2462c7b 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -46,6 +46,7 @@ def get_offload_group_idx(key: OffloadKey) -> int: @dataclass class ReqContext: + req_id: str kv_transfer_params: dict[str, Any] | None = None From 8ef126b47341545adacdff28e94b2d3282f4c302 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 13 May 2026 13:05:41 +0300 Subject: [PATCH 52/55] Address review: use req_id instead of id(req_context) for pending promotions Signed-off-by: Ronen Schaffer --- tests/v1/kv_offload/test_tiering_offloading.py | 14 ++++++++------ vllm/v1/kv_offload/tiering/base.py | 4 ++-- vllm/v1/kv_offload/tiering/manager.py | 8 ++++---- 3 files changed, 14 insertions(+), 12 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index f33474715253..bdde9f4960c4 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -30,7 +30,7 @@ TieringOffloadingManager, ) -_CTX = ReqContext() +_CTX = ReqContext(req_id="test") _MOCK_VLLM_CONFIG = MagicMock() @@ -106,6 +106,7 @@ def test_lru_eviction(self): keys=[new_block], block_ids=np.array([0], dtype=np.int64), is_promotion=False, + req_context=_CTX, ) ) @@ -138,6 +139,7 @@ def test_async_simulation(self): keys=blocks, block_ids=np.array([0, 1], dtype=np.int64), is_promotion=False, + req_context=_CTX, ) ) @@ -447,8 +449,8 @@ def test_lookup_batches_submit_load_per_request(self, manager_setup): wraps=self.secondary_tier1.submit_load ) - ctx_a = ReqContext() - ctx_b = ReqContext() + ctx_a = ReqContext(req_id="req_a") + ctx_b = ReqContext(req_id="req_b") # All lookups return None: secondary hit triggers promotion (in-flight) assert self.manager.lookup(blocks[0], ctx_a) is None @@ -485,8 +487,8 @@ def test_lookup_shared_block_no_duplicate_promotion(self, manager_setup): wraps=self.secondary_tier1.submit_load ) - ctx_a = ReqContext() - ctx_b = ReqContext() + ctx_a = ReqContext(req_id="req_a") + ctx_b = ReqContext(req_id="req_b") result_a = self.manager.lookup(shared_block, ctx_a) result_b = self.manager.lookup(shared_block, ctx_b) @@ -511,7 +513,7 @@ def test_complete_store_forwards_req_context_to_submit_store(self, manager_setup wraps=self.secondary_tier1.submit_store ) - ctx = ReqContext(kv_transfer_params={"key": "value"}) + ctx = ReqContext(req_id="req_ctx", kv_transfer_params={"key": "value"}) self.manager.prepare_store(blocks, ctx) self.manager.complete_store(blocks, ctx, success=True) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index c15d3963b577..50f0d37dd8e1 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -8,7 +8,7 @@ from abc import ABC, abstractmethod from collections.abc import Collection, Iterable -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TYPE_CHECKING import numpy as np @@ -30,7 +30,7 @@ class JobMetadata: keys: Collection[OffloadKey] block_ids: np.ndarray is_promotion: bool - req_context: ReqContext = field(default_factory=ReqContext) + req_context: ReqContext @dataclass diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 28817769c097..f83cefc9842e 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -50,9 +50,9 @@ class PendingPromotion: """Accumulator for blocks awaiting submit_load() for one (tier, request).""" + req_context: ReqContext keys: list[OffloadKey] = field(default_factory=list) block_ids: list[int] = field(default_factory=list) - req_context: ReqContext = field(default_factory=ReqContext) class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): @@ -157,10 +157,10 @@ def __init__( # Pending promotion requests accumulated during lookup() calls; flushed # as one batched submit_load() per (tier, request) in take_events(). - # Outer key: tier. Inner key: id(req_context) — the same ReqContext + # Outer key: tier. Inner key: req_context.req_id — the same ReqContext # object is reused for all block lookups of a given request per engine step. self._pending_load_submissions: dict[ - SecondaryTierManager, dict[int, PendingPromotion] + SecondaryTierManager, dict[str, PendingPromotion] ] = {} # Gate for once-per-step execution of _maybe_process_finished_jobs(). @@ -303,7 +303,7 @@ def _initiate_promotion( # Defer submit_load to take_events(). Group by (tier, request) so each # request's blocks are submitted as one batched job per tier. tier_pending = self._pending_load_submissions.setdefault(tier, {}) - ctx_id = id(req_context) + ctx_id = req_context.req_id if ctx_id not in tier_pending: tier_pending[ctx_id] = PendingPromotion( keys=[], block_ids=[], req_context=req_context From 6a2bfc600350d20a49827e4c8383ecfcd8c5ce13 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 13 May 2026 13:36:44 +0300 Subject: [PATCH 53/55] Address review: replace `from __future__ import annotations` with string literals Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/base.py | 6 ++---- vllm/v1/kv_offload/tiering/base.py | 4 +--- vllm/v1/kv_offload/tiering/example/__init__.py | 4 +--- vllm/v1/kv_offload/tiering/factory.py | 4 +--- 4 files changed, 5 insertions(+), 13 deletions(-) diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index bcf1d2462c7b..6f80ba6252be 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -4,8 +4,6 @@ Core abstractions for KV cache offloading in vLLM v1. """ -from __future__ import annotations - from abc import ABC, abstractmethod from collections.abc import Collection, Iterable, Iterator, Sequence from dataclasses import dataclass @@ -328,7 +326,7 @@ class CanonicalKVCaches: class OffloadingSpec(ABC): """Spec for an offloading connector""" - def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): + def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"): logger.warning( "Initializing OffloadingSpec. This API is experimental and " "subject to change in the future as we iterate the design." @@ -394,7 +392,7 @@ def get_manager(self) -> OffloadingManager: @abstractmethod def get_handlers( self, kv_caches: CanonicalKVCaches - ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]: + ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], "OffloadingHandler"]]: """ Get offloading handlers along with their respective src and dst types. diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 50f0d37dd8e1..ace325185ac1 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -4,8 +4,6 @@ Abstract interfaces and data types for the secondary tiering layer. """ -from __future__ import annotations - from abc import ABC, abstractmethod from collections.abc import Collection, Iterable from dataclasses import dataclass @@ -55,7 +53,7 @@ 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) -> None: self._vllm_config = vllm_config self._primary_kv_view: memoryview = primary_kv_view diff --git a/vllm/v1/kv_offload/tiering/example/__init__.py b/vllm/v1/kv_offload/tiering/example/__init__.py index 97a9a44d1617..8c3edb693f3c 100644 --- a/vllm/v1/kv_offload/tiering/example/__init__.py +++ b/vllm/v1/kv_offload/tiering/example/__init__.py @@ -9,8 +9,6 @@ requiring actual storage or network backends. """ -from __future__ import annotations - from collections import OrderedDict from collections.abc import Collection, Iterable from dataclasses import dataclass @@ -49,7 +47,7 @@ class ExampleSecondaryTier(SecondaryTierManager): def __init__( self, - vllm_config: VllmConfig, + vllm_config: "VllmConfig", primary_kv_view: memoryview, max_blocks: int = 1000, simulate_async: bool = False, diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index e24b1963fbb3..e8f4a9b32dde 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -4,8 +4,6 @@ Factory for creating secondary tier implementations. """ -from __future__ import annotations - from typing import TYPE_CHECKING from vllm.v1.kv_offload.tiering.base import SecondaryTierManager @@ -24,7 +22,7 @@ def create_secondary_tier( tier_config: dict, primary_kv_view: memoryview, - vllm_config: VllmConfig, + vllm_config: "VllmConfig", ) -> SecondaryTierManager: """ Create a secondary tier from configuration. From 100418f6bd60c09109a890f49c258a1f8e13ac17 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 13 May 2026 14:26:13 +0300 Subject: [PATCH 54/55] Address review: make mmap_region non-optional in CPUPrimaryTierOffloadingManager Signed-off-by: Ronen Schaffer --- .../v1/kv_offload/test_tiering_offloading.py | 26 +++++++++++++++---- vllm/v1/kv_offload/tiering/manager.py | 9 ++----- 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index bdde9f4960c4..a5c8d8673d25 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -34,6 +34,14 @@ _MOCK_VLLM_CONFIG = MagicMock() +def _mock_mmap_region(num_blocks: int, row_bytes: int = 16): + """Create a mock SharedOffloadRegion for testing.""" + mock = MagicMock() + view = memoryview(torch.zeros((num_blocks, row_bytes), dtype=torch.int8).numpy()) + mock.create_kv_memoryview.return_value = view + return mock + + def to_keys(int_ids: Iterable[int]) -> list[OffloadKey]: return [make_offload_key(str(i).encode(), 0) for i in int_ids] @@ -162,9 +170,12 @@ class TestTieringOffloadingManager: @pytest.fixture def manager_setup(self): # Create primary tier (CPU-based) - self.primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=5) + mock_region = _mock_mmap_region(5) + self.primary_tier = CPUPrimaryTierOffloadingManager( + num_blocks=5, mmap_region=mock_region + ) - mock_view = memoryview(torch.zeros((5, 16), dtype=torch.int8).numpy()) + mock_view = mock_region.create_kv_memoryview() # Create secondary tiers with the primary view self.secondary_tier1 = ExampleSecondaryTier( @@ -387,7 +398,8 @@ def test_failed_store_no_cascade(self, manager_setup): def test_multiple_secondary_tiers_independent_eviction(self): """Test that secondary tiers manage their own evictions.""" - mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy()) + mock_region = _mock_mmap_region(10) + mock_view = mock_region.create_kv_memoryview() # Create tier with small capacity small_tier = ExampleSecondaryTier( @@ -404,7 +416,9 @@ def test_multiple_secondary_tiers_independent_eviction(self): ) # Create a fresh primary tier for this test - primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=10) + primary_tier = CPUPrimaryTierOffloadingManager( + num_blocks=10, mmap_region=mock_region + ) manager = TieringOffloadingManager( primary_tier=primary_tier, @@ -528,7 +542,9 @@ class TestTieringOffloadingWithoutSecondaryTiers: def test_works_without_secondary_tiers(self): """Test that manager works with empty secondary_tiers list.""" - primary_tier = CPUPrimaryTierOffloadingManager(num_blocks=5) + primary_tier = CPUPrimaryTierOffloadingManager( + num_blocks=5, mmap_region=_mock_mmap_region(5) + ) # Create manager with no secondary tiers manager = TieringOffloadingManager( diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index f83cefc9842e..8ef3bfc3bb2d 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -68,9 +68,9 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): def __init__( self, num_blocks: int, + mmap_region: SharedOffloadRegion, cache_policy: str = "lru", enable_events: bool = False, - mmap_region: SharedOffloadRegion | None = None, ): super().__init__( num_blocks=num_blocks, @@ -86,9 +86,7 @@ def __init__( self.prepare_write = self.prepare_store self.complete_write = self.complete_store - self._kv_memoryview: memoryview | None = None - if mmap_region is not None: - self._kv_memoryview = mmap_region.create_kv_memoryview() + self._kv_memoryview = mmap_region.create_kv_memoryview() def get_kv_memoryview(self) -> memoryview: """Return the memoryview over the primary tier's KV cache buffer. @@ -97,9 +95,6 @@ def get_kv_memoryview(self) -> memoryview: SharedOffloadRegion mmap. Secondary tiers address block *b* as ``view[b]``. """ - assert self._kv_memoryview is not None, ( - "mmap_region must be provided to CPUPrimaryTierOffloadingManager" - ) return self._kv_memoryview def shutdown(self) -> None: From f0782e6aa1ce03ebeff61f89cc22d1efed0a282c Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 13 May 2026 14:44:17 +0300 Subject: [PATCH 55/55] Fix mypy errors Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/tiering/manager.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 8ef3bfc3bb2d..578992fdd731 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -99,12 +99,8 @@ def get_kv_memoryview(self) -> memoryview: def shutdown(self) -> None: super().shutdown() - if self._kv_memoryview is not None: - self._kv_memoryview.release() - self._kv_memoryview = None - if self._mmap_region is not None: - self._mmap_region.cleanup() - self._mmap_region = None + self._kv_memoryview.release() + self._mmap_region.cleanup() class TieringOffloadingManager(OffloadingManager):