[Scheduler] Add shortest-prefill-first scheduling (#40024)

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

@@ -99,6 +99,7 @@ class Schedule(msgspec.Struct):
                 "priority",
                 "routing-key",
                 "hrrn",
+                "shortest-prefill-first",
             ],
         ),
     ] = "fcfs"

python/sglang/srt/managers/schedule_policy.py

@@ -221,6 +221,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
+    SHORTEST_PREFILL_FIRST = "shortest-prefill-first"
 
 
 class CacheAgnosticPolicy(Enum):
@@ -249,6 +250,7 @@ class SchedulePolicy:
         self.enable_priority_scheduling = enable_priority_scheduling
         self.schedule_low_priority_values_first = schedule_low_priority_values_first
         self.priority_sign = 1 if schedule_low_priority_values_first else -1
+        self._shortest_prefill_calls = 0
 
         # It is used to find the matching prefix for in-batch prefix caching.
         self.waiting_queue_radix_tree = RadixCache.create_simulated()
@@ -294,6 +296,17 @@ class SchedulePolicy:
                 SchedulePolicy._sort_by_hrrn(
                     waiting_queue, temporary_deprioritized, processed_tokens
                 )
+            elif policy == CacheAwarePolicy.SHORTEST_PREFILL_FIRST:
+                SchedulePolicy._sort_by_shortest_prefill(
+                    waiting_queue, temporary_deprioritized
+                )
+                self._shortest_prefill_calls += 1
+                if waiting_queue and self._shortest_prefill_calls % 128 == 1:
+                    logger.info(
+                        "Experimental shortest-prefill-first: queue=%d shortest_uncached=%d",
+                        len(waiting_queue),
+                        self._shortest_prefill_work(waiting_queue[0]),
+                    )
             else:
                 raise ValueError(f"Unknown CacheAware Policy: {policy=}")
         else:
@@ -342,6 +355,10 @@ class SchedulePolicy:
         try:
             policy_enum = CacheAwarePolicy(policy)
             if getattr(tree_cache, "disable", True):
+                if policy_enum == CacheAwarePolicy.SHORTEST_PREFILL_FIRST:
+                    raise ValueError(
+                        "Experimental shortest-prefill-first requires prefix caching"
+                    )
                 # If tree_cache is disabled, using CacheAgnosticPolicy policy
                 return CacheAgnosticPolicy.FCFS
             return policy_enum
@@ -410,6 +427,50 @@ class SchedulePolicy:
                     )
         return temporary_deprioritized
 
+    @staticmethod
+    def _shortest_prefill_work(r: Req) -> int:
+        return max(
+            1,
+            len(r.origin_input_ids) + len(r.output_ids) - r.num_matched_prefix_tokens,
+        )
+
+    @staticmethod
+    def _sort_by_shortest_prefill(
+        waiting_queue: List[Req], temporary_deprioritized: Set[int]
+    ) -> None:
+        # Prioritize short uncached prefills while deferring duplicate prefixes.
+        waiting_queue.sort(
+            key=lambda r: (
+                r.rid in temporary_deprioritized,
+                SchedulePolicy._shortest_prefill_work(r),
+                r.time_stats.wait_queue_entry_time,
+            )
+        )
+
+    def shortest_prefill_chunk_limit(
+        self, chunked_req: Req, waiting_queue: List[Req], budget: int, page_size: int
+    ) -> Optional[int]:
+        """Cap the active prefill chunk to reserve tokens for shorter waiting requests."""
+        if (
+            self.policy != CacheAwarePolicy.SHORTEST_PREFILL_FIRST
+            or budget < 2 * page_size
+        ):
+            return None
+        remaining = len(chunked_req.full_untruncated_fill_ids) - len(
+            chunked_req.prefix_indices
+        )
+        reserved = 0
+        for req in waiting_queue:
+            work = self._shortest_prefill_work(req)
+            charge = _ceil_div(work, page_size) * page_size
+            if work >= remaining or reserved + charge > budget - page_size:
+                break
+            reserved += charge
+        if not reserved:
+            return None
+        # Page alignment keeps continuation boundaries allocator-compatible.
+        return (budget - reserved) // page_size * page_size
+
     @staticmethod
     def _sort_by_longest_prefix(
         waiting_queue: List[Req], temporary_deprioritized: Set[int]
@@ -583,6 +644,7 @@ class PrefillAdder:
         self.new_token_ratio = new_token_ratio
         self.rem_input_tokens = rem_input_tokens - num_mixed_decode_tokens
         self.rem_chunk_tokens = rem_chunk_tokens
+        self.chunked_req_limit: Optional[int] = None
         self.dllm_config = dllm_config
         self.exact_chunk_fill = _use_exact_chunk_fill() and dllm_config is None
 
@@ -968,6 +1030,10 @@ class PrefillAdder:
                 waiting_queue_len=self.waiting_queue_len,
             )
 
+        if self.chunked_req_limit is not None:
+            assert self.chunked_req_limit > 0
+            _rem_tokens = min(_rem_tokens, self.chunked_req_limit)
+
         cand_extend_input_len = len(req.full_untruncated_fill_ids) - len(
             req.prefix_indices
         )
@@ -1189,6 +1255,7 @@ class PrefillAdder:
                 host_hit_length=req.host_hit_length,
                 swa_host_hit_length=req.swa_host_hit_length,
                 truncation_align_size=truncation_align_size,
+                has_chunked_req=has_chunked_req,
             )
             if isinstance(admission, AddReqResult):
                 return admission
@@ -1252,6 +1319,7 @@ class PrefillAdder:
                         host_hit_length=0,
                         swa_host_hit_length=0,
                         truncation_align_size=truncation_align_size,
+                        has_chunked_req=has_chunked_req,
                     )
                     if isinstance(admission, AddReqResult):
                         return admission
@@ -1272,6 +1340,7 @@ class PrefillAdder:
         host_hit_length: int,
         swa_host_hit_length: int,
         truncation_align_size: Optional[int],
+        has_chunked_req: bool = False,
     ) -> _PrefillAdmission | AddReqResult:
         """Select a prefill shape without allocating or publishing cached KV."""
         prefix_len = len(req.prefix_indices) + host_hit_length
@@ -1314,6 +1383,12 @@ class PrefillAdder:
                 return AddReqResult.OTHER
             max_new_tokens = 0
         elif chunk_tokens_limit is not None and chunk_fit_tokens > chunk_tokens_limit:
+            if (
+                has_chunked_req
+                and get_schedule().schedule_policy == "shortest-prefill-first"
+            ):
+                # Only one unfinished chunked request can be tracked.
+                return AddReqResult.OTHER
             if self.exact_chunk_fill:
                 # Take the remainder verbatim so the batch hits exactly
                 # chunked_prefill_size. `chunk_fit_tokens > chunk_tokens_limit`

python/sglang/srt/managers/scheduler.py

@@ -3921,6 +3921,12 @@ class Scheduler(
 
         if self.chunked_req is not None:
             self.chunked_req.init_next_round_input()
+            adder.chunked_req_limit = self.policy.shortest_prefill_chunk_limit(
+                self.chunked_req,
+                self.waiting_queue,
+                adder.rem_chunk_tokens or 0,
+                self.page_size,
+            )
             self.chunked_req = adder.add_chunked_req(self.chunked_req)
 
         if self.enable_lora:

test/registered/unit/managers/test_prefill_adder.py

@@ -9,6 +9,7 @@ from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.managers.schedule_policy import (
     AddReqResult,
     PrefillAdder,
+    SchedulePolicy,
     estimate_prefill_extend_tile_metrics,
 )
 from sglang.srt.mem_cache.base_prefix_cache import (
@@ -21,6 +22,7 @@ from sglang.srt.mem_cache.prefill_budget import (
     SWAPrefillBudget,
     estimate_swa_kv_tokens,
 )
+from sglang.srt.mem_cache.radix_cache import RadixCache
 from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools
 from sglang.srt.runtime_context import get_context
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
@@ -186,6 +188,98 @@ class TestPrefillAdder(CustomTestCase):
         )
         return req
 
+    def create_shortest_prefill_adder(self, *, chunk_tokens=4096):
+        override = get_context().override_server_args(
+            schedule_policy="shortest-prefill-first"
+        )
+        override.install()
+        self.addCleanup(override.restore)
+        self.mock_tree_cache.supports_mamba.return_value = False
+        self.mock_tree_cache.is_tree_cache.return_value = False
+        self.mock_token_allocator.available_size.return_value = 32768
+        return self.create_adder(
+            self.create_running_batch(), page_size=256, rem_chunk_tokens=chunk_tokens
+        )
+
+    def test_shortest_prefill_reserves_space_for_complete_waiting_requests(self):
+        adder = self.create_shortest_prefill_adder()
+        policy = SchedulePolicy(
+            policy="shortest-prefill-first",
+            tree_cache=RadixCache.create_simulated(),
+            enable_hierarchical_cache=True,
+            enable_priority_scheduling=False,
+            schedule_low_priority_values_first=False,
+        )
+        continuation = self.create_shared_req("continuation")
+        continuation.full_untruncated_fill_ids = list(range(16384))
+        waiting = [self.create_shared_req("a"), self.create_shared_req("b")]
+        for req, length in zip(waiting, [512, 1024]):
+            req.origin_input_ids = list(range(length))
+            req.full_untruncated_fill_ids = list(range(length))
+            req.num_matched_prefix_tokens = 0
+        adder.chunked_req_limit = policy.shortest_prefill_chunk_limit(
+            continuation, waiting, adder.rem_chunk_tokens, adder.page_size
+        )
+        self.assertIs(adder.add_chunked_req(continuation), continuation)
+        self.assertEqual(continuation.extend_range.length, 2560)
+        for req in waiting:
+            adder.add_one_req(req, has_chunked_req=True, truncation_align_size=None)
+        self.assertEqual(adder.can_run_list, [continuation, *waiting])
+        self.assertIsNone(adder.new_chunked_req)
+        self.assertEqual(adder.rem_chunk_tokens, 0)
+        self.assertGreaterEqual(adder.rem_total_tokens, 0)
+
+    def test_shortest_prefill_rejects_second_unfinished_chunk(self):
+        adder = self.create_shortest_prefill_adder(chunk_tokens=512)
+        req = self.create_shared_req("second-chunk")
+        req.full_untruncated_fill_ids = list(range(1024))
+        self.assertEqual(
+            adder.add_one_req(req, has_chunked_req=True, truncation_align_size=None),
+            AddReqResult.OTHER,
+        )
+        self.assertEqual(adder.can_run_list, [])
+        self.assertIsNone(adder.new_chunked_req)
+        req.set_extend_range.assert_not_called()
+        self.mock_tree_cache.init_load_back.assert_not_called()
+
+    def test_shortest_prefill_rechecks_chunk_limit_after_host_miss(self):
+        adder = self.create_shortest_prefill_adder(chunk_tokens=512)
+        req = self.create_shared_req("host-miss")
+        req.full_untruncated_fill_ids = list(range(1024))
+        req.prefix_indices = torch.empty(0, dtype=torch.int64)
+        req.host_hit_length = 768
+        req.best_match_node = req.last_node
+        req.needs_host_load_back.return_value = True
+        self.mock_tree_cache.init_load_back.return_value = (
+            torch.empty(0, dtype=torch.int64),
+            req.last_node,
+        )
+        self.assertEqual(
+            adder.add_one_req(req, has_chunked_req=True, truncation_align_size=None),
+            AddReqResult.OTHER,
+        )
+        self.mock_tree_cache.init_load_back.assert_called_once()
+        self.assertEqual(adder.can_run_list, [])
+        req.set_extend_range.assert_not_called()
+
+    def test_shortest_prefill_preserves_memory_admission(self):
+        adder = self.create_shortest_prefill_adder()
+        self.mock_token_allocator.available_size.return_value = 256
+        req = self.create_shared_req("no-memory")
+        req.full_untruncated_fill_ids = list(range(512))
+        self.assertEqual(
+            adder.add_one_req(req, has_chunked_req=True, truncation_align_size=None),
+            AddReqResult.NO_TOKEN,
+        )
+        self.assertEqual(adder.can_run_list, [])
+
+    def test_continuation_without_limit_keeps_normal_chunk_size(self):
+        adder = self.create_shortest_prefill_adder()
+        req = self.create_shared_req("continuation")
+        req.full_untruncated_fill_ids = list(range(8192))
+        self.assertIs(adder.add_chunked_req(req), req)
+        self.assertEqual(req.extend_range.length, 4096)
+
     def test_shared_admission_reserves_all_pending_requests(self):
         adder = self.create_shared_adder()
         first, second = (

test/registered/unit/managers/test_schedule_policy.py

@@ -1,8 +1,9 @@
 import unittest
 from array import array
+from unittest.mock import patch
 
 from sglang.srt.managers.schedule_batch import Req
-from sglang.srt.managers.schedule_policy import SchedulePolicy
+from sglang.srt.managers.schedule_policy import CacheAwarePolicy, 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
@@ -124,5 +125,107 @@ class TestSchedulePolicyHRRN(CustomTestCase):
         self.assertEqual(waiting_queue[2].rid, "c")
 
 
+class TestShortestPrefillFirst(CustomTestCase):
+    def setUp(self):
+        self.policy = SchedulePolicy(
+            policy="shortest-prefill-first",
+            tree_cache=RadixCache.create_simulated(),
+            enable_hierarchical_cache=True,
+            enable_priority_scheduling=False,
+            schedule_low_priority_values_first=False,
+        )
+
+    def make_req(self, rid, uncached, *, cached=0, arrived=0):
+        req = _make_req(rid, "", list(range(uncached + cached)))
+        req.full_untruncated_fill_ids = req.origin_input_ids[:]
+        req.num_matched_prefix_tokens = cached
+        req.prefix_indices = list(range(cached))
+        req.time_stats.wait_queue_entry_time = arrived
+        return req
+
+    def test_calc_priority_uses_uncached_work(self):
+        cached = self.make_req("cached", 16, cached=4096)
+        short = self.make_req("short", 32)
+        long = self.make_req("long", 1024)
+        queue = [long, short, cached]
+        with patch.object(self.policy, "_compute_prefix_matches", return_value=set()):
+            self.policy.calc_priority(queue)
+        self.assertEqual([req.rid for req in queue], ["cached", "short", "long"])
+
+    def test_equal_work_uses_arrival_time(self):
+        older = self.make_req("z", 32, arrived=1)
+        newer = self.make_req("a", 32, arrived=2)
+        queue = [newer, older]
+        self.policy._sort_by_shortest_prefill(queue, set())
+        self.assertEqual(queue, [older, newer])
+
+    def test_duplicate_prefix_is_deprioritized(self):
+        duplicate = self.make_req("duplicate", 1)
+        other = self.make_req("other", 1024)
+        queue = [duplicate, other]
+        with patch.object(
+            self.policy, "_compute_prefix_matches", return_value={duplicate.rid}
+        ):
+            self.policy.calc_priority(queue)
+        self.assertEqual(queue, [other, duplicate])
+
+    def test_retracted_output_is_part_of_uncached_work(self):
+        replay = self.make_req("replay", 16, cached=1024)
+        replay.output_ids.extend([0] * 64)
+        short = self.make_req("short", 32)
+        queue = [replay, short]
+        self.policy._sort_by_shortest_prefill(queue, set())
+        self.assertEqual(queue, [short, replay])
+
+    def test_chunk_limit_reserves_complete_short_prefills(self):
+        continuation = self.make_req("continuation", 16384)
+        waiting = [self.make_req("a", 512), self.make_req("b", 1024)]
+        self.assertEqual(
+            self.policy.shortest_prefill_chunk_limit(continuation, waiting, 4096, 256),
+            2560,
+        )
+
+    def test_reservation_rounds_to_pages_and_keeps_continuation_progress(self):
+        continuation = self.make_req("continuation", 16384)
+        self.assertEqual(
+            self.policy.shortest_prefill_chunk_limit(
+                continuation, [self.make_req("short", 257)], 4096, 256
+            ),
+            3584,
+        )
+        self.assertEqual(
+            self.policy.shortest_prefill_chunk_limit(
+                continuation, [self.make_req("short", 3840)], 4096, 256
+            ),
+            256,
+        )
+
+    def test_no_reservation_when_request_cannot_fit_or_is_not_shorter(self):
+        continuation = self.make_req("continuation", 8192)
+        for waiting, budget in [
+            ([], 4096),
+            ([self.make_req("same", 8192)], 4096),
+            ([self.make_req("too-large", 4096)], 4096),
+            ([self.make_req("short", 1)], 256),
+        ]:
+            with self.subTest(budget=budget, waiting=[req.rid for req in waiting]):
+                self.assertIsNone(
+                    self.policy.shortest_prefill_chunk_limit(
+                        continuation, waiting, budget, 256
+                    )
+                )
+
+    def test_other_policy_keeps_normal_chunk_limit(self):
+        self.policy.policy = CacheAwarePolicy.HRRN
+        self.assertIsNone(
+            self.policy.shortest_prefill_chunk_limit(
+                self.make_req("continuation", 8192),
+                [self.make_req("short", 512)],
+                4096,
+                256,
+            )
+        )
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/test_server_args_cli_metadata.py

@@ -61,6 +61,7 @@ class TestServerArgsMigratedCliMetadata(CustomTestCase):
                 "priority",
                 "routing-key",
                 "hrrn",
+                "shortest-prefill-first",
             ],
         )
         self.assertEqual(