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

[Rust Frontend] Record Mooncake/NIXL KV-connector metrics (#52755)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Markov Ilya <markovilya197@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Ilya Markov committed Sep 4, 2026 at 08:21 UTC ae2d1ca96fa2a98f92c775acfe8a154748495a27
10 files changed +600 -62
rust/src/engine-core-client/src/metrics.rs
+181 -3
@@ -7,11 +7,14 @@ use std::time::{SystemTime, UNIX_EPOCH};
7
8 use vllm_metrics::{
9 EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames,
10 - LoraInfoLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge,
10 + LoraInfoLabels, MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily,
11 + MooncakeOperationLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge,
12 WaitingReasonLabels,
13 };
14
14 -use crate::protocol::stats::SchedulerStats;
15 +use crate::protocol::stats::{
16 + KvConnectorStats, MooncakeStats, MultiConnectorStats, NixlStats, SchedulerStats,
17 +};
18 use crate::transport::ConnectedEngine;
19
20 const WAITING_REASON_CAPACITY: &str = "capacity";
@@ -57,6 +60,24 @@ struct SchedulerStatsHandles {
60 kv_block_idle_before_evict_seconds: HistogramMetric,
61 kv_block_reuse_gap_seconds: HistogramMetric,
62
63 + // Mooncake store connector telemetry, decoded from `kv_connector_stats`.
64 + // Kept as `Family` (not pre-resolved) because `operation`/`status` are
65 + // dynamic per-record labels.
66 + mooncake_operation_time_seconds: MooncakeOperationHistogramFamily,
67 + mooncake_operation_total: MooncakeOperationCounterFamily,
68 + mooncake_operation_keys_total: MooncakeOperationCounterFamily,
69 + mooncake_operation_bytes_total: MooncakeOperationCounterFamily,
70 + mooncake_operation_failed_keys_total: MooncakeOperationCounterFamily,
71 +
72 + // NIXL connector telemetry, decoded from `kv_connector_stats`.
73 + nixl_xfer_time_seconds: HistogramMetric,
74 + nixl_post_time_seconds: HistogramMetric,
75 + nixl_bytes_transferred: HistogramMetric,
76 + nixl_num_descriptors: HistogramMetric,
77 + nixl_num_failed_transfers: U64Counter,
78 + nixl_num_failed_notifications: U64Counter,
79 + nixl_num_kv_expired_reqs: U64Counter,
80 +
81 // Non-Prometheus interval accumulator for periodic text-log helpers.
82 log_stats: SchedulerLogStatsAccumulator,
83 }
@@ -150,6 +171,20 @@ fn resolve_scheduler_stats_handles(
171 .kv_block_idle_before_evict_seconds
172 .get_or_create_owned(&labels),
173 kv_block_reuse_gap_seconds: metrics.kv_block_reuse_gap_seconds.get_or_create_owned(&labels),
174 + mooncake_operation_time_seconds: metrics.mooncake_operation_time_seconds.clone(),
175 + mooncake_operation_total: metrics.mooncake_operation_total.clone(),
176 + mooncake_operation_keys_total: metrics.mooncake_operation_keys_total.clone(),
177 + mooncake_operation_bytes_total: metrics.mooncake_operation_bytes_total.clone(),
178 + mooncake_operation_failed_keys_total: metrics.mooncake_operation_failed_keys_total.clone(),
179 + nixl_xfer_time_seconds: metrics.nixl_xfer_time_seconds.get_or_create_owned(&labels),
180 + nixl_post_time_seconds: metrics.nixl_post_time_seconds.get_or_create_owned(&labels),
181 + nixl_bytes_transferred: metrics.nixl_bytes_transferred.get_or_create_owned(&labels),
182 + nixl_num_descriptors: metrics.nixl_num_descriptors.get_or_create_owned(&labels),
183 + nixl_num_failed_transfers: metrics.nixl_num_failed_transfers.get_or_create_owned(&labels),
184 + nixl_num_failed_notifications: metrics
185 + .nixl_num_failed_notifications
186 + .get_or_create_owned(&labels),
187 + nixl_num_kv_expired_reqs: metrics.nixl_num_kv_expired_reqs.get_or_create_owned(&labels),
188 labels,
189 }
190 }
@@ -236,6 +271,79 @@ fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: &
271 }
272 }
273 }
274 +
275 + // Connector-specific KV transfer stats. A bare connector reports its own
276 + // flat payload; MultiConnector reports connector class name -> flat child
277 + // payload.
278 + if let Some(kv_connector_stats) = &stats.kv_connector_stats {
279 + match kv_connector_stats {
280 + KvConnectorStats::Nixl(stats) => record_nixl_stats(handles, stats),
281 + KvConnectorStats::Mooncake(stats) => record_mooncake_stats(handles, stats),
282 + KvConnectorStats::Multi(stats) => record_multi_connector_stats(handles, stats),
283 + KvConnectorStats::Other(_) => {}
284 + }
285 + }
286 +}
287 +
288 +fn record_multi_connector_stats(handles: &SchedulerStatsHandles, stats: &MultiConnectorStats) {
289 + for nixl in [&stats.nixl, &stats.nixl_pull, &stats.nixl_push].into_iter().flatten() {
290 + record_nixl_stats(handles, nixl);
291 + }
292 + if let Some(mooncake) = &stats.mooncake {
293 + record_mooncake_stats(handles, mooncake);
294 + }
295 +}
296 +
297 +fn record_mooncake_stats(handles: &SchedulerStatsHandles, stats: &MooncakeStats) {
298 + for (operation, records) in &stats.0 {
299 + for record in records {
300 + let labels = MooncakeOperationLabels {
301 + model_name: handles.labels.model_name.clone(),
302 + engine: handles.labels.engine,
303 + operation: operation.as_str().to_string(),
304 + status: record.status.as_str().to_string(),
305 + };
306 + handles
307 + .mooncake_operation_time_seconds
308 + .get_or_create(&labels)
309 + .observe(record.duration_seconds);
310 + handles.mooncake_operation_total.get_or_create(&labels).inc();
311 + handles
312 + .mooncake_operation_keys_total
313 + .get_or_create(&labels)
314 + .inc_by(record.num_keys);
315 + handles
316 + .mooncake_operation_bytes_total
317 + .get_or_create(&labels)
318 + .inc_by(record.num_bytes);
319 + handles
320 + .mooncake_operation_failed_keys_total
321 + .get_or_create(&labels)
322 + .inc_by(record.num_failed_keys);
323 + }
324 + }
325 +}
326 +
327 +fn record_nixl_stats(handles: &SchedulerStatsHandles, stats: &NixlStats) {
328 + for value in &stats.transfer_duration {
329 + handles.nixl_xfer_time_seconds.observe(*value);
330 + }
331 + for value in &stats.post_duration {
332 + handles.nixl_post_time_seconds.observe(*value);
333 + }
334 + for value in &stats.bytes_transferred {
335 + handles.nixl_bytes_transferred.observe(*value as f64);
336 + }
337 + for value in &stats.num_descriptors {
338 + handles.nixl_num_descriptors.observe(*value as f64);
339 + }
340 + handles
341 + .nixl_num_failed_transfers
342 + .inc_by(stats.num_failed_transfers.iter().sum());
343 + handles
344 + .nixl_num_failed_notifications
345 + .inc_by(stats.num_failed_notifications.iter().sum());
346 + handles.nixl_num_kv_expired_reqs.inc_by(stats.num_kv_expired_reqs.iter().sum());
347 }
348
349 /// Exports `vllm:lora_requests_info` as a single series covering all LoRA
@@ -286,12 +394,16 @@ fn now_unix_secs() -> f64 {
394
395 #[cfg(test)]
396 mod tests {
289 - use std::collections::BTreeSet;
397 + use std::collections::{BTreeMap, BTreeSet};
398
399 use expect_test::expect;
400 use vllm_metrics::Metrics;
401
402 use crate::metrics::LoraInfoExporter;
403 + use crate::protocol::stats::{
404 + KvConnectorStats, MooncakeOperation, MooncakeRecord, MooncakeStats, MooncakeStatus,
405 + MultiConnectorStats, NixlStats, SchedulerStats,
406 + };
407
408 fn names(values: &[&str]) -> BTreeSet<String> {
409 values.iter().map(|name| (*name).to_string()).collect()
@@ -345,4 +457,70 @@ mod tests {
457 exporter.update(&metrics.scheduler, names(&[]), names(&[]));
458 expect![[""]].assert_eq(&lora_series(&metrics.render().unwrap()));
459 }
460 +
461 + fn nixl_stats() -> NixlStats {
462 + NixlStats {
463 + transfer_duration: vec![0.01, 0.02],
464 + post_duration: vec![0.001, 0.002],
465 + bytes_transferred: vec![4096, 8192],
466 + num_descriptors: vec![2, 4],
467 + num_failed_transfers: vec![],
468 + num_failed_notifications: vec![],
469 + num_kv_expired_reqs: vec![1],
470 + }
471 + }
472 +
473 + fn mooncake_stats() -> MooncakeStats {
474 + MooncakeStats(BTreeMap::from([(
475 + MooncakeOperation::LoadGet,
476 + vec![MooncakeRecord {
477 + duration_seconds: 0.05,
478 + num_keys: 3,
479 + num_bytes: 1024,
480 + status: MooncakeStatus::Ok,
481 + num_failed_keys: 0,
482 + }],
483 + )]))
484 + }
485 +
486 + /// Records one MultiConnector payload into Mooncake and NIXL metrics.
487 + #[test]
488 + fn kv_connector_stats_are_recorded_into_mooncake_and_nixl_metrics() {
489 + let metrics = Metrics::new();
490 + let handles = super::resolve_scheduler_stats_handles(&metrics.scheduler, "model", 0);
491 +
492 + let stats = SchedulerStats {
493 + kv_connector_stats: Some(KvConnectorStats::Multi(Box::new(MultiConnectorStats {
494 + nixl: Some(nixl_stats()),
495 + mooncake: Some(mooncake_stats()),
496 + ..Default::default()
497 + }))),
498 + ..Default::default()
499 + };
500 +
501 + super::record_scheduler_stats_with_handles(&handles, &stats);
502 +
503 + let rendered = metrics.render().unwrap();
504 + assert!(rendered.contains(
505 + "vllm:mooncake_store_operation_total{model_name=\"model\",engine=\"0\",\
506 + operation=\"load_get\",status=\"ok\"} 1"
507 + ));
508 + assert!(rendered.contains(
509 + "vllm:mooncake_store_operation_keys_total{model_name=\"model\",engine=\"0\",\
510 + operation=\"load_get\",status=\"ok\"} 3"
511 + ));
512 + assert!(rendered.contains(
513 + "vllm:mooncake_store_operation_bytes_total{model_name=\"model\",engine=\"0\",\
514 + operation=\"load_get\",status=\"ok\"} 1024"
515 + ));
516 + assert!(
517 + rendered.contains(
518 + "vllm:nixl_num_kv_expired_reqs_total{model_name=\"model\",engine=\"0\"} 1"
519 + )
520 + );
521 + assert!(
522 + rendered
523 + .contains("vllm:nixl_xfer_time_seconds_count{model_name=\"model\",engine=\"0\"} 2")
524 + );
525 + }
526 }
rust/src/engine-core-client/src/protocol/stats.rs
+116 -4
@@ -5,8 +5,6 @@ use std::collections::BTreeMap;
5
6 use serde::{Deserialize, Serialize};
7
8 -use crate::protocol::OpaqueValue;
9 -
8 /// Stores cache hit statistics.
9 ///
10 /// Original Python definition:
@@ -158,6 +156,120 @@ pub struct CudagraphStats {
156 pub runtime_mode: String,
157 }
158
159 +/// KV connector telemetry DTOs carried by [`SchedulerStats`].
160 +pub mod kv_connector {
161 + use std::collections::BTreeMap;
162 +
163 + use serde::{Deserialize, Serialize};
164 +
165 + use crate::protocol::OpaqueValue;
166 +
167 + /// NIXL connector transfer telemetry.
168 + #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
169 + #[serde(deny_unknown_fields)]
170 + pub struct NixlStats {
171 + /// Time spent transferring each successful request.
172 + pub transfer_duration: Vec<f64>,
173 + /// Time spent posting each successful transfer.
174 + pub post_duration: Vec<f64>,
175 + /// Bytes transferred by each successful request.
176 + pub bytes_transferred: Vec<u64>,
177 + /// Descriptor count for each successful request.
178 + pub num_descriptors: Vec<u64>,
179 + /// Failure counter increments collected since the previous update.
180 + pub num_failed_transfers: Vec<u64>,
181 + /// Notification failure counter increments collected since the previous update.
182 + pub num_failed_notifications: Vec<u64>,
183 + /// Expired-request counter increments collected since the previous update.
184 + pub num_kv_expired_reqs: Vec<u64>,
185 + }
186 +
187 + /// Mooncake store operation name.
188 + #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
189 + #[serde(rename_all = "snake_case")]
190 + pub enum MooncakeOperation {
191 + SaveExists,
192 + SavePut,
193 + LoadGet,
194 + LookupExists,
195 + }
196 +
197 + impl MooncakeOperation {
198 + pub(crate) const fn as_str(self) -> &'static str {
199 + match self {
200 + Self::SaveExists => "save_exists",
201 + Self::SavePut => "save_put",
202 + Self::LoadGet => "load_get",
203 + Self::LookupExists => "lookup_exists",
204 + }
205 + }
206 + }
207 +
208 + /// Mooncake store operation status.
209 + #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
210 + #[serde(rename_all = "snake_case")]
211 + pub enum MooncakeStatus {
212 + Ok,
213 + Error,
214 + PartialFailure,
215 + }
216 +
217 + impl MooncakeStatus {
218 + pub(crate) const fn as_str(self) -> &'static str {
219 + match self {
220 + Self::Ok => "ok",
221 + Self::Error => "error",
222 + Self::PartialFailure => "partial_failure",
223 + }
224 + }
225 + }
226 +
227 + /// One Mooncake store operation sample.
228 + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
229 + #[serde(deny_unknown_fields)]
230 + pub struct MooncakeRecord {
231 + pub duration_seconds: f64,
232 + pub num_keys: u64,
233 + pub num_bytes: u64,
234 + pub status: MooncakeStatus,
235 + pub num_failed_keys: u64,
236 + }
237 +
238 + /// Mooncake store telemetry grouped by operation.
239 + #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
240 + #[serde(transparent)]
241 + pub struct MooncakeStats(pub BTreeMap<MooncakeOperation, Vec<MooncakeRecord>>);
242 +
243 + /// Telemetry emitted by a `MultiConnector` with a flat child payload map.
244 + #[serde_with::skip_serializing_none]
245 + #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
246 + pub struct MultiConnectorStats {
247 + #[serde(rename = "NixlConnector")]
248 + pub nixl: Option<NixlStats>,
249 + #[serde(rename = "NixlPullConnector")]
250 + pub nixl_pull: Option<NixlStats>,
251 + #[serde(rename = "NixlPushConnector")]
252 + pub nixl_push: Option<NixlStats>,
253 + #[serde(rename = "MooncakeStoreConnector")]
254 + pub mooncake: Option<MooncakeStats>,
255 + /// Child connector payloads without a Rust telemetry implementation.
256 + #[serde(flatten)]
257 + pub other: BTreeMap<String, OpaqueValue>,
258 + }
259 +
260 + /// Connector-specific scheduler telemetry.
261 + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
262 + #[serde(untagged)]
263 + pub enum KvConnectorStats {
264 + Nixl(NixlStats),
265 + Mooncake(MooncakeStats),
266 + Multi(Box<MultiConnectorStats>),
267 + Other(BTreeMap<String, OpaqueValue>),
268 + }
269 +}
270 +
271 +pub use kv_connector::*;
272 +
273 /// Stats associated with the scheduler.
274 ///
275 /// Original Python definition:
@@ -185,8 +297,8 @@ pub struct SchedulerStats {
297 pub kv_cache_eviction_events: Vec<KvCacheEvictionEvent>,
298 /// Speculative decoding scheduler stats, when enabled.
299 pub spec_decoding_stats: Option<SpecDecodingStats>,
188 - /// Connector-specific KV transfer stats, kept opaque for now.
189 - pub kv_connector_stats: Option<BTreeMap<String, OpaqueValue>>,
300 + /// Connector-specific KV transfer stats.
301 + pub kv_connector_stats: Option<KvConnectorStats>,
302 /// CUDA graph runtime stats when graph metrics are enabled.
303 pub cudagraph_stats: Option<CudagraphStats>,
304 /// Estimated MFU/performance stats, when enabled.
rust/src/engine-core-client/src/tests/client.rs
+27 -1
@@ -31,7 +31,7 @@ use crate::protocol::output::{
31 };
32 use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
33 use crate::protocol::sampling::EngineCoreSamplingParams;
34 -use crate::protocol::stats::SchedulerStats;
34 +use crate::protocol::stats::{KvConnectorStats, MooncakeOperation, SchedulerStats};
35 use crate::protocol::tensor::{WireArrayData, WireTensor};
36 use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
37 use crate::test_utils::{
@@ -2635,6 +2635,10 @@ fn python_msgpack_fixtures_match_rust_encoding() {
2635 let inline_prompt_frames = lines.next().expect("missing inline prompt logprobs fixture line");
2636 let multipart_prompt_frames =
2637 lines.next().expect("missing multipart prompt logprobs fixture line");
2638 + let nixl_stats_hex = lines.next().expect("missing NIXL stats fixture line");
2639 + let mooncake_stats_hex = lines.next().expect("missing Mooncake stats fixture line");
2640 + let multi_connector_stats_hex =
2641 + lines.next().expect("missing MultiConnector stats fixture line");
2642 let ready_response_hex = lines.next().expect("missing ready response fixture line");
2643
2644 let request_bytes = hex::decode(request_hex).unwrap();
@@ -2799,6 +2803,28 @@ fn python_msgpack_fixtures_match_rust_encoding() {
2803 .expect("multipart prompt logprobs decoded"),
2804 );
2805
2806 + let nixl_stats: KvConnectorStats =
2807 + rmp_serde::from_slice(&hex::decode(nixl_stats_hex).unwrap()).unwrap();
2808 + assert!(matches!(nixl_stats, KvConnectorStats::Nixl(_)));
2809 +
2810 + let mooncake_stats: KvConnectorStats =
2811 + rmp_serde::from_slice(&hex::decode(mooncake_stats_hex).unwrap()).unwrap();
2812 + assert!(matches!(
2813 + mooncake_stats,
2814 + KvConnectorStats::Mooncake(stats)
2815 + if stats.0.contains_key(&MooncakeOperation::LoadGet)
2816 + ));
2817 +
2818 + let multi_connector_stats: KvConnectorStats =
2819 + rmp_serde::from_slice(&hex::decode(multi_connector_stats_hex).unwrap()).unwrap();
2820 + assert!(matches!(
2821 + multi_connector_stats,
2822 + KvConnectorStats::Multi(stats)
2823 + if stats.nixl.is_some()
2824 + && stats.mooncake.is_some()
2825 + && stats.other.contains_key("UnsupportedConnector")
2826 + ));
2827 +
2828 let map_keys = |bytes: &[u8]| -> BTreeSet<String> {
2829 match decode_value(bytes) {
2830 Value::Map(entries) => entries
rust/src/engine-core-client/src/tests/python_compat.py
+29
@@ -456,6 +456,32 @@ ready_response = EngineCoreReadyResponse(
456 ),
457 )
458
459 +nixl_stats = {
460 + "transfer_duration": [0.01, 0.02],
461 + "post_duration": [0.001, 0.002],
462 + "bytes_transferred": [4096, 8192],
463 + "num_descriptors": [2, 4],
464 + "num_failed_transfers": [],
465 + "num_failed_notifications": [],
466 + "num_kv_expired_reqs": [1],
467 +}
468 +mooncake_stats = {
469 + "load_get": [
470 + {
471 + "duration_seconds": 0.05,
472 + "num_keys": 3,
473 + "num_bytes": 1024,
474 + "status": "ok",
475 + "num_failed_keys": 0,
476 + }
477 + ]
478 +}
479 +multi_connector_stats = {
480 + "NixlConnector": nixl_stats,
481 + "MooncakeStoreConnector": mooncake_stats,
482 + "UnsupportedConnector": {"sample_count": 1},
483 +}
484 +
485 print(msgspec.msgpack.encode(request).hex())
486 print(msgspec.msgpack.encode(defaults_request).hex())
487 print(msgpack.packb(multimodal_request_wire, use_bin_type=True).hex())
@@ -475,4 +501,7 @@ print(
501 for frame in encode_output_frames(multipart_prompt_logprobs, size_threshold=1)
502 )
503 )
504 +print(msgspec.msgpack.encode(nixl_stats).hex())
505 +print(msgspec.msgpack.encode(mooncake_stats).hex())
506 +print(msgspec.msgpack.encode(multi_connector_stats).hex())
507 print(msgspec.msgpack.encode(ready_response).hex())
rust/src/metrics/src/scheduler.rs
+195
@@ -29,6 +29,58 @@ fn kv_block_reuse_gap_histogram() -> Histogram {
29 Histogram::new(KV_CACHE_RESIDENCY_BUCKETS.iter().copied())
30 }
31
32 +// Buckets copied verbatim from the Python `MooncakeStorePromMetrics` /
33 +// `NixlPromMetrics` classes so dashboards built against either frontend are
34 +// interchangeable.
35 +const MOONCAKE_OPERATION_TIME_BUCKETS: [f64; 15] = [
36 + 1e-3, 5e-3, 1e-2, 5e-2, 1e-1, 2e-1, 3e-1, 4e-1, 5e-1, 7.5e-1, 1.0, 1.5, 2.0, 3.0, 4.0,
37 +];
38 +const NIXL_XFER_TIME_BUCKETS: [f64; 12] = [
39 + 0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.2, 0.3, 0.5, 0.75, 1.0, 5.0,
40 +];
41 +const NIXL_POST_TIME_BUCKETS: [f64; 13] = [
42 + 0.001, 0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.2, 0.3, 0.5, 0.75, 1.0, 5.0,
43 +];
44 +// Uniform 2KB to 8GB range: 2**(10+i) for i in range(1, 25, 2).
45 +const NIXL_BYTES_TRANSFERRED_BUCKETS: [f64; 12] = [
46 + 2048.0,
47 + 8192.0,
48 + 32768.0,
49 + 131072.0,
50 + 524288.0,
51 + 2097152.0,
52 + 8388608.0,
53 + 33554432.0,
54 + 134217728.0,
55 + 536870912.0,
56 + 2147483648.0,
57 + 8589934592.0,
58 +];
59 +const NIXL_NUM_DESCRIPTORS_BUCKETS: [f64; 14] = [
60 + 10.0, 20.0, 30.0, 50.0, 75.0, 100.0, 200.0, 400.0, 1000.0, 2000.0, 4000.0, 10000.0, 20000.0,
61 + 50000.0,
62 +];
63 +
64 +fn mooncake_operation_time_histogram() -> Histogram {
65 + Histogram::new(MOONCAKE_OPERATION_TIME_BUCKETS.iter().copied())
66 +}
67 +
68 +fn nixl_xfer_time_histogram() -> Histogram {
69 + Histogram::new(NIXL_XFER_TIME_BUCKETS.iter().copied())
70 +}
71 +
72 +fn nixl_post_time_histogram() -> Histogram {
73 + Histogram::new(NIXL_POST_TIME_BUCKETS.iter().copied())
74 +}
75 +
76 +fn nixl_bytes_transferred_histogram() -> Histogram {
77 + Histogram::new(NIXL_BYTES_TRANSFERRED_BUCKETS.iter().copied())
78 +}
79 +
80 +fn nixl_num_descriptors_histogram() -> Histogram {
81 + Histogram::new(NIXL_NUM_DESCRIPTORS_BUCKETS.iter().copied())
82 +}
83 +
84 #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
85 pub struct EngineLabels {
86 pub model_name: String,
@@ -49,6 +101,24 @@ pub struct WaitingReasonLabels {
101 pub reason: &'static str,
102 }
103
104 +/// Labels for per-operation Mooncake store connector telemetry. `operation`
105 +/// and `status` are dynamic (not known ahead of time), unlike the other
106 +/// fixed-label families in this module.
107 +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
108 +pub struct MooncakeOperationLabels {
109 + pub model_name: String,
110 + pub engine: u32,
111 + pub operation: String,
112 + pub status: String,
113 +}
114 +
115 +/// Family type aliases for the Mooncake connector telemetry, exposed so
116 +/// `engine-core-client` can hold onto the `Family` directly (rather than a
117 +/// pre-resolved handle) since `operation`/`status` vary per record.
118 +pub type MooncakeOperationHistogramFamily =
119 + Family<MooncakeOperationLabels, Histogram, fn() -> Histogram>;
120 +pub type MooncakeOperationCounterFamily = Family<MooncakeOperationLabels, U64Counter>;
121 +
122 /// Adapter names encoded as a deterministic comma-joined Prometheus label value.
123 #[derive(Clone, Debug, Hash, PartialEq, Eq)]
124 pub struct LoraAdapterNames(pub BTreeSet<String>);
@@ -184,6 +254,22 @@ pub struct SchedulerMetrics {
254 pub kv_block_idle_before_evict_seconds: HistogramFamily,
255 pub kv_block_reuse_gap_seconds: HistogramFamily,
256
257 + // Mooncake store connector telemetry. Mirrors `MooncakeStorePromMetrics`.
258 + pub mooncake_operation_time_seconds: MooncakeOperationHistogramFamily,
259 + pub mooncake_operation_total: MooncakeOperationCounterFamily,
260 + pub mooncake_operation_keys_total: MooncakeOperationCounterFamily,
261 + pub mooncake_operation_bytes_total: MooncakeOperationCounterFamily,
262 + pub mooncake_operation_failed_keys_total: MooncakeOperationCounterFamily,
263 +
264 + // NIXL connector telemetry. Mirrors `NixlPromMetrics`.
265 + pub nixl_xfer_time_seconds: HistogramFamily,
266 + pub nixl_post_time_seconds: HistogramFamily,
267 + pub nixl_bytes_transferred: HistogramFamily,
268 + pub nixl_num_descriptors: HistogramFamily,
269 + pub nixl_num_failed_transfers: Family<EngineLabels, U64Counter>,
270 + pub nixl_num_failed_notifications: Family<EngineLabels, U64Counter>,
271 + pub nixl_num_kv_expired_reqs: Family<EngineLabels, U64Counter>,
272 +
273 /// Non-Prometheus interval accumulators for periodic text-log helpers.
274 pub log_stats: Family<EngineLabels, SchedulerLogStatsAccumulator>,
275 }
@@ -336,6 +422,103 @@ impl SchedulerMetrics {
422 kv_block_reuse_gap_seconds.clone(),
423 );
424
425 + // Mooncake store connector telemetry.
426 + let mooncake_operation_time_seconds =
427 + Family::new_with_constructor(mooncake_operation_time_histogram as fn() -> Histogram);
428 + registry.register(
429 + "vllm:mooncake_store_operation_time_seconds",
430 + "Histogram of Mooncake store communication time.",
431 + mooncake_operation_time_seconds.clone(),
432 + );
433 +
434 + // NOTE: registration names below omit the `_total` suffix that the
435 + // Python metric names carry (e.g. `vllm:mooncake_store_operation_total`)
436 + // because this crate's `prometheus-client` (unlike Python's, which
437 + // strips a user-supplied `_total` before re-adding it) appends
438 + // `_total` unconditionally for counters; see the note in `lib.rs`.
439 + let mooncake_operation_total = Family::default();
440 + registry.register(
441 + "vllm:mooncake_store_operation",
442 + "Number of Mooncake store communication operations.",
443 + mooncake_operation_total.clone(),
444 + );
445 +
446 + let mooncake_operation_keys_total = Family::default();
447 + registry.register(
448 + "vllm:mooncake_store_operation_keys",
449 + "Number of Mooncake store keys touched by operations.",
450 + mooncake_operation_keys_total.clone(),
451 + );
452 +
453 + let mooncake_operation_bytes_total = Family::default();
454 + registry.register(
455 + "vllm:mooncake_store_operation_bytes",
456 + "Number of bytes transferred by Mooncake store operations.",
457 + mooncake_operation_bytes_total.clone(),
458 + );
459 +
460 + let mooncake_operation_failed_keys_total = Family::default();
461 + registry.register(
462 + "vllm:mooncake_store_operation_failed_keys",
463 + "Number of Mooncake store keys that failed in operations.",
464 + mooncake_operation_failed_keys_total.clone(),
465 + );
466 +
467 + // NIXL connector telemetry.
468 + let nixl_xfer_time_seconds =
469 + Family::new_with_constructor(nixl_xfer_time_histogram as fn() -> Histogram);
470 + registry.register(
471 + "vllm:nixl_xfer_time_seconds",
472 + "Histogram of transfer duration for NIXL KV Cache transfers.",
473 + nixl_xfer_time_seconds.clone(),
474 + );
475 +
476 + let nixl_post_time_seconds =
477 + Family::new_with_constructor(nixl_post_time_histogram as fn() -> Histogram);
478 + registry.register(
479 + "vllm:nixl_post_time_seconds",
480 + "Histogram of transfer post time for NIXL KV Cache transfers.",
481 + nixl_post_time_seconds.clone(),
482 + );
483 +
484 + let nixl_bytes_transferred =
485 + Family::new_with_constructor(nixl_bytes_transferred_histogram as fn() -> Histogram);
486 + registry.register(
487 + "vllm:nixl_bytes_transferred",
488 + "Histogram of bytes transferred per NIXL KV Cache transfers.",
489 + nixl_bytes_transferred.clone(),
490 + );
491 +
492 + let nixl_num_descriptors =
493 + Family::new_with_constructor(nixl_num_descriptors_histogram as fn() -> Histogram);
494 + registry.register(
495 + "vllm:nixl_num_descriptors",
496 + "Histogram of number of descriptors per NIXL KV Cache transfers.",
497 + nixl_num_descriptors.clone(),
498 + );
499 +
500 + let nixl_num_failed_transfers = Family::default();
501 + registry.register(
502 + "vllm:nixl_num_failed_transfers",
503 + "Number of failed NIXL KV Cache transfers.",
504 + nixl_num_failed_transfers.clone(),
505 + );
506 +
507 + let nixl_num_failed_notifications = Family::default();
508 + registry.register(
509 + "vllm:nixl_num_failed_notifications",
510 + "Number of failed NIXL KV Cache notifications.",
511 + nixl_num_failed_notifications.clone(),
512 + );
513 +
514 + let nixl_num_kv_expired_reqs = Family::default();
515 + registry.register(
516 + "vllm:nixl_num_kv_expired_reqs",
517 + "Number of requests that had their KV expire. \
518 + NOTE: This metric is tracked on the P instance.",
519 + nixl_num_kv_expired_reqs.clone(),
520 + );
521 +
522 Self {
523 scheduler_running,
524 scheduler_waiting,
@@ -356,6 +539,18 @@ impl SchedulerMetrics {
539 kv_block_lifetime_seconds,
540 kv_block_idle_before_evict_seconds,
541 kv_block_reuse_gap_seconds,
542 + mooncake_operation_time_seconds,
543 + mooncake_operation_total,
544 + mooncake_operation_keys_total,
545 + mooncake_operation_bytes_total,
546 + mooncake_operation_failed_keys_total,
547 + nixl_xfer_time_seconds,
548 + nixl_post_time_seconds,
549 + nixl_bytes_transferred,
550 + nixl_num_descriptors,
551 + nixl_num_failed_transfers,
552 + nixl_num_failed_notifications,
553 + nixl_num_kv_expired_reqs,
554 log_stats: Family::default(),
555 }
556 }
tests/v1/kv_connector/unit/test_multi_connector.py
+38 -43
@@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
23 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
24 from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
25 MultiConnector,
26 + MultiKVConnectorPromMetrics,
27 MultiKVConnectorStats,
28 MultiKVConnectorWorkerMetadata,
29 )
@@ -559,14 +560,12 @@ class TestMultiConnectorStats:
560 correct data."""
561 serialized_data = {
562 "NixlConnector": {
562 - "data": {
563 - "transfer_duration": [1.5, 2.3],
564 - "post_duration": [0.1, 0.2],
565 - "bytes_transferred": [1024, 2048],
566 - "num_descriptors": [10, 20],
567 - "num_failed_transfers": [],
568 - "num_failed_notifications": [],
569 - }
563 + "transfer_duration": [1.5, 2.3],
564 + "post_duration": [0.1, 0.2],
565 + "bytes_transferred": [1024, 2048],
566 + "num_descriptors": [10, 20],
567 + "num_failed_transfers": [],
568 + "num_failed_notifications": [],
569 }
570 }
571
@@ -584,16 +583,14 @@ class TestMultiConnectorStats:
583 """Test reconstruction with multiple connector types that have custom stats."""
584 serialized_data = {
585 "NixlConnector": {
587 - "data": {
588 - "transfer_duration": [1.5],
589 - "post_duration": [0.1],
590 - "bytes_transferred": [1024],
591 - "num_descriptors": [10],
592 - "num_failed_transfers": [],
593 - "num_failed_notifications": [],
594 - }
586 + "transfer_duration": [1.5],
587 + "post_duration": [0.1],
588 + "bytes_transferred": [1024],
589 + "num_descriptors": [10],
590 + "num_failed_transfers": [],
591 + "num_failed_notifications": [],
592 },
596 - "MockConnector": {"data": {"mock_field": [1, 2, 3]}},
593 + "MockConnector": {"mock_field": [1, 2, 3]},
594 }
595
596 stats = MultiConnector.build_kv_connector_stats(data=serialized_data)
@@ -608,20 +605,19 @@ class TestMultiConnectorStats:
605 assert isinstance(stats.data["MockConnector"], MockConnectorStats)
606 # Verify data is preserved
607 assert stats.data["MockConnector"].data == {"mock_field": [1, 2, 3]}
608 + assert stats.to_dict() == serialized_data
609
610 def test_build_kv_connector_stats_raises_error_for_unknown_connector(self):
611 """Test that unknown connectors raise an error."""
612 serialized_data = {
615 - "UnknownConnector": {"data": {"some_field": [1, 2, 3]}},
613 + "UnknownConnector": {"some_field": [1, 2, 3]},
614 "NixlConnector": {
617 - "data": {
618 - "transfer_duration": [1.5],
619 - "post_duration": [0.1],
620 - "bytes_transferred": [1024],
621 - "num_descriptors": [10],
622 - "num_failed_transfers": [],
623 - "num_failed_notifications": [],
624 - }
615 + "transfer_duration": [1.5],
616 + "post_duration": [0.1],
617 + "bytes_transferred": [1024],
618 + "num_descriptors": [10],
619 + "num_failed_transfers": [],
620 + "num_failed_notifications": [],
621 },
622 }
623
@@ -675,7 +671,7 @@ class TestMultiConnectorStats:
671
672 mixed_data = {
673 "NixlConnector": nixl_stats, # Already instantiated
678 - "MockConnector": {"data": {"mock_field": [1, 2, 3]}}, # Serialized
674 + "MockConnector": {"mock_field": [1, 2, 3]}, # Serialized
675 }
676
677 stats = MultiConnector.build_kv_connector_stats(data=mixed_data)
@@ -695,16 +691,14 @@ class TestMultiConnectorStats:
691 # so it returns None and should be skipped
692 serialized_data = {
693 "NixlConnector": {
698 - "data": {
699 - "transfer_duration": [1.5],
700 - "post_duration": [0.1],
701 - "bytes_transferred": [1024],
702 - "num_descriptors": [10],
703 - "num_failed_transfers": [],
704 - "num_failed_notifications": [],
705 - }
694 + "transfer_duration": [1.5],
695 + "post_duration": [0.1],
696 + "bytes_transferred": [1024],
697 + "num_descriptors": [10],
698 + "num_failed_transfers": [],
699 + "num_failed_notifications": [],
700 },
707 - "ExampleConnector": {"data": {"some_field": [1, 2, 3]}},
701 + "ExampleConnector": {"some_field": [1, 2, 3]},
702 }
703
704 stats = MultiConnector.build_kv_connector_stats(data=serialized_data)
@@ -718,14 +712,15 @@ class TestMultiConnectorStats:
712 # ExampleConnector should be skipped (returns None)
713 assert "ExampleConnector" not in stats.data
714
721 - def test_build_kv_connector_stats_handles_malformed_data(self):
722 - """Test that malformed data raises appropriate errors."""
723 - serialized_data = {
724 - "NixlConnector": {"wrong_field": {"transfer_duration": [1.5]}}
725 - }
715 + def test_prom_metrics_passes_flat_data_to_children(self):
716 + child_metrics = MagicMock()
717 + metrics = object.__new__(MultiKVConnectorPromMetrics)
718 + metrics._prom_metrics = {"NixlConnector": child_metrics}
719 + payload = {"NixlConnector": {"transfer_duration": [1.5]}}
720
727 - with pytest.raises(AssertionError, match="Expected a dict with a 'data' field"):
728 - MultiConnector.build_kv_connector_stats(data=serialized_data)
721 + metrics.observe(payload, engine_idx=2)
722 +
723 + child_metrics.observe.assert_called_once_with({"transfer_duration": [1.5]}, 2)
724
725 def test_aggregate_same_connector(self):
726 """Test aggregating stats from the same connector type."""
tests/v1/kv_connector/unit/test_nixl_connector.py
+2 -1
@@ -1709,8 +1709,9 @@ def test_scheduler_kv_connector_stats_aggregation():
1709 final_stats = next(
1710 iter(engine_core_outputs.values())
1711 ).scheduler_stats.kv_connector_stats
1712 + # The scheduler stats payload carries serialized per-connector dicts.
1713 nixl_stats = final_stats["NixlConnector"]
1713 - assert nixl_stats.num_successful_transfers == 2
1714 + assert len(nixl_stats["transfer_duration"]) == 2
1715
1716
1717 @pytest.mark.parametrize("distributed_executor_backend", ["ray", None])
vllm/distributed/kv_transfer/kv_connector/v1/metrics.py
+4
@@ -26,6 +26,10 @@ class KVConnectorStats:
26
27 data: dict[str, Any] = field(default_factory=dict)
28
29 + def to_dict(self) -> dict[str, Any]:
30 + """Return the serializable connector stats payload."""
31 + return self.data
32 +
33 def reset(self):
34 """Reset the stats, clear the state."""
35 raise NotImplementedError
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+7 -9
@@ -75,6 +75,11 @@ class MultiKVConnectorStats(KVConnectorStats):
75 This is used to aggregate the stats from all connectors separately.
76 """
77
78 + def to_dict(self) -> dict[str, Any]:
79 + return {
80 + connector_id: stats.to_dict() for connector_id, stats in self.data.items()
81 + }
82 +
83 def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
84 for connector_id, stats in other.data.items():
85 if connector_id not in self.data:
@@ -122,7 +127,7 @@ class MultiKVConnectorPromMetrics(KVConnectorPromMetrics):
127 f"{connector_id} is not contained in the list of registered connectors "
128 f"with Prometheus metrics support: {self._prom_metrics.keys()}"
129 )
125 - self._prom_metrics[connector_id].observe(stats_data["data"], engine_idx)
130 + self._prom_metrics[connector_id].observe(stats_data, engine_idx)
131
132
133 class MultiConnector(KVConnectorBase_V1, SupportsHMA):
@@ -617,16 +622,9 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
622 connector_name
623 )
624
620 - # stats_value is the serialized dataclass which contains {'data': {...}}
621 - # We need to extract the inner 'data' field to avoid double-nesting
622 - assert isinstance(stats_value, dict) and "data" in stats_value, (
623 - f"Expected a dict with a 'data' field, got {stats_value}"
624 - )
625 - inner_data = stats_value["data"]
626 -
625 # Use the connector's build_kv_connector_stats to reconstruct
626 if reconstructed_stats := connector_cls.build_kv_connector_stats(
629 - data=inner_data
627 + data=stats_value
628 ):
629 reconstructed_data[connector_name] = reconstructed_stats
630
vllm/v1/core/sched/scheduler.py
+1 -1
@@ -2704,7 +2704,7 @@ class Scheduler(SchedulerInterface):
2704 )
2705 spec_stats = spec_decoding_stats
2706 connector_stats_payload = (
2707 - kv_connector_stats.data if kv_connector_stats else None
2707 + kv_connector_stats.to_dict() if kv_connector_stats else None
2708 )
2709 return SchedulerStats(
2710 num_running_reqs=len(self.running),