[KV Offloading] Add basic offloading metrics (vllm-project#45958)

Signed-off-by: srinivas_oo7 <sklinkedin0120@gmail.com>
Signed-off-by: Srinivas Krovvidi <194645829+Srinivasoo7@users.noreply.github.com>
Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com>
Co-authored-by: Or Ozeri <or@ozery.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>

Author: 3 people

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

@@ -10,9 +10,11 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
     OffloadingConnectorStats,
     OffloadPromMetrics,
+    _ConnectorMetricName,
     _MetricType,
     _StatsKey,
     _TransferMetricName,
+    get_connector_metric_definitions,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
     OffloadingConnector,
@@ -37,6 +39,42 @@
 MY_LABEL = "my_label"
 
 
+def test_connector_metric_histogram_buckets():
+    metadata = get_connector_metric_definitions()
+
+    sync_delay = metadata[_ConnectorMetricName.LOOKUP_SYNC_DELAY]
+    assert isinstance(sync_delay, OffloadingHistogramMetadata)
+    assert sync_delay.buckets == (
+        0.00001,
+        0.00005,
+        0.0001,
+        0.0005,
+        0.001,
+        0.005,
+        0.01,
+        0.05,
+        0.1,
+        0.5,
+        1,
+    )
+
+    async_delay = metadata[_ConnectorMetricName.LOOKUP_ASYNC_DELAY]
+    assert isinstance(async_delay, OffloadingHistogramMetadata)
+    assert async_delay.buckets == (
+        0.0001,
+        0.0005,
+        0.001,
+        0.005,
+        0.01,
+        0.05,
+        0.1,
+        0.5,
+        1,
+        5,
+        10,
+    )
+
+
 class _FakeMetric:
     def __init__(self, **kwargs):
         self.kwargs = kwargs

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

@@ -11,6 +11,10 @@
     to_keys,
 )
 from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
+    OffloadingConnectorStats,
+    _ConnectorMetricName,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
     OffloadingConnectorScheduler,
     RequestOffloadState,
@@ -32,6 +36,72 @@
 from vllm.v1.request import RequestStatus
 
 
+def _reduce_kv_connector_stats(runner):
+    reduced: dict[str, int | float] = {}
+    for payload in runner.kv_connector_stats:
+        stats = (
+            payload
+            if hasattr(payload, "reduce")
+            else OffloadingConnectorStats(data=payload)
+        )
+        for key, value in stats.reduce().items():
+            reduced[key] = reduced.get(key, 0) + value
+    return reduced
+
+
+def test_scheduler_reports_allocation_failure(request_runner):
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+    )
+    runner.new_request(token_ids=[0] * 4)
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: None
+
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+
+    reduced = _reduce_kv_connector_stats(runner)
+    assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 1
+
+
+def test_scheduler_reports_lookup_sync_delay(request_runner):
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+    )
+    runner.new_request(token_ids=[1] * 4)
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output([])
+    )
+
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+
+    reduced = _reduce_kv_connector_stats(runner)
+    assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_count"] == 1
+    assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_sum"] > 0
+
+
+def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):
+    """A deferred lookup reports its async delay once it resolves."""
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+    )
+    runner.manager.lookup.side_effect = [LookupResult.RETRY, LookupResult.MISS]
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output([])
+    )
+
+    runner.new_request(token_ids=[1] * 4)
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+
+    reduced = _reduce_kv_connector_stats(runner)
+    assert reduced[f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_count"] == 1
+    assert reduced[f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_sum"] > 0
+
+
 @pytest.mark.parametrize("async_scheduling", [True, False])
 def test_offloading_connector(request_runner, async_scheduling: bool):
     block_size = 4

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

@@ -133,6 +133,7 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
         self.manager = MagicMock(spec=OffloadingManager)
         self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)
         self.manager.lookup.return_value = LookupResult.MISS
+        self.manager.get_stats.return_value = None
         self.manager.on_new_request.return_value = RequestOffloadingContext()
         self.handler = MockOffloadingWorker()
 
@@ -335,6 +336,7 @@ def __init__(
         self.completed_loads: list[TransferSummary] = []
         self.completed_stores: list[TransferSummary] = []
         self.flushed_gpu_blocks: set[GPUBlock] = set()
+        self.kv_connector_stats: list[Any] = []
 
         # block_id -> GPUBlock
         self.gpu_blocks: dict[int, GPUBlock] = {}
@@ -348,6 +350,12 @@ def __init__(
             slot_mapping={},
         )
 
+    def _record_kv_connector_stats(self, engine_outputs: dict[int, Any]) -> None:
+        for output in engine_outputs.values():
+            scheduler_stats = output.scheduler_stats
+            if scheduler_stats is not None and scheduler_stats.kv_connector_stats:
+                self.kv_connector_stats.append(scheduler_stats.kv_connector_stats)
+
     def new_request(
         self,
         token_ids: list[int],
@@ -524,13 +532,17 @@ def _run(
             if self.async_scheduling:
                 # in async scheduling we update the output of the previous step
                 if prev_model_runner_output is not None:
-                    self.scheduler.update_from_output(
+                    engine_outputs = self.scheduler.update_from_output(
                         prev_scheduler_output, prev_model_runner_output
                     )
+                    self._record_kv_connector_stats(engine_outputs)
                 prev_scheduler_output = scheduler_output
                 prev_model_runner_output = model_runner_output
             else:
-                self.scheduler.update_from_output(scheduler_output, model_runner_output)
+                engine_outputs = self.scheduler.update_from_output(
+                    scheduler_output, model_runner_output
+                )
+                self._record_kv_connector_stats(engine_outputs)
 
             if post_step_fn is not None:
                 post_step_fn()
@@ -546,9 +558,10 @@ def _run(
             if token_id is None:
                 if self.async_scheduling:
                     # sample last token
-                    self.scheduler.update_from_output(
+                    engine_outputs = self.scheduler.update_from_output(
                         prev_scheduler_output, prev_model_runner_output
                     )
+                    self._record_kv_connector_stats(engine_outputs)
                 break
 
         self._parse_transfers()

tests/v1/kv_offload/cpu/test_manager.py

@@ -224,6 +224,64 @@ def check_usage_stats(manager: CPUOffloadingManager, value: float):
     check_usage_stats(manager, 0.0)
 
 
+def test_cpu_manager_reports_allocation_size_histogram():
+    manager = make_cpu_manager(num_blocks=4, cache_policy="lru")
+
+    manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+    manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+    manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX)
+
+    stats = manager.get_stats()
+
+    assert stats is not None
+    reduced = stats.reduce()
+    assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count"] == 2
+    assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 3
+
+    # The cache-usage gauge is always reported, so get_stats() never returns
+    # None, but the histogram has nothing new once its samples are consumed.
+    second_stats = manager.get_stats()
+    assert second_stats is not None
+    assert f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count" not in (
+        second_stats.reduce()
+    )
+
+
+def test_cpu_manager_reports_allocation_size_on_allocation_failure(monkeypatch):
+    manager = make_cpu_manager(num_blocks=4, cache_policy="lru")
+
+    def fail_allocate_blocks(keys):
+        raise RuntimeError("allocation failed")
+
+    monkeypatch.setattr(manager, "_allocate_blocks", fail_allocate_blocks)
+
+    with pytest.raises(RuntimeError, match="allocation failed"):
+        manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX)
+
+    stats = manager.get_stats()
+
+    assert stats is not None
+    reduced = stats.reduce()
+    assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count"] == 1
+    assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 3
+
+
+def test_cpu_manager_reports_allocation_size_on_eviction_failure():
+    manager = make_cpu_manager(num_blocks=1, cache_policy="lru")
+
+    manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX)
+    manager.get_stats()
+
+    assert manager.prepare_store(to_keys([2]), _EMPTY_REQ_CTX) is None
+
+    stats = manager.get_stats()
+
+    assert stats is not None
+    reduced = stats.reduce()
+    assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count"] == 1
+    assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 1
+
+
 def test_cpu_manager():
     """
     Tests CPUOffloadingManager with lru policy.

tests/v1/kv_offload/test_factory.py

@@ -21,7 +21,7 @@
     KVCacheGroupSpec,
     KVCacheTensor,
 )
-from vllm.v1.kv_offload.base import OffloadingSpec
+from vllm.v1.kv_offload.base import OffloadingHistogramMetadata, OffloadingSpec
 from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
 from vllm.v1.kv_offload.factory import OffloadingSpecFactory
 from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
@@ -248,8 +248,8 @@ def test_duplicate_registration_raises():
 # ---------------------------------------------------------------------------
 
 
-def test_build_metric_definitions_empty_below_threshold():
-    """store_threshold < 2 → only base metric (no stores_skipped)."""
+def test_build_metric_definitions_below_threshold():
+    """store_threshold < 2 keeps stores_skipped disabled."""
     from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
 
     config = _make_vllm_config(store_threshold=1)
@@ -258,6 +258,32 @@ def test_build_metric_definitions_empty_below_threshold():
         config.kv_transfer_config.kv_connector_extra_config
     )
     assert CPUOffloadingMetrics.STORES_SKIPPED not in metrics
+    assert CPUOffloadingMetrics.CPU_ALLOCATION_SIZE in metrics
+
+
+def test_build_metric_definitions_allocation_size_histogram():
+    """CPU allocation size is always reported as a histogram."""
+    from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
+
+    config = _make_vllm_config(store_threshold=0)
+    spec_cls = OffloadingSpecFactory.get_spec_cls(config)
+    metrics = spec_cls.build_metric_definitions(
+        config.kv_transfer_config.kv_connector_extra_config
+    )
+    metadata = metrics[CPUOffloadingMetrics.CPU_ALLOCATION_SIZE]
+    assert isinstance(metadata, OffloadingHistogramMetadata)
+    assert metadata.buckets == (
+        1,
+        4,
+        16,
+        64,
+        256,
+        1024,
+        4096,
+        16384,
+        65536,
+        262144,
+    )
 
 
 def test_build_metric_definitions_returns_counter_at_threshold():

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

@@ -30,6 +30,14 @@ class _TransferMetricName:
     STORE_SIZE = "vllm:kv_offload_store_size"
 
 
+class _ConnectorMetricName:
+    """Connector-side metrics emitted by scheduler-side offloading code."""
+
+    LOOKUP_SYNC_DELAY = "vllm:kv_offload_lookup_sync_delay_seconds"
+    LOOKUP_ASYNC_DELAY = "vllm:kv_offload_lookup_async_delay_seconds"
+    ALLOCATION_FAILURE = "vllm:kv_offload_allocation_failure"
+
+
 class _TransferType:
     """Transfer direction labels for deprecated CPU offload metrics."""
 
@@ -74,6 +82,50 @@ def get_connector_metric_definitions() -> dict[str, OffloadingMetricMetadata]:
             documentation="Histogram of KV offload store operation size, in bytes.",
             buckets=TRANSFER_SIZE_BUCKETS,
         ),
+        _ConnectorMetricName.LOOKUP_SYNC_DELAY: OffloadingHistogramMetadata(
+            documentation=(
+                "Histogram of the time spent in a single offload lookup call, "
+                "in seconds."
+            ),
+            buckets=(
+                0.00001,
+                0.00005,
+                0.0001,
+                0.0005,
+                0.001,
+                0.005,
+                0.01,
+                0.05,
+                0.1,
+                0.5,
+                1,
+            ),
+        ),
+        _ConnectorMetricName.LOOKUP_ASYNC_DELAY: OffloadingHistogramMetadata(
+            documentation=(
+                "Histogram of time between a request's offload lookup first "
+                "deferring and the following lookup resolving, or request "
+                "finish, in seconds."
+            ),
+            buckets=(
+                0.0001,
+                0.0005,
+                0.001,
+                0.005,
+                0.01,
+                0.05,
+                0.1,
+                0.5,
+                1,
+                5,
+                10,
+            ),
+        ),
+        _ConnectorMetricName.ALLOCATION_FAILURE: OffloadingCounterMetadata(
+            documentation=(
+                "Number of KV offload store allocation attempts that failed."
+            ),
+        ),
     }
 
 
@@ -223,7 +275,7 @@ def is_empty(self) -> bool:
     def increase_counter(
         self,
         counter_name: str,
-        counter_increase_value: int | float,
+        counter_increase_value: int | float = 1,
         labelvalues: tuple[str, ...] = (),
     ) -> None:
         """Increase a counter on the stats payload."""