Add kv_transfer_total_mb to metrics (#15667)

python/sglang/srt/managers/scheduler_metrics_mixin.py

@@ -64,6 +64,7 @@ class SchedulerMetricsMixin:
         self.kv_transfer_latency_ms: float = 0.0
         self.kv_transfer_bootstrap_ms: float = 0.0
         self.kv_transfer_alloc_ms: float = 0.0
+        self.kv_transfer_total_mb: float = 0.0
 
         self.stats = SchedulerStats()
 
@@ -221,6 +222,7 @@ class SchedulerMetricsMixin:
                 self.stats.kv_transfer_latency_ms = self.kv_transfer_latency_ms
                 self.stats.kv_transfer_bootstrap_ms = self.kv_transfer_bootstrap_ms
                 self.stats.kv_transfer_alloc_ms = self.kv_transfer_alloc_ms
+                self.stats.kv_transfer_total_mb = self.kv_transfer_total_mb
             elif self.disaggregation_mode == DisaggregationMode.DECODE:
                 self.stats.num_decode_prealloc_queue_reqs = len(
                     self.disagg_decode_prealloc_queue.queue

python/sglang/srt/metrics/collector.py

@@ -68,6 +68,8 @@ class TimeStats:
     alloc_waiting_duration: float = 0.0
     prefill_start_time_host: float = 0.0
     prefill_end_time_host: float = 0.0
+    transfer_speed_gb_s: float = 0.0
+    transfer_total_mb: float = 0.0
 
     # Timestamp when prefill phase finishes, obtained from `time.time()`.
     # Note that this differs from the other `_time` fields tracked by the
@@ -132,7 +134,9 @@ class TimeStats:
                 f"+ other({self.format_duration(other)}); "
                 f"queue_duration={self.format_duration(queue_duration)}, "
                 f"forward_duration={self.format_duration(forward_duration)}, "
-                f"start={self.prefill_bootstrap_queue_entry_time:.3f}"
+                f"start={self.prefill_bootstrap_queue_entry_time:.3f}, "
+                f"transfer_speed={self.transfer_speed_gb_s:.2f}GB/s, "
+                f"transfer_total={self.transfer_total_mb:.2f}MB"
             )
         elif self.disagg_mode == DisaggregationMode.DECODE:
             prealloc_duration = (
@@ -221,6 +225,7 @@ class SchedulerStats:
     kv_transfer_latency_ms: float = 0.0
     kv_transfer_bootstrap_ms: float = 0.0
     kv_transfer_alloc_ms: float = 0.0
+    kv_transfer_total_mb: float = 0.0
 
     # Utilization
     utilization: float = 0.0
@@ -419,6 +424,12 @@ class SchedulerMetricsCollector:
             labelnames=labels.keys(),
             multiprocess_mode="mostrecent",
         )
+        self.kv_transfer_total_mb = Gauge(
+            name="sglang:kv_transfer_total_mb",
+            documentation="The total number of tokens transferred in the KV cache.",
+            labelnames=labels.keys(),
+            multiprocess_mode="mostrecent",
+        )
 
         # Utilization
         self.utilization = Gauge(
@@ -771,6 +782,7 @@ class SchedulerMetricsCollector:
         self._log_gauge(self.kv_transfer_latency_ms, stats.kv_transfer_latency_ms)
         self._log_gauge(self.kv_transfer_bootstrap_ms, stats.kv_transfer_bootstrap_ms)
         self._log_gauge(self.kv_transfer_alloc_ms, stats.kv_transfer_alloc_ms)
+        self._log_gauge(self.kv_transfer_total_mb, stats.kv_transfer_total_mb)
 
         # Retract
         self._log_gauge(self.num_retracted_reqs, stats.num_retracted_reqs)

python/sglang/srt/utils/common.py

@@ -3443,7 +3443,7 @@ def check_cuda_result(raw_output):
     return results
 
 
-def get_physical_device_id(pytorch_device_id: int) -> int:
+def get_physical_device_id() -> int:
     """
     Convert PyTorch logical device ID to physical device ID.
     """