[KV Offload] Add per-request offloading policy via on_new_request lifecycle hook (#43205)

Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Co-authored-by: Or Ozeri <or@ozery.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

@@ -24,7 +24,9 @@ from vllm.v1.kv_cache_interface import (
 from vllm.v1.kv_offload.base import (
     OffloadingEvent,
     OffloadingManager,
+    OffloadPolicy,
     ReqContext,
+    RequestOffloadingContext,
     get_offload_block_hash,
 )
 from vllm.v1.request import RequestStatus
@@ -748,12 +750,8 @@ class TestSlidingWindowLookup:
 
 
 @pytest.mark.parametrize("async_scheduling", [True, False])
-def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bool):
-    """With do_remote_decode=True, after loading prefix blocks from CPU,
-    all blocks must be re-stored — not just the newly computed ones.
-
-    This supports P/D disaggregation where the prefill instance offloads the
-    complete KV cache so a remote decode node can consume it."""
+def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling: bool):
+    """With REQUEST_LEVEL policy, all blocks are stored — including prefix hits."""
     gpu_block_size = 4
     block_size_factor = 3
     offloaded_block_size = gpu_block_size * block_size_factor
@@ -780,12 +778,13 @@ def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bo
     # Reset GPU prefix cache so the next request must load from CPU.
     runner.scheduler.reset_prefix_cache()
 
-    # New request with do_remote_decode=True and 2 offloaded blocks.
-    # The first offloaded block matches what we stored in CPU.
-    runner.new_request(
-        token_ids=[0] * offloaded_block_size * 2,
-        kv_transfer_params={"do_remote_decode": True},
+    # Manager returns REQUEST_LEVEL for the next request.
+    runner.manager.on_new_request.return_value = RequestOffloadingContext(
+        policy=OffloadPolicy.REQUEST_LEVEL
     )
+
+    # New request with 2 offloaded blocks; first matches what's in CPU.
+    runner.new_request(token_ids=[0] * offloaded_block_size * 2)
     runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output(keys)
@@ -957,9 +956,10 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
         )
 
     def setup(r, max_offload_tokens):
-        r.new_request(token_ids=[0] * offloaded_block_size * 3)
-        req = r.scheduler.requests[str(r.req_id)]
-        req.kv_transfer_params = {"max_offload_tokens": max_offload_tokens}
+        r.new_request(
+            token_ids=[0] * offloaded_block_size * 3,
+            kv_transfer_params={"max_offload_tokens": max_offload_tokens},
+        )
         r.manager.prepare_store.side_effect = (
             lambda keys, req_context: generate_store_output(keys)
         )

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -43,6 +43,7 @@ from vllm.v1.kv_offload.base import (
     OffloadingSpec,
     OffloadKey,
     PrepareStoreOutput,
+    RequestOffloadingContext,
     make_offload_key,
 )
 from vllm.v1.kv_offload.worker.worker import (
@@ -119,6 +120,7 @@ class MockOffloadingSpec(OffloadingSpec):
         self.manager.lookup.return_value = 0
         self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)
         self.manager.lookup.return_value = False
+        self.manager.on_new_request.return_value = RequestOffloadingContext()
         self.handler = MockOffloadingHandler()
 
     def get_manager(self) -> OffloadingManager:

tests/v1/kv_offload/test_tiering_offloading.py

@@ -19,7 +19,9 @@ import torch
 
 from vllm.v1.kv_offload.base import (
     OffloadKey,
+    OffloadPolicy,
     ReqContext,
+    RequestOffloadingContext,
     make_offload_key,
 )
 from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
@@ -412,6 +414,78 @@ class TestTieringOffloadingManager:
         job_metadata = self.secondary_tier1.submit_store.call_args.args[0]
         assert job_metadata.req_context is ctx
 
+    def test_on_new_request_lifecycle(self, manager_setup):
+        """Policy defaults to BLOCK_LEVEL, escalates when a tier requests it,
+        and is cleaned up on on_request_finished."""
+        # Default: all tiers return BLOCK_LEVEL
+        ctx = ReqContext(req_id="req_policy_lifecycle")
+        result = self.manager.on_new_request(ctx)
+        assert result.policy == OffloadPolicy.BLOCK_LEVEL
+        self.manager.on_request_finished(ctx)
+
+        # Escalate: tier1 requests REQUEST_LEVEL
+        self.secondary_tier1.on_new_request = (
+            lambda req_context: RequestOffloadingContext(
+                policy=OffloadPolicy.REQUEST_LEVEL
+            )
+        )
+
+        ctx = ReqContext(req_id="req_policy_lifecycle_2")
+        result = self.manager.on_new_request(ctx)
+        assert result.policy == OffloadPolicy.REQUEST_LEVEL
+        assert ctx.req_id in self.manager._request_level_tiers
+
+        # Cleanup
+        self.manager.on_request_finished(ctx)
+        assert ctx.req_id not in self.manager._request_level_tiers
+
+    def test_prepare_store_cascades_existing_blocks_to_request_level_tiers(
+        self, manager_setup
+    ):
+        """prepare_store cascades hit blocks to request-level tiers only."""
+        # Store some blocks to primary first
+        existing_blocks = to_keys(range(3))
+        result = self.manager.prepare_store(existing_blocks, _CTX)
+        assert result is not None
+        self.manager.complete_store(existing_blocks, _CTX, success=True)
+        # Drain cascade completions
+        list(self.manager.take_events())
+        list(self.manager.take_events())
+
+        # Make tier1 request-level, tier2 stays block-level
+        self.secondary_tier1.on_new_request = (
+            lambda req_context: RequestOffloadingContext(
+                policy=OffloadPolicy.REQUEST_LEVEL
+            )
+        )
+
+        ctx = ReqContext(req_id="req_cascade")
+        self.manager.on_new_request(ctx)
+
+        # Spy on submit_store
+        self.secondary_tier1.submit_store = MagicMock(
+            wraps=self.secondary_tier1.submit_store
+        )
+        self.secondary_tier2.submit_store = MagicMock(
+            wraps=self.secondary_tier2.submit_store
+        )
+
+        # Call prepare_store with existing + new blocks
+        new_blocks = to_keys(range(3, 5))
+        all_blocks = existing_blocks + new_blocks
+        result = self.manager.prepare_store(all_blocks, ctx)
+        assert result is not None
+        assert set(result.keys_to_store) == set(new_blocks)
+
+        # Only tier1 (request-level) should get existing blocks cascaded now.
+        # New blocks are cascaded to ALL tiers later via complete_store().
+        self.secondary_tier1.submit_store.assert_called_once()
+        job_metadata = self.secondary_tier1.submit_store.call_args.args[0]
+        assert set(job_metadata.keys) == set(existing_blocks)
+
+        # tier2 (block-level) does not get existing blocks here.
+        self.secondary_tier2.submit_store.assert_not_called()
+
 
 class TestTieringOffloadingWithoutSecondaryTiers:
     """Test TieringOffloadingManager with no secondary tiers (backward compat)."""

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

@@ -29,7 +29,9 @@ from vllm.v1.kv_offload.base import (
     OffloadingManager,
     OffloadingSpec,
     OffloadKey,
+    OffloadPolicy,
     ReqContext,
+    RequestOffloadingContext,
     get_offload_block_hash,
     make_offload_key,
 )
@@ -172,8 +174,9 @@ class RequestGroupState:
 class RequestOffloadState:
     config: SchedulerOffloadConfig
     req: Request
+    req_context: ReqContext
+    offloading_context: RequestOffloadingContext
     group_states: tuple[RequestGroupState, ...] = field(init=False)
-    req_context: ReqContext = field(init=False)
     # upper bound on tokens to offload for this request; None means no cap
     max_offload_tokens: int | None = None
     # number of hits in the GPU cache
@@ -186,10 +189,6 @@ class RequestOffloadState:
         self.group_states = tuple(
             RequestGroupState() for _ in self.config.kv_group_configs
         )
-        self.req_context = ReqContext(
-            req_id=self.req.request_id,
-            kv_transfer_params=self.req.kv_transfer_params,
-        )
         params = self.req.kv_transfer_params
 
         # NOTE: This field is experimental and subject to change in the future.
@@ -248,6 +247,13 @@ class RequestOffloadState:
             )
 
 
+def _create_req_context(req: Request) -> ReqContext:
+    return ReqContext(
+        req_id=req.request_id,
+        kv_transfer_params=req.kv_transfer_params,
+    )
+
+
 class OffloadingConnectorScheduler:
     """Implementation of Scheduler side methods"""
 
@@ -512,6 +518,18 @@ class OffloadingConnectorScheduler:
 
         return num_hit_tokens
 
+    def on_new_request(self, request: Request) -> None:
+        """Called when a new request is added to the scheduler."""
+        req_context = _create_req_context(request)
+        offloading_context = self.manager.on_new_request(req_context)
+        req_status = RequestOffloadState(
+            config=self.config,
+            req=request,
+            req_context=req_context,
+            offloading_context=offloading_context,
+        )
+        self._req_status[request.request_id] = req_status
+
     def get_num_new_matched_tokens(
         self, request: Request, num_computed_tokens: int
     ) -> tuple[int | None, bool]:
@@ -534,24 +552,15 @@ class OffloadingConnectorScheduler:
                 - `True` if tokens will be loaded asynchronously
                   (between scheduler steps).
         """
-        is_new_request = False
-        if req_status := self._req_status.get(request.request_id):
-            # make sure block IDs are cleared
-            for group_state in req_status.group_states:
-                group_state.block_ids.clear()
-        else:
-            is_new_request = True
-            req_status = RequestOffloadState(config=self.config, req=request)
-            self._req_status[request.request_id] = req_status
+        req_status = self._req_status[request.request_id]
+        for group_state in req_status.group_states:
+            group_state.block_ids.clear()
 
         req_status.update_offload_keys()
         req_status.num_locally_computed_tokens = num_computed_tokens
 
         num_hit_tokens = self._lookup(req_status)
-        if is_new_request:
-            req_status.update_num_hit_blocks(
-                num_computed_tokens + (num_hit_tokens or 0)
-            )
+        req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0))
 
         self._touch(req_status)
 
@@ -568,9 +577,6 @@ class OffloadingConnectorScheduler:
         num_locally_computed_tokens = req_status.num_locally_computed_tokens
         num_cached_tokens = num_locally_computed_tokens + num_external_tokens
 
-        params = req_status.req_context.kv_transfer_params
-        do_remote_decode = params is not None and params.get("do_remote_decode")
-
         keys_to_load: list[OffloadKey] = []
         dst_block_ids: list[int] = []
         # per group
@@ -624,10 +630,10 @@ class OffloadingConnectorScheduler:
             group_sizes.append(num_pending_gpu_blocks)
             block_indices.append(num_locally_computed_gpu_blocks)
 
-            if not do_remote_decode:
-                # For P/D prefill requests (do_remote_decode=True), we do
-                # NOT skip saving the hit prefix, as we need to stream the
-                # entire KV cache so a remote decode node can consume it.
+            # Skip prefix-hit blocks for block-level policy; for
+            # request-level, next_stored_block_idx stays at 0 so all
+            # blocks (including hits) are offloaded.
+            if req_status.offloading_context.policy == OffloadPolicy.BLOCK_LEVEL:
                 group_state.next_stored_block_idx = num_blocks
 
         # Fence dst blocks against finished-request pending stores.
@@ -950,6 +956,12 @@ class OffloadingConnectorScheduler:
         # TODO(orozery): possibly kickoff offload for last block
         # which may have been deferred due to async scheduling
         req_status = self._req_status.get(request.request_id)
+
+        req_context = (
+            req_status.req_context if req_status else _create_req_context(request)
+        )
+        self.manager.on_request_finished(req_context)
+
         if req_status is None:
             return False, None
         if not req_status.transfer_jobs:

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

@@ -124,6 +124,10 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
             return self.connector_worker.build_connector_worker_meta()
         return None
 
+    def on_new_request(self, request: "Request") -> None:
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.on_new_request(request)
+
     def get_num_new_matched_tokens(
         self, request: "Request", num_computed_tokens: int
     ) -> tuple[int | None, bool]:

vllm/v1/kv_offload/base.py

@@ -7,6 +7,7 @@ Core abstractions for KV cache offloading in vLLM v1.
 from abc import ABC, abstractmethod
 from collections.abc import Collection, Iterable, Iterator, Sequence
 from dataclasses import dataclass
+from enum import Enum
 from typing import TYPE_CHECKING, Any, NewType
 
 import numpy as np
@@ -49,6 +50,20 @@ class ReqContext:
     kv_transfer_params: dict[str, Any] | None = None
 
 
+class OffloadPolicy(Enum):
+    # Offload only newly-computed blocks as they arrive; prefix-hit
+    # blocks (already offloaded by a prior request) are skipped.
+    BLOCK_LEVEL = "block_level"
+    # Offload all blocks for the request, including prefix hits.
+    # Used by tiers that need the complete KV context for a request.
+    REQUEST_LEVEL = "request_level"
+
+
+@dataclass
+class RequestOffloadingContext:
+    policy: OffloadPolicy = OffloadPolicy.BLOCK_LEVEL
+
+
 class LoadStoreSpec(ABC):
     """
     Abstract metadata that encapsulates information allowing a worker
@@ -210,6 +225,28 @@ class OffloadingManager(ABC):
         """
         return
 
+    @abstractmethod
+    def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
+        """
+        Called when a new request is first seen by the scheduler.
+
+        Returns a RequestOffloadingContext indicating how this request's
+        blocks should be offloaded.
+
+        Args:
+            req_context: per-request context.
+        """
+        pass
+
+    def on_request_finished(self, req_context: ReqContext) -> None:
+        """
+        Called when a request has finished.
+
+        Args:
+            req_context: per-request context.
+        """
+        return
+
     def take_events(self) -> Iterable[OffloadingEvent]:
         """
         Take the offloading events from the manager.

vllm/v1/kv_offload/cpu/manager.py

@@ -11,6 +11,7 @@ from vllm.v1.kv_offload.base import (
     OffloadKey,
     PrepareStoreOutput,
     ReqContext,
+    RequestOffloadingContext,
 )
 from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
@@ -94,6 +95,9 @@ class CPUOffloadingManager(OffloadingManager):
 
     # --- OffloadingManager interface ---
 
+    def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
+        return RequestOffloadingContext()
+
     def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
         if self.counts is not None:
             if key in self.counts:

vllm/v1/kv_offload/tiering/base.py

@@ -11,7 +11,7 @@ from typing import TYPE_CHECKING
 
 import numpy as np
 
-from vllm.v1.kv_offload.base import OffloadKey, ReqContext
+from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext
 
 if TYPE_CHECKING:
     from vllm.v1.kv_offload.base import OffloadingSpec
@@ -163,6 +163,28 @@ class SecondaryTierManager(ABC):
         """
         return
 
+    @abstractmethod
+    def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
+        """
+        Called when a new request is first seen by the scheduler.
+
+        Returns a RequestOffloadingContext expressing this tier's preference
+        for how blocks should be offloaded for this request.
+
+        Args:
+            req_context: Per-request context.
+        """
+        pass
+
+    def on_request_finished(self, req_context: ReqContext) -> None:
+        """
+        Called when a request has finished.
+
+        Args:
+            req_context: per-request context.
+        """
+        return
+
     def shutdown(self) -> None:
         """Release resources held by this tier (threads, connections, etc.)."""
         return

vllm/v1/kv_offload/tiering/example/manager.py

@@ -13,7 +13,7 @@ import logging
 from collections.abc import Iterable
 from typing import TYPE_CHECKING
 
-from vllm.v1.kv_offload.base import OffloadKey, ReqContext
+from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext
 from vllm.v1.kv_offload.tiering.base import (
     JobMetadata,
     JobResult,
@@ -132,6 +132,9 @@ class ExampleSecondaryTierManager(SecondaryTierManager):
         self.completed_jobs = []
         return result
 
+    def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
+        return RequestOffloadingContext()
+
     def get_num_blocks(self) -> int:
         """Get the number of blocks currently stored in this tier."""
         return len(self.blocks)

vllm/v1/kv_offload/tiering/fs/manager.py

@@ -27,6 +27,7 @@ from vllm.v1.kv_offload.file_mapper import FileMapper
 from vllm.v1.kv_offload.tiering.base import (
     JobMetadata,
     JobResult,
+    RequestOffloadingContext,
     SecondaryTierManager,
 )
 from vllm.v1.kv_offload.tiering.fs.io import load_block, store_block
@@ -100,6 +101,9 @@ class FileSystemTierManager(SecondaryTierManager):
             thread_name_prefix="vllm_kv_py_fs",
         )
 
+    def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
+        return RequestOffloadingContext()
+
     def lookup(
         self, key: OffloadKey, req_context: ReqContext | None = None
     ) -> bool | None:

vllm/v1/kv_offload/tiering/manager.py

@@ -20,7 +20,8 @@ Key Design Principles:
    protecting blocks from eviction until complete_read() is called
 """
 
-from collections.abc import Collection, Iterable
+from collections import defaultdict
+from collections.abc import Collection, Iterable, Sequence
 from dataclasses import dataclass, field
 
 import numpy as np
@@ -31,8 +32,10 @@ from vllm.v1.kv_offload.base import (
     OffloadingEvent,
     OffloadingManager,
     OffloadKey,
+    OffloadPolicy,
     PrepareStoreOutput,
     ReqContext,
+    RequestOffloadingContext,
 )
 from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
@@ -158,6 +161,13 @@ class TieringOffloadingManager(OffloadingManager):
         # Reset at the end of each step in take_events().
         self._processed_jobs_this_step: bool = False
 
+        # Per-request set of secondary tiers that requested REQUEST_LEVEL
+        # policy. Populated in on_new_request(),
+        # cleaned up in on_request_finished().
+        self._request_level_tiers: defaultdict[str, set[SecondaryTierManager]] = (
+            defaultdict(set)
+        )
+
     def _next_job_id(self) -> JobId:
         """Generate a unique job ID for async transfer tracking."""
         job_id = self._job_id_counter
@@ -387,6 +397,9 @@ class TieringOffloadingManager(OffloadingManager):
         that any completed async transfers have their ref_cnt decremented
         before the primary tier makes eviction decisions.
 
+        For request-level tiers, blocks already present in the primary tier
+        are immediately cascaded via submit_store().
+
         Args:
             keys: Blocks to prepare for storing.
             req_context: Per-request context.
@@ -400,14 +413,63 @@ class TieringOffloadingManager(OffloadingManager):
         # successfully transferred to secondary tiers.
         self._maybe_process_finished_jobs()
 
-        # Step 2: Store to primary tier
+        # Step 2: Store to primary tier (new blocks only).
+        # Cascading of these newly-stored blocks to ALL secondary tiers
+        # happens later in complete_store(), after the GPU→Primary transfer
+        # completes.
         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.
+        if primary_result is None:
+            return None
+
+        # Step 3: For request-level tiers, cascade blocks already in primary
+        request_level_tiers = self._request_level_tiers.get(req_context.req_id)
+        if request_level_tiers is not None:
+            keys_to_store_set = set(primary_result.keys_to_store)
+            keys_already_in_primary = tuple(
+                k for k in keys if k not in keys_to_store_set
+            )
+            if keys_already_in_primary:
+                self._cascade_existing_blocks_to_request_level_tiers(
+                    keys_already_in_primary, req_context, request_level_tiers
+                )
 
         return primary_result
 
+    def _cascade_existing_blocks_to_request_level_tiers(
+        self,
+        keys: Sequence[OffloadKey],
+        req_context: ReqContext,
+        request_level_tiers: set[SecondaryTierManager],
+    ) -> None:
+        """
+        For tiers that requested request-level policy, submit_store() for
+        blocks that are already present in the primary tier.
+        """
+        # Filter out keys that are not ready in primary (e.g. in-flight)
+        ready_keys = tuple(
+            k for k in keys if self.primary_tier.lookup(k, req_context) is True
+        )
+        if not ready_keys:
+            return
+
+        for tier in request_level_tiers:
+            primary_blocks_spec = self.primary_tier.prepare_read(
+                ready_keys, req_context
+            )
+
+            job_id = self._next_job_id()
+            assert isinstance(primary_blocks_spec, CPULoadStoreSpec)
+            job_metadata = JobMetadata(
+                job_id=job_id,
+                keys=ready_keys,
+                block_ids=primary_blocks_spec.block_ids,
+                is_promotion=False,
+                req_context=req_context,
+            )
+            self._transfer_jobs[job_id] = job_metadata
+            tier.submit_store(job_metadata)
+
     def complete_store(
         self,
         keys: Collection[OffloadKey],
@@ -466,6 +528,31 @@ class TieringOffloadingManager(OffloadingManager):
         # Note: The async transfers are now in flight. Their completion is
         # tracked via get_finished_jobs() / _maybe_process_finished_jobs().
 
+    def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
+        """
+        Query each secondary tier for its offload policy preference.
+
+        Returns REQUEST_LEVEL if ANY secondary tier wants request-level.
+        Only stores REQUEST_LEVEL tier decisions for use in prepare_store.
+        """
+        for tier in self.secondary_tiers:
+            tier_ctx = tier.on_new_request(req_context)
+            if tier_ctx.policy == OffloadPolicy.REQUEST_LEVEL:
+                self._request_level_tiers[req_context.req_id].add(tier)
+
+        policy = (
+            OffloadPolicy.REQUEST_LEVEL
+            if req_context.req_id in self._request_level_tiers
+            else OffloadPolicy.BLOCK_LEVEL
+        )
+        return RequestOffloadingContext(policy=policy)
+
+    def on_request_finished(self, req_context: ReqContext) -> None:
+        self.primary_tier.on_request_finished(req_context)
+        for tier in self.secondary_tiers:
+            tier.on_request_finished(req_context)
+        self._request_level_tiers.pop(req_context.req_id, None)
+
     def take_events(self) -> Iterable[OffloadingEvent]:
         """
         End-of-step hook: flush deferred work, yield events, reset per-step state.