[Rust Frontend] Export Mooncake connector transfer metrics like Python

rust/src/engine-core-client/src/metrics.rs

@@ -13,8 +13,8 @@ use vllm_metrics::{
 };
 
 use crate::protocol::stats::{
-    CachedTokensBySource, KvConnectorStats, MooncakeStats, MultiConnectorStats, NixlStats,
-    SchedulerStats,
+    CachedTokensBySource, KvConnectorStats, MooncakeStats, MooncakeTransferStats,
+    MultiConnectorStats, NixlStats, SchedulerStats,
 };
 use crate::transport::ConnectedEngine;
 
@@ -89,6 +89,14 @@ struct SchedulerStatsHandles {
     mooncake_operation_bytes_total: MooncakeOperationCounterFamily,
     mooncake_operation_failed_keys_total: MooncakeOperationCounterFamily,
 
+    // Mooncake connector transfer telemetry, decoded from `kv_connector_stats`.
+    mooncake_xfer_time_seconds: HistogramMetric,
+    mooncake_bytes_transferred: HistogramMetric,
+    mooncake_num_descriptors: HistogramMetric,
+    mooncake_num_failed_transfers: U64Counter,
+    mooncake_num_failed_recvs: U64Counter,
+    mooncake_num_kv_expired_reqs: U64Counter,
+
     // NIXL connector telemetry, decoded from `kv_connector_stats`.
     nixl_xfer_time_seconds: HistogramMetric,
     nixl_post_time_seconds: HistogramMetric,
@@ -205,6 +213,16 @@ fn resolve_scheduler_stats_handles(
         mooncake_operation_keys_total: metrics.mooncake_operation_keys_total.clone(),
         mooncake_operation_bytes_total: metrics.mooncake_operation_bytes_total.clone(),
         mooncake_operation_failed_keys_total: metrics.mooncake_operation_failed_keys_total.clone(),
+        mooncake_xfer_time_seconds: metrics.mooncake_xfer_time_seconds.get_or_create_owned(&labels),
+        mooncake_bytes_transferred: metrics.mooncake_bytes_transferred.get_or_create_owned(&labels),
+        mooncake_num_descriptors: metrics.mooncake_num_descriptors.get_or_create_owned(&labels),
+        mooncake_num_failed_transfers: metrics
+            .mooncake_num_failed_transfers
+            .get_or_create_owned(&labels),
+        mooncake_num_failed_recvs: metrics.mooncake_num_failed_recvs.get_or_create_owned(&labels),
+        mooncake_num_kv_expired_reqs: metrics
+            .mooncake_num_kv_expired_reqs
+            .get_or_create_owned(&labels),
         nixl_xfer_time_seconds: metrics.nixl_xfer_time_seconds.get_or_create_owned(&labels),
         nixl_post_time_seconds: metrics.nixl_post_time_seconds.get_or_create_owned(&labels),
         nixl_bytes_transferred: metrics.nixl_bytes_transferred.get_or_create_owned(&labels),
@@ -313,6 +331,9 @@ fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: &
     if let Some(kv_connector_stats) = &stats.kv_connector_stats {
         match kv_connector_stats {
             KvConnectorStats::Nixl(stats) => record_nixl_stats(handles, stats),
+            KvConnectorStats::MooncakeTransfer(stats) => {
+                record_mooncake_transfer_stats(handles, stats)
+            }
             KvConnectorStats::Mooncake(stats) => record_mooncake_stats(handles, stats),
             KvConnectorStats::Multi(stats) => record_multi_connector_stats(handles, stats),
             KvConnectorStats::Other(_) => {}
@@ -324,11 +345,33 @@ fn record_multi_connector_stats(handles: &SchedulerStatsHandles, stats: &MultiCo
     for nixl in [&stats.nixl, &stats.nixl_pull, &stats.nixl_push].into_iter().flatten() {
         record_nixl_stats(handles, nixl);
     }
+    if let Some(mooncake_transfer) = &stats.mooncake_transfer {
+        record_mooncake_transfer_stats(handles, mooncake_transfer);
+    }
     if let Some(mooncake) = &stats.mooncake {
         record_mooncake_stats(handles, mooncake);
     }
 }
 
+fn record_mooncake_transfer_stats(handles: &SchedulerStatsHandles, stats: &MooncakeTransferStats) {
+    for value in &stats.transfer_duration {
+        handles.mooncake_xfer_time_seconds.observe(*value);
+    }
+    for value in &stats.bytes_transferred {
+        handles.mooncake_bytes_transferred.observe(*value as f64);
+    }
+    for value in &stats.num_descriptors {
+        handles.mooncake_num_descriptors.observe(*value as f64);
+    }
+    handles
+        .mooncake_num_failed_transfers
+        .inc_by(stats.num_failed_transfers.iter().sum());
+    handles.mooncake_num_failed_recvs.inc_by(stats.num_failed_recvs.iter().sum());
+    handles
+        .mooncake_num_kv_expired_reqs
+        .inc_by(stats.num_kv_expired_reqs.iter().sum());
+}
+
 fn record_mooncake_stats(handles: &SchedulerStatsHandles, stats: &MooncakeStats) {
     for (operation, records) in &stats.0 {
         for record in records {
@@ -437,7 +480,7 @@ mod tests {
     use crate::metrics::LoraInfoExporter;
     use crate::protocol::stats::{
         KvConnectorStats, MooncakeOperation, MooncakeRecord, MooncakeStats, MooncakeStatus,
-        MultiConnectorStats, NixlStats, SchedulerStats,
+        MooncakeTransferStats, MultiConnectorStats, NixlStats, SchedulerStats,
     };
 
     fn names(values: &[&str]) -> BTreeSet<String> {
@@ -518,6 +561,53 @@ mod tests {
         )]))
     }
 
+    /// Records a bare `MooncakeConnector` payload, in the shape Python's
+    /// `MooncakeKVConnectorStats` sends, into the `vllm:mooncake_*` transfer
+    /// metrics.
+    #[test]
+    fn mooncake_transfer_stats_are_recorded_into_mooncake_metrics() {
+        let metrics = Metrics::new();
+        let handles = super::resolve_scheduler_stats_handles(&metrics.scheduler, "model", 0);
+
+        let kv_connector_stats: KvConnectorStats = serde_json::from_value(serde_json::json!({
+            "transfer_duration": [0.01, 0.02],
+            "bytes_transferred": [4096, 8192],
+            "num_descriptors": [2, 4],
+            "num_failed_transfers": [1, 1],
+            "num_failed_recvs": [1],
+            "num_kv_expired_reqs": [],
+        }))
+        .unwrap();
+        let stats = SchedulerStats {
+            kv_connector_stats: Some(kv_connector_stats),
+            ..Default::default()
+        };
+
+        super::record_scheduler_stats_with_handles(&handles, &stats);
+
+        let rendered = metrics.render().unwrap();
+        let mooncake_transfer_series = rendered
+            .lines()
+            .filter(|line| {
+                line.starts_with("vllm:mooncake_")
+                    && !line.starts_with("vllm:mooncake_store_")
+                    && !line.contains("_bucket{")
+            })
+            .collect::<Vec<_>>()
+            .join("\n");
+        expect![[r#"
+            vllm:mooncake_xfer_time_seconds_sum{model_name="model",engine="0"} 0.03
+            vllm:mooncake_xfer_time_seconds_count{model_name="model",engine="0"} 2
+            vllm:mooncake_bytes_transferred_sum{model_name="model",engine="0"} 12288.0
+            vllm:mooncake_bytes_transferred_count{model_name="model",engine="0"} 2
+            vllm:mooncake_num_descriptors_sum{model_name="model",engine="0"} 6.0
+            vllm:mooncake_num_descriptors_count{model_name="model",engine="0"} 2
+            vllm:mooncake_num_failed_transfers_total{model_name="model",engine="0"} 2
+            vllm:mooncake_num_failed_recvs_total{model_name="model",engine="0"} 1
+            vllm:mooncake_num_kv_expired_reqs_total{model_name="model",engine="0"} 0"#]]
+        .assert_eq(&mooncake_transfer_series);
+    }
+
     /// Records one MultiConnector payload into Mooncake and NIXL metrics.
     #[test]
     fn kv_connector_stats_are_recorded_into_mooncake_and_nixl_metrics() {
@@ -527,6 +617,13 @@ mod tests {
         let stats = SchedulerStats {
             kv_connector_stats: Some(KvConnectorStats::Multi(Box::new(MultiConnectorStats {
                 nixl: Some(nixl_stats()),
+                mooncake_transfer: Some(MooncakeTransferStats {
+                    transfer_duration: vec![0.03],
+                    bytes_transferred: vec![2048],
+                    num_descriptors: vec![1],
+                    num_failed_recvs: vec![1],
+                    ..Default::default()
+                }),
                 mooncake: Some(mooncake_stats()),
                 ..Default::default()
             }))),
@@ -557,5 +654,13 @@ mod tests {
             rendered
                 .contains("vllm:nixl_xfer_time_seconds_count{model_name=\"model\",engine=\"0\"} 2")
         );
+        assert!(rendered.contains(
+            "vllm:mooncake_xfer_time_seconds_count{model_name=\"model\",engine=\"0\"} 1"
+        ));
+        assert!(
+            rendered.contains(
+                "vllm:mooncake_num_failed_recvs_total{model_name=\"model\",engine=\"0\"} 1"
+            )
+        );
     }
 }

rust/src/engine-core-client/src/protocol/stats.rs

@@ -215,6 +215,27 @@ pub mod kv_connector {
         pub num_kv_expired_reqs: Vec<u64>,
     }
 
+    /// Mooncake connector transfer telemetry.
+    ///
+    /// Original Python definition: `MooncakeKVConnectorStats` in
+    /// `vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py`.
+    #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
+    #[serde(deny_unknown_fields)]
+    pub struct MooncakeTransferStats {
+        /// Time spent transferring each successful request.
+        pub transfer_duration: Vec<f64>,
+        /// Bytes transferred by each successful request.
+        pub bytes_transferred: Vec<u64>,
+        /// Descriptor count for each successful request.
+        pub num_descriptors: Vec<u64>,
+        /// Failure counter increments collected since the previous update.
+        pub num_failed_transfers: Vec<u64>,
+        /// Receive failure counter increments collected since the previous update.
+        pub num_failed_recvs: Vec<u64>,
+        /// Expired-request counter increments collected since the previous update.
+        pub num_kv_expired_reqs: Vec<u64>,
+    }
+
     /// Mooncake store operation name.
     #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
     #[serde(rename_all = "snake_case")]
@@ -281,6 +302,8 @@ pub mod kv_connector {
         pub nixl_pull: Option<NixlStats>,
         #[serde(rename = "NixlPushConnector")]
         pub nixl_push: Option<NixlStats>,
+        #[serde(rename = "MooncakeConnector")]
+        pub mooncake_transfer: Option<MooncakeTransferStats>,
         #[serde(rename = "MooncakeStoreConnector")]
         pub mooncake: Option<MooncakeStats>,
         /// Child connector payloads without a Rust telemetry implementation.
@@ -289,10 +312,14 @@ pub mod kv_connector {
     }
 
     /// Connector-specific scheduler telemetry.
+    ///
+    /// Variants are tried in order. `Multi` accepts any map, so every
+    /// connector payload with its own variant must come before it.
     #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
     #[serde(untagged)]
     pub enum KvConnectorStats {
         Nixl(NixlStats),
+        MooncakeTransfer(MooncakeTransferStats),
         Mooncake(MooncakeStats),
         Multi(Box<MultiConnectorStats>),
         Other(BTreeMap<String, OpaqueValue>),

rust/src/engine-core-client/src/tests/client.rs

@@ -2774,6 +2774,8 @@ fn python_msgpack_fixtures_match_rust_encoding() {
         lines.next().expect("missing multipart prompt logprobs fixture line");
     let nixl_stats_hex = lines.next().expect("missing NIXL stats fixture line");
     let mooncake_stats_hex = lines.next().expect("missing Mooncake stats fixture line");
+    let mooncake_transfer_stats_hex =
+        lines.next().expect("missing Mooncake transfer stats fixture line");
     let multi_connector_stats_hex =
         lines.next().expect("missing MultiConnector stats fixture line");
     let ready_response_hex = lines.next().expect("missing ready response fixture line");
@@ -2991,12 +2993,20 @@ fn python_msgpack_fixtures_match_rust_encoding() {
             if stats.0.contains_key(&MooncakeOperation::LoadGet)
     ));
 
+    let mooncake_transfer_stats: KvConnectorStats =
+        rmp_serde::from_slice(&hex::decode(mooncake_transfer_stats_hex).unwrap()).unwrap();
+    assert!(matches!(
+        mooncake_transfer_stats,
+        KvConnectorStats::MooncakeTransfer(stats) if stats.num_failed_transfers == [1]
+    ));
+
     let multi_connector_stats: KvConnectorStats =
         rmp_serde::from_slice(&hex::decode(multi_connector_stats_hex).unwrap()).unwrap();
     assert!(matches!(
         multi_connector_stats,
         KvConnectorStats::Multi(stats)
             if stats.nixl.is_some()
+                && stats.mooncake_transfer.is_some()
                 && stats.mooncake.is_some()
                 && stats.other.contains_key("UnsupportedConnector")
     ));

rust/src/engine-core-client/src/tests/python_compat.py

@@ -543,8 +543,17 @@ class EngineCoreReadyResponse:
         }
     ]
 }
+mooncake_transfer_stats = {
+    "transfer_duration": [0.01, 0.02],
+    "bytes_transferred": [4096, 8192],
+    "num_descriptors": [2, 4],
+    "num_failed_transfers": [1],
+    "num_failed_recvs": [],
+    "num_kv_expired_reqs": [1],
+}
 multi_connector_stats = {
     "NixlConnector": nixl_stats,
+    "MooncakeConnector": mooncake_transfer_stats,
     "MooncakeStoreConnector": mooncake_stats,
     "UnsupportedConnector": {"sample_count": 1},
 }
@@ -570,6 +579,7 @@ class EngineCoreReadyResponse:
 )
 print(msgspec.msgpack.encode(nixl_stats).hex())
 print(msgspec.msgpack.encode(mooncake_stats).hex())
+print(msgspec.msgpack.encode(mooncake_transfer_stats).hex())
 print(msgspec.msgpack.encode(multi_connector_stats).hex())
 print(msgspec.msgpack.encode(ready_response).hex())
 print(extended_outputs_bytes.hex())

rust/src/metrics/src/scheduler.rs

@@ -252,6 +252,14 @@ pub struct SchedulerMetrics {
     pub mooncake_operation_bytes_total: MooncakeOperationCounterFamily,
     pub mooncake_operation_failed_keys_total: MooncakeOperationCounterFamily,
 
+    // Mooncake connector transfer telemetry. Mirrors `MooncakePromMetrics`.
+    pub mooncake_xfer_time_seconds: HistogramFamily,
+    pub mooncake_bytes_transferred: HistogramFamily,
+    pub mooncake_num_descriptors: HistogramFamily,
+    pub mooncake_num_failed_transfers: Family<EngineLabels, U64Counter>,
+    pub mooncake_num_failed_recvs: Family<EngineLabels, U64Counter>,
+    pub mooncake_num_kv_expired_reqs: Family<EngineLabels, U64Counter>,
+
     // NIXL connector telemetry. Mirrors `NixlPromMetrics`.
     pub nixl_xfer_time_seconds: HistogramFamily,
     pub nixl_post_time_seconds: HistogramFamily,
@@ -467,6 +475,54 @@ impl SchedulerMetrics {
             mooncake_operation_failed_keys_total.clone(),
         );
 
+        // Mooncake connector transfer telemetry. Python's `MooncakePromMetrics`
+        // uses the same bucket edges as `NixlPromMetrics`.
+        let mooncake_xfer_time_seconds =
+            Family::new_with_constructor(nixl_xfer_time_histogram as fn() -> Histogram);
+        registry.register(
+            "vllm:mooncake_xfer_time_seconds",
+            "Histogram of transfer duration for Mooncake KV Cache transfers.",
+            mooncake_xfer_time_seconds.clone(),
+        );
+
+        let mooncake_bytes_transferred =
+            Family::new_with_constructor(nixl_bytes_transferred_histogram as fn() -> Histogram);
+        registry.register(
+            "vllm:mooncake_bytes_transferred",
+            "Histogram of bytes transferred per Mooncake KV Cache transfer.",
+            mooncake_bytes_transferred.clone(),
+        );
+
+        let mooncake_num_descriptors =
+            Family::new_with_constructor(nixl_num_descriptors_histogram as fn() -> Histogram);
+        registry.register(
+            "vllm:mooncake_num_descriptors",
+            "Histogram of number of descriptors per Mooncake KV Cache transfer.",
+            mooncake_num_descriptors.clone(),
+        );
+
+        let mooncake_num_failed_transfers = Family::default();
+        registry.register(
+            "vllm:mooncake_num_failed_transfers",
+            "Number of failed Mooncake KV Cache transfers.",
+            mooncake_num_failed_transfers.clone(),
+        );
+
+        let mooncake_num_failed_recvs = Family::default();
+        registry.register(
+            "vllm:mooncake_num_failed_recvs",
+            "Number of failed Mooncake KV Cache receives.",
+            mooncake_num_failed_recvs.clone(),
+        );
+
+        let mooncake_num_kv_expired_reqs = Family::default();
+        registry.register(
+            "vllm:mooncake_num_kv_expired_reqs",
+            "Number of requests that had their KV expire. \
+             NOTE: This metric is tracked on the P instance.",
+            mooncake_num_kv_expired_reqs.clone(),
+        );
+
         // NIXL connector telemetry.
         let nixl_xfer_time_seconds =
             Family::new_with_constructor(nixl_xfer_time_histogram as fn() -> Histogram);
@@ -548,6 +604,12 @@ impl SchedulerMetrics {
             mooncake_operation_keys_total,
             mooncake_operation_bytes_total,
             mooncake_operation_failed_keys_total,
+            mooncake_xfer_time_seconds,
+            mooncake_bytes_transferred,
+            mooncake_num_descriptors,
+            mooncake_num_failed_transfers,
+            mooncake_num_failed_recvs,
+            mooncake_num_kv_expired_reqs,
             nixl_xfer_time_seconds,
             nixl_post_time_seconds,
             nixl_bytes_transferred,