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

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

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 01:45 UTC

[KVConnector] Add metrics to Prometheus-Grafana dashboard (#26811)

Signed-off-by: NickLucche <nlucches@redhat.com> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com>

Nicolò Lucchesi committed Oct 29, 2025 at 19:44 UTC accb8fab07de05b2df43ccd97619100526cb8695
6 files changed +365 -29
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+21 -1
@@ -50,7 +50,12 @@ if TYPE_CHECKING:
50 from vllm.attention.backends.abstract import AttentionMetadata
51 from vllm.config import VllmConfig
52 from vllm.distributed.kv_events import KVCacheEvent
53 - from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
53 + from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
54 + KVConnectorPromMetrics,
55 + KVConnectorStats,
56 + PromMetric,
57 + PromMetricT,
58 + )
59 from vllm.forward_context import ForwardContext
60 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
61 from vllm.v1.request import Request
@@ -471,3 +476,18 @@ class KVConnectorBase_V1(ABC):
476 which can implement custom aggregation logic on the data dict.
477 """
478 return None
479 +
480 + @classmethod
481 + def build_prom_metrics(
482 + cls,
483 + vllm_config: "VllmConfig",
484 + metric_types: dict[type["PromMetric"], type["PromMetricT"]],
485 + labelnames: list[str],
486 + per_engine_labelvalues: dict[int, list[str]],
487 + ) -> Optional["KVConnectorPromMetrics"]:
488 + """
489 + Create a KVConnectorPromMetrics subclass which should register
490 + per-connector Prometheus metrics and implement observe() to
491 + expose connector transfer stats via Prometheus.
492 + """
493 + return None
vllm/distributed/kv_transfer/kv_connector/v1/metrics.py
+87 -2
@@ -1,13 +1,18 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass, field
4 -from typing import Any
4 +from typing import Any, TypeAlias, TypeVar
5
6 -from vllm.config.kv_transfer import KVTransferConfig
6 +from prometheus_client import Counter, Gauge, Histogram
7 +
8 +from vllm.config import KVTransferConfig, VllmConfig
9 from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
10 from vllm.distributed.kv_transfer.kv_transfer_state import has_kv_transfer_group
11 from vllm.logger import init_logger
12
13 +PromMetric: TypeAlias = Gauge | Counter | Histogram
14 +PromMetricT = TypeVar("PromMetricT", bound=PromMetric)
15 +
16 logger = init_logger(__name__)
17
18
@@ -102,3 +107,83 @@ class KVConnectorLogging:
107
108 # Reset metrics for next interval
109 self.reset()
110 +
111 +
112 +class KVConnectorPromMetrics:
113 + """
114 + A base class for per-connector Prometheus metric registration
115 + and recording.
116 + """
117 +
118 + def __init__(
119 + self,
120 + vllm_config: VllmConfig,
121 + metric_types: dict[type[PromMetric], type[PromMetricT]],
122 + labelnames: list[str],
123 + per_engine_labelvalues: dict[int, list[str]],
124 + ):
125 + self._kv_transfer_config = vllm_config.kv_transfer_config
126 + self._gauge_cls = metric_types[Gauge]
127 + self._counter_cls = metric_types[Counter]
128 + self._histogram_cls = metric_types[Histogram]
129 + self._labelnames = labelnames
130 + self._per_engine_labelvalues = per_engine_labelvalues
131 +
132 + def make_per_engine(self, metric: PromMetric) -> PromMetric:
133 + """
134 + Create a per-engine child of a prometheus_client.Metric with
135 + the appropriate labels set. The parent metric must be created
136 + using the labelnames list.
137 + """
138 + return {
139 + idx: metric.labels(*labelvalues)
140 + for idx, labelvalues in self._per_engine_labelvalues.items()
141 + }
142 +
143 + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
144 + """
145 + Record the supplied transfer statistics to Prometheus metrics. These
146 + statistics are engine-specific, and should be recorded to a metric
147 + with the appropriate 'engine' label. These metric instances can be
148 + created using the make_per_engine() helper method.
149 + """
150 + raise NotImplementedError
151 +
152 +
153 +class KVConnectorPrometheus:
154 + """
155 + Support for registering per-connector Prometheus metrics, and
156 + recording transfer statistics to those metrics. Uses
157 + KVConnectorBase.build_prom_metrics().
158 + """
159 +
160 + _gauge_cls = Gauge
161 + _counter_cls = Counter
162 + _histogram_cls = Histogram
163 +
164 + def __init__(
165 + self,
166 + vllm_config: VllmConfig,
167 + labelnames: list[str],
168 + per_engine_labelvalues: dict[int, list[str]],
169 + ):
170 + self.prom_metrics: KVConnectorPromMetrics | None = None
171 + kv_transfer_config = vllm_config.kv_transfer_config
172 + if kv_transfer_config and kv_transfer_config.kv_connector:
173 + connector_cls = KVConnectorFactory.get_connector_class(kv_transfer_config)
174 + metric_types = {
175 + Gauge: self._gauge_cls,
176 + Counter: self._counter_cls,
177 + Histogram: self._histogram_cls,
178 + }
179 + self.prom_metrics = connector_cls.build_prom_metrics(
180 + vllm_config,
181 + metric_types,
182 + labelnames,
183 + per_engine_labelvalues,
184 + )
185 +
186 + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
187 + if self.prom_metrics is None:
188 + return
189 + self.prom_metrics.observe(transfer_stats_data, engine_idx)
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+88 -22
@@ -9,13 +9,19 @@ import torch
9
10 from vllm.config import VllmConfig
11 from vllm.config.kv_transfer import KVTransferConfig
12 +from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType
13 from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
14 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
15 KVConnectorBase_V1,
16 KVConnectorMetadata,
17 KVConnectorRole,
18 )
18 -from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
19 +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
20 + KVConnectorPromMetrics,
21 + KVConnectorStats,
22 + PromMetric,
23 + PromMetricT,
24 +)
25 from vllm.logger import init_logger
26 from vllm.v1.core.sched.output import SchedulerOutput
27 from vllm.v1.outputs import KVConnectorOutput
@@ -72,6 +78,27 @@ class MultiKVConnectorStats(KVConnectorStats):
78 self.data[connector_id] = stats
79
80
81 +class MultiKVConnectorPromMetrics(KVConnectorPromMetrics):
82 + def __init__(
83 + self,
84 + vllm_config: "VllmConfig",
85 + metric_types: dict[type[PromMetric], type[PromMetricT]],
86 + labelnames: list[str],
87 + per_engine_labelvalues: dict[int, list[str]],
88 + prom_metrics: dict[str, KVConnectorPromMetrics],
89 + ):
90 + super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
91 + self._prom_metrics = prom_metrics
92 +
93 + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
94 + for connector_id, stats_data in transfer_stats_data.items():
95 + assert connector_id in self._prom_metrics, (
96 + f"{connector_id} is not contained in the list of registered connectors "
97 + f"with Prometheus metrics support: {self._prom_metrics.keys()}"
98 + )
99 + self._prom_metrics[connector_id].observe(stats_data["data"], engine_idx)
100 +
101 +
102 class MultiConnector(KVConnectorBase_V1):
103 """
104 A wrapper for using multiple KVConnectors at the same time.
@@ -84,19 +111,13 @@ class MultiConnector(KVConnectorBase_V1):
111
112 def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
113 super().__init__(vllm_config=vllm_config, role=role)
114 +
115 self._connectors: list[KVConnectorBase_V1] = []
116 self._ktc_kv_transfer_config = []
89 - ktcs = self._kv_transfer_config.kv_connector_extra_config.get("connectors")
90 - assert ktcs is not None
91 - for ktc in ktcs:
92 - temp_config = copy.copy(vllm_config)
93 - engine_id = ktc.get("engine_id", self._kv_transfer_config.engine_id)
94 - temp_config.kv_transfer_config = KVTransferConfig(
95 - **ktc, engine_id=engine_id
96 - )
97 - self._connectors.append(
98 - KVConnectorFactory.create_connector(temp_config, role)
99 - )
117 + for connector_cls, temp_config in self._get_connector_classes_and_configs(
118 + vllm_config
119 + ):
120 + self._connectors.append(connector_cls(temp_config, role))
121 self._ktc_kv_transfer_config.append(temp_config.kv_transfer_config)
122
123 # A mapping from request id to the index of the connector chosen to
@@ -109,6 +130,32 @@ class MultiConnector(KVConnectorBase_V1):
130 # Propagated from scheduler to worker side via the connector metadata.
131 self._extra_async_saves: dict[str, int] = {}
132
133 + @classmethod
134 + def _get_connector_classes_and_configs(
135 + cls, vllm_config: "VllmConfig"
136 + ) -> list[tuple[type[KVConnectorBaseType], "VllmConfig"]]:
137 + assert vllm_config.kv_transfer_config is not None
138 + ktcs = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
139 + "connectors"
140 + )
141 + assert ktcs is not None
142 + ret: list[tuple[type[KVConnectorBaseType], VllmConfig]] = []
143 + for ktc in ktcs:
144 + temp_config = copy.copy(vllm_config)
145 + engine_id = ktc.get("engine_id", vllm_config.kv_transfer_config.engine_id)
146 + temp_config.kv_transfer_config = KVTransferConfig(
147 + **ktc, engine_id=engine_id
148 + )
149 + ret.append(
150 + (
151 + KVConnectorFactory.get_connector_class(
152 + temp_config.kv_transfer_config
153 + ),
154 + temp_config,
155 + )
156 + )
157 + return ret
158 +
159 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
160 for c in self._connectors:
161 c.register_kv_caches(kv_caches)
@@ -295,18 +342,12 @@ class MultiConnector(KVConnectorBase_V1):
342 None if the connector does not require a specific layout.
343 """
344 assert vllm_config.kv_transfer_config is not None
298 - ktcs = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
299 - "connectors"
300 - )
301 - assert ktcs is not None
345 layouts: set[str] = set()
303 - temp_vllm_config = copy.copy(vllm_config)
304 - for ktc in ktcs:
305 - kv_transfer_config = KVTransferConfig(**ktc)
306 - temp_vllm_config.kv_transfer_config = kv_transfer_config
307 - connector_cls = KVConnectorFactory.get_connector_class(kv_transfer_config)
346 + for connector_cls, temp_config in cls._get_connector_classes_and_configs(
347 + vllm_config
348 + ):
349 required_kvcache_layout = connector_cls.get_required_kvcache_layout(
309 - temp_vllm_config
350 + temp_config
351 )
352 if required_kvcache_layout is not None:
353 layouts.add(required_kvcache_layout)
@@ -372,3 +413,28 @@ class MultiConnector(KVConnectorBase_V1):
413 stats_by_connector = MultiKVConnectorStats()
414 stats_by_connector[c.__class__.__name__] = stats
415 return stats_by_connector
416 +
417 + @classmethod
418 + def build_prom_metrics(
419 + cls,
420 + vllm_config: "VllmConfig",
421 + metric_types: dict[type["PromMetric"], type["PromMetricT"]],
422 + labelnames: list[str],
423 + per_engine_labelvalues: dict[int, list[str]],
424 + ) -> KVConnectorPromMetrics:
425 + prom_metrics: dict[str, KVConnectorPromMetrics] = {}
426 + for connector_cls, temp_config in cls._get_connector_classes_and_configs(
427 + vllm_config
428 + ):
429 + connector_prom = connector_cls.build_prom_metrics(
430 + temp_config, metric_types, labelnames, per_engine_labelvalues
431 + )
432 + if connector_prom is not None:
433 + prom_metrics[connector_cls.__name__] = connector_prom
434 + return MultiKVConnectorPromMetrics(
435 + vllm_config,
436 + metric_types,
437 + labelnames,
438 + per_engine_labelvalues,
439 + prom_metrics,
440 + )
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+140 -1
@@ -30,7 +30,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
30 KVConnectorMetadata,
31 KVConnectorRole,
32 )
33 -from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
33 +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
34 + KVConnectorPromMetrics,
35 + KVConnectorStats,
36 + PromMetric,
37 + PromMetricT,
38 +)
39 from vllm.distributed.parallel_state import (
40 get_tensor_model_parallel_rank,
41 get_tensor_model_parallel_world_size,
@@ -254,6 +259,18 @@ class NixlConnector(KVConnectorBase_V1):
259 else NixlKVConnectorStats()
260 )
261
262 + @classmethod
263 + def build_prom_metrics(
264 + cls,
265 + vllm_config: VllmConfig,
266 + metric_types: dict[type[PromMetric], type[PromMetricT]],
267 + labelnames: list[str],
268 + per_engine_labelvalues: dict[int, list[str]],
269 + ) -> KVConnectorPromMetrics:
270 + return NixlPromMetrics(
271 + vllm_config, metric_types, labelnames, per_engine_labelvalues
272 + )
273 +
274 def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
275 assert self.connector_worker is not None
276 assert isinstance(self._connector_metadata, NixlConnectorMetadata)
@@ -1960,3 +1977,125 @@ class NixlKVConnectorStats(KVConnectorStats):
1977 @property
1978 def num_successful_transfers(self) -> int:
1979 return len(self.data["transfer_duration"])
1980 +
1981 +
1982 +class NixlPromMetrics(KVConnectorPromMetrics):
1983 + def __init__(
1984 + self,
1985 + vllm_config: VllmConfig,
1986 + metric_types: dict[type[PromMetric], type[PromMetricT]],
1987 + labelnames: list[str],
1988 + per_engine_labelvalues: dict[int, list[str]],
1989 + ):
1990 + super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
1991 +
1992 + buckets = [
1993 + 0.001,
1994 + 0.005,
1995 + 0.01,
1996 + 0.025,
1997 + 0.05,
1998 + 0.075,
1999 + 0.1,
2000 + 0.2,
2001 + 0.3,
2002 + 0.5,
2003 + 0.75,
2004 + 1.0,
2005 + 5.0,
2006 + ]
2007 + nixl_histogram_xfer_time = self._histogram_cls(
2008 + name="vllm:nixl_xfer_time_seconds",
2009 + documentation="Histogram of transfer duration for NIXL KV Cache transfers.",
2010 + buckets=buckets[1:],
2011 + labelnames=labelnames,
2012 + )
2013 + self.nixl_histogram_xfer_time = self.make_per_engine(nixl_histogram_xfer_time)
2014 + nixl_histogram_post_time = self._histogram_cls(
2015 + name="vllm:nixl_post_time_seconds",
2016 + documentation="Histogram of transfer post time for NIXL KV"
2017 + " Cache transfers.",
2018 + buckets=buckets,
2019 + labelnames=labelnames,
2020 + )
2021 + self.nixl_histogram_post_time = self.make_per_engine(nixl_histogram_post_time)
2022 + # uniform 2kb to 16gb range
2023 + buckets = [2 ** (10 + i) for i in range(1, 25, 2)]
2024 + nixl_histogram_bytes_transferred = self._histogram_cls(
2025 + name="vllm:nixl_bytes_transferred",
2026 + documentation="Histogram of bytes transferred per NIXL KV Cache transfers.",
2027 + buckets=buckets,
2028 + labelnames=labelnames,
2029 + )
2030 + self.nixl_histogram_bytes_transferred = self.make_per_engine(
2031 + nixl_histogram_bytes_transferred
2032 + )
2033 + buckets = [
2034 + 10,
2035 + 20,
2036 + 30,
2037 + 50,
2038 + 75,
2039 + 100,
2040 + 200,
2041 + 400,
2042 + 1000,
2043 + 2000,
2044 + 4000,
2045 + 10000,
2046 + 20000,
2047 + 50000,
2048 + ]
2049 + nixl_histogram_num_descriptors = self._histogram_cls(
2050 + name="vllm:nixl_num_descriptors",
2051 + documentation="Histogram of number of descriptors per NIXL"
2052 + " KV Cache transfers.",
2053 + buckets=buckets,
2054 + labelnames=labelnames,
2055 + )
2056 + self.nixl_histogram_num_descriptors = self.make_per_engine(
2057 + nixl_histogram_num_descriptors
2058 + )
2059 + counter_nixl_num_failed_transfers = self._counter_cls(
2060 + name="vllm:nixl_num_failed_transfers",
2061 + documentation="Number of failed NIXL KV Cache transfers.",
2062 + labelnames=labelnames,
2063 + )
2064 + self.counter_nixl_num_failed_transfers = self.make_per_engine(
2065 + counter_nixl_num_failed_transfers
2066 + )
2067 + counter_nixl_num_failed_notifications = self._counter_cls(
2068 + name="vllm:nixl_num_failed_notifications",
2069 + documentation="Number of failed NIXL KV Cache notifications.",
2070 + labelnames=labelnames,
2071 + )
2072 + self.counter_nixl_num_failed_notifications = self.make_per_engine(
2073 + counter_nixl_num_failed_notifications
2074 + )
2075 +
2076 + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
2077 + for prom_obj, list_item_key in zip(
2078 + [
2079 + self.nixl_histogram_xfer_time,
2080 + self.nixl_histogram_post_time,
2081 + self.nixl_histogram_bytes_transferred,
2082 + self.nixl_histogram_num_descriptors,
2083 + ],
2084 + [
2085 + "transfer_duration",
2086 + "post_duration",
2087 + "bytes_transferred",
2088 + "num_descriptors",
2089 + ],
2090 + ):
2091 + for list_item in transfer_stats_data[list_item_key]:
2092 + prom_obj[engine_idx].observe(list_item)
2093 + for counter_obj, counter_item_key in zip(
2094 + [
2095 + self.counter_nixl_num_failed_transfers,
2096 + self.counter_nixl_num_failed_notifications,
2097 + ],
2098 + ["num_failed_transfers", "num_failed_notifications"],
2099 + ):
2100 + for list_item in transfer_stats_data[counter_item_key]:
2101 + counter_obj[engine_idx].inc(list_item)
vllm/v1/metrics/loggers.py
+15 -3
@@ -11,7 +11,10 @@ from prometheus_client import Counter, Gauge, Histogram
11
12 import vllm.envs as envs
13 from vllm.config import SupportsMetricsInfo, VllmConfig
14 -from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorLogging
14 +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
15 + KVConnectorLogging,
16 + KVConnectorPrometheus,
17 +)
18 from vllm.logger import init_logger
19 from vllm.plugins import load_plugins_by_group
20 from vllm.v1.engine import FinishReason
@@ -339,6 +342,7 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
342 _counter_cls = Counter
343 _histogram_cls = Histogram
344 _spec_decoding_cls = SpecDecodingProm
345 + _kv_connector_cls = KVConnectorPrometheus
346
347 def __init__(
348 self, vllm_config: VllmConfig, engine_indexes: list[int] | None = None
@@ -358,12 +362,15 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
362 model_name = vllm_config.model_config.served_model_name
363 max_model_len = vllm_config.model_config.max_model_len
364
361 - spec_decode_labelvalues: dict[int, list[str]] = {
365 + per_engine_labelvalues: dict[int, list[str]] = {
366 idx: [model_name, str(idx)] for idx in engine_indexes
367 }
368
369 self.spec_decoding_prom = self._spec_decoding_cls(
366 - vllm_config.speculative_config, labelnames, spec_decode_labelvalues
370 + vllm_config.speculative_config, labelnames, per_engine_labelvalues
371 + )
372 + self.kv_connector_prom = self._kv_connector_cls(
373 + vllm_config, labelnames, per_engine_labelvalues
374 )
375
376 #
@@ -962,6 +969,11 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
969 scheduler_stats.spec_decoding_stats, engine_idx
970 )
971
972 + if scheduler_stats.kv_connector_stats is not None:
973 + self.kv_connector_prom.observe(
974 + scheduler_stats.kv_connector_stats, engine_idx
975 + )
976 +
977 if mm_cache_stats is not None:
978 self.counter_mm_cache_queries[engine_idx].inc(mm_cache_stats.queries)
979 self.counter_mm_cache_hits[engine_idx].inc(mm_cache_stats.hits)
vllm/v1/metrics/ray_wrappers.py
+14
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import time
4
5 +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorPrometheus
6 from vllm.v1.metrics.loggers import PrometheusStatLogger
7 from vllm.v1.spec_decode.metrics import SpecDecodingProm
8
@@ -141,6 +142,18 @@ class RaySpecDecodingProm(SpecDecodingProm):
142 _counter_cls = RayCounterWrapper
143
144
145 +class RayKVConnectorPrometheus(KVConnectorPrometheus):
146 + """
147 + RayKVConnectorPrometheus is used by RayMetrics to log Ray
148 + metrics. Provides the same metrics as KV connectors but
149 + uses Ray's util.metrics library.
150 + """
151 +
152 + _gauge_cls = RayGaugeWrapper
153 + _counter_cls = RayCounterWrapper
154 + _histogram_cls = RayHistogramWrapper
155 +
156 +
157 class RayPrometheusStatLogger(PrometheusStatLogger):
158 """RayPrometheusStatLogger uses Ray metrics instead."""
159
@@ -148,6 +161,7 @@ class RayPrometheusStatLogger(PrometheusStatLogger):
161 _counter_cls = RayCounterWrapper
162 _histogram_cls = RayHistogramWrapper
163 _spec_decoding_cls = RaySpecDecodingProm
164 + _kv_connector_cls = RayKVConnectorPrometheus
165
166 @staticmethod
167 def _unregister_vllm_metrics():