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, 13:30 UTC
Add token-offset based selective offload in OffloadConnector (#39983)
Signed-off-by: Angelo Ruocco <ang@zurich.ibm.com> Co-authored-by: Or Ozeri <or@ozery.com>
Angelo Ruocco committed
May 28, 2026 at 16:11 UTC
864990e8d9b3a5e058f5f77c574146963dd2df0f
2 files changed
+113
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+93
@@ -936,6 +936,99 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
936
assert runner.manager.complete_store.call_count == 0
937
938
939
+@pytest.mark.parametrize("async_scheduling", [True, False])
940
+def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
941
+ """Validates max_offload_tokens: type coercion, boundary values, and capping.
942
+
943
+ Setup: 3 offloaded blocks × 3 GPU blocks each = 9 GPU block offsets (0–8).
944
+ """
945
+ gpu_block_size = 4
946
+ block_size_factor = 3
947
+ offloaded_block_size = gpu_block_size * block_size_factor # 12
948
+ num_gpu_blocks = 100
949
+ all_offsets = (0, 1, 2, 3, 4, 5, 6, 7, 8)
950
+
951
+ def make_runner():
952
+ return request_runner(
953
+ block_size=gpu_block_size,
954
+ num_gpu_blocks=num_gpu_blocks,
955
+ async_scheduling=async_scheduling,
956
+ block_size_factor=block_size_factor,
957
+ )
958
+
959
+ def setup(r, max_offload_tokens):
960
+ r.new_request(token_ids=[0] * offloaded_block_size * 3)
961
+ req = r.scheduler.requests[str(r.req_id)]
962
+ req.kv_transfer_params = {"max_offload_tokens": max_offload_tokens}
963
+ r.manager.prepare_store.side_effect = (
964
+ lambda keys, req_context: generate_store_output(keys)
965
+ )
966
+
967
+ # With sync scheduling, the connector flushes completed stores when the
968
+ # request finishes; async scheduling defers the flush to the next step.
969
+ flushed_all = all_offsets if not async_scheduling else ()
970
+ flushed_two = (0, 1, 2, 3, 4, 5) if not async_scheduling else ()
971
+
972
+ # None -> no cap, all 9 offsets stored
973
+ r = make_runner()
974
+ setup(r, None)
975
+ r.run(
976
+ decoded_tokens=[EOS_TOKEN_ID],
977
+ expected_stored=all_offsets,
978
+ expected_flushed=flushed_all,
979
+ )
980
+
981
+ # str -> warn and fall back to no cap
982
+ r = make_runner()
983
+ setup(r, "24")
984
+ r.run(
985
+ decoded_tokens=[EOS_TOKEN_ID],
986
+ expected_stored=all_offsets,
987
+ expected_flushed=flushed_all,
988
+ )
989
+
990
+ # float -> warn and fall back to no cap
991
+ r = make_runner()
992
+ setup(r, 24.5)
993
+ r.run(
994
+ decoded_tokens=[EOS_TOKEN_ID],
995
+ expected_stored=all_offsets,
996
+ expected_flushed=flushed_all,
997
+ )
998
+
999
+ # negative -> warn and fall back to no cap
1000
+ r = make_runner()
1001
+ setup(r, -1)
1002
+ r.run(
1003
+ decoded_tokens=[EOS_TOKEN_ID],
1004
+ expected_stored=all_offsets,
1005
+ expected_flushed=flushed_all,
1006
+ )
1007
+
1008
+ # bool -> rejected (type(True) is bool, not int), falls back to no cap
1009
+ r = make_runner()
1010
+ setup(r, True)
1011
+ r.run(
1012
+ decoded_tokens=[EOS_TOKEN_ID],
1013
+ expected_stored=all_offsets,
1014
+ expected_flushed=flushed_all,
1015
+ )
1016
+
1017
+ # zero -> valid, no blocks offloaded
1018
+ r = make_runner()
1019
+ setup(r, 0)
1020
+ r.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=())
1021
+
1022
+ # positive int cap -> limits offload to first 2 offloaded blocks (offsets 0–5)
1023
+ r = make_runner()
1024
+ setup(r, 24) # 24 tokens = 2 offloaded blocks × 12 tokens each
1025
+ r.run(
1026
+ decoded_tokens=[EOS_TOKEN_ID],
1027
+ expected_stored=(0, 1, 2, 3, 4, 5),
1028
+ expected_flushed=flushed_two,
1029
+ )
1030
+
1031
+
1032
def test_flush_all_jobs_when_no_requests_remain(request_runner):
1033
"""When all tracked requests are finished, build_connector_meta flushes
1034
all pending jobs since there will be no future step to complete them."""
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+20
@@ -174,6 +174,8 @@ class RequestOffloadState:
174
req: Request
175
group_states: tuple[RequestGroupState, ...] = field(init=False)
176
req_context: ReqContext = field(init=False)
177
+ # upper bound on tokens to offload for this request; None means no cap
178
+ max_offload_tokens: int | None = None
179
# number of hits in the GPU cache
180
num_locally_computed_tokens: int = 0
181
# In-flight job IDs. Per the connector's invariant, at any given time
@@ -188,6 +190,21 @@ class RequestOffloadState:
190
req_id=self.req.request_id,
191
kv_transfer_params=self.req.kv_transfer_params,
192
)
193
+ params = self.req.kv_transfer_params
194
+
195
+ # NOTE: This field is experimental and subject to change in the future.
196
+ raw = params.get("max_offload_tokens") if params else None
197
+ if type(raw) is int and raw >= 0:
198
+ self.max_offload_tokens = raw
199
+ logger.debug(
200
+ "Request %s: max_offload_tokens set to %d",
201
+ self.req.request_id,
202
+ raw,
203
+ )
204
+ elif raw is not None:
205
+ logger.warning(
206
+ "max_offload_tokens must be a non-negative int, got %r; ignoring", raw
207
+ )
208
209
def update_offload_keys(self) -> None:
210
for group_config, group_state in zip(
@@ -683,6 +700,9 @@ class OffloadingConnectorScheduler:
700
num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens
701
# with async scheduling, some tokens may be missing
702
num_offloadable_tokens = min(num_tokens_after_batch, req.num_tokens)
703
+ max_offload_tokens = req_status.max_offload_tokens
704
+ if max_offload_tokens is not None:
705
+ num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens)
706
707
# Filter out blocks skipped due to sliding window attention / SSM
708
# or unreachable by the load path's alignment constraints.