Add metrics for having prefill and decode in different ranks (#15752)

python/sglang/srt/environ.py

@@ -386,6 +386,7 @@ class Envs:
 
     # Metrics
     SGLANG_ENABLE_METRICS_DEVICE_TIMER = EnvBool(False)
+    SGLANG_ENABLE_METRICS_DP_ATTENTION = EnvBool(False)
 
     # fmt: on
 

python/sglang/srt/managers/schedule_batch.py

@@ -74,7 +74,11 @@ from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
 from sglang.srt.mem_cache.radix_cache import RadixKey
 from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
-from sglang.srt.metrics.collector import SchedulerMetricsCollector, TimeStats
+from sglang.srt.metrics.collector import (
+    DPCooperationInfo,
+    SchedulerMetricsCollector,
+    TimeStats,
+)
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
     ForwardBatch,
@@ -1249,6 +1253,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     # Diffusion LLM
     dllm_config: Optional[DllmConfig] = None
 
+    # Metrics
+    dp_cooperation_info: Optional[DPCooperationInfo] = None
+
     @classmethod
     def init_new(
         cls,
@@ -2161,6 +2168,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             mamba_track_indices=self.mamba_track_indices,
             mamba_track_mask=self.mamba_track_mask,
             mamba_track_seqlens=self.mamba_track_seqlens,
+            dp_cooperation_info=self.dp_cooperation_info,
         )
 
     def _is_available_size_sufficient(self, num_tokens: int) -> bool:

python/sglang/srt/managers/scheduler.py

@@ -716,6 +716,7 @@ class Scheduler(
         self.last_batch: Optional[ScheduleBatch] = None
         self.forward_ct = 0
         self.last_prefill_tokens = 0
+        self.last_prefill_cache_tokens = 0
         self.return_health_check_ct = 0
         self.num_retracted_reqs: int = 0
         self.num_paused_reqs: int = 0
@@ -1830,6 +1831,8 @@ class Scheduler(
         if ret:
             trace_event_batch("schedule", ret.reqs)
 
+        self.log_prefill_stats_late(ret)
+
         return ret
 
     def get_num_allocatable_reqs(self, running_bs):

python/sglang/srt/managers/scheduler_dp_attn_mixin.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Callable, Optional
 
 import torch
 
 from sglang.srt.batch_overlap.two_batch_overlap import TboDPAttentionPreparer
 from sglang.srt.environ import envs
 from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.metrics.collector import DPCooperationInfo
 from sglang.srt.utils.common import require_mlp_tp_gather
 
 if TYPE_CHECKING:
@@ -15,6 +16,9 @@ if TYPE_CHECKING:
     from sglang.srt.managers.scheduler import Scheduler
 
 
+_ENABLE_METRICS_DP_ATTENTION = envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.get()
+
+
 @dataclass
 class MLPSyncBatchInfo:
     dp_size: int
@@ -33,6 +37,7 @@ class MLPSyncBatchInfo:
     global_num_tokens_for_logprob: list[int] = None
     tbo_split_seq_index: torch.Tensor = None
     global_forward_mode: int = None
+    dp_cooperation_info: Optional[DPCooperationInfo] = None
 
     def _get_local_tensor(self, device, dtype=torch.int64) -> torch.Tensor:
         return torch.tensor(
@@ -68,6 +73,8 @@ class MLPSyncBatchInfo:
         self.global_num_tokens_for_logprob = tp0_info[:, 1].tolist()
         self.can_cuda_graph = bool(tp0_info[:, 2].min().item())
         self.is_extend_in_batch = bool(tp0_info[:, 3].max().item())
+        if _ENABLE_METRICS_DP_ATTENTION:
+            self.dp_cooperation_info = DPCooperationInfo.create(tp0_info[:, 5].tolist())
 
 
 def _update_gather_batch(
@@ -180,6 +187,9 @@ def prepare_mlp_sync_batch_raw(
             batch_to_gather, mlp_sync_info, require_mlp_tp_gather, skip_all_gather
         )
 
+    if _ENABLE_METRICS_DP_ATTENTION and local_batch is not None:
+        local_batch.dp_cooperation_info = mlp_sync_info.dp_cooperation_info
+
     return local_batch
 
 

python/sglang/srt/managers/scheduler_metrics_mixin.py

@@ -88,7 +88,7 @@ class SchedulerMetricsMixin:
 
             if ENABLE_METRICS_DEVICE_TIMER:
                 self.forward_pass_device_timer = DeviceTimer(
-                    reporter=self.metrics_collector.increment_gpu_execution_seconds
+                    reporter=self.metrics_collector.increment_gpu_execution_seconds,
                 )
 
         if self.enable_kv_cache_events:
@@ -124,6 +124,7 @@ class SchedulerMetricsMixin:
         self.last_prefill_stats_tic = time.perf_counter()
         self.last_input_throughput = self.last_prefill_tokens / gap_latency
         self.last_prefill_tokens = adder.log_input_tokens
+        self.last_prefill_cache_tokens = adder.log_hit_tokens
 
         # TODO: generalize this for various memory pools
         if self.is_hybrid_swa:
@@ -231,23 +232,26 @@ class SchedulerMetricsMixin:
                     self.disagg_decode_transfer_queue.queue
                 )
 
-            self.metrics_collector.increment_realtime_tokens(
-                prefill_compute_tokens=adder.log_input_tokens,
-                prefill_cache_tokens=adder.log_hit_tokens,
-            )
-
             # Others
             self.calculate_utilization()
             self.metrics_collector.log_stats(self.stats)
             self._emit_kv_metrics()
         self._publish_kv_events()
 
+    def log_prefill_stats_late(self: Scheduler, batch: Optional[ScheduleBatch]):
+        """This should be called after `batch` has gathered enough metadata."""
+        if self.enable_metrics and batch is not None:
+            self.metrics_collector.increment_realtime_tokens(
+                prefill_compute_tokens=self.last_prefill_tokens,
+                prefill_cache_tokens=self.last_prefill_cache_tokens,
+                dp_cooperation_info=batch.dp_cooperation_info,
+            )
+
     def log_decode_stats(
         self: Scheduler, can_run_cuda_graph: bool, running_batch: ScheduleBatch = None
     ):
         batch = running_batch or self.running_batch
 
-        last_num_generated_tokens = self.num_generated_tokens
         gap_latency = time.perf_counter() - self.last_decode_stats_tic
         self.last_decode_stats_tic = time.perf_counter()
         self.last_gen_throughput = self.num_generated_tokens / gap_latency
@@ -388,16 +392,21 @@ class SchedulerMetricsMixin:
                     self.disagg_decode_transfer_queue.queue
                 )
 
-            self.metrics_collector.increment_realtime_tokens(
-                decode_tokens=last_num_generated_tokens
-            )
-
             # Others
             self.calculate_utilization()
             self.metrics_collector.log_stats(self.stats)
             self._emit_kv_metrics()
         self._publish_kv_events()
 
+    def log_decode_stats_every_iteration(
+        self: Scheduler, batch: ScheduleBatch, num_accepted_tokens: int
+    ):
+        self.metrics_collector.increment_realtime_tokens(
+            # TODO unify this w/ the bumping logic in `Scheduler.num_generated_tokens` accumulator
+            decode_tokens=batch.batch_size() + num_accepted_tokens,
+            dp_cooperation_info=batch.dp_cooperation_info,
+        )
+
     def log_batch_result_stats(
         self: Scheduler,
         batch: ScheduleBatch,
@@ -491,5 +500,10 @@ class SchedulerMetricsMixin:
             return
 
         category = "forward_" + batch.forward_mode.name.lower()
-        with self.forward_pass_device_timer.wrap(category=category):
+        with self.forward_pass_device_timer.wrap(
+            metadata=dict(
+                category=category,
+                dp_cooperation_info=batch.dp_cooperation_info,
+            ),
+        ):
             yield

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -445,6 +445,10 @@ class SchedulerOutputProcessorMixin:
             and self.forward_ct_decode % self.server_args.decode_log_interval == 0
         ):
             self.log_decode_stats(can_run_cuda_graph, running_batch=batch)
+        if self.enable_metrics:
+            self.log_decode_stats_every_iteration(
+                batch, num_accepted_tokens=result.num_accepted_tokens
+            )
 
     def _mamba_prefix_cache_update(
         self, req: Req, batch: ScheduleBatch, result: GenerationBatchResult, i: int

python/sglang/srt/metrics/collector.py

@@ -12,6 +12,7 @@
 # limitations under the License.
 # ==============================================================================
 """Utilities for Prometheus Metrics Collection."""
+import dataclasses
 import logging
 import os
 import time
@@ -21,6 +22,7 @@ from typing import Dict, List, Optional, Union
 from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.environ import envs
 from sglang.srt.metrics.utils import exponential_buckets, generate_buckets
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import get_bool_env_var
 
@@ -240,6 +242,24 @@ class SchedulerStats:
     is_cuda_graph: float = 0.0
 
 
+@dataclass
+class DPCooperationInfo:
+    # Users can derive that, except for cases with idle, num_decode_ranks=world_size-num_prefill_ranks
+    # We do not provide `num_decode_ranks` to avoid cardinality explosion.
+    num_prefill_ranks: int
+
+    @staticmethod
+    def create(forward_modes: List[int]):
+        return DPCooperationInfo(
+            num_prefill_ranks=sum(
+                1 for mode in forward_modes if mode == ForwardMode.EXTEND.value
+            ),
+        )
+
+    def to_labels(self):
+        return dataclasses.asdict(self)
+
+
 class SchedulerMetricsCollector:
 
     def __init__(
@@ -680,6 +700,17 @@ class SchedulerMetricsCollector:
             labelnames=list(labels.keys()) + ["category"],
         )
 
+        self.dp_cooperation_realtime_tokens_total = Counter(
+            name="sglang:dp_cooperation_realtime_tokens_total",
+            documentation="Total number of tokens processed with labels about DP cooperation.",
+            labelnames=list(labels.keys()) + ["mode", "num_prefill_ranks"],
+        )
+        self.dp_cooperation_gpu_execution_seconds_total = Counter(
+            name="sglang:dp_cooperation_gpu_execution_seconds_total",
+            documentation="Total time that GPU is busy executing a workload with labels about DP cooperation.",
+            labelnames=list(labels.keys()) + ["category", "num_prefill_ranks"],
+        )
+
     def _log_gauge(self, gauge, data: Union[int, float]) -> None:
         # Convenience function for logging to gauge.
         gauge.labels(**self.labels).set(data)
@@ -716,7 +747,11 @@ class SchedulerMetricsCollector:
         )
 
     def increment_realtime_tokens(
-        self, prefill_compute_tokens=0, prefill_cache_tokens=0, decode_tokens=0
+        self,
+        dp_cooperation_info: Optional[DPCooperationInfo],
+        prefill_compute_tokens=0,
+        prefill_cache_tokens=0,
+        decode_tokens=0,
     ):
         for mode, delta in [
             ("prefill_compute", prefill_compute_tokens),
@@ -724,10 +759,27 @@ class SchedulerMetricsCollector:
             ("decode", decode_tokens),
         ]:
             self.realtime_tokens_total.labels(**self.labels, mode=mode).inc(delta)
+            if dp_cooperation_info is not None:
+                self.dp_cooperation_realtime_tokens_total.labels(
+                    **self.labels,
+                    mode=mode,
+                    **dp_cooperation_info.to_labels(),
+                ).inc(delta)
 
-    def increment_gpu_execution_seconds(self, category: str, t: float):
+    def increment_gpu_execution_seconds(
+        self,
+        category: str,
+        t: float,
+        dp_cooperation_info: Optional[DPCooperationInfo],
+    ):
         logger.debug(f"GPU execution seconds: {category=} {t=:.3f}")
         self.gpu_execution_seconds_total.labels(**self.labels, category=category).inc(t)
+        if dp_cooperation_info is not None:
+            self.dp_cooperation_gpu_execution_seconds_total.labels(
+                **self.labels,
+                category=category,
+                **dp_cooperation_info.to_labels(),
+            ).inc(t)
 
     def log_stats(self, stats: SchedulerStats) -> None:
         self._log_gauge(self.num_running_reqs, stats.num_running_reqs)

python/sglang/srt/utils/device_timer.py

@@ -1,23 +1,23 @@
 from collections import deque
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Callable, Deque, Optional
+from typing import Callable, Deque, Dict, Optional
 
 import torch
 
 
 class DeviceTimer:
-    def __init__(self, reporter: Callable[[str, float], None]):
+    def __init__(self, reporter: Callable):
         self._intervals: Deque[_TimingInterval] = deque()
         self._reporter = reporter
 
     @contextmanager
-    def wrap(self, category: str):
+    def wrap(self, metadata: Dict):
         self._intervals.append(_TimingInterval.create())
         try:
             yield
         finally:
-            self._intervals[-1].end(category=category)
+            self._intervals[-1].end(metadata=metadata)
             self._report()
 
     def _report(self):
@@ -27,14 +27,14 @@ class DeviceTimer:
                 break
 
             self._intervals.popleft()
-            self._reporter(interval.category, interval.elapsed_time() / 1000.0)
+            self._reporter(t=interval.elapsed_time() / 1000.0, **interval.metadata)
 
 
 @dataclass
 class _TimingInterval:
     start_event: torch.cuda.Event
     end_event: Optional[torch.cuda.Event] = None
-    category: Optional[str] = None
+    metadata: Optional[Dict] = None
 
     @staticmethod
     def create():
@@ -42,13 +42,13 @@ class _TimingInterval:
         start_event.record()
         return _TimingInterval(start_event=start_event)
 
-    def end(self, category: str):
+    def end(self, metadata: Dict):
         end_event = torch.cuda.Event(enable_timing=True)
         end_event.record()
 
         assert self.end_event is None
         self.end_event = end_event
-        self.category = category
+        self.metadata = metadata
 
     def elapsed_time(self) -> float:
         return self.start_event.elapsed_time(self.end_event)

test/srt/test_metrics.py

@@ -1,25 +1,77 @@
 import unittest
+from typing import Dict, List
 
 import requests
+from prometheus_client.parser import text_string_to_metric_families
+from prometheus_client.samples import Sample
 
+from sglang.srt.environ import envs
 from sglang.srt.utils import kill_process_tree
 from sglang.test.test_utils import (
-    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
+    is_in_ci,
     popen_launch_server,
 )
 
+_MODEL_NAME = "Qwen/Qwen3-0.6B"
+
 
 class TestEnableMetrics(CustomTestCase):
-    def test_metrics_enabled(self):
+    def test_metrics_1gpu(self):
         """Test that metrics endpoint returns data when enabled"""
+        self._execute_core(
+            other_args=[],
+            verify_metrics_extra=None,
+        )
+
+    def test_metrics_2gpu(self):
+        # TODO enable when we have 2-gpu runner in nightly CI
+        if is_in_ci():
+            print("Skip test_metrics_2gpu since in 1-gpu CI")
+            return
+
+        def _verify_metrics_extra(metrics):
+            metrics_to_check = [
+                (
+                    "sglang:dp_cooperation_realtime_tokens_total",
+                    {"mode": "prefill_compute"},
+                ),
+                ("sglang:dp_cooperation_realtime_tokens_total", {"mode": "decode"}),
+                (
+                    "sglang:dp_cooperation_gpu_execution_seconds_total",
+                    {"category": "forward_prefill"},
+                ),
+                (
+                    "sglang:dp_cooperation_gpu_execution_seconds_total",
+                    {"category": "forward_decode"},
+                ),
+            ]
+            _check_metrics_positive(self, metrics, metrics_to_check)
+
+            num_prefill_ranks_values = {
+                s.labels["num_prefill_ranks"]
+                for s in metrics["sglang:dp_cooperation_realtime_tokens_total"]
+            }
+            self.assertIn("0", num_prefill_ranks_values)
+            self.assertIn("1", num_prefill_ranks_values)
+
+        self._execute_core(
+            other_args=["--tp", "2", "--dp", "2", "--enable-dp-attention"],
+            verify_metrics_extra=_verify_metrics_extra,
+        )
+
+    def _execute_core(self, other_args, verify_metrics_extra):
+        with (
+            envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.override(True),
+            envs.SGLANG_ENABLE_METRICS_DEVICE_TIMER.override(True),
+        ):
             process = popen_launch_server(
-            DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+                _MODEL_NAME,
                 DEFAULT_URL_FOR_TEST,
                 timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=["--enable-metrics", "--cuda-graph-max-bs", 2],
+                other_args=["--enable-metrics", "--cuda-graph-max-bs", 2, *other_args],
             )
 
         try:
@@ -45,11 +97,18 @@ class TestEnableMetrics(CustomTestCase):
             # Get metrics
             metrics_response = requests.get(f"{DEFAULT_URL_FOR_TEST}/metrics")
             self.assertEqual(metrics_response.status_code, 200)
-            metrics_content = metrics_response.text
+            metrics_text = metrics_response.text
 
-            print(f"metrics_content=\n{metrics_content}")
+            print(f"metrics_text=\n{metrics_text}")
 
-            # Verify essential metrics are present
+            metrics = _parse_prometheus_metrics(metrics_text)
+            self._verify_metrics_common(metrics_text, metrics)
+            if verify_metrics_extra is not None:
+                verify_metrics_extra(metrics)
+        finally:
+            kill_process_tree(process.pid)
+
+    def _verify_metrics_common(self, metrics_text, metrics):
         essential_metrics = [
             "sglang:num_running_reqs",
             "sglang:num_used_tokens",
@@ -67,21 +126,44 @@ class TestEnableMetrics(CustomTestCase):
             "sglang:inter_token_latency_seconds",
             "sglang:e2e_request_latency_seconds",
         ]
-
         for metric in essential_metrics:
-                self.assertIn(metric, metrics_content, f"Missing metric: {metric}")
+            self.assertIn(metric, metrics_text, f"Missing metric: {metric}")
 
-            # Verify model name label is present and correct
+        self.assertIn(f'model_name="{_MODEL_NAME}"', metrics_text)
-            expected_model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        self.assertIn("_sum{", metrics_text)
-            self.assertIn(f'model_name="{expected_model_name}"', metrics_content)
+        self.assertIn("_count{", metrics_text)
+        self.assertIn("_bucket{", metrics_text)
 
-            # Verify metrics have values (not empty)
+        metrics_to_check = [
-            self.assertIn("_sum{", metrics_content)
+            ("sglang:realtime_tokens_total", {"mode": "prefill_compute"}),
-            self.assertIn("_count{", metrics_content)
+            ("sglang:realtime_tokens_total", {"mode": "decode"}),
-            self.assertIn("_bucket{", metrics_content)
+            ("sglang:gpu_execution_seconds_total", {"category": "forward_extend"}),
+            ("sglang:gpu_execution_seconds_total", {"category": "forward_decode"}),
+        ]
+        _check_metrics_positive(self, metrics, metrics_to_check)
 
-        finally:
+
-            kill_process_tree(process.pid)
+def _parse_prometheus_metrics(metrics_text: str) -> Dict[str, List[Sample]]:
+    result = {}
+    for family in text_string_to_metric_families(metrics_text):
+        for sample in family.samples:
+            if sample.name not in result:
+                result[sample.name] = []
+            result[sample.name].append(sample)
+    return result
+
+
+def _get_sample_value_by_labels(samples: List[Sample], labels: Dict[str, str]) -> float:
+    for sample in samples:
+        if all(sample.labels.get(k) == v for k, v in labels.items()):
+            return sample.value
+    raise KeyError(f"No sample found with labels {labels}")
+
+
+def _check_metrics_positive(test_case, metrics, metrics_to_check):
+    for metric_name, labels in metrics_to_check:
+        value = _get_sample_value_by_labels(metrics[metric_name], labels)
+        test_case.assertGreater(value, 0, f"{metric_name} {labels}")
 
 
 if __name__ == "__main__":