[Metrics] Export scheduler stage wall time (#37636)

Co-authored-by: Pranjal Shankhdhar <pranjal.ssh@gmail.com>

python/sglang/srt/disaggregation/decode.py

@@ -96,6 +96,11 @@ from sglang.srt.observability.req_time_stats import (
     set_schedule_time_batch,
     set_time_batch,
 )
+from sglang.srt.observability.scheduler_stage_metrics import (
+    SCHEDULER_STAGE_GET_NEXT_BATCH,
+    SCHEDULER_STAGE_PROCESS_QUEUE,
+    scheduler_stage_method,
+)
 from sglang.srt.runtime_context import (
     get_disagg,
     get_memory,
@@ -103,7 +108,6 @@ from sglang.srt.runtime_context import (
 )
 from sglang.srt.utils import ceil_align, get_num_new_pages, is_npu
 from sglang.srt.utils.network import NetworkAddress
-from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 
 logger = logging.getLogger(__name__)
@@ -2581,7 +2585,7 @@ class SchedulerDisaggregationDecodeMixin:
 
         return GenerationBatchResult()
 
-    @scheduler_nvtx_method("scheduler.get_next_batch_to_run")
+    @scheduler_stage_method(SCHEDULER_STAGE_GET_NEXT_BATCH)
     def get_next_disagg_decode_batch_to_run(
         self: Scheduler, running_batch: ScheduleBatch
     ) -> NextBatchPlan:
@@ -2689,6 +2693,7 @@ class SchedulerDisaggregationDecodeMixin:
 
         return new_batch
 
+    @scheduler_stage_method(SCHEDULER_STAGE_PROCESS_QUEUE)
     def process_decode_queue(self: Scheduler):
         if self.enable_decode_hicache:
             self.tree_cache.check_hicache_events()

python/sglang/srt/disaggregation/prefill.py

@@ -72,13 +72,18 @@ from sglang.srt.mem_cache.common import (
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
 from sglang.srt.observability.req_time_stats import set_schedule_time_batch
+from sglang.srt.observability.scheduler_stage_metrics import (
+    SCHEDULER_STAGE_GET_NEXT_BATCH,
+    SCHEDULER_STAGE_PROCESS_QUEUE,
+    SchedulerStageMetricsRecorder,
+    scheduler_stage_method,
+)
 from sglang.srt.runtime_context import (
     get_disagg,
     get_parallel,
     get_schedule,
 )
 from sglang.srt.utils import is_npu
-from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
 
 if TYPE_CHECKING:
     from torch.distributed import ProcessGroup
@@ -147,6 +152,7 @@ class PrefillBootstrapQueue:
         gloo_group: ProcessGroup,
         max_total_num_tokens: int,
         scheduler: Scheduler,
+        scheduler_stage_metrics: SchedulerStageMetricsRecorder,
         pp_rank: int,
         pp_size: int,
         transfer_backend: TransferBackend,
@@ -165,6 +171,7 @@ class PrefillBootstrapQueue:
         self.queue: List[Req] = []
         self.gloo_group = gloo_group
         self.scheduler = scheduler
+        self.scheduler_stage_metrics = scheduler_stage_metrics
         self.max_total_num_tokens = (
             self.scheduler.tp_worker.model_runner.effective_max_total_num_tokens
         )
@@ -407,6 +414,7 @@ class PrefillBootstrapQueue:
         """
         req.sampling_params.max_new_tokens = 1
 
+    @scheduler_stage_method(SCHEDULER_STAGE_PROCESS_QUEUE)
     def pop_bootstrapped(
         self,
         return_failed_reqs: bool = False,
@@ -524,6 +532,7 @@ class SchedulerDisaggregationPrefillMixin:
             if room is not None and room in kv_mgr.transfer_infos:
                 prefetch(room)
 
+    @scheduler_stage_method(SCHEDULER_STAGE_PROCESS_QUEUE)
     def resolve_waiting_queue_bootstrap(self: Scheduler) -> None:
         """Resolve bootstrap status for waiting prefill requests before admission.
 
@@ -565,7 +574,7 @@ class SchedulerDisaggregationPrefillMixin:
             for req in self.waiting_queue
         )
 
-    @scheduler_nvtx_method("scheduler.get_next_batch_to_run")
+    @scheduler_stage_method(SCHEDULER_STAGE_GET_NEXT_BATCH)
     def get_next_disagg_prefill_batch_to_run(
         self: Scheduler,
         running_batch: ScheduleBatch,
@@ -875,6 +884,7 @@ class SchedulerDisaggregationPrefillMixin:
             dp_cooperation_info=batch.dp_cooperation_info,
         )
 
+    @scheduler_stage_method(SCHEDULER_STAGE_PROCESS_QUEUE)
     def process_disagg_prefill_inflight_queue(
         self: Scheduler, rids_to_check: Optional[List[str]] = None
     ) -> List[Req]:

python/sglang/srt/managers/scheduler.py

@@ -292,6 +292,15 @@ from sglang.srt.observability.req_time_stats import (
     set_schedule_time_batch,
     set_time_batch,
 )
+from sglang.srt.observability.scheduler_stage_metrics import (
+    SCHEDULER_STAGE_GET_NEXT_BATCH,
+    SCHEDULER_STAGE_IDLE,
+    SCHEDULER_STAGE_PROCESS_BATCH_RESULT,
+    SCHEDULER_STAGE_PROCESS_REQUESTS,
+    SCHEDULER_STAGE_RUN_BATCH,
+    SCHEDULER_STAGE_SANITY_CHECK_CACHE,
+    scheduler_stage_method,
+)
 from sglang.srt.observability.startup_time import build_scheduler_startup_time
 from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
 from sglang.srt.parser.reasoning_parser import ReasoningParser
@@ -339,7 +348,6 @@ from sglang.srt.utils.hf_transformers_utils import (
 )
 from sglang.srt.utils.msgspec_utils import msgspec_to_builtins
 from sglang.srt.utils.numa_utils import get_numa_node_if_available, numa_bind_to_node
-from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 from sglang.srt.utils.weight_versions import (
     compute_weight_version_spans,
@@ -642,6 +650,7 @@ class Scheduler(
         self.init_diffusion_llm()
 
         self.init_metrics_reporter(tp_rank, pp_rank, dp_rank)
+        self.scheduler_stage_metrics = self.metrics_reporter.scheduler_stage_metrics
 
         # Init schedule policy and new token estimation
         self.init_schedule_policy()
@@ -1542,6 +1551,7 @@ class Scheduler(
                 gloo_group=self.attn_tp_cpu_group,
                 max_total_num_tokens=self.max_total_num_tokens,
                 scheduler=self,
+                scheduler_stage_metrics=self.scheduler_stage_metrics,
                 pp_rank=self.ps.pp_rank,
                 pp_size=self.ps.pp_size,
                 transfer_backend=self.transfer_backend,
@@ -2022,7 +2032,7 @@ class Scheduler(
         for prev_batch, prev_result in self.result_queue:
             self.batch_result_processor.advance_grammar_fsm(prev_result, prev_batch)
 
-    @scheduler_nvtx_method("scheduler.process_input_requests")
+    @scheduler_stage_method(SCHEDULER_STAGE_PROCESS_REQUESTS)
     def process_input_requests(self, recv_reqs: List):
         now = time.monotonic()
         self.session_controller.maybe_reap(now)
@@ -2248,6 +2258,7 @@ class Scheduler(
             stream_output=lambda *a, **kw: self.output_streamer.stream_output(*a, **kw),
             get_last_batch=lambda: self.last_batch,
             scripted_scheduler_hook=self.scripted_scheduler_hook,
+            scheduler_stage_metrics=self.scheduler_stage_metrics,
         )
 
     def init_dp_attn_adapter(self) -> None:
@@ -2306,6 +2317,7 @@ class Scheduler(
             get_last_batch=lambda: self.last_batch,
             get_running_batch=lambda: self.running_batch,
             get_chunked_req=lambda: self.chunked_req,
+            scheduler_stage_metrics=self.scheduler_stage_metrics,
         )
 
     def init_rank_consensus_checker(self) -> None:
@@ -3424,7 +3436,7 @@ class Scheduler(
         # todo hisparse, maybe other info to contain for the new batch
         return batch
 
-    @scheduler_nvtx_method("scheduler.get_next_batch_to_run")
+    @scheduler_stage_method(SCHEDULER_STAGE_GET_NEXT_BATCH)
     def get_next_batch_to_run(
         self, running_batch: ScheduleBatch, last_batch: Optional[ScheduleBatch]
     ) -> NextBatchPlan:
@@ -4122,7 +4134,7 @@ class Scheduler(
             else:
                 batch.sampling_info = sched_sampling_info
 
-    @scheduler_nvtx_method("scheduler.run_batch")
+    @scheduler_stage_method(SCHEDULER_STAGE_RUN_BATCH)
     def run_batch(
         self,
         batch: ScheduleBatch,
@@ -4467,7 +4479,7 @@ class Scheduler(
         if batch_result.logits_output is not None:
             batch_result.logits_output.next_token_logits = None
 
-    @scheduler_nvtx_method("scheduler.process_batch_result")
+    @scheduler_stage_method(SCHEDULER_STAGE_PROCESS_BATCH_RESULT)
     def process_batch_result(
         self,
         batch: ScheduleBatch,
@@ -4594,6 +4606,7 @@ class Scheduler(
             if_success = False
         return ClearHiCacheReqOutput(success=if_success)
 
+    @scheduler_stage_method(SCHEDULER_STAGE_IDLE)
     def on_idle(self):
         """Idle housekeeping: guard, check, metrics, reset, sleep."""
         # Flush any health-check signal deferred while the engine was busy.
@@ -4634,6 +4647,7 @@ class Scheduler(
             self.disaggregation_mode == DisaggregationMode.DECODE
             and self.disagg_decode_transfer_queue.has_pending_deferred_releases()
         )
+        with self.scheduler_stage_metrics.record(SCHEDULER_STAGE_SANITY_CHECK_CACHE):
             if not self.enable_hisparse and not deferred_pending:
                 has_leak, messages = self.invariant_checker._check_all_pools(
                     self.pool_stats_observer.get_pool_stats(),
@@ -4643,7 +4657,9 @@ class Scheduler(
                 self.invariant_checker._check_req_pool()
                 # Byte-conservation diagnostic (allocator-owned; static pools
                 # return [] — the token identity above can't see byte leaks).
-            byte_violations = self.token_to_kv_pool_allocator.verify_byte_accounting()
+                byte_violations = (
+                    self.token_to_kv_pool_allocator.verify_byte_accounting()
+                )
                 if byte_violations:
                     self.invariant_checker._report_leak(
                         "pool-bytes", "\n".join(byte_violations)

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -26,6 +26,11 @@ from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
 from sglang.srt.mem_cache.multi_ended_allocator import (
     UnifiedMambaSWATokenToKVPoolAllocator,
 )
+from sglang.srt.observability.scheduler_stage_metrics import (
+    SCHEDULER_STAGE_SANITY_CHECK_CACHE,
+    SchedulerStageMetricsRecorder,
+    scheduler_stage_method,
+)
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils.common import (
     ceil_align,
@@ -58,6 +63,7 @@ class SchedulerInvariantChecker:
     pool_stats_observer: SchedulerPoolStatsObserver
     get_last_batch: Callable
     get_running_batch: Callable
+    scheduler_stage_metrics: SchedulerStageMetricsRecorder
     # The chunked-prefill request parked between chunks is in neither batch;
     # its uncached tokens must still be counted.
     get_chunked_req: Callable = field(default=lambda: None)
@@ -293,6 +299,7 @@ class SchedulerInvariantChecker:
 
         return full_uncached, swa_uncached
 
+    @scheduler_stage_method(SCHEDULER_STAGE_SANITY_CHECK_CACHE)
     def self_check_during_busy(self):
         if self.get_last_batch() is None:
             return

python/sglang/srt/managers/scheduler_components/metrics_reporter.py

@@ -22,6 +22,9 @@ from sglang.srt.observability.metrics_collector import (
     SchedulerStats,
     compute_routing_key_stats,
 )
+from sglang.srt.observability.scheduler_stage_metrics import (
+    SchedulerStageMetricsRecorder,
+)
 from sglang.srt.runtime_context import (
     exports_expert_balancedness_to_prometheus,
     get_context,
@@ -259,6 +262,9 @@ class SchedulerMetricsReporter:
         self._scheduler_time_accounting: Optional[_SchedulerTimeAccountingSnapshot] = (
             None
         )
+        self.scheduler_stage_metrics = SchedulerStageMetricsRecorder(
+            enabled=self.enable_metrics
+        )
 
         self.forward_pass_device_timer: Optional[DeviceTimer] = None
 
@@ -1242,9 +1248,11 @@ class SchedulerMetricsReporter:
     def start_scheduler_time_accounting(self) -> None:
         if not self.enable_metrics:
             return
+        now_wall_ns = time.monotonic_ns()
         self._scheduler_time_accounting = _SchedulerTimeAccountingSnapshot.init(
-            time.monotonic_ns(), time.process_time_ns(), True
+            now_wall_ns, time.process_time_ns(), True
         )
+        self.scheduler_stage_metrics.start(now_wall_ns)
 
     def record_scheduler_active(self) -> None:
         self._record_scheduler_time(is_idle=False)
@@ -1262,6 +1270,7 @@ class SchedulerMetricsReporter:
             self._scheduler_time_accounting = _SchedulerTimeAccountingSnapshot.init(
                 now_wall_ns, time.process_time_ns(), is_idle
             )
+            self.scheduler_stage_metrics.start(now_wall_ns)
             return
 
         accounting.sample(now_wall_ns, is_idle)
@@ -1277,6 +1286,12 @@ class SchedulerMetricsReporter:
         self.metrics_collector.increment_scheduler_process_cpu_seconds(
             elapsed_process_cpu_ns / 1e9
         )
+        for stage, elapsed_wall_ns in self.scheduler_stage_metrics.drain(
+            now_wall_ns
+        ).items():
+            self.metrics_collector.increment_scheduler_stage_seconds(
+                stage=stage, seconds=elapsed_wall_ns / 1e9
+            )
         accounting.reset(now_wall_ns, now_process_cpu_ns, is_idle)
 
     def _reset_device_timer_window(self):

python/sglang/srt/managers/scheduler_components/request_receiver.py

@@ -32,12 +32,16 @@ from sglang.srt.managers.mm_utils import (
     has_shm_features,
     unwrap_shm_features,
 )
+from sglang.srt.observability.scheduler_stage_metrics import (
+    SCHEDULER_STAGE_RECV_REQUESTS,
+    SchedulerStageMetricsRecorder,
+    scheduler_stage_method,
+)
 from sglang.srt.runtime_context import get_disagg, get_parallel, is_ep_scale_joiner
 from sglang.srt.utils import (
     broadcast_pyobj,
     point_to_point_pyobj,
 )
-from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
 
 if TYPE_CHECKING:
     from sglang.srt.configs.model_config import ModelConfig
@@ -73,13 +77,14 @@ class SchedulerRequestReceiver:
     stream_output: Callable[..., None]
     get_last_batch: Callable[[], Any]
     scripted_scheduler_hook: Optional[ScriptedSchedulerHook] = None
+    scheduler_stage_metrics: Optional[SchedulerStageMetricsRecorder] = None
 
     def recv_limit_reached(self, num_recv_reqs: int) -> bool:
         if self.max_recv_per_poll < 0:
             return False
         return num_recv_reqs >= self.max_recv_per_poll
 
-    @scheduler_nvtx_method("scheduler.recv_requests")
+    @scheduler_stage_method(SCHEDULER_STAGE_RECV_REQUESTS)
     def recv_requests(
         self,
     ) -> List[Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput, Any]]:

python/sglang/srt/observability/metrics_collector.py

@@ -25,6 +25,9 @@ from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Set, Union
 
 from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.observability.scheduler_stage_metrics import (
+    SCHEDULER_STAGE_CATEGORIES,
+)
 from sglang.srt.observability.utils import exponential_buckets, generate_buckets
 from sglang.srt.runtime_context import (
     exports_expert_balancedness_to_prometheus,
@@ -930,8 +933,17 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
             ),
             labelnames=labels.keys(),
         )
+        self.scheduler_stage_seconds_total = Counter(
+            name="sglang:scheduler_stage_seconds_total",
+            documentation=(
+                "Total scheduler-loop wall time exclusively attributed to each stage."
+            ),
+            labelnames=list(labels.keys()) + ["category"],
+        )
         self.scheduler_idle_seconds_total.labels(**labels)
         self.scheduler_process_cpu_seconds_total.labels(**labels)
+        for category in SCHEDULER_STAGE_CATEGORIES:
+            self.scheduler_stage_seconds_total.labels(**labels, category=category)
         self.estimated_flops_per_gpu_total = Counter(
             name="sglang:estimated_flops_per_gpu_total",
             documentation=(
@@ -1324,6 +1336,11 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
     def increment_scheduler_process_cpu_seconds(self, t: float) -> None:
         self.scheduler_process_cpu_seconds_total.labels(**self.labels).inc(t)
 
+    def increment_scheduler_stage_seconds(self, stage: str, seconds: float) -> None:
+        self.scheduler_stage_seconds_total.labels(**self.labels, category=stage).inc(
+            seconds
+        )
+
     def increment_estimated_perf(
         self,
         num_flops_per_gpu: float = 0.0,

python/sglang/srt/observability/scheduler_stage_metrics.py

@@ -0,0 +1,163 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+import time
+from collections import defaultdict
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from functools import wraps
+from typing import TypeVar, cast
+
+from sglang.srt.utils.nvtx_utils import (
+    NVTX_SCHEDULER_ENABLED,
+    profile_range,
+    scheduler_nvtx_method,
+)
+
+SCHEDULER_STAGE_OTHER = "other"
+SCHEDULER_STAGE_RECV_REQUESTS = "recv_requests"
+SCHEDULER_STAGE_PROCESS_REQUESTS = "process_input_requests"
+SCHEDULER_STAGE_GET_NEXT_BATCH = "get_next_batch_to_run"
+SCHEDULER_STAGE_PROCESS_QUEUE = "process_queue"
+SCHEDULER_STAGE_RUN_BATCH = "run_batch"
+SCHEDULER_STAGE_PROCESS_BATCH_RESULT = "process_batch_result"
+SCHEDULER_STAGE_SANITY_CHECK_CACHE = "sanity_check_cache"
+SCHEDULER_STAGE_IDLE = "idle"
+
+SCHEDULER_STAGE_CATEGORIES = (
+    SCHEDULER_STAGE_OTHER,
+    SCHEDULER_STAGE_RECV_REQUESTS,
+    SCHEDULER_STAGE_PROCESS_REQUESTS,
+    SCHEDULER_STAGE_GET_NEXT_BATCH,
+    SCHEDULER_STAGE_PROCESS_QUEUE,
+    SCHEDULER_STAGE_RUN_BATCH,
+    SCHEDULER_STAGE_PROCESS_BATCH_RESULT,
+    SCHEDULER_STAGE_SANITY_CHECK_CACHE,
+    SCHEDULER_STAGE_IDLE,
+)
+_SCHEDULER_STAGE_CATEGORY_SET = frozenset(SCHEDULER_STAGE_CATEGORIES)
+
+
+@dataclass(slots=True)
+class SchedulerStageMetricsRecorder:
+    """Accumulate mutually exclusive scheduler wall time by stage.
+
+    Nested stages temporarily replace their parent, and uncategorized time is
+    assigned to ``other``. Summing all categories therefore recovers elapsed
+    scheduler-loop wall time. Active torch profilers receive matching
+    ``scheduler.<stage>`` ranges without requiring Python stacks.
+    """
+
+    enabled: bool
+    _current_stage: str = SCHEDULER_STAGE_OTHER
+    _trace_stage: str | None = None
+    _last_wall_ns: int | None = None
+    _wall_ns: dict[str, int] = field(default_factory=lambda: defaultdict(int))
+
+    def start(self, wall_ns: int) -> None:
+        if not self.enabled:
+            return
+        self._current_stage = SCHEDULER_STAGE_OTHER
+        self._last_wall_ns = wall_ns
+        self._wall_ns.clear()
+
+    def enter(self, stage: str) -> str | None:
+        if (
+            not self.enabled
+            or self._last_wall_ns is None
+            or self._current_stage == stage
+        ):
+            return None
+        self._sample(time.monotonic_ns())
+        previous_stage = self._current_stage
+        self._current_stage = stage
+        return previous_stage
+
+    def exit(self, previous_stage: str | None) -> None:
+        if previous_stage is None:
+            return
+        self._sample(time.monotonic_ns())
+        self._current_stage = previous_stage
+
+    @contextmanager
+    def record(self, stage: str) -> Iterator[None]:
+        if stage not in _SCHEDULER_STAGE_CATEGORY_SET:
+            raise ValueError(f"Unknown scheduler stage: {stage}")
+        previous_stage = self.enter(stage)
+        previous_trace_stage = self._trace_stage
+        trace_stage_changed = previous_trace_stage != stage
+        if trace_stage_changed:
+            self._trace_stage = stage
+        try:
+            if trace_stage_changed:
+                with profile_range(
+                    f"scheduler.{stage}", nvtx_enabled=NVTX_SCHEDULER_ENABLED
+                ):
+                    yield
+            else:
+                yield
+        finally:
+            if trace_stage_changed:
+                self._trace_stage = previous_trace_stage
+            self.exit(previous_stage)
+
+    def drain(self, wall_ns: int) -> dict[str, int]:
+        if not self.enabled or self._last_wall_ns is None:
+            return {}
+        self._sample(wall_ns)
+        wall_ns_by_stage = dict(self._wall_ns)
+        self._wall_ns.clear()
+        return wall_ns_by_stage
+
+    def _sample(self, wall_ns: int) -> None:
+        assert self._last_wall_ns is not None
+        self._wall_ns[self._current_stage] += wall_ns - self._last_wall_ns
+        self._last_wall_ns = wall_ns
+
+
+_F = TypeVar("_F", bound=Callable)
+
+
+def scheduler_stage_method(stage: str) -> Callable[[_F], _F]:
+    if stage not in _SCHEDULER_STAGE_CATEGORY_SET:
+        raise ValueError(f"Unknown scheduler stage: {stage}")
+    trace_name = f"scheduler.{stage}"
+
+    def decorator(func: _F) -> _F:
+        profiled_func = scheduler_nvtx_method(trace_name)(func)
+
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            recorder = self.scheduler_stage_metrics
+            if recorder is None:
+                return profiled_func(self, *args, **kwargs)
+
+            previous_stage = recorder.enter(stage)
+            previous_trace_stage = recorder._trace_stage
+            trace_stage_changed = previous_trace_stage != stage
+            if trace_stage_changed:
+                recorder._trace_stage = stage
+            try:
+                return profiled_func(self, *args, **kwargs)
+            finally:
+                if trace_stage_changed:
+                    recorder._trace_stage = previous_trace_stage
+                recorder.exit(previous_stage)
+
+        return cast(_F, wrapper)
+
+    return decorator

test/registered/observability/test_metrics.py

@@ -195,6 +195,7 @@ class TestEnableMetrics(CustomTestCase):
             "sglang:startup_cuda_graph_time_seconds",
             "sglang:scheduler_idle_seconds_total",
             "sglang:scheduler_process_cpu_seconds_total",
+            "sglang:scheduler_stage_seconds_total",
         ]
         mfu_metrics = [
             "sglang:estimated_flops_per_gpu_total",
@@ -232,6 +233,7 @@ class TestEnableMetrics(CustomTestCase):
             ("sglang:forward_execution_seconds_total", {"category": "extend"}),
             ("sglang:forward_execution_seconds_total", {"category": "decode"}),
             ("sglang:scheduler_process_cpu_seconds_total", {}),
+            ("sglang:scheduler_stage_seconds_total", {"category": "other"}),
             ("sglang:process_cpu_seconds_total", {"component": "tokenizer"}),
             ("sglang:weight_memory_usage_gb", {"model_name": _MODEL_NAME}),
             ("sglang:kv_cache_memory_usage_gb", {"model_name": _MODEL_NAME}),

test/registered/unit/layers/test_minicpm_sparse_cache.py

@@ -182,6 +182,7 @@ def test_reserved_slots_are_excluded_from_full_pool_invariant():
         pool_stats_observer=SimpleNamespace(session_held_tokens=lambda: 0),
         get_last_batch=lambda: None,
         get_running_batch=lambda: None,
+        scheduler_stage_metrics=None,
     )
 
     leak, message = checker._check_full_pool(
@@ -272,6 +273,7 @@ def test_mamba_leak_diagnostic_does_not_report_reserved_slots():
         ),
         get_last_batch=lambda: None,
         get_running_batch=lambda: None,
+        scheduler_stage_metrics=None,
     )
 
     leak, message = checker._check_mamba_pool(

test/registered/unit/managers/scheduler_components/test_invariant_checker.py

@@ -42,6 +42,7 @@ class TestCheckTreeCacheGate(CustomTestCase):
             pool_stats_observer=MagicMock(),
             get_last_batch=lambda: None,
             get_running_batch=lambda: None,
+            scheduler_stage_metrics=None,
         )
 
     def test_disabled_by_default(self):

test/registered/unit/managers/test_generation_auxiliary_output.py

@@ -388,6 +388,7 @@ def test_pdmux_split_prefill_schedules_auxiliary_output_copy():
     )
     copy_done = CopyDone()
     scheduler = object.__new__(Scheduler)
+    scheduler.scheduler_stage_metrics = None
     scheduler.metrics_reporter = Mock()
     scheduler.forward_ct = 0
     scheduler._sched_idled = False

test/registered/unit/managers/test_scheduler_chunked_req_gate.py

@@ -75,6 +75,7 @@ def _make_chunk_cache(req_to_token_pool) -> ChunkCache:
 
 def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler:
     s = Scheduler.__new__(Scheduler)
+    s.scheduler_stage_metrics = None
     s._abort_on_waiting_timeout = MagicMock()
     s._abort_on_running_timeout = MagicMock()
     s.dllm_config = None

test/registered/unit/managers/test_scheduler_on_idle_load.py

@@ -22,6 +22,7 @@ register_cpu_ci(est_time=2, suite="base-a-test-cpu")
 class TestOnIdleStallPublish(CustomTestCase):
     def _stalled_scheduler(self) -> Scheduler:
         s = Scheduler.__new__(Scheduler)
+        s.scheduler_stage_metrics = None
         s.maybe_send_health_check_signal = MagicMock()
         s.is_fully_idle = MagicMock(return_value=False)  # stalled, not idle
         s.publish_load_snapshot = MagicMock(return_value=None)

test/registered/unit/multimodal/test_gpu_feature_transport.py

@@ -767,6 +767,7 @@ class TestSchedulerMmTransportBoundary(unittest.TestCase):
 
     @staticmethod
     def _prepare_scheduler(scheduler):
+        scheduler.scheduler_stage_metrics = None
         scheduler.session_controller = SimpleNamespace(maybe_reap=MagicMock())
         scheduler._request_dispatcher = MagicMock(return_value=None)
         scheduler.flush_wrapper = SimpleNamespace(check_pending=MagicMock())

test/registered/unit/observability/test_forward_pass_metrics.py

@@ -412,10 +412,15 @@ class TestSchedulerTimeAccounting(CustomTestCase):
         self.reporter = _make_reporter(self, types.SimpleNamespace())
         self.idle_seconds = []
         self.process_cpu_seconds = []
+        self.stage_seconds = []
         self.reporter.enable_metrics = True
+        self.reporter.scheduler_stage_metrics.enabled = True
         self.reporter.metrics_collector = types.SimpleNamespace(
             increment_scheduler_idle_seconds=self.idle_seconds.append,
             increment_scheduler_process_cpu_seconds=self.process_cpu_seconds.append,
+            increment_scheduler_stage_seconds=lambda **kwargs: (
+                self.stage_seconds.append(kwargs)
+            ),
         )
 
     def test_counts_idle_wall_time_and_process_cpu_time(self):
@@ -452,6 +457,10 @@ class TestSchedulerTimeAccounting(CustomTestCase):
 
         self.assertAlmostEqual(sum(self.idle_seconds), 2.6)
         self.assertAlmostEqual(sum(self.process_cpu_seconds), 1.9)
+        self.assertAlmostEqual(
+            sum(sample["seconds"] for sample in self.stage_seconds), 4.1
+        )
+        self.assertEqual({sample["stage"] for sample in self.stage_seconds}, {"other"})
 
     def test_state_transitions_accumulate_until_periodic_update(self):
         with (

test/registered/unit/observability/test_scheduler_stage_metrics.py

@@ -0,0 +1,164 @@
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.observability.scheduler_stage_metrics import (
+    SCHEDULER_STAGE_CATEGORIES,
+    SCHEDULER_STAGE_GET_NEXT_BATCH,
+    SCHEDULER_STAGE_PROCESS_QUEUE,
+    SCHEDULER_STAGE_PROCESS_REQUESTS,
+    SCHEDULER_STAGE_RUN_BATCH,
+    SchedulerStageMetricsRecorder,
+    scheduler_stage_method,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestSchedulerStageMetricsRecorder(CustomTestCase):
+    def test_category_names(self):
+        self.assertEqual(
+            set(SCHEDULER_STAGE_CATEGORIES),
+            {
+                "other",
+                "recv_requests",
+                "process_input_requests",
+                "process_batch_result",
+                "process_queue",
+                "get_next_batch_to_run",
+                "run_batch",
+                "sanity_check_cache",
+                "idle",
+            },
+        )
+
+    def test_nested_stages_are_exclusive(self):
+        recorder = SchedulerStageMetricsRecorder(enabled=True)
+        recorder.start(wall_ns=0)
+
+        with patch(
+            "sglang.srt.observability.scheduler_stage_metrics.time.monotonic_ns",
+            side_effect=[10, 30, 50, 80],
+        ):
+            outer = recorder.enter(SCHEDULER_STAGE_GET_NEXT_BATCH)
+            inner = recorder.enter(SCHEDULER_STAGE_PROCESS_QUEUE)
+            recorder.exit(inner)
+            recorder.exit(outer)
+
+        wall_ns = recorder.drain(wall_ns=100)
+
+        self.assertEqual(
+            wall_ns,
+            {
+                "other": 30,
+                "get_next_batch_to_run": 50,
+                "process_queue": 20,
+            },
+        )
+        self.assertEqual(sum(wall_ns.values()), 100)
+
+    def test_decorator_restores_stage_after_exception(self):
+        recorder = SchedulerStageMetricsRecorder(enabled=True)
+        recorder.start(wall_ns=0)
+
+        class SchedulerLike:
+            scheduler_stage_metrics = recorder
+
+            @scheduler_stage_method(SCHEDULER_STAGE_RUN_BATCH)
+            def fail(self):
+                raise RuntimeError("boom")
+
+        with (
+            patch(
+                "sglang.srt.observability.scheduler_stage_metrics.time.monotonic_ns",
+                side_effect=[10, 40],
+            ),
+            self.assertRaisesRegex(RuntimeError, "boom"),
+        ):
+            SchedulerLike().fail()
+
+        wall_ns = recorder.drain(wall_ns=50)
+        self.assertEqual(wall_ns, {"other": 20, "run_batch": 30})
+
+    def test_nested_same_stage_does_not_double_count(self):
+        recorder = SchedulerStageMetricsRecorder(enabled=True)
+        recorder.start(wall_ns=0)
+
+        with patch(
+            "sglang.srt.observability.scheduler_stage_metrics.time.monotonic_ns",
+            side_effect=[10, 40],
+        ):
+            with recorder.record(SCHEDULER_STAGE_RUN_BATCH):
+                with recorder.record(SCHEDULER_STAGE_RUN_BATCH):
+                    pass
+
+        wall_ns = recorder.drain(wall_ns=50)
+        self.assertEqual(wall_ns, {"other": 20, "run_batch": 30})
+
+    def test_trace_spans_do_not_require_python_stacks(self):
+        recorder = SchedulerStageMetricsRecorder(enabled=False)
+
+        class SchedulerLike:
+            scheduler_stage_metrics = recorder
+
+            @scheduler_stage_method(SCHEDULER_STAGE_RUN_BATCH)
+            def run(self):
+                with self.scheduler_stage_metrics.record(SCHEDULER_STAGE_PROCESS_QUEUE):
+                    pass
+
+        with torch.profiler.profile(
+            activities=[torch.profiler.ProfilerActivity.CPU],
+            with_stack=False,
+            acc_events=True,
+        ) as profiler:
+            SchedulerLike().run()
+
+        stage_events = [
+            event for event in profiler.events() if event.key.startswith("scheduler.")
+        ]
+        self.assertEqual(
+            [event.key for event in stage_events],
+            ["scheduler.run_batch", "scheduler.process_queue"],
+        )
+        self.assertTrue(all(not event.stack for event in stage_events))
+
+    def test_decorator_preserves_existing_trace_names(self):
+        recorder = SchedulerStageMetricsRecorder(enabled=False)
+
+        class SchedulerLike:
+            scheduler_stage_metrics = recorder
+
+            @scheduler_stage_method(SCHEDULER_STAGE_PROCESS_REQUESTS)
+            def process_input_requests(self):
+                pass
+
+            @scheduler_stage_method(SCHEDULER_STAGE_GET_NEXT_BATCH)
+            def get_next_batch_to_run(self):
+                pass
+
+        with torch.profiler.profile(
+            activities=[torch.profiler.ProfilerActivity.CPU],
+            with_stack=False,
+            acc_events=True,
+        ) as profiler:
+            scheduler = SchedulerLike()
+            scheduler.process_input_requests()
+            scheduler.get_next_batch_to_run()
+
+        stage_events = [
+            event for event in profiler.events() if event.key.startswith("scheduler.")
+        ]
+        self.assertEqual(
+            [event.key for event in stage_events],
+            [
+                "scheduler.process_input_requests",
+                "scheduler.get_next_batch_to_run",
+            ],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()