From b4a1d8fd714d39ef11bec6fc444443254ed04b3f Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Fri, 10 Apr 2026 16:20:33 -0700 Subject: [PATCH] [mem] Fix idle token_usage missing mamba_usage; add FIXME for naming (#22555) --- python/sglang/srt/managers/io_struct.py | 2 ++ .../sglang/srt/managers/scheduler_runtime_checker_mixin.py | 5 +++-- python/sglang/srt/observability/metrics_collector.py | 2 ++ 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 5df2dec2b..b0657777f 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1981,6 +1981,8 @@ class GetLoadsReqOutput(BaseReq): max_total_num_tokens: int = field( metadata={"metric": ("gauge", "Maximum token capacity")} ) + # FIXME: token_usage is actually max usage across all pools (KV, SWA, mamba), + # not just KV token usage. Rename requires API deprecation. token_usage: float = field(metadata={"metric": ("gauge", "Token pool usage ratio")}) gen_throughput: float = field( metadata={"metric": ("gauge", "Generation throughput tokens/sec")} diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 113073e3b..d9672ffa0 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -311,13 +311,14 @@ class SchedulerRuntimeCheckerMixin: ( num_used, _, - token_usage, - _, + full_token_usage, + mamba_usage, _, _, _, _, ) = self._get_mamba_token_info() + token_usage = max(full_token_usage, mamba_usage) else: num_used, token_usage, _, _ = self._get_token_info() diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 490d755ed..9a81ee78f 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -79,6 +79,8 @@ class SchedulerStats: # Basics num_running_reqs: QueueCount = field(default_factory=QueueCount) num_used_tokens: int = 0 + # FIXME: token_usage is actually max usage across all pools (KV, SWA, mamba), + # not just KV token usage. Rename requires API deprecation. token_usage: float = 0.0 full_token_usage: float = 0.0 pending_prealloc_token_usage: float = 0.0