Stand up SchedulerKvEventsPublisher; migrate KV-event state to it (#25625)

python/sglang/srt/managers/scheduler.py

@@ -170,6 +170,9 @@ from sglang.srt.managers.scheduler_components.dp_attn import (
 from sglang.srt.managers.scheduler_components.invariant_checker import (
     SchedulerInvariantChecker,
 )
+from sglang.srt.managers.scheduler_components.kv_events_publisher import (
+    SchedulerKvEventsPublisher,
+)
 from sglang.srt.managers.scheduler_components.pool_stats_observer import (
     SchedulerPoolStatsObserver,
 )
@@ -491,7 +494,11 @@ class Scheduler(
             tp_cpu_group=self.tp_cpu_group,
             attn_cp_cpu_group=self.attn_cp_cpu_group,
             enable_metrics=self.enable_metrics,
-            enable_kv_cache_events=self.enable_kv_cache_events,
+            enable_kv_cache_events=bool(
+                self.server_args.kv_events_config
+                and self.ps.attn_tp_rank == 0
+                and self.ps.attn_cp_rank == 0
+            ),
             ps=self.ps,
             tp_group=self.tp_group,
             enable_hierarchical_cache=self.enable_hierarchical_cache,
@@ -674,6 +681,20 @@ class Scheduler(
             get_running_batch=lambda: self.running_batch,
         )
 
+        self.kv_events_publisher = SchedulerKvEventsPublisher(
+            kv_events_config=self.server_args.kv_events_config,
+            ps=self.ps,
+            attn_tp_rank=self.ps.attn_tp_rank,
+            attn_cp_rank=self.ps.attn_cp_rank,
+            attn_dp_rank=self.ps.attn_dp_rank,
+            dp_rank=self.ps.dp_rank,
+            tree_cache=self.tree_cache,
+            send_metrics_from_scheduler=self.send_metrics_from_scheduler,
+            max_running_requests=self.max_running_requests,
+            max_total_num_tokens=self.max_total_num_tokens,
+            get_stats=lambda: self.stats,
+        )
+
         self.is_initializing = False
 
     def init_zbal_on_npu(self):
@@ -710,6 +731,7 @@ class Scheduler(
     def init_ipc_channels(self, port_args: PortArgs):
         context = zmq.Context(2)
         self.idle_sleeper = None
+        self.send_metrics_from_scheduler = None
 
         if (
             self.ps.pp_rank == 0
@@ -3118,7 +3140,7 @@ class Scheduler(
         self._maybe_log_idle_metrics()
 
         # kv event publishing
-        self._publish_kv_events()
+        self.publish_kv_events(self.kv_events_publisher)
 
         # reset token ratio
         self.new_token_ratio = self.init_new_token_ratio

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

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Optional,
+)
+
+import zmq
+
+if TYPE_CHECKING:
+    from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+
+
+class SchedulerStats: ...  # type: ignore[no-redef]
+
+
+@dataclasses.dataclass
+class KvMetrics:
+    request_active_slots: int = 0
+    request_total_slots: int = 0
+    kv_active_blocks: int = 0
+    kv_total_blocks: int = 0
+    num_requests_waiting: int = 0
+    gpu_cache_usage_perc: float = 0.0
+    gpu_prefix_cache_hit_rate: float = 0.0
+    data_parallel_rank: int = 0
+
+
+@dataclass(kw_only=True, slots=True)
+class SchedulerKvEventsPublisher:
+    kv_events_config: Optional[str]
+    ps: "ParallelState"
+    attn_tp_rank: int
+    attn_cp_rank: int
+    attn_dp_rank: int
+    dp_rank: Optional[int]
+    tree_cache: "BasePrefixCache"
+    send_metrics_from_scheduler: Optional["zmq.Socket"]
+    max_running_requests: int
+    max_total_num_tokens: int
+    get_stats: Callable
+    enable_kv_cache_events: bool = False
+    kv_event_publisher: Any = None
+
+    def __post_init__(self) -> None:
+        from sglang.srt.observability.scheduler_metrics_mixin import (
+            SchedulerMetricsMixin,
+        )
+
+        SchedulerMetricsMixin.init_kv_events(self, self.kv_events_config)

python/sglang/srt/observability/scheduler_metrics_mixin.py

@@ -20,6 +20,7 @@ from sglang.srt.managers.io_struct import (
     SpeculativeMetrics,
 )
 from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.managers.scheduler_components.kv_events_publisher import KvMetrics
 from sglang.srt.managers.utils import GenerationBatchResult
 from sglang.srt.observability.metrics_collector import (
     DPCooperationInfo,
@@ -35,6 +36,9 @@ if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
     from sglang.srt.managers.schedule_policy import PrefillAdder
     from sglang.srt.managers.scheduler import EmbeddingBatchResult, Scheduler
+    from sglang.srt.managers.scheduler_components.kv_events_publisher import (
+        SchedulerKvEventsPublisher,
+    )
 
 logger = logging.getLogger(__name__)
 
@@ -74,18 +78,6 @@ class PrefillStats:
         )
 
 
-@dataclasses.dataclass
-class KvMetrics:
-    request_active_slots: int = 0
-    request_total_slots: int = 0
-    kv_active_blocks: int = 0
-    kv_total_blocks: int = 0
-    num_requests_waiting: int = 0
-    gpu_cache_usage_perc: float = 0.0
-    gpu_prefix_cache_hit_rate: float = 0.0
-    data_parallel_rank: int = 0
-
-
 class SchedulerMetricsMixin:
     enable_fpm: bool = False
 
@@ -178,8 +170,6 @@ class SchedulerMetricsMixin:
                 reporter=_wrap_execution_reporter,
             )
 
-        self.init_kv_events(self.server_args.kv_events_config)
-
         self._init_fpm()
 
         self.scheduler_status_logger = SchedulerStatusLogger.maybe_create(
@@ -199,7 +189,10 @@ class SchedulerMetricsMixin:
                 for r in getattr(dw, "draft_runner_list", []):
                     r.device_timer = timer
 
-    def init_kv_events(self: Scheduler, kv_events_config: Optional[str]):
+    @staticmethod
+    def init_kv_events(
+        self: "SchedulerKvEventsPublisher", kv_events_config: Optional[str]
+    ):
         self.enable_kv_cache_events = bool(
             kv_events_config and self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0
         )
@@ -609,8 +602,8 @@ class SchedulerMetricsMixin:
             self.update_lora_metrics()
             self._log_hicache_stats()
             self.metrics_collector.log_stats(self.stats)
-            self._emit_kv_metrics()
+            self.emit_kv_metrics(self.kv_events_publisher)
-        self._publish_kv_events()
+        self.publish_kv_events(self.kv_events_publisher)
 
     def report_decode_stats(
         self: Scheduler,
@@ -805,8 +798,8 @@ class SchedulerMetricsMixin:
             self.update_lora_metrics()
             self._log_hicache_stats()
             self.metrics_collector.log_stats(self.stats)
-            self._emit_kv_metrics()
+            self.emit_kv_metrics(self.kv_events_publisher)
-        self._publish_kv_events()
+        self.publish_kv_events(self.kv_events_publisher)
 
     def log_batch_result_stats(
         self: Scheduler,
@@ -824,20 +817,21 @@ class SchedulerMetricsMixin:
                 balancedness=m.eplb_balancedness.item(),
             )
 
-    def _emit_kv_metrics(self: Scheduler):
+    @staticmethod
+    def emit_kv_metrics(self: "SchedulerKvEventsPublisher"):
         if not self.enable_kv_cache_events:
             return
 
         kv_metrics = KvMetrics()
-        kv_metrics.request_active_slots = self.stats.num_running_reqs.total
+        kv_metrics.request_active_slots = self.get_stats().num_running_reqs.total
         kv_metrics.request_total_slots = self.max_running_requests
         kv_metrics.kv_active_blocks = int(
-            self.stats.token_usage * self.max_total_num_tokens
+            self.get_stats().token_usage * self.max_total_num_tokens
         )
         kv_metrics.kv_total_blocks = self.max_total_num_tokens
-        kv_metrics.num_requests_waiting = self.stats.num_queue_reqs.total
+        kv_metrics.num_requests_waiting = self.get_stats().num_queue_reqs.total
-        kv_metrics.gpu_cache_usage_perc = self.stats.token_usage
+        kv_metrics.gpu_cache_usage_perc = self.get_stats().token_usage
-        kv_metrics.gpu_prefix_cache_hit_rate = self.stats.cache_hit_rate
+        kv_metrics.gpu_prefix_cache_hit_rate = self.get_stats().cache_hit_rate
         kv_metrics.data_parallel_rank = (
             self.ps.dp_rank if self.ps.dp_rank is not None else 0
         )
@@ -845,7 +839,8 @@ class SchedulerMetricsMixin:
         if not self.send_metrics_from_scheduler.closed:
             self.send_metrics_from_scheduler.send_pyobj(kv_metrics)
 
-    def _publish_kv_events(self: Scheduler):
+    @staticmethod
+    def publish_kv_events(self: "SchedulerKvEventsPublisher"):
         if not self.enable_kv_cache_events:
             return