From 65ef55e2a8c51be0523723fcb98d7138b5839ebb Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Thu, 17 Sep 2026 21:18:58 -0700 Subject: [PATCH] [Scheduler] Add shortest-prefill-first scheduling (#40024) --- .../sglang/srt/arg_groups/fields/schedule.py | 1 + python/sglang/srt/managers/schedule_policy.py | 75 +++++++++++++ python/sglang/srt/managers/scheduler.py | 6 + .../unit/managers/test_prefill_adder.py | 94 ++++++++++++++++ .../unit/managers/test_schedule_policy.py | 105 +++++++++++++++++- .../unit/test_server_args_cli_metadata.py | 1 + 6 files changed, 281 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/arg_groups/fields/schedule.py b/python/sglang/srt/arg_groups/fields/schedule.py index de4267d95..2ed71b8a3 100644 --- a/python/sglang/srt/arg_groups/fields/schedule.py +++ b/python/sglang/srt/arg_groups/fields/schedule.py @@ -99,6 +99,7 @@ class Schedule(msgspec.Struct): "priority", "routing-key", "hrrn", + "shortest-prefill-first", ], ), ] = "fcfs" diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 2ed58e55f..f4082c592 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/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` diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 7f8726896..07a1c97cd 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 1fc70b8f8..03b21a95a 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/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 = ( diff --git a/test/registered/unit/managers/test_schedule_policy.py b/test/registered/unit/managers/test_schedule_policy.py index 4fe6b3c41..45ce93734 100644 --- a/test/registered/unit/managers/test_schedule_policy.py +++ b/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() diff --git a/test/registered/unit/test_server_args_cli_metadata.py b/test/registered/unit/test_server_args_cli_metadata.py index b0fb4d952..4ebf060d6 100644 --- a/test/registered/unit/test_server_args_cli_metadata.py +++ b/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(