From 0f888442c2d0922f6de8a279657698fd1ba33750 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 18 May 2026 18:39:50 +0800 Subject: [PATCH] Stand up SchedulerKvEventsPublisher; migrate KV-event state to it (#25625) --- python/sglang/srt/managers/scheduler.py | 26 ++++++++- .../kv_events_publisher.py | 55 +++++++++++++++++++ .../observability/scheduler_metrics_mixin.py | 47 +++++++--------- 3 files changed, 100 insertions(+), 28 deletions(-) create mode 100644 python/sglang/srt/managers/scheduler_components/kv_events_publisher.py diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a8382d91f..c11394497 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py b/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py new file mode 100644 index 000000000..1f8dda3f9 --- /dev/null +++ b/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) diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index c8758dd36..44aefddb0 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/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._publish_kv_events() + self.emit_kv_metrics(self.kv_events_publisher) + 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._publish_kv_events() + self.emit_kv_metrics(self.kv_events_publisher) + 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.gpu_cache_usage_perc = self.stats.token_usage - kv_metrics.gpu_prefix_cache_hit_rate = self.stats.cache_hit_rate + kv_metrics.num_requests_waiting = self.get_stats().num_queue_reqs.total + kv_metrics.gpu_cache_usage_perc = self.get_stats().token_usage + 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