[KV Offloading] Add labeled metrics support (#45957)

Signed-off-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com>
2026-06-21 18:04:01 +00:00
parent 745bba5ea8
commit c441ad1c07
3 changed files with 384 additions and 123 deletions
@@ -22,7 +22,7 @@ from vllm.v1.kv_offload.base import (
    OffloadingGaugeMetadata,
    OffloadingHistogramMetadata,
 )
-from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
+from vllm.v1.kv_offload.factory import OffloadingSpecFactory

 LOAD_BYTES = _TransferMetricName.LOAD_BYTES
 LOAD_TIME = _TransferMetricName.LOAD_TIME
@@ -33,6 +33,8 @@ STORE_SIZE = _TransferMetricName.STORE_SIZE
 STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
 PENDING_STORES = "vllm:kv_offload_pending_stores"
 LOOKUP_LATENCY = "vllm:kv_offload_lookup_latency_seconds"
+MY_COUNTER = "my_counter"
+MY_LABEL = "my_label"


 class _FakeMetric:
@@ -67,6 +69,20 @@ class _FakeVllmConfig:
        )


+def _spec_cls_with_metric_definitions(
+    metric_definitions: dict[str, Any],
+) -> type:
+    """Build a fake offloading spec class reporting the given metric
+    definitions, so tests don't need to patch the real CPU spec."""
+
+    class _FakeOffloadingSpec:
+        @staticmethod
+        def build_metric_definitions(extra_config):
+            return metric_definitions
+
+    return _FakeOffloadingSpec
+
+
 def _metric_metadata():
    return {
        LOAD_BYTES: OffloadingCounterMetadata(
@@ -96,9 +112,17 @@ def _metric_metadata():
        LOOKUP_LATENCY: OffloadingHistogramMetadata(
            documentation="lookup latency",
        ),
+        MY_COUNTER: OffloadingCounterMetadata(
+            documentation="counter with a label",
+            labelnames=(MY_LABEL,),
+        ),
    }


+def _unlabeled(values: dict[str, Any], metric_name: str) -> Any:
+    return values[metric_name][()]
+
+
 def test_build_kv_connector_stats_with_none():
    """Test that build_kv_connector_stats returns empty stats when given None."""
    stats = OffloadingConnector.build_kv_connector_stats(data=None)
@@ -131,13 +155,13 @@ def test_build_kv_connector_stats_reconstructs_offload_stats():
            STORES_SKIPPED: _MetricType.COUNTER,
        },
        _StatsKey.DATA: {
-            LOAD_BYTES: 24,
-            LOAD_TIME: 1.5,
-            LOAD_SIZE: [16, 8],
-            STORE_BYTES: 3,
-            STORE_TIME: 0.3,
-            STORE_SIZE: [1, 2],
-            STORES_SKIPPED: 5,
+            LOAD_BYTES: {(): 24},
+            LOAD_TIME: {(): 1.5},
+            LOAD_SIZE: {(): [16, 8]},
+            STORE_BYTES: {(): 3},
+            STORE_TIME: {(): 0.3},
+            STORE_SIZE: {(): [1, 2]},
+            STORES_SKIPPED: {(): 5},
        },
    }

@@ -145,22 +169,28 @@ def test_build_kv_connector_stats_reconstructs_offload_stats():

    assert isinstance(stats, OffloadingConnectorStats)
    values = stats.data[_StatsKey.DATA]
-    assert values[LOAD_BYTES] == 24
-    assert values[LOAD_TIME] == 1.5
-    assert values[LOAD_SIZE] == [16, 8]
-    assert values[STORE_BYTES] == 3
-    assert values[STORE_TIME] == 0.3
-    assert values[STORE_SIZE] == [1, 2]
-    assert values[STORES_SKIPPED] == 5
+    assert _unlabeled(values, LOAD_BYTES) == 24
+    assert _unlabeled(values, LOAD_TIME) == 1.5
+    assert _unlabeled(values, LOAD_SIZE) == [16, 8]
+    assert _unlabeled(values, STORE_BYTES) == 3
+    assert _unlabeled(values, STORE_TIME) == 0.3
+    assert _unlabeled(values, STORE_SIZE) == [1, 2]
+    assert _unlabeled(values, STORES_SKIPPED) == 5


 def _make_stats_data(
    metric_data: dict[str, Any],
    metric_metadata: dict[str, Any],
 ) -> dict[str, Any]:
-    """Build a structured data dict from flat metric data and metadata."""
+    """Build a structured data dict from flat metric data and metadata.
+
+    Values for unlabeled metrics may be passed flat (wrapped here under the
+    empty label tuple); values for labeled metrics must already be passed as
+    a ``{labelvalues: value}`` map.
+    """
    metric_types = {}
-    for key in metric_data:
+    data = {}
+    for key, value in metric_data.items():
        md = metric_metadata[key]
        if isinstance(md, OffloadingCounterMetadata):
            metric_types[key] = _MetricType.COUNTER
@@ -168,9 +198,10 @@ def _make_stats_data(
            metric_types[key] = _MetricType.GAUGE
        elif isinstance(md, OffloadingHistogramMetadata):
            metric_types[key] = _MetricType.HISTOGRAM
+        data[key] = value if md.labelnames else {(): value}
    return {
        _StatsKey.TYPES: metric_types,
-        _StatsKey.DATA: metric_data,
+        _StatsKey.DATA: data,
    }


@@ -215,34 +246,106 @@ def test_aggregate_same_connector():

    assert result is stats1  # Should return self
    values = result.data[_StatsKey.DATA]
-    assert values[LOAD_BYTES] == 34
-    assert values[LOAD_TIME] == 2.6
-    assert values[LOAD_SIZE] == [16, 8, 3, 7]
-    assert values[STORE_BYTES] == 19
-    assert values[STORE_TIME] == 2.3
-    assert values[STORE_SIZE] == [1, 2, 16]
-    assert values[STORES_SKIPPED] == 4
-    assert values[PENDING_STORES] == 1
-    assert values[LOOKUP_LATENCY] == [0.1, 0.2, 0.3]
+    assert _unlabeled(values, LOAD_BYTES) == 34
+    assert _unlabeled(values, LOAD_TIME) == 2.6
+    assert _unlabeled(values, LOAD_SIZE) == [16, 8, 3, 7]
+    assert _unlabeled(values, STORE_BYTES) == 19
+    assert _unlabeled(values, STORE_TIME) == 2.3
+    assert _unlabeled(values, STORE_SIZE) == [1, 2, 16]
+    assert _unlabeled(values, STORES_SKIPPED) == 4
+    assert _unlabeled(values, PENDING_STORES) == 1
+    assert _unlabeled(values, LOOKUP_LATENCY) == [0.1, 0.2, 0.3]
+
+
+def test_aggregate_labeled_metrics():
+    metadata = _metric_metadata()
+    stats1 = OffloadingConnectorStats(
+        data=_make_stats_data(
+            {
+                MY_COUNTER: {
+                    ("a",): 10,
+                    ("b",): 3,
+                },
+            },
+            metadata,
+        ),
+    )
+    stats2 = OffloadingConnectorStats(
+        data=_make_stats_data(
+            {
+                MY_COUNTER: {
+                    ("a",): 7,
+                    ("c",): 5,
+                },
+            },
+            metadata,
+        ),
+    )
+
+    stats1.aggregate(stats2)
+
+    values = stats1.data[_StatsKey.DATA][MY_COUNTER]
+    assert values[("a",)] == 17
+    assert values[("b",)] == 3
+    assert values[("c",)] == 5
+
+
+def test_aggregate_labeled_metric_missing_from_self():
+    """Aggregating a labeled metric that self doesn't have at all yet."""
+    metadata = _metric_metadata()
+    stats1 = OffloadingConnectorStats()
+    stats2 = OffloadingConnectorStats(
+        data=_make_stats_data(
+            {
+                MY_COUNTER: {
+                    ("a",): 7,
+                    ("b",): 5,
+                },
+            },
+            metadata,
+        ),
+    )
+
+    stats1.aggregate(stats2)
+
+    values = stats1.data[_StatsKey.DATA][MY_COUNTER]
+    assert values[("a",)] == 7
+    assert values[("b",)] == 5
+    assert stats1.data[_StatsKey.TYPES][MY_COUNTER] == _MetricType.COUNTER
+
+
+def test_helper_methods_accept_labeled_metrics():
+    stats = OffloadingConnectorStats()
+
+    stats.increase_counter(MY_COUNTER, 3, ("a",))
+    stats.increase_counter(MY_COUNTER, 4, ("a",))
+    stats.set_gauge(PENDING_STORES, 2, ("b",))
+    stats.observe_histogram(LOOKUP_LATENCY, 0.1, ("b",))
+    stats.observe_histogram(LOOKUP_LATENCY, 0.2, ("b",))
+
+    values = stats.data[_StatsKey.DATA]
+    assert values[MY_COUNTER][("a",)] == 7
+    assert values[PENDING_STORES][("b",)] == 2
+    assert values[LOOKUP_LATENCY][("b",)] == [0.1, 0.2]


 def test_aggregate_merges_types():
    stats1 = OffloadingConnectorStats(
        data={
            _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER},
-            _StatsKey.DATA: {LOAD_BYTES: 1},
+            _StatsKey.DATA: {LOAD_BYTES: {(): 1}},
        },
    )
    stats2 = OffloadingConnectorStats(
        data={
            _StatsKey.TYPES: {PENDING_STORES: _MetricType.GAUGE},
-            _StatsKey.DATA: {PENDING_STORES: 2},
+            _StatsKey.DATA: {PENDING_STORES: {(): 2}},
        },
    )

    result = stats1.aggregate(stats2)

-    assert result.data[_StatsKey.DATA][PENDING_STORES] == 2
+    assert _unlabeled(result.data[_StatsKey.DATA], PENDING_STORES) == 2
    assert result.data[_StatsKey.TYPES][PENDING_STORES] == _MetricType.GAUGE


@@ -283,6 +386,26 @@ def test_reduce():
    assert reduced[f"{LOOKUP_LATENCY}_sum"] == sum([0.1, 0.2, 0.3])


+def test_reduce_labeled_metrics():
+    metadata = _metric_metadata()
+    stats = OffloadingConnectorStats(
+        data=_make_stats_data(
+            {
+                MY_COUNTER: {
+                    ("a",): 17,
+                    ("b",): 3,
+                },
+            },
+            metadata,
+        ),
+    )
+
+    reduced = stats.reduce()
+
+    assert reduced[f"{MY_COUNTER}:{('a',)}"] == 17
+    assert reduced[f"{MY_COUNTER}:{('b',)}"] == 3
+
+
 def test_reset():
    """Test that reset() resets all connector stats."""
    metadata = _metric_metadata()
@@ -326,11 +449,11 @@ def test_prom_metrics_observes_manager_counter():
    prom_metrics.observe(
        {
            _StatsKey.TYPES: {STORES_SKIPPED: _MetricType.COUNTER},
-            _StatsKey.DATA: {STORES_SKIPPED: 7},
+            _StatsKey.DATA: {STORES_SKIPPED: {(): 7}},
        }
    )

-    counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED)]
+    counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED, ())]
    assert counter.increments == [7]
    counter_def = prom_metrics._offloading_metric_defs[STORES_SKIPPED]
    assert counter_def.kwargs["name"] == "vllm:kv_offload_stores_skipped"
@@ -360,22 +483,22 @@ def test_prom_metrics_observes_flat_transfer_metrics_and_legacy_metrics():
                STORE_SIZE: _MetricType.HISTOGRAM,
            },
            _StatsKey.DATA: {
-                LOAD_BYTES: 24,
-                LOAD_TIME: 1.5,
-                LOAD_SIZE: [16, 8],
-                STORE_BYTES: 3,
-                STORE_TIME: 0.3,
-                STORE_SIZE: [1, 2],
+                LOAD_BYTES: {(): 24},
+                LOAD_TIME: {(): 1.5},
+                LOAD_SIZE: {(): [16, 8]},
+                STORE_BYTES: {(): 3},
+                STORE_TIME: {(): 0.3},
+                STORE_SIZE: {(): [1, 2]},
            },
        }
    )

-    assert prom_metrics.offloading_metrics[(0, LOAD_BYTES)].increments == [24]
-    assert prom_metrics.offloading_metrics[(0, LOAD_TIME)].increments == [1.5]
-    assert prom_metrics.offloading_metrics[(0, LOAD_SIZE)].observed == [16, 8]
-    assert prom_metrics.offloading_metrics[(0, STORE_BYTES)].increments == [3]
-    assert prom_metrics.offloading_metrics[(0, STORE_TIME)].increments == [0.3]
-    assert prom_metrics.offloading_metrics[(0, STORE_SIZE)].observed == [1, 2]
+    assert prom_metrics.offloading_metrics[(0, LOAD_BYTES, ())].increments == [24]
+    assert prom_metrics.offloading_metrics[(0, LOAD_TIME, ())].increments == [1.5]
+    assert prom_metrics.offloading_metrics[(0, LOAD_SIZE, ())].observed == [16, 8]
+    assert prom_metrics.offloading_metrics[(0, STORE_BYTES, ())].increments == [3]
+    assert prom_metrics.offloading_metrics[(0, STORE_TIME, ())].increments == [0.3]
+    assert prom_metrics.offloading_metrics[(0, STORE_SIZE, ())].observed == [1, 2]

    assert prom_metrics.counter_kv_bytes[(0, "CPU_to_GPU")].increments == [24]
    assert prom_metrics.counter_kv_transfer_time[(0, "CPU_to_GPU")].increments == [1.5]
@@ -396,7 +519,9 @@ def test_prom_metrics_observes_manager_gauge_and_histogram():
        ),
    }
    with patch.object(
-        CPUOffloadingSpec, "build_metric_definitions", return_value=metric_definitions
+        OffloadingSpecFactory,
+        "get_spec_cls",
+        return_value=_spec_cls_with_metric_definitions(metric_definitions),
    ):
        prom_metrics = OffloadPromMetrics(
            vllm_config=_FakeVllmConfig(store_threshold=0),  # type: ignore[arg-type]
@@ -416,20 +541,91 @@ def test_prom_metrics_observes_manager_gauge_and_histogram():
                LOOKUP_LATENCY: _MetricType.HISTOGRAM,
            },
            _StatsKey.DATA: {
-                PENDING_STORES: 5,
-                LOOKUP_LATENCY: [0.2, 0.4],
+                PENDING_STORES: {(): 5},
+                LOOKUP_LATENCY: {(): [0.2, 0.4]},
            },
        }
    )

-    gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES)]
-    histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY)]
+    gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES, ())]
+    histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY, ())]
    assert gauge.set_values == [5]
    assert histogram.observed == [0.2, 0.4]
    histogram_def = prom_metrics._offloading_metric_defs[LOOKUP_LATENCY]
    assert histogram_def.kwargs["buckets"] == (0.1, 1.0)


+def test_prom_metrics_lazily_observes_labeled_metric():
+    metric_definitions = {
+        MY_COUNTER: OffloadingCounterMetadata(
+            documentation="counter with a label",
+            labelnames=(MY_LABEL,),
+        ),
+    }
+    with patch.object(
+        OffloadingSpecFactory,
+        "get_spec_cls",
+        return_value=_spec_cls_with_metric_definitions(metric_definitions),
+    ):
+        prom_metrics = OffloadPromMetrics(
+            vllm_config=_FakeVllmConfig(store_threshold=0),  # type: ignore[arg-type]
+            metric_types={
+                Gauge: _FakeMetric,
+                Counter: _FakeMetric,
+                Histogram: _FakeMetric,
+            },
+            labelnames=["model_name", "engine"],
+            per_engine_labelvalues={0: ["model", "0"]},
+        )
+
+    assert (0, MY_COUNTER, ("a",)) not in prom_metrics.offloading_metrics
+
+    prom_metrics.observe(
+        {
+            _StatsKey.TYPES: {MY_COUNTER: _MetricType.COUNTER},
+            _StatsKey.DATA: {MY_COUNTER: {("a",): 7}},
+        }
+    )
+
+    counter = prom_metrics.offloading_metrics[(0, MY_COUNTER, ("a",))]
+    assert counter.increments == [7]
+    assert counter.labelvalues == ("model", "0", "a")
+    counter_def = prom_metrics._offloading_metric_defs[MY_COUNTER]
+    assert counter_def.kwargs["labelnames"] == ["model_name", "engine", MY_LABEL]
+
+
+def test_prom_metrics_rejects_wrong_label_count():
+    metric_definitions = {
+        MY_COUNTER: OffloadingCounterMetadata(
+            documentation="counter with a label",
+            labelnames=(MY_LABEL,),
+        ),
+    }
+    with patch.object(
+        OffloadingSpecFactory,
+        "get_spec_cls",
+        return_value=_spec_cls_with_metric_definitions(metric_definitions),
+    ):
+        prom_metrics = OffloadPromMetrics(
+            vllm_config=_FakeVllmConfig(store_threshold=0),  # type: ignore[arg-type]
+            metric_types={
+                Gauge: _FakeMetric,
+                Counter: _FakeMetric,
+                Histogram: _FakeMetric,
+            },
+            labelnames=["model_name", "engine"],
+            per_engine_labelvalues={0: ["model", "0"]},
+        )
+
+    with pytest.raises(AssertionError, match="expects 1 labels"):
+        prom_metrics.observe(
+            {
+                _StatsKey.TYPES: {MY_COUNTER: _MetricType.COUNTER},
+                _StatsKey.DATA: {MY_COUNTER: {("a", "extra"): 7}},
+            }
+        )
+
+
 def test_prom_metrics_uses_configured_manager_metrics():
    prom_metrics = OffloadPromMetrics(
        vllm_config=_FakeVllmConfig(store_threshold=0),  # type: ignore[arg-type]
@@ -458,9 +654,9 @@ def test_aggregate_into_empty_stats():
                PENDING_STORES: _MetricType.GAUGE,
            },
            _StatsKey.DATA: {
-                LOAD_BYTES: 42,
-                LOAD_SIZE: [10, 20],
-                PENDING_STORES: 3,
+                LOAD_BYTES: {(): 42},
+                LOAD_SIZE: {(): [10, 20]},
+                PENDING_STORES: {(): 3},
            },
        },
    )
@@ -469,9 +665,9 @@ def test_aggregate_into_empty_stats():

    assert result is empty
    values = result.data[_StatsKey.DATA]
-    assert values[LOAD_BYTES] == 42
-    assert values[LOAD_SIZE] == [10, 20]
-    assert values[PENDING_STORES] == 3
+    assert _unlabeled(values, LOAD_BYTES) == 42
+    assert _unlabeled(values, LOAD_SIZE) == [10, 20]
+    assert _unlabeled(values, PENDING_STORES) == 3


 def test_prom_metrics_multi_engine_routing():
@@ -490,14 +686,13 @@ def test_prom_metrics_multi_engine_routing():
    prom_metrics.observe(
        {
            _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER},
-            _StatsKey.DATA: {LOAD_BYTES: 100},
+            _StatsKey.DATA: {LOAD_BYTES: {(): 100}},
        },
        engine_idx=1,
    )

-    engine0 = prom_metrics.offloading_metrics[(0, LOAD_BYTES)]
-    engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES)]
-    assert engine0.increments == []
+    assert (0, LOAD_BYTES, ()) not in prom_metrics.offloading_metrics
+    engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES, ())]
    assert engine1.increments == [100]


@@ -518,6 +713,6 @@ def test_prom_metrics_rejects_undeclared_metric():
        prom_metrics.observe(
            {
                _StatsKey.TYPES: {"unknown:metric": _MetricType.COUNTER},
-                _StatsKey.DATA: {"unknown:metric": 1},
+                _StatsKey.DATA: {"unknown:metric": {(): 1}},
            }
        )
@@ -112,7 +112,7 @@ class _StatsKey:

    # Maps metric name -> _MetricType value
    TYPES = "types"
-    # Maps metric name -> observed value (number or list)
+    # Maps metric name -> {label values tuple -> observed value (number or list)}
    DATA = "data"


@@ -125,15 +125,17 @@ class OffloadingConnectorStats(KVConnectorStats):

        {
            _StatsKey.TYPES: {name: _MetricType.*, ...},
-            _StatsKey.DATA:  {name: value, ...},
+            _StatsKey.DATA:  {name: {labelvalues: value, ...}, ...},
        }

    This structure is self-describing: it survives IPC serialization
    without needing the full ``OffloadingMetricMetadata`` objects on the
    receiving side.

-    Counter values are aggregated by summing, gauge values use the latest
-    snapshot, and histogram values are lists of observed samples.
+    Counter values are aggregated by summing per-label-tuple, gauge values
+    use the latest snapshot per-label-tuple, and histogram values are lists of
+    observed samples per-label-tuple. Unlabeled metrics use ``()`` as their
+    labelvalues tuple.
    """

    def __post_init__(self):
@@ -160,26 +162,32 @@ class OffloadingConnectorStats(KVConnectorStats):
        assert isinstance(other, OffloadingConnectorStats)
        other_types = other._types
        other_values = other._values
-        for key, value in other_values.items():
+        for key, other_label_values in other_values.items():
            type_str = other_types.get(key)
            if type_str is None:
                raise AssertionError(f"Unknown offloading stats key: {key}")
            self._types.setdefault(key, type_str)
-            if type_str == _MetricType.HISTOGRAM:
-                assert isinstance(value, list)
-                if key not in self._values:
-                    self._values[key] = value
+            current_label_values = self._values.setdefault(key, {})
+            for labelvalues, value in other_label_values.items():
+                if type_str == _MetricType.HISTOGRAM:
+                    assert isinstance(value, list)
+                    if labelvalues not in current_label_values:
+                        current_label_values[labelvalues] = list(value)
+                    else:
+                        assert isinstance(current_label_values[labelvalues], list)
+                        current_label_values[labelvalues].extend(value)
+                elif type_str == _MetricType.COUNTER:
+                    assert isinstance(value, int | float)
+                    current_label_values[labelvalues] = (
+                        current_label_values.get(labelvalues, 0) + value
+                    )
+                elif type_str == _MetricType.GAUGE:
+                    assert isinstance(value, int | float)
+                    current_label_values[labelvalues] = value
                else:
-                    assert isinstance(self._values[key], list)
-                    self._values[key].extend(value)
-            elif type_str == _MetricType.COUNTER:
-                assert isinstance(value, int | float)
-                self._values[key] = self._values.get(key, 0) + value
-            elif type_str == _MetricType.GAUGE:
-                assert isinstance(value, int | float)
-                self._values[key] = value
-            else:
-                raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
+                    raise AssertionError(
+                        f"Unknown metric type '{type_str}' for key: {key}"
+                    )
        return self

    def reduce(self) -> dict[str, int | float]:
@@ -190,44 +198,62 @@ class OffloadingConnectorStats(KVConnectorStats):
        stats for the last time interval.
        """
        return_dict: dict[str, int | float] = {}
-        for key, value in self._values.items():
+        for key, label_value_map in self._values.items():
            type_str = self._types.get(key)
            if type_str is None:
                raise AssertionError(f"Unknown offloading stats key: {key}")
-            if type_str == _MetricType.HISTOGRAM:
-                assert isinstance(value, list)
-                return_dict[f"{key}_count"] = len(value)
-                return_dict[f"{key}_sum"] = sum(value)
-            elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE):
-                assert isinstance(value, int | float)
-                return_dict[key] = value
-            else:
-                raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
+            for labelvalues, value in label_value_map.items():
+                key_with_labels = f"{key}:{labelvalues}" if labelvalues else key
+                if type_str == _MetricType.HISTOGRAM:
+                    assert isinstance(value, list)
+                    return_dict[f"{key_with_labels}_count"] = len(value)
+                    return_dict[f"{key_with_labels}_sum"] = sum(value)
+                elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE):
+                    assert isinstance(value, int | float)
+                    return_dict[key_with_labels] = value
+                else:
+                    raise AssertionError(
+                        f"Unknown metric type '{type_str}' for key: {key}"
+                    )
        return return_dict

    def is_empty(self) -> bool:
        return not self.data.get(_StatsKey.DATA)

    def increase_counter(
-        self, counter_name: str, counter_increase_value: int | float
+        self,
+        counter_name: str,
+        counter_increase_value: int | float,
+        labelvalues: tuple[str, ...] = (),
    ) -> None:
        """Increase a counter on the stats payload."""
        self._types.setdefault(counter_name, _MetricType.COUNTER)
-        self._values[counter_name] = (
-            self._values.get(counter_name, 0) + counter_increase_value
+        counter_values = self._values.setdefault(counter_name, {})
+        counter_values[labelvalues] = (
+            counter_values.get(labelvalues, 0) + counter_increase_value
        )

-    def set_gauge(self, gauge_name: str, gauge_value: int | float) -> None:
+    def set_gauge(
+        self,
+        gauge_name: str,
+        gauge_value: int | float,
+        labelvalues: tuple[str, ...] = (),
+    ) -> None:
        """Set a gauge snapshot on the stats payload."""
        self._types.setdefault(gauge_name, _MetricType.GAUGE)
-        self._values[gauge_name] = gauge_value
+        gauge_values = self._values.setdefault(gauge_name, {})
+        gauge_values[labelvalues] = gauge_value

    def observe_histogram(
-        self, histogram_name: str, histogram_value: int | float
+        self,
+        histogram_name: str,
+        histogram_value: int | float,
+        labelvalues: tuple[str, ...] = (),
    ) -> None:
        """Record a histogram observation on the stats payload."""
        self._types.setdefault(histogram_name, _MetricType.HISTOGRAM)
-        self._values.setdefault(histogram_name, []).append(histogram_value)
+        histogram_values = self._values.setdefault(histogram_name, {})
+        histogram_values.setdefault(labelvalues, []).append(histogram_value)


 class OffloadPromMetrics(KVConnectorPromMetrics):
@@ -255,7 +281,10 @@ class OffloadPromMetrics(KVConnectorPromMetrics):

        self._observe_deprecated_metrics = issubclass(spec_cls, CPUOffloadingSpec)
        self._offloading_metric_defs: dict[str, PromMetricT] = {}
-        self.offloading_metrics: dict[tuple[int, str], PromMetricT] = {}
+        # (engine_idx, metric_name, labelvalues) -> metric with bound labels
+        self.offloading_metrics: dict[
+            tuple[int, str, tuple[str, ...]], PromMetricT
+        ] = {}

        self._counter_kv_bytes = self._counter_cls(
            name=_DEPRECATED_TOTAL_BYTES,
@@ -301,10 +330,6 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
            self._offloading_metric_defs[metric_name] = self._create_metric(
                metric_name, metadata
            )
-            for engine_idx, labelvalues in per_engine_labelvalues.items():
-                self.offloading_metrics[(engine_idx, metric_name)] = (
-                    self._offloading_metric_defs[metric_name].labels(*labelvalues)
-                )

    def _create_metric(
        self, metric_name: str, metadata: OffloadingMetricMetadata
@@ -312,7 +337,7 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
        kwargs: dict[str, Any] = {
            "name": metric_name,
            "documentation": metadata.documentation,
-            "labelnames": self._labelnames,
+            "labelnames": self._labelnames + list(metadata.labelnames),
        }
        if isinstance(metadata, OffloadingCounterMetadata):
            metric_cls = self._counter_cls
@@ -326,11 +351,37 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
            raise AssertionError(f"Unknown offloading metric metadata: {metadata}")
        return metric_cls(**kwargs)

+    def _get_prometheus_metric(
+        self,
+        metric_name: str,
+        labelvalues: tuple[str, ...],
+        engine_idx: int,
+    ) -> PromMetric:
+        metadata = self._offloading_metric_metadata[metric_name]
+        if len(labelvalues) != len(metadata.labelnames):
+            raise AssertionError(
+                f"Metric {metric_name} expects {len(metadata.labelnames)} labels, "
+                f"got {len(labelvalues)}"
+            )
+        key = (engine_idx, metric_name, labelvalues)
+        prom_metric = self.offloading_metrics.get(key)
+        if prom_metric is None:
+            engine_labelvalues = self.per_engine_labelvalues[engine_idx]
+            prom_metric = self._offloading_metric_defs[metric_name].labels(
+                *(engine_labelvalues + list(labelvalues))
+            )
+            self.offloading_metrics[key] = prom_metric
+        return prom_metric
+
    def _increase_counter(
-        self, metric_name: str, value: int | float, engine_idx: int
+        self,
+        metric_name: str,
+        value: int | float,
+        labelvalues: tuple[str, ...],
+        engine_idx: int,
    ) -> None:
-        self.offloading_metrics[(engine_idx, metric_name)].inc(value)
-        if not self._observe_deprecated_metrics:
+        self._get_prometheus_metric(metric_name, labelvalues, engine_idx).inc(value)
+        if labelvalues or not self._observe_deprecated_metrics:
            return
        # Keep deprecated CPU offload transfer metrics updated during the
        # transition to flat metric names.
@@ -343,15 +394,26 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
        elif metric_name == _TransferMetricName.STORE_TIME:
            self.counter_kv_transfer_time[(engine_idx, _TransferType.STORE)].inc(value)

-    def _set_gauge(self, metric_name: str, value: int | float, engine_idx: int) -> None:
-        self.offloading_metrics[(engine_idx, metric_name)].set(value)
+    def _set_gauge(
+        self,
+        metric_name: str,
+        value: int | float,
+        labelvalues: tuple[str, ...],
+        engine_idx: int,
+    ) -> None:
+        self._get_prometheus_metric(metric_name, labelvalues, engine_idx).set(value)

    def _observe_histogram(
-        self, metric_name: str, value: list[int | float], engine_idx: int
+        self,
+        metric_name: str,
+        value: list[int | float],
+        labelvalues: tuple[str, ...],
+        engine_idx: int,
    ) -> None:
+        prom_metric = self._get_prometheus_metric(metric_name, labelvalues, engine_idx)
        for observation in value:
-            self.offloading_metrics[(engine_idx, metric_name)].observe(observation)
-            if not self._observe_deprecated_metrics:
+            prom_metric.observe(observation)
+            if labelvalues or not self._observe_deprecated_metrics:
                continue
            # Keep deprecated CPU offload transfer metrics updated during the
            # transition to flat metric names.
@@ -368,20 +430,23 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
        """Observe transfer statistics."""
        metric_types = transfer_stats_data.get(_StatsKey.TYPES, {})
        metric_data = transfer_stats_data.get(_StatsKey.DATA, {})
-        for key, value in metric_data.items():
+        for key, label_value_map in metric_data.items():
            type_str = metric_types.get(key)
            if type_str is None:
                raise AssertionError(f"Unknown offloading stats key: {key}")
            assert key in self._offloading_metric_defs
-            if type_str == _MetricType.COUNTER:
-                assert isinstance(value, int | float)
-                self._increase_counter(key, value, engine_idx)
-            elif type_str == _MetricType.GAUGE:
-                assert isinstance(value, int | float)
-                self._set_gauge(key, value, engine_idx)
-            elif type_str == _MetricType.HISTOGRAM:
-                assert isinstance(value, list)
-                assert all(isinstance(v, int | float) for v in value)
-                self._observe_histogram(key, value, engine_idx)
-            else:
-                raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
+            for labelvalues, value in label_value_map.items():
+                if type_str == _MetricType.COUNTER:
+                    assert isinstance(value, int | float)
+                    self._increase_counter(key, value, labelvalues, engine_idx)
+                elif type_str == _MetricType.GAUGE:
+                    assert isinstance(value, int | float)
+                    self._set_gauge(key, value, labelvalues, engine_idx)
+                elif type_str == _MetricType.HISTOGRAM:
+                    assert isinstance(value, list)
+                    assert all(isinstance(v, int | float) for v in value)
+                    self._observe_histogram(key, value, labelvalues, engine_idx)
+                else:
+                    raise AssertionError(
+                        f"Unknown metric type '{type_str}' for key: {key}"
+                    )
@@ -129,6 +129,7 @@ The class provides the following primitives:
@dataclass(frozen=True)
 class OffloadingMetricMetadata:
    documentation: str
+    labelnames: tuple[str, ...] = ()


@dataclass(frozen=True)