[KV Offload] Add per-request max_load_tokens control (#55885)

Signed-off-by: Alberto Perdomo <aperdomo@redhat.com>

docs/features/kv_offloading_usage.md

@@ -22,6 +22,35 @@ flowchart LR
     CPU <--> SN["..."]
 ```
 
+## Per-request load control
+
+Individual requests can cap how many tokens are loaded from offloaded storage
+by setting `max_load_tokens` in `kv_transfer_params`. The cap applies to tokens
+beyond those already available in the GPU prefix cache. Set it to `0` to
+disable external loading for the request:
+
+```json
+{
+  "kv_transfer_params": {
+    "max_load_tokens": 0
+  }
+}
+```
+
+GPU prefix-cache reuse remains enabled, and tokens beyond the aligned load cap
+are recomputed. Omitting the field leaves loading uncapped. The value must be a
+non-negative integer; invalid values are ignored. Positive caps are rounded
+down to a boundary supported by the configured KV cache groups. Offloading
+remains enabled unless it is controlled separately with `max_offload_tokens`.
+
+!!! warning
+    `max_load_tokens` is experimental and subject to change.
+
+`kv_load_tiers` continues to select secondary tiers. CPU is always included:
+it can satisfy a resident hit directly and is the required staging tier for
+secondary-to-GPU loads. Consequently, an empty tier list allows CPU hits but
+does not query any secondary tier.
+
 ## Terminology: Chunks
 
 The unit of operation is a **chunk** — a fixed-size piece of KV data covering a group of tokens. By default, a chunk maps to a single accelerator block. A configurable `blocks_per_chunk` parameter allows larger chunks, yielding larger I/Os to the host and secondary tiers.

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

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import MagicMock, call
 
 import pytest
@@ -195,15 +196,17 @@ def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler:
 
 def _make_partial_tail_request(
     scheduler: OffloadingConnectorScheduler,
+    kv_transfer_params: dict[str, Any] | None = None,
 ) -> MagicMock:
     request = MagicMock()
     request.request_id = "req"
-    request.kv_transfer_params = None
+    request.kv_transfer_params = kv_transfer_params
     request.num_prompt_tokens = 30
     request.num_tokens = 30
     request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(7)]
     request.all_token_ids = list(range(30))
     request.lora_request = None
+    request.skip_reading_prefix_cache = False
     request.is_finished.return_value = False
     scheduler.on_new_request(request)
     return request
@@ -427,6 +430,53 @@ def test_lookup_cap_stops_at_authoritative_prefix_boundary():
     assert scheduler._req_status["req"].partial_tail_boundary == 20
 
 
+def test_max_load_tokens_zero_skips_partial_tail_lookup():
+    scheduler = _make_partial_tail_scheduler()
+    request = _make_partial_tail_request(
+        scheduler, kv_transfer_params={"max_load_tokens": 0}
+    )
+    scheduler.manager.lookup.return_value = LookupResult.HIT
+
+    assert scheduler.get_num_new_matched_tokens(request, 16) == (0, False)
+    scheduler.manager.lookup.assert_not_called()
+    assert all(
+        state.num_hit_chunks == 1 for state in scheduler._req_status["req"].group_states
+    )
+
+
+def test_max_load_tokens_caps_tokens_beyond_gpu_prefix():
+    scheduler = _make_partial_tail_scheduler()
+    request = _make_partial_tail_request(
+        scheduler, kv_transfer_params={"max_load_tokens": 8}
+    )
+    scheduler.manager.lookup.return_value = LookupResult.HIT
+
+    assert scheduler.get_num_new_matched_tokens(request, 16) == (8, True)
+    assert scheduler._req_status["req"].partial_tail_boundary == 24
+
+
+def test_max_load_tokens_rounds_down_without_partial_tail():
+    scheduler = _make_partial_tail_scheduler()
+    scheduler.config = scheduler.config._replace(supports_partial_tail=False)
+    request = _make_partial_tail_request(
+        scheduler, kv_transfer_params={"max_load_tokens": 20}
+    )
+    scheduler.manager.lookup.return_value = LookupResult.HIT
+
+    assert scheduler.get_num_new_matched_tokens(request, 0) == (16, True)
+
+
+@pytest.mark.parametrize("value", ["8", 8.5, -1, True])
+def test_invalid_max_load_tokens_is_ignored(value):
+    scheduler = _make_partial_tail_scheduler()
+    request = _make_partial_tail_request(
+        scheduler, kv_transfer_params={"max_load_tokens": value}
+    )
+    scheduler.manager.lookup.return_value = LookupResult.HIT
+
+    assert scheduler.get_num_new_matched_tokens(request, 0) == (28, True)
+
+
 def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary():
     """A recurrent group may hold unhashed blocks inside the computed prefix.
 
@@ -2964,6 +3014,47 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool):
     runner.manager.lookup.assert_not_called()
 
 
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_max_load_tokens_limits_external_load(request_runner, async_scheduling: bool):
+    """The load cap limits external reads without changing the store path."""
+    block_size = 4
+    blocks_per_chunk = 3
+    tokens_per_chunk = block_size * blocks_per_chunk
+
+    runner = request_runner(
+        block_size=block_size,
+        num_gpu_blocks=100,
+        async_scheduling=async_scheduling,
+        blocks_per_chunk=blocks_per_chunk,
+    )
+
+    runner.new_request(token_ids=[0] * tokens_per_chunk)
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=(0, 1, 2),
+    )
+
+    runner.scheduler.reset_prefix_cache()
+    runner.manager.lookup.reset_mock()
+    runner.new_request(
+        token_ids=[0] * tokens_per_chunk,
+        kv_transfer_params={"max_load_tokens": 0},
+    )
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_loaded=(),
+        expected_stored=(0, 1, 2),
+    )
+
+    runner.manager.lookup.assert_not_called()
+
+
 # ---------------------------------------------------------------------------
 # Eagle/MTP test class
 # ---------------------------------------------------------------------------

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

@@ -354,6 +354,8 @@ class RequestOffloadState:
     req_context: ReqContext
     offloading_context: RequestOffloadingContext
     group_states: tuple[RequestGroupState, ...] = field(init=False)
+    # upper bound on tokens to load for this request; None means no cap
+    max_load_tokens: int | None = None
     # 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
@@ -376,6 +378,20 @@ class RequestOffloadState:
         )
         params = self.req.kv_transfer_params
 
+        # NOTE: This field is experimental and subject to change in the future.
+        raw = params.get("max_load_tokens") if params else None
+        if type(raw) is int and raw >= 0:
+            self.max_load_tokens = raw
+            logger.debug(
+                "Request %s: max_load_tokens set to %d",
+                self.req.request_id,
+                raw,
+            )
+        elif raw is not None:
+            logger.warning(
+                "max_load_tokens must be a non-negative int, got %r; ignoring", raw
+            )
+
         # 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:
@@ -767,6 +783,11 @@ class OffloadingConnectorScheduler:
             max_hit_size_tokens = min(
                 max_hit_size_tokens, num_computed_tokens + max_num_new_tokens
             )
+        if req_status.max_load_tokens is not None:
+            max_hit_size_tokens = min(
+                max_hit_size_tokens,
+                num_computed_tokens + req_status.max_load_tokens,
+            )
         if self._sliding_window_groups:
             # the last prompt token has to be recomputed to get the logprobs
             # for sliding window attention, we must reduce by 1 to make sure
@@ -810,6 +831,10 @@ class OffloadingConnectorScheduler:
                 max_hit_size_tokens = min(
                     max_hit_size_tokens, len(offload_keys) * tokens_per_chunk
                 )
+                if req_status.max_load_tokens is not None:
+                    max_hit_size_tokens = round_down(
+                        max_hit_size_tokens, tokens_per_chunk
+                    )
                 if max_hit_size_tokens - num_computed_tokens < tokens_per_chunk:
                     # We can only load less than a chunk, so skip.
                     return 0
@@ -965,6 +990,11 @@ class OffloadingConnectorScheduler:
         max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1)
         if max_num_new_tokens is not None:
             max_boundary = min(max_boundary, local_tokens + max_num_new_tokens)
+        if req_status.max_load_tokens is not None:
+            max_boundary = min(
+                max_boundary,
+                local_tokens + req_status.max_load_tokens,
+            )
         max_boundary = round_down(max_boundary, tokens_per_hash)
         if max_boundary <= complete_boundary:
             return complete_hit