minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Move pool-stats sampling to SchedulerPoolStatsObserver (#25621)

Browse Source
This commit is contained in:
fzyzcjy
2026-05-18 18:37:50 +08:00
committed by GitHub
parent 9fdf73f393
commit ee392a1e14
4 changed files with 176 additions and 227 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+4 -6
View File
@@ -2343,9 +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(
self.pool_stats_observer,
).get_max_pool_usage()
max_pool_usage = (
self.pool_stats_observer.get_pool_stats().get_max_pool_usage()
)
prefill_delayer_single_pass = PrefillDelayerSinglePassExecutor(
self.prefill_delayer, token_usage=max_pool_usage
)
@@ -3061,9 +3061,7 @@ class Scheduler(
# 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(
self.pool_stats_observer,
)
self.pool_stats_observer.get_pool_stats()
)
if has_leak:
self._report_leak("pool", "\n".join(messages))
python/sglang/srt/managers/scheduler_components/pool_stats_observer.py
+150
View File
@@ -153,3 +153,153 @@ class SchedulerPoolStatsObserver:
max_total_num_tokens: int
get_last_batch: Callable
get_running_batch: Callable
def streaming_session_count(self) -> int:
return sum(
1
for session in self.session_controller.sessions.values()
if session.streaming
)
def active_pool_idxs(self) -> 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.get_last_batch(), self.get_running_batch()]:
if batch is None or batch.is_empty():
continue
for req in batch.reqs:
if req.req_pool_idx is not None:
idxs.add(req.req_pool_idx)
return idxs
def session_held_tokens(self) -> int:
return self.tree_cache.session_held_tokens(self.active_pool_idxs())
def session_held_full_tokens(self) -> int:
return self.tree_cache.session_held_full_tokens(self.active_pool_idxs())
def session_held_swa_tokens(self) -> int:
return self.tree_cache.session_held_swa_tokens(self.active_pool_idxs())
def session_held_req_count(self) -> int:
return self.tree_cache.session_held_req_count()
def session_held_mamba_slots(self) -> int:
return self.tree_cache.session_held_mamba_slots(self.active_pool_idxs())
def get_pool_stats(self) -> PoolStats:
if self.is_hybrid_swa:
pool_stats = self._get_swa_token_info()
elif self.is_hybrid_ssm:
pool_stats = self._get_mamba_token_info()
else:
pool_stats = self._get_token_info()
if self.enable_hisparse:
pool_stats = self._get_hisparse_token_info(pool_stats)
# swa + ssm can coexist: overlay mamba fields onto swa stats
if self.is_hybrid_ssm:
mamba_stats = self._get_mamba_token_info()
pool_stats.is_hybrid_ssm = True
pool_stats.mamba_num_used = mamba_stats.mamba_num_used
pool_stats.mamba_usage = mamba_stats.mamba_usage
pool_stats.mamba_available_size = mamba_stats.mamba_available_size
pool_stats.mamba_evictable_size = mamba_stats.mamba_evictable_size
return pool_stats
def _get_token_info(self) -> 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)
token_usage = num_used / self.max_total_num_tokens
return PoolStats(
full_num_used=num_used,
full_token_usage=token_usage,
full_available_size=available_size,
full_evictable_size=evictable_size,
)
def _get_hisparse_token_info(self, 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(
pool_stats,
is_hisparse=True,
hisparse_device_tokens=h.device_tokens,
hisparse_device_token_usage=h.device_token_usage,
hisparse_host_tokens=h.host_tokens,
hisparse_host_token_usage=h.host_token_usage,
)
return pool_stats
def _get_mamba_token_info(self):
is_mamba_radix_cache = (
self.tree_cache.supports_mamba() and self.tree_cache.is_tree_cache()
)
full_available_size = self.token_to_kv_pool_allocator.available_size()
full_evictable_size = (
self.tree_cache.full_evictable_size() if is_mamba_radix_cache else 0
)
mamba_available_size = self.req_to_token_pool.mamba_pool.available_size()
mamba_evictable_size = (
self.tree_cache.mamba_evictable_size() if is_mamba_radix_cache else 0
)
full_num_used = self.token_to_kv_pool_allocator.size - (
full_available_size + full_evictable_size
)
mamba_num_used = self.req_to_token_pool.mamba_pool.size - (
mamba_available_size + mamba_evictable_size
)
full_token_usage = full_num_used / self.token_to_kv_pool_allocator.size
mamba_usage = mamba_num_used / self.req_to_token_pool.mamba_pool.size
return PoolStats(
is_hybrid_ssm=True,
full_num_used=full_num_used,
full_token_usage=full_token_usage,
full_available_size=full_available_size,
full_evictable_size=full_evictable_size,
mamba_num_used=mamba_num_used,
mamba_usage=mamba_usage,
mamba_available_size=mamba_available_size,
mamba_evictable_size=mamba_evictable_size,
)
def _get_swa_token_info(self) -> 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()
swa_evictable_size = self.tree_cache.swa_evictable_size()
full_num_used = self.full_tokens_per_layer - (
full_available_size + full_evictable_size
)
swa_num_used = self.swa_tokens_per_layer - (
swa_available_size + swa_evictable_size
)
# FIXME(hisparse): host-backup transiently over-releases the device pool
# counter, producing negative full_num_used / swa_num_used. We clamp to 0
# to keep token_usage / leak checks sane, but the underlying accounting
# bug should be fixed so the clamp can go away.
if self.enable_hisparse:
full_num_used = max(0, full_num_used)
swa_num_used = max(0, swa_num_used)
full_token_usage = full_num_used / self.full_tokens_per_layer
swa_token_usage = swa_num_used / self.swa_tokens_per_layer
return PoolStats(
is_hybrid_swa=True,
full_num_used=full_num_used,
full_token_usage=full_token_usage,
full_available_size=full_available_size,
full_evictable_size=full_evictable_size,
swa_num_used=swa_num_used,
swa_token_usage=swa_token_usage,
swa_available_size=swa_available_size,
swa_evictable_size=swa_evictable_size,
)
python/sglang/srt/managers/scheduler_runtime_checker_mixin.py
+13 -208
View File
@@ -1,6 +1,5 @@
from __future__ import annotations
import dataclasses
import logging
import time
import warnings
@@ -8,195 +7,17 @@ from typing import TYPE_CHECKING, List, Tuple
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.environ import envs
from sglang.srt.managers.scheduler_components.pool_stats_observer import PoolStats
from sglang.srt.observability.metrics_collector import QueueCount
from sglang.srt.utils.common import ceil_align, raise_error_or_warn
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:
@staticmethod
def streaming_session_count(self: "SchedulerPoolStatsObserver") -> int:
return sum(
1
for session in self.session_controller.sessions.values()
if session.streaming
)
@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.get_last_batch(), self.get_running_batch()]:
if batch is None or batch.is_empty():
continue
for req in batch.reqs:
if req.req_pool_idx is not None:
idxs.add(req.req_pool_idx)
return idxs
@staticmethod
def session_held_tokens(self: "SchedulerPoolStatsObserver") -> int:
return self.tree_cache.session_held_tokens(
SchedulerRuntimeCheckerMixin.active_pool_idxs(self)
)
@staticmethod
def session_held_full_tokens(self: "SchedulerPoolStatsObserver") -> int:
return self.tree_cache.session_held_full_tokens(
SchedulerRuntimeCheckerMixin.active_pool_idxs(self)
)
@staticmethod
def session_held_swa_tokens(self: "SchedulerPoolStatsObserver") -> int:
return self.tree_cache.session_held_swa_tokens(
SchedulerRuntimeCheckerMixin.active_pool_idxs(self)
)
@staticmethod
def session_held_req_count(self: "SchedulerPoolStatsObserver") -> int:
return self.tree_cache.session_held_req_count()
@staticmethod
def session_held_mamba_slots(self: "SchedulerPoolStatsObserver") -> int:
return self.tree_cache.session_held_mamba_slots(
SchedulerRuntimeCheckerMixin.active_pool_idxs(self)
)
@staticmethod
def get_pool_stats(self: "SchedulerPoolStatsObserver") -> PoolStats:
if self.is_hybrid_swa:
pool_stats = SchedulerRuntimeCheckerMixin._get_swa_token_info(self)
elif self.is_hybrid_ssm:
pool_stats = SchedulerRuntimeCheckerMixin._get_mamba_token_info(self)
else:
pool_stats = SchedulerRuntimeCheckerMixin._get_token_info(self)
if self.enable_hisparse:
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 = 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
pool_stats.mamba_available_size = mamba_stats.mamba_available_size
pool_stats.mamba_evictable_size = mamba_stats.mamba_evictable_size
return pool_stats
@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)
token_usage = num_used / self.max_total_num_tokens
return PoolStats(
full_num_used=num_used,
full_token_usage=token_usage,
full_available_size=available_size,
full_evictable_size=evictable_size,
)
@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(
pool_stats,
is_hisparse=True,
hisparse_device_tokens=h.device_tokens,
hisparse_device_token_usage=h.device_token_usage,
hisparse_host_tokens=h.host_tokens,
hisparse_host_token_usage=h.host_token_usage,
)
return pool_stats
@staticmethod
def _get_mamba_token_info(self: "SchedulerPoolStatsObserver"):
is_mamba_radix_cache = (
self.tree_cache.supports_mamba() and self.tree_cache.is_tree_cache()
)
full_available_size = self.token_to_kv_pool_allocator.available_size()
full_evictable_size = (
self.tree_cache.full_evictable_size() if is_mamba_radix_cache else 0
)
mamba_available_size = self.req_to_token_pool.mamba_pool.available_size()
mamba_evictable_size = (
self.tree_cache.mamba_evictable_size() if is_mamba_radix_cache else 0
)
full_num_used = self.token_to_kv_pool_allocator.size - (
full_available_size + full_evictable_size
)
mamba_num_used = self.req_to_token_pool.mamba_pool.size - (
mamba_available_size + mamba_evictable_size
)
full_token_usage = full_num_used / self.token_to_kv_pool_allocator.size
mamba_usage = mamba_num_used / self.req_to_token_pool.mamba_pool.size
return PoolStats(
is_hybrid_ssm=True,
full_num_used=full_num_used,
full_token_usage=full_token_usage,
full_available_size=full_available_size,
full_evictable_size=full_evictable_size,
mamba_num_used=mamba_num_used,
mamba_usage=mamba_usage,
mamba_available_size=mamba_available_size,
mamba_evictable_size=mamba_evictable_size,
)
@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()
swa_evictable_size = self.tree_cache.swa_evictable_size()
full_num_used = self.full_tokens_per_layer - (
full_available_size + full_evictable_size
)
swa_num_used = self.swa_tokens_per_layer - (
swa_available_size + swa_evictable_size
)
# FIXME(hisparse): host-backup transiently over-releases the device pool
# counter, producing negative full_num_used / swa_num_used. We clamp to 0
# to keep token_usage / leak checks sane, but the underlying accounting
# bug should be fixed so the clamp can go away.
if self.enable_hisparse:
full_num_used = max(0, full_num_used)
swa_num_used = max(0, swa_num_used)
full_token_usage = full_num_used / self.full_tokens_per_layer
swa_token_usage = swa_num_used / self.swa_tokens_per_layer
return PoolStats(
is_hybrid_swa=True,
full_num_used=full_num_used,
full_token_usage=full_token_usage,
full_available_size=full_available_size,
full_evictable_size=full_evictable_size,
swa_num_used=swa_num_used,
swa_token_usage=swa_token_usage,
swa_available_size=swa_available_size,
swa_evictable_size=swa_evictable_size,
)
@staticmethod
def _check_pool_invariant(
pool_name: str,
@@ -221,21 +42,15 @@ class SchedulerRuntimeCheckerMixin:
) -> Tuple[bool, str]:
if self.is_hybrid_swa:
protected = self.tree_cache.full_protected_size()
session_held = self.session_held_full_tokens(
self.pool_stats_observer,
)
session_held = self.pool_stats_observer.session_held_full_tokens()
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(
self.pool_stats_observer,
)
session_held = self.pool_stats_observer.session_held_tokens()
total = self.token_to_kv_pool_allocator.size
else:
protected = self.tree_cache.protected_size()
session_held = self.session_held_tokens(
self.pool_stats_observer,
)
session_held = self.pool_stats_observer.session_held_tokens()
total = self.max_total_num_tokens
return self._check_pool_invariant(
"full",
@@ -255,9 +70,7 @@ class SchedulerRuntimeCheckerMixin:
ps.swa_available_size,
ps.swa_evictable_size,
self.tree_cache.swa_protected_size(),
self.session_held_swa_tokens(
self.pool_stats_observer,
),
self.pool_stats_observer.session_held_swa_tokens(),
self.swa_tokens_per_layer,
uncached,
)
@@ -268,9 +81,7 @@ class SchedulerRuntimeCheckerMixin:
ps.mamba_available_size,
ps.mamba_evictable_size,
self.tree_cache.mamba_protected_size(),
self.session_held_mamba_slots(
self.pool_stats_observer,
),
self.pool_stats_observer.session_held_mamba_slots(),
self.req_to_token_pool.mamba_pool.size,
)
if leak:
@@ -351,9 +162,7 @@ class SchedulerRuntimeCheckerMixin:
)
return
ps = self.get_pool_stats(
self.pool_stats_observer,
)
ps = self.pool_stats_observer.get_pool_stats()
full_uncached, swa_uncached = self._get_total_uncached_sizes()
full_leak, full_msg = self._check_full_pool(ps, uncached=full_uncached)
@@ -377,9 +186,7 @@ class SchedulerRuntimeCheckerMixin:
else:
req_total_size = self.req_to_token_pool.size
session_req_count = self.session_held_req_count(
self.pool_stats_observer,
)
session_req_count = self.pool_stats_observer.session_held_req_count()
if len(self.req_to_token_pool.free_slots) + session_req_count != req_total_size:
msg = (
"req_to_token_pool memory leak detected!"
@@ -434,12 +241,12 @@ class SchedulerRuntimeCheckerMixin:
):
return
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.pool_stats_observer.get_pool_stats().update_scheduler_stats(self.stats)
self.stats.num_streaming_sessions = (
self.pool_stats_observer.streaming_session_count()
)
self.stats.streaming_session_held_tokens = self.session_held_tokens(
self.pool_stats_observer,
self.stats.streaming_session_held_tokens = (
self.pool_stats_observer.session_held_tokens()
)
priority_enabled = self.enable_priority_scheduling
@@ -483,9 +290,7 @@ def create_scheduler_watchdog(
if scheduler.is_initializing:
return ""
_, messages = scheduler._check_all_pools(
scheduler.get_pool_stats(
scheduler.pool_stats_observer,
)
scheduler.pool_stats_observer.get_pool_stats()
)
return (
f"{scheduler.cur_batch.batch_size()=}\n"
python/sglang/srt/observability/scheduler_metrics_mixin.py
+9 -13
View File
@@ -499,9 +499,7 @@ class SchedulerMetricsMixin:
prefill_stats.log_input_tokens / gap_latency if gap_latency > 0 else 0.0
)
pool_stats = self.get_pool_stats(
self.pool_stats_observer,
)
pool_stats = self.pool_stats_observer.get_pool_stats()
token_usage_msg = ", ".join(pool_stats.get_prefill_usage_msg_parts()) + ", "
self.stats.new_token_ratio = prefill_stats.new_token_ratio
@@ -662,9 +660,7 @@ class SchedulerMetricsMixin:
self.num_generated_tokens = 0
num_running_reqs = len(batch.reqs)
pool_stats = self.get_pool_stats(
self.pool_stats_observer,
)
pool_stats = self.pool_stats_observer.get_pool_stats()
token_usage_msg = ", ".join(pool_stats.get_decode_usage_msg_parts()) + ", "
if RECORD_STEP_TIME:
@@ -783,11 +779,11 @@ class SchedulerMetricsMixin:
)
# Streaming session metrics
self.stats.num_streaming_sessions = self.streaming_session_count(
self.pool_stats_observer,
self.stats.num_streaming_sessions = (
self.pool_stats_observer.streaming_session_count()
)
self.stats.streaming_session_held_tokens = self.session_held_tokens(
self.pool_stats_observer,
self.stats.streaming_session_held_tokens = (
self.pool_stats_observer.session_held_tokens()
)
# Routing key metrics
@@ -1019,9 +1015,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(
self.pool_stats_observer,
).get_kv_token_stats()
num_used_tokens, kv_token_usage = (
self.pool_stats_observer.get_pool_stats().get_kv_token_stats()
)
num_total_tokens = num_used_tokens + sum(
req.seqlen for queue in waiting_queues for req in queue
)