@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 75dc588269

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

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

alberto committed Sep 16, 2026 at 04:57 UTC 75dc5882699be65eaee144cb7d0b06c1f86d3cec

3 files changed +151 -1

docs/features/kv_offloading_usage.md

+29

@@ -22,6 +22,35 @@ flowchart LR

22 CPU <--> SN["..."]

23 ```

25 +## Per-request load control

26 +

27 +Individual requests can cap how many tokens are loaded from offloaded storage

28 +by setting `max_load_tokens` in `kv_transfer_params`. The cap applies to tokens

29 +beyond those already available in the GPU prefix cache. Set it to `0` to

30 +disable external loading for the request:

31 +

32 +```json

33 +{

34 + "kv_transfer_params": {

35 + "max_load_tokens": 0

36 + }

37 +}

38 +```

39 +

40 +GPU prefix-cache reuse remains enabled, and tokens beyond the aligned load cap

41 +are recomputed. Omitting the field leaves loading uncapped. The value must be a

42 +non-negative integer; invalid values are ignored. Positive caps are rounded

43 +down to a boundary supported by the configured KV cache groups. Offloading

44 +remains enabled unless it is controlled separately with `max_offload_tokens`.

45 +

46 +!!! warning

47 + `max_load_tokens` is experimental and subject to change.

48 +

49 +`kv_load_tiers` continues to select secondary tiers. CPU is always included:

50 +it can satisfy a resident hit directly and is the required staging tier for

51 +secondary-to-GPU loads. Consequently, an empty tier list allows CPU hits but

52 +does not query any secondary tier.

53 +

54 ## Terminology: Chunks

56 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

+92 -1

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from types import SimpleNamespace

4 +from typing import Any

5 from unittest.mock import MagicMock, call

7 import pytest

@@ -195,15 +196,17 @@ def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler:

196

197 def _make_partial_tail_request(

198 scheduler: OffloadingConnectorScheduler,

199 + kv_transfer_params: dict[str, Any] | None = None,

200 ) -> MagicMock:

201 request = MagicMock()

202 request.request_id = "req"

201 - request.kv_transfer_params = None

203 + request.kv_transfer_params = kv_transfer_params

204 request.num_prompt_tokens = 30

205 request.num_tokens = 30

206 request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(7)]

207 request.all_token_ids = list(range(30))

208 request.lora_request = None

209 + request.skip_reading_prefix_cache = False

210 request.is_finished.return_value = False

211 scheduler.on_new_request(request)

212 return request

@@ -427,6 +430,53 @@ def test_lookup_cap_stops_at_authoritative_prefix_boundary():

430 assert scheduler._req_status["req"].partial_tail_boundary == 20

431

432

433 +def test_max_load_tokens_zero_skips_partial_tail_lookup():

434 + scheduler = _make_partial_tail_scheduler()

435 + request = _make_partial_tail_request(

436 + scheduler, kv_transfer_params={"max_load_tokens": 0}

437 + )

438 + scheduler.manager.lookup.return_value = LookupResult.HIT

439 +

440 + assert scheduler.get_num_new_matched_tokens(request, 16) == (0, False)

441 + scheduler.manager.lookup.assert_not_called()

442 + assert all(

443 + state.num_hit_chunks == 1 for state in scheduler._req_status["req"].group_states

444 + )

445 +

446 +

447 +def test_max_load_tokens_caps_tokens_beyond_gpu_prefix():

448 + scheduler = _make_partial_tail_scheduler()

449 + request = _make_partial_tail_request(

450 + scheduler, kv_transfer_params={"max_load_tokens": 8}

451 + )

452 + scheduler.manager.lookup.return_value = LookupResult.HIT

453 +

454 + assert scheduler.get_num_new_matched_tokens(request, 16) == (8, True)

455 + assert scheduler._req_status["req"].partial_tail_boundary == 24

456 +

457 +

458 +def test_max_load_tokens_rounds_down_without_partial_tail():

459 + scheduler = _make_partial_tail_scheduler()

460 + scheduler.config = scheduler.config._replace(supports_partial_tail=False)

461 + request = _make_partial_tail_request(

462 + scheduler, kv_transfer_params={"max_load_tokens": 20}

463 + )

464 + scheduler.manager.lookup.return_value = LookupResult.HIT

465 +

466 + assert scheduler.get_num_new_matched_tokens(request, 0) == (16, True)

467 +

468 +

469 +@pytest.mark.parametrize("value", ["8", 8.5, -1, True])

470 +def test_invalid_max_load_tokens_is_ignored(value):

471 + scheduler = _make_partial_tail_scheduler()

472 + request = _make_partial_tail_request(

473 + scheduler, kv_transfer_params={"max_load_tokens": value}

474 + )

475 + scheduler.manager.lookup.return_value = LookupResult.HIT

476 +

477 + assert scheduler.get_num_new_matched_tokens(request, 0) == (28, True)

478 +

479 +

480 def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary():

481 """A recurrent group may hold unhashed blocks inside the computed prefix.

482

@@ -2964,6 +3014,47 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool):

3014 runner.manager.lookup.assert_not_called()

3015

3016

3017 +@pytest.mark.parametrize("async_scheduling", [True, False])

3018 +def test_max_load_tokens_limits_external_load(request_runner, async_scheduling: bool):

3019 + """The load cap limits external reads without changing the store path."""

3020 + block_size = 4

3021 + blocks_per_chunk = 3

3022 + tokens_per_chunk = block_size * blocks_per_chunk

3023 +

3024 + runner = request_runner(

3025 + block_size=block_size,

3026 + num_gpu_blocks=100,

3027 + async_scheduling=async_scheduling,

3028 + blocks_per_chunk=blocks_per_chunk,

3029 + )

3030 +

3031 + runner.new_request(token_ids=[0] * tokens_per_chunk)

3032 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

3033 + generate_store_output(keys)

3034 + )

3035 + runner.run(

3036 + decoded_tokens=[EOS_TOKEN_ID],

3037 + expected_stored=(0, 1, 2),

3038 + )

3039 +

3040 + runner.scheduler.reset_prefix_cache()

3041 + runner.manager.lookup.reset_mock()

3042 + runner.new_request(

3043 + token_ids=[0] * tokens_per_chunk,

3044 + kv_transfer_params={"max_load_tokens": 0},

3045 + )

3046 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

3047 + generate_store_output(keys)

3048 + )

3049 + runner.run(

3050 + decoded_tokens=[EOS_TOKEN_ID],

3051 + expected_loaded=(),

3052 + expected_stored=(0, 1, 2),

3053 + )

3054 +

3055 + runner.manager.lookup.assert_not_called()

3056 +

3057 +

3058 # ---------------------------------------------------------------------------

3059 # Eagle/MTP test class

3060 # ---------------------------------------------------------------------------

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

+30

@@ -354,6 +354,8 @@ class RequestOffloadState:

354 req_context: ReqContext

355 offloading_context: RequestOffloadingContext

356 group_states: tuple[RequestGroupState, ...] = field(init=False)

357 + # upper bound on tokens to load for this request; None means no cap

358 + max_load_tokens: int | None = None

359 # upper bound on tokens to offload for this request; None means no cap

360 max_offload_tokens: int | None = None

361 # number of hits in the GPU cache

@@ -376,6 +378,20 @@ class RequestOffloadState:

378 )

379 params = self.req.kv_transfer_params

380

381 + # NOTE: This field is experimental and subject to change in the future.

382 + raw = params.get("max_load_tokens") if params else None

383 + if type(raw) is int and raw >= 0:

384 + self.max_load_tokens = raw

385 + logger.debug(

386 + "Request %s: max_load_tokens set to %d",

387 + self.req.request_id,

388 + raw,

389 + )

390 + elif raw is not None:

391 + logger.warning(

392 + "max_load_tokens must be a non-negative int, got %r; ignoring", raw

393 + )

394 +

395 # NOTE: This field is experimental and subject to change in the future.

396 raw = params.get("max_offload_tokens") if params else None

397 if type(raw) is int and raw >= 0:

@@ -767,6 +783,11 @@ class OffloadingConnectorScheduler:

783 max_hit_size_tokens = min(

784 max_hit_size_tokens, num_computed_tokens + max_num_new_tokens

785 )

786 + if req_status.max_load_tokens is not None:

787 + max_hit_size_tokens = min(

788 + max_hit_size_tokens,

789 + num_computed_tokens + req_status.max_load_tokens,

790 + )

791 if self._sliding_window_groups:

792 # the last prompt token has to be recomputed to get the logprobs

793 # for sliding window attention, we must reduce by 1 to make sure

@@ -810,6 +831,10 @@ class OffloadingConnectorScheduler:

831 max_hit_size_tokens = min(

832 max_hit_size_tokens, len(offload_keys) * tokens_per_chunk

833 )

834 + if req_status.max_load_tokens is not None:

835 + max_hit_size_tokens = round_down(

836 + max_hit_size_tokens, tokens_per_chunk

837 + )

838 if max_hit_size_tokens - num_computed_tokens < tokens_per_chunk:

839 # We can only load less than a chunk, so skip.

840 return 0

@@ -965,6 +990,11 @@ class OffloadingConnectorScheduler:

990 max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1)

991 if max_num_new_tokens is not None:

992 max_boundary = min(max_boundary, local_tokens + max_num_new_tokens)

993 + if req_status.max_load_tokens is not None:

994 + max_boundary = min(

995 + max_boundary,

996 + local_tokens + req_status.max_load_tokens,

997 + )

998 max_boundary = round_down(max_boundary, tokens_per_hash)

999 if max_boundary <= complete_boundary:

1000 return complete_hit