[KV offload] Add max_recompute_pending_tokens to skip waiting on short in-flight offloads

docs/features/kv_offloading_usage.md

@@ -109,6 +109,7 @@ vllm serve <model> \
 | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
 | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
 | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
+| `max_recompute_pending_tokens` | no | `0` | both | A lookup that hits chunks whose offload is still in flight (`HIT_PENDING`) normally waits for them. If waiting would gain at most this many tokens (rounded down to whole chunks), the request instead loads the ready prefix and recomputes the rest. Pending chunks inside a sliding-window group's required window and uncertain (`RETRY`) results still wait. `0` always waits. |
 | `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Partial recurrent tails emit the hash-aligned portion from the physical block start through the tail boundary. Other sliding-window/SSM chunks keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
 | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
 

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

@@ -122,6 +122,7 @@ def test_swa_offload_window_covers_unaligned_hit(
         tokens_per_hash=8,
         blocks_per_chunk=4,
         offload_prompt_only=True,
+        max_recompute_pending_tokens=0,
         kv_events_config=OffloadingKVEventsConfig(
             enable_kv_cache_events=False, self_describing_kv_events=False
         ),
@@ -736,6 +737,60 @@ def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):
     assert reduced[f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_sum"] > 0
 
 
+@pytest.mark.parametrize(
+    ("budget_blocks", "expected_loaded", "deferred"),
+    [(2, (0,), False), (1, (), True)],
+)
+def test_recompute_budget_decides_whether_pending_chunks_are_awaited(
+    request_runner, budget_blocks: int, expected_loaded: tuple[int, ...], deferred
+):
+    """With chunks [HIT, HIT_PENDING, HIT], waiting would gain two chunks. A
+    budget covering them loads the ready first chunk at once and recomputes
+    the rest; a smaller budget defers the request as before."""
+    block_size = 4
+    runner = request_runner(
+        block_size=block_size,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+        extra_config_overrides={
+            "max_recompute_pending_tokens": budget_blocks * block_size
+        },
+    )
+    token_ids = [1] * (3 * block_size + 1)
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+    runner.new_request(token_ids=token_ids)
+    runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0, 1, 2))
+
+    runner.scheduler.reset_prefix_cache()
+    results = iter([LookupResult.HIT, LookupResult.HIT_PENDING, LookupResult.HIT])
+    runner.manager.lookup.side_effect = lambda key, ctx: next(
+        results, LookupResult.MISS
+    )
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output([])
+    )
+    runner.new_request(token_ids=token_ids)
+    runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=expected_loaded)
+
+    reduced = _reduce_kv_connector_stats(runner)
+    async_delay_count = reduced.get(
+        f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_count", 0
+    )
+    assert async_delay_count == int(deferred)
+
+
+def test_negative_recompute_budget_is_rejected(request_runner):
+    with pytest.raises(ValueError, match="max_recompute_pending_tokens"):
+        request_runner(
+            block_size=4,
+            num_gpu_blocks=10,
+            async_scheduling=False,
+            extra_config_overrides={"max_recompute_pending_tokens": -1},
+        )
+
+
 def test_max_offload_tokens_zero_does_not_record_pending_lookups(request_runner):
     runner = request_runner(
         block_size=4,
@@ -1558,13 +1613,14 @@ def lookup(key, req_context):
 _LOOKUP_GROUP_CONFIG = MagicMock()
 
 
-def _maximal_lookup(sched, keys, start_chunk_idx: int = 0):
+def _maximal_lookup(sched, keys, start_chunk_idx: int = 0, max_recompute: int = 0):
     return sched._maximal_prefix_lookup(
         keys,
         _EMPTY_REQ_CTX,
         _LOOKUP_REQ,
         _LOOKUP_GROUP_CONFIG,
         start_chunk_idx,
+        max_recompute,
     )
 
 
@@ -1734,6 +1790,45 @@ def test_retry_stops_at_miss(self):
         assert sched.manager.lookup.call_count == 2
         sched._events_tracker.record_lookup.assert_not_called()
 
+    def test_short_pending_run_is_cut_before_first_pending(self):
+        """Pending chunks worth at most the recompute budget are not awaited:
+        the hit ends before the first pending chunk, ready hits after it
+        included, so the request recomputes that run instead of waiting."""
+        sched = _make_scheduler_with_lookup(
+            {
+                1: LookupResult.HIT,
+                2: LookupResult.HIT_PENDING,
+                3: LookupResult.HIT,
+                4: LookupResult.HIT_PENDING,
+            }
+        )
+        assert _maximal_lookup(sched, to_keys([1, 2, 3, 4]), 0, max_recompute=3) == 1
+
+    def test_pending_run_over_budget_defers(self):
+        """The budget counts every chunk from the first pending one to the end
+        of the hit, ready or not; one over and the lookup still defers."""
+        sched = _make_scheduler_with_lookup(
+            {
+                1: LookupResult.HIT,
+                2: LookupResult.HIT_PENDING,
+                3: LookupResult.HIT,
+                4: LookupResult.HIT_PENDING,
+            }
+        )
+        assert _maximal_lookup(sched, to_keys([1, 2, 3, 4]), 0, max_recompute=2) is None
+
+    def test_leading_pending_chunk_within_budget_is_a_miss(self):
+        sched = _make_scheduler_with_lookup({1: LookupResult.HIT_PENDING})
+        assert _maximal_lookup(sched, to_keys([1]), 0, max_recompute=1) == 0
+
+    def test_retry_defers_regardless_of_budget(self):
+        """RETRY means the chunk's location is unknown, so a recompute budget
+        must not cut the lookup short."""
+        sched = _make_scheduler_with_lookup(
+            {1: LookupResult.HIT, 2: LookupResult.RETRY}
+        )
+        assert _maximal_lookup(sched, to_keys([1, 2]), 0, max_recompute=10) is None
+
 
 class TestSlidingWindowLookup:
     def test_pending_chunk_behind_gap_does_not_defer(self):

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

@@ -184,6 +184,12 @@ class SchedulerOffloadConfig(NamedTuple):
     alignment_tokens: int | None = None
     retention_interval: int | None = None
     dcp_world_size: int = 1
+    # Wait for in-flight offloads only when they would gain more than this
+    # many tokens; otherwise recompute them. 0 always waits.
+    max_recompute_pending_tokens: int = 0
+
+    def max_recompute_pending_chunks(self, tokens_per_chunk: int) -> int:
+        return self.max_recompute_pending_tokens // tokens_per_chunk
 
     @classmethod
     def from_spec(
@@ -342,6 +348,7 @@ def from_spec(
             alignment_tokens=alignment_tokens,
             retention_interval=retention_interval,
             dcp_world_size=vllm_config.parallel_config.decode_context_parallel_size,
+            max_recompute_pending_tokens=spec.max_recompute_pending_tokens,
         )
 
 
@@ -674,11 +681,19 @@ def _maximal_prefix_lookup(
         req: Request,
         group_config: GroupOffloadConfig,
         start_chunk_idx: int,
+        max_recompute_pending_chunks: int = 0,
     ) -> int | None:
         """Return the number of consecutive offloaded chunks from the start,
-        or None if the backend deferred a lookup."""
+        or None if the backend deferred a lookup.
+
+        A HIT_PENDING chunk defers the lookup unless the run of chunks from the
+        first pending one to the end of the hit is at most
+        `max_recompute_pending_chunks` long, in which case the hit is cut
+        before it so the request recomputes that run instead of waiting.
+        """
         hit_count = 0
         defer_lookup = False
+        first_pending_idx: int | None = None
         for local_idx, key in enumerate(keys):
             result = self.manager.lookup(key, req_context)
             match result:
@@ -691,15 +706,27 @@ def _maximal_prefix_lookup(
                     )
                     hit_count += 1
                 case LookupResult.HIT_PENDING:
-                    defer_lookup = True
+                    if first_pending_idx is None:
+                        first_pending_idx = local_idx
                     hit_count += 1
                 case LookupResult.RETRY:
                     # Don't break: keep scanning to let manager kick off
                     # async lookups (until a miss is detected).
                     defer_lookup = True
                 case LookupResult.MISS:
                     break
-        return hit_count if not defer_lookup else None
+        if defer_lookup:
+            return None
+        if first_pending_idx is None:
+            return hit_count
+        if hit_count - first_pending_idx > max_recompute_pending_chunks:
+            return None
+        logger.debug(
+            "Request %s recomputing %d chunks pending offload",
+            req.request_id,
+            hit_count - first_pending_idx,
+        )
+        return first_pending_idx
 
     def _sliding_window_lookup(
         self,
@@ -850,6 +877,7 @@ def _lookup_complete_chunks(
                         req_status.req,
                         group_config,
                         start_chunk_idx,
+                        self.config.max_recompute_pending_chunks(tokens_per_chunk),
                     )
                 else:
                     required_window = sliding_window_size_in_chunks
@@ -1016,7 +1044,12 @@ def _lookup(
                 req_status.partial_tail_boundary = boundary
                 return boundary - local_tokens
 
-        if pending and complete_hit == 0:
+        if (
+            pending
+            and complete_hit == 0
+            and max_boundary - complete_boundary
+            > self.config.max_recompute_pending_tokens
+        ):
             return None
         return complete_hit
 

vllm/v1/kv_offload/base.py

@@ -606,6 +606,19 @@ def __init__(self, config: OffloadingConfig):
             self.extra_config.get("offload_prompt_only", True)
         )
 
+        # A lookup that hits chunks whose offload is still in flight waits for
+        # them only if waiting would gain more than this many tokens. Below
+        # that, the request loads the ready prefix and recomputes the rest.
+        # 0 (the default) always waits.
+        self.max_recompute_pending_tokens: int = int(
+            self.extra_config.get("max_recompute_pending_tokens", 0)
+        )
+        if self.max_recompute_pending_tokens < 0:
+            raise ValueError(
+                "max_recompute_pending_tokens must be non-negative, got "
+                f"{self.max_recompute_pending_tokens}"
+            )
+
         self.tokens_per_block = tuple(group.tokens_per_block for group in config.groups)
         self.tokens_per_hash = config.cache.tokens_per_hash
         self.blocks_per_chunk = config.cache.blocks_per_chunk