[Scheduler] Add HRRN schedule policy to significantly reduce TTFT (#32911)

Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

python/sglang/srt/arg_groups/fields/schedule.py

@@ -95,6 +95,7 @@ class Schedule:
                 "lof",
                 "priority",
                 "routing-key",
+                "hrrn",
             ],
         ),
     ] = "fcfs"

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)

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."""

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

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)
 

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,
     )

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()

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,