From d63108fb18bf7c2d295c5fed3aedb64e77a46fba Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Fri, 29 May 2026 06:39:43 +0300 Subject: [PATCH] [kv_offload] Skip decode-phase blocks in CPU offload (#43797) Signed-off-by: Itay Etelis Co-authored-by: Itay Etelis --- .../offloading_connector/test_scheduler.py | 50 +++++++++++++++++++ .../unit/offloading_connector/utils.py | 7 +++ .../kv_connector/v1/offloading/scheduler.py | 11 ++++ vllm/v1/kv_offload/base.py | 8 +++ 4 files changed, 76 insertions(+) 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 ff30a6f42fa..5cc7065c1a4 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/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.""" diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 4248a89bf6c..2dc1071a711 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/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 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 e336ed94ffe..6560bcf9def 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/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] = [] diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index d932fd2d1be..ec168de8917 100644 --- a/vllm/v1/kv_offload/base.py +++ b/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