[kv_offload] Add multi-tier KV cache offloading framework

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
@@ -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,24 @@ 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:
@@ -174,7 +196,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 +217,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

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

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",
+)

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",
+]