@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e040899a00

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

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

Srinivas Krovvidi committed Jul 7, 2026 at 01:26 UTC e040899a006ffd07c51151b32d29649f30cb21e9

10 files changed +320 -16

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

+38

@@ -10,9 +10,11 @@ from prometheus_client import Counter, Gauge, Histogram

10 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (

11 OffloadingConnectorStats,

12 OffloadPromMetrics,

13 + _ConnectorMetricName,

14 _MetricType,

15 _StatsKey,

16 _TransferMetricName,

17 + get_connector_metric_definitions,

18 )

19 from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (

20 OffloadingConnector,

@@ -37,6 +39,42 @@ MY_COUNTER = "my_counter"

39 MY_LABEL = "my_label"

42 +def test_connector_metric_histogram_buckets():

43 + metadata = get_connector_metric_definitions()

44 +

45 + sync_delay = metadata[_ConnectorMetricName.LOOKUP_SYNC_DELAY]

46 + assert isinstance(sync_delay, OffloadingHistogramMetadata)

47 + assert sync_delay.buckets == (

48 + 0.00001,

49 + 0.00005,

50 + 0.0001,

51 + 0.0005,

52 + 0.001,

53 + 0.005,

54 + 0.01,

55 + 0.05,

56 + 0.1,

57 + 0.5,

58 + 1,

59 + )

60 +

61 + async_delay = metadata[_ConnectorMetricName.LOOKUP_ASYNC_DELAY]

62 + assert isinstance(async_delay, OffloadingHistogramMetadata)

63 + assert async_delay.buckets == (

64 + 0.0001,

65 + 0.0005,

66 + 0.001,

67 + 0.005,

68 + 0.01,

69 + 0.05,

70 + 0.1,

71 + 0.5,

72 + 1,

73 + 5,

74 + 10,

75 + )

76 +

77 +

78 class _FakeMetric:

79 def __init__(self, **kwargs):

80 self.kwargs = kwargs

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

+70

@@ -11,6 +11,10 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (

11 to_keys,

12 )

13 from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID

14 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (

15 + OffloadingConnectorStats,

16 + _ConnectorMetricName,

17 +)

18 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (

19 OffloadingConnectorScheduler,

20 RequestOffloadState,

@@ -32,6 +36,72 @@ from vllm.v1.kv_offload.base import (

36 from vllm.v1.request import RequestStatus

39 +def _reduce_kv_connector_stats(runner):

40 + reduced: dict[str, int | float] = {}

41 + for payload in runner.kv_connector_stats:

42 + stats = (

43 + payload

44 + if hasattr(payload, "reduce")

45 + else OffloadingConnectorStats(data=payload)

46 + )

47 + for key, value in stats.reduce().items():

48 + reduced[key] = reduced.get(key, 0) + value

49 + return reduced

50 +

51 +

52 +def test_scheduler_reports_allocation_failure(request_runner):

53 + runner = request_runner(

54 + block_size=4,

55 + num_gpu_blocks=10,

56 + async_scheduling=False,

57 + )

58 + runner.new_request(token_ids=[0] * 4)

59 + runner.manager.prepare_store.side_effect = lambda keys, req_context: None

60 +

61 + runner.run(decoded_tokens=[EOS_TOKEN_ID])

62 +

63 + reduced = _reduce_kv_connector_stats(runner)

64 + assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 1

65 +

66 +

67 +def test_scheduler_reports_lookup_sync_delay(request_runner):

68 + runner = request_runner(

69 + block_size=4,

70 + num_gpu_blocks=10,

71 + async_scheduling=False,

72 + )

73 + runner.new_request(token_ids=[1] * 4)

74 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

75 + generate_store_output([])

76 + )

77 +

78 + runner.run(decoded_tokens=[EOS_TOKEN_ID])

79 +

80 + reduced = _reduce_kv_connector_stats(runner)

81 + assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_count"] == 1

82 + assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_sum"] > 0

83 +

84 +

85 +def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):

86 + """A deferred lookup reports its async delay once it resolves."""

87 + runner = request_runner(

88 + block_size=4,

89 + num_gpu_blocks=10,

90 + async_scheduling=False,

91 + )

92 + runner.manager.lookup.side_effect = [LookupResult.RETRY, LookupResult.MISS]

93 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

94 + generate_store_output([])

95 + )

96 +

97 + runner.new_request(token_ids=[1] * 4)

98 + runner.run(decoded_tokens=[EOS_TOKEN_ID])

99 +

100 + reduced = _reduce_kv_connector_stats(runner)

101 + assert reduced[f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_count"] == 1

102 + assert reduced[f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_sum"] > 0

103 +

104 +

105 @pytest.mark.parametrize("async_scheduling", [True, False])

106 def test_offloading_connector(request_runner, async_scheduling: bool):

107 block_size = 4

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

+16 -3

@@ -133,6 +133,7 @@ class MockOffloadingSpec(OffloadingSpec):

133 self.manager = MagicMock(spec=OffloadingManager)

134 self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)

135 self.manager.lookup.return_value = LookupResult.MISS

136 + self.manager.get_stats.return_value = None

137 self.manager.on_new_request.return_value = RequestOffloadingContext()

138 self.handler = MockOffloadingWorker()

139

@@ -335,6 +336,7 @@ class RequestRunner:

336 self.completed_loads: list[TransferSummary] = []

337 self.completed_stores: list[TransferSummary] = []

338 self.flushed_gpu_blocks: set[GPUBlock] = set()

339 + self.kv_connector_stats: list[Any] = []

340

341 # block_id -> GPUBlock

342 self.gpu_blocks: dict[int, GPUBlock] = {}

@@ -348,6 +350,12 @@ class RequestRunner:

350 slot_mapping={},

351 )

352

353 + def _record_kv_connector_stats(self, engine_outputs: dict[int, Any]) -> None:

354 + for output in engine_outputs.values():

355 + scheduler_stats = output.scheduler_stats

356 + if scheduler_stats is not None and scheduler_stats.kv_connector_stats:

357 + self.kv_connector_stats.append(scheduler_stats.kv_connector_stats)

358 +

359 def new_request(

360 self,

361 token_ids: list[int],

@@ -524,13 +532,17 @@ class RequestRunner:

532 if self.async_scheduling:

533 # in async scheduling we update the output of the previous step

534 if prev_model_runner_output is not None:

527 - self.scheduler.update_from_output(

535 + engine_outputs = self.scheduler.update_from_output(

536 prev_scheduler_output, prev_model_runner_output

537 )

538 + self._record_kv_connector_stats(engine_outputs)

539 prev_scheduler_output = scheduler_output

540 prev_model_runner_output = model_runner_output

541 else:

533 - self.scheduler.update_from_output(scheduler_output, model_runner_output)

542 + engine_outputs = self.scheduler.update_from_output(

543 + scheduler_output, model_runner_output

544 + )

545 + self._record_kv_connector_stats(engine_outputs)

546

547 if post_step_fn is not None:

548 post_step_fn()

@@ -546,9 +558,10 @@ class RequestRunner:

558 if token_id is None:

559 if self.async_scheduling:

560 # sample last token

549 - self.scheduler.update_from_output(

561 + engine_outputs = self.scheduler.update_from_output(

562 prev_scheduler_output, prev_model_runner_output

563 )

564 + self._record_kv_connector_stats(engine_outputs)

565 break

566

567 self._parse_transfers()

tests/v1/kv_offload/cpu/test_manager.py

+58

@@ -224,6 +224,64 @@ def test_cpu_manager_reports_cache_usage_gauge():

224 check_usage_stats(manager, 0.0)

225

226

227 +def test_cpu_manager_reports_allocation_size_histogram():

228 + manager = make_cpu_manager(num_blocks=4, cache_policy="lru")

229 +

230 + manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)

231 + manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)

232 + manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX)

233 +

234 + stats = manager.get_stats()

235 +

236 + assert stats is not None

237 + reduced = stats.reduce()

238 + assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count"] == 2

239 + assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 3

240 +

241 + # The cache-usage gauge is always reported, so get_stats() never returns

242 + # None, but the histogram has nothing new once its samples are consumed.

243 + second_stats = manager.get_stats()

244 + assert second_stats is not None

245 + assert f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count" not in (

246 + second_stats.reduce()

247 + )

248 +

249 +

250 +def test_cpu_manager_reports_allocation_size_on_allocation_failure(monkeypatch):

251 + manager = make_cpu_manager(num_blocks=4, cache_policy="lru")

252 +

253 + def fail_allocate_blocks(keys):

254 + raise RuntimeError("allocation failed")

255 +

256 + monkeypatch.setattr(manager, "_allocate_blocks", fail_allocate_blocks)

257 +

258 + with pytest.raises(RuntimeError, match="allocation failed"):

259 + manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX)

260 +

261 + stats = manager.get_stats()

262 +

263 + assert stats is not None

264 + reduced = stats.reduce()

265 + assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count"] == 1

266 + assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 3

267 +

268 +

269 +def test_cpu_manager_reports_allocation_size_on_eviction_failure():

270 + manager = make_cpu_manager(num_blocks=1, cache_policy="lru")

271 +

272 + manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX)

273 + manager.get_stats()

274 +

275 + assert manager.prepare_store(to_keys([2]), _EMPTY_REQ_CTX) is None

276 +

277 + stats = manager.get_stats()

278 +

279 + assert stats is not None

280 + reduced = stats.reduce()

281 + assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count"] == 1

282 + assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 1

283 +

284 +

285 def test_cpu_manager():

286 """

287 Tests CPUOffloadingManager with lru policy.

tests/v1/kv_offload/test_factory.py

+29 -3

@@ -21,7 +21,7 @@ from vllm.v1.kv_cache_interface import (

21 KVCacheGroupSpec,

22 KVCacheTensor,

23 )

24 -from vllm.v1.kv_offload.base import OffloadingSpec

24 +from vllm.v1.kv_offload.base import OffloadingHistogramMetadata, OffloadingSpec

25 from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec

26 from vllm.v1.kv_offload.factory import OffloadingSpecFactory

27 from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec

@@ -248,8 +248,8 @@ def test_duplicate_registration_raises():

248 # ---------------------------------------------------------------------------

249

250

251 -def test_build_metric_definitions_empty_below_threshold():

252 - """store_threshold < 2 → only base metric (no stores_skipped)."""

251 +def test_build_metric_definitions_below_threshold():

252 + """store_threshold < 2 keeps stores_skipped disabled."""

253 from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics

254

255 config = _make_vllm_config(store_threshold=1)

@@ -258,6 +258,32 @@ def test_build_metric_definitions_empty_below_threshold():

258 config.kv_transfer_config.kv_connector_extra_config

259 )

260 assert CPUOffloadingMetrics.STORES_SKIPPED not in metrics

261 + assert CPUOffloadingMetrics.CPU_ALLOCATION_SIZE in metrics

262 +

263 +

264 +def test_build_metric_definitions_allocation_size_histogram():

265 + """CPU allocation size is always reported as a histogram."""

266 + from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics

267 +

268 + config = _make_vllm_config(store_threshold=0)

269 + spec_cls = OffloadingSpecFactory.get_spec_cls(config)

270 + metrics = spec_cls.build_metric_definitions(

271 + config.kv_transfer_config.kv_connector_extra_config

272 + )

273 + metadata = metrics[CPUOffloadingMetrics.CPU_ALLOCATION_SIZE]

274 + assert isinstance(metadata, OffloadingHistogramMetadata)

275 + assert metadata.buckets == (

276 + 1,

277 + 4,

278 + 16,

279 + 64,

280 + 256,

281 + 1024,

282 + 4096,

283 + 16384,

284 + 65536,

285 + 262144,

286 + )

287

288

289 def test_build_metric_definitions_returns_counter_at_threshold():

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

+53 -1

@@ -30,6 +30,14 @@ class _TransferMetricName:

30 STORE_SIZE = "vllm:kv_offload_store_size"

33 +class _ConnectorMetricName:

34 + """Connector-side metrics emitted by scheduler-side offloading code."""

35 +

36 + LOOKUP_SYNC_DELAY = "vllm:kv_offload_lookup_sync_delay_seconds"

37 + LOOKUP_ASYNC_DELAY = "vllm:kv_offload_lookup_async_delay_seconds"

38 + ALLOCATION_FAILURE = "vllm:kv_offload_allocation_failure"

39 +

40 +

41 class _TransferType:

42 """Transfer direction labels for deprecated CPU offload metrics."""

@@ -74,6 +82,50 @@ def get_connector_metric_definitions() -> dict[str, OffloadingMetricMetadata]:

82 documentation="Histogram of KV offload store operation size, in bytes.",

83 buckets=TRANSFER_SIZE_BUCKETS,

84 ),

85 + _ConnectorMetricName.LOOKUP_SYNC_DELAY: OffloadingHistogramMetadata(

86 + documentation=(

87 + "Histogram of the time spent in a single offload lookup call, "

88 + "in seconds."

89 + ),

90 + buckets=(

91 + 0.00001,

92 + 0.00005,

93 + 0.0001,

94 + 0.0005,

95 + 0.001,

96 + 0.005,

97 + 0.01,

98 + 0.05,

99 + 0.1,

100 + 0.5,

101 + 1,

102 + ),

103 + ),

104 + _ConnectorMetricName.LOOKUP_ASYNC_DELAY: OffloadingHistogramMetadata(

105 + documentation=(

106 + "Histogram of time between a request's offload lookup first "

107 + "deferring and the following lookup resolving, or request "

108 + "finish, in seconds."

109 + ),

110 + buckets=(

111 + 0.0001,

112 + 0.0005,

113 + 0.001,

114 + 0.005,

115 + 0.01,

116 + 0.05,

117 + 0.1,

118 + 0.5,

119 + 1,

120 + 5,

121 + 10,

122 + ),

123 + ),

124 + _ConnectorMetricName.ALLOCATION_FAILURE: OffloadingCounterMetadata(

125 + documentation=(

126 + "Number of KV offload store allocation attempts that failed."

127 + ),

128 + ),

129 }

130

131

@@ -223,7 +275,7 @@ class OffloadingConnectorStats(KVConnectorStats):

275 def increase_counter(

276 self,

277 counter_name: str,

226 - counter_increase_value: int | float,

278 + counter_increase_value: int | float = 1,

279 labelvalues: tuple[str, ...] = (),

280 ) -> None:

281 """Increase a counter on the stats payload."""

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

+37 -8

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import time

4 from collections.abc import Iterable, Sequence

5 from dataclasses import dataclass, field

6 from itertools import islice

@@ -21,6 +22,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import (

22 )

23 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (

24 OffloadingConnectorStats,

25 + _ConnectorMetricName,

26 _TransferMetricName,

27 )

28 from vllm.logger import init_logger

@@ -246,6 +248,9 @@ class RequestOffloadState:

248 # In-flight job IDs. Per the connector's invariant, at any given time

249 # this contains either a single load job, or one or more store jobs.

250 transfer_jobs: set[int] = field(default_factory=set)

251 + # time.monotonic() of this request's first deferred offload lookup;

252 + # None once consumed (observed) or while no lookup is pending.

253 + deferred_lookup_start_time: float | None = None

254

255 def __post_init__(self) -> None:

256 self.group_states = tuple(

@@ -325,7 +330,7 @@ class OffloadingConnectorScheduler:

330 ):

331 self.config = SchedulerOffloadConfig.from_spec(spec)

332 self.manager: OffloadingManager = spec.get_manager()

328 - self._connector_stats: OffloadingConnectorStats | None = None

333 + self._connector_stats = OffloadingConnectorStats()

334

335 full_attention_groups: list[int] = []

336 sliding_window_groups: list[int] = []

@@ -375,6 +380,18 @@ class OffloadingConnectorScheduler:

380

381 self._events_tracker = OffloadingEventsTracker(spec.kv_events_config)

382

383 + def _maybe_observe_lookup_async_delay(

384 + self, req_status: RequestOffloadState

385 + ) -> None:

386 + start_time = req_status.deferred_lookup_start_time

387 + if start_time is None:

388 + return

389 + req_status.deferred_lookup_start_time = None

390 + self._connector_stats.observe_histogram(

391 + _ConnectorMetricName.LOOKUP_ASYNC_DELAY,

392 + time.monotonic() - start_time,

393 + )

394 +

395 def _generate_job_id(self) -> int:

396 job_id = self._job_counter

397 self._job_counter += 1

@@ -685,7 +702,17 @@ class OffloadingConnectorScheduler:

702 if request.skip_reading_prefix_cache:

703 num_hit_tokens = 0

704 else:

705 + lookup_start = time.monotonic()

706 num_hit_tokens = self._lookup(req_status)

707 + self._connector_stats.observe_histogram(

708 + _ConnectorMetricName.LOOKUP_SYNC_DELAY,

709 + time.monotonic() - lookup_start,

710 + )

711 + if num_hit_tokens is None:

712 + if req_status.deferred_lookup_start_time is None:

713 + req_status.deferred_lookup_start_time = lookup_start

714 + else:

715 + self._maybe_observe_lookup_async_delay(req_status)

716 req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0))

717

718 self._touch(req_status)

@@ -925,6 +952,9 @@ class OffloadingConnectorScheduler:

952 new_offload_keys, req_status.req_context

953 )

954 if store_output is None:

955 + self._connector_stats.increase_counter(

956 + _ConnectorMetricName.ALLOCATION_FAILURE

957 + )

958 logger.warning("Request %s: cannot store blocks", req_id)

959 continue

960

@@ -1115,10 +1145,7 @@ class OffloadingConnectorScheduler:

1145 transfer_stats.observe_histogram(

1146 _TransferMetricName.STORE_SIZE, size

1147 )

1118 - if self._connector_stats is None:

1119 - self._connector_stats = transfer_stats

1120 - else:

1121 - self._connector_stats.aggregate(transfer_stats)

1148 + self._connector_stats.aggregate(transfer_stats)

1149

1150 for job_id, count in meta.completed_jobs.items():

1151 assert count > 0

@@ -1159,8 +1186,10 @@ class OffloadingConnectorScheduler:

1186 del self._req_status[job_status.req_id]

1187

1188 def get_stats(self) -> OffloadingConnectorStats | None:

1162 - stats = self._connector_stats

1163 - self._connector_stats = None

1189 + stats: OffloadingConnectorStats | None = None

1190 + if not self._connector_stats.is_empty():

1191 + stats = self._connector_stats

1192 + self._connector_stats = OffloadingConnectorStats()

1193

1194 manager_stats = self.manager.get_stats()

1195 if manager_stats is not None:

@@ -1198,7 +1227,7 @@ class OffloadingConnectorScheduler:

1227 return False, None

1228

1229 self.manager.on_request_finished(req_status.req_context)

1201 -

1230 + self._maybe_observe_lookup_async_delay(req_status)

1231 if not req_status.transfer_jobs:

1232 # No in-flight jobs: no later complete_store()/complete_load() calls

1233 # need this request's state.

vllm/v1/kv_offload/cpu/common.py

@@ -8,6 +8,7 @@ from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec

8 class CPUOffloadingMetrics:

9 STORES_SKIPPED = "vllm:kv_offload_stores_skipped"

10 CPU_CACHE_USAGE_PERC = "vllm:kv_offload_cpu_cache_usage_perc"

11 + CPU_ALLOCATION_SIZE = "vllm:kv_offload_cpu_allocation_size"

14 class CPULoadStoreSpec(BlockIDsLoadStoreSpec):

vllm/v1/kv_offload/cpu/manager.py

@@ -69,6 +69,7 @@ class CPUOffloadingManager(OffloadingManager):

69 self.store_threshold: int = store_threshold

70 self.max_tracker_size: int = max_tracker_size

71 self.stores_skipped_in_current_batch: int = 0

72 + self.allocation_sizes_in_current_batch: list[int] = []

74 # Number of block references. It is ordered so can evict the LRU entry in O(1).

75 self.counts: OrderedDict[OffloadKey, int] | None = (

@@ -185,6 +186,7 @@ class CPUOffloadingManager(OffloadingManager):

186 evicted_keys=[],

187 )

188

189 + self.allocation_sizes_in_current_batch.append(len(keys_to_store))

190 num_blocks_to_evict = len(keys_to_store) - self._get_num_free_blocks()

191

192 to_evict: list[OffloadKey] = []

@@ -300,10 +302,17 @@ class CPUOffloadingManager(OffloadingManager):

302 usage = num_used / self._num_blocks if self._num_blocks > 0 else 0.0

303 stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC, usage)

304

305 + for allocation_size in self.allocation_sizes_in_current_batch:

306 + stats.observe_histogram(

307 + CPUOffloadingMetrics.CPU_ALLOCATION_SIZE, allocation_size

308 + )

309 + self.allocation_sizes_in_current_batch.clear()

310 +

311 if self.store_threshold >= 2:

312 stats.increase_counter(

313 CPUOffloadingMetrics.STORES_SKIPPED,

314 self.stores_skipped_in_current_batch,

315 )

316 self.stores_skipped_in_current_batch = 0

317 +

318 return stats

vllm/v1/kv_offload/cpu/spec.py

+9 -1

@@ -12,6 +12,7 @@ from vllm.v1.kv_offload.base import (

12 CanonicalKVCaches,

13 OffloadingCounterMetadata,

14 OffloadingGaugeMetadata,

15 + OffloadingHistogramMetadata,

16 OffloadingManager,

17 OffloadingMetricMetadata,

18 OffloadingSpec,

@@ -37,7 +38,14 @@ class CPUOffloadingSpec(OffloadingSpec):

38 "values indicate transfers (stores or promotions) may be "

39 "dropped due to insufficient capacity."

40 ),

40 - )

41 + ),

42 + CPUOffloadingMetrics.CPU_ALLOCATION_SIZE: OffloadingHistogramMetadata(

43 + documentation=(

44 + "Histogram of the number of CPU blocks requested by each "

45 + "KV offload prepare_store call."

46 + ),

47 + buckets=(1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144),

48 + ),

49 }

50 store_threshold = int(extra_config.get("store_threshold", 0))

51 if store_threshold >= 2: