diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index a5c8684b42..22a2350d80 100644 --- a/docs/features/kv_offloading_usage.md +++ b/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. 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 e7a83cb7b5..968474dd74 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/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 # --------------------------------------------------------------------------- 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 2fb3c5bb90..8cd5e90b6f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/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