[Bugfix][KV Offload] Bound primary HIT_PENDING waits

docs/features/kv_offloading_usage.md

@@ -75,6 +75,7 @@ vllm serve <model> \
 | `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
 | `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). |
+| `hit_pending_timeout_s` | no | `5.0` | multi-tier | Maximum seconds a request waits for an in-flight CPU-primary write before treating that block as a miss and recomputing it locally. Must be positive. |
 | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
 | `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_offload/tiering/test_tiering_offloading.py

@@ -11,12 +11,14 @@
 5. Eviction coordination between tiers
 """
 
+import logging
 from collections.abc import Iterable
 from unittest.mock import MagicMock
 
 import pytest
 import torch
 
+import vllm.v1.kv_offload.tiering.manager as tiering_manager_module
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
     OffloadingConnectorStats,
 )
@@ -170,6 +172,10 @@ def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch):
         metrics[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES],
         OffloadingCounterMetadata,
     )
+    assert isinstance(
+        metrics[TieringOffloadingMetrics.HIT_PENDING_TIMEOUTS],
+        OffloadingCounterMetadata,
+    )
     for metric_name in (
         TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC,
         TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC,
@@ -218,6 +224,15 @@ def test_tiering_manager_aggregates_secondary_stats():
     assert MetricsSecondaryTierManager.MY_TIER_METRIC not in second_stats.data["data"]
 
 
+@pytest.mark.parametrize("timeout", [0.0, -1.0, float("inf"), float("nan")])
+def test_tiering_manager_rejects_invalid_hit_pending_timeout(timeout):
+    with pytest.raises(ValueError, match="must be a positive number") as exc_info:
+        TieringOffloadingManager(
+            primary_tier=MagicMock(), hit_pending_timeout_s=timeout
+        )
+    assert repr(timeout) in str(exc_info.value)
+
+
 class TestExampleSecondaryTierManager:
     """Tests for ExampleSecondaryTierManager implementation."""
 
@@ -317,6 +332,12 @@ def _start_request(self, req_context: ReqContext = _CTX):
         if req_context.req_id not in self.manager._req_state:
             self.manager.on_new_request(req_context)
 
+    @staticmethod
+    def _install_clock(monkeypatch) -> list[float]:
+        now = [100.0]
+        monkeypatch.setattr(tiering_manager_module.time, "monotonic", lambda: now[0])
+        return now
+
     def test_failed_promotion_finalizes_primary_with_failure(self, manager_setup):
         """A failed promotion still finalizes the primary slots with
         success=False; the tier corrects its own verdict."""
@@ -499,6 +520,157 @@ def test_lookup_from_primary(self, manager_setup):
         # Lookup should find all blocks in primary
         assert count_hits(self.manager, blocks) == 3
 
+    def test_primary_hit_pending_times_out_and_stays_miss(
+        self, manager_setup, monkeypatch
+    ):
+        """A stalled primary write cannot defer one request indefinitely."""
+        now = self._install_clock(monkeypatch)
+
+        ctx = ReqContext(req_id="stalled-primary")
+        self._start_request(ctx)
+        block = to_keys([10])[0]
+        self.primary_tier.lookup = MagicMock(return_value=LookupResult.HIT_PENDING)
+
+        assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING
+        now[0] += 4.999
+        assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING
+        now[0] += 0.001
+        assert self.manager.lookup(block, ctx) is LookupResult.MISS
+
+        stats = self.manager.get_stats()
+        assert stats is not None
+        assert (
+            stats.data["data"][TieringOffloadingMetrics.HIT_PENDING_TIMEOUTS][()] == 1
+        )
+        assert TieringOffloadingMetrics.BLOCK_QUERIES not in stats.data["data"]
+
+        self.primary_tier.lookup.return_value = LookupResult.HIT
+        assert self.manager.lookup(block, ctx) is LookupResult.MISS
+
+        self.manager.on_request_finished(ctx)
+        assert ctx.req_id not in self.manager._primary_pending_lookups
+
+    def test_primary_hit_pending_timeout_warning_is_rate_limited(
+        self, manager_setup, monkeypatch, caplog
+    ):
+        now = self._install_clock(monkeypatch)
+
+        ctx = ReqContext(req_id="warning-rate-limit")
+        self._start_request(ctx)
+        first_block, second_block = to_keys([10, 11])
+        self.primary_tier.lookup = MagicMock(return_value=LookupResult.HIT_PENDING)
+
+        assert self.manager.lookup(first_block, ctx) is LookupResult.HIT_PENDING
+        assert self.manager.lookup(second_block, ctx) is LookupResult.HIT_PENDING
+
+        with caplog.at_level(logging.WARNING, logger=tiering_manager_module.__name__):
+            now[0] += 5.0
+            assert self.manager.lookup(first_block, ctx) is LookupResult.MISS
+            assert self.manager.lookup(second_block, ctx) is LookupResult.MISS
+
+        messages = [
+            record.getMessage()
+            for record in caplog.records
+            if "HIT_PENDING timeout" in record.getMessage()
+        ]
+        assert len(messages) == 1
+        assert repr(first_block) in messages[0]
+
+    def test_primary_hit_pending_deadline_is_per_request(
+        self, manager_setup, monkeypatch
+    ):
+        """A later request gets its own full wait window for the same key."""
+        now = self._install_clock(monkeypatch)
+
+        block = to_keys([10])[0]
+        first_ctx = ReqContext(req_id="first-waiter")
+        second_ctx = ReqContext(req_id="second-waiter")
+        self._start_request(first_ctx)
+        self._start_request(second_ctx)
+        self.primary_tier.lookup = MagicMock(return_value=LookupResult.HIT_PENDING)
+
+        assert self.manager.lookup(block, first_ctx) is LookupResult.HIT_PENDING
+        now[0] += 4.0
+        assert self.manager.lookup(block, second_ctx) is LookupResult.HIT_PENDING
+        now[0] += 1.0
+        assert self.manager.lookup(block, first_ctx) is LookupResult.MISS
+        assert self.manager.lookup(block, second_ctx) is LookupResult.HIT_PENDING
+
+    def test_primary_hit_pending_deadline_is_per_key(self, manager_setup, monkeypatch):
+        """A later key gets its own full wait window within one request."""
+        now = self._install_clock(monkeypatch)
+
+        ctx = ReqContext(req_id="multi-key-waiter")
+        self._start_request(ctx)
+        first_block, second_block = to_keys([10, 11])
+        self.primary_tier.lookup = MagicMock(return_value=LookupResult.HIT_PENDING)
+
+        assert self.manager.lookup(first_block, ctx) is LookupResult.HIT_PENDING
+        now[0] += 4.0
+        assert self.manager.lookup(second_block, ctx) is LookupResult.HIT_PENDING
+        now[0] += 1.0
+        assert self.manager.lookup(first_block, ctx) is LookupResult.MISS
+        assert self.manager.lookup(second_block, ctx) is LookupResult.HIT_PENDING
+
+    def test_primary_hit_pending_deadline_clears_on_resolution(
+        self, manager_setup, monkeypatch
+    ):
+        """A resolved primary lookup does not retain its old deadline."""
+        now = self._install_clock(monkeypatch)
+
+        ctx = ReqContext(req_id="resolving-primary")
+        self._start_request(ctx)
+        block = to_keys([10])[0]
+        self.primary_tier.lookup = MagicMock(return_value=LookupResult.HIT_PENDING)
+
+        assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING
+        now[0] += 4.0
+        self.primary_tier.lookup.return_value = LookupResult.HIT
+        assert self.manager.lookup(block, ctx) is LookupResult.HIT
+
+        now[0] += 2.0
+        self.primary_tier.lookup.return_value = LookupResult.HIT_PENDING
+        assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING
+
+    def test_promotion_starts_hit_pending_deadline(self, manager_setup, monkeypatch):
+        """The lookup that initiates a promotion starts the wait window."""
+        now = self._install_clock(monkeypatch)
+
+        ctx = ReqContext(req_id="stalled-promotion")
+        self._start_request(ctx)
+        block = to_keys([10])[0]
+        self.secondary_tier1.blocks[block] = True
+
+        assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING
+        now[0] += 5.0
+        assert self.manager.lookup(block, ctx) is LookupResult.MISS
+
+    def test_failed_promotions_do_not_reset_hit_pending_deadline(
+        self, manager_setup, monkeypatch
+    ):
+        """Repeated failed promotions share one total deferral budget."""
+        now = self._install_clock(monkeypatch)
+
+        ctx = ReqContext(req_id="repeated-promotion-failures")
+        self._start_request(ctx)
+        block = to_keys([10])[0]
+        self.secondary_tier1.blocks[block] = True
+
+        def fail_promotion(job: TransferJob) -> None:
+            self.secondary_tier1.completed_jobs.append(
+                JobResult(job_id=job.job_id, success=False)
+            )
+
+        self.secondary_tier1.submit_load = fail_promotion
+
+        for _ in range(2):
+            assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING
+            self._simulate_on_schedule_end()
+            now[0] += 3.0
+            self._simulate_on_schedule_end()
+
+        assert self.manager.lookup(block, ctx) is LookupResult.MISS
+
     def test_promotion_from_secondary(self, manager_setup):
         """Test promotion of blocks from secondary to primary tier."""
         blocks = to_keys(range(3))
@@ -1242,6 +1414,7 @@ def test_reset_cache_clears_orchestrator_state(self, manager_setup):
             is LookupResult.HIT_PENDING
         )
         assert self.manager._pending_load_submissions
+        assert self.manager._primary_pending_lookups
 
         # Request-level tier registration.
         self.secondary_tier1.on_new_request = lambda req_context: (
@@ -1264,6 +1437,7 @@ def test_reset_cache_clears_orchestrator_state(self, manager_setup):
         # Orchestrator state cleared.
         assert self.manager._jobs == {}
         assert self.manager._pending_load_submissions == {}
+        assert self.manager._primary_pending_lookups == {}
         assert set(self.manager._req_state) == {_CTX.req_id, rl_ctx.req_id}
         assert self.manager._processed_jobs_this_step is False
 

vllm/v1/kv_offload/tiering/base.py

@@ -52,6 +52,7 @@ class TieringOffloadingMetrics:
     PROMOTION_ALLOCATION_FAILURES = (
         "vllm:kv_offload_tiering_promotion_allocation_failures"
     )
+    HIT_PENDING_TIMEOUTS = "vllm:kv_offload_tiering_hit_pending_timeouts"
     ACTIVE_PROMOTION_JOBS = "vllm:kv_offload_tiering_active_promotion_jobs"
     ACTIVE_CASCADE_JOBS = "vllm:kv_offload_tiering_active_cascade_jobs"