@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9dfc313bdc
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Feature/offloading manager stats (#35669)

Signed-off-by: Sriusa4414@gmail.com Signed-off-by: srinivas_oo7 <Sriusa4414@gmail.com> Signed-off-by: srinivas_oo7 <sklinkedin0120@gmail.com> Signed-off-by: Srinivasoo7 <158864704+Srinivasoo7@users.noreply.github.com> Signed-off-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: Srinivasoo7 <158864704+Srinivasoo7@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

Srinivas Krovvidi committed Jun 10, 2026 at 07:44 UTC 9dfc313bdce57ac7062204ace419a7f10dac0399
14 files changed +1200 -213
tests/v1/kv_connector/unit/offloading_connector/test_metrics.py
+453 -81
@@ -1,11 +1,102 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from types import SimpleNamespace
4 +from typing import Any
5 +from unittest.mock import patch
6 +
7 +import pytest
8 +from prometheus_client import Counter, Gauge, Histogram
9 +
10 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
11 OffloadingConnectorStats,
12 + OffloadPromMetrics,
13 + _MetricType,
14 + _StatsKey,
15 + _TransferMetricName,
16 )
17 from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
18 OffloadingConnector,
19 )
20 +from vllm.v1.kv_offload.base import (
21 + OffloadingCounterMetadata,
22 + OffloadingGaugeMetadata,
23 + OffloadingHistogramMetadata,
24 +)
25 +from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
26 +
27 +LOAD_BYTES = _TransferMetricName.LOAD_BYTES
28 +LOAD_TIME = _TransferMetricName.LOAD_TIME
29 +LOAD_SIZE = _TransferMetricName.LOAD_SIZE
30 +STORE_BYTES = _TransferMetricName.STORE_BYTES
31 +STORE_TIME = _TransferMetricName.STORE_TIME
32 +STORE_SIZE = _TransferMetricName.STORE_SIZE
33 +STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
34 +PENDING_STORES = "vllm:kv_offload_pending_stores"
35 +LOOKUP_LATENCY = "vllm:kv_offload_lookup_latency_seconds"
36 +
37 +
38 +class _FakeMetric:
39 + def __init__(self, **kwargs):
40 + self.kwargs = kwargs
41 + self.children: list[_FakeMetric] = []
42 + self.observed: list[int | float] = []
43 + self.increments: list[int | float] = []
44 + self.set_values: list[int | float] = []
45 + self.labelvalues: tuple[object, ...] = ()
46 +
47 + def labels(self, *labelvalues):
48 + child = _FakeMetric(**self.kwargs)
49 + child.labelvalues = labelvalues
50 + self.children.append(child)
51 + return child
52 +
53 + def observe(self, value):
54 + self.observed.append(value)
55 +
56 + def inc(self, value):
57 + self.increments.append(value)
58 +
59 + def set(self, value):
60 + self.set_values.append(value)
61 +
62 +
63 +class _FakeVllmConfig:
64 + def __init__(self, store_threshold: int = 2):
65 + self.kv_transfer_config = SimpleNamespace(
66 + kv_connector_extra_config={"store_threshold": store_threshold}
67 + )
68 +
69 +
70 +def _metric_metadata():
71 + return {
72 + LOAD_BYTES: OffloadingCounterMetadata(
73 + documentation="load bytes",
74 + ),
75 + LOAD_TIME: OffloadingCounterMetadata(
76 + documentation="load time",
77 + ),
78 + LOAD_SIZE: OffloadingHistogramMetadata(
79 + documentation="load size",
80 + ),
81 + STORE_BYTES: OffloadingCounterMetadata(
82 + documentation="store bytes",
83 + ),
84 + STORE_TIME: OffloadingCounterMetadata(
85 + documentation="store time",
86 + ),
87 + STORE_SIZE: OffloadingHistogramMetadata(
88 + documentation="store size",
89 + ),
90 + STORES_SKIPPED: OffloadingCounterMetadata(
91 + documentation="stores skipped",
92 + ),
93 + PENDING_STORES: OffloadingGaugeMetadata(
94 + documentation="pending stores",
95 + ),
96 + LOOKUP_LATENCY: OffloadingHistogramMetadata(
97 + documentation="lookup latency",
98 + ),
99 + }
100
101
102 def test_build_kv_connector_stats_with_none():
@@ -14,7 +105,6 @@ def test_build_kv_connector_stats_with_none():
105
106 assert stats is not None
107 assert isinstance(stats, OffloadingConnectorStats)
17 - assert len(stats.data) == 0
108 assert stats.is_empty()
109
110
@@ -24,7 +114,6 @@ def test_build_kv_connector_stats_with_empty_dict():
114
115 assert stats is not None
116 assert isinstance(stats, OffloadingConnectorStats)
27 - assert len(stats.data) == 0
117 assert stats.is_empty()
118
119
@@ -32,114 +121,186 @@ def test_build_kv_connector_stats_reconstructs_offload_stats():
121 """Test that OffloadingConnector stats are properly reconstructed with
122 correct data."""
123 serialized_data = {
35 - "CPU_to_GPU": [
36 - {"op_size": 16, "op_time": 1.0},
37 - {"op_size": 8, "op_time": 0.5},
38 - ],
39 - "GPU_to_CPU": [
40 - {"op_size": 1, "op_time": 0.1},
41 - {"op_size": 2, "op_time": 0.2},
42 - ],
124 + _StatsKey.TYPES: {
125 + LOAD_BYTES: _MetricType.COUNTER,
126 + LOAD_TIME: _MetricType.COUNTER,
127 + LOAD_SIZE: _MetricType.HISTOGRAM,
128 + STORE_BYTES: _MetricType.COUNTER,
129 + STORE_TIME: _MetricType.COUNTER,
130 + STORE_SIZE: _MetricType.HISTOGRAM,
131 + STORES_SKIPPED: _MetricType.COUNTER,
132 + },
133 + _StatsKey.DATA: {
134 + LOAD_BYTES: 24,
135 + LOAD_TIME: 1.5,
136 + LOAD_SIZE: [16, 8],
137 + STORE_BYTES: 3,
138 + STORE_TIME: 0.3,
139 + STORE_SIZE: [1, 2],
140 + STORES_SKIPPED: 5,
141 + },
142 }
143
144 stats = OffloadingConnector.build_kv_connector_stats(data=serialized_data)
145
47 - offload_connector_stats = stats
48 - assert isinstance(offload_connector_stats, OffloadingConnectorStats)
49 - assert offload_connector_stats.data["CPU_to_GPU"] == [
50 - {"op_size": 16, "op_time": 1.0},
51 - {"op_size": 8, "op_time": 0.5},
52 - ]
53 - assert offload_connector_stats.data["GPU_to_CPU"] == [
54 - {"op_size": 1, "op_time": 0.1},
55 - {"op_size": 2, "op_time": 0.2},
56 - ]
146 + assert isinstance(stats, OffloadingConnectorStats)
147 + values = stats.data[_StatsKey.DATA]
148 + assert values[LOAD_BYTES] == 24
149 + assert values[LOAD_TIME] == 1.5
150 + assert values[LOAD_SIZE] == [16, 8]
151 + assert values[STORE_BYTES] == 3
152 + assert values[STORE_TIME] == 0.3
153 + assert values[STORE_SIZE] == [1, 2]
154 + assert values[STORES_SKIPPED] == 5
155 +
156 +
157 +def _make_stats_data(
158 + metric_data: dict[str, Any],
159 + metric_metadata: dict[str, Any],
160 +) -> dict[str, Any]:
161 + """Build a structured data dict from flat metric data and metadata."""
162 + metric_types = {}
163 + for key in metric_data:
164 + md = metric_metadata[key]
165 + if isinstance(md, OffloadingCounterMetadata):
166 + metric_types[key] = _MetricType.COUNTER
167 + elif isinstance(md, OffloadingGaugeMetadata):
168 + metric_types[key] = _MetricType.GAUGE
169 + elif isinstance(md, OffloadingHistogramMetadata):
170 + metric_types[key] = _MetricType.HISTOGRAM
171 + return {
172 + _StatsKey.TYPES: metric_types,
173 + _StatsKey.DATA: metric_data,
174 + }
175
176
177 def test_aggregate_same_connector():
178 """Test aggregating stats from the same connector type."""
179 + metadata = _metric_metadata()
180 stats1 = OffloadingConnectorStats(
62 - data={
63 - "CPU_to_GPU": [
64 - {"op_size": 16, "op_time": 1.0},
65 - {"op_size": 8, "op_time": 0.5},
66 - ],
67 - "GPU_to_CPU": [
68 - {"op_size": 1, "op_time": 0.1},
69 - {"op_size": 2, "op_time": 0.2},
70 - ],
71 - }
181 + data=_make_stats_data(
182 + {
183 + LOAD_BYTES: 24,
184 + LOAD_TIME: 1.5,
185 + LOAD_SIZE: [16, 8],
186 + STORE_BYTES: 3,
187 + STORE_TIME: 0.3,
188 + STORE_SIZE: [1, 2],
189 + STORES_SKIPPED: 1,
190 + PENDING_STORES: 3,
191 + LOOKUP_LATENCY: [0.1],
192 + },
193 + metadata,
194 + ),
195 )
196
197 stats2 = OffloadingConnectorStats(
75 - data={
76 - "CPU_to_GPU": [
77 - {"op_size": 3, "op_time": 0.2},
78 - {"op_size": 7, "op_time": 0.9},
79 - ],
80 - "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
81 - }
198 + data=_make_stats_data(
199 + {
200 + LOAD_BYTES: 10,
201 + LOAD_TIME: 1.1,
202 + LOAD_SIZE: [3, 7],
203 + STORE_BYTES: 16,
204 + STORE_TIME: 2,
205 + STORE_SIZE: [16],
206 + STORES_SKIPPED: 3,
207 + PENDING_STORES: 1,
208 + LOOKUP_LATENCY: [0.2, 0.3],
209 + },
210 + metadata,
211 + ),
212 )
213
214 result = stats1.aggregate(stats2)
215
216 assert result is stats1 # Should return self
87 - offload_connector_stats = result
88 - assert offload_connector_stats.data["CPU_to_GPU"] == [
89 - {"op_size": 16, "op_time": 1.0},
90 - {"op_size": 8, "op_time": 0.5},
91 - {"op_size": 3, "op_time": 0.2},
92 - {"op_size": 7, "op_time": 0.9},
93 - ]
94 - assert offload_connector_stats.data["GPU_to_CPU"] == [
95 - {"op_size": 1, "op_time": 0.1},
96 - {"op_size": 2, "op_time": 0.2},
97 - {"op_size": 16, "op_time": 2},
98 - ]
217 + values = result.data[_StatsKey.DATA]
218 + assert values[LOAD_BYTES] == 34
219 + assert values[LOAD_TIME] == 2.6
220 + assert values[LOAD_SIZE] == [16, 8, 3, 7]
221 + assert values[STORE_BYTES] == 19
222 + assert values[STORE_TIME] == 2.3
223 + assert values[STORE_SIZE] == [1, 2, 16]
224 + assert values[STORES_SKIPPED] == 4
225 + assert values[PENDING_STORES] == 1
226 + assert values[LOOKUP_LATENCY] == [0.1, 0.2, 0.3]
227 +
228 +
229 +def test_aggregate_merges_types():
230 + stats1 = OffloadingConnectorStats(
231 + data={
232 + _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER},
233 + _StatsKey.DATA: {LOAD_BYTES: 1},
234 + },
235 + )
236 + stats2 = OffloadingConnectorStats(
237 + data={
238 + _StatsKey.TYPES: {PENDING_STORES: _MetricType.GAUGE},
239 + _StatsKey.DATA: {PENDING_STORES: 2},
240 + },
241 + )
242 +
243 + result = stats1.aggregate(stats2)
244 +
245 + assert result.data[_StatsKey.DATA][PENDING_STORES] == 2
246 + assert result.data[_StatsKey.TYPES][PENDING_STORES] == _MetricType.GAUGE
247
248
249 def test_reduce():
102 - """Test that reduce() correctly reduces all nested connector stats."""
250 + """Test that reduce() correctly reduces connector stats."""
251 + metadata = _metric_metadata()
252 stats = OffloadingConnectorStats(
104 - data={
105 - "CPU_to_GPU": [
106 - {"op_size": 16, "op_time": 1.0},
107 - {"op_size": 8, "op_time": 0.5},
108 - {"op_size": 3, "op_time": 0.2},
109 - {"op_size": 7, "op_time": 0.9},
110 - ],
111 - "GPU_to_CPU": [
112 - {"op_size": 1, "op_time": 0.1},
113 - {"op_size": 2, "op_time": 0.2},
114 - {"op_size": 16, "op_time": 2},
115 - ],
116 - }
253 + data=_make_stats_data(
254 + {
255 + LOAD_BYTES: 34,
256 + LOAD_TIME: 2.6,
257 + LOAD_SIZE: [16, 8, 3, 7],
258 + STORE_BYTES: 19,
259 + STORE_TIME: 2.3,
260 + STORE_SIZE: [1, 2, 16],
261 + STORES_SKIPPED: 11,
262 + PENDING_STORES: 2,
263 + LOOKUP_LATENCY: [0.1, 0.2, 0.3],
264 + },
265 + metadata,
266 + ),
267 )
268
269 reduced = stats.reduce()
270
271 assert isinstance(reduced, dict)
122 - # Check that the stats were reduced (should have aggregated values)
123 - assert "CPU_to_GPU_total_bytes" in reduced
124 - assert "CPU_to_GPU_total_time" in reduced
125 - assert "GPU_to_CPU_total_bytes" in reduced
126 - assert "GPU_to_CPU_total_time" in reduced
127 - assert reduced["CPU_to_GPU_total_bytes"] == 34
128 - assert reduced["CPU_to_GPU_total_time"] == 2.6
129 - assert reduced["GPU_to_CPU_total_time"] == 2.3
130 - assert reduced["GPU_to_CPU_total_bytes"] == 19
272 + assert reduced[LOAD_BYTES] == 34
273 + assert reduced[LOAD_TIME] == 2.6
274 + assert reduced[f"{LOAD_SIZE}_count"] == 4
275 + assert reduced[f"{LOAD_SIZE}_sum"] == 34
276 + assert reduced[STORE_BYTES] == 19
277 + assert reduced[STORE_TIME] == 2.3
278 + assert reduced[f"{STORE_SIZE}_count"] == 3
279 + assert reduced[f"{STORE_SIZE}_sum"] == 19
280 + assert reduced[STORES_SKIPPED] == 11
281 + assert reduced[PENDING_STORES] == 2
282 + assert reduced[f"{LOOKUP_LATENCY}_count"] == 3
283 + assert reduced[f"{LOOKUP_LATENCY}_sum"] == sum([0.1, 0.2, 0.3])
284
285
286 def test_reset():
134 - """Test that reset() resets all nested connector stats."""
287 + """Test that reset() resets all connector stats."""
288 + metadata = _metric_metadata()
289 offload_connector_stats = OffloadingConnectorStats(
136 - data={
137 - "CPU_to_GPU": [
138 - {"op_size": 3, "op_time": 0.2},
139 - {"op_size": 7, "op_time": 0.9},
140 - ],
141 - "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
142 - }
290 + data=_make_stats_data(
291 + {
292 + LOAD_BYTES: 10,
293 + LOAD_TIME: 1.1,
294 + LOAD_SIZE: [3, 7],
295 + STORE_BYTES: 16,
296 + STORE_TIME: 2,
297 + STORE_SIZE: [16],
298 + STORES_SKIPPED: 4,
299 + PENDING_STORES: 2,
300 + LOOKUP_LATENCY: [0.1],
301 + },
302 + metadata,
303 + ),
304 )
305
306 assert not offload_connector_stats.is_empty()
@@ -148,4 +309,215 @@ def test_reset():
309
310 # After reset, stats should be empty
311 assert offload_connector_stats.is_empty()
151 - assert len(offload_connector_stats.data) == 0
312 +
313 +
314 +def test_prom_metrics_observes_manager_counter():
315 + prom_metrics = OffloadPromMetrics(
316 + vllm_config=_FakeVllmConfig(), # type: ignore[arg-type]
317 + metric_types={
318 + Gauge: _FakeMetric,
319 + Counter: _FakeMetric,
320 + Histogram: _FakeMetric,
321 + },
322 + labelnames=["model_name", "engine"],
323 + per_engine_labelvalues={0: ["model", "0"]},
324 + )
325 +
326 + prom_metrics.observe(
327 + {
328 + _StatsKey.TYPES: {STORES_SKIPPED: _MetricType.COUNTER},
329 + _StatsKey.DATA: {STORES_SKIPPED: 7},
330 + }
331 + )
332 +
333 + counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED)]
334 + assert counter.increments == [7]
335 + counter_def = prom_metrics._offloading_metric_defs[STORES_SKIPPED]
336 + assert counter_def.kwargs["name"] == "vllm:kv_offload_stores_skipped"
337 + assert counter.labelvalues == ("model", "0")
338 +
339 +
340 +def test_prom_metrics_observes_flat_transfer_metrics_and_legacy_metrics():
341 + prom_metrics = OffloadPromMetrics(
342 + vllm_config=_FakeVllmConfig(), # type: ignore[arg-type]
343 + metric_types={
344 + Gauge: _FakeMetric,
345 + Counter: _FakeMetric,
346 + Histogram: _FakeMetric,
347 + },
348 + labelnames=["model_name", "engine"],
349 + per_engine_labelvalues={0: ["model", "0"]},
350 + )
351 +
352 + prom_metrics.observe(
353 + {
354 + _StatsKey.TYPES: {
355 + LOAD_BYTES: _MetricType.COUNTER,
356 + LOAD_TIME: _MetricType.COUNTER,
357 + LOAD_SIZE: _MetricType.HISTOGRAM,
358 + STORE_BYTES: _MetricType.COUNTER,
359 + STORE_TIME: _MetricType.COUNTER,
360 + STORE_SIZE: _MetricType.HISTOGRAM,
361 + },
362 + _StatsKey.DATA: {
363 + LOAD_BYTES: 24,
364 + LOAD_TIME: 1.5,
365 + LOAD_SIZE: [16, 8],
366 + STORE_BYTES: 3,
367 + STORE_TIME: 0.3,
368 + STORE_SIZE: [1, 2],
369 + },
370 + }
371 + )
372 +
373 + assert prom_metrics.offloading_metrics[(0, LOAD_BYTES)].increments == [24]
374 + assert prom_metrics.offloading_metrics[(0, LOAD_TIME)].increments == [1.5]
375 + assert prom_metrics.offloading_metrics[(0, LOAD_SIZE)].observed == [16, 8]
376 + assert prom_metrics.offloading_metrics[(0, STORE_BYTES)].increments == [3]
377 + assert prom_metrics.offloading_metrics[(0, STORE_TIME)].increments == [0.3]
378 + assert prom_metrics.offloading_metrics[(0, STORE_SIZE)].observed == [1, 2]
379 +
380 + assert prom_metrics.counter_kv_bytes[(0, "CPU_to_GPU")].increments == [24]
381 + assert prom_metrics.counter_kv_transfer_time[(0, "CPU_to_GPU")].increments == [1.5]
382 + assert prom_metrics.histogram_transfer_size[(0, "CPU_to_GPU")].observed == [16, 8]
383 + assert prom_metrics.counter_kv_bytes[(0, "GPU_to_CPU")].increments == [3]
384 + assert prom_metrics.counter_kv_transfer_time[(0, "GPU_to_CPU")].increments == [0.3]
385 + assert prom_metrics.histogram_transfer_size[(0, "GPU_to_CPU")].observed == [1, 2]
386 +
387 +
388 +def test_prom_metrics_observes_manager_gauge_and_histogram():
389 + metric_definitions = {
390 + PENDING_STORES: OffloadingGaugeMetadata(
391 + documentation="Number of currently pending KV offload stores.",
392 + ),
393 + LOOKUP_LATENCY: OffloadingHistogramMetadata(
394 + documentation="KV offload lookup latency.",
395 + buckets=(0.1, 1.0),
396 + ),
397 + }
398 + with patch.object(
399 + CPUOffloadingSpec, "build_metric_definitions", return_value=metric_definitions
400 + ):
401 + prom_metrics = OffloadPromMetrics(
402 + vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type]
403 + metric_types={
404 + Gauge: _FakeMetric,
405 + Counter: _FakeMetric,
406 + Histogram: _FakeMetric,
407 + },
408 + labelnames=["model_name", "engine"],
409 + per_engine_labelvalues={0: ["model", "0"]},
410 + )
411 +
412 + prom_metrics.observe(
413 + {
414 + _StatsKey.TYPES: {
415 + PENDING_STORES: _MetricType.GAUGE,
416 + LOOKUP_LATENCY: _MetricType.HISTOGRAM,
417 + },
418 + _StatsKey.DATA: {
419 + PENDING_STORES: 5,
420 + LOOKUP_LATENCY: [0.2, 0.4],
421 + },
422 + }
423 + )
424 +
425 + gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES)]
426 + histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY)]
427 + assert gauge.set_values == [5]
428 + assert histogram.observed == [0.2, 0.4]
429 + histogram_def = prom_metrics._offloading_metric_defs[LOOKUP_LATENCY]
430 + assert histogram_def.kwargs["buckets"] == (0.1, 1.0)
431 +
432 +
433 +def test_prom_metrics_uses_configured_manager_metrics():
434 + prom_metrics = OffloadPromMetrics(
435 + vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type]
436 + metric_types={
437 + Gauge: _FakeMetric,
438 + Counter: _FakeMetric,
439 + Histogram: _FakeMetric,
440 + },
441 + labelnames=["model_name", "engine"],
442 + per_engine_labelvalues={0: ["model", "0"]},
443 + )
444 +
445 + assert STORES_SKIPPED not in prom_metrics._offloading_metric_metadata
446 +
447 +
448 +def test_aggregate_into_empty_stats():
449 + """Aggregating non-empty stats into a fresh (empty) stats object works."""
450 + empty = OffloadingConnectorStats()
451 + assert empty.is_empty()
452 +
453 + non_empty = OffloadingConnectorStats(
454 + data={
455 + _StatsKey.TYPES: {
456 + LOAD_BYTES: _MetricType.COUNTER,
457 + LOAD_SIZE: _MetricType.HISTOGRAM,
458 + PENDING_STORES: _MetricType.GAUGE,
459 + },
460 + _StatsKey.DATA: {
461 + LOAD_BYTES: 42,
462 + LOAD_SIZE: [10, 20],
463 + PENDING_STORES: 3,
464 + },
465 + },
466 + )
467 +
468 + result = empty.aggregate(non_empty)
469 +
470 + assert result is empty
471 + values = result.data[_StatsKey.DATA]
472 + assert values[LOAD_BYTES] == 42
473 + assert values[LOAD_SIZE] == [10, 20]
474 + assert values[PENDING_STORES] == 3
475 +
476 +
477 +def test_prom_metrics_multi_engine_routing():
478 + """Metrics are routed to the correct engine index."""
479 + prom_metrics = OffloadPromMetrics(
480 + vllm_config=_FakeVllmConfig(), # type: ignore[arg-type]
481 + metric_types={
482 + Gauge: _FakeMetric,
483 + Counter: _FakeMetric,
484 + Histogram: _FakeMetric,
485 + },
486 + labelnames=["model_name", "engine"],
487 + per_engine_labelvalues={0: ["model", "0"], 1: ["model", "1"]},
488 + )
489 +
490 + prom_metrics.observe(
491 + {
492 + _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER},
493 + _StatsKey.DATA: {LOAD_BYTES: 100},
494 + },
495 + engine_idx=1,
496 + )
497 +
498 + engine0 = prom_metrics.offloading_metrics[(0, LOAD_BYTES)]
499 + engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES)]
500 + assert engine0.increments == []
501 + assert engine1.increments == [100]
502 +
503 +
504 +def test_prom_metrics_rejects_undeclared_metric():
505 + """observe() asserts if a metric was never declared in metadata."""
506 + prom_metrics = OffloadPromMetrics(
507 + vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type]
508 + metric_types={
509 + Gauge: _FakeMetric,
510 + Counter: _FakeMetric,
511 + Histogram: _FakeMetric,
512 + },
513 + labelnames=["model_name", "engine"],
514 + per_engine_labelvalues={0: ["model", "0"]},
515 + )
516 +
517 + with pytest.raises(AssertionError):
518 + prom_metrics.observe(
519 + {
520 + _StatsKey.TYPES: {"unknown:metric": _MetricType.COUNTER},
521 + _StatsKey.DATA: {"unknown:metric": 1},
522 + }
523 + )
tests/v1/kv_connector/unit/offloading_connector/test_worker_metadata.py
+21
@@ -4,7 +4,9 @@
4 import pytest
5
6 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
7 + DirectionalTransferStats,
8 OffloadingWorkerMetadata,
9 + TransferStats,
10 )
11
12 pytestmark = pytest.mark.cpu_test
@@ -30,3 +32,22 @@ def test_aggregate_multiple_workers():
32 meta3 = OffloadingWorkerMetadata(completed_jobs={42: 1, 43: 1, 8: 1})
33 result = meta1.aggregate(meta2).aggregate(meta3)
34 assert result.completed_jobs == {42: 3, 43: 2, 7: 2, 8: 2}
35 +
36 +
37 +def test_aggregate_transfer_stats():
38 + meta1 = OffloadingWorkerMetadata(
39 + transfer_stats=TransferStats(
40 + load=DirectionalTransferStats(bytes=10, time=0.5, sizes=[10])
41 + )
42 + )
43 + meta2 = OffloadingWorkerMetadata(
44 + transfer_stats=TransferStats(
45 + load=DirectionalTransferStats(bytes=20, time=1.0, sizes=[20, 30])
46 + )
47 + )
48 +
49 + result = meta1.aggregate(meta2)
50 +
51 + assert result.transfer_stats.load.bytes == 30
52 + assert result.transfer_stats.load.time == 1.5
53 + assert result.transfer_stats.load.sizes == [10, 20, 30]
tests/v1/kv_connector/unit/test_offloading_connector.py
+162
@@ -5,6 +5,7 @@ import time
5
6 import msgspec
7 import msgspec.msgpack
8 +import prometheus_client
9 import pytest
10 import zmq
11 from tqdm import tqdm
@@ -303,6 +304,167 @@ def test_cpu_offloading(
304 del llm
305
306
307 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
308 +def test_cpu_offloading_metrics() -> None:
309 + """Verify that offloading Prometheus metrics (new flat and deprecated
310 + labeled) are emitted after stores and loads."""
311 + extra_config: dict = {
312 + "cpu_bytes_to_use": 500 << 20,
313 + "block_size": CPU_BLOCK_SIZES,
314 + }
315 + kv_transfer_config = KVTransferConfig(
316 + kv_connector="OffloadingConnector",
317 + kv_role="kv_both",
318 + kv_connector_extra_config=extra_config,
319 + )
320 +
321 + llm = LLM(
322 + model="meta-llama/Llama-3.2-1B-Instruct",
323 + max_model_len=4096,
324 + gpu_memory_utilization=0.5,
325 + kv_transfer_config=kv_transfer_config,
326 + disable_log_stats=False,
327 + )
328 +
329 + try:
330 + prompt_token_ids = list(range(500))
331 +
332 + # First generate: cold run, triggers a store to CPU.
333 + # Use max_tokens>1 so the request is still producing output
334 + # tokens when the async store completes and stats get drained.
335 + # (The LLMEngine only records stats on steps with request outputs.)
336 + llm.generate(
337 + [TokensPrompt(prompt_token_ids=prompt_token_ids)],
338 + SamplingParams(max_tokens=10),
339 + use_tqdm=False,
340 + )
341 +
342 + # Wait for the async offload to finish, then reset GPU prefix cache
343 + # so the next generate must load from CPU.
344 + _wait_for_prefix_cache_reset(llm)
345 +
346 + # Second generate: triggers a load from CPU.
347 + # Send a short filler alongside the load prompt so there's always
348 + # a request producing output tokens when the load stats get drained
349 + # (the LLMEngine only records stats on steps with request outputs).
350 + filler = TokensPrompt(prompt_token_ids=[0])
351 + llm.generate(
352 + [filler, TokensPrompt(prompt_token_ids=prompt_token_ids)],
353 + SamplingParams(max_tokens=50),
354 + use_tqdm=False,
355 + )
356 +
357 + # Metric helpers.
358 + registry = prometheus_client.REGISTRY
359 +
360 + def _get_counter_value(
361 + name: str, labels: dict[str, str] | None = None
362 + ) -> float:
363 + total = 0.0
364 + for metric in registry.collect():
365 + if metric.name == name:
366 + for sample in metric.samples:
367 + if sample.name != name + "_total":
368 + continue
369 + if labels and not all(
370 + sample.labels.get(k) == v for k, v in labels.items()
371 + ):
372 + continue
373 + total += sample.value
374 + return total
375 +
376 + def _get_histogram_count(
377 + name: str, labels: dict[str, str] | None = None
378 + ) -> float:
379 + total = 0.0
380 + for metric in registry.collect():
381 + if metric.name == name:
382 + for sample in metric.samples:
383 + if sample.name != name + "_count":
384 + continue
385 + if labels and not all(
386 + sample.labels.get(k) == v for k, v in labels.items()
387 + ):
388 + continue
389 + total += sample.value
390 + return total
391 +
392 + # Stats are drained asynchronously — if the transfer finishes
393 + # after the last engine step for that generate() call, the metrics
394 + # won't appear until a subsequent step. Retry with dummy generates
395 + # to force additional stats drains.
396 + deadline = time.monotonic() + _RESET_CACHE_TIMEOUT
397 + while time.monotonic() < deadline:
398 + store_bytes = _get_counter_value("vllm:kv_offload_store_bytes")
399 + load_bytes = _get_counter_value("vllm:kv_offload_load_bytes")
400 + if store_bytes > 0 and load_bytes > 0:
401 + break
402 + llm.generate(
403 + [TokensPrompt(prompt_token_ids=[0])],
404 + SamplingParams(max_tokens=1),
405 + use_tqdm=False,
406 + )
407 +
408 + # New flat counter metrics
409 + store_bytes = _get_counter_value("vllm:kv_offload_store_bytes")
410 + assert store_bytes > 0, f"Expected store_bytes > 0, got {store_bytes}"
411 + load_bytes = _get_counter_value("vllm:kv_offload_load_bytes")
412 + assert load_bytes > 0, f"Expected load_bytes > 0, got {load_bytes}"
413 + store_time = _get_counter_value("vllm:kv_offload_store_time")
414 + assert store_time > 0, f"Expected store_time > 0, got {store_time}"
415 + load_time = _get_counter_value("vllm:kv_offload_load_time")
416 + assert load_time > 0, f"Expected load_time > 0, got {load_time}"
417 +
418 + # New flat histogram metrics
419 + store_size_count = _get_histogram_count("vllm:kv_offload_store_size")
420 + assert store_size_count > 0, (
421 + f"Expected store_size histogram observations > 0, got {store_size_count}"
422 + )
423 + load_size_count = _get_histogram_count("vllm:kv_offload_load_size")
424 + assert load_size_count > 0, (
425 + f"Expected load_size histogram observations > 0, got {load_size_count}"
426 + )
427 +
428 + # Deprecated labeled metrics — verify per transfer_type label.
429 + load_label = {"transfer_type": "CPU_to_GPU"}
430 + store_label = {"transfer_type": "GPU_to_CPU"}
431 +
432 + dep_load_bytes = _get_counter_value("vllm:kv_offload_total_bytes", load_label)
433 + assert dep_load_bytes > 0, (
434 + f"Expected deprecated load bytes > 0, got {dep_load_bytes}"
435 + )
436 + dep_store_bytes = _get_counter_value("vllm:kv_offload_total_bytes", store_label)
437 + assert dep_store_bytes > 0, (
438 + f"Expected deprecated store bytes > 0, got {dep_store_bytes}"
439 + )
440 + dep_load_time = _get_counter_value("vllm:kv_offload_total_time", load_label)
441 + assert dep_load_time > 0, (
442 + f"Expected deprecated load time > 0, got {dep_load_time}"
443 + )
444 + dep_store_time = _get_counter_value("vllm:kv_offload_total_time", store_label)
445 + assert dep_store_time > 0, (
446 + f"Expected deprecated store time > 0, got {dep_store_time}"
447 + )
448 + dep_load_size = _get_histogram_count("vllm:kv_offload_size", load_label)
449 + assert dep_load_size > 0, (
450 + f"Expected deprecated load size observations > 0, got {dep_load_size}"
451 + )
452 + dep_store_size = _get_histogram_count("vllm:kv_offload_size", store_label)
453 + assert dep_store_size > 0, (
454 + f"Expected deprecated store size observations > 0, got {dep_store_size}"
455 + )
456 +
457 + # Flat and deprecated metrics must be consistent (dual-write).
458 + assert store_bytes == dep_store_bytes
459 + assert load_bytes == dep_load_bytes
460 + assert store_time == dep_store_time
461 + assert load_time == dep_load_time
462 + assert store_size_count == dep_store_size
463 + assert load_size_count == dep_load_size
464 + finally:
465 + del llm
466 +
467 +
468 def test_tiering_offloading() -> None:
469 """Tests OffloadingConnector with TieringOffloadingSpec."""
470 extra_config: dict = {
tests/v1/kv_offload/cpu/test_manager.py
+48 -7
@@ -18,6 +18,8 @@ from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
18 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
19 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
20
21 +STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
22 +
23
24 def make_req_context(
25 req_id: str = "", kv_transfer_params: dict | None = None
@@ -29,6 +31,22 @@ def make_req_context(
31 _EMPTY_REQ_CTX = make_req_context()
32
33
34 +def make_cpu_manager(
35 + num_blocks: int = 4,
36 + cache_policy: str = "lru",
37 + enable_events: bool = False,
38 + store_threshold: int = 0,
39 + max_tracker_size: int = 64_000,
40 +) -> CPUOffloadingManager:
41 + return CPUOffloadingManager(
42 + num_blocks=num_blocks,
43 + cache_policy=cache_policy,
44 + enable_events=enable_events,
45 + store_threshold=store_threshold,
46 + max_tracker_size=max_tracker_size,
47 + )
48 +
49 +
50 @dataclass
51 class ExpectedPrepareStoreOutput:
52 keys_to_store: list[int]
@@ -110,7 +128,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
128 candidate to make room for [3, 4, 5]
129 - After complete_store([2, 3, 4, 5]), block 2 must still be present.
130 """
113 - manager = CPUOffloadingManager(
131 + manager = make_cpu_manager(
132 num_blocks=4,
133 cache_policy=eviction_policy,
134 enable_events=True,
@@ -144,14 +162,37 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
162 assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True
163
164
165 +def test_filter_reused_manager_reports_stores_skipped_counter():
166 + manager = make_cpu_manager(
167 + num_blocks=4,
168 + cache_policy="lru",
169 + store_threshold=2,
170 + )
171 +
172 + prepare_store_output = manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX)
173 +
174 + verify_store_output(
175 + prepare_store_output,
176 + ExpectedPrepareStoreOutput(
177 + keys_to_store=[],
178 + store_block_ids=[],
179 + evicted_keys=[],
180 + ),
181 + )
182 + stats = manager.get_stats()
183 + assert stats is not None
184 + assert stats.reduce()[STORES_SKIPPED] == 3
185 + stats = manager.get_stats()
186 + assert stats is not None
187 + assert stats.reduce()[STORES_SKIPPED] == 0
188 +
189 +
190 def test_cpu_manager():
191 """
192 Tests CPUOffloadingManager with lru policy.
193 """
194 # initialize a CPU manager with a capacity of 4 blocks
152 - cpu_manager = CPUOffloadingManager(
153 - num_blocks=4, cache_policy="lru", enable_events=True
154 - )
195 + cpu_manager = make_cpu_manager(num_blocks=4, cache_policy="lru", enable_events=True)
196
197 # prepare store [1, 2]
198 prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
@@ -264,7 +305,7 @@ def test_cpu_manager():
305
306 def test_prepare_load_preserves_key_order():
307 """block_ids[i] must correspond to keys[i] (co-indexed invariant)."""
267 - manager = CPUOffloadingManager(num_blocks=4, cache_policy="lru")
308 + manager = make_cpu_manager(num_blocks=4, cache_policy="lru")
309
310 key_a, key_b, key_c = to_key(0), to_key(1), to_key(2)
311
@@ -305,7 +346,7 @@ class TestARCPolicy:
346 def _make_manager(
347 self, num_blocks: int = 4, enable_events: bool = True
348 ) -> tuple[CPUOffloadingManager, ARCCachePolicy]:
308 - manager = CPUOffloadingManager(
349 + manager = make_cpu_manager(
350 num_blocks=num_blocks,
351 cache_policy="arc",
352 enable_events=enable_events,
@@ -605,7 +646,7 @@ def test_filter_reused_manager():
646 """
647 Tests CPUOffloadingManager reuse filtering (store_threshold=2).
648 """
608 - manager = CPUOffloadingManager(
649 + manager = make_cpu_manager(
650 num_blocks=4,
651 cache_policy="lru",
652 enable_events=True,
vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py
+44 -1
@@ -11,6 +11,45 @@ from vllm.v1.kv_offload.worker.worker import TransferSpec
11 ReqId = str
12
13
14 +@dataclass(slots=True)
15 +class DirectionalTransferStats:
16 + bytes: int = 0
17 + time: float = 0.0
18 + sizes: list[int | float] = field(default_factory=list)
19 +
20 + def aggregate(
21 + self, other: "DirectionalTransferStats"
22 + ) -> "DirectionalTransferStats":
23 + return DirectionalTransferStats(
24 + bytes=self.bytes + other.bytes,
25 + time=self.time + other.time,
26 + sizes=[*self.sizes, *other.sizes],
27 + )
28 +
29 + def record(self, num_bytes: int, time: float) -> None:
30 + self.bytes += num_bytes
31 + self.time += time
32 + self.sizes.append(num_bytes)
33 +
34 + def is_empty(self) -> bool:
35 + return self.bytes == 0 and self.time == 0.0 and not self.sizes
36 +
37 +
38 +@dataclass(slots=True)
39 +class TransferStats:
40 + load: DirectionalTransferStats = field(default_factory=DirectionalTransferStats)
41 + store: DirectionalTransferStats = field(default_factory=DirectionalTransferStats)
42 +
43 + def aggregate(self, other: "TransferStats") -> "TransferStats":
44 + return TransferStats(
45 + load=self.load.aggregate(other.load),
46 + store=self.store.aggregate(other.store),
47 + )
48 +
49 + def is_empty(self) -> bool:
50 + return self.load.is_empty() and self.store.is_empty()
51 +
52 +
53 @dataclass
54 class TransferJob:
55 """A transfer job bundling request context with transfer spec.
@@ -43,6 +82,7 @@ class OffloadingWorkerMetadata(KVConnectorWorkerMetadata):
82 """
83
84 completed_jobs: dict[int, int] = field(default_factory=dict)
85 + transfer_stats: TransferStats = field(default_factory=TransferStats)
86
87 def mark_completed(self, job_id: int) -> None:
88 """Record a transfer job completion from this worker."""
@@ -57,4 +97,7 @@ class OffloadingWorkerMetadata(KVConnectorWorkerMetadata):
97 for job_id, v in other.completed_jobs.items():
98 merged[job_id] = merged.get(job_id, 0) + v
99
60 - return OffloadingWorkerMetadata(completed_jobs=merged)
100 + return OffloadingWorkerMetadata(
101 + completed_jobs=merged,
102 + transfer_stats=self.transfer_stats.aggregate(other.transfer_stats),
103 + )
vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py
+310 -88
@@ -10,37 +10,176 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
10 PromMetric,
11 PromMetricT,
12 )
13 -from vllm.logger import init_logger
14 -from vllm.v1.kv_offload.worker.worker import TransferType
13 +from vllm.v1.kv_offload.base import (
14 + OffloadingCounterMetadata,
15 + OffloadingGaugeMetadata,
16 + OffloadingHistogramMetadata,
17 + OffloadingMetricMetadata,
18 +)
19 +from vllm.v1.kv_offload.factory import OffloadingSpecFactory
20
16 -logger = init_logger(__name__)
21
22 +class _TransferMetricName:
23 + """Flat metric names for GPU↔offload-medium transfer operations."""
24 +
25 + LOAD_BYTES = "vllm:kv_offload_load_bytes"
26 + LOAD_TIME = "vllm:kv_offload_load_time"
27 + LOAD_SIZE = "vllm:kv_offload_load_size"
28 + STORE_BYTES = "vllm:kv_offload_store_bytes"
29 + STORE_TIME = "vllm:kv_offload_store_time"
30 + STORE_SIZE = "vllm:kv_offload_store_size"
31 +
32 +
33 +class _TransferType:
34 + """Transfer direction labels for deprecated CPU offload metrics."""
35 +
36 + LOAD = "CPU_to_GPU"
37 + STORE = "GPU_to_CPU"
38 + ALL = (LOAD, STORE)
39 +
40 +
41 +TRANSFER_SIZE_BUCKETS = (
42 + 1e6,
43 + 5e6,
44 + 10e6,
45 + 20e6,
46 + 40e6,
47 + 60e6,
48 + 80e6,
49 + 100e6,
50 + 150e6,
51 + 200e6,
52 +)
53 +
54 +
55 +def get_connector_metric_definitions() -> dict[str, OffloadingMetricMetadata]:
56 + return {
57 + _TransferMetricName.LOAD_BYTES: OffloadingCounterMetadata(
58 + documentation="Total bytes loaded from offload storage to GPU.",
59 + ),
60 + _TransferMetricName.LOAD_TIME: OffloadingCounterMetadata(
61 + documentation="Total load time from offload storage to GPU, in seconds.",
62 + ),
63 + _TransferMetricName.LOAD_SIZE: OffloadingHistogramMetadata(
64 + documentation="Histogram of KV offload load operation size, in bytes.",
65 + buckets=TRANSFER_SIZE_BUCKETS,
66 + ),
67 + _TransferMetricName.STORE_BYTES: OffloadingCounterMetadata(
68 + documentation="Total bytes stored from GPU to offload storage.",
69 + ),
70 + _TransferMetricName.STORE_TIME: OffloadingCounterMetadata(
71 + documentation="Total store time from GPU to offload storage, in seconds.",
72 + ),
73 + _TransferMetricName.STORE_SIZE: OffloadingHistogramMetadata(
74 + documentation="Histogram of KV offload store operation size, in bytes.",
75 + buckets=TRANSFER_SIZE_BUCKETS,
76 + ),
77 + }
78
19 -@dataclass
20 -class OffloadingOperationMetrics:
21 - op_size: int
22 - op_time: float
79 +
80 +_DEPRECATED_TOTAL_BYTES = "vllm:kv_offload_total_bytes"
81 +_DEPRECATED_TOTAL_TIME = "vllm:kv_offload_total_time"
82 +_DEPRECATED_SIZE = "vllm:kv_offload_size"
83 +
84 +# Deprecated legacy transfer metrics, kept during the migration to the flat
85 +# metric names above. These stay in a separate definition block because they
86 +# use a transfer_type label, but are emitted from the same flat stats payload
87 +# for compatibility.
88 +_DEPRECATED_CONNECTOR_METRIC_DEFINITIONS: dict[str, OffloadingMetricMetadata] = {
89 + _DEPRECATED_TOTAL_BYTES: OffloadingCounterMetadata(
90 + documentation="Number of bytes offloaded by KV connector",
91 + ),
92 + _DEPRECATED_TOTAL_TIME: OffloadingCounterMetadata(
93 + documentation="Total time measured by all KV offloading operations",
94 + ),
95 + _DEPRECATED_SIZE: OffloadingHistogramMetadata(
96 + documentation="Histogram of KV offload transfer size, in bytes.",
97 + buckets=TRANSFER_SIZE_BUCKETS,
98 + ),
99 +}
100 +
101 +
102 +class _MetricType:
103 + """Type tags embedded in the serialized stats payload."""
104 +
105 + COUNTER = "counter"
106 + GAUGE = "gauge"
107 + HISTOGRAM = "histogram"
108 +
109 +
110 +class _StatsKey:
111 + """Top-level keys in the serialized stats dict."""
112 +
113 + # Maps metric name -> _MetricType value
114 + TYPES = "types"
115 + # Maps metric name -> observed value (number or list)
116 + DATA = "data"
117
118
119 @dataclass
120 class OffloadingConnectorStats(KVConnectorStats):
121 + """
122 + Offloading connector stats use flat metric names as keys.
123 +
124 + The ``data`` dict is structured using ``_StatsKey`` / ``_MetricType``::
125 +
126 + {
127 + _StatsKey.TYPES: {name: _MetricType.*, ...},
128 + _StatsKey.DATA: {name: value, ...},
129 + }
130 +
131 + This structure is self-describing: it survives IPC serialization
132 + without needing the full ``OffloadingMetricMetadata`` objects on the
133 + receiving side.
134 +
135 + Counter values are aggregated by summing, gauge values use the latest
136 + snapshot, and histogram values are lists of observed samples.
137 + """
138 +
139 def __post_init__(self):
28 - if not self.data:
29 - # Empty container init, no data is passed in.
140 + if _StatsKey.DATA not in self.data:
141 self.reset()
142
143 def reset(self):
33 - self.data: dict[str, list[OffloadingOperationMetrics]] = {}
144 + self.data: dict[str, Any] = {
145 + _StatsKey.TYPES: {},
146 + _StatsKey.DATA: {},
147 + }
148 +
149 + @property
150 + def _types(self) -> dict[str, str]:
151 + return self.data[_StatsKey.TYPES]
152
35 - def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
36 - if not other.is_empty():
37 - for k, v in other.data.items():
38 - if k not in self.data:
39 - self.data[k] = v
153 + @property
154 + def _values(self) -> dict[str, Any]:
155 + return self.data[_StatsKey.DATA]
156 +
157 + def aggregate(self, other: "KVConnectorStats") -> "KVConnectorStats":
158 + if other.is_empty():
159 + return self
160 + assert isinstance(other, OffloadingConnectorStats)
161 + other_types = other._types
162 + other_values = other._values
163 + for key, value in other_values.items():
164 + type_str = other_types.get(key)
165 + if type_str is None:
166 + raise AssertionError(f"Unknown offloading stats key: {key}")
167 + self._types.setdefault(key, type_str)
168 + if type_str == _MetricType.HISTOGRAM:
169 + assert isinstance(value, list)
170 + if key not in self._values:
171 + self._values[key] = value
172 else:
41 - accumulator = self.data[k]
42 - assert isinstance(accumulator, list)
43 - accumulator.extend(v)
173 + assert isinstance(self._values[key], list)
174 + self._values[key].extend(value)
175 + elif type_str == _MetricType.COUNTER:
176 + assert isinstance(value, int | float)
177 + self._values[key] = self._values.get(key, 0) + value
178 + elif type_str == _MetricType.GAUGE:
179 + assert isinstance(value, int | float)
180 + self._values[key] = value
181 + else:
182 + raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
183 return self
184
185 def reduce(self) -> dict[str, int | float]:
@@ -51,29 +190,44 @@ class OffloadingConnectorStats(KVConnectorStats):
190 stats for the last time interval.
191 """
192 return_dict: dict[str, int | float] = {}
54 - for transfer_type, ops_list in self.data.items():
55 - assert isinstance(ops_list, list)
56 - total_bytes = 0
57 - total_time = 0.0
58 - for op in ops_list:
59 - assert isinstance(op, dict)
60 - total_bytes += op["op_size"]
61 - total_time += op["op_time"]
62 - return_dict[f"{transfer_type}_total_bytes"] = total_bytes
63 - return_dict[f"{transfer_type}_total_time"] = total_time
193 + for key, value in self._values.items():
194 + type_str = self._types.get(key)
195 + if type_str is None:
196 + raise AssertionError(f"Unknown offloading stats key: {key}")
197 + if type_str == _MetricType.HISTOGRAM:
198 + assert isinstance(value, list)
199 + return_dict[f"{key}_count"] = len(value)
200 + return_dict[f"{key}_sum"] = sum(value)
201 + elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE):
202 + assert isinstance(value, int | float)
203 + return_dict[key] = value
204 + else:
205 + raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
206 return return_dict
207
208 def is_empty(self) -> bool:
67 - return not self.data
68 -
69 - def record_transfer(self, num_bytes: int, time: float, transfer_type: TransferType):
70 - src, dst = transfer_type
71 - transfer_type_key = src + "_to_" + dst
72 - op = OffloadingOperationMetrics(num_bytes, time)
73 - if transfer_type_key in self.data:
74 - self.data[transfer_type_key].append(op)
75 - else:
76 - self.data[transfer_type_key] = [op]
209 + return not self.data.get(_StatsKey.DATA)
210 +
211 + def increase_counter(
212 + self, counter_name: str, counter_increase_value: int | float
213 + ) -> None:
214 + """Increase a counter on the stats payload."""
215 + self._types.setdefault(counter_name, _MetricType.COUNTER)
216 + self._values[counter_name] = (
217 + self._values.get(counter_name, 0) + counter_increase_value
218 + )
219 +
220 + def set_gauge(self, gauge_name: str, gauge_value: int | float) -> None:
221 + """Set a gauge snapshot on the stats payload."""
222 + self._types.setdefault(gauge_name, _MetricType.GAUGE)
223 + self._values[gauge_name] = gauge_value
224 +
225 + def observe_histogram(
226 + self, histogram_name: str, histogram_value: int | float
227 + ) -> None:
228 + """Record a histogram observation on the stats payload."""
229 + self._types.setdefault(histogram_name, _MetricType.HISTOGRAM)
230 + self._values.setdefault(histogram_name, []).append(histogram_value)
231
232
233 class OffloadPromMetrics(KVConnectorPromMetrics):
@@ -89,77 +243,145 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
243 self.histogram_transfer_size: dict[tuple[int, str], PromMetricT] = {}
244 self.counter_kv_bytes: dict[tuple[int, str], PromMetricT] = {}
245 self.counter_kv_transfer_time: dict[tuple[int, str], PromMetricT] = {}
92 - buckets = [ # In bytes
93 - 1e6,
94 - 5e6,
95 - 10e6,
96 - 20e6,
97 - 40e6,
98 - 60e6,
99 - 80e6,
100 - 100e6,
101 - 150e6,
102 - 200e6,
103 - ]
246 + spec_cls = OffloadingSpecFactory.get_spec_cls(vllm_config)
247 + kv_transfer_config = vllm_config.kv_transfer_config
248 + assert kv_transfer_config is not None
249 + extra_config = kv_transfer_config.kv_connector_extra_config
250 + self._offloading_metric_metadata: dict[str, OffloadingMetricMetadata] = {
251 + **spec_cls.build_metric_definitions(extra_config),
252 + **get_connector_metric_definitions(),
253 + }
254 + from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
255 +
256 + self._observe_deprecated_metrics = issubclass(spec_cls, CPUOffloadingSpec)
257 + self._offloading_metric_defs: dict[str, PromMetricT] = {}
258 + self.offloading_metrics: dict[tuple[int, str], PromMetricT] = {}
259
260 self._counter_kv_bytes = self._counter_cls(
106 - name="vllm:kv_offload_total_bytes",
107 - documentation="Number of bytes offloaded by KV connector",
261 + name=_DEPRECATED_TOTAL_BYTES,
262 + documentation=_DEPRECATED_CONNECTOR_METRIC_DEFINITIONS[
263 + _DEPRECATED_TOTAL_BYTES
264 + ].documentation,
265 labelnames=labelnames + ["transfer_type"],
266 )
267
268 self._counter_kv_transfer_time = self._counter_cls(
112 - name="vllm:kv_offload_total_time",
113 - documentation="Total time measured by all KV offloading operations",
269 + name=_DEPRECATED_TOTAL_TIME,
270 + documentation=_DEPRECATED_CONNECTOR_METRIC_DEFINITIONS[
271 + _DEPRECATED_TOTAL_TIME
272 + ].documentation,
273 labelnames=labelnames + ["transfer_type"],
274 )
275
276 + deprecated_size_metadata = _DEPRECATED_CONNECTOR_METRIC_DEFINITIONS[
277 + _DEPRECATED_SIZE
278 + ]
279 + assert isinstance(deprecated_size_metadata, OffloadingHistogramMetadata)
280 self._histogram_transfer_size = self._histogram_cls(
118 - name="vllm:kv_offload_size",
119 - documentation="Histogram of KV offload transfer size, in bytes.",
120 - buckets=buckets[:],
281 + name=_DEPRECATED_SIZE,
282 + documentation=deprecated_size_metadata.documentation,
283 + buckets=deprecated_size_metadata.buckets,
284 labelnames=labelnames + ["transfer_type"],
285 )
286
124 - def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
125 - """
126 - Observe transfer statistics from the new data structure.
127 - transfer_stats_data is expected to be a dict where:
128 - - keys are transfer type strings (e.g., "cpu_to_gpu", "gpu_to_cpu")
129 - - values are lists of OffloadingOperationMetrics objects
130 - """
131 -
132 - for transfer_type, ops in transfer_stats_data.items():
133 - # Cache:
134 - if (engine_idx, transfer_type) not in self.histogram_transfer_size:
287 + for engine_idx, labelvalues in per_engine_labelvalues.items():
288 + for transfer_type in _TransferType.ALL:
289 + bounded_labelvalues = labelvalues + [transfer_type]
290 self.histogram_transfer_size[(engine_idx, transfer_type)] = (
136 - self._histogram_transfer_size.labels(
137 - *(self.per_engine_labelvalues[engine_idx] + [transfer_type])
138 - )
291 + self._histogram_transfer_size.labels(*bounded_labelvalues)
292 )
293 self.counter_kv_bytes[(engine_idx, transfer_type)] = (
141 - self._counter_kv_bytes.labels(
142 - *(self.per_engine_labelvalues[engine_idx] + [transfer_type])
143 - )
294 + self._counter_kv_bytes.labels(*bounded_labelvalues)
295 )
296 self.counter_kv_transfer_time[(engine_idx, transfer_type)] = (
146 - self._counter_kv_transfer_time.labels(
147 - *(self.per_engine_labelvalues[engine_idx] + [transfer_type])
148 - )
297 + self._counter_kv_transfer_time.labels(*bounded_labelvalues)
298 )
299
151 - # Process ops:
152 - assert isinstance(ops, list)
153 - for op in ops: # ops is a list of serialized OffloadingOperationMetrics
154 - assert isinstance(op, dict)
155 - # Observe size histogram
156 - self.histogram_transfer_size[(engine_idx, transfer_type)].observe(
157 - op["op_size"]
300 + for metric_name, metadata in self._offloading_metric_metadata.items():
301 + self._offloading_metric_defs[metric_name] = self._create_metric(
302 + metric_name, metadata
303 + )
304 + for engine_idx, labelvalues in per_engine_labelvalues.items():
305 + self.offloading_metrics[(engine_idx, metric_name)] = (
306 + self._offloading_metric_defs[metric_name].labels(*labelvalues)
307 )
308
160 - # Increment byte and time counters
161 - self.counter_kv_bytes[(engine_idx, transfer_type)].inc(op["op_size"])
309 + def _create_metric(
310 + self, metric_name: str, metadata: OffloadingMetricMetadata
311 + ) -> Any:
312 + kwargs: dict[str, Any] = {
313 + "name": metric_name,
314 + "documentation": metadata.documentation,
315 + "labelnames": self._labelnames,
316 + }
317 + if isinstance(metadata, OffloadingCounterMetadata):
318 + metric_cls = self._counter_cls
319 + elif isinstance(metadata, OffloadingGaugeMetadata):
320 + metric_cls = self._gauge_cls
321 + elif isinstance(metadata, OffloadingHistogramMetadata):
322 + metric_cls = self._histogram_cls
323 + if metadata.buckets is not None:
324 + kwargs["buckets"] = metadata.buckets
325 + else:
326 + raise AssertionError(f"Unknown offloading metric metadata: {metadata}")
327 + return metric_cls(**kwargs)
328
163 - self.counter_kv_transfer_time[(engine_idx, transfer_type)].inc(
164 - op["op_time"]
329 + def _increase_counter(
330 + self, metric_name: str, value: int | float, engine_idx: int
331 + ) -> None:
332 + self.offloading_metrics[(engine_idx, metric_name)].inc(value)
333 + if not self._observe_deprecated_metrics:
334 + return
335 + # Keep deprecated CPU offload transfer metrics updated during the
336 + # transition to flat metric names.
337 + if metric_name == _TransferMetricName.LOAD_BYTES:
338 + self.counter_kv_bytes[(engine_idx, _TransferType.LOAD)].inc(value)
339 + elif metric_name == _TransferMetricName.LOAD_TIME:
340 + self.counter_kv_transfer_time[(engine_idx, _TransferType.LOAD)].inc(value)
341 + elif metric_name == _TransferMetricName.STORE_BYTES:
342 + self.counter_kv_bytes[(engine_idx, _TransferType.STORE)].inc(value)
343 + elif metric_name == _TransferMetricName.STORE_TIME:
344 + self.counter_kv_transfer_time[(engine_idx, _TransferType.STORE)].inc(value)
345 +
346 + def _set_gauge(self, metric_name: str, value: int | float, engine_idx: int) -> None:
347 + self.offloading_metrics[(engine_idx, metric_name)].set(value)
348 +
349 + def _observe_histogram(
350 + self, metric_name: str, value: list[int | float], engine_idx: int
351 + ) -> None:
352 + for observation in value:
353 + self.offloading_metrics[(engine_idx, metric_name)].observe(observation)
354 + if not self._observe_deprecated_metrics:
355 + continue
356 + # Keep deprecated CPU offload transfer metrics updated during the
357 + # transition to flat metric names.
358 + if metric_name == _TransferMetricName.LOAD_SIZE:
359 + self.histogram_transfer_size[(engine_idx, _TransferType.LOAD)].observe(
360 + observation
361 + )
362 + elif metric_name == _TransferMetricName.STORE_SIZE:
363 + self.histogram_transfer_size[(engine_idx, _TransferType.STORE)].observe(
364 + observation
365 )
366 +
367 + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
368 + """Observe transfer statistics."""
369 + metric_types = transfer_stats_data.get(_StatsKey.TYPES, {})
370 + metric_data = transfer_stats_data.get(_StatsKey.DATA, {})
371 + for key, value in metric_data.items():
372 + type_str = metric_types.get(key)
373 + if type_str is None:
374 + raise AssertionError(f"Unknown offloading stats key: {key}")
375 + assert key in self._offloading_metric_defs
376 + if type_str == _MetricType.COUNTER:
377 + assert isinstance(value, int | float)
378 + self._increase_counter(key, value, engine_idx)
379 + elif type_str == _MetricType.GAUGE:
380 + assert isinstance(value, int | float)
381 + self._set_gauge(key, value, engine_idx)
382 + elif type_str == _MetricType.HISTOGRAM:
383 + assert isinstance(value, list)
384 + assert all(isinstance(v, int | float) for v in value)
385 + self._observe_histogram(key, value, engine_idx)
386 + else:
387 + raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+55 -1
@@ -14,6 +14,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
14 ReqId,
15 TransferJob,
16 )
17 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
18 + OffloadingConnectorStats,
19 + _TransferMetricName,
20 +)
21 from vllm.logger import init_logger
22 from vllm.utils.math_utils import cdiv
23 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
@@ -259,9 +263,13 @@ def _create_req_context(req: Request) -> ReqContext:
263 class OffloadingConnectorScheduler:
264 """Implementation of Scheduler side methods"""
265
262 - def __init__(self, spec: OffloadingSpec):
266 + def __init__(
267 + self,
268 + spec: OffloadingSpec,
269 + ):
270 self.config = SchedulerOffloadConfig.from_spec(spec)
271 self.manager: OffloadingManager = spec.get_manager()
272 + self._connector_stats: OffloadingConnectorStats | None = None
273
274 full_attention_groups: list[int] = []
275 sliding_window_groups: list[int] = []
@@ -952,6 +960,39 @@ class OffloadingConnectorScheduler:
960 if not isinstance(meta, OffloadingWorkerMetadata):
961 assert meta is None
962 meta = OffloadingWorkerMetadata()
963 + if not meta.transfer_stats.is_empty():
964 + transfer_stats = OffloadingConnectorStats()
965 + if not meta.transfer_stats.load.is_empty():
966 + transfer_stats.increase_counter(
967 + _TransferMetricName.LOAD_BYTES,
968 + meta.transfer_stats.load.bytes,
969 + )
970 + transfer_stats.increase_counter(
971 + _TransferMetricName.LOAD_TIME,
972 + meta.transfer_stats.load.time,
973 + )
974 + for size in meta.transfer_stats.load.sizes:
975 + transfer_stats.observe_histogram(
976 + _TransferMetricName.LOAD_SIZE, size
977 + )
978 + if not meta.transfer_stats.store.is_empty():
979 + transfer_stats.increase_counter(
980 + _TransferMetricName.STORE_BYTES,
981 + meta.transfer_stats.store.bytes,
982 + )
983 + transfer_stats.increase_counter(
984 + _TransferMetricName.STORE_TIME,
985 + meta.transfer_stats.store.time,
986 + )
987 + for size in meta.transfer_stats.store.sizes:
988 + transfer_stats.observe_histogram(
989 + _TransferMetricName.STORE_SIZE, size
990 + )
991 + if self._connector_stats is None:
992 + self._connector_stats = transfer_stats
993 + else:
994 + self._connector_stats.aggregate(transfer_stats)
995 +
996 for job_id, count in meta.completed_jobs.items():
997 assert count > 0
998 if job_id < self._stale_job_threshold:
@@ -990,6 +1031,19 @@ class OffloadingConnectorScheduler:
1031 if not req_status.transfer_jobs and req_status.req.is_finished():
1032 del self._req_status[job_status.req_id]
1033
1034 + def get_stats(self) -> OffloadingConnectorStats | None:
1035 + stats = self._connector_stats
1036 + self._connector_stats = None
1037 +
1038 + manager_stats = self.manager.get_stats()
1039 + if manager_stats is not None:
1040 + if stats is None:
1041 + stats = manager_stats
1042 + else:
1043 + stats.aggregate(manager_stats)
1044 +
1045 + return stats
1046 +
1047 def request_finished(
1048 self,
1049 request: Request,
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+9 -25
@@ -5,17 +5,11 @@ from dataclasses import replace
5
6 import torch
7
8 -from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
9 - KVConnectorStats,
10 -)
8 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
9 OffloadingConnectorMetadata,
10 OffloadingWorkerMetadata,
11 ReqId,
12 )
16 -from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
17 - OffloadingConnectorStats,
18 -)
13 from vllm.logger import init_logger
14 from vllm.v1.attention.backend import AttentionBackend
15 from vllm.v1.kv_cache_interface import (
@@ -44,7 +38,6 @@ class OffloadingConnectorWorker:
38 self.spec = spec
39 self.worker = OffloadingWorker()
40
47 - self.kv_connector_stats = OffloadingConnectorStats()
41 # job_id -> req_id for in-flight loads.
42 self._load_jobs: dict[int, ReqId] = {}
43 self._unsubmitted_store_jobs: list[tuple[int, TransferSpec]] = []
@@ -271,15 +264,18 @@ class OffloadingConnectorWorker:
264 # we currently do not support job failures
265 job_id = transfer_result.job_id
266 assert transfer_result.success
267 + is_load = job_id in self._load_jobs
268 if (
275 - transfer_result.transfer_time
269 + transfer_result.transfer_time is not None
270 and transfer_result.transfer_size is not None
277 - and transfer_result.transfer_type is not None
271 ):
279 - self.kv_connector_stats.record_transfer(
280 - num_bytes=transfer_result.transfer_size,
281 - time=transfer_result.transfer_time,
282 - transfer_type=transfer_result.transfer_type,
272 + if is_load:
273 + stats = self._connector_worker_meta.transfer_stats.load
274 + else:
275 + stats = self._connector_worker_meta.transfer_stats.store
276 + stats.record(
277 + transfer_result.transfer_size,
278 + transfer_result.transfer_time,
279 )
280
281 self._connector_worker_meta.mark_completed(job_id)
@@ -297,18 +293,6 @@ class OffloadingConnectorWorker:
293 self._connector_worker_meta = OffloadingWorkerMetadata()
294 return meta
295
300 - def get_kv_connector_stats(self) -> KVConnectorStats | None:
301 - """
302 - Get the KV transfer stats for the connector.
303 - """
304 -
305 - if self.kv_connector_stats.is_empty():
306 - return None
307 - # Clear stats for next iteration
308 - kv_connector_stats = self.kv_connector_stats
309 - self.kv_connector_stats = OffloadingConnectorStats()
310 - return kv_connector_stats
311 -
296 def shutdown(self) -> None:
297 self._unsubmitted_store_jobs.clear()
298 self._load_jobs.clear()
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+8 -3
@@ -184,9 +184,14 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
184 return True
185
186 def get_kv_connector_stats(self) -> KVConnectorStats | None:
187 - if self.connector_worker is None:
188 - return None # We only emit stats from the worker-side
189 - return self.connector_worker.get_kv_connector_stats()
187 + if self.connector_scheduler is not None:
188 + return self.connector_scheduler.get_stats()
189 +
190 + # TODO(orozery): Remove once PR #43877 lands
191 + if self.connector_worker is not None:
192 + return OffloadingConnectorStats()
193 +
194 + return None
195
196 @classmethod
197 def build_kv_connector_stats(
vllm/v1/kv_offload/base.py
+34
@@ -19,6 +19,9 @@ from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
19
20 if TYPE_CHECKING:
21 from vllm.config import VllmConfig
22 + from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
23 + OffloadingConnectorStats,
24 + )
25 from vllm.v1.kv_cache_interface import KVCacheConfig
26 from vllm.v1.kv_offload.worker.worker import OffloadingHandler
27
@@ -123,6 +126,26 @@ The class provides the following primitives:
126 """
127
128
129 +@dataclass(frozen=True)
130 +class OffloadingMetricMetadata:
131 + documentation: str
132 +
133 +
134 +@dataclass(frozen=True)
135 +class OffloadingCounterMetadata(OffloadingMetricMetadata):
136 + pass
137 +
138 +
139 +@dataclass(frozen=True)
140 +class OffloadingGaugeMetadata(OffloadingMetricMetadata):
141 + pass
142 +
143 +
144 +@dataclass(frozen=True)
145 +class OffloadingHistogramMetadata(OffloadingMetricMetadata):
146 + buckets: tuple[float, ...] | None = None
147 +
148 +
149 class OffloadingManager(ABC):
150 @abstractmethod
151 def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
@@ -269,6 +292,10 @@ class OffloadingManager(ABC):
292 """Evict all tracked blocks and reset internal state."""
293 return
294
295 + def get_stats(self) -> "OffloadingConnectorStats | None":
296 + """Return collected metrics since last call, or None if disabled."""
297 + return None
298 +
299 def shutdown(self) -> None:
300 """Shutdown the manager and release any resources."""
301 return
@@ -378,6 +405,13 @@ class CanonicalKVCaches:
405 class OffloadingSpec(ABC):
406 """Spec for an offloading connector"""
407
408 + @classmethod
409 + def build_metric_definitions(
410 + cls, extra_config: dict[str, Any]
411 + ) -> dict[str, "OffloadingMetricMetadata"]:
412 + """Return Prometheus metric definitions emitted by this spec."""
413 + return {}
414 +
415 def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"):
416 logger.warning(
417 "Initializing OffloadingSpec. This API is experimental and "
vllm/v1/kv_offload/cpu/common.py
+2
@@ -4,6 +4,8 @@ from typing_extensions import override
4
5 from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec
6
7 +METRIC_STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
8 +
9
10 class CPULoadStoreSpec(BlockIDsLoadStoreSpec):
11 """
vllm/v1/kv_offload/cpu/manager.py
+19 -1
@@ -6,6 +6,9 @@ from typing import Literal
6
7 from typing_extensions import override
8
9 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
10 + OffloadingConnectorStats,
11 +)
12 from vllm.v1.kv_offload.base import (
13 LoadStoreSpec,
14 OffloadingEvent,
@@ -15,7 +18,7 @@ from vllm.v1.kv_offload.base import (
18 ReqContext,
19 RequestOffloadingContext,
20 )
18 -from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
21 +from vllm.v1.kv_offload.cpu.common import METRIC_STORES_SKIPPED, CPULoadStoreSpec
22 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
23 from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
24 from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
@@ -58,6 +61,7 @@ class CPUOffloadingManager(OffloadingManager):
61 self._policy: CachePolicy = policy_cls(cache_capacity=num_blocks)
62 self.store_threshold: int = store_threshold
63 self.max_tracker_size: int = max_tracker_size
64 + self.stores_skipped_in_current_batch: int = 0
65
66 # Number of block references. It is ordered so can evict the LRU entry in O(1).
67 self.counts: OrderedDict[OffloadKey, int] | None = (
@@ -154,7 +158,9 @@ class CPUOffloadingManager(OffloadingManager):
158 req_context: ReqContext,
159 ) -> PrepareStoreOutput | None:
160 if self.counts is not None:
161 + num_keys = len(keys)
162 keys = [k for k in keys if self.counts.get(k, 0) >= self.store_threshold]
163 + self.stores_skipped_in_current_batch += num_keys - len(keys)
164 # filter out blocks that are already stored
165 keys_to_store = [k for k in keys if self._policy.get(k) is None]
166
@@ -253,3 +259,15 @@ class CPUOffloadingManager(OffloadingManager):
259 if self.events is not None:
260 yield from self.events
261 self.events.clear()
262 +
263 + def get_stats(self) -> OffloadingConnectorStats | None:
264 + if self.store_threshold < 2:
265 + return None
266 +
267 + stats = OffloadingConnectorStats()
268 + stats.increase_counter(
269 + METRIC_STORES_SKIPPED,
270 + self.stores_skipped_in_current_batch,
271 + )
272 + self.stores_skipped_in_current_batch = 0
273 + return stats
vllm/v1/kv_offload/cpu/spec.py
+20 -1
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Iterator
4 +from typing import Any
5
6 from typing_extensions import override
7
@@ -12,10 +13,12 @@ from vllm.v1.kv_offload.base import (
13 CanonicalKVCaches,
14 GPULoadStoreSpec,
15 LoadStoreSpec,
16 + OffloadingCounterMetadata,
17 OffloadingManager,
18 + OffloadingMetricMetadata,
19 OffloadingSpec,
20 )
18 -from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
21 +from vllm.v1.kv_offload.cpu.common import METRIC_STORES_SKIPPED, CPULoadStoreSpec
22 from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers
23 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
24 from vllm.v1.kv_offload.worker.worker import OffloadingHandler
@@ -24,6 +27,22 @@ from vllm.v1.kv_offload.worker.worker import OffloadingHandler
27 class CPUOffloadingSpec(OffloadingSpec):
28 BLOCK_SIZE_ALIGNMENT = 1
29
30 + @classmethod
31 + def build_metric_definitions(
32 + cls, extra_config: dict[str, Any]
33 + ) -> dict[str, OffloadingMetricMetadata]:
34 + store_threshold = int(extra_config.get("store_threshold", 0))
35 + if store_threshold < 2:
36 + return {}
37 + return {
38 + METRIC_STORES_SKIPPED: OffloadingCounterMetadata(
39 + documentation=(
40 + "Number of KV offload stores skipped because the reuse "
41 + "threshold was not reached."
42 + ),
43 + )
44 + }
45 +
46 def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
47 super().__init__(vllm_config, kv_cache_config)
48
vllm/v1/kv_offload/factory.py
+15 -5
@@ -30,11 +30,7 @@ class OffloadingSpecFactory:
30 cls._registry[name] = loader
31
32 @classmethod
33 - def create_spec(
34 - cls,
35 - config: "VllmConfig",
36 - kv_cache_config: "KVCacheConfig",
37 - ) -> OffloadingSpec:
33 + def get_spec_cls(cls, config: "VllmConfig") -> type[OffloadingSpec]:
34 kv_transfer_config = config.kv_transfer_config
35 assert kv_transfer_config is not None
36 extra_config = kv_transfer_config.kv_connector_extra_config
@@ -48,6 +44,20 @@ class OffloadingSpecFactory:
44 spec_module = importlib.import_module(spec_module_path)
45 spec_cls = getattr(spec_module, spec_name)
46 assert issubclass(spec_cls, OffloadingSpec)
47 + return spec_cls
48 +
49 + @classmethod
50 + def create_spec(
51 + cls,
52 + config: "VllmConfig",
53 + kv_cache_config: "KVCacheConfig",
54 + ) -> OffloadingSpec:
55 + kv_transfer_config = config.kv_transfer_config
56 + assert kv_transfer_config is not None
57 + spec_name = kv_transfer_config.kv_connector_extra_config.get(
58 + "spec_name", "CPUOffloadingSpec"
59 + )
60 + spec_cls = cls.get_spec_cls(config)
61 logger.info("Creating offloading spec with name: %s", spec_name)
62 return spec_cls(config, kv_cache_config)
63