From 4df5df911ba9f1d1ffdeb1516e106f2df691fbdf Mon Sep 17 00:00:00 2001 From: SovietPower <73899105+SovietPower@users.noreply.github.com> Date: Tue, 8 Sep 2026 22:14:03 +0800 Subject: [PATCH] [Scheduler] Add HRRN schedule policy to significantly reduce TTFT (#32911) Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- .../sglang/srt/arg_groups/fields/schedule.py | 1 + python/sglang/srt/disaggregation/prefill.py | 5 + python/sglang/srt/managers/schedule_batch.py | 3 + python/sglang/srt/managers/schedule_policy.py | 61 ++++++++- python/sglang/srt/managers/scheduler.py | 13 +- .../test_generation_auxiliary_output.py | 2 + .../unit/managers/test_schedule_policy.py | 128 ++++++++++++++++++ .../unit/test_server_args_cli_metadata.py | 11 +- 8 files changed, 220 insertions(+), 4 deletions(-) create mode 100644 test/registered/unit/managers/test_schedule_policy.py diff --git a/python/sglang/srt/arg_groups/fields/schedule.py b/python/sglang/srt/arg_groups/fields/schedule.py index 7ee9aca2e..bee8e21ee 100644 --- a/python/sglang/srt/arg_groups/fields/schedule.py +++ b/python/sglang/srt/arg_groups/fields/schedule.py @@ -95,6 +95,7 @@ class Schedule: "lof", "priority", "routing-key", + "hrrn", ], ), ] = "fcfs" diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 13348fbf7..b91ea8161 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -486,6 +486,9 @@ class PrefillBootstrapQueue: bootstrapped_reqs.append(req) indices_to_remove.add(i) req.time_stats.set_wait_queue_entry_time() + req.arrival_processed_tokens = ( + self.scheduler.processed_tokens_counter + ) elif poll == KVPoll.WaitingForInput: if should_force_retry(req): # skip checking for testing if not self.ensure_metadata_buffer(req): @@ -496,6 +499,7 @@ class PrefillBootstrapQueue: bootstrapped_reqs.append(req) indices_to_remove.add(i) req.time_stats.set_wait_queue_entry_time() + req.arrival_processed_tokens = self.scheduler.processed_tokens_counter else: raise RuntimeError( f"Unexpected poll state {poll} for req {req.rid} in pop_bootstrapped" @@ -1457,4 +1461,5 @@ class SchedulerDisaggregationPrefillMixin: if self.metrics_reporter.enable_metrics: self.metrics_collector.increment_prefill_retries(1) req.time_stats.set_wait_queue_entry_time() + req.arrival_processed_tokens = self.processed_tokens_counter self.waiting_queue.insert(0, req) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 1fcc3eacc..d5683500e 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1323,6 +1323,9 @@ class Req(ReqDllmMixin): # For hisparse self.hisparse_staging = False + # Snapshot of the scheduler prefill-token counter taken at waiting_queue entry; used by HRRN aging. + self.arrival_processed_tokens: int = 0 + @property def seqlen(self) -> int: """Get the current sequence length of the request.""" diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index ceb9d56ae..1b884b6ec 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -207,6 +207,7 @@ class CacheAwarePolicy(Enum): LPM = "lpm" # longest prefix match DFS_WEIGHT = "dfs-weight" # depth-first search weighting + HRRN = "hrrn" # highest response ratio next, token-based aging class CacheAgnosticPolicy(Enum): @@ -240,7 +241,10 @@ class SchedulePolicy: self.waiting_queue_radix_tree = RadixCache.create_simulated() def calc_priority( - self, waiting_queue: List[Req], running_batch: Optional[ScheduleBatch] = None + self, + waiting_queue: List[Req], + running_batch: Optional[ScheduleBatch] = None, + processed_tokens: int = 0, ) -> None: policy = self._determine_active_policy(waiting_queue) @@ -273,6 +277,10 @@ class SchedulePolicy: ) elif policy == CacheAwarePolicy.DFS_WEIGHT: SchedulePolicy._sort_by_dfs_weight(waiting_queue, self.tree_cache) + elif policy == CacheAwarePolicy.HRRN: + SchedulePolicy._sort_by_hrrn( + waiting_queue, temporary_deprioritized, processed_tokens + ) else: raise ValueError(f"Unknown CacheAware Policy: {policy=}") else: @@ -293,7 +301,14 @@ class SchedulePolicy: raise ValueError(f"Unknown CacheAgnostic Policy: {policy=}") def _determine_active_policy(self, waiting_queue: List[Req]) -> Policy: - if self.policy == CacheAwarePolicy.LPM and len(waiting_queue) > 128: + if ( + self.policy + in ( + CacheAwarePolicy.LPM, + CacheAwarePolicy.HRRN, + ) + and len(waiting_queue) > 128 + ): # Turn off the expensive prefix matching and sorting when the #queue is large. return CacheAgnosticPolicy.FCFS return self.policy @@ -395,6 +410,48 @@ class SchedulePolicy: ) ) + @staticmethod + def _uncached_len(r: Req) -> int: + """Number of tokens that must actually be prefilled for this req + (all cache levels — device + host via hicache — counted as cached).""" + return max(0, len(r.origin_input_ids) - r.num_matched_prefix_tokens) + + @staticmethod + def _sort_by_hrrn( + waiting_queue: List[Req], + temporary_deprioritized: Set[int], + processed_tokens: int, + ) -> None: + """Highest Response Ratio Next, with token-based aging. + + Equivalence with classic HRRN when throughput is constant: + ratio = 1 + wait_sec / est_prefill_time + = 1 + (processed_tokens - arrival_processed_tokens) / uncached + + Caller (Scheduler) contract: + - Maintain a monotonically increasing counter of prefill tokens processed so far + (accumulate batch.extend_num_tokens per forward). Pass it in as `processed_tokens`. + - Snapshot `req.arrival_processed_tokens = counter` when the req enters waiting_queue + (pop_bootstrapped for disagg prefill, _add_request_to_queue for unified). + + Call sites that omit `processed_tokens` (dllm, disagg decode) + degrade to rid-lexicographic order; those queues carry no prefill work. + """ + + def _key(r: Req): + rid = r.rid + if rid in temporary_deprioritized: + return (float("inf"), rid) + uncached = SchedulePolicy._uncached_len(r) + if uncached <= 0: + # No prefill work; drain immediately. + return (-float("inf"), rid) + waited_tokens = max(0, processed_tokens - r.arrival_processed_tokens) + ratio_delta = waited_tokens / uncached + return (-ratio_delta, rid) + + waiting_queue.sort(key=_key) + @staticmethod def _sort_by_dfs_weight( waiting_queue: List[Req], tree_cache: BasePrefixCache diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index cf361cc8d..461c28a65 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -456,6 +456,8 @@ class Scheduler( self.is_initializing = True # init_soft_watchdog starts a daemon thread that reads these on its first tick. self.forward_ct: int = 0 + # Prefill tokens processed so far; used as the aging axis for the HRRN scheduling policy. Reqs snapshot this at waiting_queue entry. + self.processed_tokens_counter: int = 0 self.cur_batch_for_debug: Optional[ScheduleBatch] = None self.init_soft_watchdog() @@ -3152,6 +3154,7 @@ class Scheduler( self._prefetch_kvcache(req) self.waiting_queue.append(req) req.time_stats.set_wait_queue_entry_time() + req.arrival_processed_tokens = self.processed_tokens_counter elif self.disaggregation_mode == DisaggregationMode.PREFILL: self._prefetch_kvcache(req) self.disagg_prefill_bootstrap_queue.add( @@ -3745,7 +3748,11 @@ class Scheduler( return None, running_batch # Get priority queue - self.policy.calc_priority(self.waiting_queue, running_batch) + self.policy.calc_priority( + self.waiting_queue, + running_batch, + processed_tokens=self.processed_tokens_counter, + ) if TEST_RETRACT and running_bs > TEST_RETRACT_NO_PREFILL_BS: # If we are testing retraction and the running batch size exceeds @@ -4203,6 +4210,10 @@ class Scheduler( batch.after_idle_gap = self._sched_idled self._sched_idled = False + # Accumulate the prefill-token counter used by the HRRN scheduling policy. Decode / prebuilt batches contribute 0. + if batch.extend_num_tokens: + self.processed_tokens_counter += batch.extend_num_tokens + if self.scripted_scheduler_hook is not None: self.scripted_scheduler_hook.on_run_batch(batch) diff --git a/test/registered/unit/managers/test_generation_auxiliary_output.py b/test/registered/unit/managers/test_generation_auxiliary_output.py index e35154f18..9a4df9f2c 100644 --- a/test/registered/unit/managers/test_generation_auxiliary_output.py +++ b/test/registered/unit/managers/test_generation_auxiliary_output.py @@ -407,6 +407,7 @@ def test_pdmux_split_prefill_schedules_auxiliary_output_copy(): scheduler.scheduler_stage_metrics = None scheduler.metrics_reporter = Mock() scheduler.forward_ct = 0 + scheduler.processed_tokens_counter = 0 scheduler._sched_idled = False scheduler.scripted_scheduler_hook = None scheduler.profiler_manager = SimpleNamespace(_profile_batch_predicate=Mock()) @@ -431,6 +432,7 @@ def test_pdmux_split_prefill_schedules_auxiliary_output_copy(): reqs=[], req_pool_indices=torch.tensor([3]), input_ids=torch.tensor([5]), + extend_num_tokens=1, return_logprob=False, return_hidden_states=False, ) diff --git a/test/registered/unit/managers/test_schedule_policy.py b/test/registered/unit/managers/test_schedule_policy.py new file mode 100644 index 000000000..4fe6b3c41 --- /dev/null +++ b/test/registered/unit/managers/test_schedule_policy.py @@ -0,0 +1,128 @@ +import unittest +from array import array + +from sglang.srt.managers.schedule_batch import Req +from sglang.srt.managers.schedule_policy import SchedulePolicy +from sglang.srt.mem_cache.radix_cache import RadixCache +from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _make_req(rid, origin_input_text, origin_input_ids, sampling_params=None, **kwargs): + if sampling_params is None: + sampling_params = SamplingParams() + return Req( + rid, + origin_input_text, + array("q", origin_input_ids), + sampling_params, + **kwargs, + ) + + +class TestSchedulePolicyHRRN(CustomTestCase): + def test_calc_priority_hrrn(self): + """HRRN sorts by response ratio (waited_tokens / uncached_tokens). + + With three fresh reqs (waited_tokens = 0 for all), the ratio is 0 / uncached for each, + and stable sort keeps original order among equal keys -- effectively pure SUF via the + rid tie-breaker is avoided, so we set a non-zero arrival snapshot on some reqs to exercise + the aging half of the formula. + """ + tree_cache = RadixCache.create_simulated() + + # r_short: small uncached, just arrived (waited=0) + # r_long: large uncached, just arrived (waited=0) + # r_aged: medium uncached, arrived long ago (waited>>0) + r_short = _make_req("short", "a", [1]) + r_long = _make_req("long", "a" * 10, list(range(10))) + r_aged = _make_req("aged", "a" * 3, [1, 2, 3]) + + # Fresh reqs arrived with the counter still at 0. + r_short.arrival_processed_tokens = 0 + r_long.arrival_processed_tokens = 0 + r_aged.arrival_processed_tokens = 0 + + waiting_queue = [r_long, r_aged, r_short] + + policy = SchedulePolicy( + policy="hrrn", + tree_cache=tree_cache, + enable_hierarchical_cache=True, + enable_priority_scheduling=False, + schedule_low_priority_values_first=False, + ) + # processed_tokens=1000 -> waited_tokens is 1000 for every req here. + # Ratios: short = 1000/1 = 1000; aged = 1000/3 ~= 333; long = 1000/10 = 100. + # Highest ratio first -> short, aged, long. + policy.calc_priority(waiting_queue, processed_tokens=1000) + + self.assertEqual(waiting_queue[0].rid, "short") + self.assertEqual(waiting_queue[1].rid, "aged") + self.assertEqual(waiting_queue[2].rid, "long") + + def test_calc_priority_hrrn_aging_overtakes_short(self): + """A long request that has waited enough should overtake a + just-arrived short request.""" + tree_cache = RadixCache.create_simulated() + + r_long_old = _make_req("long_old", "a" * 100, list(range(100))) + r_short_new = _make_req("short_new", "a", [1]) + + # long_old arrived at counter=0 and has been waiting; short_new + # just arrived (its snapshot equals the current counter). + r_long_old.arrival_processed_tokens = 0 + r_short_new.arrival_processed_tokens = 100000 + + waiting_queue = [r_short_new, r_long_old] + + policy = SchedulePolicy( + policy="hrrn", + tree_cache=tree_cache, + enable_hierarchical_cache=True, + enable_priority_scheduling=False, + schedule_low_priority_values_first=False, + ) + # processed_tokens=100000 -> long_old.waited = 100000, ratio = 100000/100 = 1000. + # short_new.waited = 0, ratio = 0. + # long_old should now be first. + policy.calc_priority(waiting_queue, processed_tokens=100000) + + self.assertEqual(waiting_queue[0].rid, "long_old") + self.assertEqual(waiting_queue[1].rid, "short_new") + + def test_calc_priority_hrrn_cached_length_affects_order(self): + """Cached prefix length shortens uncached, so reqs with the same input + length can sort differently by HRRN. Also verifies rid tie-break for + equal-ratio reqs. + + Uses _sort_by_hrrn directly to bypass the prefix-match pass inside + calc_priority, which would overwrite num_matched_prefix_tokens. + """ + r_more_cached = _make_req("a", "x" * 100, list(range(100))) + r_less_cached = _make_req("b", "x" * 100, list(range(100))) + r_tie = _make_req("c", "x" * 100, list(range(100))) + + # r_more_cached: 90 cached -> uncached = 10 -> ratio = 1000 / 10 = 100. + # r_less_cached: 0 cached -> uncached = 100 -> ratio = 1000 / 100 = 10. + # r_tie: 0 cached -> uncached = 100 -> ratio = 1000 / 100 = 10 (ties with r_less_cached; rid "b" < "c"). + r_more_cached.num_matched_prefix_tokens = 90 + r_less_cached.num_matched_prefix_tokens = 0 + r_tie.num_matched_prefix_tokens = 0 + r_more_cached.arrival_processed_tokens = 0 + r_less_cached.arrival_processed_tokens = 0 + r_tie.arrival_processed_tokens = 0 + + waiting_queue = [r_tie, r_less_cached, r_more_cached] + SchedulePolicy._sort_by_hrrn(waiting_queue, set(), processed_tokens=1000) + + self.assertEqual(waiting_queue[0].rid, "a") + self.assertEqual(waiting_queue[1].rid, "b") + self.assertEqual(waiting_queue[2].rid, "c") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/test_server_args_cli_metadata.py b/test/registered/unit/test_server_args_cli_metadata.py index 8e34487d9..8708c245a 100644 --- a/test/registered/unit/test_server_args_cli_metadata.py +++ b/test/registered/unit/test_server_args_cli_metadata.py @@ -45,7 +45,16 @@ class TestServerArgsMigratedCliMetadata(CustomTestCase): ) self.assertEqual( self.actions_by_option["--schedule-policy"].choices, - ["lpm", "random", "fcfs", "dfs-weight", "lof", "priority", "routing-key"], + [ + "lpm", + "random", + "fcfs", + "dfs-weight", + "lof", + "priority", + "routing-key", + "hrrn", + ], ) self.assertEqual( self.actions_by_option["--load-balance-method"].choices,