[kv_offload] Skip decode-phase blocks in CPU offload (#43797)

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
Co-authored-by: Itay Etelis <itay.etelis@ibm.com>

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

@@ -1029,6 +1029,56 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
     )
 
 
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_offload_prompt_only(request_runner, async_scheduling: bool):
+    """offload_prompt_only=True offloads prompt blocks but never decode blocks.
+
+    Setup: a 2-offloaded-block prompt followed by enough decode tokens to fill
+    4 more offloaded blocks. The flag clamps the offloadable token count to the
+    prompt length, so only the prompt's blocks (GPU offsets 0-5) are ever
+    eligible for store; the decode blocks (offsets >= 6) are skipped.
+
+    The request is intentionally not terminated (no EOS): a store is only
+    flushed when a request finishes (or is preempted), so without a finish
+    there is nothing to flush and the assertion stays free of flush-timing
+    subtleties. The decode steps are still enough for the prompt store to
+    complete and show up in expected_stored.
+    """
+    gpu_block_size = 4
+    block_size_factor = 3
+    offloaded_block_size = gpu_block_size * block_size_factor  # 12
+    num_prompt_blocks = 2
+    num_decode_blocks = 4
+    prompt_offsets = (0, 1, 2, 3, 4, 5)
+
+    runner = request_runner(
+        block_size=gpu_block_size,
+        num_gpu_blocks=100,
+        async_scheduling=async_scheduling,
+        block_size_factor=block_size_factor,
+        extra_config_overrides={"offload_prompt_only": True},
+    )
+
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output(keys)
+    )
+
+    runner.new_request(token_ids=[0] * offloaded_block_size * num_prompt_blocks)
+    runner.run(
+        decoded_tokens=[0] * (offloaded_block_size * num_decode_blocks),
+        expected_stored=prompt_offsets,
+    )
+
+    # Timing-independent guard: only the prompt's blocks were ever offered for
+    # store. If decode blocks leaked through, more keys would appear here.
+    offered_keys = {
+        key
+        for call in runner.manager.prepare_store.call_args_list
+        for key in call.args[0]
+    }
+    assert len(offered_keys) == num_prompt_blocks
+
+
 def test_flush_all_jobs_when_no_requests_remain(request_runner):
     """When all tracked requests are finished, build_connector_meta flushes
     all pending jobs since there will be no future step to complete them."""

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

@@ -171,6 +171,7 @@ class RequestRunner:
         block_size_factor: int = 1,
         async_scheduling: bool = True,
         kv_cache_groups: list[KVCacheGroupSpec] | None = None,
+        extra_config_overrides: dict[str, Any] | None = None,
     ):
         assert block_size_factor == 1 or kv_cache_groups is None, (
             "block_size_factor > 1 requires all groups to have the same "
@@ -194,9 +195,13 @@ class RequestRunner:
         extra_config: dict[str, Any] = {
             "spec_name": "MockOffloadingSpec",
             "spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils",  # noqa: E501
+            # Preserve legacy behavior for tests; new opt-in tests override.
+            "offload_prompt_only": False,
         }
         if block_size_factor > 1:
             extra_config["block_size"] = block_size * block_size_factor
+        if extra_config_overrides:
+            extra_config.update(extra_config_overrides)
 
         vllm_config.kv_transfer_config = KVTransferConfig(
             kv_connector="OffloadingConnector",
@@ -593,6 +598,7 @@ def request_runner():
         async_scheduling,
         block_size_factor=1,
         kv_cache_groups=None,
+        extra_config_overrides=None,
     ):
         runner = RequestRunner(
             block_size=block_size,
@@ -600,6 +606,7 @@ def request_runner():
             block_size_factor=block_size_factor,
             async_scheduling=async_scheduling,
             kv_cache_groups=kv_cache_groups,
+            extra_config_overrides=extra_config_overrides,
         )
         runners.append(runner)
         return runner

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

@@ -94,6 +94,7 @@ class SchedulerOffloadConfig(NamedTuple):
     kv_group_configs: tuple[GroupOffloadConfig, ...]
     block_size_factor: int
     num_workers: int
+    offload_prompt_only: bool
 
     @classmethod
     def from_spec(cls, spec: OffloadingSpec) -> "SchedulerOffloadConfig":
@@ -156,6 +157,7 @@ class SchedulerOffloadConfig(NamedTuple):
                 for idx, gpu_block_size in enumerate(spec.gpu_block_size)
             ),
             block_size_factor=spec.block_size_factor,
+            offload_prompt_only=spec.offload_prompt_only,
         )
 
 
@@ -710,6 +712,15 @@ class OffloadingConnectorScheduler:
             if max_offload_tokens is not None:
                 num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens)
 
+            # Skip decode-phase blocks: clamp to the prompt length so only
+            # prefill (prompt) blocks become eligible for store. next_stored_idx
+            # never advances past this boundary, so decode blocks are never
+            # queued in this or any later step.
+            if self.config.offload_prompt_only:
+                num_offloadable_tokens = min(
+                    num_offloadable_tokens, req.num_prompt_tokens
+                )
+
             # Filter out blocks skipped due to sliding window attention / SSM
             # or unreachable by the load path's alignment constraints.
             new_offload_keys: list[OffloadKey] = []

vllm/v1/kv_offload/base.py

@@ -380,6 +380,14 @@ class OffloadingSpec(ABC):
         assert kv_transfer_config is not None
         self.extra_config = kv_transfer_config.kv_connector_extra_config
 
+        # When True, only prompt (prefill) blocks are offloaded; decode-phase
+        # blocks (KV generated after the prompt) are skipped. Useful when prior
+        # turns' generated tokens are dropped before the next turn (e.g.
+        # reasoning models that strip thinking).
+        self.offload_prompt_only: bool = bool(
+            self.extra_config.get("offload_prompt_only", True)
+        )
+
         parallel_config = vllm_config.parallel_config
         context_parallel_factor = (
             parallel_config.decode_context_parallel_size