minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Yuwei An
2026-09-17 21:18:58 -07:00
committed by GitHub
parent c055dc6ff6
commit 65ef55e2a8
6 changed files with 281 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/fields/schedule.py
+1
View File
@@ -99,6 +99,7 @@ class Schedule(msgspec.Struct):
"priority",
"routing-key",
"hrrn",
"shortest-prefill-first",
],
),
] = "fcfs"
python/sglang/srt/managers/schedule_policy.py
+75
View File
@@ -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
+6
View File
@@ -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
+94
View File
@@ -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
+104 -1
View File
@@ -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
+1
View File
@@ -61,6 +61,7 @@ class TestServerArgsMigratedCliMetadata(CustomTestCase):
"priority",
"routing-key",
"hrrn",
"shortest-prefill-first",
],
)
self.assertEqual(