A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 12:15 UTC
[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>
Itay Etelis committed
May 29, 2026 at 06:39 UTC
d63108fb18bf7c2d295c5fed3aedb64e77a46fba
4 files changed
+76
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+50
@@ -1029,6 +1029,56 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
1029
)
1030
1031
1032
+@pytest.mark.parametrize("async_scheduling", [True, False])
1033
+def test_offload_prompt_only(request_runner, async_scheduling: bool):
1034
+ """offload_prompt_only=True offloads prompt blocks but never decode blocks.
1035
+
1036
+ Setup: a 2-offloaded-block prompt followed by enough decode tokens to fill
1037
+ 4 more offloaded blocks. The flag clamps the offloadable token count to the
1038
+ prompt length, so only the prompt's blocks (GPU offsets 0-5) are ever
1039
+ eligible for store; the decode blocks (offsets >= 6) are skipped.
1040
+
1041
+ The request is intentionally not terminated (no EOS): a store is only
1042
+ flushed when a request finishes (or is preempted), so without a finish
1043
+ there is nothing to flush and the assertion stays free of flush-timing
1044
+ subtleties. The decode steps are still enough for the prompt store to
1045
+ complete and show up in expected_stored.
1046
+ """
1047
+ gpu_block_size = 4
1048
+ block_size_factor = 3
1049
+ offloaded_block_size = gpu_block_size * block_size_factor # 12
1050
+ num_prompt_blocks = 2
1051
+ num_decode_blocks = 4
1052
+ prompt_offsets = (0, 1, 2, 3, 4, 5)
1053
+
1054
+ runner = request_runner(
1055
+ block_size=gpu_block_size,
1056
+ num_gpu_blocks=100,
1057
+ async_scheduling=async_scheduling,
1058
+ block_size_factor=block_size_factor,
1059
+ extra_config_overrides={"offload_prompt_only": True},
1060
+ )
1061
+
1062
+ runner.manager.prepare_store.side_effect = (
1063
+ lambda keys, req_context: generate_store_output(keys)
1064
+ )
1065
+
1066
+ runner.new_request(token_ids=[0] * offloaded_block_size * num_prompt_blocks)
1067
+ runner.run(
1068
+ decoded_tokens=[0] * (offloaded_block_size * num_decode_blocks),
1069
+ expected_stored=prompt_offsets,
1070
+ )
1071
+
1072
+ # Timing-independent guard: only the prompt's blocks were ever offered for
1073
+ # store. If decode blocks leaked through, more keys would appear here.
1074
+ offered_keys = {
1075
+ key
1076
+ for call in runner.manager.prepare_store.call_args_list
1077
+ for key in call.args[0]
1078
+ }
1079
+ assert len(offered_keys) == num_prompt_blocks
1080
+
1081
+
1082
def test_flush_all_jobs_when_no_requests_remain(request_runner):
1083
"""When all tracked requests are finished, build_connector_meta flushes
1084
all pending jobs since there will be no future step to complete them."""
tests/v1/kv_connector/unit/offloading_connector/utils.py
+7
@@ -171,6 +171,7 @@ class RequestRunner:
171
block_size_factor: int = 1,
172
async_scheduling: bool = True,
173
kv_cache_groups: list[KVCacheGroupSpec] | None = None,
174
+ extra_config_overrides: dict[str, Any] | None = None,
175
):
176
assert block_size_factor == 1 or kv_cache_groups is None, (
177
"block_size_factor > 1 requires all groups to have the same "
@@ -194,9 +195,13 @@ class RequestRunner:
195
extra_config: dict[str, Any] = {
196
"spec_name": "MockOffloadingSpec",
197
"spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils", # noqa: E501
198
+ # Preserve legacy behavior for tests; new opt-in tests override.
199
+ "offload_prompt_only": False,
200
}
201
if block_size_factor > 1:
202
extra_config["block_size"] = block_size * block_size_factor
203
+ if extra_config_overrides:
204
+ extra_config.update(extra_config_overrides)
205
206
vllm_config.kv_transfer_config = KVTransferConfig(
207
kv_connector="OffloadingConnector",
@@ -593,6 +598,7 @@ def request_runner():
598
async_scheduling,
599
block_size_factor=1,
600
kv_cache_groups=None,
601
+ extra_config_overrides=None,
602
):
603
runner = RequestRunner(
604
block_size=block_size,
@@ -600,6 +606,7 @@ def request_runner():
606
block_size_factor=block_size_factor,
607
async_scheduling=async_scheduling,
608
kv_cache_groups=kv_cache_groups,
609
+ extra_config_overrides=extra_config_overrides,
610
)
611
runners.append(runner)
612
return runner
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+11
@@ -94,6 +94,7 @@ class SchedulerOffloadConfig(NamedTuple):
94
kv_group_configs: tuple[GroupOffloadConfig, ...]
95
block_size_factor: int
96
num_workers: int
97
+ offload_prompt_only: bool
98
99
@classmethod
100
def from_spec(cls, spec: OffloadingSpec) -> "SchedulerOffloadConfig":
@@ -156,6 +157,7 @@ class SchedulerOffloadConfig(NamedTuple):
157
for idx, gpu_block_size in enumerate(spec.gpu_block_size)
158
),
159
block_size_factor=spec.block_size_factor,
160
+ offload_prompt_only=spec.offload_prompt_only,
161
)
162
163
@@ -710,6 +712,15 @@ class OffloadingConnectorScheduler:
712
if max_offload_tokens is not None:
713
num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens)
714
715
+ # Skip decode-phase blocks: clamp to the prompt length so only
716
+ # prefill (prompt) blocks become eligible for store. next_stored_idx
717
+ # never advances past this boundary, so decode blocks are never
718
+ # queued in this or any later step.
719
+ if self.config.offload_prompt_only:
720
+ num_offloadable_tokens = min(
721
+ num_offloadable_tokens, req.num_prompt_tokens
722
+ )
723
+
724
# Filter out blocks skipped due to sliding window attention / SSM
725
# or unreachable by the load path's alignment constraints.
726
new_offload_keys: list[OffloadKey] = []
vllm/v1/kv_offload/base.py
+8
@@ -380,6 +380,14 @@ class OffloadingSpec(ABC):
380
assert kv_transfer_config is not None
381
self.extra_config = kv_transfer_config.kv_connector_extra_config
382
383
+ # When True, only prompt (prefill) blocks are offloaded; decode-phase
384
+ # blocks (KV generated after the prompt) are skipped. Useful when prior
385
+ # turns' generated tokens are dropped before the next turn (e.g.
386
+ # reasoning models that strip thinking).
387
+ self.offload_prompt_only: bool = bool(
388
+ self.extra_config.get("offload_prompt_only", True)
389
+ )
390
+
391
parallel_config = vllm_config.parallel_config
392
context_parallel_factor = (
393
parallel_config.decode_context_parallel_size