From 526af158452c7b9323a14b7fa5016aeb51b49efd Mon Sep 17 00:00:00 2001 From: Hanming Lu <69857889+hanming-lu@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:27:21 -0700 Subject: [PATCH] [Metrics] Discount queued prefill load by recent cache hits when waiting-queue matching is off (#35248) --- python/sglang/srt/environ.py | 3 + python/sglang/srt/managers/schedule_policy.py | 7 ++ python/sglang/srt/managers/scheduler.py | 4 ++ .../scheduler_components/load_inquirer.py | 12 +++- .../scheduler_components/metrics_reporter.py | 31 ++++++++ .../unit/managers/test_load_inquirer.py | 70 +++++++++++++++++++ .../test_forward_pass_metrics.py | 7 ++ 7 files changed, 131 insertions(+), 3 deletions(-) create mode 100644 test/registered/unit/managers/test_load_inquirer.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 3abbabe13..e9acf9578 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -522,6 +522,9 @@ class Envs: SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION = EnvInt(4096) SGLANG_MAX_NEW_TOKENS_LIMIT = EnvInt(None) SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR = EnvFloat(0.75) + # Window for the token-weighted recent cache-hit rate used to estimate + # waiting-queue prefill load. + SGLANG_CACHE_HIT_RATE_WINDOW_SECONDS = EnvFloat(15.0) SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES = EnvInt(None) SGLANG_PREFILL_DELAYER_TOKEN_USAGE_LOW_WATERMARK = EnvFloat(None) SGLANG_DATA_PARALLEL_BUDGET_INTERVAL = EnvInt(1) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 6780d914a..34472fce2 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -293,6 +293,13 @@ class SchedulePolicy: return CacheAgnosticPolicy.FCFS return self.policy + def waiting_queue_prefix_matched(self, waiting_queue: List[Req]) -> bool: + policy = self._determine_active_policy(waiting_queue) + return ( + isinstance(policy, CacheAwarePolicy) + or self.tree_cache.supports_fast_match_prefix() + ) + def _validate_and_adjust_policy( self, policy: str, tree_cache: BasePrefixCache ) -> Policy: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ecf79dc20..b60db6484 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2123,6 +2123,10 @@ class Scheduler( spec_algorithm=self.spec_algorithm, get_running_batch=lambda: self.running_batch, get_waiting_queue=lambda: self.waiting_queue, + waiting_queue_prefix_matched=lambda: self.policy.waiting_queue_prefix_matched( + self.waiting_queue + ), + get_recent_cache_hit_rate=lambda: self.metrics_reporter.recent_cache_hit_rate, get_stats=lambda: self.metrics_reporter.stats, get_chunked_req=lambda: self.chunked_req, get_disagg_prefill_bootstrap_queue=lambda: self.disagg_prefill_bootstrap_queue, diff --git a/python/sglang/srt/managers/scheduler_components/load_inquirer.py b/python/sglang/srt/managers/scheduler_components/load_inquirer.py index 988b41c67..e470a40ae 100644 --- a/python/sglang/srt/managers/scheduler_components/load_inquirer.py +++ b/python/sglang/srt/managers/scheduler_components/load_inquirer.py @@ -43,6 +43,8 @@ class SchedulerLoadInquirer: spec_algorithm: SpeculativeAlgorithm get_running_batch: Callable get_waiting_queue: Callable + waiting_queue_prefix_matched: Callable + get_recent_cache_hit_rate: Callable get_stats: Callable get_chunked_req: Callable get_disagg_prefill_bootstrap_queue: Callable @@ -76,13 +78,17 @@ class SchedulerLoadInquirer: return num_pending_tokens def get_num_waiting_uncached_tokens(self) -> int: - """Get uncached input tokens waiting for prefill compute.""" + """Estimate input tokens waiting for prefill compute.""" if self.disaggregation_mode == DisaggregationMode.DECODE: return 0 + waiting_queue_prefix_matched = self.waiting_queue_prefix_matched() + cache_miss_rate = 1.0 - self.get_recent_cache_hit_rate() 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) + if waiting_queue_prefix_matched: + num_tokens += max(0, req.seqlen - req.num_matched_prefix_tokens) + else: + num_tokens += int(req.seqlen * cache_miss_rate) cr = self.get_chunked_req() if cr is not None: num_tokens += max(0, cr.seqlen - len(cr.prefix_indices)) diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 0f4ad921b..b6a12ca66 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -39,6 +39,28 @@ logger = logging.getLogger(__name__) RECORD_STEP_TIME = envs.SGLANG_RECORD_STEP_TIME.get() LOG_FORWARD_ITERS = envs.SGLANG_LOG_FORWARD_ITERS.get() ENABLE_METRICS_DEVICE_TIMER = envs.SGLANG_ENABLE_METRICS_DEVICE_TIMER.get() +CACHE_HIT_RATE_WINDOW_SECONDS = envs.SGLANG_CACHE_HIT_RATE_WINDOW_SECONDS.get() + + +class _CacheHitRateWindow: + def __init__(self) -> None: + self.samples = deque() + self.hit_tokens = 0 + self.total_tokens = 0 + + def add(self, hit_tokens: int, total_tokens: int, now: float) -> float: + if total_tokens > 0: + self.samples.append((now, hit_tokens, total_tokens)) + self.hit_tokens += hit_tokens + self.total_tokens += total_tokens + + cutoff = now - CACHE_HIT_RATE_WINDOW_SECONDS + while self.samples and self.samples[0][0] <= cutoff: + _, expired_hit_tokens, expired_total_tokens = self.samples.popleft() + self.hit_tokens -= expired_hit_tokens + self.total_tokens -= expired_total_tokens + + return self.hit_tokens / self.total_tokens if self.total_tokens > 0 else 0.0 def _decode_total_seq_lens(batch: ScheduleBatch) -> int: @@ -118,6 +140,10 @@ class SchedulerMetricsReporter: self._eplb_balancedness_history = [ deque(maxlen=window_size) for window_size in EPLB_BALANCEDNESS_WINDOW_SIZES ] + self.cache_hit_rate_window = _CacheHitRateWindow() + # Windowed rate for waiting-queue load estimation only; the exported + # cache_hit_rate stats keep their per-report semantics. + self.recent_cache_hit_rate = 0.0 def _init_metrics( self, @@ -654,6 +680,11 @@ class SchedulerMetricsReporter: cache_hit_rate = ( effective_hit_tokens / total_tokens if total_tokens > 0 else 0.0 ) + self.recent_cache_hit_rate = self.cache_hit_rate_window.add( + effective_hit_tokens, + total_tokens, + now, + ) self.metrics_collector.increment_effective_prefill_tokens( input_tokens=effective_input_tokens, device_hit_tokens=prefill_stats.log_device_hit_tokens, diff --git a/test/registered/unit/managers/test_load_inquirer.py b/test/registered/unit/managers/test_load_inquirer.py new file mode 100644 index 000000000..73f3f1478 --- /dev/null +++ b/test/registered/unit/managers/test_load_inquirer.py @@ -0,0 +1,70 @@ +import unittest +from types import SimpleNamespace + +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.managers.schedule_policy import ( + CacheAgnosticPolicy, + CacheAwarePolicy, + SchedulePolicy, +) +from sglang.srt.managers.scheduler_components.load_inquirer import ( + SchedulerLoadInquirer, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestSchedulePolicyWaitingQueueMatching(unittest.TestCase): + def make_policy(self, policy, supports_fast_match_prefix): + schedule_policy = object.__new__(SchedulePolicy) + schedule_policy.policy = policy + schedule_policy.tree_cache = SimpleNamespace( + supports_fast_match_prefix=lambda: supports_fast_match_prefix + ) + return schedule_policy + + def test_cache_agnostic_policy_requires_fast_matching(self): + policy = self.make_policy(CacheAgnosticPolicy.FCFS, False) + self.assertFalse(policy.waiting_queue_prefix_matched([])) + + policy.tree_cache = SimpleNamespace(supports_fast_match_prefix=lambda: True) + self.assertTrue(policy.waiting_queue_prefix_matched([])) + + def test_lpm_queue_limit_can_disable_matching(self): + policy = self.make_policy(CacheAwarePolicy.LPM, False) + self.assertTrue(policy.waiting_queue_prefix_matched([None] * 128)) + self.assertFalse(policy.waiting_queue_prefix_matched([None] * 129)) + + +class TestSchedulerLoadInquirer(unittest.TestCase): + def make_inquirer(self, waiting_queue_prefix_matched): + waiting_req = SimpleNamespace(seqlen=100, num_matched_prefix_tokens=20) + chunked_req = SimpleNamespace(seqlen=50, prefix_indices=range(10)) + return SimpleNamespace( + disaggregation_mode=DisaggregationMode.NULL, + get_waiting_queue=lambda: [waiting_req], + waiting_queue_prefix_matched=lambda: waiting_queue_prefix_matched, + get_chunked_req=lambda: chunked_req, + get_recent_cache_hit_rate=lambda: 0.75, + ) + + def test_waiting_tokens_are_estimated_when_prefix_matching_is_skipped(self): + inquirer = self.make_inquirer(waiting_queue_prefix_matched=False) + + self.assertEqual( + SchedulerLoadInquirer.get_num_waiting_uncached_tokens(inquirer), + 65, + ) + + def test_waiting_tokens_use_exact_match_when_prefix_matching_is_done(self): + inquirer = self.make_inquirer(waiting_queue_prefix_matched=True) + + self.assertEqual( + SchedulerLoadInquirer.get_num_waiting_uncached_tokens(inquirer), + 120, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/observability/test_forward_pass_metrics.py b/test/registered/unit/observability/test_forward_pass_metrics.py index 6c458fc17..57dd4018e 100644 --- a/test/registered/unit/observability/test_forward_pass_metrics.py +++ b/test/registered/unit/observability/test_forward_pass_metrics.py @@ -13,6 +13,7 @@ from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.managers.scheduler_components.metrics_reporter import ( PrefillStats, SchedulerMetricsReporter, + _CacheHitRateWindow, ) from sglang.test.test_utils import CustomTestCase @@ -136,6 +137,12 @@ class TestForwardPassMetrics(unittest.TestCase): self.reporter = _make_reporter(self, self.scheduler) self.scheduler.enable_fpm = True + def test_cache_hit_rate_window_keeps_last_15s_of_tokens(self): + window = _CacheHitRateWindow() + self.assertEqual(window.add(hit_tokens=20, total_tokens=100, now=0.0), 0.2) + self.assertEqual(window.add(hit_tokens=80, total_tokens=100, now=10.0), 0.5) + self.assertEqual(window.add(hit_tokens=90, total_tokens=100, now=15.0), 0.85) + def _make_batch(self, **overrides): defaults = dict( forward_mode=_FakeForwardMode(),