From f6cd1a982249a9c37586c68cebc890b3e2c4bd22 Mon Sep 17 00:00:00 2001 From: cctry Date: Wed, 3 Jun 2026 18:29:49 -0700 Subject: [PATCH] Add num_waiting_uncached_tokens load metric (#27174) Co-authored-by: cctry Co-authored-by: Lianmin Zheng --- python/sglang/srt/managers/io_struct.py | 8 +++++++ python/sglang/srt/managers/load_snapshot.py | 3 +++ python/sglang/srt/managers/schedule_batch.py | 6 +++++ python/sglang/srt/managers/schedule_policy.py | 24 +++++++++++++++---- .../scheduler_components/load_inquirer.py | 15 ++++++++++++ .../sglang/srt/mem_cache/base_prefix_cache.py | 3 +++ .../unit/mem_cache/test_radix_force_miss.py | 4 ++++ 7 files changed, 59 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 8b6251171..439908e49 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -2071,6 +2071,14 @@ class GetLoadsReqOutput(BaseReq): num_waiting_reqs: int = field( metadata={"metric": ("gauge", "Number of waiting requests")} ) + num_waiting_uncached_tokens: int = field( + metadata={ + "metric": ( + "gauge", + "Number of uncached input tokens waiting for prefill compute", + ) + } + ) num_used_tokens: int = field( metadata={"metric": ("gauge", "Number of tokens in use")} ) diff --git a/python/sglang/srt/managers/load_snapshot.py b/python/sglang/srt/managers/load_snapshot.py index ade1ed737..ff0ce49bb 100644 --- a/python/sglang/srt/managers/load_snapshot.py +++ b/python/sglang/srt/managers/load_snapshot.py @@ -144,6 +144,7 @@ CORE_METRIC_FIELDS = ( "dp_rank", "num_running_reqs", "num_waiting_reqs", + "num_waiting_uncached_tokens", "num_used_tokens", "num_total_tokens", "max_total_num_tokens", @@ -218,6 +219,7 @@ class LoadSnapshot(msgspec.Struct, omit_defaults=True): dp_rank: int = 0 num_running_reqs: int = 0 num_waiting_reqs: int = 0 + num_waiting_uncached_tokens: int = 0 num_used_tokens: int = 0 num_total_tokens: int = 0 max_total_num_tokens: int = 0 @@ -292,6 +294,7 @@ class LoadSnapshot(msgspec.Struct, omit_defaults=True): "dp_rank": self.dp_rank, "num_running_reqs": self.num_running_reqs, "num_waiting_reqs": self.num_waiting_reqs, + "num_waiting_uncached_tokens": self.num_waiting_uncached_tokens, "num_used_tokens": self.num_used_tokens, "num_total_tokens": self.num_total_tokens, "max_total_num_tokens": self.max_total_num_tokens, diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 678aea43a..c3886d3b5 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -816,6 +816,11 @@ class Req(ReqDllmMixin): self.last_host_node: Any = None self.best_match_node: Any = None self.host_hit_length = 0 + # Total cached prefix length (on-device prefix_indices + host_hit_length), + # capped at the max allowed prefix. Set during prefix matching at schedule + # time and used to estimate uncached tokens / sort by longest prefix for + # load reporting. + self.num_matched_prefix_tokens = 0 # Tokens loaded from storage backend (L3) during prefetch for this request self.storage_hit_length = 0 # The node to lock until for swa radix tree lock ref @@ -1315,6 +1320,7 @@ class Req(ReqDllmMixin): self.indexer_topk = None self.last_node = None self.cache_protected_len = 0 + self.num_matched_prefix_tokens = 0 self.swa_uuid_for_lock = None self.swa_prefix_lock_released = False self.extend_input_len = 0 diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index cb9857455..43975c57e 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -50,7 +50,7 @@ from sglang.srt.mem_cache.hisparse_memory_pool import ( ) from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator -from sglang.srt.server_args import ServerArgs +from sglang.srt.server_args import ServerArgs, get_global_server_args if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator @@ -115,6 +115,10 @@ def match_prefix_for_req( match_result.best_match_node, match_result.host_hit_length, ) + max_len = req._compute_max_prefix_len(len(token_ids)) + req.num_matched_prefix_tokens = min( + len(req.prefix_indices) + req.host_hit_length, max_len + ) if match_result.mamba_branching_seqlen is not None: req.mamba_branching_seqlen = match_result.mamba_branching_seqlen if match_result.cache_protected_len is not None: @@ -162,6 +166,20 @@ class SchedulePolicy: def calc_priority( self, waiting_queue: List[Req], running_batch: Optional[ScheduleBatch] = None ) -> None: + policy = self._determine_active_policy(waiting_queue) + + # Populate req.num_matched_prefix_tokens at schedule time. Cache-aware policies + # set it in _compute_prefix_matches; do the same full match for + # cache-agnostic policies when the radix supports it, so the load + # snapshot has it. Skip on decode (never prefills). + if ( + not isinstance(policy, CacheAwarePolicy) + and self.tree_cache.supports_fast_match_prefix() + and get_global_server_args().disaggregation_mode != "decode" + ): + for r in waiting_queue: + match_prefix_for_req(self.tree_cache, r) + if self.policy == CacheAgnosticPolicy.FCFS: if self.enable_priority_scheduling: SchedulePolicy._sort_by_priority_and_fcfs( @@ -169,8 +187,6 @@ class SchedulePolicy: ) return - policy = self._determine_active_policy(waiting_queue) - if isinstance(policy, CacheAwarePolicy): temporary_deprioritized = self._compute_prefix_matches( waiting_queue, policy @@ -279,7 +295,7 @@ class SchedulePolicy: """Sorts the waiting queue based on the longest prefix match.""" waiting_queue.sort( key=lambda r: ( - -len(r.prefix_indices) + -r.num_matched_prefix_tokens if r.rid not in temporary_deprioritized else float("inf") ) diff --git a/python/sglang/srt/managers/scheduler_components/load_inquirer.py b/python/sglang/srt/managers/scheduler_components/load_inquirer.py index 32acaa44e..3f10d7eda 100644 --- a/python/sglang/srt/managers/scheduler_components/load_inquirer.py +++ b/python/sglang/srt/managers/scheduler_components/load_inquirer.py @@ -72,6 +72,19 @@ class SchedulerLoadInquirer: num_pending_tokens += req.seqlen - len(req.prefix_indices) - chunk_deduct return num_pending_tokens + def get_num_waiting_uncached_tokens(self) -> int: + """Get uncached input tokens waiting for prefill compute.""" + if self.disaggregation_mode == DisaggregationMode.DECODE: + return 0 + num_tokens = 0 + for req in self.get_waiting_queue(): + # if match-in-waiting-queue disabled, this metric returns seq_lens + num_tokens += max(0, req.seqlen - req.num_matched_prefix_tokens) + cr = self.get_chunked_req() + if cr is not None: + num_tokens += max(0, cr.seqlen - len(cr.prefix_indices)) + return num_tokens + def get_loads(self, req: GetLoadsReqInput = None) -> GetLoadsReqOutput: """ Get comprehensive load metrics for /v1/loads endpoint. @@ -101,6 +114,7 @@ class SchedulerLoadInquirer: ) num_waiting_reqs = sum(len(queue) for queue in waiting_queues) + num_waiting_uncached_tokens = self.get_num_waiting_uncached_tokens() num_used_tokens, kv_token_usage = ( self.pool_stats_observer.get_pool_stats().get_kv_token_stats() ) @@ -193,6 +207,7 @@ class SchedulerLoadInquirer: timestamp=time.time(), num_running_reqs=num_running_reqs, num_waiting_reqs=num_waiting_reqs, + num_waiting_uncached_tokens=num_waiting_uncached_tokens, num_used_tokens=num_used_tokens, num_total_tokens=num_total_tokens, max_total_num_tokens=self.max_total_num_tokens, diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index e0ba5b4a9..1514e601d 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -237,6 +237,9 @@ class BasePrefixCache(ABC, PrefixCacheTrait): def match_prefix(self, params: MatchPrefixParams) -> MatchResult: pass + def supports_fast_match_prefix(self) -> bool: + return False + @abstractmethod def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs): pass diff --git a/test/registered/unit/mem_cache/test_radix_force_miss.py b/test/registered/unit/mem_cache/test_radix_force_miss.py index 77d2be962..82c77401f 100644 --- a/test/registered/unit/mem_cache/test_radix_force_miss.py +++ b/test/registered/unit/mem_cache/test_radix_force_miss.py @@ -36,9 +36,13 @@ class _StubReq: self.last_host_node = None self.best_match_node = None self.host_hit_length = None + self.num_matched_prefix_tokens = 0 self.mamba_branching_seqlen = None self.cache_protected_len = None + def _compute_max_prefix_len(self, input_len): + return max(input_len - 1, 0) + class TestZeroMatchResult(unittest.TestCase): def test_zero_replaces_indices_and_nodes(self):