[KV Offloading] Add tiering metric plumbing (#45959)
Signed-off-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com>
Srinivas Krovvidi committed
Jun 23, 2026 at 07:10 UTC
091bc1026ea4e5120893ae9a730fdb1bf5e873e2
5 files changed
+184
-13
tests/v1/kv_offload/tiering/test_tiering_offloading.py
+107
@@ -17,18 +17,29 @@ from unittest.mock import MagicMock
17
import pytest
18
import torch
19
20
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
21
+ OffloadingConnectorStats,
22
+)
23
from vllm.v1.kv_offload.base import (
24
+ OffloadingCounterMetadata,
25
OffloadKey,
26
OffloadPolicy,
27
ReqContext,
28
RequestOffloadingContext,
29
make_offload_key,
30
)
31
+from vllm.v1.kv_offload.tiering.base import (
32
+ JobMetadata,
33
+ JobResult,
34
+ SecondaryTierManager,
35
+)
36
from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
37
+from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory
38
from vllm.v1.kv_offload.tiering.manager import (
39
CPUPrimaryTierOffloadingManager,
40
TieringOffloadingManager,
41
)
42
+from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
43
44
_CTX = ReqContext(req_id="test")
45
_MOCK_OFFLOADING_SPEC = MagicMock()
@@ -63,6 +74,102 @@ def count_hits(manager, keys: list[OffloadKey]) -> int | None:
74
return count
75
76
77
+class MetricsSecondaryTierManager(SecondaryTierManager):
78
+ """Test-only secondary tier that declares and emits one labeled metric."""
79
+
80
+ MY_TIER_METRIC = "my_tier_metric"
81
+
82
+ @classmethod
83
+ def build_metric_definitions(cls, extra_config):
84
+ return {
85
+ cls.MY_TIER_METRIC: OffloadingCounterMetadata(
86
+ documentation="Number of bytes served by the test tier.",
87
+ labelnames=("tier",),
88
+ )
89
+ }
90
+
91
+ def __init__(self, *args, **kwargs):
92
+ super().__init__(*args, **kwargs)
93
+ self.stats: OffloadingConnectorStats | None = None
94
+
95
+ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
96
+ return False
97
+
98
+ def submit_store(self, job_metadata: JobMetadata) -> None:
99
+ return
100
+
101
+ def submit_load(self, job_metadata: JobMetadata) -> None:
102
+ return
103
+
104
+ def get_finished_jobs(self) -> Iterable[JobResult]:
105
+ return ()
106
+
107
+ def drain_jobs(self) -> None:
108
+ return
109
+
110
+ def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
111
+ return RequestOffloadingContext()
112
+
113
+ def get_stats(self) -> OffloadingConnectorStats | None:
114
+ stats = self.stats
115
+ self.stats = None
116
+ return stats
117
+
118
+
119
+def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch):
120
+ monkeypatch.setitem(
121
+ SecondaryTierFactory._registry,
122
+ "test_metrics",
123
+ lambda: MetricsSecondaryTierManager,
124
+ )
125
+
126
+ metrics = TieringOffloadingSpec.build_metric_definitions(
127
+ {"secondary_tiers": [{"type": "test_metrics"}]}
128
+ )
129
+
130
+ metadata = metrics[MetricsSecondaryTierManager.MY_TIER_METRIC]
131
+ assert metadata.documentation == "Number of bytes served by the test tier."
132
+ assert metadata.labelnames == ("tier",)
133
+
134
+
135
+def test_tiering_manager_aggregates_secondary_stats():
136
+ mock_region = _mock_mmap_region(5)
137
+ primary_tier = CPUPrimaryTierOffloadingManager(
138
+ num_blocks=5, mmap_region=mock_region
139
+ )
140
+ secondary_tier = MetricsSecondaryTierManager(
141
+ offloading_spec=_MOCK_OFFLOADING_SPEC,
142
+ primary_kv_view=mock_region.create_kv_memoryview(),
143
+ tier_type="test_metrics",
144
+ )
145
+ secondary_stats = OffloadingConnectorStats()
146
+ secondary_stats.increase_counter(
147
+ MetricsSecondaryTierManager.MY_TIER_METRIC, 7, ("test_metrics",)
148
+ )
149
+ secondary_tier.stats = secondary_stats
150
+ manager = TieringOffloadingManager(
151
+ primary_tier=primary_tier,
152
+ secondary_tiers=[secondary_tier],
153
+ )
154
+
155
+ stats = manager.get_stats()
156
+
157
+ assert stats is not None
158
+ assert (
159
+ stats.data["data"][MetricsSecondaryTierManager.MY_TIER_METRIC][
160
+ ("test_metrics",)
161
+ ]
162
+ == 7
163
+ )
164
+
165
+ # The primary tier's cache-usage gauge is always reported, so get_stats()
166
+ # never returns None, but the secondary tier has nothing new to report
167
+ # once its stats have been consumed.
168
+ second_stats = manager.get_stats()
169
+ assert second_stats is not None
170
+ assert MetricsSecondaryTierManager.MY_TIER_METRIC not in second_stats.data["data"]
171
+
172
+
173
class TestExampleSecondaryTierManager:
174
"""Tests for ExampleSecondaryTierManager implementation."""
175
vllm/v1/kv_offload/tiering/base.py
+21
-2
@@ -7,13 +7,21 @@ Abstract interfaces and data types for the secondary tiering layer.
7
from abc import ABC, abstractmethod
8
from collections.abc import Collection, Iterable
9
from dataclasses import dataclass
10
-from typing import TYPE_CHECKING
10
+from typing import TYPE_CHECKING, Any
11
12
import numpy as np
13
14
-from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext
14
+from vllm.v1.kv_offload.base import (
15
+ OffloadingMetricMetadata,
16
+ OffloadKey,
17
+ ReqContext,
18
+ RequestOffloadingContext,
19
+)
20
21
if TYPE_CHECKING:
22
+ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
23
+ OffloadingConnectorStats,
24
+ )
25
from vllm.v1.kv_offload.base import OffloadingSpec
26
27
# Type alias for job IDs used in async transfer tracking
@@ -228,3 +236,14 @@ class SecondaryTierManager(ABC):
236
def shutdown(self) -> None:
237
"""Release resources held by this tier (threads, connections, etc.)."""
238
return
239
+
240
+ @classmethod
241
+ def build_metric_definitions(
242
+ cls, extra_config: dict[str, Any]
243
+ ) -> dict[str, OffloadingMetricMetadata]:
244
+ """Return Prometheus metric definitions emitted by this tier."""
245
+ return {}
246
+
247
+ def get_stats(self) -> "OffloadingConnectorStats | None":
248
+ """Return and reset metric observations collected by this tier."""
249
+ return None
vllm/v1/kv_offload/tiering/factory.py
+12
-10
@@ -31,25 +31,27 @@ class SecondaryTierFactory:
31
primary_kv_view: memoryview,
32
offloading_spec: "OffloadingSpec",
33
) -> SecondaryTierManager:
34
+ tier_cls = cls.get_tier_class(tier_config)
35
config = tier_config.copy()
36
+ tier_type = config.pop("type")
37
+ return tier_cls(
38
+ offloading_spec=offloading_spec,
39
+ primary_kv_view=primary_kv_view,
40
+ tier_type=tier_type,
41
+ **config,
42
+ )
43
36
- tier_type = config.pop("type", None)
44
+ @classmethod
45
+ def get_tier_class(cls, tier_config: dict) -> type[SecondaryTierManager]:
46
+ tier_type = tier_config.get("type")
47
if not tier_type:
48
raise ValueError("Secondary tier configuration must include 'type'")
39
-
49
if tier_type not in cls._registry:
50
raise ValueError(
51
f"Unknown secondary tier type: {tier_type!r}. "
52
f"Supported types: {list(cls._registry)}"
53
)
45
-
46
- tier_cls = cls._registry[tier_type]()
47
- return tier_cls(
48
- offloading_spec=offloading_spec,
49
- primary_kv_view=primary_kv_view,
50
- tier_type=tier_type,
51
- **config,
52
- )
54
+ return cls._registry[tier_type]()
55
56
57
SecondaryTierFactory.register_tier(
vllm/v1/kv_offload/tiering/manager.py
+21
@@ -27,6 +27,9 @@ from dataclasses import dataclass, field
27
import numpy as np
28
from typing_extensions import override
29
30
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
31
+ OffloadingConnectorStats,
32
+)
33
from vllm.logger import init_logger
34
from vllm.v1.kv_offload.base import (
35
LoadStoreSpec,
@@ -628,6 +631,24 @@ class TieringOffloadingManager(OffloadingManager):
631
self._request_level_tiers.clear()
632
self._processed_jobs_this_step = False
633
634
+ @override
635
+ def get_stats(self) -> OffloadingConnectorStats | None:
636
+ stats = self.primary_tier.get_stats()
637
+
638
+ if stats is not None and stats.is_empty():
639
+ stats = None
640
+
641
+ for tier in self.secondary_tiers:
642
+ tier_stats = tier.get_stats()
643
+ if tier_stats is None or tier_stats.is_empty():
644
+ continue
645
+ if stats is None:
646
+ stats = tier_stats
647
+ else:
648
+ stats.aggregate(tier_stats)
649
+
650
+ return stats
651
+
652
@override
653
def shutdown(self) -> None:
654
"""Shutdown all tiers and release resources."""
vllm/v1/kv_offload/tiering/spec.py
+23
-1
@@ -31,13 +31,19 @@ Example configuration:
31
}
32
"""
33
34
+from typing import Any
35
+
36
import torch
37
from typing_extensions import override
38
39
from vllm.config import VllmConfig
40
from vllm.logger import init_logger
41
from vllm.v1.kv_cache_interface import KVCacheConfig
40
-from vllm.v1.kv_offload.base import CanonicalKVCaches, OffloadingManager
42
+from vllm.v1.kv_offload.base import (
43
+ CanonicalKVCaches,
44
+ OffloadingManager,
45
+ OffloadingMetricMetadata,
46
+)
47
from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers
48
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
49
from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
@@ -65,6 +71,22 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
71
72
BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
73
74
+ @classmethod
75
+ @override
76
+ def build_metric_definitions(
77
+ cls, extra_config: dict[str, Any]
78
+ ) -> dict[str, OffloadingMetricMetadata]:
79
+ metrics = super().build_metric_definitions(extra_config)
80
+ secondary_tier_configs = extra_config.get("secondary_tiers", [])
81
+ if not isinstance(secondary_tier_configs, list):
82
+ raise ValueError("secondary_tiers must be a list of tier configurations")
83
+
84
+ for tier_config in secondary_tier_configs:
85
+ assert isinstance(tier_config, dict)
86
+ tier_cls = SecondaryTierFactory.get_tier_class(tier_config)
87
+ metrics.update(tier_cls.build_metric_definitions(tier_config))
88
+ return metrics
89
+
90
def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
91
super().__init__(vllm_config, kv_cache_config)
92
# Redeclare for mypy: parent sets this but `--follow-imports skip` hides it