diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 15e1ff6d9..2892ba071 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -167,6 +167,9 @@ from sglang.srt.managers.schedule_policy import ( from sglang.srt.managers.scheduler_components.dp_attn import ( SchedulerDPAttnAdapter, ) +from sglang.srt.managers.scheduler_components.pool_stats_observer import ( + SchedulerPoolStatsObserver, +) from sglang.srt.managers.scheduler_components.profiler_manager import ( SchedulerProfilerManager, ) @@ -611,6 +614,22 @@ class Scheduler( get_require_mlp_sync=lambda: self.require_mlp_sync, ) + self.pool_stats_observer = SchedulerPoolStatsObserver( + tree_cache=self.tree_cache, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + req_to_token_pool=self.req_to_token_pool, + session_controller=self.session_controller, + hisparse_coordinator=self.hisparse_coordinator, + is_hybrid_swa=self.is_hybrid_swa, + is_hybrid_ssm=self.is_hybrid_ssm, + enable_hisparse=self.enable_hisparse, + full_tokens_per_layer=self.full_tokens_per_layer, + swa_tokens_per_layer=self.swa_tokens_per_layer, + max_total_num_tokens=self.max_total_num_tokens, + get_last_batch=lambda: self.last_batch, + get_running_batch=lambda: self.running_batch, + ) + self.is_initializing = False def init_zbal_on_npu(self): @@ -2324,7 +2343,9 @@ class Scheduler( prefill_delayer_single_pass = None if self.prefill_delayer: # Get max usage across all pools for prefill delay decision - max_pool_usage = self.get_pool_stats().get_max_pool_usage() + max_pool_usage = self.get_pool_stats( + self.pool_stats_observer, + ).get_max_pool_usage() prefill_delayer_single_pass = PrefillDelayerSinglePassExecutor( self.prefill_delayer, token_usage=max_pool_usage ) @@ -3039,7 +3060,11 @@ class Scheduler( # memory leak check (skipped for hisparse — pool counters intentionally # diverge during host-backup, see _get_swa_token_info clamp). if not self.enable_hisparse: - has_leak, messages = self._check_all_pools(self.get_pool_stats()) + has_leak, messages = self._check_all_pools( + self.get_pool_stats( + self.pool_stats_observer, + ) + ) if has_leak: self._report_leak("pool", "\n".join(messages)) self._check_req_pool() diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 331e22a47..986d473da 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -1,7 +1,20 @@ from __future__ import annotations import dataclasses -from typing import List, Optional, Tuple +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Any, + Callable, + List, + Optional, + Tuple, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator + from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache + from sglang.srt.mem_cache.memory_pool import ReqToTokenPool class SchedulerStats: ... # type: ignore[no-redef] @@ -123,3 +136,20 @@ class PoolStats: stats.kv_available_tokens = self.full_available_size stats.kv_evictable_tokens = self.full_evictable_size stats.kv_used_tokens = self.full_num_used + + +@dataclass(kw_only=True, slots=True, frozen=True) +class SchedulerPoolStatsObserver: + tree_cache: "BasePrefixCache" + token_to_kv_pool_allocator: "BaseTokenToKVPoolAllocator" + req_to_token_pool: "ReqToTokenPool" + session_controller: Any + hisparse_coordinator: Any + is_hybrid_swa: bool + is_hybrid_ssm: bool + enable_hisparse: bool + full_tokens_per_layer: Any + swa_tokens_per_layer: Any + max_total_num_tokens: int + get_last_batch: Callable + get_running_batch: Callable diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 419bca683..f343318e7 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -15,26 +15,31 @@ from sglang.srt.utils.watchdog import WatchdogRaw if TYPE_CHECKING: from sglang.srt.managers.scheduler import Scheduler + from sglang.srt.managers.scheduler_components.pool_stats_observer import ( + SchedulerPoolStatsObserver, + ) logger = logging.getLogger(__name__) class SchedulerRuntimeCheckerMixin: - def _streaming_session_count(self: Scheduler) -> int: + @staticmethod + def streaming_session_count(self: "SchedulerPoolStatsObserver") -> int: return sum( 1 for session in self.session_controller.sessions.values() if session.streaming ) - def _active_pool_idxs(self: Scheduler) -> set: + @staticmethod + def active_pool_idxs(self: "SchedulerPoolStatsObserver") -> set: """Pool idxs currently owned by reqs in last_batch / running_batch. Used to decide which session slots' KV is owned by batch reqs (and thus counted via uncached_size, not session_held). """ idxs = set() - for batch in [self.last_batch, self.running_batch]: + for batch in [self.get_last_batch(), self.get_running_batch()]: if batch is None or batch.is_empty(): continue for req in batch.reqs: @@ -42,35 +47,51 @@ class SchedulerRuntimeCheckerMixin: idxs.add(req.req_pool_idx) return idxs - def _session_held_tokens(self: Scheduler) -> int: - return self.tree_cache.session_held_tokens(self._active_pool_idxs()) + @staticmethod + def session_held_tokens(self: "SchedulerPoolStatsObserver") -> int: + return self.tree_cache.session_held_tokens( + SchedulerRuntimeCheckerMixin.active_pool_idxs(self) + ) - def _session_held_full_tokens(self: Scheduler) -> int: - return self.tree_cache.session_held_full_tokens(self._active_pool_idxs()) + @staticmethod + def session_held_full_tokens(self: "SchedulerPoolStatsObserver") -> int: + return self.tree_cache.session_held_full_tokens( + SchedulerRuntimeCheckerMixin.active_pool_idxs(self) + ) - def _session_held_swa_tokens(self: Scheduler) -> int: - return self.tree_cache.session_held_swa_tokens(self._active_pool_idxs()) + @staticmethod + def session_held_swa_tokens(self: "SchedulerPoolStatsObserver") -> int: + return self.tree_cache.session_held_swa_tokens( + SchedulerRuntimeCheckerMixin.active_pool_idxs(self) + ) - def _session_held_req_count(self: Scheduler) -> int: + @staticmethod + def session_held_req_count(self: "SchedulerPoolStatsObserver") -> int: return self.tree_cache.session_held_req_count() - def _session_held_mamba_slots(self: Scheduler) -> int: - return self.tree_cache.session_held_mamba_slots(self._active_pool_idxs()) + @staticmethod + def session_held_mamba_slots(self: "SchedulerPoolStatsObserver") -> int: + return self.tree_cache.session_held_mamba_slots( + SchedulerRuntimeCheckerMixin.active_pool_idxs(self) + ) - def get_pool_stats(self: Scheduler) -> PoolStats: + @staticmethod + def get_pool_stats(self: "SchedulerPoolStatsObserver") -> PoolStats: if self.is_hybrid_swa: - pool_stats = self._get_swa_token_info() + pool_stats = SchedulerRuntimeCheckerMixin._get_swa_token_info(self) elif self.is_hybrid_ssm: - pool_stats = self._get_mamba_token_info() + pool_stats = SchedulerRuntimeCheckerMixin._get_mamba_token_info(self) else: - pool_stats = self._get_token_info() + pool_stats = SchedulerRuntimeCheckerMixin._get_token_info(self) if self.enable_hisparse: - pool_stats = self._get_hisparse_token_info(pool_stats) + pool_stats = SchedulerRuntimeCheckerMixin._get_hisparse_token_info( + self, pool_stats + ) # swa + ssm can coexist: overlay mamba fields onto swa stats if self.is_hybrid_ssm: - mamba_stats = self._get_mamba_token_info() + mamba_stats = SchedulerRuntimeCheckerMixin._get_mamba_token_info(self) pool_stats.is_hybrid_ssm = True pool_stats.mamba_num_used = mamba_stats.mamba_num_used pool_stats.mamba_usage = mamba_stats.mamba_usage @@ -79,7 +100,8 @@ class SchedulerRuntimeCheckerMixin: return pool_stats - def _get_token_info(self: Scheduler) -> PoolStats: + @staticmethod + def _get_token_info(self: "SchedulerPoolStatsObserver") -> PoolStats: available_size = self.token_to_kv_pool_allocator.available_size() evictable_size = self.tree_cache.evictable_size() num_used = self.max_total_num_tokens - (available_size + evictable_size) @@ -91,7 +113,10 @@ class SchedulerRuntimeCheckerMixin: full_evictable_size=evictable_size, ) - def _get_hisparse_token_info(self: Scheduler, pool_stats: PoolStats) -> PoolStats: + @staticmethod + def _get_hisparse_token_info( + self: "SchedulerPoolStatsObserver", pool_stats: PoolStats + ) -> PoolStats: if self.enable_hisparse and self.hisparse_coordinator is not None: h = self.hisparse_coordinator.get_token_stats() return dataclasses.replace( @@ -104,7 +129,8 @@ class SchedulerRuntimeCheckerMixin: ) return pool_stats - def _get_mamba_token_info(self: Scheduler): + @staticmethod + def _get_mamba_token_info(self: "SchedulerPoolStatsObserver"): is_mamba_radix_cache = ( self.tree_cache.supports_mamba() and self.tree_cache.is_tree_cache() ) @@ -137,7 +163,8 @@ class SchedulerRuntimeCheckerMixin: mamba_evictable_size=mamba_evictable_size, ) - def _get_swa_token_info(self: Scheduler) -> PoolStats: + @staticmethod + def _get_swa_token_info(self: "SchedulerPoolStatsObserver") -> PoolStats: full_available_size = self.token_to_kv_pool_allocator.full_available_size() full_evictable_size = self.tree_cache.full_evictable_size() swa_available_size = self.token_to_kv_pool_allocator.swa_available_size() @@ -194,15 +221,21 @@ class SchedulerRuntimeCheckerMixin: ) -> Tuple[bool, str]: if self.is_hybrid_swa: protected = self.tree_cache.full_protected_size() - session_held = self._session_held_full_tokens() + session_held = self.session_held_full_tokens( + self.pool_stats_observer, + ) total = self.full_tokens_per_layer elif self.is_hybrid_ssm and self.tree_cache.supports_mamba(): protected = self.tree_cache.full_protected_size() - session_held = self._session_held_tokens() + session_held = self.session_held_tokens( + self.pool_stats_observer, + ) total = self.token_to_kv_pool_allocator.size else: protected = self.tree_cache.protected_size() - session_held = self._session_held_tokens() + session_held = self.session_held_tokens( + self.pool_stats_observer, + ) total = self.max_total_num_tokens return self._check_pool_invariant( "full", @@ -222,7 +255,9 @@ class SchedulerRuntimeCheckerMixin: ps.swa_available_size, ps.swa_evictable_size, self.tree_cache.swa_protected_size(), - self._session_held_swa_tokens(), + self.session_held_swa_tokens( + self.pool_stats_observer, + ), self.swa_tokens_per_layer, uncached, ) @@ -233,7 +268,9 @@ class SchedulerRuntimeCheckerMixin: ps.mamba_available_size, ps.mamba_evictable_size, self.tree_cache.mamba_protected_size(), - self._session_held_mamba_slots(), + self.session_held_mamba_slots( + self.pool_stats_observer, + ), self.req_to_token_pool.mamba_pool.size, ) if leak: @@ -314,7 +351,9 @@ class SchedulerRuntimeCheckerMixin: ) return - ps = self.get_pool_stats() + ps = self.get_pool_stats( + self.pool_stats_observer, + ) full_uncached, swa_uncached = self._get_total_uncached_sizes() full_leak, full_msg = self._check_full_pool(ps, uncached=full_uncached) @@ -338,7 +377,9 @@ class SchedulerRuntimeCheckerMixin: else: req_total_size = self.req_to_token_pool.size - session_req_count = self._session_held_req_count() + session_req_count = self.session_held_req_count( + self.pool_stats_observer, + ) if len(self.req_to_token_pool.free_slots) + session_req_count != req_total_size: msg = ( "req_to_token_pool memory leak detected!" @@ -393,9 +434,13 @@ class SchedulerRuntimeCheckerMixin: ): return - self.get_pool_stats().update_scheduler_stats(self.stats) - self.stats.num_streaming_sessions = self._streaming_session_count() - self.stats.streaming_session_held_tokens = self._session_held_tokens() + self.get_pool_stats(self.pool_stats_observer).update_scheduler_stats(self.stats) + self.stats.num_streaming_sessions = self.streaming_session_count( + self.pool_stats_observer, + ) + self.stats.streaming_session_held_tokens = self.session_held_tokens( + self.pool_stats_observer, + ) priority_enabled = self.enable_priority_scheduling self.stats.num_running_reqs = QueueCount.from_reqs( @@ -437,7 +482,11 @@ def create_scheduler_watchdog( def dump_info() -> str: if scheduler.is_initializing: return "" - _, messages = scheduler._check_all_pools(scheduler.get_pool_stats()) + _, messages = scheduler._check_all_pools( + scheduler.get_pool_stats( + scheduler.pool_stats_observer, + ) + ) return ( f"{scheduler.cur_batch.batch_size()=}\n" f"{scheduler.cur_batch.reqs=}\n" + "\n".join(messages) diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index 9a2ab6f64..ce91ca89e 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -499,7 +499,9 @@ class SchedulerMetricsMixin: prefill_stats.log_input_tokens / gap_latency if gap_latency > 0 else 0.0 ) - pool_stats = self.get_pool_stats() + pool_stats = self.get_pool_stats( + self.pool_stats_observer, + ) token_usage_msg = ", ".join(pool_stats.get_prefill_usage_msg_parts()) + ", " self.stats.new_token_ratio = prefill_stats.new_token_ratio @@ -660,7 +662,9 @@ class SchedulerMetricsMixin: self.num_generated_tokens = 0 num_running_reqs = len(batch.reqs) - pool_stats = self.get_pool_stats() + pool_stats = self.get_pool_stats( + self.pool_stats_observer, + ) token_usage_msg = ", ".join(pool_stats.get_decode_usage_msg_parts()) + ", " if RECORD_STEP_TIME: @@ -779,8 +783,12 @@ class SchedulerMetricsMixin: ) # Streaming session metrics - self.stats.num_streaming_sessions = self._streaming_session_count() - self.stats.streaming_session_held_tokens = self._session_held_tokens() + self.stats.num_streaming_sessions = self.streaming_session_count( + self.pool_stats_observer, + ) + self.stats.streaming_session_held_tokens = self.session_held_tokens( + self.pool_stats_observer, + ) # Routing key metrics # (to reduce the overhead, we only compute this when all requests have routing_key) @@ -1011,7 +1019,9 @@ class SchedulerMetricsMixin: waiting_queues.append(self.disagg_decode_prealloc_queue.retracted_queue) num_waiting_reqs = sum(len(queue) for queue in waiting_queues) - num_used_tokens, kv_token_usage = self.get_pool_stats().get_kv_token_stats() + num_used_tokens, kv_token_usage = self.get_pool_stats( + self.pool_stats_observer, + ).get_kv_token_stats() num_total_tokens = num_used_tokens + sum( req.seqlen for queue in waiting_queues for req in queue )