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>

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

@@ -936,6 +936,99 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
     assert runner.manager.complete_store.call_count == 0
 
 
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
+    """Validates max_offload_tokens: type coercion, boundary values, and capping.
+
+    Setup: 3 offloaded blocks × 3 GPU blocks each = 9 GPU block offsets (0–8).
+    """
+    gpu_block_size = 4
+    block_size_factor = 3
+    offloaded_block_size = gpu_block_size * block_size_factor  # 12
+    num_gpu_blocks = 100
+    all_offsets = (0, 1, 2, 3, 4, 5, 6, 7, 8)
+
+    def make_runner():
+        return request_runner(
+            block_size=gpu_block_size,
+            num_gpu_blocks=num_gpu_blocks,
+            async_scheduling=async_scheduling,
+            block_size_factor=block_size_factor,
+        )
+
+    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.manager.prepare_store.side_effect = (
+            lambda keys, req_context: generate_store_output(keys)
+        )
+
+    # With sync scheduling, the connector flushes completed stores when the
+    # request finishes; async scheduling defers the flush to the next step.
+    flushed_all = all_offsets if not async_scheduling else ()
+    flushed_two = (0, 1, 2, 3, 4, 5) if not async_scheduling else ()
+
+    # None -> no cap, all 9 offsets stored
+    r = make_runner()
+    setup(r, None)
+    r.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=all_offsets,
+        expected_flushed=flushed_all,
+    )
+
+    # str -> warn and fall back to no cap
+    r = make_runner()
+    setup(r, "24")
+    r.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=all_offsets,
+        expected_flushed=flushed_all,
+    )
+
+    # float -> warn and fall back to no cap
+    r = make_runner()
+    setup(r, 24.5)
+    r.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=all_offsets,
+        expected_flushed=flushed_all,
+    )
+
+    # negative -> warn and fall back to no cap
+    r = make_runner()
+    setup(r, -1)
+    r.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=all_offsets,
+        expected_flushed=flushed_all,
+    )
+
+    # bool -> rejected (type(True) is bool, not int), falls back to no cap
+    r = make_runner()
+    setup(r, True)
+    r.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=all_offsets,
+        expected_flushed=flushed_all,
+    )
+
+    # zero -> valid, no blocks offloaded
+    r = make_runner()
+    setup(r, 0)
+    r.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=())
+
+    # positive int cap -> limits offload to first 2 offloaded blocks (offsets 0–5)
+    r = make_runner()
+    setup(r, 24)  # 24 tokens = 2 offloaded blocks × 12 tokens each
+    r.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=(0, 1, 2, 3, 4, 5),
+        expected_flushed=flushed_two,
+    )
+
+
 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."""

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

@@ -174,6 +174,8 @@ class RequestOffloadState:
     req: Request
     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
     num_locally_computed_tokens: int = 0
     # In-flight job IDs. Per the connector's invariant, at any given time
@@ -188,6 +190,21 @@ class RequestOffloadState:
             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.
+        raw = params.get("max_offload_tokens") if params else None
+        if type(raw) is int and raw >= 0:
+            self.max_offload_tokens = raw
+            logger.debug(
+                "Request %s: max_offload_tokens set to %d",
+                self.req.request_id,
+                raw,
+            )
+        elif raw is not None:
+            logger.warning(
+                "max_offload_tokens must be a non-negative int, got %r; ignoring", raw
+            )
 
     def update_offload_keys(self) -> None:
         for group_config, group_state in zip(
@@ -683,6 +700,9 @@ class OffloadingConnectorScheduler:
             num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens
             # with async scheduling, some tokens may be missing
             num_offloadable_tokens = min(num_tokens_after_batch, req.num_tokens)
+            max_offload_tokens = req_status.max_offload_tokens
+            if max_offload_tokens is not None:
+                num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens)
 
             # Filter out blocks skipped due to sliding window attention / SSM
             # or unreachable by the load path's alignment constraints.